From f7e375df211f4e51c0ef5822e2c79ac0eeb95386 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Tue, 14 Jul 2026 21:52:11 +0200 Subject: [PATCH 01/20] Add budgeted host memory safety primitives --- tests/test_dxgi_meminfo.py | 186 ++++ tests/test_nvml_free_signal.py | 105 ++ tests/test_pin_manager.py | 134 +++ toolkit/memory_management/allocator_cap.py | 122 +++ toolkit/memory_management/dxgi_meminfo.py | 493 +++++++++ toolkit/memory_management/manager.py | 40 +- toolkit/memory_management/manager_modules.py | 72 +- toolkit/memory_management/nvml_meminfo.py | 234 +++++ toolkit/memory_management/pin_manager.py | 752 ++++++++++++++ toolkit/memory_management/vram_budget.py | 988 +++++++++++++++++++ 10 files changed, 3104 insertions(+), 22 deletions(-) create mode 100644 tests/test_dxgi_meminfo.py create mode 100644 tests/test_nvml_free_signal.py create mode 100644 tests/test_pin_manager.py create mode 100644 toolkit/memory_management/allocator_cap.py create mode 100644 toolkit/memory_management/dxgi_meminfo.py create mode 100644 toolkit/memory_management/nvml_meminfo.py create mode 100644 toolkit/memory_management/pin_manager.py create mode 100644 toolkit/memory_management/vram_budget.py diff --git a/tests/test_dxgi_meminfo.py b/tests/test_dxgi_meminfo.py new file mode 100644 index 0000000000..a9033f2cbc --- /dev/null +++ b/tests/test_dxgi_meminfo.py @@ -0,0 +1,186 @@ +import os +import unittest +from unittest import mock + +from toolkit.memory_management import dxgi_meminfo, pin_manager + +GIB = 1024 ** 3 + + +class _FakeVM: + def __init__(self, total): + self.total = total + + +class _FakePsutil: + def __init__(self, total): + self._vm = _FakeVM(total) + + def virtual_memory(self): + return self._vm + + +class DxgiHeadroomTests(unittest.TestCase): + def setUp(self): + self._saved_pinned = pin_manager.pinned_bytes_by_kind() + pin_manager._LEDGER.clear() + + def _restore(): + pin_manager._LEDGER.clear() + pin_manager._LEDGER.update(self._saved_pinned) + + self.addCleanup(_restore) + + def test_compute_non_local_headroom_normal_case(self): + self.assertEqual( + dxgi_meminfo.compute_non_local_headroom_bytes(16 * GIB, 10 * GIB, 2 * GIB), + 4 * GIB, + ) + + def test_compute_non_local_headroom_clamps_to_zero(self): + self.assertEqual( + dxgi_meminfo.compute_non_local_headroom_bytes(16 * GIB, 15 * GIB, 2 * GIB), + 0, + ) + + def test_compute_non_local_headroom_allows_zero_reserve(self): + self.assertEqual( + dxgi_meminfo.compute_non_local_headroom_bytes(16 * GIB, 10 * GIB, 0), + 6 * GIB, + ) + + def test_pinned_bytes_headroom_falls_back_to_legacy_proxy(self): + with mock.patch.object( + dxgi_meminfo, "query_non_local_video_memory_info", return_value=None + ): + with mock.patch.object(pin_manager, "_psutil", _FakePsutil(32 * GIB)): + with mock.patch.dict( + os.environ, + { + "AI_TOOLKIT_PINNED_WEIGHT_WDDM_FRACTION": "0.25", + "AI_TOOLKIT_WDDM_DXGI_DISABLE": "0", + "AI_TOOLKIT_WDDM_DXGI_CONTROL_DISABLE": "0", + }, + clear=False, + ): + pin_manager.register_pinned_bytes(3 * GIB) + self.assertEqual(pin_manager.pinned_bytes_headroom(0), 5 * GIB) + + def test_dxgi_headroom_does_not_subtract_pinned_ledger(self): + reading = dxgi_meminfo.DxgiMemoryInfo( + budget_bytes=16 * GIB, + current_usage_bytes=10 * GIB, + available_for_reservation_bytes=0, + current_reservation_bytes=0, + ) + with mock.patch.object( + dxgi_meminfo, "query_non_local_video_memory_info", return_value=reading + ): + with mock.patch.dict( + os.environ, + { + "AI_TOOLKIT_WDDM_SPILL_RESERVE_GIB": "1.0", + "AI_TOOLKIT_WDDM_SPILL_RESERVE_PCT": "0", + "AI_TOOLKIT_WDDM_DXGI_DISABLE": "0", + "AI_TOOLKIT_WDDM_DXGI_CONTROL_DISABLE": "0", + }, + clear=False, + ): + pin_manager.register_pinned_bytes(4 * GIB) + # pct pinned to 0 -> reserve is exactly the 1 GiB floor, so the + # focus is the double-subtract invariant (ledger not subtracted). + self.assertEqual(pin_manager.pinned_bytes_headroom(0), 5 * GIB) + + def test_control_disable_keeps_legacy_proxy_for_control(self): + reading = dxgi_meminfo.DxgiMemoryInfo( + budget_bytes=64 * GIB, + current_usage_bytes=1 * GIB, + available_for_reservation_bytes=0, + current_reservation_bytes=0, + ) + with mock.patch.object( + dxgi_meminfo, "query_non_local_video_memory_info", return_value=reading + ): + with mock.patch.object(pin_manager, "_psutil", _FakePsutil(32 * GIB)): + with mock.patch.dict( + os.environ, + { + "AI_TOOLKIT_PINNED_WEIGHT_WDDM_FRACTION": "0.25", + "AI_TOOLKIT_WDDM_DXGI_CONTROL_DISABLE": "1", + }, + clear=False, + ): + self.assertEqual(pin_manager.pinned_bytes_headroom(0), 8 * GIB) + + +class SpillReserveMarginTests(unittest.TestCase): + def setUp(self): + def _clear(): + pin_manager._SPILL_RESERVE_FLOOR_GIB_OVERRIDE = None + pin_manager._SPILL_RESERVE_PCT_OVERRIDE = None + _clear() + self.addCleanup(_clear) + + def test_margin_is_pct_of_budget_when_pct_dominates(self): + with mock.patch.dict( + os.environ, + {"AI_TOOLKIT_WDDM_SPILL_RESERVE_FLOOR_GIB": "2.0", + "AI_TOOLKIT_WDDM_SPILL_RESERVE_PCT": "0.20"}, + clear=False, + ): + # 0.20 * 16 GiB = 3.2 GiB > 2 GiB floor. + self.assertEqual( + pin_manager.dxgi_spill_reserve_bytes(16 * GIB), + int(0.20 * 16 * GIB), + ) + + def test_margin_is_floor_when_floor_dominates(self): + with mock.patch.dict( + os.environ, + {"AI_TOOLKIT_WDDM_SPILL_RESERVE_FLOOR_GIB": "2.0", + "AI_TOOLKIT_WDDM_SPILL_RESERVE_PCT": "0.20"}, + clear=False, + ): + # 0.20 * 8 GiB = 1.6 GiB < 2 GiB floor. + self.assertEqual(pin_manager.dxgi_spill_reserve_bytes(8 * GIB), 2 * GIB) + + def test_margin_falls_back_to_floor_without_budget(self): + with mock.patch.dict( + os.environ, + {"AI_TOOLKIT_WDDM_SPILL_RESERVE_FLOOR_GIB": "2.0", + "AI_TOOLKIT_WDDM_SPILL_RESERVE_PCT": "0.20"}, + clear=False, + ): + self.assertEqual(pin_manager.dxgi_spill_reserve_bytes(None), 2 * GIB) + self.assertEqual(pin_manager.dxgi_spill_reserve_bytes(0), 2 * GIB) + + def test_config_override_supersedes_env(self): + with mock.patch.dict( + os.environ, + {"AI_TOOLKIT_WDDM_SPILL_RESERVE_FLOOR_GIB": "2.0", + "AI_TOOLKIT_WDDM_SPILL_RESERVE_PCT": "0.20"}, + clear=False, + ): + pin_manager.set_spill_reserve_policy(floor_gib=3.0, pct=0.10) + # config floor 3 GiB dominates 0.10 * 16 = 1.6 GiB. + self.assertEqual(pin_manager.dxgi_spill_reserve_bytes(16 * GIB), 3 * GIB) + + +class SafeForControlTests(unittest.TestCase): + def test_confident_auto_detect_is_safe(self): + self.assertTrue(dxgi_meminfo.safe_for_control("single_nvidia")) + self.assertTrue(dxgi_meminfo.safe_for_control("luid")) + + def test_manual_override_is_safe_but_flagged_manual(self): + self.assertTrue(dxgi_meminfo.safe_for_control("env_override")) + self.assertTrue(dxgi_meminfo.is_manual_control("env_override")) + self.assertFalse(dxgi_meminfo.is_manual_control("single_nvidia")) + + def test_fallback_methods_are_not_safe(self): + for method in ("sole_hardware_adapter", "global_conservative", "", None): + self.assertFalse(dxgi_meminfo.safe_for_control(method)) + self.assertFalse(dxgi_meminfo.is_manual_control(method)) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_nvml_free_signal.py b/tests/test_nvml_free_signal.py new file mode 100644 index 0000000000..46686c5a43 --- /dev/null +++ b/tests/test_nvml_free_signal.py @@ -0,0 +1,105 @@ +"""The governing device-free signal must reflect the whole card, not this process. + +Background: torch.cuda.mem_get_info reports a per-process promise, not physical +availability. Measured on a 4070 with a second process holding 9 GiB, it claimed +4.70 GiB free while the card physically had 0.55 GiB. Planning residency against +that number silently crosses the WDDM dedicated cliff (no error, no allocator +retry -- just ~4x slower steps). These tests pin the reconciliation policy. +""" + +import unittest + +import torch + +from toolkit.memory_management import nvml_meminfo, vram_budget + +GIB = 1024 ** 3 + + +class ReconcileFreeBytesTests(unittest.TestCase): + """Pure policy: which free value governs.""" + + def test_prefers_physical_when_driver_over_reports(self): + # The production failure: driver promises 4.70 GiB, card has 0.55 GiB. + governing = vram_budget.reconcile_free_bytes( + int(4.70 * GIB), int(0.55 * GIB) + ) + self.assertAlmostEqual(governing / GIB, 0.55, places=2) + + def test_falls_back_to_driver_when_nvml_unavailable(self): + # No NVML (non-NVIDIA, no driver): must not crash, must not invent a + # number -- prior behavior is preserved exactly. + governing = vram_budget.reconcile_free_bytes(int(4.70 * GIB), None) + self.assertAlmostEqual(governing / GIB, 4.70, places=2) + + def test_takes_the_pessimistic_signal_when_driver_is_lower(self): + # The driver's promise is a real constraint too; never let NVML's larger + # number talk us into exceeding it. + governing = vram_budget.reconcile_free_bytes(int(1.0 * GIB), int(6.0 * GIB)) + self.assertAlmostEqual(governing / GIB, 1.0, places=2) + + def test_agrees_with_driver_when_uncontended(self): + # Single-process case: the two signals match, so this policy costs no + # residency versus the old behavior. + governing = vram_budget.reconcile_free_bytes(int(10.6 * GIB), int(10.6 * GIB)) + self.assertAlmostEqual(governing / GIB, 10.6, places=2) + + def test_never_negative(self): + self.assertEqual(vram_budget.reconcile_free_bytes(-5, None), 0) + self.assertEqual(vram_budget.reconcile_free_bytes(GIB, -5), 0) + + +@unittest.skipUnless(torch.cuda.is_available(), "needs CUDA") +class NvmlSensorTests(unittest.TestCase): + """The sensor itself, against the real driver.""" + + def test_reports_plausible_physical_totals(self): + info = nvml_meminfo.query_device_memory_info(0) + if info is None: + self.skipTest("NVML unavailable on this box") + self.assertGreater(info.total_bytes, 0) + self.assertEqual(info.used_bytes + info.free_bytes, info.total_bytes) + card_total = int(torch.cuda.get_device_properties(0).total_memory) + # Same card: NVML total is within a few percent of torch's (they differ + # slightly on what they count as reserved by the driver). + self.assertLess(abs(info.total_bytes - card_total) / card_total, 0.05) + + def test_physical_free_never_exceeds_the_card(self): + free = nvml_meminfo.physical_free_bytes(0) + if free is None: + self.skipTest("NVML unavailable on this box") + card_total = int(torch.cuda.get_device_properties(0).total_memory) + self.assertGreaterEqual(free, 0) + self.assertLessEqual(free, card_total) + + def test_governing_free_is_never_more_optimistic_than_the_driver(self): + # The invariant that keeps us off the cliff: whatever the sensors say, + # the governing number can only be <= what torch would have used. + device = torch.device("cuda", 0) + driver_free = int(torch.cuda.mem_get_info(device)[0]) + governing = vram_budget.device_free_bytes(device) + self.assertLessEqual(governing, driver_free) + self.assertGreaterEqual(governing, 0) + + def test_device_mem_info_is_a_drop_in_for_mem_get_info(self): + device = torch.device("cuda", 0) + free, total = vram_budget.device_mem_info(device) + _driver_free, driver_total = torch.cuda.mem_get_info(device) + self.assertEqual(total, int(driver_total)) + self.assertLessEqual(free, total) + + +@unittest.skipUnless(torch.cuda.is_available(), "needs CUDA") +class DeviceSnapshotTests(unittest.TestCase): + def test_snapshot_free_is_the_governed_value(self): + device = torch.device("cuda", 0) + snap = vram_budget.DeviceSnapshot.capture(device) + self.assertIsNotNone(snap) + self.assertEqual(snap.free, vram_budget.device_free_bytes(device)) + # non_torch is what the docstring claims: everything on the card that + # isn't torch's allocator -- which now genuinely includes other processes. + self.assertGreaterEqual(snap.non_torch, 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pin_manager.py b/tests/test_pin_manager.py new file mode 100644 index 0000000000..a0767c149d --- /dev/null +++ b/tests/test_pin_manager.py @@ -0,0 +1,134 @@ +import re +import unittest +from pathlib import Path +from unittest import mock + +from toolkit.memory_management import pin_manager + + +GIB = 1024 ** 3 +REPO_ROOT = Path(__file__).resolve().parents[1] + + +class PinManagerTests(unittest.TestCase): + def setUp(self): + pin_manager.reset_for_tests() + pin_manager.set_host_cache_reserve_bytes(None) + + def tearDown(self): + pin_manager.reset_for_tests() + pin_manager.set_host_cache_reserve_bytes(None) + + def test_ledger_tracks_kinds(self): + pin_manager.register_pinned_bytes(100, "weights") + pin_manager.register_pinned_bytes(50, "bounce") + pin_manager.release_pinned_bytes(25, "weights") + + self.assertEqual(pin_manager.total_pinned_bytes(), 125) + self.assertEqual( + pin_manager.pinned_bytes_by_kind(), + {"weights": 75, "bounce": 50}, + ) + + def test_reserve_reduces_available_grant(self): + pin_manager.set_host_cache_reserve_bytes(2 * GIB) + with mock.patch.object(pin_manager, "pinned_bytes_headroom", return_value=8 * GIB): + self.assertEqual(pin_manager.available_for_pin(mode="training"), 6 * GIB) + self.assertTrue(pin_manager.can_pin(6 * GIB, mode="training")) + self.assertFalse(pin_manager.can_pin(7 * GIB, mode="training")) + + def test_plan_full_pin_disables_bounce(self): + pin_manager.set_host_cache_reserve_bytes(1 * GIB) + with mock.patch.object(pin_manager, "pinned_bytes_headroom", return_value=10 * GIB): + plan = pin_manager.plan_budgets( + offloaded_weight_bytes=8 * GIB, + requested_bounce_bytes=4 * GIB, + mode="training", + ) + self.assertEqual(plan["strategy"], "full_pin_no_bounce") + self.assertEqual(plan["weight_budget_bytes"], 8 * GIB) + self.assertEqual(plan["bounce_budget_bytes"], 0) + + def test_plan_partial_prioritizes_bounce_then_weights(self): + pin_manager.set_host_cache_reserve_bytes(1 * GIB) + with mock.patch.object(pin_manager, "pinned_bytes_headroom", return_value=10 * GIB): + plan = pin_manager.plan_budgets( + offloaded_weight_bytes=12 * GIB, + requested_bounce_bytes=4 * GIB, + mode="training", + ) + self.assertEqual(plan["strategy"], "partial_bounce_first") + self.assertEqual(plan["bounce_budget_bytes"], 4 * GIB) + self.assertEqual(plan["weight_budget_bytes"], 5 * GIB) + + + def test_release_clamps_within_kind_only(self): + # An unmatched release (consumer releasing bytes it never registered, + # e.g. pinning was disabled) must not drain other consumers' entries. + pin_manager.register_pinned_bytes(100, "weights") + pin_manager.release_pinned_bytes(50, "bounce") + self.assertEqual(pin_manager.pinned_bytes_by_kind(), {"weights": 100}) + + def test_weight_tier_reconcile_cannot_shrink_evictables(self): + # Priority guard: weights are the lowest pin tier, so a weights-tier + # shortfall may empty the host cache but must never evict the bounce + # pool to make room for itself. + calls = [] + pin_manager.register_evictable(lambda need: calls.append(need) or 0) + pin_manager.reconcile(1 * GIB, allow_shrink=False) + self.assertEqual(calls, []) + pin_manager.reconcile(1 * GIB, allow_shrink=True) + self.assertEqual(calls, [1 * GIB]) + + def test_release_handle_is_idempotent(self): + pin_manager.register_pinned_bytes(64, "save_stager") + handle = pin_manager.PinHandle( + tensor=None, nbytes=64, kind="save_stager", pinned=True + ) + pin_manager.release(handle) + self.assertEqual(pin_manager.total_pinned_bytes(), 0) + pin_manager.release(handle) # second release must be a no-op + self.assertEqual(pin_manager.total_pinned_bytes(), 0) + + +class PinConformanceTests(unittest.TestCase): + """No direct pinning outside the pin manager (PIN_MANAGER_PLAN S2). + + Every page-lock in the memory subsystem must route through pin_manager so + the ledger stays authoritative. Scope is the offload subsystem + async + save; upstream dataloader pin_memory usage is out of scope.""" + + SCOPED_FILES = ( + "toolkit/async_save.py", + "toolkit/memory_management/bounce_pool.py", + "toolkit/memory_management/ingraph_stream.py", + "toolkit/memory_management/manager.py", + "toolkit/memory_management/manager_modules.py", + "toolkit/memory_management/checkpoint_autotuner.py", + "toolkit/memory_management/pinned_arena.py", + ) + PATTERN = re.compile(r"pin_memory\s*=\s*True|\.pin_memory\(\)") + + def test_no_direct_pinning_outside_pin_manager(self): + offenders = [] + for rel in self.SCOPED_FILES: + path = REPO_ROOT / rel + if not path.exists(): + continue + for lineno, line in enumerate( + path.read_text(encoding="utf-8").splitlines(), 1 + ): + stripped = line.lstrip() + if stripped.startswith("#"): + continue + if self.PATTERN.search(line) and "noqa: pin-manager" not in line: + offenders.append(f"{rel}:{lineno}: {line.strip()}") + self.assertEqual( + offenders, [], + "direct pinning outside pin_manager (route through pin_alloc / " + "pin_tensor_in_place):\n" + "\n".join(offenders), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/toolkit/memory_management/allocator_cap.py b/toolkit/memory_management/allocator_cap.py new file mode 100644 index 0000000000..47707e73bd --- /dev/null +++ b/toolkit/memory_management/allocator_cap.py @@ -0,0 +1,122 @@ +"""Shared WDDM allocator-cap mechanism for both memory backends.""" + +from __future__ import annotations + +import sys + +import torch + +from . import vram_budget + +GIB = 1024 ** 3 +APPLIED_FRACTIONS: dict[int, float] = {} +RELIEF_BYTES: dict[int, int] = {} + + + + +def wddm_cliff_cap_bytes(device, wddm_hard_gib=None) -> int: + """Return the governing allocator cap at the current WDDM cliff.""" + dev = torch.device(device if device is not None else "cuda") + index = dev.index if dev.index is not None else torch.cuda.current_device() + try: + hard_gib = float(wddm_hard_gib) if wddm_hard_gib is not None else 1.0 + except (TypeError, ValueError): + hard_gib = 1.0 + if hard_gib <= 0: + hard_gib = 1.0 + total = vram_budget.device_total_bytes(index) + free_bytes, _ = vram_budget.device_mem_info(index) + reserved_bytes = torch.cuda.memory_reserved(index) + return int( + vram_budget.cap_fraction( + total, free_bytes, reserved_bytes, hard_gib + ) + * total + ) + + + + +def applied_cap_bytes(device) -> int | None: + """Return the last cap bound on this device in physical allocator bytes.""" + if not torch.cuda.is_available(): + return None + dev = torch.device(device if device is not None else "cuda") + if dev.type != "cuda": + return None + index = dev.index if dev.index is not None else torch.cuda.current_device() + fraction = APPLIED_FRACTIONS.get(index) + if fraction is None: + return None + return int(float(fraction) * vram_budget.real_device_total_bytes(index)) + + +def apply_wddm_hard_allocator_cap( + device, + wddm_hard_gib=None, + *, + target_cap_bytes=None, + log_prefix="[MemoryManager]", +): + """Bind torch's allocator below the WDDM dedicated-memory cliff. + + Call only at a phase boundary. The governing capacity may be a simulated + smaller card, but torch's fraction is always converted against the physical + card total. + """ + if sys.platform != "win32" or not torch.cuda.is_available(): + return None + dev = torch.device(device if device is not None else "cuda") + if dev.type != "cuda": + return None + index = dev.index if dev.index is not None else torch.cuda.current_device() + try: + hard_gib = float(wddm_hard_gib) if wddm_hard_gib is not None else 1.0 + except (TypeError, ValueError): + hard_gib = 1.0 + if hard_gib <= 0: + hard_gib = 1.0 + total = vram_budget.device_total_bytes(index) + real_total = vram_budget.real_device_total_bytes(index) + free_bytes, _governing_total = vram_budget.device_mem_info(index) + reserved_bytes = torch.cuda.memory_reserved(index) + cliff_fraction = vram_budget.cap_fraction( + total, free_bytes, reserved_bytes, hard_gib + ) + fraction = cliff_fraction + reclaimed = False + if target_cap_bytes is not None: + target_fraction = float(target_cap_bytes) / float(total) + fraction = max(0.1, min(cliff_fraction, target_fraction)) + reclaimed = fraction < cliff_fraction - 1e-9 + + relief_bytes = RELIEF_BYTES.get(index, 0) + if relief_bytes: + fraction = min(1.0, fraction + relief_bytes / float(total)) + applied = fraction * total / float(real_total) + previous = APPLIED_FRACTIONS.get(index) + tolerance = (64 * 1024**2) / real_total + if previous is not None and abs(previous - applied) < tolerance: + return previous + + torch.cuda.set_per_process_memory_fraction(applied, index) + APPLIED_FRACTIONS[index] = applied + non_torch = max(0, (total - free_bytes) - reserved_bytes) + source = ( + f"reclaim target, cliff {cliff_fraction * total / GIB:.2f} GiB" + if reclaimed + else "cliff bound" + ) + if relief_bytes: + source += f"; +{relief_bytes / GIB:.2f} GiB post-violation relief" + if total != real_total: + source += f"; SIMULATED {total / GIB:.2f} GiB card" + print( + f"{log_prefix} WDDM hard allocator cap: " + f"{fraction * total / GIB:.2f}/{total / GIB:.2f} GiB " + f"({source}; margin {hard_gib:.2f} GiB, " + f"non_torch {non_torch / GIB:.2f} GiB; allocation beyond this " + "recycles cache or raises OOM instead of silently paging)" + ) + return applied diff --git a/toolkit/memory_management/dxgi_meminfo.py b/toolkit/memory_management/dxgi_meminfo.py new file mode 100644 index 0000000000..802fb14e4c --- /dev/null +++ b/toolkit/memory_management/dxgi_meminfo.py @@ -0,0 +1,493 @@ +"""Windows DXGI video-memory budget sensor. + +This module is intentionally ctypes-only and Windows-only. On non-Windows, or +when DXGI is unavailable, public query functions return None so callers can use +their existing host-RAM proxy. +""" + +from __future__ import annotations + +import ctypes +import os +import threading +import time +from typing import NamedTuple, Optional + + +_DXGI_ERROR_NOT_FOUND = 0x887A0002 +_DXGI_ADAPTER_FLAG_SOFTWARE = 2 +_DXGI_MEMORY_SEGMENT_GROUP_LOCAL = 0 +_DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL = 1 +_NVIDIA_VENDOR_ID = 0x10DE + + +class GUID(ctypes.Structure): + _fields_ = [ + ("Data1", ctypes.c_uint32), + ("Data2", ctypes.c_uint16), + ("Data3", ctypes.c_uint16), + ("Data4", ctypes.c_ubyte * 8), + ] + + +IID_IDXGIFactory1 = GUID( + 0x770AAE78, + 0xF26F, + 0x4DBA, + (ctypes.c_ubyte * 8)(0xA8, 0x29, 0x25, 0x3C, 0x83, 0xD1, 0xB3, 0x87), +) +IID_IDXGIAdapter3 = GUID( + 0x645967A4, + 0x1392, + 0x4310, + (ctypes.c_ubyte * 8)(0xA7, 0x98, 0x80, 0x53, 0xCE, 0x3E, 0x93, 0xFD), +) + + +class LUID(ctypes.Structure): + _fields_ = [ + ("LowPart", ctypes.c_uint32), + ("HighPart", ctypes.c_int32), + ] + + +class DXGI_ADAPTER_DESC1(ctypes.Structure): + _fields_ = [ + ("Description", ctypes.c_wchar * 128), + ("VendorId", ctypes.c_uint32), + ("DeviceId", ctypes.c_uint32), + ("SubSysId", ctypes.c_uint32), + ("Revision", ctypes.c_uint32), + ("DedicatedVideoMemory", ctypes.c_size_t), + ("DedicatedSystemMemory", ctypes.c_size_t), + ("SharedSystemMemory", ctypes.c_size_t), + ("AdapterLuid", LUID), + ("Flags", ctypes.c_uint32), + ] + + +class _DXGI_QUERY_VIDEO_MEMORY_INFO(ctypes.Structure): + _fields_ = [ + ("Budget", ctypes.c_uint64), + ("CurrentUsage", ctypes.c_uint64), + ("AvailableForReservation", ctypes.c_uint64), + ("CurrentReservation", ctypes.c_uint64), + ] + + +class DxgiMemoryInfo(NamedTuple): + budget_bytes: int + current_usage_bytes: int + available_for_reservation_bytes: int + current_reservation_bytes: int + + +# Which adapter-match methods are trustworthy enough to *drive control* +# (aggressive pin-for-speed), versus only telemetry. Auto-detected confident +# matches (single_nvidia, and the Stage-B luid match) are safe. An explicit +# env override is honored for control but is a human-forced adapter, so it is +# logged distinctly as "manual" -- a misconfiguration must be visible, not +# silently trusted as if auto-verified. Everything else (sole_hardware_adapter, +# global_conservative, or no adapter at all) stays conservative. +_CONTROL_SAFE_METHODS = frozenset({"single_nvidia", "luid"}) +_CONTROL_MANUAL_METHODS = frozenset({"env_override"}) + + +def safe_for_control(match_method: Optional[str]) -> bool: + """Is a reading from this match method trustworthy enough to drive control? + + Pure classifier (unit-testable, no ctypes). True for confidently + auto-detected adapters and for an explicit manual override; False for + ambiguous/fallback selections and when the sensor is unavailable. + """ + if not match_method: + return False + return match_method in _CONTROL_SAFE_METHODS or match_method in _CONTROL_MANUAL_METHODS + + +def is_manual_control(match_method: Optional[str]) -> bool: + """True when control is driven off a human-forced (env-override) adapter.""" + return bool(match_method) and match_method in _CONTROL_MANUAL_METHODS + + +class DxgiAdapterInfo(NamedTuple): + index: int + description: str + vendor_id: int + device_id: int + luid: str + match_method: str + safe_for_control: bool + manual_control: bool + + +class DxgiAdapterRecord(NamedTuple): + index: int + description: str + vendor_id: int + device_id: int + luid: str + flags: int + dedicated_video_memory_bytes: int + dedicated_system_memory_bytes: int + shared_system_memory_bytes: int + is_software: bool + + +class _AdapterSelection(NamedTuple): + ptr: ctypes.c_void_p + info: DxgiAdapterInfo + + +_selection_lock = threading.Lock() +_selection: Optional[_AdapterSelection] = None +_selection_unavailable = False +_query_lock = threading.Lock() +_query_cache: dict[tuple[int, int], tuple[float, DxgiMemoryInfo]] = {} +_warned: set[str] = set() + + +def _log_once(key: str, message: str) -> None: + if key in _warned: + return + _warned.add(key) + print(message) + + +def _failed(hr: int) -> bool: + return int(hr) < 0 + + +def _hr_u32(hr: int) -> int: + return int(hr) & 0xFFFFFFFF + + +def _luid_text(luid: LUID) -> str: + return f"{int(luid.HighPart) & 0xFFFFFFFF:08x}:{int(luid.LowPart):08x}" + + +def _com_call(ptr, vtbl_index, restype, argtypes, *args): + """Dispatch a COM vtable call against ``ptr``.""" + if not ptr: + raise RuntimeError("null COM pointer") + winfunctype = getattr(ctypes, "WINFUNCTYPE", ctypes.CFUNCTYPE) + raw_ptr = ctypes.c_void_p(ptr) if isinstance(ptr, int) else ptr + vtbl = ctypes.cast(raw_ptr, ctypes.POINTER(ctypes.POINTER(ctypes.c_void_p))).contents + address = vtbl[vtbl_index] + prototype = winfunctype(restype, ctypes.c_void_p, *argtypes) + return prototype(address)(raw_ptr, *args) + + +def _create_dxgi_factory1(): + if os.name != "nt": + return None + factory = ctypes.c_void_p() + hr = ctypes.windll.dxgi.CreateDXGIFactory1( + ctypes.byref(IID_IDXGIFactory1), + ctypes.byref(factory), + ) + if _failed(hr): + raise OSError(f"CreateDXGIFactory1 failed hr=0x{_hr_u32(hr):08x}") + return factory + + +def _get_desc1(adapter) -> DXGI_ADAPTER_DESC1: + desc = DXGI_ADAPTER_DESC1() + # IDXGIAdapter1::GetDesc1: IUnknown 0-2, IDXGIObject 3-6, + # IDXGIAdapter 7-9, IDXGIAdapter1 slot 10. + hr = _com_call( + adapter, + 10, + ctypes.c_long, + [ctypes.POINTER(DXGI_ADAPTER_DESC1)], + ctypes.byref(desc), + ) + if _failed(hr): + raise OSError(f"IDXGIAdapter1::GetDesc1 failed hr=0x{_hr_u32(hr):08x}") + return desc + + +def _record_from_desc(index: int, desc: DXGI_ADAPTER_DESC1) -> DxgiAdapterRecord: + return DxgiAdapterRecord( + index=index, + description=str(desc.Description).rstrip("\x00"), + vendor_id=int(desc.VendorId), + device_id=int(desc.DeviceId), + luid=_luid_text(desc.AdapterLuid), + flags=int(desc.Flags), + dedicated_video_memory_bytes=int(desc.DedicatedVideoMemory), + dedicated_system_memory_bytes=int(desc.DedicatedSystemMemory), + shared_system_memory_bytes=int(desc.SharedSystemMemory), + is_software=bool(int(desc.Flags) & _DXGI_ADAPTER_FLAG_SOFTWARE), + ) + + +def _enum_adapters1(factory) -> list[tuple[ctypes.c_void_p, DxgiAdapterRecord]]: + adapters = [] + index = 0 + while True: + adapter = ctypes.c_void_p() + # IDXGIFactory1::EnumAdapters1: IUnknown 0-2, IDXGIObject 3-6, + # IDXGIFactory 7-11, IDXGIFactory1 slot 12. + hr = _com_call( + factory, + 12, + ctypes.c_long, + [ctypes.c_uint32, ctypes.POINTER(ctypes.c_void_p)], + ctypes.c_uint32(index), + ctypes.byref(adapter), + ) + if _hr_u32(hr) == _DXGI_ERROR_NOT_FOUND: + break + if _failed(hr): + raise OSError(f"IDXGIFactory1::EnumAdapters1 failed hr=0x{_hr_u32(hr):08x}") + desc = _get_desc1(adapter) + adapters.append((adapter, _record_from_desc(index, desc))) + index += 1 + return adapters + + +def enumerate_adapters() -> list[DxgiAdapterRecord]: + """Return all DXGI adapters for diagnostics/probes, or [] if unavailable.""" + if os.name != "nt": + return [] + try: + factory = _create_dxgi_factory1() + if factory is None: + return [] + return [record for _adapter, record in _enum_adapters1(factory)] + except Exception as exc: + _log_once( + "enumerate_failed", + f"[DXGI] adapter enumeration unavailable: {exc}", + ) + return [] + + +def _query_interface_adapter3(adapter) -> ctypes.c_void_p: + adapter3 = ctypes.c_void_p() + # IUnknown::QueryInterface slot 0. + hr = _com_call( + adapter, + 0, + ctypes.c_long, + [ctypes.POINTER(GUID), ctypes.POINTER(ctypes.c_void_p)], + ctypes.byref(IID_IDXGIAdapter3), + ctypes.byref(adapter3), + ) + if _failed(hr): + raise OSError(f"IDXGIAdapter3 QueryInterface failed hr=0x{_hr_u32(hr):08x}") + return adapter3 + + +def _select_hardware_adapter3() -> Optional[_AdapterSelection]: + if os.name != "nt": + return None + + factory = _create_dxgi_factory1() + adapters = _enum_adapters1(factory) + if not adapters: + _log_once("no_adapters", "[DXGI] no adapters found; using legacy pin proxy") + return None + + override = os.environ.get("AI_TOOLKIT_WDDM_DXGI_ADAPTER_INDEX", "").strip() + if override: + try: + override_index = int(override) + for adapter, record in adapters: + if record.index == override_index: + return _selection_from_record( + adapter, record, match_method="env_override" + ) + except Exception: + pass + _log_once( + "bad_override", + "[DXGI] AI_TOOLKIT_WDDM_DXGI_ADAPTER_INDEX did not match a DXGI " + "adapter; using legacy pin proxy", + ) + return None + + hardware = [ + (adapter, record) + for adapter, record in adapters + if not record.is_software + ] + nvidia = [ + (adapter, record) + for adapter, record in hardware + if record.vendor_id == _NVIDIA_VENDOR_ID + ] + if len(nvidia) == 1: + adapter, record = nvidia[0] + return _selection_from_record(adapter, record, match_method="single_nvidia") + if len(nvidia) == 0 and len(hardware) == 1: + adapter, record = hardware[0] + _log_once( + "sole_hardware_adapter", + "[DXGI] no NVIDIA adapter found; using sole hardware adapter for " + "shared-memory pin budget", + ) + return _selection_from_record( + adapter, record, match_method="sole_hardware_adapter" + ) + _log_once( + "ambiguous_adapter", + "[DXGI] adapter selection is ambiguous in Stage A; using legacy pin proxy", + ) + return None + + +def _selection_from_record(adapter, record: DxgiAdapterRecord, match_method: str): + adapter3 = _query_interface_adapter3(adapter) + manual = is_manual_control(match_method) + if manual: + # A human forced the adapter via AI_TOOLKIT_WDDM_DXGI_ADAPTER_INDEX. It + # still drives control, but log it distinctly so a misconfiguration is + # visible rather than lumped in with confident auto-detection. + _log_once( + "manual_control_override", + f"[DXGI] manual adapter override in effect (index={record.index} " + f"{record.description!r}); driving control as 'manual' -- verify " + "this is the training GPU", + ) + info = DxgiAdapterInfo( + index=record.index, + description=record.description, + vendor_id=record.vendor_id, + device_id=record.device_id, + luid=record.luid, + match_method=match_method, + safe_for_control=safe_for_control(match_method), + manual_control=manual, + ) + return _AdapterSelection(adapter3, info) + + +def _selected_adapter() -> Optional[_AdapterSelection]: + global _selection, _selection_unavailable + if os.name != "nt": + return None + with _selection_lock: + if _selection is not None: + return _selection + if _selection_unavailable: + return None + try: + _selection = _select_hardware_adapter3() + except Exception as exc: + _log_once( + "selection_failed", + f"[DXGI] adapter selection unavailable: {exc}; using legacy pin proxy", + ) + _selection = None + if _selection is None: + _selection_unavailable = True + return _selection + + +def selected_adapter_info() -> Optional[DxgiAdapterInfo]: + selection = _selected_adapter() + return selection.info if selection is not None else None + + +def control_is_eligible(cuda_device_index: int = 0) -> bool: + """True when the resolved adapter is trustworthy enough to drive control. + + Consumed by the pin-for-speed policy: aggressive full-pin engages only when + this is True (confident auto-detect or an explicit manual override). When + DXGI is unavailable/ambiguous this is False and the policy stays + conservative. ``cuda_device_index`` is accepted for the Stage-B per-device + match; Stage A resolves the single cached adapter. + """ + del cuda_device_index # Stage A: single cached adapter. + info = selected_adapter_info() + return bool(info is not None and info.safe_for_control) + + +def query_video_memory_info( + cuda_device_index: int = 0, + segment_group: int = _DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL, + min_interval_s: float = 0.5, +) -> Optional[DxgiMemoryInfo]: + """Query LOCAL or NON_LOCAL DXGI video-memory info for the selected adapter.""" + del cuda_device_index # Stage A uses one cached adapter; Stage B matches per CUDA device. + if os.name != "nt": + return None + selection = _selected_adapter() + if selection is None: + return None + cache_key = (0, int(segment_group)) + now = time.monotonic() + with _query_lock: + cached = _query_cache.get(cache_key) + if cached is not None and min_interval_s > 0: + ts, info = cached + if now - ts < min_interval_s: + return info + try: + raw = _DXGI_QUERY_VIDEO_MEMORY_INFO() + # IDXGIAdapter3::QueryVideoMemoryInfo: IUnknown 0-2, IDXGIObject 3-6, + # IDXGIAdapter 7-9, IDXGIAdapter1 10, IDXGIAdapter2 11-13, + # IDXGIAdapter3 slot 14. + hr = _com_call( + selection.ptr, + 14, + ctypes.c_long, + [ctypes.c_uint32, ctypes.c_int, ctypes.POINTER(_DXGI_QUERY_VIDEO_MEMORY_INFO)], + ctypes.c_uint32(0), + ctypes.c_int(int(segment_group)), + ctypes.byref(raw), + ) + if _failed(hr): + raise OSError( + f"IDXGIAdapter3::QueryVideoMemoryInfo failed hr=0x{_hr_u32(hr):08x}" + ) + info = DxgiMemoryInfo( + budget_bytes=int(raw.Budget), + current_usage_bytes=int(raw.CurrentUsage), + available_for_reservation_bytes=int(raw.AvailableForReservation), + current_reservation_bytes=int(raw.CurrentReservation), + ) + except Exception as exc: + _log_once( + "query_failed", + f"[DXGI] video-memory query unavailable: {exc}; using legacy pin proxy", + ) + return None + with _query_lock: + _query_cache[cache_key] = (now, info) + return info + + +def query_non_local_video_memory_info( + cuda_device_index: int = 0, + min_interval_s: float = 0.5, +) -> Optional[DxgiMemoryInfo]: + return query_video_memory_info( + cuda_device_index=cuda_device_index, + segment_group=_DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL, + min_interval_s=min_interval_s, + ) + + +def query_local_video_memory_info( + cuda_device_index: int = 0, + min_interval_s: float = 0.5, +) -> Optional[DxgiMemoryInfo]: + return query_video_memory_info( + cuda_device_index=cuda_device_index, + segment_group=_DXGI_MEMORY_SEGMENT_GROUP_LOCAL, + min_interval_s=min_interval_s, + ) + + +def compute_non_local_headroom_bytes( + budget_bytes: int, + current_usage_bytes: int, + spill_reserve_bytes: int, +) -> int: + return max( + 0, + int(budget_bytes) - int(current_usage_bytes) - max(0, int(spill_reserve_bytes)), + ) diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index c048e940d5..3ea6aa9f13 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -4,7 +4,9 @@ ConvLayerMemoryManager, OstrisLinearLayerMemoryManager, _DEVICE_STATE, + _release_cpu_pin, ) +from . import allocator_cap, pin_manager import random LINEAR_MODULES = [ @@ -71,6 +73,34 @@ def memory_managed_to(self, *args, **kwargs): return self.module._mm_to(dtype=dtype) return self.module + @staticmethod + def offload_shape_key_from_batch(batch_list, **flags): + """Build a stable policy key from batch tensor shapes and flags.""" + shapes = [] + + def visit(value): + if torch.is_tensor(value): + shape = tuple(int(dim) for dim in value.shape) + if len(shape) >= 2: + shapes.append((str(value.dtype), shape)) + return + if isinstance(value, dict): + for item in value.values(): + visit(item) + return + if isinstance(value, (list, tuple)): + for item in value: + visit(item) + return + for name in ("tensor", "latents", "images", "control_tensor"): + if hasattr(value, name): + visit(getattr(value, name)) + + visit(batch_list) + shape_key = tuple(sorted(set(shapes)))[:16] + policy_key = tuple(sorted(flags.items())) + return shape_key, policy_key + @classmethod def attach( cls, @@ -79,6 +109,7 @@ def attach( offload_percent: float = 1.0, ignore_modules: list[torch.nn.Module] = [] ): + allocator_cap.apply_wddm_hard_allocator_cap(device) if hasattr(module, "_memory_manager"): # already attached return @@ -212,12 +243,13 @@ def detach(cls, module: torch.nn.Module): if param is None or not isinstance(param, torch.nn.Parameter): continue try: - if param.data.is_pinned(): + released = _release_cpu_pin(param.data) + if released is not param.data: object.__setattr__( child, param_name, torch.nn.Parameter( - param.data.clone(), + released, requires_grad=param.requires_grad, ), ) @@ -232,8 +264,7 @@ def detach(cls, module: torch.nn.Module): try: if buf.device.type != "cpu": buf = buf.to("cpu") - if buf.is_pinned(): - buf = buf.clone() + buf = _release_cpu_pin(buf) child._buffers[buf_name] = buf except Exception: pass @@ -251,4 +282,5 @@ def detach(cls, module: torch.nn.Module): for key in keys_to_delete: del _DEVICE_STATE[key] + pin_manager.reconcile(allow_shrink=False) torch.cuda.empty_cache() diff --git a/toolkit/memory_management/manager_modules.py b/toolkit/memory_management/manager_modules.py index 44f3b28880..5217e7b1ba 100644 --- a/toolkit/memory_management/manager_modules.py +++ b/toolkit/memory_management/manager_modules.py @@ -14,6 +14,8 @@ from typing import TYPE_CHECKING, Optional, Tuple from torch.overrides import has_torch_function_unary # (ADD) torchao detection +from . import pin_manager + if TYPE_CHECKING: from .manager import MemoryManager @@ -190,15 +192,46 @@ def _pin_inner_tensors(t: torch.Tensor) -> None: continue if hasattr(inner, "__tensor_flatten__"): _pin_inner_tensors(inner) # recurse: AQT -> tensor_impl -> data/scale - elif ( - isinstance(inner, torch.Tensor) - and inner.device.type == "cpu" - and not inner.is_pinned() - ): - try: - setattr(t, name, inner.pin_memory()) - except Exception: - pass + elif isinstance(inner, torch.Tensor) and inner.device.type == "cpu": + if pin_manager.is_host_pinned(inner): + continue + if pin_manager.pin_tensor_in_place(inner, kind="weights"): + continue + size = int(inner.numel() * inner.element_size()) + handle = pin_manager.pin_alloc( + size, "weights", required=False + ) + if handle.pinned: + view = handle.tensor.view(inner.dtype).reshape(inner.shape) + view.copy_(inner) + setattr(t, name, view) + + +def _release_cpu_pin(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Release pin-manager ownership and return pageable storage when needed.""" + if t is None or not isinstance(t, torch.Tensor) or t.device.type != "cpu": + return t + if hasattr(t, "__tensor_flatten__"): + try: + names, _ = t.__tensor_flatten__() + except Exception: + names = () + for name in names: + inner = getattr(t, name, None) + if not isinstance(inner, torch.Tensor): + continue + replacement = _release_cpu_pin(inner) + if replacement is not inner: + setattr(t, name, replacement) + return t + if pin_manager.unpin_tensor_in_place(t, kind="weights"): + return t + if t.is_pinned(): + pin_manager.release_pinned_bytes( + int(t.numel() * t.element_size()), kind="weights" + ) + return t.clone() + return t def _ensure_cpu_pinned(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: @@ -216,10 +249,16 @@ def _ensure_cpu_pinned(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: _pin_inner_tensors(t) return t if torch.cuda.is_available(): - try: - t = t.pin_memory() - except RuntimeError: - pass + if not pin_manager.is_host_pinned(t): + if not pin_manager.pin_tensor_in_place(t, kind="weights"): + size = int(t.numel() * t.element_size()) + handle = pin_manager.pin_alloc( + size, "weights", required=False + ) + if handle.pinned: + pinned = handle.tensor.view(t.dtype).reshape(t.shape) + pinned.copy_(t) + t = pinned return t @@ -664,11 +703,8 @@ def __init__( continue if buf.device.type != "cpu": buf = buf.to("cpu") - if torch.cuda.is_available() and not buf.is_pinned(): - try: - buf = buf.pin_memory() - except RuntimeError: - pass + if torch.cuda.is_available(): + buf = _ensure_cpu_pinned(buf) module._buffers[name] = buf bias = module._parameters.get("bias", None) if bias is not None: diff --git a/toolkit/memory_management/nvml_meminfo.py b/toolkit/memory_management/nvml_meminfo.py new file mode 100644 index 0000000000..550cafff6a --- /dev/null +++ b/toolkit/memory_management/nvml_meminfo.py @@ -0,0 +1,234 @@ +"""Cross-platform NVML sensor for TRUE physical VRAM occupancy. + +Why this exists +--------------- +``torch.cuda.mem_get_info`` does NOT report physical card availability. It +reports what the driver is willing to promise *this process*, which on WDDM is +a budget the OS intends to satisfy by paging other processes out. Measured on +an RTX 4070 (11.99 GiB) while a second process held 9 GiB: + + nvidia-smi / NVML physical free : 0.55 GiB <- the truth + torch.cuda.mem_get_info free : 4.70 GiB <- over-reports ~8x + DXGI LOCAL Budget - CurrentUsage: 4.58 GiB <- also over-reports + +Planning residency against the optimistic numbers is exactly how a run silently +crosses the dedicated-VRAM cliff: the allocation "succeeds", WDDM pages to +system RAM, and throughput collapses with no error and no allocator retry to +detect it (observed: a training job at 304 s/it instead of 73 s/it because an +orphaned process was squatting on the card). + +NVML reports device-wide physical occupancy across *every* process, so it stays +correct when a game, a ComfyUI server, or a leftover job shares the GPU. It is +also the only such sensor that works on both Windows and Linux -- DXGI is +Windows-only and, per the numbers above, measures a permission rather than an +availability. + +This module is ctypes-only and adds no dependency: it binds ``nvml.dll`` on +Windows and ``libnvidia-ml.so.1`` on Linux, both shipped with the NVIDIA driver. +Every public function returns ``None`` when NVML is unavailable (no NVIDIA +driver, non-NVIDIA GPU) so callers can fall back to their previous behavior. + +Cost: ``nvmlDeviceGetMemoryInfo`` measures ~3 us/call, ~25x cheaper than +``torch.cuda.mem_get_info`` (~80 us), so it is safe to call on hot paths. +""" + +from __future__ import annotations + +import ctypes +import os +import threading +from typing import NamedTuple + +_NVML_SUCCESS = 0 + + +class NvmlMemoryInfo(NamedTuple): + total_bytes: int + free_bytes: int + used_bytes: int + + +class _NvmlMemory(ctypes.Structure): + _fields_ = [ + ("total", ctypes.c_ulonglong), + ("free", ctypes.c_ulonglong), + ("used", ctypes.c_ulonglong), + ] + + +_lock = threading.Lock() +_lib = None +_init_failed = False +_handles: dict[int, ctypes.c_void_p] = {} +_warned: set[str] = set() + + +def _log_once(key: str, message: str) -> None: + if key in _warned: + return + _warned.add(key) + print(message) + + +def _load_library(): + names = ("nvml.dll",) if os.name == "nt" else ("libnvidia-ml.so.1", "libnvidia-ml.so") + for name in names: + try: + return ctypes.CDLL(name) + except OSError: + continue + return None + + +def _library(): + """Load + initialize NVML once. None when unavailable.""" + global _lib, _init_failed + if _lib is not None: + return _lib + if _init_failed: + return None + lib = _load_library() + if lib is None: + _init_failed = True + _log_once( + "nvml_missing", + "[NVML] library not found; falling back to torch.cuda.mem_get_info " + "(device-free readings will not see other processes)", + ) + return None + try: + rc = lib.nvmlInit_v2() + except Exception as exc: # pragma: no cover - driver-level failure + _init_failed = True + _log_once("nvml_init_raise", f"[NVML] init raised: {exc}; using mem_get_info") + return None + if rc != _NVML_SUCCESS: + _init_failed = True + _log_once("nvml_init_failed", f"[NVML] nvmlInit_v2 failed rc={rc}; using mem_get_info") + return None + _lib = lib + return _lib + + +def _torch_device_uuid(cuda_device_index: int) -> str | None: + """UUID of a CUDA device as torch sees it, normalized. None if unavailable.""" + try: + import torch + + raw = getattr(torch.cuda.get_device_properties(cuda_device_index), "uuid", None) + except Exception: + return None + if raw is None: + return None + return str(raw).replace("GPU-", "").replace("-", "").lower() + + +def _nvml_device_uuid(lib, handle: ctypes.c_void_p) -> str | None: + buf = ctypes.create_string_buffer(96) + try: + rc = lib.nvmlDeviceGetUUID(handle, buf, ctypes.c_uint(96)) + except Exception: # pragma: no cover - symbol missing on ancient drivers + return None + if rc != _NVML_SUCCESS: + return None + return buf.value.decode(errors="replace").replace("GPU-", "").replace("-", "").lower() + + +def _resolve_handle(lib, cuda_device_index: int) -> ctypes.c_void_p | None: + """Map a CUDA device index to an NVML handle. + + Prefer UUID matching: NVML enumerates *physical* devices, while CUDA indices + are affected by CUDA_VISIBLE_DEVICES and CUDA_DEVICE_ORDER, so index-to-index + identity is not guaranteed on a multi-GPU box. Fall back to the index only + when UUIDs are unavailable. + """ + count = ctypes.c_uint() + if lib.nvmlDeviceGetCount_v2(ctypes.byref(count)) != _NVML_SUCCESS: + return None + + want = _torch_device_uuid(cuda_device_index) + if want: + for idx in range(int(count.value)): + handle = ctypes.c_void_p() + if lib.nvmlDeviceGetHandleByIndex_v2(ctypes.c_uint(idx), ctypes.byref(handle)) != _NVML_SUCCESS: + continue + if _nvml_device_uuid(lib, handle) == want: + return handle + _log_once( + "uuid_no_match", + f"[NVML] no NVML device matched CUDA device {cuda_device_index} by UUID; " + "falling back to index mapping", + ) + + if cuda_device_index >= int(count.value): + return None + handle = ctypes.c_void_p() + if lib.nvmlDeviceGetHandleByIndex_v2( + ctypes.c_uint(cuda_device_index), ctypes.byref(handle) + ) != _NVML_SUCCESS: + return None + return handle + + +def _handle_for(cuda_device_index: int) -> ctypes.c_void_p | None: + with _lock: + cached = _handles.get(cuda_device_index) + if cached is not None: + return cached + lib = _library() + if lib is None: + return None + try: + handle = _resolve_handle(lib, cuda_device_index) + except Exception as exc: # pragma: no cover + _log_once("handle_raise", f"[NVML] device resolution raised: {exc}") + return None + if handle is None: + _log_once( + "handle_missing", + f"[NVML] could not resolve NVML handle for CUDA device {cuda_device_index}", + ) + return None + _handles[cuda_device_index] = handle + return handle + + +def query_device_memory_info(cuda_device_index: int = 0) -> NvmlMemoryInfo | None: + """Physical VRAM totals for a device, across ALL processes on the GPU. + + ``free_bytes`` is the real bytes left on the card -- unlike + ``torch.cuda.mem_get_info``, it accounts for other processes (a game, a + ComfyUI server, an orphaned training job). Returns None when NVML is + unavailable, so callers must keep a fallback. + """ + lib = _library() + if lib is None: + return None + handle = _handle_for(cuda_device_index) + if handle is None: + return None + mem = _NvmlMemory() + try: + rc = lib.nvmlDeviceGetMemoryInfo(handle, ctypes.byref(mem)) + except Exception as exc: # pragma: no cover + _log_once("query_raise", f"[NVML] memory query raised: {exc}") + return None + if rc != _NVML_SUCCESS: + _log_once("query_failed", f"[NVML] nvmlDeviceGetMemoryInfo failed rc={rc}") + return None + return NvmlMemoryInfo( + total_bytes=int(mem.total), + free_bytes=int(mem.free), + used_bytes=int(mem.used), + ) + + +def physical_free_bytes(cuda_device_index: int = 0) -> int | None: + """True physical free bytes on the card, or None when NVML is unavailable.""" + info = query_device_memory_info(cuda_device_index) + return None if info is None else info.free_bytes + + +def is_available() -> bool: + """True when NVML can be used to read physical occupancy.""" + return _library() is not None diff --git a/toolkit/memory_management/pin_manager.py b/toolkit/memory_management/pin_manager.py new file mode 100644 index 0000000000..23242d8a21 --- /dev/null +++ b/toolkit/memory_management/pin_manager.py @@ -0,0 +1,752 @@ +"""Single authority for process-local pinned host memory. + +Pinned host memory is a shared budget on Windows/WDDM. This module keeps a +per-consumer ledger and centralizes the DXGI/RAM headroom checks that used to +live in individual consumers. +""" + +from __future__ import annotations + +import os +import threading +from dataclasses import dataclass +from typing import Callable, Optional + +import torch + + +try: + import psutil as _psutil +except Exception: + _psutil = None + + +GIB = 1024 ** 3 + + +class PinBudgetExceeded(RuntimeError): + """Raised when a must-pin allocation cannot fit the current host-pin budget.""" + + +@dataclass +class PinHandle: + tensor: torch.Tensor + nbytes: int + kind: str + pinned: bool + # "alloc" = cudaHostAlloc via torch's caching host allocator (bytes only + # return DXGI budget after _empty_host_pin_cache, and the allocator + # rounds requests up to power-of-two buckets -- the DXGI cost can be up + # to 2x nbytes). "register" = exact-size pageable tensor pinned with + # cudaHostRegister (DXGI cost == nbytes, returned immediately on + # release). + mechanism: str = "alloc" + + +_LOCK = threading.RLock() +_LEDGER: dict[str, int] = {} +_EVICTABLES: list[Callable[[int], int]] = [] + +# Weight-tier consumers are the LOWEST pin priority (PIN_MANAGER_PLAN +# allocation strategy): they may reclaim the torch host cache during +# reconcile, but must never shrink evictable higher-priority consumers +# (the bounce pool) to make room for themselves. +_WEIGHT_TIER_KINDS = ("weights", "ingraph_pack") + +_SPILL_RESERVE_FLOOR_GIB_OVERRIDE: Optional[float] = None +_SPILL_RESERVE_PCT_OVERRIDE: Optional[float] = None +_HOST_CACHE_RESERVE_BYTES_OVERRIDE: Optional[int] = None + +dxgi_meminfo = None + + +def _env_bool(name: str, default: bool = False) -> bool: + raw = os.environ.get(name) + if raw is None: + return default + return raw.strip().lower() not in ("0", "false", "no", "off", "") + + +def _cuda_device_index(device=None) -> int: + if device is None: + return 0 + try: + dev = torch.device(device) + except Exception: + return 0 + if dev.type != "cuda": + return 0 + if dev.index is not None: + return int(dev.index) + try: + return int(torch.cuda.current_device()) + except Exception: + return 0 + + +def total_pinned_bytes() -> int: + with _LOCK: + return int(sum(_LEDGER.values())) + + +def pinned_bytes_by_kind() -> dict[str, int]: + with _LOCK: + return dict(_LEDGER) + + +def register_pinned_bytes(n: int, kind: str = "unknown") -> None: + n = int(n) + if n <= 0: + return + kind = str(kind or "unknown") + with _LOCK: + _LEDGER[kind] = _LEDGER.get(kind, 0) + n + + +def release_pinned_bytes(n: int, kind: str = "unknown") -> None: + n = int(n) + if n <= 0: + return + kind = str(kind or "unknown") + with _LOCK: + # Clamp within the kind only: an unmatched release (e.g. a consumer + # releasing bytes it never registered because pinning was disabled) + # must not drain other consumers' ledger entries. + if kind in _LEDGER: + _LEDGER[kind] = max(0, _LEDGER[kind] - n) + if _LEDGER[kind] == 0: + del _LEDGER[kind] + + +def reset_for_tests() -> None: + with _LOCK: + _LEDGER.clear() + _EVICTABLES.clear() + + +def set_spill_reserve_policy( + floor_gib: Optional[float] = None, pct: Optional[float] = None +) -> None: + global _SPILL_RESERVE_FLOOR_GIB_OVERRIDE, _SPILL_RESERVE_PCT_OVERRIDE + if floor_gib is not None: + _SPILL_RESERVE_FLOOR_GIB_OVERRIDE = max(0.0, float(floor_gib)) + if pct is not None: + _SPILL_RESERVE_PCT_OVERRIDE = max(0.0, float(pct)) + + +def _spill_reserve_floor_gib() -> float: + if _SPILL_RESERVE_FLOOR_GIB_OVERRIDE is not None: + return _SPILL_RESERVE_FLOOR_GIB_OVERRIDE + for name, default in ( + ("AI_TOOLKIT_WDDM_SPILL_RESERVE_FLOOR_GIB", None), + ("AI_TOOLKIT_WDDM_SPILL_RESERVE_GIB", "1.0"), + ): + raw = os.environ.get(name) + if raw is None: + if default is None: + continue + raw = default + try: + return max(0.0, float(raw)) + except (TypeError, ValueError): + continue + return 2.0 + + +def _spill_reserve_pct() -> float: + if _SPILL_RESERVE_PCT_OVERRIDE is not None: + return _SPILL_RESERVE_PCT_OVERRIDE + try: + return max(0.0, float(os.environ.get("AI_TOOLKIT_WDDM_SPILL_RESERVE_PCT", "0.10"))) + except (TypeError, ValueError): + return 0.10 + + +def dxgi_spill_reserve_bytes(budget_bytes: Optional[int] = None) -> int: + floor_bytes = int(_spill_reserve_floor_gib() * GIB) + if budget_bytes and int(budget_bytes) > 0: + return max(floor_bytes, int(_spill_reserve_pct() * float(budget_bytes))) + return floor_bytes + + +def _spill_reserve_for_kind(kind: str, budget_bytes: Optional[int]) -> int: + # Weight-tier pins are one-shot STATIC commitments sized at attach/enable + # and fail-closed (strict mode raises rather than degrade): the per-tensor + # weight pins, the ingraph packs, and the pinned arena (kind="weights"). + # The pct-based reserve exists as slack for the *dynamic* streaming + # consumer (the bounce pool), which grows at runtime -- a static commitment + # does not need it and keeps only the floor. The all-28 ingraph proof ran + # at 14.07/15.13 GiB committed, which the pct reserve would have refused; + # the pinned arena feeding the same all-streamed trunk (Phase 3 Slice B) is + # the identical commitment and must get the same floor, or a full-model + # training arena loses its last block to the pct reserve -> non_pinned_pack. + if kind in _WEIGHT_TIER_KINDS: + return int(_spill_reserve_floor_gib() * GIB) + return dxgi_spill_reserve_bytes(budget_bytes) + + +def get_dxgi_meminfo(): + if _env_bool("AI_TOOLKIT_WDDM_DXGI_DISABLE"): + return None + global dxgi_meminfo + if dxgi_meminfo is None: + try: + from . import dxgi_meminfo as _dxgi_meminfo + except Exception: + return None + dxgi_meminfo = _dxgi_meminfo + return dxgi_meminfo + + +def dxgi_pinned_headroom( + cuda_device_index: Optional[int] = None, *, kind: str = "unknown" +) -> Optional[int]: + if _env_bool("AI_TOOLKIT_WDDM_DXGI_CONTROL_DISABLE"): + return None + dxgi = get_dxgi_meminfo() + if dxgi is None: + return None + info = dxgi.query_non_local_video_memory_info( + cuda_device_index=0 if cuda_device_index is None else int(cuda_device_index), + min_interval_s=0.0, + ) + if info is None: + return None + return dxgi.compute_non_local_headroom_bytes( + info.budget_bytes, + info.current_usage_bytes, + _spill_reserve_for_kind(kind, info.budget_bytes), + ) + + +def pinned_bytes_headroom( + cuda_device_index: Optional[int] = None, *, kind: str = "unknown" +) -> Optional[int]: + headroom = dxgi_pinned_headroom(cuda_device_index, kind=kind) + if headroom is not None: + return headroom + if _psutil is None: + return None + try: + total = _psutil.virtual_memory().total + except Exception: + return None + try: + fraction = float(os.environ.get("AI_TOOLKIT_PINNED_WEIGHT_WDDM_FRACTION", "0.25")) + except (TypeError, ValueError): + fraction = 0.25 + if fraction <= 0: + return None + return max(0, int(total * fraction) - total_pinned_bytes()) + + +def set_host_cache_reserve_bytes(nbytes: Optional[int]) -> None: + global _HOST_CACHE_RESERVE_BYTES_OVERRIDE + _HOST_CACHE_RESERVE_BYTES_OVERRIDE = None if nbytes is None else max(0, int(nbytes)) + + +def host_cache_reserve_bytes(mode: Optional[str] = None) -> int: + if _HOST_CACHE_RESERVE_BYTES_OVERRIDE is not None: + return _HOST_CACHE_RESERVE_BYTES_OVERRIDE + if mode == "sampling": + default = "0.0" + else: + default = os.environ.get("AI_TOOLKIT_PIN_HOST_CACHE_RESERVE_GIB", "1.0") + try: + return int(max(0.0, float(default)) * GIB) + except (TypeError, ValueError): + return 0 + + +def available_for_pin( + *, + kind: str = "unknown", + nbytes: int = 0, + device=None, + reserve_bytes: int = 0, + mode: Optional[str] = None, +) -> Optional[int]: + reserve = max(0, int(reserve_bytes or 0)) + host_cache_reserve_bytes(mode) + headroom = pinned_bytes_headroom(_cuda_device_index(device), kind=kind) + if headroom is None: + return None + return max(0, int(headroom) - reserve) + + +def _empty_host_pin_cache() -> None: + for name in ("_host_emptyCache", "_accelerator_emptyHostCache"): + try: + fn = getattr(torch._C, name, None) + if fn is not None: + fn() + except Exception: + pass + + +def register_evictable(shrink: Callable[[int], int]) -> None: + with _LOCK: + if shrink not in _EVICTABLES: + _EVICTABLES.append(shrink) + + +def unregister_evictable(shrink: Callable[[int], int]) -> None: + with _LOCK: + if shrink in _EVICTABLES: + _EVICTABLES.remove(shrink) + + +def reconcile(required_bytes: int = 0, *, device=None, allow_shrink: bool = True) -> int: + """Escalation before failing a pin request: empty the torch host-pin cache, + then (for non-weight-tier requests) ask evictable consumers to shrink. + + ``allow_shrink=False`` is the priority guard: weight-tier requests may not + evict the bounce pool -- eviction runs in reverse priority order, and + weights are already the lowest tier.""" + _empty_host_pin_cache() + if not allow_shrink: + return 0 + freed = 0 + need = max(0, int(required_bytes or 0)) + with _LOCK: + evictables = list(_EVICTABLES) + for shrink in evictables: + try: + freed += max(0, int(shrink(max(0, need - freed)))) + except Exception: + pass + if need and freed >= need: + break + return freed + + +_REGISTERED_HOST_PIN_LOCK = threading.Lock() +_REGISTERED_HOST_PINS: dict[int, tuple[int, str]] = {} + + +def pin_tensor_in_place(t: torch.Tensor, kind: str = "weights", *, device=None) -> bool: + """Pin an existing CPU tensor storage with cudaHostRegister when possible.""" + if not isinstance(t, torch.Tensor) or t.device.type != "cpu" or t.is_pinned(): + return False + size = int(t.numel() * t.element_size()) + if size <= 0 or not torch.cuda.is_available(): + return False + available = available_for_pin(kind=kind, nbytes=size, device=device) + if available is not None and size > available: + reconcile( + size - available, + device=device, + allow_shrink=kind not in _WEIGHT_TIER_KINDS, + ) + available = available_for_pin(kind=kind, nbytes=size, device=device) + if available is not None and size > available: + return False + try: + from torch.cuda import _pin_memory_utils as pin_memory_utils + ptr = int(t.data_ptr()) + if ptr == 0: + return False + pin_memory_utils.pin_memory(ptr, size) + except Exception: + return False + with _REGISTERED_HOST_PIN_LOCK: + _REGISTERED_HOST_PINS[int(t.data_ptr())] = (size, str(kind or "unknown")) + register_pinned_bytes(size, kind) + return True + + +def is_host_pinned(t: torch.Tensor) -> bool: + """True if this tensor's storage is usable as pinned host memory. + + torch's ``is_pinned()`` only recognizes buffers allocated by its own + caching host allocator; memory pinned in place with cudaHostRegister + (``pin_tensor_in_place`` / ``pin_register`` -- the weight/arena tier) + reports ``is_pinned() == False`` even though CUDA treats it as pinned for + transfer purposes. Consult the registration table too so consumers that + gate on "is this flat pinned" (e.g. borrowed ingraph packs) see registered + arena flats as pinned rather than falsely rejecting them as pageable. + """ + if not isinstance(t, torch.Tensor): + return False + try: + if t.is_pinned(): + return True + except Exception: + pass + with _REGISTERED_HOST_PIN_LOCK: + return int(t.data_ptr()) in _REGISTERED_HOST_PINS + + +def unpin_tensor_in_place(t: torch.Tensor, kind: Optional[str] = None) -> bool: + if not isinstance(t, torch.Tensor): + return False + ptr = int(t.data_ptr()) + with _REGISTERED_HOST_PIN_LOCK: + entry = _REGISTERED_HOST_PINS.pop(ptr, None) + if entry is None: + return False + size, registered_kind = entry + try: + from torch.cuda import _pin_memory_utils as pin_memory_utils + pin_memory_utils.unpin_memory(ptr) + except Exception: + with _REGISTERED_HOST_PIN_LOCK: + _REGISTERED_HOST_PINS[ptr] = (size, registered_kind) + return False + release_pinned_bytes(size, kind or registered_kind) + return True + + +# Storage-base data_ptrs of pinned arena flats. A streamed leaf is a VIEW into +# one of these flats at an offset, so its OWN data_ptr misses the exact-ptr +# _REGISTERED_HOST_PINS table (which keys on the registered flat ptr, not the +# view). But every such view shares the flat's untyped storage, whose base ptr +# is recorded here. The eager/pre-compile streaming pinned-bypass +# (manager_modules._profile_is_pinned / bounce_pool._is_pinned) consults this so +# register-pinned arena views are recognized as pinned and skip bounce staging. +# O(1) storage-base lookup, NOT the rejected per-forward range scan. Refcounted +# so a rebuild that recycles the same storage base ptr (release old flat, alloc +# new) never leaves a transient gap. +_ARENA_BACKED_STORAGE_LOCK = threading.Lock() +_ARENA_BACKED_STORAGE_PTRS: dict[int, int] = {} + + +def _storage_base_ptr(t: torch.Tensor) -> Optional[int]: + if not isinstance(t, torch.Tensor): + return None + try: + if t.device.type != "cpu": + return None + ptr = int(t.untyped_storage().data_ptr()) + except Exception: + return None + return ptr if ptr != 0 else None + + +def register_arena_storage(t: torch.Tensor) -> None: + """Mark a pinned arena flat's storage so views into it read as pinned.""" + ptr = _storage_base_ptr(t) + if ptr is None: + return + with _ARENA_BACKED_STORAGE_LOCK: + _ARENA_BACKED_STORAGE_PTRS[ptr] = _ARENA_BACKED_STORAGE_PTRS.get(ptr, 0) + 1 + + +def unregister_arena_storage(t: torch.Tensor) -> None: + ptr = _storage_base_ptr(t) + if ptr is None: + return + with _ARENA_BACKED_STORAGE_LOCK: + count = _ARENA_BACKED_STORAGE_PTRS.get(ptr) + if count is None: + return + if count <= 1: + _ARENA_BACKED_STORAGE_PTRS.pop(ptr, None) + else: + _ARENA_BACKED_STORAGE_PTRS[ptr] = count - 1 + + +def is_arena_backed(t: torch.Tensor) -> bool: + """True if this CPU tensor is a view into a pinned arena flat.""" + ptr = _storage_base_ptr(t) + if ptr is None: + return False + with _ARENA_BACKED_STORAGE_LOCK: + return ptr in _ARENA_BACKED_STORAGE_PTRS + + +def release(handle: PinHandle) -> None: + """Return a pin_alloc grant to the ledger. + + Accounting is explicit and deterministic: every pinned PinHandle must be + released exactly once by its owner (a GC finalizer on the handle's tensor + was tried and rejected -- callers immediately re-view the tensor, so the + original Python object dies while the pinned storage lives on).""" + if handle is None or not getattr(handle, "pinned", False): + return + if getattr(handle, "mechanism", "alloc") == "register": + # unpin_tensor_in_place does the ledger release itself (and the + # cudaHostUnregister returns DXGI budget immediately). + unpin_tensor_in_place(handle.tensor, getattr(handle, "kind", None)) + else: + release_pinned_bytes(int(handle.nbytes), getattr(handle, "kind", "unknown")) + handle.pinned = False + handle.nbytes = 0 + +def pin_alloc( + nbytes: int, + kind: str, + *, + device=None, + required: bool = True, + reserve_bytes: int = 0, + mode: Optional[str] = None, +) -> PinHandle: + nbytes = int(nbytes) + kind = str(kind or "unknown") + if nbytes <= 0: + tensor = torch.empty(max(0, nbytes), dtype=torch.uint8) + return PinHandle(tensor=tensor, nbytes=0, kind=kind, pinned=False) + if not torch.cuda.is_available(): + tensor = torch.empty(nbytes, dtype=torch.uint8) + return PinHandle(tensor=tensor, nbytes=nbytes, kind=kind, pinned=False) + + allow_shrink = kind not in _WEIGHT_TIER_KINDS + available = available_for_pin( + kind=kind, nbytes=nbytes, device=device, reserve_bytes=reserve_bytes, mode=mode + ) + if available is not None and nbytes > available: + reconcile(nbytes - available, device=device, allow_shrink=allow_shrink) + available = available_for_pin( + kind=kind, nbytes=nbytes, device=device, reserve_bytes=reserve_bytes, mode=mode + ) + if available is not None and nbytes > available: + if not required: + print( + f"[PinManager] pin refused ({kind}): {nbytes / GIB:.2f} GiB " + f"> {available / GIB:.2f} GiB available (ledger/DXGI budget); " + "returning pageable" + ) + tensor = torch.empty(nbytes, dtype=torch.uint8) + return PinHandle(tensor=tensor, nbytes=nbytes, kind=kind, pinned=False) + raise PinBudgetExceeded(_budget_message(kind, nbytes, available, device=device)) + try: + tensor = torch.empty(nbytes, dtype=torch.uint8, pin_memory=True) + except RuntimeError: + reconcile(nbytes, device=device, allow_shrink=allow_shrink) + try: + tensor = torch.empty(nbytes, dtype=torch.uint8, pin_memory=True) + except RuntimeError as error: + if not required: + print( + f"[PinManager] pin failed at OS level ({kind}): " + f"{nbytes / GIB:.2f} GiB cudaHostAlloc/pin raised " + f"'{error}' (ledger said {'' if available is None else f'{available / GIB:.2f} GiB '}" + "available -- likely system RAM pressure); returning pageable" + ) + tensor = torch.empty(nbytes, dtype=torch.uint8) + return PinHandle(tensor=tensor, nbytes=nbytes, kind=kind, pinned=False) + raise PinBudgetExceeded(_budget_message(kind, nbytes, available, device=device)) from error + register_pinned_bytes(nbytes, kind) + return PinHandle(tensor=tensor, nbytes=nbytes, kind=kind, pinned=True) + + +def pin_register_prepare(nbytes: int) -> tuple[torch.Tensor, int]: + """Allocate the page-aligned pageable buffer a register-mechanism pin + will need, WITHOUT pinning it yet. + + Split out of ``pin_register`` so callers who need to populate the buffer + (e.g. copying leaf tensors into a block flat) can do so on ordinary + pageable memory -- a plain memcpy that faults in pages at normal RAM + bandwidth -- before ``cudaHostRegister`` runs. Registering a virgin, + never-touched buffer forces the OS to commit+pin every page during the + syscall itself, which is measurably slower than registering pages that + are already resident (I1, ~1-1.5s per full arena build). + + Returns ``(candidate, padded_nbytes)``; ``candidate`` is an untouched + pageable view, exactly the layout ``pin_register`` used to build inline. + """ + nbytes = int(nbytes) + if nbytes <= 0: + return torch.empty(0, dtype=torch.uint8), 0 + page = 4096 + padded = (nbytes + page - 1) // page * page + # cudaHostRegister works at PAGE granularity: it registers every 4096-byte + # page the range touches. Two buffers that share a page (small buffers + # from the same allocator arena, or the boundary page between adjacent + # mallocs) collide -- registering the second raises CUDA "resource already + # mapped" (the 763bb75 root cause). Guarantee the registered range's pages + # are exclusive to THIS allocation: over-allocate with a full slack page on + # each side and register only the page-aligned interior, so no neighbor + # allocation can own a page we register. The slack bases stay alive with + # the returned tensor (the view keeps the base storage referenced). + base = torch.empty(padded + 3 * page, dtype=torch.uint8) + base_ptr = base.data_ptr() + # First page boundary at least one full page into the allocation. + aligned_start = ((base_ptr + page + page - 1) // page) * page + offset = aligned_start - base_ptr + candidate = base[offset:offset + padded] + return candidate, padded + + +def pin_register_commit( + candidate: torch.Tensor, + nbytes: int, + kind: str, + *, + device=None, + required: bool = False, +) -> PinHandle: + """Pin an already-prepared (and optionally already-populated) buffer + from :func:`pin_register_prepare` with cudaHostRegister.""" + nbytes = int(nbytes) + kind = str(kind or "unknown") + if nbytes <= 0 or not torch.cuda.is_available(): + return PinHandle(tensor=candidate, nbytes=0, kind=kind, + pinned=False, mechanism="register") + padded = candidate.numel() * candidate.element_size() + # pin_tensor_in_place budget-checks (with reconcile) and does the ledger + # accounting (of the padded size -- the true DXGI cost). + if pin_tensor_in_place(candidate, kind, device=device): + return PinHandle(tensor=candidate, nbytes=padded, kind=kind, + pinned=True, mechanism="register") + if required: + raise PinBudgetExceeded( + _budget_message( + kind, nbytes, + available_for_pin(kind=kind, nbytes=nbytes, device=device), + device=device, + ) + ) + print( + f"[PinManager] pin refused ({kind}): {nbytes / GIB:.2f} GiB " + "cudaHostRegister denied (ledger/DXGI budget); returning pageable" + ) + return PinHandle(tensor=candidate, nbytes=nbytes, kind=kind, pinned=False, + mechanism="register") + + +def pin_register( + nbytes: int, + kind: str, + *, + device=None, + required: bool = False, +) -> PinHandle: + """Exact-size host buffer pinned with cudaHostRegister. + + Unlike pin_alloc, this never touches torch's caching host allocator, so + the DXGI shared-budget cost is exactly ``nbytes`` (the caching allocator + rounds up to power-of-two buckets: observed live, 8.86 GiB of pin_alloc + flats committed 12.70 GiB of DXGI usage -- ~40% invisible overhead) and + release returns the budget immediately. Intended for large long-lived + buffers (the pinned weight arena); small/churny consumers should keep + using pin_alloc. + + Convenience wrapper over :func:`pin_register_prepare` + + :func:`pin_register_commit` for callers with no data to populate before + pinning (e.g. tests). Callers that populate a leaf-carrying flat should + call the two steps directly with the copy in between (see + ``ingraph_stream.pack_block_host``). + """ + nbytes = int(nbytes) + kind = str(kind or "unknown") + if nbytes <= 0 or not torch.cuda.is_available(): + tensor = torch.empty(max(0, nbytes), dtype=torch.uint8) + return PinHandle(tensor=tensor, nbytes=max(0, nbytes) if nbytes > 0 else 0, + kind=kind, pinned=False, mechanism="register") + candidate, _padded = pin_register_prepare(nbytes) + return pin_register_commit(candidate, nbytes, kind, device=device, required=required) + + +def pin_empty(shape, dtype, kind: str, *, device=None, required: bool = False): + element_size = torch.empty((), dtype=dtype).element_size() + n = 1 + for dim in tuple(shape): + n *= int(dim) + handle = pin_alloc(n * element_size, kind, device=device, required=required) + return handle.tensor.view(dtype).reshape(tuple(shape)), handle.pinned + + +def can_pin( + nbytes: int, + *, + kind: str = "unknown", + device=None, + reserve_bytes: int = 0, + mode: Optional[str] = None, +) -> bool: + available = available_for_pin( + kind=kind, + nbytes=nbytes, + device=device, + reserve_bytes=reserve_bytes, + mode=mode, + ) + return available is None or int(nbytes) <= available + + + +def plan_budgets( + *, + offloaded_weight_bytes: int, + requested_bounce_bytes: int, + device=None, + mode: str = "training", +) -> dict: + """Return pin budgets using the fixed priority policy from PIN_MANAGER_PLAN.""" + weights = max(0, int(offloaded_weight_bytes or 0)) + requested_bounce = max(0, int(requested_bounce_bytes or 0)) + reserve = host_cache_reserve_bytes(mode) + headroom = pinned_bytes_headroom(_cuda_device_index(device)) + if headroom is None: + # No authoritative probe: keep the old requested shape, but still report + # the reserve so diagnostics show the implicit consumer exists. + return { + "mode": mode, + "strategy": "fallback_no_probe", + "headroom_bytes": None, + "reserve_bytes": reserve, + "bounce_budget_bytes": requested_bounce, + "weight_budget_bytes": weights, + } + usable = max(0, int(headroom) - reserve) + if weights and weights <= usable: + return { + "mode": mode, + "strategy": "full_pin_no_bounce", + "headroom_bytes": int(headroom), + "reserve_bytes": reserve, + "bounce_budget_bytes": 0, + "weight_budget_bytes": weights, + } + bounce_budget = min(requested_bounce, usable) + weight_budget = max(0, usable - bounce_budget) + return { + "mode": mode, + "strategy": "partial_bounce_first", + "headroom_bytes": int(headroom), + "reserve_bytes": reserve, + "bounce_budget_bytes": bounce_budget, + "weight_budget_bytes": min(weights, weight_budget), + } + +def snapshot(device=None, mode: Optional[str] = None) -> dict: + headroom = pinned_bytes_headroom(_cuda_device_index(device)) + reserve = host_cache_reserve_bytes(mode) + by_kind = pinned_bytes_by_kind() + total = sum(by_kind.values()) + return { + "pinned_total_bytes": total, + "pinned_by_kind": by_kind, + "headroom_bytes": headroom, + "host_cache_reserve_bytes": reserve, + "available_bytes": None if headroom is None else max(0, int(headroom) - reserve), + } + + +def format_snapshot(device=None, mode: Optional[str] = None) -> str: + snap = snapshot(device=device, mode=mode) + by_kind = " ".join( + f"{key}={value / GIB:.2f}GiB" + for key, value in sorted(snap["pinned_by_kind"].items()) + ) + headroom = snap["headroom_bytes"] + avail = snap["available_bytes"] + return ( + "pin ledger: " + f"total={snap['pinned_total_bytes'] / GIB:.2f}GiB " + f"headroom={'n/a' if headroom is None else f'{headroom / GIB:.2f}GiB'} " + f"reserve={snap['host_cache_reserve_bytes'] / GIB:.2f}GiB " + f"free={'n/a' if avail is None else f'{avail / GIB:.2f}GiB'} " + f"{by_kind}" + ).strip() + + +def _budget_message(kind: str, nbytes: int, available: Optional[int], *, device=None) -> str: + snap = snapshot(device=device) + return ( + f"pinned host memory budget exceeded for {kind}: " + f"request={nbytes / GIB:.2f} GiB " + f"available={'unknown' if available is None else f'{available / GIB:.2f} GiB'} " + f"ledger={snap['pinned_by_kind']}" + ) diff --git a/toolkit/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py new file mode 100644 index 0000000000..b58f2e95a4 --- /dev/null +++ b/toolkit/memory_management/vram_budget.py @@ -0,0 +1,988 @@ +"""Dedicated-VRAM (WDDM) budget arithmetic: one source of truth, one name per meaning. + +This module owns the *dedicated-cliff* quantities. Windows/WDDM has two distinct +memory cliffs with different failure modes: + +* **Dedicated ceiling** (this module): crossing the card's physical VRAM makes + WDDM silently page GPU memory to system RAM -- catastrophic slowdown, no + error. Governed by ``device_free_bytes`` below -- NOT by + ``torch.cuda.mem_get_info``, which does *not* see other processes (see the + note on the free signal). Quantities: ``hard_gib`` (the never-cross + device-free floor, default 1.0), ``margin_gib`` (the planning headroom, + >= hard). +* **Shared / DXGI NON_LOCAL budget** (NOT this module): pinned host memory + commits against it and exhausting it is a hard cudaErrorMemoryAllocation. + That reserve lives in ``pin_manager`` / ``bounce_pool`` + (``dxgi_spill_reserve_bytes``); do not conflate the two. + +The free signal (do not regress this) +------------------------------------- +``torch.cuda.mem_get_info`` free is NOT physical availability. It is what the +driver will promise *this process*; on WDDM that promise is backed by paging +other processes out. Measured on the 4070 with a second process holding 9 GiB: + + NVML physical free : 0.55 GiB <- the truth + torch.cuda.mem_get_info free : 4.70 GiB <- over-reports ~8x + DXGI LOCAL Budget - CurrentUsage: 4.58 GiB <- also over-reports + +Planning against the optimistic number is precisely how a run silently crosses +the dedicated cliff (a 73 s/it job ran at 304 s/it with an orphan process on the +card, and no allocator retry fired). So ``free`` here means NVML physical free +whenever NVML is available -- see ``device_free_bytes``. DXGI stays where it +belongs: the *shared* NON_LOCAL pin budget, not the dedicated cliff. + +Vocabulary (each quantity has exactly one name): + +* ``total`` / ``free`` -- physical card bytes, via ``device_free_bytes`` + (NVML-backed; sees every process on the GPU). +* ``torch_reserved`` / ``torch_allocated`` -- torch's caching allocator. +* ``non_torch`` -- ``(total - free) - torch_reserved``: CUDA context, VAE/TE, + the Windows desktop, other processes. Torch is never the card's only tenant. +* ``hard`` -- device-free floor the card must keep (WDDM spill guard). +* ``margin`` -- planning headroom subtracted from budgets; ``>= hard``. + +Everything here is pure (CPU-testable) except ``DeviceSnapshot.capture`` and +``device_free_bytes``. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import Optional + +import torch + +from . import nvml_meminfo + +GIB = 1024 ** 3 + + +def _env(name: str, default: str) -> str: + value = os.environ.get(name) + return default if value is None or value == "" else value + + +def reconcile_free_bytes(driver_free_bytes, physical_free_bytes) -> int: + """Pick the governing device-free value (pure; CPU-testable). + + ``physical_free_bytes`` is NVML's device-wide free (None when NVML is + unavailable). We take the MIN of the two signals rather than trusting NVML + blindly: the driver number is a real constraint too, and whichever sensor is + more pessimistic is the one that keeps us off the cliff. In the uncontended + single-process case the two agree, so this does not cost residency. + """ + driver = max(0, int(driver_free_bytes)) + if physical_free_bytes is None: + return driver + return min(driver, max(0, int(physical_free_bytes))) + + +# -------------------------------------------------------------------------- +# Simulated smaller card (validation knob) +# -------------------------------------------------------------------------- +# Pretend the GPU is smaller than it is, so an 8 GB / 6 GB card's residency, +# streaming and cap behaviour can be exercised on a 12 GB one. Implemented as a +# *phantom ballast*: a fixed number of bytes subtracted from BOTH total and free +# everywhere this module reports them. That keeps every derived quantity +# self-consistent (non_torch, margins, promotion checks, the allocator cap) with +# no per-call-site special cases -- an unfittable model is unfittable for the +# same arithmetic reason it would be on the real small card. +# +# It does not shrink the physical card: the allocator cap is what actually makes +# an over-plan fail, and `_apply_wddm_hard_allocator_cap` converts the simulated +# cap bytes back into a fraction of the REAL total before handing it to torch. +_SIMULATED_CARD_BYTES: int | None = None + + +def set_simulated_card_bytes(total_bytes) -> None: + """Pretend the card has ``total_bytes`` of VRAM (``None`` disables).""" + global _SIMULATED_CARD_BYTES + if total_bytes is None: + _SIMULATED_CARD_BYTES = None + return + value = int(total_bytes) + if value <= 0: + raise ValueError(f"simulated card size must be positive, got {total_bytes}") + _SIMULATED_CARD_BYTES = value + + +def simulated_card_bytes() -> int | None: + return _SIMULATED_CARD_BYTES + + +def apply_simulated_card(simulated_vram_gib, *, device=None) -> int | None: + """Install (or clear) the simulated card size; logs once when active. + + ``simulated_vram_gib`` of ``None``/``0`` clears the simulation. A size at or + above the real card is a no-op with a warning -- the ballast can only hide + VRAM, never invent it. + """ + if not simulated_vram_gib: + set_simulated_card_bytes(None) + return None + + wanted = int(float(simulated_vram_gib) * GIB) + if device is None or not torch.cuda.is_available(): + set_simulated_card_bytes(wanted) + return wanted + + real = real_device_total_bytes(device) + if wanted >= real: + print( + "[MemoryManager] simulated VRAM " + f"{wanted / GIB:.2f} GiB >= real card {real / GIB:.2f} GiB; " + "ignoring (a simulation can only shrink the card)" + ) + set_simulated_card_bytes(None) + return None + + set_simulated_card_bytes(wanted) + print( + "[MemoryManager] SIMULATED CARD: reporting " + f"{wanted / GIB:.2f} GiB total (real {real / GIB:.2f} GiB); " + f"{(real - wanted) / GIB:.2f} GiB is hidden from both total and free, " + "and the allocator cap enforces it. Planning, residency and OOMs now " + "behave as they would on the smaller card." + ) + return wanted + + +def real_device_total_bytes(device) -> int: + """Physical card size, ignoring any simulation.""" + return int(torch.cuda.mem_get_info(device)[1]) + + +def simulated_ballast_bytes(device) -> int: + """Bytes hidden from total AND free to fake a smaller card (0 when off).""" + if _SIMULATED_CARD_BYTES is None: + return 0 + return max(0, real_device_total_bytes(device) - _SIMULATED_CARD_BYTES) + + +def device_total_bytes(device) -> int: + """Governing card size: the simulated one when a simulation is active.""" + return max(0, real_device_total_bytes(device) - simulated_ballast_bytes(device)) + + +def device_free_bytes(device) -> int: + """Governing physical free bytes on ``device``, across ALL processes. + + Prefer this over ``torch.cuda.mem_get_info(...)[0]`` anywhere a decision is + made (residency, promotion, reserve, cap). ``mem_get_info`` alone is blind to + other GPU tenants and will happily plan a run onto memory that is not there. + Delta/probe loops that measure this process's *own* allocation deltas may + keep using the raw driver number -- they are measuring differences, not + availability. + """ + driver_free = int(torch.cuda.mem_get_info(device)[0]) + index = torch.device(device).index + if index is None: + index = torch.cuda.current_device() + free = reconcile_free_bytes( + driver_free, nvml_meminfo.physical_free_bytes(index) + ) + return max(0, free - simulated_ballast_bytes(device)) + + +# A quiet box still has non-torch bytes on the card: the CUDA context, cuDNN / +# cuBLAS workspaces, and the Windows desktop compositor. Measured ~1.15 GiB on +# the 4070 with a bare context, so treat everything under this as the cost of +# doing business rather than "contention". +FOREIGN_NOMINAL_BYTES = int(1.5 * GIB) +# Do not cry about a browser tab. Only excess above this is worth a line in the +# log; below it, the residency we lose is not what makes or breaks a run. +FOREIGN_WARN_FLOOR_BYTES = int(1.0 * GIB) +# Above this, a foreign tenant is doing real damage even if the model could +# never have been fully resident anyway: every GiB it holds is a GiB we stream +# from the CPU on every single step. Big models (Krea2 on a 12 GB card) never +# fit fully, so "would it have fit?" must NOT be the thing that decides whether +# we raise our voice -- lost residency is. +FOREIGN_SEVERE_BYTES = int(2.0 * GIB) + + +@dataclass(frozen=True) +class ForeignVramReport: + """Is someone ELSE on this GPU, and is that why we cannot fit? + + ``foreign`` is every byte on the card that is not our torch allocator -- + other processes, plus our own context/desktop baseline. ``excess`` is what + remains after allowing a nominal baseline, i.e. the part plausibly caused by + another tenant (a game, a ComfyUI server, an orphaned training job). + """ + + foreign_bytes: int + excess_bytes: int + want_bytes: int + have_bytes: int + fits_now: bool + fits_without_excess: bool + severity: str # "none" | "benign" | "contributing" | "blocking" + + @property + def is_blocking(self) -> bool: + """The model would have fit; another tenant is the reason it does not.""" + return self.severity == "blocking" + + def format(self) -> str: + return ( + f"foreign={self.foreign_bytes / GIB:.2f} GiB " + f"(excess={self.excess_bytes / GIB:.2f} GiB) " + f"want={self.want_bytes / GIB:.2f} GiB " + f"have={self.have_bytes / GIB:.2f} GiB " + f"severity={self.severity}" + ) + + +def assess_foreign_vram( + *, + total_bytes, + free_bytes, + torch_reserved_bytes, + want_bytes, + have_bytes, + nominal_bytes: int = FOREIGN_NOMINAL_BYTES, + warn_floor_bytes: int = FOREIGN_WARN_FLOOR_BYTES, +) -> ForeignVramReport: + """Classify external VRAM pressure at a phase boundary (pure, CPU-testable). + + ``want_bytes`` -- bytes to hold the model fully resident. + ``have_bytes`` -- residency budget we actually got. + ``free_bytes`` -- must be the NVML-backed physical free (``device_free_bytes``); + with ``mem_get_info`` the foreign bytes are invisible, which + is the whole reason this check exists. + + The verdict is a counterfactual, not a threshold on usage alone: contention + only earns a warning when it *changed the outcome*. + + blocking -- we are streaming, but would have fit without the excess. + Someone else's memory is costing us real throughput. + contributing -- the excess costs residency, but the model would not have + fit anyway. Worth a note; not the root cause. + benign -- another tenant is present, but we fit regardless. + none -- nothing meaningful beyond our own baseline. + """ + used = max(0, int(total_bytes) - int(free_bytes)) + foreign = max(0, used - max(0, int(torch_reserved_bytes))) + excess = max(0, foreign - max(0, int(nominal_bytes))) + + want = max(0, int(want_bytes)) + have = max(0, int(have_bytes)) + fits_now = have >= want + fits_without_excess = (have + excess) >= want + + if excess < max(0, int(warn_floor_bytes)): + severity = "none" + elif fits_now: + severity = "benign" + elif fits_without_excess: + severity = "blocking" + else: + severity = "contributing" + + return ForeignVramReport( + foreign_bytes=foreign, + excess_bytes=excess, + want_bytes=want, + have_bytes=have, + fits_now=fits_now, + fits_without_excess=fits_without_excess, + severity=severity, + ) + + +def format_foreign_vram_warning(report: ForeignVramReport, *, phase: str) -> str | None: + """Operator-facing line for a phase boundary. None when there is nothing to say. + + Volume is set by HARM, not by the fit counterfactual: a model too big to ever + be fully resident (Krea2 on 12 GB) still loses a GiB of residency for every + GiB a foreign tenant holds, and streams it from the CPU every step. + """ + if report.severity in ("none", "benign"): + return None + + gib = GIB + excess = report.excess_bytes / gib + culprits = ( + "Check for another training job, a leftover/orphaned run of this same job, " + "a ComfyUI server, or a game." + ) + head = ( + f"{excess:.2f} GiB of this card is held outside this process " + f"({report.foreign_bytes / gib:.2f} GiB total non-torch), costing an equal " + f"amount of resident weights during {phase} -- those are streamed from the " + f"CPU on every step instead." + ) + + if report.severity == "blocking": + return ( + f"[MemoryManager] WARNING: VRAM contention is why {phase} is streaming. " + f"{head} The model needs {report.want_bytes / gib:.2f} GiB resident and " + f"only {report.have_bytes / gib:.2f} GiB is available -- it WOULD fit " + f"entirely if that memory were free. {culprits}" + ) + if report.excess_bytes >= FOREIGN_SEVERE_BYTES: + return ( + f"[MemoryManager] WARNING: heavy VRAM contention during {phase}. {head} " + f"(The model, {report.want_bytes / gib:.2f} GiB, would not be fully " + f"resident even on an idle card, but this is still costing real " + f"throughput.) {culprits}" + ) + return ( + f"[MemoryManager] note: {head} {culprits}" + ) + + +def device_mem_info(device) -> tuple[int, int]: + """Drop-in for ``torch.cuda.mem_get_info``: ``(free, total)`` in bytes. + + Identical shape to the torch call, but ``free`` is the NVML-backed physical + free (sees other processes) and ``total`` honours a simulated smaller card. + """ + return device_free_bytes(device), device_total_bytes(device) + + +@dataclass(frozen=True) +class DeviceSnapshot: + """Point-in-time physical view of a CUDA device (bytes).""" + + total: int + free: int + torch_reserved: int + torch_allocated: int + + @property + def used(self) -> int: + return max(0, self.total - self.free) + + @property + def non_torch(self) -> int: + """Device bytes held by anyone but torch's caching allocator.""" + return max(0, self.used - self.torch_reserved) + + @staticmethod + def capture(device) -> Optional["DeviceSnapshot"]: + if device is None or not torch.cuda.is_available(): + return None + dev = torch.device(device) + if dev.type != "cuda": + return None + return DeviceSnapshot( + total=device_total_bytes(dev), + # NVML-backed: sees other processes on the card. See the module + # docstring -- mem_get_info free would over-report here. + free=device_free_bytes(dev), + torch_reserved=int(torch.cuda.memory_reserved(dev)), + torch_allocated=int(torch.cuda.memory_allocated(dev)), + ) + + def format(self) -> str: + return ( + f"torch_allocated={self.torch_allocated / GIB:.2f} GiB " + f"torch_reserved={self.torch_reserved / GIB:.2f} GiB " + f"device_used={self.used / GIB:.2f}/{self.total / GIB:.2f} GiB " + f"device_free={self.free / GIB:.2f} GiB " + f"non_torch={self.non_torch / GIB:.2f} GiB" + ) + + +@dataclass(frozen=True) +class WddmMargins: + """Dedicated-cliff margins for one phase (training attach / sampling start). + + Resolve ONCE at the phase boundary and pass by value; do not re-read env + vars mid-phase (they cannot change mid-run, and re-reads hide which value + actually governed a decision). + """ + + hard_gib: float + margin_gib: float + source: str # "config" | "env" | "auto" + + @property + def hard_bytes(self) -> int: + return int(self.hard_gib * GIB) + + @property + def margin_bytes(self) -> int: + return int(self.margin_gib * GIB) + + def format(self) -> str: + return ( + f"wddm_hard={self.hard_gib:.2f} GiB " + f"wddm_margin={self.margin_gib:.2f} GiB ({self.source})" + ) + + +def auto_margin_gib(device, pct: float = 0.10, floor_gib: float = 1.0) -> float: + """Auto planning margin: max(floor, pct * card size).""" + try: + total_bytes = device_total_bytes(device) + except Exception: + total_bytes = 0 + total_gib = max(0.0, float(total_bytes) / GIB) + return max(float(floor_gib), float(pct) * total_gib) + + +def resolve_margins( + device, + margin_value, + hard_value, + *, + margin_env: str, + hard_env: str, +) -> WddmMargins: + """Resolve the phase's margins from config value > env > auto. + + ``margin_value`` / ``hard_value`` are the config-supplied values (``None`` + means "consult the env var"; a negative margin or "auto" means auto). + ``margin`` is clamped to at least ``hard``. + """ + hard_gib = float(_env(hard_env, "1.0")) if hard_value is None else float(hard_value) + raw = _env(margin_env, "-1.0") if margin_value is None else margin_value + source = "env" if margin_value is None else "config" + try: + margin_gib = float(raw) + auto = margin_gib < 0 + except (TypeError, ValueError): + auto = str(raw).strip().lower() == "auto" + margin_gib = -1.0 + if auto: + margin_gib = auto_margin_gib(device) + source = "auto" + return WddmMargins( + hard_gib=hard_gib, + margin_gib=max(margin_gib, hard_gib or 0.0), + source=source, + ) + + +def cap_fraction(total_bytes, free_bytes, reserved_bytes, hard_gib) -> float: + """Allocator-cap fraction so device_used stays <= total - hard (pure). + + The cap governs torch's own reserved pool, but torch is not the card's + only tenant (``non_torch``). Capping torch at ``total - hard`` alone lets + device_used reach ``total - hard + non_torch`` (observed: reserved 10.97 + + non_torch 1.02 = 11.99/11.99 GiB, device_free 0, silent WDDM paging). + Subtract the measured non-torch share so the whole device, not just torch, + keeps the hard margin free. + """ + total = float(max(1, int(total_bytes))) + non_torch = max(0.0, (total - float(free_bytes)) - float(reserved_bytes)) + cap_bytes = total - float(hard_gib) * GIB - non_torch + return max(0.1, min(1.0, cap_bytes / total)) + + +def sampling_allocator_budget_free_bytes( + total_bytes, + allocated_bytes, + cap_fraction, + hard_bytes, + *, + gc_threshold=0.95, +): + """Allocated-side equivalent of driver-free for the sampling planner (pure). + + Driver-free counts torch's idle cached segments as *used*, so a plan built + from it refuses residency that the allocator cap's GC would reclaim on + demand. The allocator-side capacity is governed by the gc target + (``gc_threshold * cap``): live allocations may safely grow to it, and the + planner's margin beyond the hard floor (which is already inside the cap) + stays free below the target as the fragmentation/allowance pad. + + Returned in the same units/meaning as ``mem_get_info`` free so the + downstream ``usable = free - working_reserve - margin`` keeps its shape: + + usable = threshold*cap - allocated - working_reserve - (margin - hard) + + Returns ``None`` when no cap fraction is known (non-Windows / cap not + applied); callers should then stay on the driver-free number. + """ + if cap_fraction is None: + return None + cap_bytes = float(cap_fraction) * float(max(1, int(total_bytes))) + return int( + float(gc_threshold) * cap_bytes + - float(max(0, int(allocated_bytes))) + + float(max(0, int(hard_bytes))) + ) + + +def sampling_guard_predicted_peak_free(total_b, free_b, reserved_b, peak_reserved_b): + """Predicted free VRAM at the next forward's peak (pure). + + ``non_torch = (total - free) - reserved`` plus the worst forward's reserved + high-water is what the next peak will occupy; the prediction is ``total`` + minus that. It shrinks one-for-one as external use grows -- which is the + cohabitation guard's trigger. Forward-only sampling never OOMs at the cliff + (it pages silently), so the guard watches this instead of an exception. + """ + other_b = max(0, (total_b - free_b) - reserved_b) + return total_b - (peak_reserved_b + other_b) + + +def training_cliff_predicted_peak_free_gib( + total_gib, device_free_gib, torch_reserved_gib, peak_allocated_gib +): + """Driver free expected when the next step rebuilds its live peak (pure). + + ``empty_cache`` can make step-end free look healthy by dropping idle cached + blocks, but the next forward/backward will recreate the live peak. Keep + non-allocator residents (``non_torch``) from the current snapshot and ask + whether peak allocated memory itself clears the WDDM hard floor. + """ + device_used_gib = max(0.0, total_gib - device_free_gib) + non_torch_gib = max(0.0, device_used_gib - torch_reserved_gib) + return total_gib - (max(0.0, peak_allocated_gib) + non_torch_gib) + + +def training_promotion_worst_shape_free_gib( + *, + resident_gib, + added_block_gib, + ring_gib, + worst_working_reserve_gib, + other_gib, + total_gib, +): + """Predicted device-free margin on the WORST measured resolution after + promoting one resident block (pure, CPU-testable). + + Residency is global -- a block promoted to resident stays resident for every + resolution bucket -- but the activation working set is not: the largest + measured resolution defines the tightest cohabitation peak. A promotion + decided on a roomy low-res step still has to leave room for the worst measured + resolution's working set, or that high-res step silently pages across the WDDM + dedicated cliff (which no allocator retry counter catches). So gate the + promotion on the worst measured working reserve, not the current step's free. + + ``worst_working_reserve_gib`` is the reserve the plan actually applies to every + shape (already the max across measured buckets in the training controller). + Conservative / from-below: assumes the promoted block adds its full size to the + resident baseline and the ring does not shrink to compensate. Returns the + predicted worst-case device-free GiB; the caller vetoes the promotion when it + would fall below the promote floor. + """ + predicted_used = ( + max(0.0, float(resident_gib)) + + max(0.0, float(added_block_gib)) + + max(0.0, float(ring_gib)) + + max(0.0, float(worst_working_reserve_gib)) + + max(0.0, float(other_gib)) + ) + return float(total_gib) - predicted_used + + +def training_eager_promote_blocks( + *, + resident_gib, + block_gib, + ring_gib, + worst_working_reserve_gib, + other_gib, + total_gib, + promote_floor_gib, + max_blocks, +): + """How many equal-sized blocks may be promoted at once while keeping the + predicted worst-shape free margin at or above ``promote_floor_gib`` (pure). + + This is the eager-fill counterpart of the one-block-at-a-time climb: a roomy + card leaves GiBs idle if residency only ever grows one block per cadence + window. The prediction is the same conservative worst-measured-resolution + model as ``training_promotion_worst_shape_free_gib`` -- the blocks are assumed + to add their full size and the ring is assumed not to shrink -- so the floor is + what the run actually keeps free on its tightest measured shape. Returns 0 when + not even one block fits, which the caller reports as a worst-shape veto. + """ + block = float(block_gib) + limit = int(max_blocks) + if block <= 0.0 or limit <= 0: + return 0 + free_now = training_promotion_worst_shape_free_gib( + resident_gib=resident_gib, + added_block_gib=0.0, + ring_gib=ring_gib, + worst_working_reserve_gib=worst_working_reserve_gib, + other_gib=other_gib, + total_gib=total_gib, + ) + room = free_now - float(promote_floor_gib) + if room < block: + return 0 + return min(limit, int(room // block)) + + +def sampling_step_should_trim(free_before_b, trim_margin_b) -> bool: + """Whether realized device-free warrants a per-step cache trim (pure). + + WDDM pages on the committed footprint silently, so the trigger is realized + free, not an allocated-side or peak signal. Trim (empty_cache) is cheap and + non-destructive, so the bar is just "free has dropped into the margin." + """ + return free_before_b < trim_margin_b + + +def sampling_step_should_demote(free_after_b, hard_floor_b) -> bool: + """Whether to escalate to a block demote after a trim (pure). + + Only when trimming left free still under the hard floor -- i.e. there was + no idle cache to reclaim, so the pressure is real (external) and the only + relief is giving back resident weights. + """ + return free_after_b < hard_floor_b + + +def estimate_sampling_working_reserve_bytes( + image_tokens: int, + text_tokens: int = 512, + *, + batch_cfg: bool = False, + fp8_native: bool = True, + base_bytes: int = int(2.2 * GIB), + per_token_bytes: int = 40 * 1024, + dequant_pad_bytes: int = int(1.4 * GIB), + safety: float = 1.15, + headroom_bytes: int = GIB, +) -> int: + """Cold-start estimate of the sampling working set (pure, CPU-testable). + + Used before any measured peak exists, so a high-resolution first sample + plans enough streaming up front instead of discovering the working set + via mid-denoise OOM demotions (every demote invalidates compiled state + and, under strict ingraph, changes the pack set). + + Linear-in-tokens model calibrated on Krea2 RTX 4070 smoke runs + (2026-07-07/08, fp8 + cutlass attention, sequential CFG, partial + residency; ``sampling_extra`` = peak allocated minus resident weights): + + 512x512 -> L = 1024 + 512 = 1536 tokens, extra ~= 2.2 GiB + 2000x2000-> L = 15625 + 512 = 16137 tokens, extra ~= 2.8 GiB + => per_token ~= 40 KiB, base ~= 2.2 GiB (streaming/dequant churn + + fixed workspaces dominate; per-token activations are small) + + 512x512 dequant fallback (fp8 sampling off) -> extra ~= 3.6 GiB + => dequant_pad ~= 1.4 GiB (torchao fp32 dequant transients) + + Batched CFG scales the token-dependent share by 2.5, not 2.0: besides the + batch-2 doubling, the fp32 intermediates (observed as (2, L, features) + fp32 allocations, 740 MiB each at 2000px) are underweighted in the + batch-1-calibrated per-token constant -- the x2.0 estimate ran ~1 GiB + short at 2000px (two demote rounds). ``safety`` covers model-to-model + variation, and + ``headroom_bytes`` (flat +1 GiB) deliberately overestimates: streaming + one extra block costs a little bandwidth, while underestimating costs a + mid-denoise demote -- which invalidates compiled state, mutates the + strict-ingraph pack set, and (observed at 2000px) can cascade into a + full streamed transition. The learned per-run reserve replaces this + estimate after the first measured sample. + """ + tokens = max(0, int(image_tokens)) + max(0, int(text_tokens)) + token_bytes = int(tokens * per_token_bytes * (2.5 if batch_cfg else 1.0)) + estimate = int(base_bytes) + token_bytes + if not fp8_native: + estimate += int(dequant_pad_bytes) + return int(estimate * float(safety)) + int(headroom_bytes) + + +def estimate_training_working_reserve_bytes( + image_tokens: int, + text_tokens: int = 512, + *, + base_bytes: int = int(5.2 * GIB), + per_token_bytes: int = 610 * 1024, + safety: float = 1.15, + headroom_bytes: int = GIB, +) -> int: + """Cold-start estimate of the training working set (pure, CPU-testable). + + Training's cold-start reserve was a flat constant (planner.py's + ``DEFAULT_AUTO_WORKING_RESERVE_GIB = 5.0``) with no resolution awareness at + all, unlike sampling's shape-aware ``estimate_sampling_working_reserve_bytes`` + above. At low resolution that flat reserve is generous; at high resolution + it is not enough, so the attach-time residency plan keeps too many blocks + resident, leaves activations too little headroom, and the run discovers the + shortfall only via a cold-start WDDM-cap-violation storm -- each violation + widens the allocator cap by a fixed ``WDDM_CAP_RELIEF_BYTES`` (0.5 GiB), so + a large resolution jump can cost several wasted/skipped steps before the + cap finally catches up (observed: Krea2 LoKr at 1024x1024 skipped 5/5 fake + steps under the flat default, never reaching a real step). + + Linear-in-tokens model calibrated on Krea2 LoKr RTX 4070 smoke runs + (2026-07-14, ``--block-stream-only`` so zero blocks are resident and + ``torch_max_allocated`` is purely the forward+backward+optimizer + footprint, uncontaminated by the residency split this estimate feeds): + + 512x512 -> 1024 tokens, torch_max_allocated ~= 5.76 GiB + 1024x1024 -> 4096 tokens, torch_max_allocated ~= 7.50 GiB + => per_token ~= 595 KiB, base ~= 5.17 GiB (rounded to 610 KiB / 5.2 GiB) + + Only two points, one adapter variant, one card -- weaker calibration than + the sampling estimator above. ``text_tokens`` is folded in at the same + per-token rate by symmetry with the sampling model; it was held constant + across both calibration runs, not independently measured. ``safety`` and + the flat ``headroom_bytes`` deliberately overestimate: streaming one extra + block is cheap, under-reserving costs the cap-violation storm above. + """ + tokens = max(0, int(image_tokens)) + max(0, int(text_tokens)) + estimate = int(base_bytes) + int(tokens * per_token_bytes) + return int(estimate * float(safety)) + int(headroom_bytes) + + +def sampling_overshoot_margin_bytes( + overshoot_gib: float = 0.86, + safety_gib: float = 0.375, + hard_bytes: int = 0, +) -> int: + """Auto sampling margin in the allocator-cap era (pure, CPU-testable). + + With the reclaim allocator cap guarding the WDDM cliff (a capped allocation + recycles cache or raises a loud OOM, it never silently pages), the sampling + margin's only remaining job is to cover the caching allocator's + reserved-over-allocated overshoot. Measured on Krea2 fp8 512px that overshoot + is ~0.86 GiB and rock-steady (std ~0.05 GiB across 8 seeds), so the auto + margin is that measured overshoot plus one safety block -- NOT the old + ``0.10 * card`` cushion, which was sized for a chaotic allocator that kept + jumping over the limit and no longer misbehaves. Narrowing it hands the + difference straight to resident weights (fewer streamed blocks). Floored at + the hard margin so it can never drop below the WDDM device-free floor. + """ + return int( + max(float(overshoot_gib) + float(safety_gib), float(max(0, int(hard_bytes))) / GIB) + * GIB + ) + + +def training_guard_pressure(dxgi: dict, physical: dict) -> dict: + """Combine the DXGI LOCAL and physical cliff signals (pure). + + Pressure if EITHER signal predicts the next step's peak crosses its floor. + The DXGI LOCAL budget is a per-process OS grant and its usage counter + excludes other processes, so it can bless a layout the physical + (mem_get_info) view already knows will overfill the card -- and vice versa + when the OS shrinks the budget early. The merged dict keeps the DXGI + fields at the top level (``source`` compatibility) and carries the + physical signal under ``physical_*``. + """ + merged = dict(dxgi) + merged["pressure"] = bool(dxgi.get("pressure")) or bool(physical.get("pressure")) + merged["physical_predicted_peak_free_gib"] = physical.get("predicted_peak_free_gib") + merged["physical_target_free_gib"] = physical.get("target_free_gib") + merged["pressure_sources"] = [ + src["source"] for src in (dxgi, physical) if src.get("pressure") + ] + return merged + + +# --------------------------------------------------------------------------- +# Two-timescale residency control (see tasks/done/RESIDENCY_TWO_TIMESCALE_PLAN.md) +# +# Allowance lives in *target-space* (0.95*cap - live); the allocator cap is set +# in *cap-space*. The two differ by the gc_threshold factor: a cap raise of ``d`` +# only adds ``gc_threshold * d`` of GC target / allowance. Every conversion below +# carries the ``/ gc_threshold`` so no call site open-codes it (that missing +# divisor silently under-reserves and licenses a promotion that immediately binds). +# +# All functions here are pure/CPU-testable. The cap and residency levers both act +# only at phase boundaries, and a cap change is realized lazily -- on the next +# fresh cudaMalloc, i.e. the next forward()/step -- so the controller reads +# counters that lag its move by one window (the FSM's verify phases absorb this). +# --------------------------------------------------------------------------- + +GC_THRESHOLD = 0.95 + + +def allocator_allowance_bytes(cap_bytes, live_bytes, *, gc_threshold=GC_THRESHOLD) -> int: + """Idle-cache allowance under the cap: ``gc_threshold*cap - live`` (pure). + + The caching allocator sweeps idle segments when reserved would cross the GC + target ``gc_threshold * cap``; live bytes (residents + ring + activations) + count against that target but cannot be freed. So the room left for reusable + idle cache -- the buffer between smooth reuse and a fresh-cudaMalloc sweep -- + is ``gc_threshold*cap - live``. Negative means live alone exceeds the target: + every sweep dumps all cache and every reuse re-mallocs (self-sustaining + thrash), so callers must keep this positive at the live peak. + """ + return int(float(gc_threshold) * float(max(0, int(cap_bytes))) - float(max(0, int(live_bytes)))) + + +def cap_bytes_for_live( + planned_live_bytes, + cache_budget_bytes, + cliff_cap_bytes, + *, + floor_cap_bytes=0, + gc_threshold=GC_THRESHOLD, +) -> int: + """Cap that hosts ``planned_live`` plus an idle-cache budget (pure). + + Inverse of :func:`allocator_allowance_bytes`: to let live grow to + ``planned_live`` while keeping ``cache_budget`` of reusable idle cache under + the GC target, the cap must be ``(planned_live + cache_budget) / gc_threshold``. + Clamped to the WDDM cliff bound above (never license silent paging; see + :func:`cap_fraction`) and an optional floor below. + """ + want = (float(max(0, int(planned_live_bytes))) + float(max(0, int(cache_budget_bytes)))) / float(gc_threshold) + want = min(want, float(int(cliff_cap_bytes))) + want = max(want, float(max(0, int(floor_cap_bytes)))) + return int(want) + + +def cap_can_host_promotion( + live_bytes, + block_bytes, + slack_pad_bytes, + cliff_cap_bytes, + *, + gc_threshold=GC_THRESHOLD, +) -> bool: + """Can the cheap cap lever (tier 1) absorb one more resident block? (pure). + + Promoting a streamed block to resident raises live by ``block_bytes``. To + keep ``slack_pad_bytes`` of allowance afterward, the GC target must reach + ``live + block + slack``, i.e. the cap must reach + ``(live + block + slack) / gc_threshold``. The cap lever can do this only if + that target cap is still under the WDDM cliff bound; otherwise the cap is + pinned at the cliff and the allowance must come from lowering live -- an + expensive resident demote (tier 2). + + The ``/ gc_threshold`` is load-bearing: a naive ``cliff - cap >= block`` test + under-reserves by the 0.95 factor. + """ + need_cap = ( + float(max(0, int(live_bytes))) + + float(max(0, int(block_bytes))) + + float(max(0, int(slack_pad_bytes))) + ) / float(gc_threshold) + return need_cap <= float(int(cliff_cap_bytes)) + + +def residency_promote_ok( + num_alloc_retries, + allocator_slack_bytes, + block_bytes, + slack_pad_bytes, +) -> bool: + """Sampling climb gate: convert one streamed block to resident? (pure). + + Approach residency from below (undershoot-and-climb): only add a block when + the telemetry proves the room is really there -- + + * ``num_alloc_retries == 0`` over the window (nothing cap-binding), AND + * worst-shape allocator slack (``0.95 * cap - predicted_live``) + exceeds one block plus the pad, so the promotion still leaves + ``slack_pad`` of reusable-cache allowance. + + Both must hold: retries can be zero simply because residency is too low, so + the worst-shape allocator-slack test is what proves there is room to spend. + """ + if int(num_alloc_retries or 0) > 0: + return False + return float(allocator_slack_bytes or 0) > float(block_bytes) + float(slack_pad_bytes) + + +# --- Hysteresis FSM (one transition per phase boundary) --------------------- +# +# States mirror the plan's state machine. DEMOTE_REQUIRED is folded into the +# transition (emit "demote" and land in COLD) since the ring resize is a +# synchronous boundary transaction, not a state that waits a window. + +FSM_COLD = "cold" # measurements invalid (post-compile/retrace/demote) +FSM_STABLE = "stable" # clean + eligible for the from-below climb +FSM_CAP_VERIFY = "cap_verify" # cap raised, confirming it took +FSM_PROMOTION_VERIFY = "promotion_verify" # one block promoted, confirming clean +FSM_COOLDOWN = "cooldown" # re-promotion barred N windows after a rollback + +ACT_HOLD = "hold" +ACT_RAISE_CAP = "raise_cap" +ACT_PROMOTE = "promote" +ACT_DEMOTE = "demote" +ACT_ROLLBACK = "rollback" + + +@dataclass(frozen=True) +class ResidencyFsmState: + name: str = FSM_COLD + windows_in_state: int = 0 + + +def residency_fsm_step( + state: ResidencyFsmState, + signals: dict, + *, + k_clean: int = 2, + k_verify: int = 2, + cooldown_n: int = 4, +) -> tuple[ResidencyFsmState, str]: + """Advance the residency controller one phase boundary (pure, CPU-testable). + + ``signals`` (all read as bools unless noted): + * ``measurements_invalid`` -- a recompile / retrace / layout change happened; + every counter read across it is meaningless -> force COLD. + * ``binding`` -- retries or external pressure this window (allowance too low). + * ``cap_can_relieve`` -- a cap raise can restore the pad at current live + (cliff has room); if False under pressure, only a demote can. + * ``promote_gate`` -- :func:`residency_promote_ok` verdict (there is slack). + * ``cap_covers_promo`` -- the cliff already hosts the promotion with no raise + (:func:`cap_can_host_promotion` at the current cap headroom). + + Returns ``(next_state, action)`` with ``action`` in the ``ACT_*`` set. The + verify phases each span ``k_verify`` windows because a cap/residency move only + shows its signal on the *next* step (the one-window GC lag). + """ + name = state.name + w = state.windows_in_state + 1 + + invalid = bool(signals.get("measurements_invalid")) + binding = bool(signals.get("binding")) + cap_relieve = bool(signals.get("cap_can_relieve")) + promote_gate = bool(signals.get("promote_gate")) + cap_covers = bool(signals.get("cap_covers_promo")) + + def stay(action=ACT_HOLD): + return ResidencyFsmState(name, w), action + + def enter(new_name, action=ACT_HOLD): + return ResidencyFsmState(new_name, 0), action + + # A demote is a synchronous transaction that invalidates the layout -> COLD. + def demote(): + return ResidencyFsmState(FSM_COLD, 0), ACT_DEMOTE + + if invalid and name not in (FSM_COLD,): + return enter(FSM_COLD) + + if name == FSM_COLD: + if invalid or binding: + return ResidencyFsmState(FSM_COLD, 0 if invalid else w), ACT_HOLD + return enter(FSM_STABLE) if w >= k_clean else stay() + + if name == FSM_STABLE: + if binding: + return enter(FSM_CAP_VERIFY, ACT_RAISE_CAP) if cap_relieve else demote() + # Eligible to climb only once the state has held clean for k_clean windows. + if promote_gate and w >= k_clean: + if cap_covers: + return enter(FSM_PROMOTION_VERIFY, ACT_PROMOTE) + return enter(FSM_CAP_VERIFY, ACT_RAISE_CAP) # pre-fund, promote after verify + return stay() + + if name == FSM_CAP_VERIFY: + if binding: + return demote() # the raise didn't relieve -> shed live + return enter(FSM_STABLE) if w >= k_verify else stay() + + if name == FSM_PROMOTION_VERIFY: + if binding: + return enter(FSM_COOLDOWN, ACT_ROLLBACK) + if w == 1: + return stay() # ignore only layout-contaminated non-binding signals + return enter(FSM_STABLE) if w >= k_verify + 1 else stay() + + if name == FSM_COOLDOWN: + if binding: + return demote() # demote still allowed during cooldown + return enter(FSM_STABLE) if w >= cooldown_n else stay() + + # Unknown state: fail safe to COLD. + return enter(FSM_COLD) From 26a5f052f43cf5029fb4c8376b11d8f7d3891838 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Tue, 14 Jul 2026 21:55:26 +0200 Subject: [PATCH 02/20] Add generic arena offload runtime --- requirements_base.txt | 2 +- tests/test_arena_canonical_transaction.py | 296 ++++ tests/test_arena_load_session.py | 87 ++ tests/test_arena_offload_api.py | 241 ++++ tests/test_arena_offload_planner.py | 104 ++ tests/test_arena_offload_policy.py | 789 +++++++++++ tests/test_canonical_arena.py | 173 +++ tests/test_generic_block_dispatcher.py | 482 +++++++ tests/test_residency.py | 207 +++ tests/test_residency_two_timescale.py | 218 +++ tests/test_transfer_plan.py | 141 ++ tests/test_transfer_runtime.py | 165 +++ .../arena_offload/__init__.py | 57 + .../memory_management/arena_offload/api.py | 358 +++++ .../arena_offload/construction.py | 590 ++++++++ .../arena_offload/discovery.py | 265 ++++ .../arena_offload/dispatcher.py | 333 +++++ .../memory_management/arena_offload/errors.py | 39 + .../memory_management/arena_offload/fp8.py | 113 ++ .../memory_management/arena_offload/layout.py | 769 +++++++++++ .../arena_offload/load_session.py | 180 +++ .../arena_offload/ownership.py | 62 + .../arena_offload/planner.py | 221 +++ .../memory_management/arena_offload/policy.py | 529 +++++++ .../arena_offload/resources.py | 165 +++ .../arena_offload/runtime.py | 1217 +++++++++++++++++ .../arena_offload/transfer.py | 816 +++++++++++ toolkit/memory_management/canonical_arena.py | 276 ++++ .../memory_management/immutable_runtime.py | 880 ++++++++++++ toolkit/memory_management/residency.py | 338 +++++ toolkit/memory_management/runtime.py | 82 ++ toolkit/memory_management/transfer_plan.py | 216 +++ toolkit/quantization/__init__.py | 1 + toolkit/quantization/fp8_linear.py | 752 ++++++++++ toolkit/quantization/fp8_transpose.py | 87 ++ toolkit/quantization/storage.py | 206 +++ toolkit/util/quantize.py | 247 +++- 37 files changed, 11692 insertions(+), 12 deletions(-) create mode 100644 tests/test_arena_canonical_transaction.py create mode 100644 tests/test_arena_load_session.py create mode 100644 tests/test_arena_offload_api.py create mode 100644 tests/test_arena_offload_planner.py create mode 100644 tests/test_arena_offload_policy.py create mode 100644 tests/test_canonical_arena.py create mode 100644 tests/test_generic_block_dispatcher.py create mode 100644 tests/test_residency.py create mode 100644 tests/test_residency_two_timescale.py create mode 100644 tests/test_transfer_plan.py create mode 100644 tests/test_transfer_runtime.py create mode 100644 toolkit/memory_management/arena_offload/__init__.py create mode 100644 toolkit/memory_management/arena_offload/api.py create mode 100644 toolkit/memory_management/arena_offload/construction.py create mode 100644 toolkit/memory_management/arena_offload/discovery.py create mode 100644 toolkit/memory_management/arena_offload/dispatcher.py create mode 100644 toolkit/memory_management/arena_offload/errors.py create mode 100644 toolkit/memory_management/arena_offload/fp8.py create mode 100644 toolkit/memory_management/arena_offload/layout.py create mode 100644 toolkit/memory_management/arena_offload/load_session.py create mode 100644 toolkit/memory_management/arena_offload/ownership.py create mode 100644 toolkit/memory_management/arena_offload/planner.py create mode 100644 toolkit/memory_management/arena_offload/policy.py create mode 100644 toolkit/memory_management/arena_offload/resources.py create mode 100644 toolkit/memory_management/arena_offload/runtime.py create mode 100644 toolkit/memory_management/arena_offload/transfer.py create mode 100644 toolkit/memory_management/canonical_arena.py create mode 100644 toolkit/memory_management/immutable_runtime.py create mode 100644 toolkit/memory_management/residency.py create mode 100644 toolkit/memory_management/runtime.py create mode 100644 toolkit/memory_management/transfer_plan.py create mode 100644 toolkit/quantization/__init__.py create mode 100644 toolkit/quantization/fp8_linear.py create mode 100644 toolkit/quantization/fp8_transpose.py create mode 100644 toolkit/quantization/storage.py diff --git a/requirements_base.txt b/requirements_base.txt index afbff335f5..0c0472491d 100644 --- a/requirements_base.txt +++ b/requirements_base.txt @@ -1,4 +1,4 @@ -torchao==0.10.0 +torchao==0.17.0 safetensors git+https://github.com/huggingface/diffusers.git@c943837899b16cbae2f619b8dd4f7bb6f07dd81a #pip install git+https://github.com/huggingface/diffusers.git@refs/pull/13432/head diff --git a/tests/test_arena_canonical_transaction.py b/tests/test_arena_canonical_transaction.py new file mode 100644 index 0000000000..53252abffe --- /dev/null +++ b/tests/test_arena_canonical_transaction.py @@ -0,0 +1,296 @@ +import gc +import unittest +import weakref +from unittest import mock + +import torch + +from toolkit.memory_management import pin_manager +from toolkit.memory_management.arena_offload.construction import CanonicalBuildError +from toolkit.memory_management.canonical_arena import CanonicalArena + + +def frozen_linear(): + layer = torch.nn.Linear(4, 4) + layer.requires_grad_(False) + return layer + + +class FailingLinear(torch.nn.Linear): + fail_publication = False + + def __setattr__(self, name, value): + if name == "weight" and getattr(self, "fail_publication", False): + raise RuntimeError("injected_commit_failure") + super().__setattr__(name, value) + + +class CanonicalTransactionTests(unittest.TestCase): + def test_prepare_does_not_mutate_and_direct_population_uses_final_views(self): + model = torch.nn.Sequential(frozen_linear()) + layer = model[0] + original = layer.weight + arena = CanonicalArena() + build = arena.prepare({"blocks.0": [("linear", layer)]}, model=model) + self.assertIs(layer.weight, original) + destination = build.destinations[("blocks.0", "linear", "weight")] + expected = torch.full_like(destination, 3) + build.populate(lambda destinations: destinations[("blocks.0", "linear", "weight")].copy_(expected)) + build.commit() + try: + self.assertEqual(layer.weight.data_ptr(), destination.data_ptr()) + torch.testing.assert_close(layer.weight, expected) + finally: + CanonicalArena.unguard_whole_model_to(model) + arena.release() + + def test_model_source_leaves_match_destination_keys(self): + model = torch.nn.Sequential(frozen_linear()) + layer = model[0] + arena = CanonicalArena() + build = arena.prepare({"blocks.0": [("linear", layer)]}, model=model) + try: + sources = dict(build.model_source_leaves()) + self.assertEqual(set(sources), set(build.destinations)) + torch.testing.assert_close( + sources[("blocks.0", "linear", "weight")], layer.weight + ) + torch.testing.assert_close( + sources[("blocks.0", "linear", "bias")], layer.bias + ) + finally: + build.rollback() + + def test_incremental_direct_block_releases_bounded_source_before_commit(self): + model = torch.nn.Sequential(frozen_linear()) + layer = model[0] + expected = layer.weight.detach().clone() + source = layer.weight + source_ref = weakref.ref(source) + arena = CanonicalArena() + build = arena.prepare({}, model=model) + build.add_block("blocks.0", [("linear", layer)]) + build.populate_block_from_model("blocks.0") + build.release_block_sources_to_meta("blocks.0") + source = None + gc.collect() + + self.assertIsNone(source_ref()) + self.assertEqual(layer.weight.device.type, "meta") + + build.finish_population() + build.commit() + try: + torch.testing.assert_close(layer.weight, expected) + finally: + CanonicalArena.unguard_whole_model_to(model) + arena.release() + + def test_state_dict_population_consumes_each_source_before_next_block(self): + layers = [frozen_linear(), frozen_linear()] + model = torch.nn.Module() + model.blocks = torch.nn.ModuleList() + for layer in layers: + block = torch.nn.Module() + block.linear = layer + model.blocks.append(block) + state = { + f"blocks.{index}.linear.weight": torch.full_like(layer.weight, index + 1) + for index, layer in enumerate(layers) + } + state.update({ + f"blocks.{index}.linear.bias": torch.full_like(layer.bias, index + 3) + for index, layer in enumerate(layers) + }) + expected = {key: value.clone() for key, value in state.items()} + first_refs = ( + weakref.ref(state["blocks.0.linear.weight"]), + weakref.ref(state["blocks.0.linear.bias"]), + ) + residual = torch.tensor([9.0]) + state["head.weight"] = residual + + def blocks(): + yield "blocks.0", [("linear", layers[0])] + gc.collect() + self.assertNotIn("blocks.0.linear.weight", state) + self.assertNotIn("blocks.0.linear.bias", state) + self.assertTrue(all(ref() is None for ref in first_refs)) + self.assertIn("blocks.1.linear.weight", state) + yield "blocks.1", [("linear", layers[1])] + + arena = CanonicalArena() + build = arena.prepare({}, model=model) + consumed = build.populate_from_state_dict_consuming( + state, + blocks=blocks(), + ) + + self.assertEqual( + set(consumed), + { + "blocks.0.linear.weight", + "blocks.0.linear.bias", + "blocks.1.linear.weight", + "blocks.1.linear.bias", + }, + ) + self.assertEqual(set(state), {"head.weight"}) + self.assertIs(state["head.weight"], residual) + + build.commit() + try: + for index, layer in enumerate(layers): + torch.testing.assert_close( + layer.weight, + expected[f"blocks.{index}.linear.weight"], + ) + torch.testing.assert_close( + layer.bias, + expected[f"blocks.{index}.linear.bias"], + ) + finally: + arena.release() + + def test_direct_population_failure_rolls_back_without_marker_or_pin_leak(self): + model = torch.nn.Sequential(frozen_linear()) + layer = model[0] + original = layer.weight + before = pin_manager.pinned_bytes_by_kind().get("weights", 0) + build = CanonicalArena().prepare({"blocks.0": [("linear", layer)]}, model=model) + with self.assertRaisesRegex(RuntimeError, "injected_population_failure"): + build.populate(lambda _destinations: (_ for _ in ()).throw(RuntimeError("injected_population_failure"))) + self.assertIs(layer.weight, original) + self.assertFalse(hasattr(model, "_arena_offload_runtime")) + self.assertEqual(pin_manager.pinned_bytes_by_kind().get("weights", 0), before) + model(torch.randn(1, 4)) + + def test_pin_failure_after_an_earlier_block_releases_everything(self): + layers = [frozen_linear(), frozen_linear()] + originals = [layer.weight for layer in layers] + arena = CanonicalArena() + build = arena.prepare({ + "blocks.0": [("linear", layers[0])], + "blocks.1": [("linear", layers[1])], + }) + real_commit = pin_manager.pin_register_commit + calls = 0 + + def fail_second(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected_pin_failure") + return real_commit(*args, **kwargs) + + with mock.patch.object(pin_manager, "pin_register_commit", side_effect=fail_second): + with self.assertRaisesRegex(RuntimeError, "injected_pin_failure"): + build.populate_from_model() + self.assertEqual([layer.weight is original for layer, original in zip(layers, originals)], [True, True]) + self.assertFalse(arena.canonicalized) + + def test_commit_failure_after_first_repoint_restores_parameter_identity(self): + first = frozen_linear() + second = FailingLinear(4, 4) + second.requires_grad_(False) + originals = (first.weight, second.weight) + arena = CanonicalArena() + build = arena.prepare({"blocks.0": [("first", first), ("second", second)]}) + build.populate_from_model() + second.fail_publication = True + with self.assertRaisesRegex(RuntimeError, "injected_commit_failure"): + build.commit() + second.fail_publication = False + self.assertIs(first.weight, originals[0]) + self.assertIs(second.weight, originals[1]) + self.assertFalse(arena.canonicalized) + first(torch.randn(1, 4)) + second(torch.randn(1, 4)) + + + def test_unsupported_layout_mid_stack_leaves_originals_untouched(self): + layers = [frozen_linear(), frozen_linear()] + originals = [layer.weight for layer in layers] + arena = CanonicalArena() + from toolkit.memory_management.arena_offload import construction + + real_inspect = construction.inspect_block + calls = 0 + + def fail_second(key, entries): + nonlocal calls + calls += 1 + if calls == 2: + raise ValueError("unsupported_quant_wrapper:blocks.1:linear") + return real_inspect(key, entries) + + with mock.patch.object(construction, "inspect_block", side_effect=fail_second): + with self.assertRaisesRegex(ValueError, "unsupported_quant_wrapper"): + arena.prepare({ + "blocks.0": [("linear", layers[0])], + "blocks.1": [("linear", layers[1])], + }) + self.assertEqual( + [layer.weight is original for layer, original in zip(layers, originals)], + [True, True], + ) + self.assertFalse(arena.canonicalized) + + def test_allocation_failure_after_first_block_prepared_is_atomic(self): + layers = [frozen_linear(), frozen_linear()] + originals = [layer.weight for layer in layers] + arena = CanonicalArena() + real_prepare = pin_manager.pin_register_prepare + calls = 0 + + def fail_second(nbytes): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("injected_allocation_failure") + return real_prepare(nbytes) + + with mock.patch.object(pin_manager, "pin_register_prepare", side_effect=fail_second): + with self.assertRaisesRegex(RuntimeError, "injected_allocation_failure"): + arena.prepare({ + "blocks.0": [("linear", layers[0])], + "blocks.1": [("linear", layers[1])], + }) + self.assertEqual( + [layer.weight is original for layer, original in zip(layers, originals)], + [True, True], + ) + self.assertFalse(arena.canonicalized) + + def test_wrapper_validation_failure_releases_committed_pin(self): + model = torch.nn.Sequential(frozen_linear()) + layer = model[0] + original = layer.weight + before = pin_manager.pinned_bytes_by_kind().get("weights", 0) + build = CanonicalArena().prepare( + {"blocks.0": [("linear", layer)]}, + model=model, + ) + with mock.patch( + "toolkit.memory_management.arena_offload.construction.linear_views", + side_effect=RuntimeError("injected_wrapper_validation_failure"), + ): + with self.assertRaisesRegex(RuntimeError, "injected_wrapper_validation_failure"): + build.populate_from_model() + self.assertIs(layer.weight, original) + self.assertEqual(pin_manager.pinned_bytes_by_kind().get("weights", 0), before) + model(torch.randn(1, 4)) + + def test_commit_without_population_rolls_back_prepared_build(self): + layer = frozen_linear() + original = layer.weight + arena = CanonicalArena() + build = arena.prepare({"blocks.0": [("linear", layer)]}) + with self.assertRaisesRegex(CanonicalBuildError, "canonical_build_not_populated"): + build.commit() + self.assertIs(layer.weight, original) + self.assertFalse(arena.canonicalized) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_arena_load_session.py b/tests/test_arena_load_session.py new file mode 100644 index 0000000000..bb0c143927 --- /dev/null +++ b/tests/test_arena_load_session.py @@ -0,0 +1,87 @@ +from types import SimpleNamespace + +import torch + +from toolkit.memory_management.arena_offload import model_load_arena_session +from toolkit.memory_management.arena_offload.load_session import ( + PENDING_CANONICAL_BUILD_ATTR, + claim_pending_canonical_build, +) +from toolkit.memory_management.runtime import close_memory_runtime_preparation + + +class Block(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + + +class Transformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.blocks = torch.nn.ModuleList([Block(), Block()]) + self.head = torch.nn.Linear(4, 4) + + +def base_model(): + return SimpleNamespace( + model_config=SimpleNamespace( + layer_offloading=True, + layer_offloading_smart=True, + ), + device_torch=torch.device("cpu"), + te_only=False, + get_transformer_block_names=lambda: ["blocks"], + ) + + +def frozen_transformer(): + model = Transformer() + model.requires_grad_(False) + return model + + +def test_common_load_state_dict_uses_generic_arena_session(): + source = frozen_transformer() + expected_head = source.head.weight.detach().clone() + state = {key: value.clone() for key, value in source.state_dict().items()} + target = frozen_transformer() + + with model_load_arena_session(base_model()): + incompatible = target.load_state_dict(state, strict=True, assign=True) + + assert not incompatible.missing_keys + assert not incompatible.unexpected_keys + assert "blocks.0.linear.weight" not in state + torch.testing.assert_close(target.head.weight, expected_head) + build = claim_pending_canonical_build(target) + assert build is not None + build.rollback() + + +def test_unfinished_generic_load_is_released_by_shared_cleanup(): + owner = base_model() + source = frozen_transformer() + state = {key: value.clone() for key, value in source.state_dict().items()} + target = frozen_transformer() + + with model_load_arena_session(owner): + target.load_state_dict(state, strict=True, assign=True) + + assert hasattr(target, PENDING_CANONICAL_BUILD_ATTR) + close_memory_runtime_preparation(owner) + assert not hasattr(target, PENDING_CANONICAL_BUILD_ATTR) + + +def test_trainable_target_falls_back_to_normal_assignment(): + source = Transformer() + expected = source.blocks[0].linear.weight.detach().clone() + state = {key: value.clone() for key, value in source.state_dict().items()} + target = Transformer() + + with model_load_arena_session(base_model()) as session: + target.load_state_dict(state, strict=True, assign=True) + + assert session.unsupported_reason is not None + assert not hasattr(target, PENDING_CANONICAL_BUILD_ATTR) + torch.testing.assert_close(target.blocks[0].linear.weight, expected) diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py new file mode 100644 index 0000000000..03a8c702cb --- /dev/null +++ b/tests/test_arena_offload_api.py @@ -0,0 +1,241 @@ +"""Facade-level coverage for `toolkit.memory_management.arena_offload`. + +These test the seam, not the machine: the helpers shared code now relies on +(`get_arena_runtime`, `is_memory_managed`, `memory_runtime_owns_compile`) and +the config mapping. Building a real arena needs CUDA and a real model; lifecycle +is covered by the arena contract tests and the Krea2 train smoke. +""" + +import ast +from dataclasses import fields +from pathlib import Path +import unittest + +import torch + +from toolkit.memory_management.arena_offload import ( + ArenaOffloadConfig, + get_arena_runtime, + is_arena_offloaded, + is_memory_managed, + memory_runtime_owns_compile, +) +from toolkit.memory_management.arena_offload.api import RUNTIME_ATTR, unwrap +from toolkit.memory_management.arena_offload.runtime import _fixed_working_bytes +from toolkit.memory_management.arena_offload.runtime import ArenaOffloadRuntime + +GIB = 1024**3 + + +def test_arena_runtime_excludes_legacy_training_policy_calls(): + source_path = ( + Path(__file__).parents[1] / "jobs" / "process" / "BaseSDTrainProcess.py" + ) + tree = ast.parse(source_path.read_text(encoding="utf-8")) + parents = {} + for parent in ast.walk(tree): + for child in ast.iter_child_nodes(parent): + parents[child] = parent + + guarded_calls = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + name = func.attr if isinstance(func, ast.Attribute) else None + if name not in { + "prepare_training_memory_for_shape", + "auto_tune_training_memory", + }: + continue + ancestor = parents.get(node) + guarded = False + while ancestor is not None: + if isinstance(ancestor, ast.If): + condition = ast.unparse(ancestor.test) + if "arena_runtime is None" in condition: + guarded = True + break + ancestor = parents.get(ancestor) + guarded_calls.append((name, guarded)) + + # Upstream's legacy backend has no fork-local autotune calls. If those + # calls are added later, they must be explicitly excluded for arena runs. + assert all(guarded for _name, guarded in guarded_calls) + + +class _Wrapper(torch.nn.Module): + """Stands in for Accelerate/DDP, which expose the real model at `.module`.""" + + def __init__(self, inner): + super().__init__() + self.module = inner + + +class _FakeModelConfig: + quantize = True + qtype = "qfloat8" + layer_offloading = True + layer_offloading_smart = True + layer_offloading_fp8_forward = True + layer_offloading_fp8_grad_input = True + layer_offloading_fp8_sampling = True + compile = False + compile_sample = True + train_compile_blocks = False + layer_offloading_smart_working_reserve_gb = -1.0 + layer_offloading_smart_wddm_margin_gb = None + layer_offloading_smart_wddm_hard_gb = 1.0 + layer_offloading_wddm_spill_reserve_pct = 0.10 + layer_offloading_block_stream_only = False + layer_offloading_checkpoint_keep_last = 2 + layer_offloading_prefetch_depth = 3 + layer_offloading_smart_sampling_working_reserve_gb = -1.0 + layer_offloading_smart_sampling_wddm_margin_gb = -1.0 + layer_offloading_smart_sampling_wddm_hard_gb = 1.0 + + +class ArenaOffloadHelpersTest(unittest.TestCase): + def test_helpers_are_none_safe(self): + self.assertIsNone(get_arena_runtime(None)) + self.assertFalse(is_arena_offloaded(None)) + self.assertFalse(is_memory_managed(None)) + self.assertFalse(memory_runtime_owns_compile(None)) + + def test_plain_module_is_not_managed(self): + model = torch.nn.Linear(4, 4) + self.assertFalse(is_arena_offloaded(model)) + self.assertFalse(is_memory_managed(model)) + self.assertFalse(memory_runtime_owns_compile(model)) + + def test_runtime_found_through_wrappers(self): + inner = torch.nn.Linear(4, 4) + runtime = object() + setattr(inner, RUNTIME_ATTR, runtime) + wrapped = _Wrapper(_Wrapper(inner)) + + self.assertIs(unwrap(wrapped), inner) + self.assertIs(get_arena_runtime(wrapped), runtime) + self.assertTrue(is_arena_offloaded(wrapped)) + self.assertTrue(is_memory_managed(wrapped)) + self.assertTrue(memory_runtime_owns_compile(wrapped)) + + def test_legacy_backend_is_managed_but_does_not_own_compile(self): + """The distinction generic block compile depends on.""" + model = torch.nn.Linear(4, 4) + model._memory_manager = object() + + self.assertTrue(is_memory_managed(model)) + self.assertFalse(is_arena_offloaded(model)) + self.assertFalse(memory_runtime_owns_compile(model)) + + def test_unwrap_terminates_on_self_referential_wrapper(self): + model = torch.nn.Linear(4, 4) + model.module = model # a module that is its own `.module` + self.assertIs(unwrap(model), model) + + def test_place_permanent_modules_does_not_move_canonical_leaf(self): + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.root_token = torch.nn.Parameter( + torch.ones(4), requires_grad=False + ) + self.register_buffer("root_buffer", torch.ones(4)) + self.canonical = torch.nn.Linear(4, 4) + self.permanent = torch.nn.Linear(4, 4) + + model = Model() + runtime = object.__new__(ArenaOffloadRuntime) + runtime._model = model + runtime._canonical_modules = (model.canonical,) + runtime._closed = False + + runtime.place_permanent_modules("cpu", torch.float64) + + self.assertEqual(model.canonical.weight.dtype, torch.float32) + self.assertEqual(model.permanent.weight.dtype, torch.float64) + self.assertEqual(model.root_token.dtype, torch.float64) + self.assertEqual(model.root_buffer.dtype, torch.float64) + + +class ArenaOffloadConfigTest(unittest.TestCase): + def test_from_model_config_maps_the_public_surface(self): + config = ArenaOffloadConfig.from_model_config(_FakeModelConfig()) + + self.assertTrue(config.enabled) + self.assertTrue(config.fp8_forward) + self.assertTrue(config.fp8_backward) + self.assertTrue(config.fp8_sampling) + # compile_blocks is derived, not its own public knob. + self.assertTrue(config.compile_blocks) + self.assertEqual(config._policy.prefetch_depth, 3) + self.assertEqual(config._policy.checkpoint_keep_last, 2) + + def test_public_surface_is_narrow(self): + public = {field.name for field in fields(ArenaOffloadConfig) if not field.name.startswith("_")} + self.assertEqual( + public, + { + "enabled", + "fp8_forward", + "fp8_backward", + "fp8_sampling", + "compile_blocks", + }, + ) + + def test_fp8_flags_require_fp8_weights(self): + """An fp8_* toggle on a non-fp8 model is a no-op, not a crash.""" + + class NoQuant(_FakeModelConfig): + quantize = False + + with self.assertWarnsRegex(RuntimeWarning, "ignored irrelevant FP8 options"): + config = ArenaOffloadConfig.from_model_config(NoQuant()) + self.assertFalse(config.fp8_forward) + self.assertFalse(config.fp8_backward) + self.assertFalse(config.fp8_sampling) + self.assertTrue(config.enabled) + + def test_missing_attributes_fall_back_to_defaults(self): + config = ArenaOffloadConfig.from_model_config(object()) + self.assertFalse(config.enabled) + self.assertFalse(config.compile_blocks) + self.assertEqual(config._policy.prefetch_depth, 3) + + def test_compatibility_aliases_map_to_internal_policy(self): + class Aliases: + layer_offloading_smart_headroom_gb = 4.0 + layer_offloading_smart_buffer_gb = 1.5 + layer_offloading_smart_hard_buffer_gb = 0.75 + + policy = ArenaOffloadConfig.from_model_config(Aliases())._policy + self.assertEqual(policy.working_reserve_gib, 4.0) + self.assertEqual(policy.wddm_margin_gib, 1.5) + self.assertEqual(policy.wddm_hard_gib, 0.75) + + def test_backward_without_fp8_forward_is_ignored_once(self): + class Invalid: + quantize = True + qtype = "qfloat8" + layer_offloading_fp8_grad_input = True + + with self.assertWarnsRegex(RuntimeWarning, "fp8_backward_without_fp8_forward"): + config = ArenaOffloadConfig.from_model_config(Invalid()) + self.assertFalse(config.fp8_backward) + + +class SamplingReserveTest(unittest.TestCase): + def test_auto_working_reserve_is_none(self): + """Unset / negative / 'auto' all mean 'let the runtime size it'.""" + for value in (None, -1.0, "auto"): + self.assertIsNone(_fixed_working_bytes(value)) + + def test_explicit_working_reserve_is_bytes(self): + self.assertEqual(_fixed_working_bytes(2.0), 2 * GIB) + self.assertEqual(_fixed_working_bytes(0), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_arena_offload_planner.py b/tests/test_arena_offload_planner.py new file mode 100644 index 0000000000..68a9dac11b --- /dev/null +++ b/tests/test_arena_offload_planner.py @@ -0,0 +1,104 @@ +from types import SimpleNamespace +from unittest import mock + +from toolkit.memory_management.arena_offload.planner import GIB, build_training_plan + + +class _Arena: + def __init__(self, records): + self._records = {record.block_key: record for record in records} + + def block_keys(self): + return tuple(self._records) + + def block_record(self, key): + return self._records[key] + + +def _record(key, committed_gib): + return SimpleNamespace( + block_key=key, + committed_bytes=int(committed_gib * GIB), + modules=(object(),), + leaf_names=("weight",), + ) + + +def _config(): + return SimpleNamespace( + _policy=SimpleNamespace( + working_reserve_gib=-1, + wddm_hard_gib=1.0, + wddm_margin_gib=1.0, + checkpoint_keep_last=0, + prefetch_depth=2, + ) + ) + + +def test_auto_plan_uses_all_resident_fast_path_when_complete_model_fits(): + records = [_record(f"blocks.{index}", 2.0) for index in range(3)] + with ( + mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(12 * GIB, 12 * GIB), + ), + mock.patch( + "toolkit.memory_management.arena_offload.planner._singleton_stats", + return_value=(1 * GIB, 0, set()), + ), + ): + plan = build_training_plan( + SimpleNamespace(), _Arena(records), (), "cuda", _config() + ) + + assert plan["all_resident_fit"] + assert plan["offloaded_layers"] == 0 + assert plan["generic_resident_bytes"] == 6 * GIB + assert plan["ring_bytes"] == 0 + assert plan["working_reserve_bytes"] == 4 * GIB + + +def test_auto_plan_keeps_streaming_when_complete_model_does_not_fit(): + records = [_record(f"blocks.{index}", 2.0) for index in range(3)] + with ( + mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(10 * GIB, 12 * GIB), + ), + mock.patch( + "toolkit.memory_management.arena_offload.planner._singleton_stats", + return_value=(1 * GIB, 0, set()), + ), + ): + plan = build_training_plan( + SimpleNamespace(), _Arena(records), (), "cuda", _config() + ) + + assert not plan["all_resident_fit"] + assert plan["offloaded_layers"] > 0 + assert plan["ring_bytes"] > 0 + assert plan["working_reserve_bytes"] == 5 * GIB + + +def test_explicit_working_reserve_controls_all_resident_fit(): + records = [_record(f"blocks.{index}", 2.0) for index in range(3)] + config = _config() + config._policy.working_reserve_gib = 2.0 + with ( + mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(10 * GIB, 12 * GIB), + ), + mock.patch( + "toolkit.memory_management.arena_offload.planner._singleton_stats", + return_value=(1 * GIB, 0, set()), + ), + ): + plan = build_training_plan( + SimpleNamespace(), _Arena(records), (), "cuda", config + ) + + assert plan["all_resident_fit"] + assert plan["offloaded_layers"] == 0 + assert plan["working_reserve_bytes"] == 2 * GIB diff --git a/tests/test_arena_offload_policy.py b/tests/test_arena_offload_policy.py new file mode 100644 index 0000000000..f425beb096 --- /dev/null +++ b/tests/test_arena_offload_policy.py @@ -0,0 +1,789 @@ +import contextlib +import json +from types import SimpleNamespace + +import pytest + +from toolkit.memory_management.arena_offload.policy import ( + ArenaResidencyController, + TrainingSignalWindow, + transfer_benefits_from_residency, +) +from toolkit.memory_management.arena_offload.runtime import ArenaOffloadRuntime + + +def observe(window, **overrides): + values = { + "shape_key": (512, 512), + "step_num": 1, + "allocator_counters": { + "num_alloc_retries": 10, + "num_device_alloc": 20, + "num_device_free": 30, + }, + "peak_allocated_bytes": 100, + "peak_reserved_bytes": 140, + "device_free_bytes": 500, + "resident_bytes": 200, + "ring_bytes": 80, + "compile_counters": None, + "transfer_counters": None, + "step_wall_ms": 10.0, + } + values.update(overrides) + return window.observe(**values) + + +def test_allocator_deltas_tolerate_counter_reset(): + window = TrainingSignalWindow() + first = observe(window) + assert first["allocator"] == { + "alloc_retries_delta": 10, + "alloc_count_delta": 20, + "free_count_delta": 30, + } + second = observe( + window, + allocator_counters={ + "num_alloc_retries": 12, + "num_device_alloc": 25, + "num_device_free": 37, + }, + ) + assert second["allocator"] == { + "alloc_retries_delta": 2, + "alloc_count_delta": 5, + "free_count_delta": 7, + } + reset = observe( + window, + allocator_counters={ + "num_alloc_retries": 1, + "num_device_alloc": 2, + "num_device_free": 3, + }, + ) + assert reset["allocator"] == { + "alloc_retries_delta": 1, + "alloc_count_delta": 2, + "free_count_delta": 3, + } + + +def test_per_shape_peaks_skip_warmup_and_track_independently(): + window = TrainingSignalWindow() + observe(window, peak_allocated_bytes=100, peak_reserved_bytes=140) + observe(window, peak_allocated_bytes=110, peak_reserved_bytes=150) + observe( + window, + shape_key=(768, 768), + peak_allocated_bytes=200, + peak_reserved_bytes=260, + ) + peaks = window.shape_peaks + assert peaks[(512, 512)].steps == 1 + assert peaks[(512, 512)].peak_allocated_bytes == 110 + assert peaks[(512, 512)].peak_reserved_bytes == 150 + assert peaks[(768, 768)].steps == 0 + + +def test_compile_activity_invalidates_shape_measurements(): + window = TrainingSignalWindow() + observe(window) + observe(window, peak_allocated_bytes=110) + assert window.shape_peaks[(512, 512)].steps == 1 + + invalid = observe( + window, + compile_counters={"frames": 4, "graphs": 2, "graph_breaks": 0}, + ) + assert invalid["compile_invalid"] is True + assert window.shape_peaks == {} + + stable = observe( + window, + compile_counters={"frames": 4, "graphs": 2, "graph_breaks": 0}, + ) + assert stable["compile_invalid"] is False + assert window.shape_peaks[(512, 512)].warmup_steps == 1 + + +def test_transfer_metrics_require_settled_multi_step_window_and_handle_reset(): + window = TrainingSignalWindow(transfer_window_steps=3) + assert observe(window)["transfer"] is None + assert observe(window)["transfer"] is None + settled = observe( + window, + transfer_counters={"h2d_ms": 12.0, "bytes": 96_000_000}, + )["transfer"] + assert settled["steps"] == 3 + assert settled["h2d_duty_pct"] == pytest.approx(40.0) + assert settled["achieved_gbps"] == pytest.approx(8.0) + assert settled["h2d_duty_overflow"] is False + + observe(window) + observe(window) + after_reset = observe( + window, + transfer_counters={"h2d_ms": 6.0, "bytes": 42_000_000}, + )["transfer"] + assert after_reset["h2d_duty_pct"] == pytest.approx(20.0) + assert after_reset["achieved_gbps"] == pytest.approx(7.0) + + + + +def test_runtime_preserves_shape_peaks_after_layout_change(): + signals = TrainingSignalWindow() + observe(signals) + observe(signals, peak_allocated_bytes=110) + assert signals.shape_peaks + + next_plan = object() + executor = SimpleNamespace( + transition_training_block=lambda key, resident: { + "changed": True, + "block_key": key, + "resident": resident, + "plan": next_plan, + } + ) + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._executor = executor + runtime._signals = signals + runtime._training_plan = object() + + result = runtime.transition_training_block("blocks.3", resident=True) + assert result["plan"] is next_plan + assert runtime._training_plan is next_plan + assert signals.shape_peaks + + +def test_worst_shape_allocator_slack_reconstructs_current_layout(): + signals = TrainingSignalWindow() + observe( + signals, + shape_key=(512, 512), + peak_allocated_bytes=700, + resident_bytes=100, + ring_bytes=50, + ) + observe( + signals, + shape_key=(512, 512), + peak_allocated_bytes=750, + resident_bytes=100, + ring_bytes=50, + ) + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._signals = signals + runtime._residency = SimpleNamespace(resident_bytes=lambda: 120) + runtime._smart_plan = {"singleton_resident_bytes": 30} + runtime._training_ring_bytes = lambda: 50 + + # working=600, current layout=150 resident + 50 ring => live=800. + assert runtime._worst_shape_allocator_slack_bytes(1000) == 150 + + +def test_aggressive_capacity_counts_exact_smallest_blocks_under_both_budgets(): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._policy = SimpleNamespace(slack_pad_bytes=5) + runtime._promotion_candidates = lambda: tuple( + {"block_key": f"blocks.{index}", "block_bytes": size} + for index, size in enumerate((10, 20, 30, 40, 50)) + ) + runtime._worst_shape_allocator_slack_bytes = lambda _cap: 106 + runtime._worst_shape_candidate_margin_bytes = ( + lambda candidate: 120 - candidate["block_bytes"] + ) + + # Four blocks consume 100 bytes and leave the 5-byte allocator pad. The + # fifth would exceed both the allocator and physical budgets. + assert runtime._aggressive_promotion_capacity(1000) == 4 + + +def test_training_cap_binding_uses_configured_phase_margin(monkeypatch): + calls = [] + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime." + "allocator_cap.apply_wddm_hard_allocator_cap", + lambda device, hard, **kwargs: calls.append((device, hard, kwargs)), + ) + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._device = "cuda:1" + runtime._config = SimpleNamespace( + _policy=SimpleNamespace(wddm_hard_gib=1.25) + ) + + runtime._bind_training_cap() + assert calls == [ + ("cuda:1", 1.25, {"log_prefix": "[ArenaOffload]"}) + ] +def test_shape_working_peak_excludes_residency(): + window = TrainingSignalWindow() + observe(window, peak_allocated_bytes=900, resident_bytes=200) + observe(window, peak_allocated_bytes=1000, resident_bytes=300) + peak = window.shape_peaks[(512, 512)] + assert peak.working_peak_bytes == 620 + + +def test_signal_contains_memory_accounting(): + signal = observe(TrainingSignalWindow()) + assert signal["reclaimable_at_peak_bytes"] == 40 + assert signal["device_free_bytes"] == 500 + assert signal["resident_bytes"] == 200 + assert signal["ring_bytes"] == 80 + assert signal["live_bytes"] == 280 + json.dumps(TrainingSignalWindow().diagnostics()) + +def test_transfer_benefit_gate_requires_valid_nonzero_streaming(): + assert not transfer_benefits_from_residency(None) + assert not transfer_benefits_from_residency({"bytes": 0, "h2d_ms": 20.0}) + assert not transfer_benefits_from_residency({"bytes": 100, "h2d_ms": 0.0}) + assert not transfer_benefits_from_residency( + {"bytes": 100, "h2d_ms": 20.0, "h2d_duty_overflow": True} + ) + assert transfer_benefits_from_residency( + {"bytes": 100, "h2d_ms": 20.0, "h2d_duty_pct": 1.0} + ) + + +def test_controller_promotes_exact_candidate_then_rolls_it_back(): + controller = ArenaResidencyController(slack_pad_bytes=10) + candidate = {"block_key": "blocks.3", "block_bytes": 20} + clean = { + "allocator": {"alloc_retries_delta": 0}, + "peak_allocated_bytes": 100, + "resident_bytes": 200, + "reclaimable_at_peak_bytes": 100, + "compile_invalid": False, + "transfer": { + "bytes": 100, + "h2d_ms": 20.0, + "h2d_duty_pct": 80.0, + "achieved_gbps": 10.0, + }, + } + actions = [] + for _ in range(4): + decision = controller.step( + clean, + candidate=candidate, + demote_candidate=None, + cliff_cap_bytes=1000, + worst_shape_free_bytes=100, + ) + actions.append(decision) + promoted = next(item for item in actions if item.action == "promote") + assert promoted.block_key == "blocks.3" + diagnostics = controller.diagnostics() + assert diagnostics["last_block_key"] == "blocks.3" + assert diagnostics["last_block_bytes"] == 20 + assert diagnostics["last_target_cap_bytes"] is None + + controller.step( + clean, + candidate=None, + demote_candidate=None, + cliff_cap_bytes=1000, + worst_shape_free_bytes=0, + ) + dirty = { + **clean, + "allocator": { + "alloc_retries_delta": 0, + "free_count_delta": 1, + }, + } + rollback = controller.step( + dirty, + candidate=None, + demote_candidate=None, + cliff_cap_bytes=1000, + worst_shape_free_bytes=0, + ) + assert rollback.action == "rollback" + assert rollback.block_key == "blocks.3" + assert controller.last_safe_residency_bytes == 200 + assert controller.last_rejected_residency_bytes == 220 + + for _ in range(8): + held = controller.step( + clean, + candidate=candidate, + demote_candidate=None, + cliff_cap_bytes=1000, + worst_shape_free_bytes=100, + worst_shape_allocator_slack_bytes=20, + ) + assert held.action != "promote" + assert held.reason == "allocator_headband" + + promoted_again = None + for _ in range(4): + promoted_again = controller.step( + clean, + candidate=candidate, + demote_candidate=None, + cliff_cap_bytes=1000, + worst_shape_free_bytes=100, + worst_shape_allocator_slack_bytes=31, + ) + if promoted_again.action == "promote": + break + assert promoted_again.action == "promote" + + +def test_controller_cold_starts_one_whole_block_below(): + controller = ArenaResidencyController() + decision = controller.step( + None, + candidate={"block_key": "blocks.4", "block_bytes": 20}, + demote_candidate={"block_key": "blocks.2", "block_bytes": 30}, + cliff_cap_bytes=1000, + worst_shape_free_bytes=100, + ) + assert decision.action == "demote" + assert decision.block_key == "blocks.2" + assert decision.reason == "approach_from_below" + + +def test_controller_promotes_each_step_with_four_block_headroom(): + controller = ArenaResidencyController(slack_pad_bytes=10) + controller.bootstrapped = True + clean = { + "allocator": { + "alloc_retries_delta": 0, + "free_count_delta": 0, + }, + "resident_bytes": 200, + "compile_invalid": False, + "transfer": None, + } + + first = controller.step( + clean, + candidate={"block_key": "blocks.3", "block_bytes": 20}, + demote_candidate=None, + cliff_cap_bytes=1000, + current_cap_bytes=1000, + worst_shape_free_bytes=100, + worst_shape_allocator_slack_bytes=200, + aggressive_promotion_capacity=4, + ) + assert first.action == "promote" + assert first.reason == "abundant_four_block_headroom" + + second = controller.step( + {**clean, "resident_bytes": 220}, + candidate={"block_key": "blocks.4", "block_bytes": 20}, + demote_candidate=None, + cliff_cap_bytes=1000, + current_cap_bytes=1000, + worst_shape_free_bytes=80, + worst_shape_allocator_slack_bytes=180, + aggressive_promotion_capacity=4, + ) + assert second.action == "promote" + assert second.block_key == "blocks.4" + assert controller.pending_promotion["block_key"] == "blocks.4" + + +def test_controller_four_block_fast_lane_keeps_safety_vetoes(): + controller = ArenaResidencyController(slack_pad_bytes=10) + controller.bootstrapped = True + dirty = { + "allocator": { + "alloc_retries_delta": 1, + "free_count_delta": 0, + }, + "resident_bytes": 200, + "compile_invalid": False, + "transfer": None, + } + decision = controller.step( + dirty, + candidate={"block_key": "blocks.3", "block_bytes": 20}, + demote_candidate={"block_key": "blocks.1", "block_bytes": 20}, + cliff_cap_bytes=1000, + current_cap_bytes=1000, + worst_shape_free_bytes=100, + worst_shape_allocator_slack_bytes=200, + aggressive_promotion_capacity=4, + ) + assert decision.action != "promote" + assert controller.diagnostics()["last_aggressive_gate"] is False + + +def test_controller_does_not_bypass_bootstrap_verification(): + controller = ArenaResidencyController(slack_pad_bytes=10) + controller.bootstrapped = True + controller.begin_bootstrap_promotion( + ("blocks.0", "blocks.1", "blocks.2", "blocks.3"), + 80, + 200, + 1000, + ) + clean = { + "allocator": { + "alloc_retries_delta": 0, + "free_count_delta": 0, + }, + "resident_bytes": 280, + "compile_invalid": False, + "transfer": None, + } + decision = controller.step( + clean, + candidate={"block_key": "blocks.4", "block_bytes": 20}, + demote_candidate=None, + cliff_cap_bytes=1000, + current_cap_bytes=1000, + worst_shape_free_bytes=100, + worst_shape_allocator_slack_bytes=200, + aggressive_promotion_capacity=4, + ) + assert decision.action == "hold" + assert controller.pending_promotion["block_keys"] == ( + "blocks.0", + "blocks.1", + "blocks.2", + "blocks.3", + ) + assert controller.diagnostics()["last_aggressive_gate"] is False + + +def test_controller_raises_cap_by_fixed_fsm_increment(): + controller = ArenaResidencyController(slack_pad_bytes=10) + controller.bootstrapped = True + controller.state = type(controller.state)("stable", 2) + signal = { + "allocator": {"alloc_retries_delta": 1}, + "peak_allocated_bytes": 400, + "reclaimable_at_peak_bytes": 0, + "compile_invalid": False, + "transfer": None, + } + decision = controller.step( + signal, + candidate=None, + demote_candidate={"block_key": "blocks.0", "block_bytes": 100}, + cliff_cap_bytes=1000, + current_cap_bytes=700, + worst_shape_free_bytes=100, + ) + assert decision.action == "raise_cap" + assert decision.target_cap_bytes == 710 + + +def test_arena_sampling_binds_fp8_to_canonical_and_singletons(monkeypatch): + model = SimpleNamespace() + canonical = SimpleNamespace() + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._model = model + runtime._device = "cuda" + runtime._canonical_modules = (canonical,) + runtime._config = SimpleNamespace(fp8_sampling=True) + runtime._sampling_fp8_canonical = 0 + runtime._sampling_fp8_singletons = 0 + runtime._smart_plan = {"singleton_runtime_ids": {11, 22}} + runtime._training_plan = object() + runtime._executor = SimpleNamespace( + TRAIN="train", + activate=lambda *_args: None, + ) + runtime._bind_training_cap = lambda: None + calls = [] + + def fake_enable( + module, include_ids=None, live_ids=None, training=False, device=None + ): + ids = set(include_ids) + calls.append(("enable", module, ids, set(live_ids), training, device)) + return [(None, None, None, None, value) for value in ids] + + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime.enable_fp8", + fake_enable, + ) + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime.disable_fp8", + lambda restores: calls.append(("disable", list(restores))), + ) + + with runtime.sampling_session(): + assert runtime._sampling_fp8_canonical == 1 + assert runtime._sampling_fp8_singletons == 2 + + expected_ids = {id(canonical), 11, 22} + assert calls[0] == ( + "enable", model, expected_ids, expected_ids, False, "cuda" + ) + assert calls[1][0] == "disable" + assert {restore[4] for restore in calls[1][1]} == expected_ids + + +def test_arena_close_releases_all_owned_resources_after_executor_error(): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._resources = SimpleNamespace( + release=lambda: (_ for _ in ()).throw(RuntimeError("executor boom")) + ) + + import pytest + with pytest.raises(RuntimeError, match="executor boom"): + runtime.close() + + assert not runtime._closed + + +def test_bf16_sampling_reserves_largest_singleton_dequant(monkeypatch): + gib = 1024 ** 3 + dequant = 864 * 1024 ** 2 + captured = {} + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._device = "cpu" + runtime._smart_plan = { + "largest_singleton_bf16_dequant_bytes": dequant + } + runtime._config = SimpleNamespace( + fp8_sampling=False, + _policy=SimpleNamespace( + sampling_working_reserve_gib="auto", + sampling_wddm_hard_gib=1.0, + sampling_wddm_margin_gib=1.0, + ), + ) + + @contextlib.contextmanager + def sampling(**kwargs): + captured.update(kwargs) + yield + + runtime._executor = SimpleNamespace(sampling=sampling) + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime." + "allocator_cap.apply_wddm_hard_allocator_cap", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime.resolve_margin_gib", + lambda *_args, **_kwargs: 1.0, + ) + + with runtime.sampling_image( + shape_key=(768, 768), cold_working_bytes=2 * gib + ): + pass + + assert captured["cold_floor_bytes"] == gib + dequant + assert captured["hot_floor_bytes"] == int(1.25 * gib) + dequant + + +def test_bootstrap_uses_min_physical_free_and_one_gib_margin(): + gib = 1024 ** 3 + block_bytes = 200 * 1024 ** 2 + records = { + f"blocks.{index}": SimpleNamespace( + committed_bytes=block_bytes, + leaf_names=("linear",), + ) + for index in range(3) + } + plan = SimpleNamespace(resident_leaf_keys=frozenset()) + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._bootstrap_complete = False + runtime._bootstrap_min_free_bytes = 2 * gib + 450 * 1024 ** 2 + runtime._bootstrap_budget_bytes = 0 + runtime._bootstrap_block_keys = () + runtime._last_step_num = 50_000 + runtime._successful_training_steps = 1 + runtime._config = SimpleNamespace( + _policy=SimpleNamespace(wddm_hard_gib=1.0) + ) + runtime._model = SimpleNamespace() + runtime._arena = SimpleNamespace( + block_keys=lambda: tuple(records), + block_record=lambda key: records[key], + ) + runtime._residency = SimpleNamespace( + plan=plan, + resident_bytes=lambda: 0, + ) + runtime._training_plan = plan + runtime._smart_plan = {"singleton_resident_bytes": 100} + runtime._policy = ArenaResidencyController() + transitions = [] + runtime.transition_training_blocks = lambda keys, resident: ( + transitions.append((tuple(keys), resident)) + or { + "changed": True, + "block_keys": tuple(keys), + "plan": object(), + } + ) + + assert runtime._bootstrap_training_residency(10 * gib) is False + assert runtime._bootstrap_complete is False + + runtime._successful_training_steps = 2 + assert runtime._bootstrap_training_residency(10 * gib) is True + assert runtime._bootstrap_budget_bytes == 450 * 1024 ** 2 + assert transitions == [(("blocks.0", "blocks.1"), True)] + assert runtime._policy.pending_promotion["block_keys"] == ( + "blocks.0", + "blocks.1", + ) + assert runtime._policy.pending_promotion["resident_bytes_before"] == 100 + + +def test_bootstrap_ignores_first_runtime_sample_after_checkpoint_resume(): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._bootstrap_complete = False + runtime._bootstrap_min_free_bytes = None + runtime._successful_training_steps = 1 + + runtime.record_training_physical_free_min(123) + assert runtime._bootstrap_min_free_bytes is None + + runtime._successful_training_steps = 2 + runtime.record_training_physical_free_min(456) + assert runtime._bootstrap_min_free_bytes == 456 + + +def test_bootstrap_keeps_priority_over_four_block_fast_lane(): + gib = 1024 ** 3 + block_bytes = 100 * 1024 ** 2 + records = { + f"blocks.{index}": SimpleNamespace( + committed_bytes=block_bytes, + leaf_names=("linear",), + ) + for index in range(5) + } + plan = SimpleNamespace(resident_leaf_keys=frozenset()) + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._bootstrap_complete = False + runtime._bootstrap_min_free_bytes = 3 * gib + runtime._bootstrap_budget_bytes = 0 + runtime._bootstrap_block_keys = () + runtime._successful_training_steps = 2 + runtime._config = SimpleNamespace( + _policy=SimpleNamespace(wddm_hard_gib=1.0) + ) + runtime._arena = SimpleNamespace( + block_keys=lambda: tuple(records), + block_record=lambda key: records[key], + ) + runtime._residency = SimpleNamespace( + plan=plan, + resident_bytes=lambda: 0, + ) + runtime._training_plan = plan + runtime._smart_plan = {"singleton_resident_bytes": 0} + runtime._policy = ArenaResidencyController() + transitions = [] + runtime.transition_training_blocks = lambda keys, resident: ( + transitions.append((tuple(keys), resident)) + or {"changed": True, "block_keys": tuple(keys), "plan": object()} + ) + + assert runtime._bootstrap_training_residency(10 * gib) is True + assert transitions == [ + (("blocks.0", "blocks.1", "blocks.2", "blocks.3", "blocks.4"), True) + ] + assert runtime._bootstrap_complete is True + + +def test_arena_allocation_failure_drains_and_rolls_back(monkeypatch): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._device = "cpu" + runtime._config = SimpleNamespace( + _policy=SimpleNamespace(wddm_hard_gib=1.0) + ) + runtime._last_training_cap_target_bytes = None + runtime._signals = TrainingSignalWindow() + runtime._policy = ArenaResidencyController() + runtime._policy.pending_promotion = { + "block_key": "blocks.7", + "block_bytes": 20, + "resident_bytes_before": 200, + "resident_bytes_after": 220, + "previous_cap_target_bytes": 1000, + } + transitions = [] + runtime.transition_training_blocks = ( + lambda keys, resident: transitions.append((tuple(keys), resident)) + ) + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.transfer.drain_fetch_runtime", + lambda: 2, + ) + cap_calls = [] + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime." + "allocator_cap.apply_wddm_hard_allocator_cap", + lambda *args, **kwargs: cap_calls.append((args, kwargs)), + ) + + runtime._handle_training_failure( + __import__("torch").cuda.OutOfMemoryError("synthetic allocator OOM"), + shape_key=(768, 768), + step_num=65, + ) + + assert transitions == [(("blocks.7",), False)] + assert runtime._policy.last_safe_residency_bytes == 200 + assert runtime._policy.last_rejected_residency_bytes == 220 + assert runtime._last_failure_event["rollback_block"] == ["blocks.7"] + assert runtime._last_failure_event["abandoned_fetch_tickets"] == 2 + assert cap_calls[0][1]["target_cap_bytes"] == 1000 + + +def test_failed_training_step_preserves_original_error_if_cleanup_fails(): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._last_shape_key = None + runtime._last_step_num = None + runtime._device = "cpu" + runtime._last_policy_error = None + runtime._apply_training_policy = lambda: None + runtime._handle_training_failure = ( + lambda *_args, **_kwargs: (_ for _ in ()).throw( + RuntimeError("cleanup failed") + ) + ) + runtime._executor = SimpleNamespace( + TRAIN="train", + execution=lambda _mode: contextlib.nullcontext(), + ) + + with pytest.raises(ValueError, match="original"): + with runtime.training_step(shape_key=(768, 768), step_num=2): + raise ValueError("original") + assert "cleanup failed" in runtime._last_policy_error + + +def test_failed_training_step_does_not_publish_partial_peak(monkeypatch): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._last_shape_key = None + runtime._last_step_num = None + runtime._device = "cpu" + runtime._apply_training_policy = lambda: None + runtime._handle_training_failure = lambda *_args, **_kwargs: None + runtime._executor = SimpleNamespace( + TRAIN="train", + execution=lambda _mode: contextlib.nullcontext(), + ) + observed = [] + runtime._observe_training_step = lambda **kwargs: observed.append(kwargs) + + with pytest.raises(RuntimeError, match="synthetic OOM"): + with runtime.training_step(shape_key=(768, 768), step_num=2): + raise RuntimeError("synthetic OOM") + + assert observed == [] diff --git a/tests/test_canonical_arena.py b/tests/test_canonical_arena.py new file mode 100644 index 0000000000..7ca914a7aa --- /dev/null +++ b/tests/test_canonical_arena.py @@ -0,0 +1,173 @@ +import io +import unittest +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from toolkit.memory_management import pin_manager +from toolkit.memory_management.canonical_arena import ( + CanonicalArena, + CanonicalArenaError, +) +def _linear(in_f=8, out_f=4, bias=True): + layer = nn.Linear(in_f, out_f, bias=bias) + layer.weight.requires_grad_(False) + if bias: + layer.bias.requires_grad_(False) + return layer + + +class CanonicalizeTests(unittest.TestCase): + def test_canonicalize_repoints_params_into_one_flat_per_block(self): + a = _linear() + b = _linear() + arena = CanonicalArena() + try: + stats = arena.canonicalize({"blocks.0": [("lin_a", a), ("lin_b", b)]}) + self.assertEqual(stats.blocks, 1) + flat_ptr = arena.block_pack("blocks.0").host_flat.untyped_storage().data_ptr() + self.assertEqual(a.weight.untyped_storage().data_ptr(), flat_ptr) + self.assertEqual(b.weight.untyped_storage().data_ptr(), flat_ptr) + self.assertEqual(a.bias.untyped_storage().data_ptr(), flat_ptr) + self.assertTrue(arena.canonicalized) + finally: + arena.release() + + def test_double_canonicalize_is_rejected(self): + layer = _linear() + arena = CanonicalArena() + try: + arena.canonicalize({"blocks.0": [("lin", layer)]}) + with self.assertRaises(CanonicalArenaError): + arena.canonicalize({"blocks.0": [("lin", layer)]}) + finally: + arena.release() + + def test_trainable_leaf_is_rejected_and_nothing_repointed(self): + trainable = _linear() + trainable.weight.requires_grad_(True) + frozen = _linear() + original_frozen_ptr = frozen.weight.untyped_storage().data_ptr() + arena = CanonicalArena() + with self.assertRaises(CanonicalArenaError): + arena.canonicalize( + {"blocks.0": [("a", trainable)], "blocks.1": [("b", frozen)]} + ) + # Frozen block's leaf must never have been repointed: the frozen + # check runs over EVERY block before ANY block is built. + self.assertEqual(frozen.weight.untyped_storage().data_ptr(), original_frozen_ptr) + self.assertFalse(arena.canonicalized) + + def test_trainable_bias_is_rejected(self): + layer = _linear() + layer.bias.requires_grad_(True) + arena = CanonicalArena() + with self.assertRaises(CanonicalArenaError): + arena.canonicalize({"blocks.0": [("lin", layer)]}) + + def test_entries_without_module_are_rejected(self): + layer = _linear() + arena = CanonicalArena() + with self.assertRaises(CanonicalArenaError): + arena.canonicalize({"blocks.0": [("lin", layer.weight, layer.bias)]}) + + def test_state_dict_round_trip_after_canonicalize(self): + layer = _linear() + expected = {k: v.detach().clone() for k, v in layer.state_dict().items()} + arena = CanonicalArena() + try: + arena.canonicalize({"blocks.0": [("lin", layer)]}) + buffer = io.BytesIO() + torch.save(layer.state_dict(), buffer) + buffer.seek(0) + loaded = torch.load(buffer, weights_only=True) + for key, value in expected.items(): + self.assertTrue(torch.equal(value, loaded[key]), key) + # load_state_dict must copy IN PLACE, preserving the arena view + # (Parameter identity/storage unchanged -- Invariant 4). + flat_ptr = arena.block_pack("blocks.0").host_flat.untyped_storage().data_ptr() + layer.load_state_dict(loaded) + self.assertEqual(layer.weight.untyped_storage().data_ptr(), flat_ptr) + finally: + arena.release() + + def test_release_returns_pin_ledger_bytes(self): + layer = _linear(in_f=512, out_f=512, bias=True) + arena = CanonicalArena() + arena.canonicalize({"blocks.0": [("lin", layer)]}) + pinned = arena.committed_pinned_bytes() + self.assertGreater(pinned, 0) + before = pin_manager.pinned_bytes_by_kind().get("weights", 0) + arena.release() + after = pin_manager.pinned_bytes_by_kind().get("weights", 0) + # Ledger tracks page-rounded committed bytes (cudaHostRegister is + # page-granular); committed_pinned_bytes() reports the unpadded + # logical total, so release only guarantees the delta covers at + # LEAST the logical bytes, not an exact match. + self.assertGreaterEqual(before - after, pinned) + self.assertEqual(arena.committed_pinned_bytes(), 0) + self.assertEqual(arena.block_keys(), ()) + + def test_unknown_block_lookup_returns_none(self): + arena = CanonicalArena() + try: + arena.canonicalize({"blocks.0": [("lin", _linear())]}) + self.assertIsNone(arena.block_pack("blocks.1")) + self.assertIsNone(arena.block_record("blocks.1")) + finally: + arena.release() + + +class WholeModelToGuardTests(unittest.TestCase): + def test_guarded_to_raises(self): + model = nn.Sequential(_linear()) + arena = CanonicalArena() + try: + arena.canonicalize({"blocks.0": [("lin", model[0])]}) + CanonicalArena.guard_whole_model_to(model) + with self.assertRaises(CanonicalArenaError): + model.to(torch.device("cpu")) + finally: + CanonicalArena.unguard_whole_model_to(model) + arena.release() + + def test_guard_is_idempotent(self): + model = nn.Sequential(_linear()) + CanonicalArena.guard_whole_model_to(model) + original = model.to + CanonicalArena.guard_whole_model_to(model) + self.assertIs(model.to, original) + CanonicalArena.unguard_whole_model_to(model) + + def test_guarded_to_allows_only_idempotent_runtime_placement(self): + model = nn.Sequential(_linear()) + model._arena_offload_runtime = SimpleNamespace( + _permanent_placement=(torch.device("cpu"), torch.float32) + ) + CanonicalArena.guard_whole_model_to(model) + try: + self.assertIs(model.to(torch.device("cpu")), model) + self.assertIs( + model.to(device=torch.device("cpu"), dtype=torch.float32), model + ) + with self.assertRaises(CanonicalArenaError): + model.to(device=torch.device("cpu"), dtype=torch.float64) + finally: + CanonicalArena.unguard_whole_model_to(model) + del model._arena_offload_runtime + + def test_unguard_restores_normal_to(self): + model = nn.Sequential(_linear()) + CanonicalArena.guard_whole_model_to(model) + CanonicalArena.unguard_whole_model_to(model) + # Ordinary .to() must work again (no canonicalized leaves here). + model.to(torch.device("cpu")) + + def test_unguarded_model_never_raises(self): + model = nn.Sequential(_linear()) + model.to(torch.device("cpu")) # sanity: no guard installed, no raise + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_generic_block_dispatcher.py b/tests/test_generic_block_dispatcher.py new file mode 100644 index 0000000000..e23617a068 --- /dev/null +++ b/tests/test_generic_block_dispatcher.py @@ -0,0 +1,482 @@ +from types import MethodType +from unittest import mock +from dataclasses import replace + +import pytest +import torch +from torch.utils.checkpoint import checkpoint + +from toolkit.memory_management.arena_offload import ( + ArenaOffloadConfig, + close_arena_offload, + discover_blocks, + prepare_arena_offload, +) +from toolkit.memory_management.arena_offload.discovery import BlockDiscoveryError +from toolkit.memory_management.arena_offload.ownership import active_process_owner +from toolkit.memory_management.residency import ResidencyPlan +from toolkit.memory_management.runtime import get_memory_runtime + + +class _Block(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(4, 4) + + def forward(self, value): + return torch.nn.functional.silu(self.proj(value)) + + +class _Transformer(torch.nn.Module): + def __init__(self, count=3): + super().__init__() + self.blocks = torch.nn.ModuleList([_Block() for _ in range(count)]) + self.gradient_checkpointing = False + self._checkpoint_keep_last = 0 + + def enable_gradient_checkpointing(self, keep_last=0): + self.gradient_checkpointing = True + self._checkpoint_keep_last = int(keep_last) + + def forward(self, value): + cutoff = len(self.blocks) - self._checkpoint_keep_last + for index, block in enumerate(self.blocks): + if self.gradient_checkpointing and torch.is_grad_enabled() and index < cutoff: + value = checkpoint(block, value, use_reentrant=False) + else: + value = block(value) + return value + + +def _frozen_transformer(): + model = _Transformer() + model.requires_grad_(False) + return model + + +def _fp8_transformer(device, count=3, width=32): + from optimum.quanto import freeze + + from toolkit.util.quantize import get_qtype, quantize + + class Fp8Block(torch.nn.Module): + def __init__(self): + super().__init__() + self.proj = torch.nn.Linear(width, width, bias=False) + + def forward(self, value): + return torch.nn.functional.silu(self.proj(value)) + + class Fp8Transformer(torch.nn.Module): + def __init__(self): + super().__init__() + self.first = torch.nn.Linear(width, width, bias=False) + self.blocks = torch.nn.ModuleList(Fp8Block() for _ in range(count)) + self.gradient_checkpointing = True + self._checkpoint_keep_last = 1 + + def forward(self, value): + value = self.first(value) + cutoff = len(self.blocks) - self._checkpoint_keep_last + for index, block in enumerate(self.blocks): + if torch.is_grad_enabled() and index < cutoff: + value = checkpoint(block, value, use_reentrant=False) + else: + value = block(value) + return value + + model = Fp8Transformer().to(device=device, dtype=torch.bfloat16) + quantize(model, weights=get_qtype("float8")) + freeze(model) + return model + + +def _fp8_runtime(model, device, *, forward, backward, compile_blocks): + config = ArenaOffloadConfig( + enabled=True, + fp8_forward=forward, + fp8_backward=backward, + compile_blocks=compile_blocks, + _compile_dynamic=False, + ) + config = replace( + config, + _policy=replace( + config._policy, + working_reserve_gib=0.0, + wddm_margin_gib=0.0, + wddm_hard_gib=1.0, + checkpoint_keep_last=1, + ), + ) + with mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(1 * 1024**3, 12 * 1024**3), + ): + return prepare_arena_offload( + model, + device=device, + block_names=("blocks",), + config=config, + ) + + +def _fp8_train_once(model, runtime, device, step, network=None): + import contextlib + + value = torch.randn( + 2, 3, 32, device=device, dtype=torch.bfloat16, requires_grad=True + ) + network_context = network if network is not None else contextlib.nullcontext() + with runtime.training_step(shape_key=(2, 3, 32), step_num=step), network_context: + output = model(value) + output.float().square().mean().backward() + assert value.grad is not None + if network is not None: + assert all(parameter.grad is not None for parameter in network.parameters()) + return output.detach() + + +class _AdapterBaseModel: + arch = "synthetic" + use_old_lokr_format = False + + def __init__(self, device, dtype): + self.device_torch = torch.device(device) + self.torch_dtype = dtype + + def get_transformer_block_names(self): + return None + + +def _apply_lora(model, device, dtype): + from toolkit.config_modules import NetworkConfig + from toolkit.lora_special import LoRASpecialNetwork + + config = NetworkConfig( + type="lora", linear=8, linear_alpha=8.0, transformer_only=True + ) + network = LoRASpecialNetwork( + text_encoder=None, + unet=model, + lora_dim=8, + multiplier=1.0, + alpha=8.0, + train_unet=True, + train_text_encoder=False, + network_config=config, + network_type=config.type, + transformer_only=True, + is_transformer=True, + target_lin_modules=[model.__class__.__name__], + base_model=_AdapterBaseModel(device, dtype), + ) + network.force_to(device, dtype=torch.float32) + network._update_torch_multiplier() + network.apply_to(None, model, apply_text_encoder=False, apply_unet=True) + network.can_merge_in = False + network.prepare_grad_etc(None, model) + return network + + +def test_declared_container_discovery_accounts_all_block_state(): + model = _frozen_transformer() + selection = discover_blocks(model, container_paths=("blocks",)) + assert selection.container_paths == ("blocks",) + assert selection.block_keys == ("blocks.0", "blocks.1", "blocks.2") + assert all(len(entries) == 1 for entries in selection.entries_by_block.values()) + assert selection.accounting.managed_entries == 6 + assert selection.accounting.managed_bytes > 0 + + +def test_shared_managed_state_is_rejected_before_construction(): + model = _frozen_transformer() + shared = model.blocks[0].proj.weight + model.blocks[1].proj.weight = shared + with pytest.raises(BlockDiscoveryError, match="shared_managed"): + discover_blocks(model, container_paths=("blocks",)) + + +def test_checkpointing_rejection_precedes_canonical_commit(): + model = _frozen_transformer() + original = model.blocks[0].proj.weight + with pytest.raises(ValueError, match="requires model gradient checkpointing"): + prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=ArenaOffloadConfig(enabled=True), + ) + assert model.blocks[0].proj.weight is original + assert active_process_owner() is None + assert not hasattr(model, "_arena_offload_runtime") + + +def test_saved_installed_forward_checkpoint_backward_and_teardown(): + torch.manual_seed(17) + model = _frozen_transformer() + model.enable_gradient_checkpointing(keep_last=1) + reference_input = torch.randn(2, 4) + reference = model(reference_input).detach() + with mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(8 * 1024**3, 12 * 1024**3), + ), mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.auto_margin_gib", + return_value=1.0, + ): + config = ArenaOffloadConfig(enabled=True, compile_blocks=False) + config = replace( + config, + _policy=replace(config._policy, checkpoint_keep_last=1), + ) + runtime = prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=config, + ) + installed = [] + adapters = [] + for block in model.blocks: + saved = block.forward + block.adapter_gain = torch.nn.Parameter(torch.zeros(())) + + def installed_forward(self, value, _saved=saved): + return _saved(value) + self.adapter_gain * value + + bound = MethodType(installed_forward, block) + block.forward = bound + installed.append(bound) + adapters.append(block.adapter_gain) + + runtime.finalize() + diagnostics = runtime.diagnostics() + accounting = diagnostics["accounting"] + assert diagnostics["checkpoint_owner"] == "model" + assert diagnostics["state_audit"]["managed_entries"] == 6 + assert accounting["payload_reconciled"] + assert accounting["canonical_payload_bytes"] == ( + accounting["canonical_resident_payload_bytes"] + + accounting["streamed_payload_bytes"] + ) + assert accounting["protected_training_blocks"] == ("blocks.2",) + assert accounting["protected_training_blocks_resident"] + with pytest.raises(RuntimeError, match="outside_transformer_execution"): + model.blocks[0](torch.randn(2, 4)) + + value = reference_input.detach().clone().requires_grad_(True) + with runtime.training_step(shape_key=(2, 4), step_num=1): + output = model(value) + output.sum().backward() + torch.testing.assert_close(output.detach(), reference) + assert value.grad is not None + assert all(parameter.grad is not None for parameter in adapters) + + protected = runtime._executor.protected_training_leaf_keys + assert any(block == "blocks.2" for block, _leaf in protected) + with pytest.raises(RuntimeError, match="protected_training_block"): + runtime.transition_training_block("blocks.2", resident=False) + + close_arena_offload(model) + assert all( + block.forward is saved + for block, saved in zip(model.blocks, installed, strict=True) + ) + assert get_memory_runtime(model) is None + assert model._arena_offload_disposed + assert active_process_owner() is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cuda_streamed_compiled_train_sample_train(): + from toolkit.memory_management.arena_offload import transfer + + torch.manual_seed(23) + device = torch.device("cuda") + model = _frozen_transformer().to(device) + model.enable_gradient_checkpointing(keep_last=1) + config = ArenaOffloadConfig( + enabled=True, + compile_blocks=True, + _compile_dynamic=False, + ) + config = replace( + config, + _policy=replace( + config._policy, + working_reserve_gib=0.0, + wddm_margin_gib=0.0, + wddm_hard_gib=1.0, + checkpoint_keep_last=1, + ), + ) + with mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(1 * 1024**3, 12 * 1024**3), + ): + runtime = prepare_arena_offload( + model, + device=device, + block_names=("blocks",), + config=config, + ) + adapters = [] + for block in model.blocks: + saved = block.forward + block.adapter_gain = torch.nn.Parameter(torch.zeros((), device=device)) + + def installed_forward(self, value, _saved=saved): + return _saved(value) + self.adapter_gain * value + + block.forward = MethodType(installed_forward, block) + adapters.append(block.adapter_gain) + runtime.finalize() + accounting = runtime.diagnostics()["accounting"] + assert accounting["payload_reconciled"] + assert accounting["mixed_residency"] + assert accounting["resident_blocks"] >= 1 + assert accounting["streamed_blocks"] >= 1 + assert accounting["planned_training_h2d_bytes"] == ( + 2 * accounting["planned_forward_h2d_bytes"] + ) + assert accounting["protected_training_blocks_resident"] + streamed = [ + runtime._executor.source(index).transfer is not None + for index in range(runtime.block_count) + ] + assert any(streamed) + assert not streamed[-1] + + def train_once(step): + transfer_before = transfer.lifetime_fetch_stats()["bytes"] + planned = runtime.diagnostics()["accounting"][ + "planned_training_h2d_bytes" + ] + value = torch.randn(2, 4, device=device, requires_grad=True) + with runtime.training_step(shape_key=(2, 4), step_num=step): + output = model(value) + output.square().mean().backward() + assert value.grad is not None + assert all(parameter.grad is not None for parameter in adapters) + for parameter in adapters: + parameter.grad = None + assert transfer.lifetime_fetch_stats()["bytes"] - transfer_before == planned + return output.detach() + + first = train_once(1) + sample_plan = ResidencyPlan.build("sample", ()) + runtime._executor.activate(runtime._executor.SAMPLE, sample_plan) + sample_accounting = runtime.diagnostics()["accounting"] + sample_transfer_before = transfer.lifetime_fetch_stats()["bytes"] + with torch.no_grad(), runtime._executor.execution(runtime._executor.SAMPLE): + sampled = model(torch.randn(2, 4, device=device)) + assert transfer.lifetime_fetch_stats()["bytes"] - sample_transfer_before == ( + sample_accounting["planned_forward_h2d_bytes"] + ) + assert torch.isfinite(sampled).all() + runtime._executor.activate(runtime._executor.TRAIN, runtime._training_plan) + second = train_once(2) + assert torch.isfinite(first).all() + assert torch.isfinite(second).all() + close_arena_offload(model) + assert active_process_owner() is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cuda_fp8_gates_select_distinct_canonical_arena_paths(): + from toolkit.quantization import fp8_linear + + device = torch.device("cuda") + real_scaled_mm = torch._scaled_mm + real_grad_input = fp8_linear._grad_input_compute + results = {} + + for name, forward, backward in ( + ("baseline", False, False), + ("forward", True, False), + ("forward_backward", True, True), + ): + scaled_calls = [] + grad_input_calls = [] + + def counted_scaled_mm(*args, **kwargs): + scaled_calls.append(1) + return real_scaled_mm(*args, **kwargs) + + def counted_grad_input(*args, **kwargs): + grad_input_calls.append(1) + return real_grad_input(*args, **kwargs) + + model = _fp8_transformer(device) + runtime = _fp8_runtime( + model, + device, + forward=forward, + backward=backward, + compile_blocks=False, + ) + try: + with mock.patch.object(torch, "_scaled_mm", counted_scaled_mm), mock.patch.object( + fp8_linear, "_grad_input_compute", counted_grad_input + ): + runtime.finalize() + _fp8_train_once(model, runtime, device, 1) + diagnostics = runtime.diagnostics() + results[name] = { + "scaled": len(scaled_calls), + "grad_input": len(grad_input_calls), + "canonical": diagnostics["training_fp8_canonical"], + "singletons": diagnostics["training_fp8_singletons"], + } + finally: + close_arena_offload(model) + + assert results["baseline"] == { + "scaled": 0, + "grad_input": 0, + "canonical": 0, + "singletons": 0, + } + assert results["forward"]["scaled"] > 0 + assert results["forward"]["grad_input"] == 0 + assert results["forward"]["canonical"] == 3 + assert results["forward"]["singletons"] == 1 + assert results["forward_backward"]["scaled"] > 0 + assert results["forward_backward"]["grad_input"] > 0 + assert results["forward_backward"]["canonical"] == 3 + assert results["forward_backward"]["singletons"] == 1 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_cuda_compiled_fp8_canonical_arena_emits_scaled_mm(): + device = torch.device("cuda") + model = _fp8_transformer(device) + runtime = _fp8_runtime( + model, + device, + forward=True, + backward=True, + compile_blocks=True, + ) + network = _apply_lora(model, device, torch.bfloat16) + try: + runtime.finalize(network) + _fp8_train_once(model, runtime, device, 1, network) + for parameter in network.parameters(): + parameter.grad = None + with torch.profiler.profile( + activities=[torch.profiler.ProfilerActivity.CUDA] + ) as profiler: + _fp8_train_once(model, runtime, device, 2, network) + scaled_mm = [ + event + for event in profiler.key_averages() + if "_scaled_mm" in event.key + ] + assert sum(event.count for event in scaled_mm) > 0 + assert runtime.diagnostics()["training_fp8_canonical"] == 3 + assert runtime.diagnostics()["training_fp8_singletons"] == 1 + finally: + close_arena_offload(model) + torch.cuda.empty_cache() diff --git a/tests/test_residency.py b/tests/test_residency.py new file mode 100644 index 0000000000..13d2b9f963 --- /dev/null +++ b/tests/test_residency.py @@ -0,0 +1,207 @@ +from types import SimpleNamespace + +import pytest +import torch + +from toolkit.memory_management import pin_manager +from toolkit.memory_management.canonical_arena import CanonicalArena +from toolkit.memory_management.immutable_runtime import ImmutableTransformerRuntime +from toolkit.memory_management.residency import ( + ResidencyError, + ResidencyPlan, + ResidencyState, +) + + + + +def _linear(seed=0, *, device="cpu", dtype=torch.float32): + torch.manual_seed(seed) + layer = torch.nn.Linear(8, 8, bias=True, device=device, dtype=dtype) + layer.requires_grad_(False) + return layer + + +@pytest.fixture +def arena_layers(): + layers = {"a": _linear(1), "b": _linear(2), "c": _linear(3)} + arena = CanonicalArena() + arena.canonicalize({"blocks.0": list(layers.items())}) + try: + yield arena, layers + finally: + arena.release() +def test_phase_plan_and_existing_planner_seam(arena_layers): + arena, layers = arena_layers + smart = {"offload_ids": {id(layers["a"]), id(layers["c"])}} + plan = ResidencyPlan.from_smart_plan(arena, smart, phase="train") + # Immutable policy normalizes any partially offloaded canonical block to + # fully streamed. Source snapshots may still represent mixed layouts, but + # controller/planner output is whole-block. + assert plan.resident_leaf_keys == frozenset() + assert plan == ResidencyPlan.from_smart_plan(arena, smart, phase="train") + assert plan.fingerprint != ResidencyPlan.build("sample", plan.resident_leaf_keys).fingerprint + + +def test_runtime_training_transitions_are_whole_block(arena_layers): + arena, layers = arena_layers + block = SimpleNamespace(entries=tuple(layers.items())) + model = SimpleNamespace(blocks=(block,)) + state = ResidencyState(arena, "cpu") + state.reconcile(ResidencyPlan.build("train", ())) + runtime = ImmutableTransformerRuntime( + model, + state, + blocks=model.blocks, + block_keys=("blocks.0",), + entries_by_block={"blocks.0": tuple(layers.items())}, + compile_blocks=False, + ) + + growth = runtime.increase_training_residency( + arena.block_record("blocks.0").committed_bytes, + ) + expected = frozenset( + ("blocks.0", leaf_name) for leaf_name in layers + ) + assert state.plan.resident_leaf_keys == expected + assert growth["added_blocks"] == ("blocks.0",) + + relief = runtime.reduce_training_residency(1) + assert relief["removed_blocks"] == ("blocks.0",) + assert state.plan.resident_leaf_keys == frozenset() + +def test_exact_training_block_transaction_uses_stable_key(arena_layers): + arena, layers = arena_layers + block = SimpleNamespace(entries=tuple(layers.items())) + model = SimpleNamespace(blocks=(block,)) + state = ResidencyState(arena, "cpu") + state.reconcile(ResidencyPlan.build("train", ())) + runtime = ImmutableTransformerRuntime( + model, + state, + blocks=model.blocks, + block_keys=("blocks.0",), + entries_by_block={"blocks.0": tuple(layers.items())}, + compile_blocks=False, + ) + + promoted = runtime.transition_training_block("blocks.0", resident=True) + expected = frozenset(("blocks.0", name) for name in layers) + assert promoted["changed"] is True + assert state.plan.resident_leaf_keys == expected + + unchanged = runtime.transition_training_block("blocks.0", resident=True) + assert unchanged["changed"] is False + + demoted = runtime.transition_training_block("blocks.0", resident=False) + assert demoted["changed"] is True + assert state.plan.resident_leaf_keys == frozenset() + + +def test_cpu_reconcile_never_mutates_parameters_or_pin_ledger(arena_layers): + arena, layers = arena_layers + state = ResidencyState(arena, "cpu") + identities = {name: id(layer.weight) for name, layer in layers.items()} + pointers = { + name: layer.weight.untyped_storage().data_ptr() for name, layer in layers.items() + } + pins = (pin_manager.total_pinned_bytes(), pin_manager.pinned_bytes_by_kind()) + + delta = state.reconcile( + ResidencyPlan.build("train", (("blocks.0", "a"), ("blocks.0", "c"))) + ) + assert delta.promoted == (("blocks.0", "a"), ("blocks.0", "c")) + assert state.streamed_leaf_names("blocks.0") == ("b",) + state.reconcile(ResidencyPlan.build("sample", (("blocks.0", "b"),))) + state.clear() + + assert pins == (pin_manager.total_pinned_bytes(), pin_manager.pinned_bytes_by_kind()) + assert identities == {name: id(layer.weight) for name, layer in layers.items()} + assert pointers == { + name: layer.weight.untyped_storage().data_ptr() for name, layer in layers.items() + } + + +def test_failed_multi_promotion_rolls_back_atomically(arena_layers, monkeypatch): + arena, _layers = arena_layers + state = ResidencyState(arena, "cpu") + state.promote(("blocks.0", "c"), phase="seed") + before_plan = state.plan + before_sidecar = state.resident_leaf(("blocks.0", "c")) + before_pins = (pin_manager.total_pinned_bytes(), pin_manager.pinned_bytes_by_kind()) + original = state._build_sidecar + calls = 0 + + def fail_second(key): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("synthetic promotion failure") + return original(key) + + monkeypatch.setattr(state, "_build_sidecar", fail_second) + target = ResidencyPlan.build( + "train", (("blocks.0", "a"), ("blocks.0", "b"), ("blocks.0", "c")) + ) + with pytest.raises(RuntimeError, match="synthetic promotion failure"): + state.reconcile(target) + + assert state.plan is before_plan + assert state.resident_leaf(("blocks.0", "c")) is before_sidecar + assert state.resident_leaf(("blocks.0", "a")) is None + assert state.resident_leaf(("blocks.0", "b")) is None + assert before_pins == (pin_manager.total_pinned_bytes(), pin_manager.pinned_bytes_by_kind()) + + +def test_unknown_leaf_fails_before_state_change(arena_layers): + arena, _layers = arena_layers + state = ResidencyState(arena, "cpu") + with pytest.raises(ResidencyError, match="unknown_residency_leaf"): + state.reconcile(ResidencyPlan.build("train", (("blocks.0", "missing"),))) + assert state.resident_bytes() == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_cuda_sidecars_match_canonical_values_and_demote_without_writeback(arena_layers): + arena, layers = arena_layers + state = ResidencyState(arena, "cuda") + host_before = {name: layer.weight.detach().clone() for name, layer in layers.items()} + state.reconcile( + ResidencyPlan.build("train", (("blocks.0", "a"), ("blocks.0", "c"))) + ) + for name in ("a", "c"): + sidecar = state.resident_leaf(("blocks.0", name)) + assert sidecar.weight.device.type == "cuda" + torch.testing.assert_close(sidecar.weight.cpu(), host_before[name]) + state.demote(("blocks.0", "a")) + torch.cuda.synchronize() + # Frozen demotion drops the device copy and performs no D2H write-back. + torch.testing.assert_close(layers["a"].weight, host_before["a"]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_fp8_sidecar_preserves_wrapper_type_and_values(): + from optimum.quanto import freeze + + from toolkit.util.quantize import get_qtype, quantize + + model = torch.nn.Sequential(torch.nn.Linear(16, 16, bias=False).to(torch.bfloat16)) + quantize(model, weights=get_qtype("qfloat8")) + freeze(model) + layer = model[0] + layer.weight.requires_grad_(False) + expected_type = type(layer.weight.data) + expected = layer.weight.data.dequantize().clone() + arena = CanonicalArena() + arena.canonicalize({"blocks.0": [("q", layer)]}) + try: + state = ResidencyState(arena, "cuda") + state.promote(("blocks.0", "q"), phase="train") + weight = state.resident_tensor(("blocks.0", "q")) + assert type(weight) is expected_type + assert all(leaf.device.type == "cuda" for leaf in weight.__tensor_flatten__()[0] + for leaf in [getattr(weight, leaf)]) + torch.testing.assert_close(weight.dequantize().cpu(), expected) + finally: + arena.release() diff --git a/tests/test_residency_two_timescale.py b/tests/test_residency_two_timescale.py new file mode 100644 index 0000000000..1cdb644f00 --- /dev/null +++ b/tests/test_residency_two_timescale.py @@ -0,0 +1,218 @@ +"""Pure-CPU coverage for the two-timescale residency controller. + +Validates the cap/allowance arithmetic against the measured Krea2 cap-descent +numbers (7.35 GiB floor, 6.77 GiB live, 0.21 GiB knee slack, ~0.375 GiB block) +and the hysteresis FSM's transitions. No CUDA -- these are the reason the +controller can be trusted without GPU CI (see the plan / cuda-testing policy). +""" + +import pytest + +from toolkit.memory_management import vram_budget as vb + +GIB = vb.GIB + + +def gib(x): + return int(x * GIB) + + +# --- Arithmetic vs the measured cap-descent ------------------------------- + +def test_allowance_matches_measured_knee(): + # 0.95*7.35 - 6.77 = +0.21 GiB (clean, the floor cap); 0.95*7.10 - 6.77 = + # -0.02 GiB (dirty). The model validated to within one 0.25 notch. + clean = vb.allocator_allowance_bytes(gib(7.35), gib(6.77)) + dirty = vb.allocator_allowance_bytes(gib(7.10), gib(6.77)) + assert clean == pytest.approx(gib(0.2125), abs=gib(0.01)) + assert dirty < 0 + + +def test_cap_for_live_is_allowance_inverse(): + # To host 6.77 live + 0.21 cache budget the cap must be ~7.35 (the floor). + cap = vb.cap_bytes_for_live(gib(6.77), gib(0.21), cliff_cap_bytes=gib(9.85)) + assert cap == pytest.approx(gib(7.35), abs=gib(0.02)) + # Round-trips: that cap yields ~the requested cache budget back. + assert vb.allocator_allowance_bytes(cap, gib(6.77)) == pytest.approx(gib(0.21), abs=gib(0.01)) + + +def test_cap_for_live_clamped_to_cliff(): + cap = vb.cap_bytes_for_live(gib(11.0), gib(2.0), cliff_cap_bytes=gib(9.85)) + assert cap == gib(9.85) + + +def test_promotion_precheck_uses_the_0p95_divisor(): + # need_cap = (6.77 + 0.375 + 0.21) / 0.95 = 7.742 GiB. + live, block, slack = gib(6.77), gib(0.375), gib(0.21) + # Sampling: cliff ~9.85 has room -> cap lever can fund the block. + assert vb.cap_can_host_promotion(live, block, slack, gib(9.85)) is True + # Training-like: cap pinned at the 7.35 floor/cliff -> must demote instead. + assert vb.cap_can_host_promotion(live, block, slack, gib(7.35)) is False + # The naive (no /0.95) test would wrongly pass at cliff = 7.36 + # (6.77+0.375+0.21 = 7.355 < 7.36); the real need_cap 7.742 rejects it. + assert vb.cap_can_host_promotion(live, block, slack, gib(7.36)) is False + + +def test_promote_gate_needs_zero_retries_and_a_block_of_slack(): + block, slack = gib(0.375), gib(0.21) # need > 0.585 GiB reclaimable + assert vb.residency_promote_ok(0, gib(0.6), block, slack) is True + assert vb.residency_promote_ok(0, gib(0.5), block, slack) is False # not enough slack + assert vb.residency_promote_ok(1, gib(0.6), block, slack) is False # retries bind + + +# --- FSM transitions ------------------------------------------------------- + +def drive(state, signal, **kw): + return vb.residency_fsm_step(state, signal, **kw) + + +CLEAN = {"binding": False} +BIND = {"binding": True} + + +def test_cold_settles_to_stable_after_k_clean(): + s = vb.ResidencyFsmState() # COLD + s, a = drive(s, CLEAN, k_clean=2) + assert s.name == vb.FSM_COLD and a == vb.ACT_HOLD + s, a = drive(s, CLEAN, k_clean=2) + assert s.name == vb.FSM_STABLE + + +def _stable(): + s = vb.ResidencyFsmState() + for _ in range(2): + s, _ = drive(s, CLEAN, k_clean=2) + assert s.name == vb.FSM_STABLE + return s + + +def test_stable_pressure_raises_cap_when_it_can_relieve(): + s = _stable() + s, a = drive(s, {"binding": True, "cap_can_relieve": True}) + assert s.name == vb.FSM_CAP_VERIFY and a == vb.ACT_RAISE_CAP + + +def test_stable_pressure_demotes_when_cap_pinned(): + s = _stable() + s, a = drive(s, {"binding": True, "cap_can_relieve": False}) + assert s.name == vb.FSM_COLD and a == vb.ACT_DEMOTE + + +def test_promotion_direct_when_cap_already_covers(): + s = _stable() + # eligibility needs k_clean clean windows in STABLE first + s, _ = drive(s, {"binding": False, "promote_gate": True, "cap_covers_promo": True}, k_clean=2) + assert s.name == vb.FSM_STABLE # w=1 < k_clean + s, a = drive(s, {"binding": False, "promote_gate": True, "cap_covers_promo": True}, k_clean=2) + assert s.name == vb.FSM_PROMOTION_VERIFY and a == vb.ACT_PROMOTE + + +def test_promotion_prefunds_cap_when_needed(): + s = _stable() + sig = {"binding": False, "promote_gate": True, "cap_covers_promo": False} + s, _ = drive(s, sig, k_clean=2) + s, a = drive(s, sig, k_clean=2) + assert s.name == vb.FSM_CAP_VERIFY and a == vb.ACT_RAISE_CAP + + +def test_cap_verify_clean_returns_to_stable(): + s = vb.ResidencyFsmState(vb.FSM_CAP_VERIFY, 0) + s, _ = drive(s, CLEAN, k_verify=2) + assert s.name == vb.FSM_CAP_VERIFY + s, _ = drive(s, CLEAN, k_verify=2) + assert s.name == vb.FSM_STABLE + + +def test_cap_verify_pressure_escalates_to_demote(): + s = vb.ResidencyFsmState(vb.FSM_CAP_VERIFY, 0) + s, a = drive(s, BIND) + assert s.name == vb.FSM_COLD and a == vb.ACT_DEMOTE + + +def test_promotion_verify_rolls_back_binding_first_window(): + s = vb.ResidencyFsmState(vb.FSM_PROMOTION_VERIFY, 0) + s, a = drive(s, BIND, k_verify=2) + assert s.name == vb.FSM_COOLDOWN and a == vb.ACT_ROLLBACK + + +def test_promotion_verify_clean_commits(): + s = vb.ResidencyFsmState(vb.FSM_PROMOTION_VERIFY, 0) + for _ in range(3): # first window ignored, then k_verify clean + s, _ = drive(s, CLEAN, k_verify=2) + assert s.name == vb.FSM_STABLE + + +def test_cooldown_bars_repromotion_then_releases(): + s = vb.ResidencyFsmState(vb.FSM_COOLDOWN, 0) + for _ in range(3): + s, _ = drive(s, {"binding": False, "promote_gate": True, "cap_covers_promo": True}, cooldown_n=4) + assert s.name == vb.FSM_COOLDOWN # still barred, no promote + s, _ = drive(s, CLEAN, cooldown_n=4) + assert s.name == vb.FSM_STABLE + + +def test_measurements_invalid_forces_cold(): + s = _stable() + s, a = drive(s, {"measurements_invalid": True, "binding": False}) + assert s.name == vb.FSM_COLD and a == vb.ACT_HOLD + + +def test_no_oscillation_rollback_then_cooldown_holds(): + # A failed promotion must not immediately re-promote: rollback -> cooldown + # bars promotion for cooldown_n windows even if the gate keeps firing. + s = vb.ResidencyFsmState(vb.FSM_PROMOTION_VERIFY, 1) # past the cold window + s, a = drive(s, BIND) + assert (s.name, a) == (vb.FSM_COOLDOWN, vb.ACT_ROLLBACK) + promotes = 0 + for _ in range(3): + s, a = drive(s, {"binding": False, "promote_gate": True, "cap_covers_promo": True}, cooldown_n=4) + promotes += a == vb.ACT_PROMOTE + assert promotes == 0 + + +# --- Training worst-resolution promotion guard ----------------------------- + +def test_worst_shape_free_subtracts_all_cohabitants_plus_block(): + # resident 5.0 + block 0.375 + ring 0.5 + worst reserve 2.5 + other 1.15 + # = 9.525 used; on a 12 GiB card that leaves 2.475 free. + free = vb.training_promotion_worst_shape_free_gib( + resident_gib=5.0, + added_block_gib=0.375, + ring_gib=0.5, + worst_working_reserve_gib=2.5, + other_gib=1.15, + total_gib=12.0, + ) + assert free == pytest.approx(2.475, abs=1e-6) + + +def test_worst_shape_gate_vetoes_when_high_res_reserve_would_page(): + # Same layout, but the worst measured resolution needs a 3.6 GiB reserve. + # A promotion decided on a roomy low-res step (that only needed ~1.7) would + # push the high-res cohabitation peak past the promote floor -> veto. + hold_high = 2.0 + roomy_lowres = vb.training_promotion_worst_shape_free_gib( + resident_gib=5.0, added_block_gib=0.375, ring_gib=0.5, + worst_working_reserve_gib=1.7, other_gib=1.15, total_gib=12.0, + ) + worst_highres = vb.training_promotion_worst_shape_free_gib( + resident_gib=5.0, added_block_gib=0.375, ring_gib=0.5, + worst_working_reserve_gib=3.6, other_gib=1.15, total_gib=12.0, + ) + # The current (low-res) step looks safe; the worst measured shape does not. + assert roomy_lowres >= hold_high + assert worst_highres < hold_high + + +def test_worst_shape_free_is_conservative_about_the_block(): + # Adding the block can only lower the predicted free (never hidden by a + # shrinking ring) -- the from-below assumption. + without_block = vb.training_promotion_worst_shape_free_gib( + resident_gib=5.0, added_block_gib=0.0, ring_gib=0.5, + worst_working_reserve_gib=2.5, other_gib=1.15, total_gib=12.0, + ) + with_block = vb.training_promotion_worst_shape_free_gib( + resident_gib=5.0, added_block_gib=0.375, ring_gib=0.5, + worst_working_reserve_gib=2.5, other_gib=1.15, total_gib=12.0, + ) + assert with_block == pytest.approx(without_block - 0.375, abs=1e-6) diff --git a/tests/test_transfer_plan.py b/tests/test_transfer_plan.py new file mode 100644 index 0000000000..5d7358019a --- /dev/null +++ b/tests/test_transfer_plan.py @@ -0,0 +1,141 @@ +import unittest +from itertools import pairwise + +import torch +import torch.nn as nn + +from toolkit.memory_management.canonical_arena import CanonicalArena +from toolkit.memory_management.transfer_plan import ( + TransferPlanError, + build_transfer_plan, +) + + +def _linear(in_f=8, out_f=4, bias=True): + layer = nn.Linear(in_f, out_f, bias=bias) + layer.weight.requires_grad_(False) + if bias: + layer.bias.requires_grad_(False) + return layer + + +@unittest.skipUnless(torch.cuda.is_available(), "canonical arena pinning requires CUDA") +class BuildTransferPlanTests(unittest.TestCase): + def setUp(self): + self.arena = CanonicalArena() + self.a = _linear() + self.b = _linear() + self.c = _linear() + self.arena.canonicalize( + {"blocks.0": [("a", self.a), ("b", self.b), ("c", self.c)]} + ) + self.addCleanup(self.arena.release) + + def test_fully_streamed_block_is_one_coalesced_range(self): + record = self.arena.block_record("blocks.0") + plan = build_transfer_plan(record, ["a", "b", "c"]) + self.assertTrue(plan.fully_streamed) + self.assertEqual(plan.num_ranges, 1) + self.assertEqual(plan.ranges[0].nbytes, plan.compact_nbytes) + # Weight+bias for all 3 leaves must be present in the compact layout. + for name in ("a", "b", "c"): + self.assertIn("weight", plan.leaf_specs[name]) + self.assertIn("bias", plan.leaf_specs[name]) + + def test_partial_streaming_breaks_coalescing_at_resident_leaf(self): + record = self.arena.block_record("blocks.0") + plan = build_transfer_plan(record, ["a", "c"]) # b resident, breaks the run + self.assertFalse(plan.fully_streamed) + self.assertEqual(plan.num_ranges, 2) + self.assertEqual(set(plan.leaf_specs.keys()), {"a", "c"}) + + def test_compact_offsets_are_exact_and_non_overlapping(self): + record = self.arena.block_record("blocks.0") + plan = build_transfer_plan(record, ["a", "b", "c"]) + spans = [] + for name, roles in plan.leaf_specs.items(): + for role, spec in roles.items(): + spans.append((spec.dst_offset, spec.dst_offset + spec.nbytes, name, role)) + spans.sort() + for prev, nxt in pairwise(spans): + self.assertLessEqual(prev[1], nxt[0]) + self.assertLessEqual(spans[-1][1], plan.compact_nbytes) + + def test_range_bytes_equal_compact_total_and_cover_every_leaf(self): + record = self.arena.block_record("blocks.0") + plan = build_transfer_plan(record, ["a", "b"]) + # Coalescing a small alignment gap (< LEAF_ALIGN) into one range is + # allowed to copy a few extra padding bytes -- ranges always sum + # exactly to compact_nbytes, and every real leaf byte must fall + # within that budget (no leaf is ever short-changed or truncated). + self.assertEqual(sum(r.nbytes for r in plan.ranges), plan.compact_nbytes) + leaf_bytes = sum( + spec.nbytes + for name in ("a", "b") + for spec in plan.leaf_specs[name].values() + ) + self.assertLessEqual(leaf_bytes, plan.compact_nbytes) + # Padding slack introduced by coalescing must stay under one + # alignment quantum per streamed leaf sub-tensor (weight/bias each + # contribute at most one internal gap). + n_sub_tensors = sum(len(roles) for roles in plan.leaf_specs.values()) + self.assertLess(plan.compact_nbytes - leaf_bytes, 256 * n_sub_tensors) + + def test_fingerprint_stable_across_rebuilds_same_selection(self): + record = self.arena.block_record("blocks.0") + p1 = build_transfer_plan(record, ["a", "b"]) + p2 = build_transfer_plan(record, ["b", "a"]) # order-independent input + self.assertEqual(p1.fingerprint, p2.fingerprint) + + def test_fingerprint_differs_for_different_selection(self): + record = self.arena.block_record("blocks.0") + p1 = build_transfer_plan(record, ["a", "b"]) + p2 = build_transfer_plan(record, ["a", "c"]) + self.assertNotEqual(p1.fingerprint, p2.fingerprint) + + def test_empty_streamed_set_fails_closed(self): + record = self.arena.block_record("blocks.0") + with self.assertRaises(TransferPlanError): + build_transfer_plan(record, []) + + def test_ranges_tensor_shape_and_dtype(self): + record = self.arena.block_record("blocks.0") + plan = build_transfer_plan(record, ["a", "c"]) + t = plan.ranges_tensor() + self.assertEqual(t.dtype, torch.int64) + self.assertEqual(tuple(t.shape), (plan.num_ranges, 3)) + self.assertEqual(t.device.type, "cpu") + + def test_compact_leaf_view_reads_correct_bytes(self): + record = self.arena.block_record("blocks.0") + plan = build_transfer_plan(record, ["a", "b", "c"]) + # Simulate the compact device buffer directly on the (CPU-readable) + # host flat's own bytes to check the view math without a real H2D + # copy: build a compact CPU buffer by literally performing the + # plan's ranges as CPU-to-CPU copies, then verify per-leaf values. + compact = torch.empty(plan.compact_nbytes, dtype=torch.uint8) + flat = record.host_flat + for r in plan.ranges: + compact[r.dst_offset:r.dst_offset + r.nbytes].copy_( + flat[r.src_offset:r.src_offset + r.nbytes] + ) + for name, layer in (("a", self.a), ("b", self.b), ("c", self.c)): + view = plan.compact_leaf_view(compact, name, "weight") + torch.testing.assert_close(view, layer.weight.data) + view_b = plan.compact_leaf_view(compact, name, "bias") + torch.testing.assert_close(view_b, layer.bias.data) + + + + def test_unknown_streamed_leaf_fails_closed(self): + record = self.arena.block_record("blocks.0") + with self.assertRaisesRegex(TransferPlanError, "transfer_plan_unknown_leaf"): + build_transfer_plan(record, ["a", "missing"]) + + def test_leaf_metadata_is_immutable(self): + record = self.arena.block_record("blocks.0") + plan = build_transfer_plan(record, ["a"]) + with self.assertRaises(TypeError): + plan.leaf_specs["a"]["weight"] = None +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_transfer_runtime.py b/tests/test_transfer_runtime.py new file mode 100644 index 0000000000..b35821530f --- /dev/null +++ b/tests/test_transfer_runtime.py @@ -0,0 +1,165 @@ +import pytest +import torch +import torch.nn.functional as F + +from toolkit.memory_management.arena_offload import transfer as ingraph_stream +from toolkit.memory_management.canonical_arena import CanonicalArena +from toolkit.memory_management.transfer_plan import build_transfer_plan + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required"), +] + + +def _linear(seed): + torch.manual_seed(seed) + layer = torch.nn.Linear(8, 8, bias=True) + layer.requires_grad_(False) + return layer + + +@pytest.fixture +def arena_block(): + layers = {name: _linear(seed) for seed, name in enumerate(("a", "b", "c"), 1)} + arena = CanonicalArena() + arena.canonicalize({"blocks.0": list(layers.items())}) + ingraph_stream.configure_fetch_runtime(depth=2) + ingraph_stream.fetch_stats(reset=True) + try: + yield arena.block_record("blocks.0"), layers + finally: + ingraph_stream.drain_fetch_runtime() + arena.release() + + +def _fetch(record, plan): + token = torch.ops.mm.fetch_start_multi( + record.host_flat, plan.ranges_tensor(), plan.compact_nbytes + ) + flat = torch.ops.mm.fetch_wait(token, plan.compact_nbytes) + return token, flat + + +def test_mixed_resident_streamed_numerical_parity_and_exact_stats(arena_block): + record, layers = arena_block + plan = build_transfer_plan(record, ("a", "c")) + token, flat = _fetch(record, plan) + x = torch.randn(4, 8, device="cuda") + + streamed = 0 + for name in ("a", "c"): + weight = plan.compact_leaf_view(flat, name, "weight") + bias = plan.compact_leaf_view(flat, name, "bias") + streamed = streamed + F.linear(x, weight, bias) + resident = F.linear( + x, layers["b"].weight.to("cuda"), layers["b"].bias.to("cuda") + ) + actual = streamed + resident + expected = sum(F.linear(x, layer.weight.to("cuda"), layer.bias.to("cuda")) + for layer in layers.values()) + torch.ops.mm.fetch_free_after(token, actual) + torch.cuda.synchronize() + + torch.testing.assert_close(actual, expected) + stats = ingraph_stream.fetch_stats() + assert stats["fetches"] == 1 + assert stats["bytes"] == plan.compact_nbytes + assert stats["copies"] == plan.num_ranges + + +def test_fully_streamed_plan_uses_single_copy_fast_path(arena_block): + record, _layers = arena_block + plan = build_transfer_plan(record, ("a", "b", "c")) + assert plan.fully_streamed and plan.num_ranges == 1 + token, flat = _fetch(record, plan) + span = plan.ranges[0] + expected = record.host_flat[span.src_offset:span.src_offset + span.nbytes].to("cuda") + torch.testing.assert_close(flat, expected) + torch.ops.mm.fetch_free(token) + torch.cuda.synchronize() + assert ingraph_stream.fetch_stats()["copies"] == 1 + + +def test_depth_two_reuse_hammer(arena_block): + record, _layers = arena_block + plan = build_transfer_plan(record, ("a", "c")) + expected = torch.empty(plan.compact_nbytes, dtype=torch.uint8) + for span in plan.ranges: + expected[span.dst_offset:span.dst_offset + span.nbytes].copy_( + record.host_flat[span.src_offset:span.src_offset + span.nbytes] + ) + expected = expected.to("cuda") + + for _ in range(50): + first_token, first = _fetch(record, plan) + second_token, second = _fetch(record, plan) + torch.testing.assert_close(first, expected) + torch.testing.assert_close(second, expected) + torch.ops.mm.fetch_free(first_token) + torch.ops.mm.fetch_free(second_token) + torch.cuda.synchronize() + stats = ingraph_stream.fetch_stats() + assert stats["fetches"] == 100 + assert stats["copies"] == 100 * plan.num_ranges + + +@pytest.mark.parametrize( + "ranges,compact_nbytes,match", + [ + (torch.tensor([[-1, 0, 4]], dtype=torch.int64), 4, "invalid range"), + (torch.tensor([[0, 0, 10**9]], dtype=torch.int64), 10**9, "out of bounds"), + (torch.tensor([[0, 1, 4]], dtype=torch.int64), 5, "not compact"), + (torch.tensor([[0, 0, 4], [2, 4, 4]], dtype=torch.int64), 8, "overlaps"), + ], +) +def test_invalid_ranges_fail_closed(arena_block, ranges, compact_nbytes, match): + record, _layers = arena_block + with pytest.raises(RuntimeError, match=match): + torch.ops.mm.fetch_start_multi(record.host_flat, ranges, compact_nbytes) + + +def test_pageable_non_arena_source_and_unknown_ticket_fail_closed(): + host = torch.empty(64, dtype=torch.uint8) + ranges = torch.tensor([[0, 0, 64]], dtype=torch.int64) + with pytest.raises(RuntimeError, match="registered canonical arena"): + torch.ops.mm.fetch_start_multi(host, ranges, 64) + with pytest.raises(RuntimeError, match="unknown ticket"): + torch.ops.mm.fetch_wait(torch.tensor([987654], dtype=torch.int64), 64) + + +def test_same_plan_new_storage_has_zero_recompiles(arena_block): + record, _layers = arena_block + plan = build_transfer_plan(record, ("a", "c")) + ranges = plan.ranges_tensor() + + other_layers = {name: _linear(seed) for seed, name in enumerate(("a", "b", "c"), 11)} + other_arena = CanonicalArena() + other_arena.canonicalize({"blocks.0": list(other_layers.items())}) + other_record = other_arena.block_record("blocks.0") + other_plan = build_transfer_plan(other_record, ("a", "c")) + assert other_plan.ranges == plan.ranges + + def consume(host, static_ranges, guard): + token = torch.ops.mm.fetch_start_multi_after( + host, static_ranges, plan.compact_nbytes, guard + ) + flat = torch.ops.mm.fetch_wait(token, plan.compact_nbytes) + out = flat.float().sum() + guard + torch.ops.mm.fetch_free_after(token, out) + return out + + torch._dynamo.reset() + compiled = torch.compile(consume, fullgraph=True, dynamic=False) + guard = torch.ones((), device="cuda") + try: + compiled(record.host_flat, ranges, guard) + torch.cuda.synchronize() + graphs = torch._dynamo.utils.counters["stats"].get("unique_graphs", 0) + assert graphs >= 1 + compiled(other_record.host_flat, other_plan.ranges_tensor(), guard) + torch.cuda.synchronize() + assert torch._dynamo.utils.counters["stats"].get("unique_graphs", 0) == graphs + assert sum(torch._dynamo.utils.counters["graph_break"].values()) == 0 + finally: + ingraph_stream.drain_fetch_runtime() + other_arena.release() diff --git a/toolkit/memory_management/arena_offload/__init__.py b/toolkit/memory_management/arena_offload/__init__.py new file mode 100644 index 0000000000..92e4e45666 --- /dev/null +++ b/toolkit/memory_management/arena_offload/__init__.py @@ -0,0 +1,57 @@ +"""Arena offload: the block-native weight-streaming backend. + +This package is the only supported integration surface for arena offload. +Model integrations and the shared trainer must go through `api`; they must not +construct `CanonicalArena`, `ResidencyState`, `ResidencyPlan`, or the immutable +runtime directly. + +Dependency rule (three tiers): + + host_memory (pin_manager, vram_budget, nvml_meminfo, dxgi_meminfo) + imports neither backend + + arena_offload -> may import host_memory; must NOT import MemoryManager + MemoryManager -> may import host_memory; must NOT import arena_offload + +Arena planning, transfer, FP8 transforms, and lifecycle cleanup are owned here; +the legacy manager remains a separate backend. +""" + +from .api import ( + ArenaOffloadConfig, + close_arena_offload, + get_arena_runtime, + is_arena_offloaded, + is_memory_managed, + memory_runtime_owns_compile, + prepare_canonical_storage, + prepare_canonical_storage_from_state_dict, + prepare_arena_offload, +) +from .runtime import ArenaOffloadRuntime +from .dispatcher import DISPATCHER_GENERATION +from .discovery import BlockDiscoveryError, discover_blocks +from .errors import ArenaCleanupError, ArenaSetupFatalError +from .load_session import model_load_arena_session +from ..runtime import close_memory_runtime, get_memory_runtime + +__all__ = [ + "ArenaOffloadConfig", + "ArenaCleanupError", + "ArenaOffloadRuntime", + "ArenaSetupFatalError", + "BlockDiscoveryError", + "DISPATCHER_GENERATION", + "close_arena_offload", + "close_memory_runtime", + "get_arena_runtime", + "get_memory_runtime", + "discover_blocks", + "is_arena_offloaded", + "is_memory_managed", + "memory_runtime_owns_compile", + "model_load_arena_session", + "prepare_canonical_storage", + "prepare_canonical_storage_from_state_dict", + "prepare_arena_offload", +] diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py new file mode 100644 index 0000000000..3c2b028326 --- /dev/null +++ b/toolkit/memory_management/arena_offload/api.py @@ -0,0 +1,358 @@ +"""Public integration surface for arena offload. + +Everything a model integration or the shared trainer is allowed to touch lives +here. The rule the rest of the codebase must follow: + + from toolkit.memory_management.arena_offload import ( + prepare_arena_offload, get_arena_runtime, ... + ) + +and nothing else. In particular, no `CanonicalArena`, `ResidencyState`, +`ResidencyPlan` or dispatcher-controller imports outside this package. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass, field +from typing import Any +import warnings + +from ..runtime import ( + RUNTIME_ATTR, + close_memory_runtime, + get_memory_runtime, + is_memory_managed, + memory_runtime_owns_compile, + unwrap_memory_model, +) +from .runtime import ArenaOffloadRuntime + +GIB = 1024**3 +_FP8_QTYPES = ("qfloat8", "float8") +_COMPATIBILITY_ALIASES = { + "layer_offloading_smart_working_reserve_gb": ( + "layer_offloading_smart_headroom_gb", + ), + "layer_offloading_smart_wddm_margin_gb": ( + "layer_offloading_smart_buffer_gb", + ), + "layer_offloading_smart_wddm_hard_gb": ( + "layer_offloading_smart_hard_buffer_gb", + ), + "layer_offloading_smart_sampling_working_reserve_gb": ( + "layer_offloading_smart_sampling_headroom_gb", + ), + "layer_offloading_smart_sampling_wddm_margin_gb": ( + "layer_offloading_smart_sampling_buffer_gb", + ), + "layer_offloading_smart_sampling_wddm_hard_gb": ( + "layer_offloading_smart_sampling_hard_buffer_gb", + ), +} + + +def unwrap(model): + """Peel Accelerate / DDP / torch.compile wrappers without importing them. + + The arena package must stay importable from a bare CPU test process, so this + does not go through `toolkit.accelerator.unwrap_model` (which constructs a + global `Accelerator`). + """ + return unwrap_memory_model(model) + + +@dataclass(frozen=True) +class _ArenaPolicyOptions: + """Internal policy inputs retained while fork job aliases are migrated.""" + + working_reserve_gib: float | None = None + wddm_margin_gib: float | None = None + wddm_hard_gib: float | None = None + checkpoint_keep_last: int = 0 + prefetch_depth: int = 3 + + sampling_working_reserve_gib: float | None = None + sampling_wddm_margin_gib: float | None = None + sampling_wddm_hard_gib: float | None = 1.0 + + +@dataclass(frozen=True) +class ArenaOffloadConfig: + """The narrow public configuration surface of arena offload. + + Fields prefixed with ``_`` are derived integration details, not additional + user-facing arena controls. + """ + + enabled: bool = False + fp8_forward: bool = False + fp8_backward: bool = False + fp8_sampling: bool = False + compile_blocks: bool = False + _compile_dynamic: bool | None = True + _compile_dynamic_hints: tuple[tuple[int, int | None, int | None], ...] = () + # Validation knob: pretend the card is this many GiB, so small-card + # behaviour (deeper streaming, tighter caps, a residency plan that cannot + # fit) is exercisable on a bigger one. 0/None = use the real card. + _simulated_vram_gib: float | None = None + _policy: _ArenaPolicyOptions = field( + default_factory=_ArenaPolicyOptions, repr=False + ) + + @classmethod + def from_model_config( + cls, model_config, *, training_working_reserve_hint_bytes: int | None = None + ) -> ArenaOffloadConfig: + def get(name: str, default: Any = None) -> Any: + if hasattr(model_config, name): + return getattr(model_config, name) + for alias in _COMPATIBILITY_ALIASES.get(name, ()): + if hasattr(model_config, alias): + return getattr(model_config, alias) + return default + + raw_working_reserve_gib = get("layer_offloading_smart_working_reserve_gb") + working_reserve_gib = raw_working_reserve_gib + if training_working_reserve_hint_bytes: + try: + is_auto = raw_working_reserve_gib is None or float(raw_working_reserve_gib) < 0 + except (TypeError, ValueError): + is_auto = str(raw_working_reserve_gib).strip().lower() == "auto" + if is_auto: + # A caller-supplied, resolution-aware hint (see + # vram_budget.estimate_training_working_reserve_bytes) beats the + # planner's flat DEFAULT_AUTO_WORKING_RESERVE_GIB fallback, but + # never overrides an explicit user value. + working_reserve_gib = float(training_working_reserve_hint_bytes) / float(GIB) + + fp8_weights = bool(get("quantize", False)) and get("qtype") in _FP8_QTYPES + requested_forward = bool(get("layer_offloading_fp8_forward", False)) + requested_backward = bool(get("layer_offloading_fp8_grad_input", False)) + requested_sampling = bool(get("layer_offloading_fp8_sampling", False)) + ignored = [] + if not fp8_weights: + ignored.extend( + name + for name, requested in ( + ("fp8_forward", requested_forward), + ("fp8_backward", requested_backward), + ("fp8_sampling", requested_sampling), + ) + if requested + ) + elif requested_backward and not requested_forward: + ignored.append("fp8_backward_without_fp8_forward") + if ignored: + warnings.warn( + "arena offload ignored irrelevant FP8 options: " + + ", ".join(ignored), + RuntimeWarning, + stacklevel=2, + ) + + return cls( + enabled=bool( + get("layer_offloading", False) + and get("layer_offloading_smart", False) + ), + fp8_forward=fp8_weights and requested_forward, + fp8_backward=fp8_weights + and requested_forward + and requested_backward, + fp8_sampling=fp8_weights + and requested_sampling, + compile_blocks=bool( + get("compile", False) + or get("compile_sample", False) + or get("train_compile_blocks", False) + ), + _compile_dynamic=( + None + if get("compile_dynamic", True) is None + else bool(get("compile_dynamic", True)) + ), + _compile_dynamic_hints=tuple( + tuple(hint) for hint in (get("compile_dynamic_hints", ()) or ()) + ), + _simulated_vram_gib=( + float(get("layer_offloading_simulated_vram_gb") or 0.0) or None + ), + _policy=_ArenaPolicyOptions( + working_reserve_gib=working_reserve_gib, + wddm_margin_gib=get("layer_offloading_smart_wddm_margin_gb"), + wddm_hard_gib=get("layer_offloading_smart_wddm_hard_gb"), + checkpoint_keep_last=max( + 0, int(get("layer_offloading_checkpoint_keep_last", 0) or 0) + ), + prefetch_depth=int(get("layer_offloading_prefetch_depth", 3) or 3), + sampling_working_reserve_gib=get( + "layer_offloading_smart_sampling_working_reserve_gb" + ), + sampling_wddm_margin_gib=get( + "layer_offloading_smart_sampling_wddm_margin_gb" + ), + sampling_wddm_hard_gib=get( + "layer_offloading_smart_sampling_wddm_hard_gb", 1.0 + ), + ), + ) + + + +def prepare_canonical_storage( + transformer, + *, + block_names: Sequence[str] | None = None, + device=None, + defer_blocks: bool = False, +): + """Prepare final arena destinations without publishing model Parameters.""" + from ..canonical_arena import CanonicalArena + from .discovery import discover_blocks + from .resources import ArenaRuntimeResources + + selection = discover_blocks( + transformer, container_paths=tuple(block_names or ()) + ) + resources = None + if device is not None: + resources = ArenaRuntimeResources(transformer, device) + resources.acquire_process_owner() + try: + entries = ( + {} + if defer_blocks + else { + key: list(selection.entries_by_block[key]) + for key in selection.block_keys + } + ) + arena = CanonicalArena() + build = arena.prepare(entries, model=transformer) + if resources is not None: + resources.adopt_canonical_build(build) + return build + except BaseException: + if resources is not None: + resources.release() + raise + + +def prepare_canonical_storage_from_state_dict( + transformer, + state_dict, + *, + block_names: Sequence[str] | None = None, + device=None, +): + """Build canonical storage incrementally from a mutable state mapping. + + Serialized source keys and physical tensor leaves are inferred from the + transformer's module paths, each module's own state-dict surface, and its + quantization storage declaration. The complete mapping is validated before + allocation or destructive consumption begins. + """ + from .discovery import discover_blocks + + selection = discover_blocks( + transformer, container_paths=tuple(block_names or ()) + ) + entries_by_block = selection.entries_by_block + + from .construction import infer_state_dict_schema, validate_state_dict_schema + + schema = infer_state_dict_schema(transformer, entries_by_block) + validate_state_dict_schema(state_dict, schema) + build = prepare_canonical_storage( + transformer, + block_names=block_names, + device=device, + defer_blocks=True, + ) + build.populate_from_state_dict_consuming( + state_dict, blocks=entries_by_block.items() + ) + return build + + +def prepare_arena_offload( + transformer, + *, + device, + config: ArenaOffloadConfig, + block_names: Sequence[str] | None = None, + ignore_modules: Sequence[Any] | None = None, + canonical_build=None, +) -> ArenaOffloadRuntime: + """Canonicalize the model's execution blocks and prepare the arena runtime. + + Call after the base weights are final and BEFORE the training network is + applied. Loaders that populate final arena destinations directly pass their + populated ``canonical_build``; other models use the compatibility source, + which copies from the already-materialized model. The runtime comes back + unfinalized; the trainer calls ``finalize()`` once the network is installed. + + The runtime is published on `transformer._arena_offload_runtime`. + """ + if not config.enabled: + raise ValueError("arena_offload_not_enabled") + from .load_session import claim_pending_canonical_build + + pending_build = claim_pending_canonical_build(transformer) + if canonical_build is None: + canonical_build = pending_build + elif pending_build is not None: + pending_build.rollback() + canonical_build.rollback() + raise ValueError("arena_multiple_pending_canonical_builds") + try: + from .discovery import discover_blocks + + if not bool(getattr(transformer, "gradient_checkpointing", False)): + raise ValueError( + "arena training requires model gradient checkpointing before " + "canonical storage is committed" + ) + configured_keep_last = int(config._policy.checkpoint_keep_last) + model_keep_last = getattr(transformer, "_checkpoint_keep_last", None) + if model_keep_last is None: + if configured_keep_last: + raise ValueError( + "arena checkpoint keep-last requires a model-owned " + "keep-last declaration" + ) + elif int(model_keep_last) != configured_keep_last: + raise ValueError( + "arena checkpoint keep-last does not match the model-owned " + f"value: config={configured_keep_last} model={int(model_keep_last)}" + ) + selection = discover_blocks( + transformer, + container_paths=tuple(block_names or ()), + ) + except BaseException: + if canonical_build is not None: + canonical_build.rollback() + raise + return ArenaOffloadRuntime._prepare( + transformer, + device=device, + selection=selection, + config=config, + ignore_modules=ignore_modules, + canonical_build=canonical_build, + ) + + +def get_arena_runtime(model) -> ArenaOffloadRuntime | None: + """The arena runtime for `model`, or None. Unwraps Accelerate/DDP/compile.""" + return get_memory_runtime(model) + + +def is_arena_offloaded(model) -> bool: + return get_arena_runtime(model) is not None + + +def close_arena_offload(model) -> None: + close_memory_runtime(model) diff --git a/toolkit/memory_management/arena_offload/construction.py b/toolkit/memory_management/arena_offload/construction.py new file mode 100644 index 0000000000..1b5995e3b9 --- /dev/null +++ b/toolkit/memory_management/arena_offload/construction.py @@ -0,0 +1,590 @@ +"""Destination-first transactional construction of canonical arena storage.""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +import torch + +from toolkit.memory_management import pin_manager +from toolkit.quantization.storage import module_storage_binding, named_tensor_storage +from .layout import ( + BlockPack, + LeafSpec, + LinearSpec, + inspect_block, + LayerStorageView, + linear_views, + make_block_view_maker, + substitution_views, + typed_view, +) + + +class CanonicalBuildError(RuntimeError): + pass + + +class CanonicalStateInferenceError(CanonicalBuildError): + pass + + +@dataclass(frozen=True) +class _StateDestination: + key: tuple[str, str, str] + shape: tuple[int, ...] + dtype: torch.dtype + + +def _module_paths(model) -> dict[int, str]: + if model is None: + raise CanonicalStateInferenceError("canonical_state_inference_requires_model") + paths = {} + duplicates = set() + try: + modules = model.named_modules(remove_duplicate=False) + except TypeError: + modules = model.named_modules() + for path, module in modules: + previous = paths.get(id(module)) + if previous is not None and previous != path: + duplicates.add(id(module)) + else: + paths[id(module)] = str(path) + for identity in duplicates: + paths.pop(identity, None) + return paths + + +def _same_storage_view(left: torch.Tensor, right: torch.Tensor) -> bool: + """Whether two tensor objects name the same physical/meta storage view.""" + try: + same_storage = ( + left.untyped_storage()._cdata == right.untyped_storage()._cdata + ) + except (AttributeError, RuntimeError): + return False + return bool( + same_storage + and left.storage_offset() == right.storage_offset() + and tuple(left.shape) == tuple(right.shape) + and tuple(left.stride()) == tuple(right.stride()) + and left.dtype == right.dtype + ) + + +def infer_state_dict_schema( + model, entries_by_block, *, allow_unserialized_storage: bool = False +) -> dict: + """Infer serialized source entries for managed physical destinations. + + Direct model-source construction may include declared execution buffers + that are intentionally non-persistent. State-dict consumers keep the + strict default because those destinations cannot be populated from an + absent serialized source. + """ + module_paths = _module_paths(model) + schema = {} + for block_key, raw_entries in entries_by_block.items(): + block_schema = {} + for entry_name, module in tuple(raw_entries): + module_path = module_paths.get(id(module)) + if module_path is None: + raise CanonicalStateInferenceError( + f"canonical_state_module_path_ambiguous:{block_key}:{entry_name}" + ) + binding = module_storage_binding(module) + serialized = [] + groups = [] + for relative_key, value in module.state_dict(keep_vars=True).items(): + leaves = named_tensor_storage(value) + start = len(serialized) + serialized.extend(leaves) + groups.append((str(relative_key), start, len(serialized))) + matched = [] + claimed = set() + for declared in binding.tensors: + candidates = [ + index + for index, serialized_leaf in enumerate(serialized) + if _same_storage_view( + serialized_leaf.tensor, declared.tensor + ) + ] + if not candidates: + if allow_unserialized_storage: + continue + raise CanonicalStateInferenceError( + "canonical_state_storage_not_serialized:" + f"{block_key}:{entry_name}:{declared.name}" + ) + if len(candidates) != 1 or candidates[0] in claimed: + raise CanonicalStateInferenceError( + "canonical_state_storage_match_ambiguous:" + f"{block_key}:{entry_name}:{declared.name}" + ) + match = candidates[0] + matched.append((match, declared)) + claimed.add(match) + matched_by_index = {index: declared for index, declared in matched} + for relative_key, start, end in groups: + selected = [ + (index, matched_by_index[index]) + for index in range(start, end) + if index in matched_by_index + ] + if not selected: + continue + if len(selected) != end - start: + raise CanonicalStateInferenceError( + f"canonical_state_partial_serialized_entry:{module_path}.{relative_key}" + ) + source_key = ( + f"{module_path}.{relative_key}" + if module_path + else relative_key + ) + if source_key in block_schema: + raise CanonicalStateInferenceError( + f"canonical_state_duplicate_source:{source_key}" + ) + destinations = [] + for index, declared in selected: + serialized_leaf = serialized[index] + tensor = serialized_leaf.tensor + if ( + tuple(tensor.shape) != tuple(declared.tensor.shape) + or tensor.dtype != declared.tensor.dtype + ): + raise CanonicalStateInferenceError( + "canonical_state_target_schema_mismatch:" + f"{source_key}:{declared.name}" + ) + destinations.append( + _StateDestination( + (str(block_key), str(entry_name), declared.name), + tuple(tensor.shape), + tensor.dtype, + ) + ) + block_schema[source_key] = tuple(destinations) + schema[str(block_key)] = block_schema + return schema + + +def validate_state_dict_schema(state_dict, schema) -> None: + """Prove inference for every managed source before destructive copying.""" + for block_schema in schema.values(): + for source_key, destinations in block_schema.items(): + if source_key not in state_dict: + raise CanonicalStateInferenceError( + f"canonical_state_missing_source:{source_key}" + ) + leaves = named_tensor_storage(state_dict[source_key]) + if len(leaves) != len(destinations): + raise CanonicalStateInferenceError( + "canonical_state_source_count_mismatch:" + f"{source_key}:source={len(leaves)}:expected={len(destinations)}" + ) + for leaf, destination in zip(leaves, destinations, strict=True): + if ( + tuple(leaf.tensor.shape) != destination.shape + or leaf.tensor.dtype != destination.dtype + ): + raise CanonicalStateInferenceError( + f"canonical_state_source_schema_mismatch:{source_key}" + ) + + +@dataclass +class _PreparedBlock: + key: str + entries: tuple + layout: object + flat: torch.Tensor + pending: object + handle: object | None = None + pack: BlockPack | None = None + + +class PreparedCanonicalBuild: + """A prepared arena build whose model publication is atomic.""" + + def __init__(self, arena, entries_by_block, *, model=None, kind="weights"): + self.arena = arena + self.model = model + self.kind = kind + self.blocks = [] + self.destinations = {} + self.entries_by_block = {} + self.state_schema = {} + self._originals = [] + self._populated = False + self._committed = False + if arena.canonicalized: + raise CanonicalBuildError("canonical_arena_double_canonicalize") + try: + for key, raw_entries in entries_by_block.items(): + self.add_block(key, raw_entries) + except Exception: + self.rollback() + raise + + def add_block(self, key, raw_entries) -> None: + """Prepare one final block layout for a bounded direct loader.""" + if self._populated or self._committed: + raise CanonicalBuildError("canonical_build_already_populated") + if key in self.entries_by_block: + raise CanonicalBuildError(f"canonical_build_duplicate_block:{key}") + entries = tuple(raw_entries) + block_schema = ( + infer_state_dict_schema( + self.model, + {str(key): entries}, + allow_unserialized_storage=True, + )[str(key)] + if self.model is not None + else {} + ) + layout = inspect_block(key, entries) + for name, module in entries: + binding = module_storage_binding(module) + for substitution in binding.substitutions: + target = substitution.name + if target in module._parameters: + value = module._parameters[target] + if value is not None and value.requires_grad: + raise CanonicalBuildError( + f"canonical_arena_trainable_leaf:{key}:{name}.{target}" + ) + self._originals.append((module, target, "parameter", value)) + elif target in module._buffers: + self._originals.append( + (module, target, "buffer", module._buffers[target]) + ) + else: + raise CanonicalBuildError( + f"canonical_arena_missing_target:{key}:{name}.{target}" + ) + flat, padded = pin_manager.pin_register_prepare(layout.nbytes) + block = _PreparedBlock(key, entries, layout, flat, padded) + self.blocks.append(block) + self.entries_by_block[key] = entries + self.state_schema[str(key)] = block_schema + for linear in layout.linears: + for leaf in linear.leaf_descriptors: + self.destinations[(key, linear.name, leaf.role)] = typed_view(flat, leaf) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + if exc_type is not None or not self._committed: + self.rollback() + return False + + def populate(self, source) -> None: + try: + source(self.destinations) + self._finish_population() + except Exception: + self.rollback() + raise + + def populate_from_model(self) -> None: + try: + for destination_key, source in self.model_source_leaves(): + self.destinations[destination_key].copy_(source) + self._finish_population() + except Exception: + self.rollback() + raise + + def populate_block_from_model(self, block_key: str) -> None: + """Copy one bounded loaded block without finalizing the whole build.""" + for destination_key, source in self.model_source_leaves(block_key=block_key): + self.destinations[destination_key].copy_(source) + + def populate_from_state_dict_consuming( + self, + state_dict, + *, + blocks=None, + ) -> tuple[str, ...]: + """Populate final flats one block at a time and consume their sources. + + Source keys and physical leaves are inferred from each managed + module's own ``state_dict`` surface and storage declaration. When + ``blocks`` is supplied it yields ``(block_key, entries)`` pairs and + each block is allocated only when its turn begins. Otherwise the + already-prepared blocks are populated in their existing order. + + Every fully copied source entry is removed from ``state_dict`` before + the next block is requested. The mapping is therefore intentionally + destructive even if a later block fails; callers must discard it on + failure, while the canonical build still rolls its own resources back. + """ + if self._populated or self._committed: + raise CanonicalBuildError("canonical_build_already_populated") + if blocks is not None and self.blocks: + raise CanonicalBuildError("canonical_build_blocks_already_prepared") + + consumed = [] + try: + if blocks is None: + block_items = ((block.key, None) for block in tuple(self.blocks)) + else: + block_items = iter(blocks) + + for block_key, entries in block_items: + if entries is not None: + self.add_block(block_key, entries) + block_schema = self.state_schema.get(str(block_key), {}) + if not block_schema: + raise CanonicalBuildError( + f"canonical_build_unknown_block:{block_key}" + ) + source_keys = tuple(block_schema) + for source_key in source_keys: + self.copy_state_entry(source_key, state_dict[source_key]) + for source_key in source_keys: + state_dict.pop(source_key) + consumed.extend(source_keys) + + self._finish_population() + return tuple(consumed) + except Exception: + self.rollback() + raise + + def copy_state_entry(self, source_key: str, value) -> bool: + """Copy one serialized state entry when it belongs to this build.""" + destinations = None + for block_schema in self.state_schema.values(): + destinations = block_schema.get(str(source_key)) + if destinations is not None: + break + if destinations is None: + return False + leaves = named_tensor_storage(value) + if len(leaves) != len(destinations): + raise CanonicalBuildError( + f"canonical_state_source_count_mismatch:{source_key}" + ) + for leaf, destination in zip(leaves, destinations, strict=True): + source = leaf.tensor + if ( + tuple(source.shape) != destination.shape + or source.dtype != destination.dtype + ): + raise CanonicalBuildError( + f"canonical_state_source_schema_mismatch:{source_key}" + ) + self.destinations[destination.key].copy_(source) + del source + return True + + def release_block_sources_to_meta(self, block_key: str) -> None: + """Drop a direct loader's bounded source after its final flat is populated.""" + block = next(item for item in self.blocks if item.key == block_key) + meta_flat = torch.empty(block.layout.nbytes, dtype=torch.uint8, device="meta") + replacements = {} + modules = dict(block.entries) + meta_linears = [] + for linear in block.layout.linears: + module = modules[linear.name] + for target, value in substitution_views(meta_flat, linear).items(): + requires_grad = ( + linear.weight_requires_grad + if target == "weight" + else linear.bias_requires_grad if target == "bias" else False + ) + replacements[(module, target)] = self._publish_state( + module, target, value, requires_grad=requires_grad + ) + meta_weight = module._parameters.get("weight") + meta_linears.append( + replace( + linear, + weight_template=( + meta_weight.data if meta_weight is not None else None + ), + ) + ) + block.layout = replace(block.layout, linears=tuple(meta_linears)) + self._originals = [ + (module, target, kind, replacements.get((module, target), value)) + for module, target, kind, value in self._originals + ] + + def finish_population(self) -> None: + try: + self._finish_population() + except Exception: + self.rollback() + raise + + def storage_views(self, block_key: str) -> tuple[LayerStorageView, ...]: + """Return populated immutable storage declarations before commit.""" + if not self._populated: + raise CanonicalBuildError("canonical_build_not_populated") + try: + block = next(item for item in self.blocks if item.key == block_key) + except StopIteration as error: + raise CanonicalBuildError( + f"canonical_build_unknown_block:{block_key}" + ) from error + return tuple( + LayerStorageView( + spec=LinearSpec( + name=linear.name, + tensors=tuple( + LeafSpec(**leaf.__dict__) + for leaf in linear.leaf_descriptors + ), + execution_key=linear.execution_key, + weight_leaf_count=linear.weight_leaf_count, + weight_template=linear.weight_template, + weight_requires_grad=linear.weight_requires_grad, + bias_requires_grad=linear.bias_requires_grad, + substitutions=linear.substitutions, + ), + tensors=tuple( + typed_view(block.flat, leaf) + for leaf in linear.leaf_descriptors + ), + ) + for linear in block.layout.linears + ) + + def model_source_leaves(self, *, block_key: str | None = None): + """Yield loaded model leaves keyed exactly like final destinations.""" + for block in self.blocks: + if block_key is not None and block.key != block_key: + continue + by_name = dict(block.entries) + for linear in block.layout.linears: + module = by_name[linear.name] + binding = module_storage_binding(module) + for descriptor, declared in zip( + linear.leaf_descriptors, binding.tensors, strict=True + ): + yield ( + block.key, + linear.name, + descriptor.role, + ), declared.tensor + + def _finish_population(self) -> None: + for block in self.blocks: + handle = pin_manager.pin_register_commit( + block.flat, block.layout.nbytes, self.kind, required=False + ) + if not handle.pinned: + pin_manager.release(handle) + raise CanonicalBuildError(f"canonical_arena_pin_budget_exceeded:{block.key}") + block.handle = handle + block.flat = handle.tensor + # Validate every supported reconstruction before publication. + for linear in block.layout.linears: + if linear.weight_leaf_count: + # Preserve the established wrapper-validation seam while + # executable substitutions become the source of truth. + linear_views(block.flat, linear) + substitution_views(block.flat, linear) + self._populated = True + + @staticmethod + def _publish_state(module, target, value, *, requires_grad=False): + if target in module._parameters: + published = torch.nn.Parameter(value, requires_grad=requires_grad) + setattr(module, target, published) + return published + if target in module._buffers: + setattr(module, target, value) + return value + raise CanonicalBuildError(f"canonical_arena_missing_target:{target}") + + def commit(self): + if not self._populated: + self.rollback() + raise CanonicalBuildError("canonical_build_not_populated") + from toolkit.memory_management.canonical_arena import BlockRecord, CanonicalArenaStats + published = [] + try: + for block in self.blocks: + specs = [] + for linear in block.layout.linears: + module = dict(block.entries)[linear.name] + for target, value in substitution_views( + block.flat, linear + ).items(): + requires_grad = ( + linear.weight_requires_grad + if target == "weight" + else linear.bias_requires_grad if target == "bias" else False + ) + self._publish_state( + module, + target, + value, + requires_grad=requires_grad, + ) + published.append(module) + specs.append(LinearSpec( + name=linear.name, + tensors=tuple( + LeafSpec(**leaf.__dict__) + for leaf in linear.leaf_descriptors + ), + execution_key=linear.execution_key, + weight_leaf_count=linear.weight_leaf_count, + weight_template=linear.weight_template, + weight_requires_grad=linear.weight_requires_grad, + bias_requires_grad=linear.bias_requires_grad, + substitutions=linear.substitutions, + )) + pack = BlockPack( + block.key, + block.flat, + tuple(specs), + block.layout.nbytes, + True, + pin_handle=block.handle, + ) + pack.view_maker = make_block_view_maker(pack) + block.pack = pack + names = tuple(name for name, _ in block.entries) + modules = tuple(module for _, module in block.entries) + self.arena._blocks[block.key] = BlockRecord(block.key, pack, names, modules) + pin_manager.register_arena_storage(block.flat) + self.arena._canonicalized = True + if self.model is not None: + self.arena.guard_whole_model_to(self.model) + self._committed = True + return CanonicalArenaStats(len(self.blocks), sum(b.layout.nbytes for b in self.blocks)) + except Exception: + self.rollback() + raise + + def rollback(self) -> None: + for module, target, kind, value in reversed(self._originals): + # Bypass user/module publication hooks: rollback must remain valid + # even when the injected/real failure was an attribute assignment. + if kind == "parameter": + module._parameters[target] = value + else: + module._buffers[target] = value + if self.model is not None: + self.arena.unguard_whole_model_to(self.model) + for block in self.blocks: + if block.pack is not None: + pin_manager.unregister_arena_storage(block.flat) + pin_manager.release(block.handle) + block.handle = None + self.arena._blocks.clear() + self.arena._canonicalized = False + self._populated = False + resources = getattr(self, "_arena_resources", None) + if resources is not None and not resources.canonical_committed: + resources.release() diff --git a/toolkit/memory_management/arena_offload/discovery.py b/toolkit/memory_management/arena_offload/discovery.py new file mode 100644 index 0000000000..e5f1f4daf6 --- /dev/null +++ b/toolkit/memory_management/arena_offload/discovery.py @@ -0,0 +1,265 @@ +"""Generic repeated-block discovery and pre-commit state accounting.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from toolkit.quantization.storage import module_storage_binding + + +class BlockDiscoveryError(RuntimeError): + pass + + +@dataclass(frozen=True) +class BlockStateAccounting: + managed_entries: int + managed_bytes: int + trainable_entries: int + resident_entries: int + resident_bytes: int + + +@dataclass(frozen=True) +class BlockSelection: + container_paths: tuple[str, ...] + blocks: tuple[torch.nn.Module, ...] + block_keys: tuple[str, ...] + entries_by_block: dict[str, tuple[tuple[str, torch.nn.Module], ...]] + accounting: BlockStateAccounting + + +def _resolve_path(model, path: str): + value = model + for component in str(path).split("."): + value = getattr(value, component, None) + if value is None: + raise BlockDiscoveryError(f"block_container_not_found:{path}") + return value + + +def managed_entries(block: torch.nn.Module) -> tuple[tuple[str, torch.nn.Module], ...]: + entries = [] + try: + modules = block.named_modules(remove_duplicate=False) + except TypeError: + modules = block.named_modules() + for path, module in modules: + if not path or not isinstance(module, torch.nn.Linear): + continue + try: + module_storage_binding(module) + except Exception as error: + raise BlockDiscoveryError( + f"unsupported_managed_module:{path}:{type(module).__qualname__}" + ) from error + entries.append((path, module)) + return tuple(entries) + + +def _container_candidate(path, container): + if not isinstance(container, (torch.nn.ModuleList, torch.nn.Sequential)): + return None + blocks = tuple(container) + if len(blocks) < 2 or any(not isinstance(block, torch.nn.Module) for block in blocks): + return None + entries = tuple(managed_entries(block) for block in blocks) + populated = sum(bool(items) for items in entries) + if populated < 2: + return None + payload = 0 + for items in entries: + for _name, module in items: + binding = module_storage_binding(module) + payload += sum( + int(item.tensor.numel() * item.tensor.element_size()) + for item in binding.tensors + ) + return str(path), blocks, entries, payload + + +def _select_containers(model, container_paths): + if container_paths: + selected = [] + for path in tuple(container_paths): + candidate = _container_candidate(path, _resolve_path(model, path)) + if candidate is None: + raise BlockDiscoveryError(f"invalid_block_container:{path}") + selected.append(candidate) + return tuple(selected) + + candidates = [] + for path, module in model.named_modules(): + if not path: + continue + candidate = _container_candidate(path, module) + if candidate is not None: + candidates.append(candidate) + if not candidates: + raise BlockDiscoveryError("no_repeated_block_container") + candidates.sort(key=lambda item: (-item[3], item[0])) + if len(candidates) > 1 and candidates[0][3] == candidates[1][3]: + paths = ",".join(item[0] for item in candidates if item[3] == candidates[0][3]) + raise BlockDiscoveryError(f"ambiguous_block_container:{paths}") + return (candidates[0],) + + +def _named_buffers(block): + for module_path, module in block.named_modules(): + for name, value in module._buffers.items(): + if value is None: + continue + path = f"{module_path}.{name}" if module_path else name + yield path, value + + +def _tensor_identity(tensor): + if tensor.device.type == "meta": + return ("meta", id(tensor)) + if tensor.numel() == 0: + return ("empty", tensor.device.type, id(tensor)) + try: + return ( + tensor.device.type, + tensor.untyped_storage().data_ptr(), + tensor.storage_offset(), + tensor.numel(), + ) + except Exception: + return ("object", id(tensor)) + + +def _state_bytes(value) -> int: + if not isinstance(value, torch.Tensor) or value.device.type == "meta": + return 0 + return int(value.numel() * value.element_size()) + + +def _audit_state(blocks, block_keys, entries_by_block) -> BlockStateAccounting: + managed_objects = {} + managed_storage = {} + managed_entries_count = 0 + managed_bytes = 0 + trainable_entries = 0 + resident_entries = 0 + resident_bytes = 0 + resident_seen = set() + + for block, block_key in zip(blocks, block_keys, strict=True): + declared = {} + for module_path, module in entries_by_block[block_key]: + binding = module_storage_binding(module) + for substitution in binding.substitutions: + target = substitution.name + state_path = f"{module_path}.{target}" + if state_path in declared: + raise BlockDiscoveryError( + f"conflicting_replacement_declaration:{block_key}.{state_path}" + ) + if torch.nn.utils.parametrize.is_parametrized(module, target): + raise BlockDiscoveryError( + f"parametrized_managed_state:{block_key}.{state_path}" + ) + if target in module._parameters: + value = module._parameters[target] + elif target in module._buffers: + value = module._buffers[target] + else: + raise BlockDiscoveryError( + f"missing_replacement_target:{block_key}.{state_path}" + ) + if value is None: + raise BlockDiscoveryError( + f"missing_replacement_target:{block_key}.{state_path}" + ) + declared[state_path] = value + previous = managed_objects.get(id(value)) + if previous is not None: + raise BlockDiscoveryError( + f"shared_managed_state:{previous}:{block_key}.{state_path}" + ) + managed_objects[id(value)] = f"{block_key}.{state_path}" + managed_entries_count += 1 + for tensor in binding.tensors: + identity = _tensor_identity(tensor.tensor) + previous = managed_storage.get(identity) + if previous is not None and previous[0] != block_key: + raise BlockDiscoveryError( + f"shared_managed_storage:{previous[1]}:{block_key}.{module_path}" + ) + if previous is None: + managed_bytes += _state_bytes(tensor.tensor) + managed_storage[identity] = ( + block_key, + f"{block_key}.{module_path}.{tensor.name}", + tensor.tensor, + ) + + parameters = tuple(block.named_parameters(recurse=True, remove_duplicate=False)) + # Non-persistent buffers still participate in execution and must be + # present in the complete state audit. Persistence only controls the + # default state_dict serialization surface. + buffers = tuple(_named_buffers(block)) + enumerated = {name for name, _value in parameters + buffers} + missing = set(declared) - enumerated + if missing: + raise BlockDiscoveryError( + f"managed_state_not_enumerated:{block_key}.{sorted(missing)[0]}" + ) + for name, value in parameters + buffers: + if name in declared: + continue + if isinstance(value, torch.nn.Parameter) and value.requires_grad: + trainable_entries += 1 + continue + resident_entries += 1 + identity = _tensor_identity(value) + if identity not in resident_seen: + resident_seen.add(identity) + resident_bytes += _state_bytes(value) + + return BlockStateAccounting( + managed_entries=managed_entries_count, + managed_bytes=managed_bytes, + trainable_entries=trainable_entries, + resident_entries=resident_entries, + resident_bytes=resident_bytes, + ) + + +def discover_blocks(model, *, container_paths=()) -> BlockSelection: + """Select repeated blocks and audit all state before canonical commit.""" + containers = _select_containers(model, tuple(container_paths or ())) + paths = tuple(item[0] for item in containers) + blocks = tuple(block for item in containers for block in item[1]) + if len({id(block) for block in blocks}) != len(blocks): + raise BlockDiscoveryError("duplicate_selected_block") + + block_keys = tuple( + f"{path}.{index}" + for path, _blocks, _entries, _payload in containers + for index in range(len(_blocks)) + ) + entry_groups = tuple( + entries + for _path, _blocks, groups, _payload in containers + for entries in groups + ) + entries_by_block = dict(zip(block_keys, entry_groups, strict=True)) + managed_modules = [ + module + for entries in entry_groups + for _name, module in entries + ] + if len({id(module) for module in managed_modules}) != len(managed_modules): + raise BlockDiscoveryError("overlapping_managed_module") + accounting = _audit_state(blocks, block_keys, entries_by_block) + return BlockSelection( + container_paths=paths, + blocks=blocks, + block_keys=block_keys, + entries_by_block=entries_by_block, + accounting=accounting, + ) diff --git a/toolkit/memory_management/arena_offload/dispatcher.py b/toolkit/memory_management/arena_offload/dispatcher.py new file mode 100644 index 0000000000..6cb66ab0c6 --- /dev/null +++ b/toolkit/memory_management/arena_offload/dispatcher.py @@ -0,0 +1,333 @@ +"""Generic saved-forward dispatcher over canonical arena storage.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from toolkit.memory_management.immutable_runtime import ( + ImmutableProgram, + ImmutableRuntimeError, + ImmutableTransformerRuntime, + build_program_fingerprint, +) +from toolkit.memory_management.arena_offload.transfer import ( + configure_fetch_runtime, + free_on_backward, +) + + +DISPATCHER_GENERATION = "generic-block-dispatcher-v1" + + +def _in_backward_graph_task() -> bool: + """True for non-reentrant checkpoint replay inside autograd backward.""" + try: + return int(torch._C._current_graph_task_id()) >= 0 + except (AttributeError, RuntimeError): + return False + + +class OriginalBlockInvoker(torch.nn.Module): + """Own a selected block while calling its preserved installed forward.""" + + def __init__(self, block, saved_forward): + super().__init__() + self.block = block + self._saved_forward = saved_forward + + def forward(self, *args, **kwargs): + return self._saved_forward(*args, **kwargs) + + +@dataclass(frozen=True) +class _Replacement: + state_name: str + leaf_index: int + tensor_indices: tuple[int, ...] + reconstruct: object + + +class _InstalledDispatcher: + def __init__(self, executor, index): + self.executor = executor + self.index = int(index) + + @torch.compiler.disable + def __call__(self, *args, **kwargs): + return self.executor.dispatch(self.index, args, kwargs) + + +class GenericBlockDispatcherRuntime(ImmutableTransformerRuntime): + """Reuse residency/transfer policy while preserving ordinary block math.""" + + def __init__( + self, + model, + residency, + *, + selection, + depth=3, + compile_blocks=True, + compile_dynamic=True, + compile_dynamic_hints=(), + protected_training_leaf_keys=(), + owner_token=None, + ): + self.selection = selection + super().__init__( + model, + residency, + blocks=selection.blocks, + block_keys=selection.block_keys, + entries_by_block=selection.entries_by_block, + depth=depth, + compile_blocks=compile_blocks, + compile_dynamic=compile_dynamic, + compile_dynamic_hints=compile_dynamic_hints, + protected_training_leaf_keys=protected_training_leaf_keys, + owner_token=owner_token, + ) + self._invokers = () + self._dispatchers = () + self._saved_forwards = () + self._replacements = () + + def _replacement_plan(self, index, invoker): + abi = self._block_abis[index] + record = self.residency.arena.block_record(abi.block_key) + replacements = [] + for leaf_index, leaf_name in enumerate(record.leaf_names): + spec = record.leaf_spec(leaf_name) + for substitution in spec.substitutions: + replacements.append( + _Replacement( + state_name=( + f"block.{leaf_name}.{substitution.name}" + ), + leaf_index=leaf_index, + tensor_indices=substitution.tensor_indices, + reconstruct=substitution.reconstruct, + ) + ) + try: + parameter_names = { + name for name, _value in invoker.named_parameters(remove_duplicate=False) + } + except TypeError: + parameter_names = {name for name, _value in invoker.named_parameters()} + try: + buffer_names = { + name for name, _value in invoker.named_buffers(remove_duplicate=False) + } + except TypeError: + buffer_names = {name for name, _value in invoker.named_buffers()} + available = parameter_names | buffer_names + for replacement in replacements: + if replacement.state_name not in available: + raise ImmutableRuntimeError( + f"installed_replacement_target_missing:{abi.block_key}." + f"{replacement.state_name}" + ) + if len({item.state_name for item in replacements}) != len(replacements): + raise ImmutableRuntimeError( + f"installed_replacement_target_conflict:{abi.block_key}" + ) + return tuple(replacements) + + @staticmethod + def _build_state(replacements, leaf_args): + state = {} + for replacement in replacements: + tensors = leaf_args[replacement.leaf_index] + selected = tuple(tensors[index] for index in replacement.tensor_indices) + state[replacement.state_name] = replacement.reconstruct(selected) + return state + + def _get_dispatch_kernel(self, index): + key = (DISPATCHER_GENERATION, int(index)) + existing = self._block_kernels.get(key) + if existing is not None: + return existing + invoker = self._invokers[index] + replacements = self._replacements[index] + + def kernel(leaf_args, args, kwargs): + state = self._build_state(replacements, leaf_args) + return torch.func.functional_call( + invoker, + state, + args, + kwargs, + strict=False, + tie_weights=False, + ) + + if self.compile_blocks: + kernel = torch.compile( + kernel, + mode="default", + fullgraph=False, + dynamic=self.compile_dynamic, + ) + self._block_kernels[key] = kernel + return kernel + + def _dispatcher_program(self, mode): + fingerprint = build_program_fingerprint( + mode, + self._block_abis, + architecture_key=DISPATCHER_GENERATION, + depth=self.depth, + checkpoint_mode="model" if mode == self.TRAIN else "none", + ) + return ImmutableProgram(mode=mode, fingerprint=fingerprint, trunk=None) + + def finalize_execution(self): + if self._finalized: + return self + configure_fetch_runtime(depth=self.depth, owner_token=self.owner_token) + saved = tuple(block.forward for block in self._blocks) + invokers = tuple( + OriginalBlockInvoker(block, saved_forward) + for block, saved_forward in zip(self._blocks, saved, strict=True) + ) + replacements = tuple( + self._replacement_plan(index, invoker) + for index, invoker in enumerate(invokers) + ) + dispatchers = tuple( + _InstalledDispatcher(self, index) + for index in range(len(self._blocks)) + ) + self._saved_forwards = saved + self._invokers = invokers + self._replacements = replacements + self._dispatchers = dispatchers + installed = [] + try: + for block, dispatcher in zip(self._blocks, dispatchers, strict=True): + block.forward = dispatcher + installed.append(block) + except BaseException: + for index, block in enumerate(installed): + block.forward = saved[index] + raise + self._programs = { + self.TRAIN: self._dispatcher_program(self.TRAIN), + self.SAMPLE: self._dispatcher_program(self.SAMPLE), + } + self._finalization_signature = (DISPATCHER_GENERATION,) + self._finalized = True + return self + + def _mark_dispatch_dynamic(self, tensor): + if not self.compile_blocks or not self.compile_dynamic_hints: + return + for dim, lo, hi in self.compile_dynamic_hints: + size = int(tensor.shape[dim]) + if (lo is not None and size < lo) or (hi is not None and size > hi): + self._warn_hint_out_of_range(dim, size, lo, hi) + continue + torch._dynamo.maybe_mark_dynamic(tensor, int(dim)) + + def dispatch(self, index, args, kwargs): + self._require_finalized() + if self._active_token is None: + raise ImmutableRuntimeError( + f"block_dispatch_outside_transformer_execution:{self._block_abis[index].block_key}" + ) + if self._active_mode == self.SAMPLE and torch.is_grad_enabled(): + raise ImmutableRuntimeError( + "immutable_execution_mode_mismatch:active=sample:call=train" + ) + if not args or not isinstance(args[0], torch.Tensor): + raise ImmutableRuntimeError( + f"unsupported_block_arguments:{self._block_abis[index].block_key}" + ) + + source = self._sources.source(index) + transfer = source.transfer + token = None + compact_flat = None + first = args[0] + training = self._active_mode == self.TRAIN + if transfer is not None: + if training and any( + (source.block_key, leaf) in self.protected_training_leaf_keys + for leaf in source.leaf_names + ): + raise ImmutableRuntimeError( + f"uncheckpointed_block_not_resident:{source.block_key}" + ) + host = self.residency.arena.block_record(source.block_key).host_flat + nbytes = int(transfer.compact_nbytes) + guard = first.reshape(-1)[:1].clone() if training else first + token = torch.ops.mm.fetch_start_multi_after( + host, + source.ranges, + nbytes, + guard, + ) + compact_flat = torch.ops.mm.fetch_wait(token, nbytes) + if training and torch.is_grad_enabled(): + first = free_on_backward(first, token) + args = (first, *args[1:]) + + leaf_args = source.assemble_leaf_args(self.residency, compact_flat) + self._mark_dispatch_dynamic(first) + output = self._get_dispatch_kernel(index)(leaf_args, args, kwargs) + if token is not None: + # The first checkpoint pass discards its fetched views, so return + # that slot after forward. Replay runs inside an autograd graph + # task; its token is instead released by free_on_backward after + # the compiled block backward has consumed the substituted state. + if ( + not training + or not torch.is_grad_enabled() + or not _in_backward_graph_task() + ): + torch.ops.mm.fetch_free_after(token, output) + return output + + def close(self): + if self._sources.active_executions: + raise ImmutableRuntimeError("cannot_close_during_execution") + for block, dispatcher, saved in zip( + self._blocks, + self._dispatchers, + self._saved_forwards, + ): + if block.forward is dispatcher: + block.forward = saved + self._dispatchers = () + self._saved_forwards = () + self._invokers = () + self._replacements = () + super().close() + + +def prepare_block_dispatcher_runtime( + transformer, + residency, + *, + selection, + depth=3, + compile_blocks=True, + compile_dynamic=True, + compile_dynamic_hints=(), + protected_training_leaf_keys=(), + owner_token=None, +): + return GenericBlockDispatcherRuntime( + transformer, + residency, + selection=selection, + depth=depth, + compile_blocks=compile_blocks, + compile_dynamic=compile_dynamic, + compile_dynamic_hints=compile_dynamic_hints, + protected_training_leaf_keys=protected_training_leaf_keys, + owner_token=owner_token, + ) diff --git a/toolkit/memory_management/arena_offload/errors.py b/toolkit/memory_management/arena_offload/errors.py new file mode 100644 index 0000000000..84da9d2926 --- /dev/null +++ b/toolkit/memory_management/arena_offload/errors.py @@ -0,0 +1,39 @@ +"""Arena lifecycle error classifications.""" + + +class ArenaSetupFatalError(RuntimeError): + """Arena setup failed after destructive canonical commit.""" + + +class ArenaCleanupError(RuntimeError): + """One or more best-effort arena cleanup operations failed.""" + + def __init__(self, failures): + self.failures = tuple(failures) + detail = "; ".join( + f"{label}: {type(error).__name__}: {error}" + for label, error in self.failures + ) + super().__init__(f"arena runtime cleanup failed: {detail}") + + +def is_fatal_arena_setup(error) -> bool: + """Preserve fatal classification through exception wrapper layers.""" + seen = set() + pending = [error] + while pending: + current = pending.pop() + if current is None or id(current) in seen: + continue + seen.add(id(current)) + if isinstance(current, ArenaSetupFatalError): + return True + pending.extend( + (getattr(current, "__cause__", None), getattr(current, "__context__", None)) + ) + return False + + +def recover_allows_next_job(error, recover: bool) -> bool: + """Whether the configured sequential runner may continue after error.""" + return bool(recover) and not is_fatal_arena_setup(error) diff --git a/toolkit/memory_management/arena_offload/fp8.py b/toolkit/memory_management/arena_offload/fp8.py new file mode 100644 index 0000000000..391061cdf0 --- /dev/null +++ b/toolkit/memory_management/arena_offload/fp8.py @@ -0,0 +1,113 @@ +"""Arena integration for the neutral row-wise FP8 execution policy.""" + +from __future__ import annotations + +import torch + +from toolkit.quantization.fp8_linear import ( + bind_parameter_operation, + declare_fp8_linear, + set_fp8_grad_input_enabled, +) + + +LINEAR_MODULES = {"Linear", "LoRACompatibleLinear", "QLinear"} + + +def _container(child): + container, attribute = child, "forward" + owner_ref = getattr(child, "ara_lora_ref", None) + owner = owner_ref() if callable(owner_ref) else None + if owner is None: + candidate = getattr(getattr(child, "forward", None), "__self__", None) + if candidate is not None and candidate is not child: + owner = candidate + if owner is not None and hasattr(owner, "org_forward"): + container, attribute = owner, "org_forward" + return container, attribute + + +def _live_tensors(child, operation): + """Read the state currently installed by ``functional_call``. + + Canonical arena modules are repointed to host storage between calls. The + generic dispatcher replaces their weight/bias state with resident or + freshly-streamed CUDA tensors for the duration of a block call, so a + compiled FP8 forward must read the module state here instead of closing + over the host tensors seen during setup. + """ + declaration = declare_fp8_linear(child.weight) + if declaration is None: + raise RuntimeError("arena_fp8_live_weight_lost_declaration") + tensors = [declaration.qdata, declaration.scale] + if operation.bias_index is not None: + tensors.append(child.bias) + return tuple(tensors) + + +def enable( + model, + *, + include_ids=None, + live_ids=None, + training: bool, + device=None, +): + """Install bound FP8 operations without owning their execution policy. + + ``live_ids`` identifies arena-managed modules whose tensors may change + device or storage after setup. Their wrappers read the current module + state at call time instead of retaining stale setup-time tensors. + """ + restores = [] + include_ids = None if include_ids is None else set(include_ids) + live_ids = set(live_ids or ()) + for child in model.modules(): + if child.__class__.__name__ not in LINEAR_MODULES: + continue + if include_ids is not None and id(child) not in include_ids: + continue + weight = getattr(child, "weight", None) + if not isinstance(weight, torch.nn.Parameter) or weight.requires_grad: + continue + bias = getattr(child, "bias", None) + operation, tensors = bind_parameter_operation( + weight, + bias, + device=weight.device if device is None else device, + ) + if operation.format_key != "rowwise_fp8" or not operation.native: + continue + container, attribute = _container(child) + original = getattr(container, attribute) + live = id(child) in live_ids + + def installed( + x, + *args, + _child=child, + _live=live, + _tensors=tensors, + _operation=operation, + _original=original, + **kwargs, + ): + if args or kwargs: + return _original(x, *args, **kwargs) + tensors_now = ( + _live_tensors(_child, _operation) if _live else _tensors + ) + if training: + return _operation.forward_train(x, tensors_now) + return _operation.forward_sample(x, tensors_now) + + setattr(container, attribute, installed) + restores.append((container, attribute, original, installed, id(child))) + return restores + + +def disable(restores) -> None: + for restore in reversed(restores): + container, attribute, original, installed = restore[:4] + if getattr(container, attribute, None) is installed: + setattr(container, attribute, original) diff --git a/toolkit/memory_management/arena_offload/layout.py b/toolkit/memory_management/arena_offload/layout.py new file mode 100644 index 0000000000..8d8ac77474 --- /dev/null +++ b/toolkit/memory_management/arena_offload/layout.py @@ -0,0 +1,769 @@ +"""Static host layout and canonical-arena packing primitives. + +This module owns no CUDA stream, execution hook, trace, queue, or transfer +lifetime. It describes host storage, wrapper reconstruction, and typed views. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Iterable + +import torch + +from toolkit.quantization.storage import linear_storage_binding, module_storage_binding +from toolkit.memory_management import pin_manager + +LEAF_ALIGN = 256 + + +@dataclass(frozen=True) +class LeafSpec: + offset: int + nbytes: int + dtype: torch.dtype + shape: tuple[int, ...] + role: str + + +@dataclass(frozen=True) +class LinearSpec: + name: str + tensors: tuple[LeafSpec, ...] + execution_key: tuple + weight_leaf_count: int + weight_template: torch.Tensor + weight_requires_grad: bool + bias_requires_grad: bool + substitutions: tuple = () + + def tensor(self, name: str) -> LeafSpec | None: + return next((item for item in self.tensors if item.role == name), None) + + +@dataclass +class BlockPack: + block_key: str + host_flat: torch.Tensor + linears: tuple[LinearSpec, ...] + required_pin_bytes: int + pinned: bool + view_maker: object | None = None + # Ownership of ``host_flat``'s pin grant. ``pin_handle`` is the PinHandle + # returned by pin_manager.pin_alloc for this pack's OWN flat allocation + # (None when the pack didn't allocate -- e.g. it borrows another owner's + # storage). ``owns_flat`` gates release_pack: a borrowed pack (arena-backed, + # borrowed_from_arena=True) must never release someone else's handle. + pin_handle: object | None = None + owns_flat: bool = True + borrowed_from_arena: bool = False + + +@dataclass(frozen=True) +class LayerStorageView: + """Opaque ordered tensors for one logical layer. + + Storage movement deliberately does not attach an execution operation or + interpret tensor roles. The architecture adapter binds execution after the + immutable storage ABI has been finalized. + """ + + spec: LinearSpec + tensors: tuple[torch.Tensor, ...] + + +def layer_storage_views(pack: BlockPack) -> tuple[LayerStorageView, ...]: + """Expose a block's immutable execution declarations in leaf order.""" + return tuple( + LayerStorageView( + spec=spec, + tensors=tuple(typed_view(pack.host_flat, leaf) for leaf in spec.tensors), + ) + for spec in pack.linears + ) + + +def _rebuild_from_leaves(src, leaves_iter): + try: + names, ctx = src.__tensor_flatten__() + except Exception: + return next(leaves_iter) + moved = {} + for name in names: + inner = getattr(src, name, None) + moved[name] = None if inner is None else _rebuild_from_leaves(inner, leaves_iter) + return type(src).__tensor_unflatten__(moved, ctx, src.size(), src.stride()) + + +def _aligned_offsets(leaves: Iterable[torch.Tensor], align: int = LEAF_ALIGN): + offsets = [] + total = 0 + for leaf in leaves: + total = (total + align - 1) // align * align + offsets.append(total) + total += leaf.numel() * leaf.element_size() + return offsets, total + + +def _empty_host_flat( + nbytes: int, + *, + pin: bool = True, + kind: str = "ingraph_pack", + pin_mechanism: str = "alloc", +) -> tuple[torch.Tensor, bool, object | None]: + if not pin: + return torch.empty(nbytes, dtype=torch.uint8), False, None + if pin_mechanism == "register": + # I1: prepare the page-aligned buffer WITHOUT registering it yet -- + # pack_block_host copies the leaves into it (ordinary pageable + # memcpy) before the caller commits the cudaHostRegister pin. See + # pin_register_prepare's docstring for why population-before-pin is + # faster than registering a virgin buffer. + candidate, padded = pin_manager.pin_register_prepare(nbytes) + return candidate, False, ("register_pending", padded, kind) + else: + handle = pin_manager.pin_alloc( + nbytes, + kind, + required=False, + mode="sampling", + ) + return handle.tensor, bool(handle.pinned), handle + + +def release_pack(pack: "BlockPack | None") -> None: + """Release a pack's own pin grant. + + A borrowed pack (``owns_flat=False``, e.g. arena-backed) must never + release someone else's handle -- the owner (the arena) is responsible for + its own flat's lifetime.""" + if pack is None or not pack.owns_flat: + return + pin_manager.release(pack.pin_handle) + pack.pin_handle = None + + +def pack_block_host( + block_key: str, + linears, + *, + repoint: bool = True, + pin: bool = True, + kind: str = "ingraph_pack", + pin_mechanism: str = "alloc", +) -> BlockPack: + """Pack declared Linear storage tuples into one aligned host buffer.""" + normalized = [] + leaves = [] + for entry in linears: + if len(entry) == 2: + name, module = entry + weight = module.weight + bias = getattr(module, "bias", None) + else: + name, weight, bias = entry + module = None + binding = linear_storage_binding(weight, bias) + tensors = tuple(item.tensor for item in binding.tensors) + normalized.append((name, module, weight, bias, binding, tensors)) + leaves.extend(tensors) + + offsets, total = _aligned_offsets(leaves) + host, pinned, pin_handle = _empty_host_flat( + total, pin=pin, kind=kind, pin_mechanism=pin_mechanism + ) + register_pending = isinstance(pin_handle, tuple) and pin_handle[:1] == ( + "register_pending", + ) + try: + for leaf, offset in zip(leaves, offsets): + nbytes = leaf.numel() * leaf.element_size() + host[offset:offset + nbytes].view(leaf.dtype).reshape(leaf.shape).copy_(leaf) + + if register_pending: + _, _padded, register_kind = pin_handle + pin_handle = pin_manager.pin_register_commit( + host, total, register_kind, required=False + ) + host = pin_handle.tensor + pinned = bool(pin_handle.pinned) + + cursor = 0 + specs = [] + for name, module, weight, bias, binding, tensors in normalized: + tensor_specs = [] + views = [] + for declared, leaf in zip(binding.tensors, tensors): + offset = offsets[cursor] + nbytes = leaf.numel() * leaf.element_size() + tensor_specs.append( + LeafSpec( + offset=offset, + nbytes=nbytes, + dtype=leaf.dtype, + shape=tuple(leaf.shape), + role=declared.name, + ) + ) + views.append( + host[offset:offset + nbytes].view(leaf.dtype).reshape(leaf.shape) + ) + cursor += 1 + + if repoint and module is not None: + weight_views = views[:binding.weight_leaf_count] + weight_view = ( + weight_views[0] + if binding.weight_leaf_count == 1 + else _rebuild_from_leaves(binding.weight_template, iter(weight_views)) + ) + module.weight = torch.nn.Parameter( + weight_view, + requires_grad=getattr(weight, "requires_grad", False), + ) + if bias is not None: + module.bias = torch.nn.Parameter( + views[binding.weight_leaf_count], + requires_grad=getattr(bias, "requires_grad", False), + ) + + specs.append( + LinearSpec( + name=name, + tensors=tuple(tensor_specs), + execution_key=binding.execution_key, + weight_leaf_count=binding.weight_leaf_count, + weight_template=binding.weight_template, + weight_requires_grad=getattr(weight, "requires_grad", False), + bias_requires_grad=( + getattr(bias, "requires_grad", False) + if bias is not None + else False + ), + substitutions=binding.substitutions, + ) + ) + except Exception: + pin_manager.release(pin_handle) + raise + + pack = BlockPack( + block_key=block_key, + host_flat=host, + linears=tuple(specs), + required_pin_bytes=int(total), + pinned=bool(pinned), + pin_handle=pin_handle, + owns_flat=True, + borrowed_from_arena=False, + ) + pack.view_maker = make_block_view_maker(pack) + return pack + + +class ArenaBorrowError(ValueError): + """A block's live params don't actually live in the flat they were + expected to borrow from -- caller must fall back to an owned pack.""" + + +def pack_block_host_from_flat(block_key: str, linears, flat: torch.Tensor) -> "BlockPack": + """Describe declared storage tuples already resident in an owned flat.""" + flat_storage = flat.untyped_storage() + flat_ptr = flat.data_ptr() + flat_end = flat_ptr + flat.numel() * flat.element_size() + + def offset_of(leaf: torch.Tensor) -> int: + if leaf.untyped_storage().data_ptr() != flat_storage.data_ptr(): + raise ArenaBorrowError(f"arena_layout_mismatch:{block_key}:not_in_flat") + offset = leaf.data_ptr() - flat_ptr + nbytes = leaf.numel() * leaf.element_size() + if offset < 0 or offset + nbytes > flat_end - flat_ptr: + raise ArenaBorrowError(f"arena_layout_mismatch:{block_key}:out_of_range") + return offset + + specs = [] + for entry in linears: + if len(entry) == 2: + name, module = entry + weight = module.weight + bias = getattr(module, "bias", None) + else: + name, weight, bias = entry + binding = linear_storage_binding(weight, bias) + tensor_specs = [] + for declared in binding.tensors: + leaf = declared.tensor + tensor_specs.append( + LeafSpec( + offset=offset_of(leaf), + nbytes=leaf.numel() * leaf.element_size(), + dtype=leaf.dtype, + shape=tuple(leaf.shape), + role=declared.name, + ) + ) + specs.append( + LinearSpec( + name=name, + tensors=tuple(tensor_specs), + execution_key=binding.execution_key, + weight_leaf_count=binding.weight_leaf_count, + weight_template=binding.weight_template, + weight_requires_grad=getattr(weight, "requires_grad", False), + bias_requires_grad=( + getattr(bias, "requires_grad", False) + if bias is not None + else False + ), + substitutions=binding.substitutions, + ) + ) + + pack = BlockPack( + block_key=block_key, + host_flat=flat, + linears=tuple(specs), + required_pin_bytes=int(flat.numel() * flat.element_size()), + pinned=bool(pin_manager.is_host_pinned(flat)), + pin_handle=None, + owns_flat=False, + borrowed_from_arena=True, + ) + pack.view_maker = make_block_view_maker(pack) + return pack + + +class IngraphPackError(RuntimeError): + """A block pack could not be built or borrowed. ``reasons`` carries the + stable fail-closed tokens callers surface as ``_ingraph_unavailable_reasons`` + (``non_pinned_pack``, ``unsupported_quant_wrapper``, ``wrapper_pack_missing``, + ``arena_borrow_required``).""" + + def __init__(self, reasons, message: str = ""): + self.reasons = tuple(dict.fromkeys(reasons)) + super().__init__(message or ",".join(self.reasons)) + + +@dataclass +class PackBuildResult: + # Keyed by the model's STABLE block_key string, never a positional index -- + # the caller maps its own indices back locally. + packs: "dict[str, BlockPack]" + borrowed: int + owned: int + pageable: int # always 0 on success (a pageable pack raises non_pinned_pack) + reasons: tuple = () + + +def build_or_borrow_block_packs( + arena, + entries_by_block: dict, + *, + repoint: bool = False, + pin_mechanism: str = "register", + allow_owned_fallback: bool = True, +) -> PackBuildResult: + """Borrow each block's pack from the pinned arena, else build an owned one. + + The single place the in-graph pack policy lives, shared by every model's + ``enable_ingraph_sampling`` / ``enable_ingraph_training`` glue (see the + "in-graph arena protocol" in ``pinned_arena``). Nothing here knows about any + particular model: ``entries_by_block`` maps a stable ``block_key`` to that + block's ``(name, module)`` linear entries, and ``arena`` is duck-typed (any + object exposing ``try_borrow_pack(block_key, entries)``), so this module + never imports ``pinned_arena`` -- which imports it. + + Policy, centralized so callers cannot re-implement it inconsistently: + + * Borrow when the arena holds a current, pinned flat for the block: zero + alloc, zero copy, no second pin of the same bytes. + * Otherwise build an owned pack, but only if ``allow_owned_fallback``. + Under strict pinned-arena validation the caller passes False so a silent + fall back to owned packs cannot make a run "pass" without proving a + single borrow. + * Every streamed pack must be pinned; a pageable one fails the whole set + closed (``non_pinned_pack``) -- strict in-graph is all-or-nothing. + * On any failure, release ONLY packs we own. ``release_pack`` no-ops on a + borrowed pack (``owns_flat=False``), so the arena's flats are never freed + out from under it. + """ + packs: "dict[str, BlockPack]" = {} + borrowed = 0 + owned = 0 + try: + for block_key, raw_entries in entries_by_block.items(): + entries = list(raw_entries) + pack = arena.try_borrow_pack(block_key, entries) if arena is not None else None + if pack is not None: + borrowed += 1 + else: + if not allow_owned_fallback: + raise IngraphPackError( + ("arena_borrow_required",), + f"arena_borrow_required: block {block_key!r} is not " + "borrowable from the pinned arena", + ) + try: + pack = pack_block_host( + block_key, + entries, + repoint=repoint, + pin_mechanism=pin_mechanism, + ) + except ValueError as error: + message = str(error) + reason = ( + "wrapper_pack_missing" + if "wrapper packing" in message + else "unsupported_quant_wrapper" + ) + raise IngraphPackError( + (reason,), f"{reason} ({message})" + ) from error + owned += 1 + packs[block_key] = pack + if any(not pack.pinned for pack in packs.values()): + raise IngraphPackError(("non_pinned_pack",)) + except BaseException: + for pack in packs.values(): + release_pack(pack) + raise + return PackBuildResult(packs=packs, borrowed=borrowed, owned=owned, pageable=0) + + +def is_streamed_module(module) -> bool: + """The memory manager's marker for "this Linear's weights live on the host + and are fetched per call". + + Read it BEFORE stripping compile contaminants -- the strip deletes the + attribute, after which every leaf looks resident. + """ + return hasattr(module, "_layer_memory_manager") + + +def resident_linear_tensors(module) -> tuple: + """Return a Linear's declared opaque storage tuple in stable order.""" + binding = linear_storage_binding(module.weight, getattr(module, "bias", None)) + return tuple(item.tensor for item in binding.tensors) + + +@dataclass(frozen=True) +class BlockLeafPlan: + """Where each of a block's Linear leaves gets its weights this phase. + + The pack is a transfer-coalescing device (one H2D for N leaves), NOT a + residency decision. The memory planner splits residency per-Linear, so a + block is routinely part streamed / part resident. ``sources`` records, in + the caller's canonical leaf order, whether each leaf reads from the fetched + flat (``(True, i)`` -> ``streamed_views[i]``) or straight off its resident + Parameter (``(False, i)`` -> ``resident_args[i]``). Both are trace-time + constants, so the compiled block specializes on its residency pattern. + + ``pack is None`` means every leaf is resident: no flat, no fetch, no token. + """ + + block_key: str + pack: "BlockPack | None" + sources: tuple + resident_args: tuple + borrowed_from_arena: bool = False + + @property + def streams(self) -> bool: + return self.pack is not None + + +def assemble_leaf_args(plan: BlockLeafPlan, streamed_views: tuple = ()) -> tuple: + """Interleave fetched views and resident Parameters back into the block's + canonical leaf order. Pure Python over trace-time constants.""" + return tuple( + streamed_views[index] if from_pack else plan.resident_args[index] + for from_pack, index in plan.sources + ) + + +@dataclass +class BlockPlanResult: + plans: "dict[str, BlockLeafPlan]" + borrowed: int + owned: int + fully_resident: int + streamed_leaves: int + resident_leaves: int + reasons: tuple = () + + +def build_block_leaf_plans( + arena, + entries_by_block: dict, + *, + is_streamed=is_streamed_module, + repoint: bool = False, + pin_mechanism: str = "register", + allow_owned_fallback: bool = True, +) -> BlockPlanResult: + """Plan every block's leaves, packing only the ones the manager streams. + + ``entries_by_block`` maps a stable ``block_key`` to that block's FULL + ``(name, module)`` leaf list in canonical order. This splits each block by + ``is_streamed``, builds/borrows a pack over the streamed subset only, and + reads the resident leaves straight off their Parameters. + + Packing only the streamed subset is what lets the trunk coexist with the + planner's per-Linear residency: a partially-resident block yields a smaller + flat (so a smaller prefetch ring) and skips the fetch entirely for leaves + already on the device. Asking the arena for leaves it never offloaded is + what produced ``borrow refused: stale_modules=3/8``. + """ + streamed_by_block: dict = {} + for block_key, raw_entries in entries_by_block.items(): + streamed = [(name, module) for name, module in raw_entries if is_streamed(module)] + if streamed: + streamed_by_block[block_key] = streamed + + result = build_or_borrow_block_packs( + arena, + streamed_by_block, + repoint=repoint, + pin_mechanism=pin_mechanism, + allow_owned_fallback=allow_owned_fallback, + ) + try: + plans: "dict[str, BlockLeafPlan]" = {} + streamed_leaves = 0 + resident_leaves = 0 + for block_key, raw_entries in entries_by_block.items(): + pack = result.packs.get(block_key) + stream_index = { + name: index + for index, (name, _) in enumerate(streamed_by_block.get(block_key, ())) + } + sources = [] + resident_args = [] + for name, module in raw_entries: + index = stream_index.get(name) + if index is not None: + sources.append((True, index)) + streamed_leaves += 1 + continue + try: + tensors = resident_linear_tensors(module) + except ValueError as error: + raise IngraphPackError( + ("unsupported_quant_wrapper",), + f"unsupported_quant_wrapper ({block_key}.{name}: {error})", + ) from error + sources.append((False, len(resident_args))) + resident_args.append(tensors) + resident_leaves += 1 + plans[block_key] = BlockLeafPlan( + block_key=block_key, + pack=pack, + sources=tuple(sources), + resident_args=tuple(resident_args), + borrowed_from_arena=bool(pack is not None and pack.borrowed_from_arena), + ) + except BaseException: + for pack in result.packs.values(): + release_pack(pack) + raise + return BlockPlanResult( + plans=plans, + borrowed=result.borrowed, + owned=result.owned, + fully_resident=sum(1 for plan in plans.values() if not plan.streams), + streamed_leaves=streamed_leaves, + resident_leaves=resident_leaves, + ) + + +def _flat_view( + flat: torch.Tensor, + offset: int, + nbytes: int, + dtype: torch.dtype, + shape: tuple[int, ...], +) -> torch.Tensor: + return flat[offset:offset + nbytes].view(dtype).reshape(shape) + + +def _flat_clone_view( + flat: torch.Tensor, + offset: int, + nbytes: int, + dtype: torch.dtype, + shape: tuple[int, ...], +) -> torch.Tensor: + return flat[offset:offset + nbytes].clone().view(dtype).reshape(shape) + +def leaf_view(flat: torch.Tensor, spec: LeafSpec) -> torch.Tensor: + return _flat_view(flat, spec.offset, spec.nbytes, spec.dtype, spec.shape) + + +def block_storage_views( + flat: torch.Tensor, + pack: BlockPack, +) -> dict[str, LayerStorageView]: + """Return opaque ordered tensor views without binding execution.""" + out = {} + for spec in pack.linears: + tensors = tuple(leaf_view(flat, item) for item in spec.tensors) + out[spec.name] = LayerStorageView( + spec=spec, + tensors=tensors, + ) + return out + + +def make_block_view_maker(pack: BlockPack): + """Return a flat-buffer view maker that yields only tensor tuples.""" + entries = [] + for spec in pack.linears: + entries.append(tuple( + (item.offset, item.nbytes, item.dtype, item.shape) + for item in spec.tensors + )) + entries = tuple(entries) + + def view_maker(flat: torch.Tensor, _entries=entries): + out = [] + for tensor_entries in _entries: + tensors = [] + for index, item in enumerate(tensor_entries): + view = _flat_view(flat, item[0], item[1], item[2], item[3]) + tensors.append(view if index == 0 else view.clone()) + out.append(tuple(tensors)) + return tuple(out) + + return view_maker + + +def block_tensor_views(flat: torch.Tensor, pack: BlockPack) -> tuple: + maker = pack.view_maker + if maker is None: + maker = make_block_view_maker(pack) + pack.view_maker = maker + return maker(flat) + + + + +@dataclass(frozen=True) +class LeafDescriptor: + role: str + offset: int + nbytes: int + dtype: torch.dtype + shape: tuple[int, ...] + + +@dataclass(frozen=True) +class LinearLayout: + name: str + leaf_descriptors: tuple[LeafDescriptor, ...] + weight_leaf_count: int + weight_requires_grad: bool + bias_requires_grad: bool + execution_key: tuple + weight_template: torch.Tensor | None + substitutions: tuple = () + + def leaf(self, role: str) -> LeafDescriptor | None: + return next((leaf for leaf in self.leaf_descriptors if leaf.role == role), None) + + +@dataclass(frozen=True) +class BlockLayout: + block_key: str + linears: tuple[LinearLayout, ...] + nbytes: int + + +def flatten_leaves(value: torch.Tensor) -> list[torch.Tensor]: + try: + names, _ = value.__tensor_flatten__() + except Exception: + return [value] + leaves = [] + for name in names: + child = getattr(value, name, None) + if child is not None: + leaves.extend(flatten_leaves(child)) + return leaves + + +# Compatibility for legacy manager/residency imports during extraction. +_flatten_leaves = flatten_leaves + + +def rebuild_from_leaves(template: torch.Tensor, leaves: Iterable[torch.Tensor]): + iterator = iter(leaves) + + def rebuild(value): + try: + names, context = value.__tensor_flatten__() + except Exception: + return next(iterator) + children = {} + for name in names: + child = getattr(value, name, None) + children[name] = None if child is None else rebuild(child) + return type(value).__tensor_unflatten__(children, context, value.size(), value.stride()) + + return rebuild(template) + + +def inspect_block(block_key: str, entries) -> BlockLayout: + cursor = 0 + linears = [] + for name, module in entries: + weight = module._parameters.get("weight") + bias = module._parameters.get("bias") + binding = module_storage_binding(module) + leaves = [item.tensor for item in binding.tensors] + roles = [item.name for item in binding.tensors] + descriptors = [] + for role, leaf in zip(roles, leaves): + cursor = (cursor + LEAF_ALIGN - 1) // LEAF_ALIGN * LEAF_ALIGN + nbytes = leaf.numel() * leaf.element_size() + descriptors.append(LeafDescriptor(role, cursor, nbytes, leaf.dtype, tuple(leaf.shape))) + cursor += nbytes + linears.append(LinearLayout( + name=name, + leaf_descriptors=tuple(descriptors), + weight_leaf_count=binding.weight_leaf_count, + weight_requires_grad=bool(weight.requires_grad) if weight is not None else False, + bias_requires_grad=bool(bias.requires_grad) if bias is not None else False, + execution_key=binding.execution_key, + weight_template=binding.weight_template, + substitutions=binding.substitutions, + )) + return BlockLayout(block_key, tuple(linears), cursor) + + +def typed_view(flat: torch.Tensor, leaf: LeafDescriptor) -> torch.Tensor: + return flat[leaf.offset:leaf.offset + leaf.nbytes].view(leaf.dtype).reshape(leaf.shape) + + +def linear_views(flat: torch.Tensor, layout: LinearLayout): + views = tuple(typed_view(flat, leaf) for leaf in layout.leaf_descriptors) + weight_views = views[:layout.weight_leaf_count] + weight = ( + weight_views[0] + if layout.weight_leaf_count == 1 + else rebuild_from_leaves(layout.weight_template, weight_views) + ) + bias = views[layout.weight_leaf_count] if len(views) > layout.weight_leaf_count else None + return weight, bias + + +def substitution_views(flat: torch.Tensor, layout: LinearLayout) -> dict[str, torch.Tensor]: + """Reconstruct every declared module-state target from one flat.""" + views = tuple(typed_view(flat, leaf) for leaf in layout.leaf_descriptors) + return { + substitution.name: substitution.reconstruct( + tuple(views[index] for index in substitution.tensor_indices) + ) + for substitution in layout.substitutions + } diff --git a/toolkit/memory_management/arena_offload/load_session.py b/toolkit/memory_management/arena_offload/load_session.py new file mode 100644 index 0000000000..52f96f68f8 --- /dev/null +++ b/toolkit/memory_management/arena_offload/load_session.py @@ -0,0 +1,180 @@ +"""Generic production checkpoint-to-arena load session.""" + +from __future__ import annotations + +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass, field + +import torch + + +PENDING_CANONICAL_BUILD_ATTR = "_arena_pending_canonical_build" +_CURRENT_SESSION = ContextVar("arena_direct_load_session", default=None) + + +@dataclass +class _DirectLoadSession: + device: object + block_names: tuple[str, ...] + pending: dict[int, tuple[object, object]] = field(default_factory=dict) + unsupported_reason: str | None = None + + def publish(self, model, build) -> None: + if getattr(model, PENDING_CANONICAL_BUILD_ATTR, None) is not None: + raise RuntimeError("arena_direct_load_duplicate_pending_build") + setattr(model, PENDING_CANONICAL_BUILD_ATTR, build) + self.pending[id(model)] = (model, build) + + def claim(self, model): + build = getattr(model, PENDING_CANONICAL_BUILD_ATTR, None) + if build is None: + return None + delattr(model, PENDING_CANONICAL_BUILD_ATTR) + self.pending.pop(id(model), None) + return build + + def rollback_pending(self) -> None: + for model, build in tuple(self.pending.values()): + if getattr(model, PENDING_CANONICAL_BUILD_ATTR, None) is build: + delattr(model, PENDING_CANONICAL_BUILD_ATTR) + build.rollback() + self.pending.clear() + + +def _arena_load_enabled(base_model, enabled) -> bool: + if enabled is not None: + return bool(enabled) + config = getattr(base_model, "model_config", None) + return bool( + config is not None + and getattr(config, "layer_offloading", False) + and getattr(config, "layer_offloading_smart", False) + and not getattr(base_model, "te_only", False) + ) + + +def _managed_source_keys(build) -> set[str]: + return { + source_key + for block_schema in build.state_schema.values() + for source_key in block_schema + } + + +@contextmanager +def model_load_arena_session(base_model, *, enabled=None): + """Offer generic direct arena ingestion during ``base_model.load_model``. + + Unsupported or non-inferable state schemas leave the state mapping intact + and use the ordinary assignment path. A consumed mapping remains owned by + the model until the shared trainer claims it through + ``prepare_arena_offload`` or releases the unfinished preparation. + """ + if not _arena_load_enabled(base_model, enabled): + yield None + return + if _CURRENT_SESSION.get() is not None: + raise RuntimeError("nested_arena_direct_load_session") + block_names = () + provider = getattr(base_model, "get_transformer_block_names", None) + if callable(provider): + block_names = tuple(provider() or ()) + session = _DirectLoadSession( + device=getattr(base_model, "device_torch", None), + block_names=block_names, + ) + token = _CURRENT_SESSION.set(session) + original_load_state_dict = torch.nn.Module.load_state_dict + + def load_state_dict(module, state_dict, strict=True, assign=False): + build = try_prepare_canonical_from_state_dict(module, state_dict) + if build is None: + return original_load_state_dict( + module, state_dict, strict=strict, assign=assign + ) + expected_missing = _managed_source_keys(build) + try: + incompatible = original_load_state_dict( + module, state_dict, strict=False, assign=assign + ) + missing = set(incompatible.missing_keys) + unexpected = set(incompatible.unexpected_keys) + if missing != expected_missing or unexpected: + raise RuntimeError( + "arena_direct_load_residual_mismatch:" + f"missing={sorted(missing - expected_missing)[:5]}:" + f"unconsumed={sorted(expected_missing - missing)[:5]}:" + f"unexpected={sorted(unexpected)[:5]}" + ) + return type(incompatible)([], []) + except BaseException: + discard_pending_canonical_build(module) + raise + + torch.nn.Module.load_state_dict = load_state_dict + try: + yield session + if session.pending: + setattr( + base_model, + "_arena_pending_load_models", + tuple(model for model, _build in session.pending.values()), + ) + except BaseException: + session.rollback_pending() + raise + finally: + torch.nn.Module.load_state_dict = original_load_state_dict + _CURRENT_SESSION.reset(token) + + +def try_prepare_canonical_from_state_dict(model, state_dict): + """Consume inferable managed state into a pending canonical build.""" + session = _CURRENT_SESSION.get() + if session is None: + return None + from .api import prepare_canonical_storage_from_state_dict + from .construction import CanonicalBuildError, CanonicalStateInferenceError + from .discovery import BlockDiscoveryError + from ..canonical_arena import CanonicalArenaError + + try: + build = prepare_canonical_storage_from_state_dict( + model, + state_dict, + block_names=session.block_names, + device=session.device, + ) + except ( + BlockDiscoveryError, + CanonicalArenaError, + CanonicalBuildError, + CanonicalStateInferenceError, + ) as error: + session.unsupported_reason = str(error) + return None + try: + session.publish(model, build) + except BaseException: + build.rollback() + raise + return build + + +def claim_pending_canonical_build(model): + """Take a generic loader build at the normal arena-attach boundary.""" + session = _CURRENT_SESSION.get() + if session is not None: + return session.claim(model) + build = getattr(model, PENDING_CANONICAL_BUILD_ATTR, None) + if build is not None: + delattr(model, PENDING_CANONICAL_BUILD_ATTR) + return build + + +def discard_pending_canonical_build(model) -> None: + """Rollback a pending build after residual assignment fails.""" + build = claim_pending_canonical_build(model) + if build is not None: + build.rollback() diff --git a/toolkit/memory_management/arena_offload/ownership.py b/toolkit/memory_management/arena_offload/ownership.py new file mode 100644 index 0000000000..33f81af6fd --- /dev/null +++ b/toolkit/memory_management/arena_offload/ownership.py @@ -0,0 +1,62 @@ +"""Minimal process-global ownership for the arena transfer runtime.""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class ArenaOwnerToken: + value: object + device: torch.device + + +_LOCK = threading.Lock() +_ACTIVE: ArenaOwnerToken | None = None + + +def normalize_device(device) -> torch.device: + normalized = torch.device(device) + if normalized.type == "cuda" and normalized.index is None: + index = torch.cuda.current_device() if torch.cuda.is_available() else 0 + normalized = torch.device("cuda", index) + return normalized + + +def acquire_process_owner(device) -> ArenaOwnerToken: + global _ACTIVE + requested = normalize_device(device) + with _LOCK: + if _ACTIVE is not None: + raise RuntimeError( + "arena_runtime_already_active:\n" + f"active_device={_ACTIVE.device} requested_device={requested}" + ) + token = ArenaOwnerToken(object(), requested) + _ACTIVE = token + return token + + +def validate_process_owner(token: ArenaOwnerToken | None) -> ArenaOwnerToken: + with _LOCK: + if token is None or _ACTIVE is not token: + raise RuntimeError("arena_runtime_owner_mismatch") + return token + + +def release_process_owner(token: ArenaOwnerToken | None) -> None: + global _ACTIVE + with _LOCK: + if token is None: + return + if _ACTIVE is not token: + raise RuntimeError("arena_runtime_owner_mismatch") + _ACTIVE = None + + +def active_process_owner() -> ArenaOwnerToken | None: + with _LOCK: + return _ACTIVE diff --git a/toolkit/memory_management/arena_offload/planner.py b/toolkit/memory_management/arena_offload/planner.py new file mode 100644 index 0000000000..dfa28038fa --- /dev/null +++ b/toolkit/memory_management/arena_offload/planner.py @@ -0,0 +1,221 @@ +"""Cold-start, whole-block residency planning owned by arena offload.""" + +from __future__ import annotations + +import torch + +from toolkit.quantization.storage import temporary_materialization_bytes + +from .. import vram_budget +from .layout import flatten_leaves + +GIB = 1024**3 +DEFAULT_AUTO_WORKING_RESERVE_GIB = 5.0 +# Full residency removes the transfer ring, but training still needs activation, +# adapter, dequantization, and allocator-fragmentation headroom. A production +# Orbit4 full-model smoke exhausted a 2 GiB reserve during checkpoint backward; +# 4 GiB is the narrowest evidence-backed automatic bound. Explicit policy values +# remain authoritative for workloads with measured tighter requirements. +DEFAULT_ALL_RESIDENT_WORKING_RESERVE_GIB = 4.0 +DEFAULT_RESIDENT_FLOOR_GIB = 2.0 + + +def resolve_margin_gib(device, value, *, hard_gib=0.0) -> float: + try: + margin = float(value) + automatic = margin < 0 + except (TypeError, ValueError): + automatic = value is None or str(value).strip().lower() == "auto" + margin = -1.0 + if automatic: + margin = vram_budget.auto_margin_gib(device) + return max(float(margin), float(hard_gib or 0.0)) + + +def training_pinned_keys_for_keep_last(model, keep_last, block_keys=None) -> set[str]: + keep_last = max(0, int(keep_last or 0)) + if block_keys is not None: + ordered = tuple(str(key) for key in block_keys) + return set(ordered[max(0, len(ordered) - keep_last):]) + blocks = getattr(model, "blocks", None) + if blocks is None or keep_last <= 0: + return set() + start = max(0, len(blocks) - keep_last) + return {f"blocks.{index}" for index in range(start, len(blocks))} + + +def _tensor_storage_bytes(value) -> int: + try: + leaves = flatten_leaves(value) + except Exception: + leaves = (value,) + return sum( + int(leaf.numel() * leaf.element_size()) + for leaf in leaves + if isinstance(leaf, torch.Tensor) and leaf.device.type != "meta" + ) + + +def _singleton_stats(model, canonical_modules) -> tuple[int, int, set[int]]: + canonical_ids = {id(module) for module in canonical_modules} + seen_parameters = set() + total = 0 + largest_materialization = 0 + runtime_ids = set() + for module in model.modules(): + if id(module) in canonical_ids: + continue + direct = tuple(module.parameters(recurse=False)) + if direct: + runtime_ids.add(id(module)) + for parameter in direct: + if id(parameter) in seen_parameters: + continue + seen_parameters.add(id(parameter)) + total += _tensor_storage_bytes(parameter.data) + largest_materialization = max( + largest_materialization, + temporary_materialization_bytes(parameter.data), + ) + return total, largest_materialization, runtime_ids + + +def build_training_plan( + model, arena, canonical_modules, device, config, *, block_keys=None +) -> dict: + """Choose an initial whole-block layout without the legacy manager.""" + device = torch.device(device) + policy = config._policy + try: + working_value = float(policy.working_reserve_gib) + automatic = working_value < 0 + except (TypeError, ValueError): + automatic = policy.working_reserve_gib is None or str( + policy.working_reserve_gib + ).strip().lower() == "auto" + working_value = DEFAULT_AUTO_WORKING_RESERVE_GIB + if automatic: + working_value = DEFAULT_AUTO_WORKING_RESERVE_GIB + + hard_gib = float(policy.wddm_hard_gib or 1.0) + margin_gib = resolve_margin_gib( + device, policy.wddm_margin_gib, hard_gib=hard_gib + ) + free_bytes, total_bytes = vram_budget.device_mem_info(device) + working_bytes = int(max(0.0, working_value) * GIB) + margin_bytes = int(margin_gib * GIB) + hard_bytes = int(hard_gib * GIB) + + singleton_bytes, largest_singleton_dequant, runtime_ids = _singleton_stats( + model, canonical_modules + ) + records = [arena.block_record(key) for key in arena.block_keys()] + records = [record for record in records if record is not None] + block_bytes = sum(record.committed_bytes for record in records) + all_resident_working_bytes = working_bytes + if automatic: + all_resident_working_bytes = min( + working_bytes, + int(DEFAULT_ALL_RESIDENT_WORKING_RESERVE_GIB * GIB), + ) + all_resident_fit = ( + singleton_bytes + block_bytes + all_resident_working_bytes + <= max(0, int(free_bytes) - margin_bytes) + ) + if all_resident_fit: + # A transfer ring and the generic 5 GiB cold-start reserve are both + # counterproductive when the complete compressed model plus the + # evidence-backed execution reserve fits below the physical WDDM margin. + working_bytes = all_resident_working_bytes + + pinned_keys = training_pinned_keys_for_keep_last( + model, policy.checkpoint_keep_last, block_keys=block_keys + ) + resident_keys = ( + {record.block_key for record in records} + if all_resident_fit + else { + record.block_key + for record in records + if record.block_key in pinned_keys + } + ) + streamed = [record for record in records if record.block_key not in resident_keys] + largest_stream = max((record.committed_bytes for record in streamed), default=0) + ring_bytes = largest_stream * max(1, int(policy.prefetch_depth)) + usable = max(0, int(free_bytes) - margin_bytes - working_bytes) + resident_budget = max(0, usable - singleton_bytes - ring_bytes) + resident_bytes = sum( + record.committed_bytes for record in records if record.block_key in resident_keys + ) + floor = min(block_bytes, int(DEFAULT_RESIDENT_FLOOR_GIB * GIB)) + + candidates = sorted( + (record for record in records if record.block_key not in resident_keys), + key=lambda record: (record.committed_bytes, record.block_key), + ) + target = ( + block_bytes + if all_resident_fit + else resident_budget if not automatic else min(resident_budget, floor) + ) + for record in candidates: + if resident_bytes >= target: + break + if resident_bytes + record.committed_bytes > resident_budget: + continue + resident_keys.add(record.block_key) + resident_bytes += record.committed_bytes + + offload_ids = { + id(module) + for record in records + if record.block_key not in resident_keys + for module in record.modules + } + protected = frozenset( + (record.block_key, leaf_name) + for record in records + if record.block_key in pinned_keys + for leaf_name in record.leaf_names + ) + try: + device_used = torch.cuda.device_memory_used(device) + torch_reserved = torch.cuda.memory_reserved(device) + system_reserve = max(0, int(device_used) - int(torch_reserved)) + except Exception: + system_reserve = max(0, int(total_bytes) - int(free_bytes)) + + return { + "offload_ids": offload_ids, + "offloaded_layers": len(offload_ids), + "candidate_layers": sum(len(record.modules) for record in records), + "model_bytes": singleton_bytes + block_bytes, + "resident_bytes": singleton_bytes + resident_bytes, + "must_resident_bytes": singleton_bytes + sum( + record.committed_bytes for record in records if record.block_key in pinned_keys + ), + "must_resident_layer_keys": set(), + "pinned_resident_bytes": sum( + record.committed_bytes for record in records if record.block_key in pinned_keys + ), + "pinned_resident_keys": set(pinned_keys), + "protected_training_leaf_keys": protected, + "generic_resident_bytes": max(0, resident_bytes), + "ring_bytes": ring_bytes, + "gpu_stream_need_bytes": ring_bytes, + "gpu_stream_budget_bytes": ring_bytes, + "working_reserve_bytes": working_bytes, + "wddm_margin_bytes": margin_bytes, + "wddm_hard_bytes": hard_bytes, + "system_reserve_bytes": system_reserve, + "usable_bytes": usable, + "free_bytes": int(free_bytes), + "fits": singleton_bytes + resident_bytes + ring_bytes <= usable, + "singleton_resident_bytes": singleton_bytes, + "largest_singleton_bf16_dequant_bytes": largest_singleton_dequant, + "singleton_runtime_ids": runtime_ids, + "auto_working_reserve": automatic, + "all_resident_fit": all_resident_fit, + "all_resident_working_reserve_bytes": all_resident_working_bytes, + } diff --git a/toolkit/memory_management/arena_offload/policy.py b/toolkit/memory_management/arena_offload/policy.py new file mode 100644 index 0000000000..616932bece --- /dev/null +++ b/toolkit/memory_management/arena_offload/policy.py @@ -0,0 +1,529 @@ +"""Arena-native training signals and two-timescale residency policy.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .. import vram_budget + +_ALLOC = ("num_alloc_retries", "num_device_alloc", "num_device_free") +_COMPILE = ("frames", "graphs", "graph_breaks") + +DEFAULT_SLACK_PAD_BYTES = 256 * 1024**2 +AGGRESSIVE_PROMOTION_MIN_CAPACITY = 4 + + +def transfer_benefits_from_residency(transfer) -> bool: + """Return whether a valid window proves that weights are still streaming.""" + if not transfer or transfer.get("h2d_duty_overflow"): + return False + return ( + int(transfer.get("bytes", 0) or 0) > 0 + and float(transfer.get("h2d_ms", 0.0) or 0.0) > 0.0 + ) + + +@dataclass(frozen=True) +class PolicyDecision: + action: str + block_key: str | None = None + block_bytes: int = 0 + target_cap_bytes: int | None = None + reason: str = "" + block_keys: tuple[str, ...] = () + + +class ArenaResidencyController: + """Stateful wiring around the pure two-timescale residency FSM.""" + + def __init__(self, *, slack_pad_bytes=DEFAULT_SLACK_PAD_BYTES): + self.state = vram_budget.ResidencyFsmState() + self.slack_pad_bytes = max(0, int(slack_pad_bytes)) + self.last_action = "hold" + self.last_reason = "cold_start" + self.last_promoted_key = None + self.last_block_key = None + self.last_block_bytes = 0 + self.last_target_cap_bytes = None + self.last_worst_shape_margin_bytes = None + self.last_throughput_gate = None + self.last_promote_gate = None + self.last_cap_covers_promo = None + self.last_worst_shape_allocator_slack_bytes = None + self.last_aggressive_capacity = 0 + self.last_aggressive_gate = False + self.pending_promotion = None + self.last_safe_residency_bytes = None + self.last_rejected_residency_bytes = None + self.bootstrapped = False + + def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, + worst_shape_free_bytes, worst_shape_allocator_slack_bytes=None, + current_cap_bytes=None, aggressive_promotion_capacity=0): + if not self.bootstrapped: + self.bootstrapped = True + if demote_candidate is not None: + return self._decision( + "demote", + demote_candidate, + reason="approach_from_below", + ) + if signal is None: + return self._hold("awaiting_signal") + + block_bytes = 0 if candidate is None else int(candidate["block_bytes"]) + allocator = signal.get("allocator") or {} + retries = int(allocator.get("alloc_retries_delta", 0) or 0) + device_frees = int(allocator.get("free_count_delta", 0) or 0) + allocator_slack = int( + signal.get("reclaimable_at_peak_bytes", 0) + if worst_shape_allocator_slack_bytes is None + else worst_shape_allocator_slack_bytes + ) + throughput_ok = transfer_benefits_from_residency(signal.get("transfer")) + worst_ok = candidate is not None and int(worst_shape_free_bytes) >= 0 + promote_ok = ( + candidate is not None + and throughput_ok + and worst_ok + and vram_budget.residency_promote_ok( + retries, allocator_slack, block_bytes, self.slack_pad_bytes + ) + ) + active_cap = ( + int(cliff_cap_bytes) + if current_cap_bytes is None + else int(current_cap_bytes) + ) + cap_covers = ( + candidate is not None + and allocator_slack > block_bytes + self.slack_pad_bytes + ) + binding = retries > 0 or int(worst_shape_free_bytes) < 0 + aggressive_capacity = max(0, int(aggressive_promotion_capacity or 0)) + bootstrap_pending = bool( + self.pending_promotion is not None + and self.pending_promotion.get("block_keys") + ) + aggressive_ok = ( + candidate is not None + and aggressive_capacity >= AGGRESSIVE_PROMOTION_MIN_CAPACITY + and not bootstrap_pending + and not bool(signal.get("compile_invalid")) + and retries == 0 + and device_frees == 0 + and worst_ok + and allocator_slack > block_bytes + self.slack_pad_bytes + ) + self.last_worst_shape_margin_bytes = int(worst_shape_free_bytes) + self.last_worst_shape_allocator_slack_bytes = allocator_slack + self.last_throughput_gate = bool(throughput_ok) + self.last_promote_gate = bool(promote_ok) + self.last_cap_covers_promo = bool(cap_covers) + self.last_aggressive_capacity = aggressive_capacity + self.last_aggressive_gate = bool(aggressive_ok) + cap_raise_bytes = ( + block_bytes + if promote_ok and block_bytes > 0 + else self.slack_pad_bytes + ) + needed_cap = min( + int(cliff_cap_bytes), + active_cap + max(0, int(cap_raise_bytes)), + ) + if ( + self.pending_promotion is not None + and (retries > 0 or device_frees > 0) + ): + return self.reject_pending_promotion("promotion_allocator_gc") + + # Abundant measured headroom does not need the multi-window transfer + # proof. Spend only one block per step, leaving at least three blocks of + # measured capacity in reserve, and make the new block the rollback + # candidate for the next boundary. + if aggressive_ok: + self._begin_promotion( + candidate, + resident_bytes_before=int(signal.get("resident_bytes", 0) or 0), + active_cap_bytes=active_cap, + ) + return self._decision( + "promote", candidate, reason="abundant_four_block_headroom" + ) + + previous_state = self.state.name + self.state, action = vram_budget.residency_fsm_step( + self.state, + { + "measurements_invalid": bool(signal.get("compile_invalid")), + "binding": binding, + "cap_can_relieve": ( + active_cap < needed_cap <= int(cliff_cap_bytes) + ), + "promote_gate": promote_ok, + "cap_covers_promo": cap_covers, + }, + ) + if ( + previous_state == vram_budget.FSM_PROMOTION_VERIFY + and self.state.name == vram_budget.FSM_STABLE + ): + self.pending_promotion = None + self.last_promoted_key = None + if action == vram_budget.ACT_PROMOTE and candidate is not None: + self._begin_promotion( + candidate, + resident_bytes_before=int(signal.get("resident_bytes", 0) or 0), + active_cap_bytes=active_cap, + ) + return self._decision( + action, candidate, reason="safe_transfer_benefit" + ) + if ( + action == vram_budget.ACT_ROLLBACK + and self.pending_promotion is not None + ): + return self.reject_pending_promotion("promotion_bound") + if action == vram_budget.ACT_DEMOTE and demote_candidate is not None: + return self._decision( + action, demote_candidate, reason="live_pressure" + ) + if action == vram_budget.ACT_RAISE_CAP: + self.last_action = action + self.last_reason = "prefund_or_relieve" + self.last_block_key = None + self.last_block_bytes = 0 + self.last_target_cap_bytes = needed_cap + return PolicyDecision( + action, target_cap_bytes=needed_cap, reason=self.last_reason + ) + reason = ( + "worst_shape_veto" + if candidate is not None and not worst_ok + else ( + "throughput_gate" + if candidate is not None and not throughput_ok + else ( + "allocator_headband" + if candidate is not None + and allocator_slack <= block_bytes + self.slack_pad_bytes + else "fsm_hold" + ) + ) + ) + return self._hold(reason, candidate=candidate) + + def reject_pending_promotion(self, reason): + pending = self.pending_promotion + if pending is None: + return self._hold(reason) + self.last_safe_residency_bytes = int( + pending["resident_bytes_before"] + ) + self.last_rejected_residency_bytes = int( + pending["resident_bytes_after"] + ) + self.pending_promotion = None + self.last_promoted_key = None + self.state = vram_budget.ResidencyFsmState( + vram_budget.FSM_COOLDOWN, 0 + ) + self.last_action = "rollback" + self.last_reason = reason + self.last_block_key = pending["block_key"] + self.last_block_bytes = int(pending["block_bytes"]) + self.last_target_cap_bytes = int(pending["previous_cap_target_bytes"]) + return PolicyDecision( + "rollback", + pending["block_key"], + int(pending["block_bytes"]), + target_cap_bytes=int(pending["previous_cap_target_bytes"]), + reason=reason, + block_keys=tuple( + pending.get("block_keys") or (pending["block_key"],) + ), + ) + + def begin_bootstrap_promotion( + self, block_keys, block_bytes, resident_bytes_before, active_cap_bytes + ): + keys = tuple(str(key) for key in block_keys) + self.last_promoted_key = keys[-1] if keys else None + self.pending_promotion = { + "block_key": self.last_promoted_key, + "block_keys": keys, + "block_bytes": int(block_bytes), + "resident_bytes_before": int(resident_bytes_before), + "resident_bytes_after": int(resident_bytes_before) + int(block_bytes), + "previous_cap_target_bytes": int(active_cap_bytes), + } + self.state = vram_budget.ResidencyFsmState( + vram_budget.FSM_PROMOTION_VERIFY, 0 + ) + self.last_action = "promote" + self.last_reason = "bootstrap_physical_free" + self.last_block_key = self.last_promoted_key + self.last_block_bytes = int(block_bytes) + + def _begin_promotion( + self, candidate, *, resident_bytes_before, active_cap_bytes + ): + block_bytes = int(candidate["block_bytes"]) + self.last_promoted_key = candidate["block_key"] + self.pending_promotion = { + "block_key": candidate["block_key"], + "block_bytes": block_bytes, + "resident_bytes_before": int(resident_bytes_before), + "resident_bytes_after": int(resident_bytes_before) + block_bytes, + "previous_cap_target_bytes": int(active_cap_bytes), + } + self.state = vram_budget.ResidencyFsmState( + vram_budget.FSM_PROMOTION_VERIFY, 0 + ) + + def allocation_failure(self): + return self.reject_pending_promotion("promotion_allocation_failure") + + def _decision(self, action, candidate, *, reason): + self.last_action = action + self.last_reason = reason + self.last_block_key = candidate["block_key"] + self.last_block_bytes = int(candidate["block_bytes"]) + self.last_target_cap_bytes = None + return PolicyDecision( + action, + candidate["block_key"], + int(candidate["block_bytes"]), + reason=reason, + ) + + def _hold(self, reason, *, candidate=None): + self.last_action = "hold" + self.last_reason = reason + self.last_block_key = ( + None if candidate is None else candidate["block_key"] + ) + self.last_block_bytes = ( + 0 if candidate is None else int(candidate["block_bytes"]) + ) + self.last_target_cap_bytes = None + return PolicyDecision("hold", reason=reason) + + def diagnostics(self): + return { + "state": self.state.name, + "windows_in_state": self.state.windows_in_state, + "last_action": self.last_action, + "last_reason": self.last_reason, + "last_promoted_key": self.last_promoted_key, + "last_block_key": self.last_block_key, + "last_block_bytes": self.last_block_bytes, + "last_target_cap_bytes": self.last_target_cap_bytes, + "last_worst_shape_margin_bytes": self.last_worst_shape_margin_bytes, + "last_throughput_gate": self.last_throughput_gate, + "last_promote_gate": self.last_promote_gate, + "last_cap_covers_promo": self.last_cap_covers_promo, + "last_worst_shape_allocator_slack_bytes": ( + self.last_worst_shape_allocator_slack_bytes + ), + "last_aggressive_capacity": self.last_aggressive_capacity, + "last_aggressive_gate": self.last_aggressive_gate, + "pending_promotion": self.pending_promotion, + "last_safe_residency_bytes": self.last_safe_residency_bytes, + "last_rejected_residency_bytes": ( + self.last_rejected_residency_bytes + ), + "slack_pad_bytes": self.slack_pad_bytes, + } + + +def _deltas(previous, current, keys, cast=int): + result = {} + for key in keys: + now = cast((current or {}).get(key, 0) or 0) + before = cast((previous or {}).get(key, 0) or 0) + result[key] = now - before if now >= before else now + return result + + +@dataclass(frozen=True) +class ShapePeak: + steps: int = 0 + warmup_steps: int = 1 + peak_allocated_bytes: int = 0 + peak_reserved_bytes: int = 0 + working_peak_bytes: int = 0 + + +class TrainingSignalWindow: + """Runtime-owned, CPU-testable training policy observations.""" + + def __init__(self, *, transfer_window_steps=4): + self.transfer_window_steps = max(2, int(transfer_window_steps)) + self._allocator_previous = None + self._compile_previous = None + self._transfer_previous = None + self._transfer_steps = 0 + self._transfer_wall_ms = 0.0 + self._transfer_h2d_ms = 0.0 + self._transfer_bytes = 0 + self._shape_peaks = {} + self._last_signal = None + + @property + def shape_peaks(self): + return dict(self._shape_peaks) + + @property + def last_signal(self): + return None if self._last_signal is None else dict(self._last_signal) + + @property + def transfer_snapshot_due(self): + return self._transfer_steps + 1 >= self.transfer_window_steps + + def invalidate_shape_peaks(self): + self._shape_peaks.clear() + self._last_signal = None + + def observe( + self, *, shape_key, step_num, allocator_counters, + peak_allocated_bytes, peak_reserved_bytes, device_free_bytes, + resident_bytes, ring_bytes, compile_counters=None, + transfer_counters=None, step_wall_ms=0.0, + ): + alloc_delta = _deltas(self._allocator_previous, allocator_counters, _ALLOC) + self._allocator_previous = { + key: int((allocator_counters or {}).get(key, 0) or 0) for key in _ALLOC + } + compile_delta = _deltas( + self._compile_previous, compile_counters, _COMPILE + ) + compile_invalid = compile_counters is not None and compile_delta["frames"] > 0 + if compile_counters is not None: + self._compile_previous = { + key: int(compile_counters.get(key, 0) or 0) for key in _COMPILE + } + if compile_invalid: + self.invalidate_shape_peaks() + + key = _shape_key(shape_key) + if not compile_invalid: + self._record_shape_peak( + key, + int(peak_allocated_bytes or 0), + int(peak_reserved_bytes or 0), + int(resident_bytes or 0), + int(ring_bytes or 0), + ) + transfer = self._observe_transfer(transfer_counters, float(step_wall_ms or 0.0)) + allocated = int(peak_allocated_bytes or 0) + reserved = int(peak_reserved_bytes or 0) + resident = int(resident_bytes or 0) + ring = int(ring_bytes or 0) + signal = { + "shape_key": key, + "step_num": None if step_num is None else int(step_num), + "allocator": { + "alloc_retries_delta": alloc_delta["num_alloc_retries"], + "alloc_count_delta": alloc_delta["num_device_alloc"], + "free_count_delta": alloc_delta["num_device_free"], + }, + "peak_allocated_bytes": allocated, + "peak_reserved_bytes": reserved, + "reclaimable_at_peak_bytes": max(0, reserved - allocated), + "device_free_bytes": int(device_free_bytes or 0), + "resident_bytes": resident, + "ring_bytes": ring, + "live_bytes": resident + ring, + "compile_invalid": bool(compile_invalid), + "compile_delta": compile_delta, + "transfer": transfer, + } + self._last_signal = signal + return dict(signal) + + def diagnostics(self): + return { + "last_signal": self.last_signal, + "shape_peaks": [ + { + "shape_key": key, + "steps": peak.steps, + "warmup_steps": peak.warmup_steps, + "peak_allocated_bytes": peak.peak_allocated_bytes, + "peak_reserved_bytes": peak.peak_reserved_bytes, + "working_peak_bytes": peak.working_peak_bytes, + } + for key, peak in self._shape_peaks.items() + ], + "transfer_window_steps": self.transfer_window_steps, + "transfer_pending_steps": self._transfer_steps, + } + + def _record_shape_peak(self, key, allocated, reserved, resident, ring): + previous = self._shape_peaks.get(key) + if previous is None: + self._shape_peaks[key] = ShapePeak() + return + self._shape_peaks[key] = ShapePeak( + steps=previous.steps + 1, + warmup_steps=previous.warmup_steps, + peak_allocated_bytes=max(previous.peak_allocated_bytes, allocated), + peak_reserved_bytes=max(previous.peak_reserved_bytes, reserved), + working_peak_bytes=max( + previous.working_peak_bytes, + max(0, allocated - resident - ring), + ), + ) + + def _observe_transfer(self, counters, wall_ms): + self._transfer_steps += 1 + self._transfer_wall_ms += max(0.0, wall_ms) + if counters is not None: + h2d = _deltas( + self._transfer_previous, counters, ("h2d_ms",), float + )["h2d_ms"] + byte_count = _deltas( + self._transfer_previous, counters, ("bytes",) + )["bytes"] + self._transfer_previous = { + "h2d_ms": float(counters.get("h2d_ms", 0.0) or 0.0), + "bytes": int(counters.get("bytes", 0) or 0), + } + self._transfer_h2d_ms += h2d + self._transfer_bytes += byte_count + if self._transfer_steps < self.transfer_window_steps: + return None + duty = ( + None if self._transfer_wall_ms <= 0.0 + else 100.0 * self._transfer_h2d_ms / self._transfer_wall_ms + ) + gbps = ( + None if self._transfer_h2d_ms <= 0.0 + else self._transfer_bytes / (self._transfer_h2d_ms * 1_000_000.0) + ) + result = { + "steps": self._transfer_steps, + "step_wall_ms": self._transfer_wall_ms, + "h2d_ms": self._transfer_h2d_ms, + "bytes": self._transfer_bytes, + "h2d_duty_pct": duty, + "h2d_duty_overflow": bool(duty is not None and duty > 100.0), + "achieved_gbps": gbps, + } + self._transfer_steps = 0 + self._transfer_wall_ms = 0.0 + self._transfer_h2d_ms = 0.0 + self._transfer_bytes = 0 + return result + + +def _shape_key(value): + if value is None: + return ("unknown",) + if isinstance(value, tuple): + return value + if isinstance(value, list): + return tuple(value) + return (value,) diff --git a/toolkit/memory_management/arena_offload/resources.py b/toolkit/memory_management/arena_offload/resources.py new file mode 100644 index 0000000000..0f41cd1be8 --- /dev/null +++ b/toolkit/memory_management/arena_offload/resources.py @@ -0,0 +1,165 @@ +"""Preparation-scoped ownership for one arena runtime.""" + +from __future__ import annotations + +from .errors import ArenaCleanupError +from .ownership import acquire_process_owner, release_process_owner + + +class ArenaRuntimeResources: + """Own resources from pre-commit preparation through runtime close.""" + + def __init__(self, model, device) -> None: + self.model = model + self.device = device + self.owner_token = None + self.canonical_build = None + self.arena = None + self.residency = None + self.executor = None + self.runtime = None + self.fp8_restores = [] + self.published_attributes = [] + self.canonical_modules = () + self.canonical_committed = False + self.closing = False + self.disposed = False + self.released = False + self._releasing = False + self._original_movement = { + name: getattr(model, name, None) for name in ("to", "cuda", "cpu") + } + self._original_forward = getattr(model, "forward", None) + + def acquire_process_owner(self) -> None: + if self.owner_token is None: + self.owner_token = acquire_process_owner(self.device) + + def adopt_canonical_build(self, build) -> None: + self.canonical_build = build + self.arena = build.arena + build._arena_resources = self + + def mark_canonical_committed(self) -> None: + self.canonical_committed = True + if self.canonical_build is not None: + # Rollback is invalid after the destructive boundary. Drop its + # original Parameter references immediately so the resource owner + # does not retain a second model-sized payload for runtime life. + self.canonical_build._arena_resources = None + self.canonical_build._originals.clear() + self.canonical_build = None + + def adopt_residency(self, residency) -> None: + self.residency = residency + + def adopt_executor(self, executor) -> None: + self.executor = executor + + def adopt_runtime(self, runtime) -> None: + self.runtime = runtime + + def record_fp8_restore(self, label, restore) -> None: + self.fp8_restores.append((label, restore)) + + def record_published_attribute(self, owner, name, value) -> None: + setattr(owner, name, value) + self.published_attributes.append((owner, name, value)) + + def _install_disposed_movement_guard(self) -> None: + def disposed(*_args, **_kwargs): + raise RuntimeError("arena_offload_transformer_disposed") + + for name in self._original_movement: + if self._original_movement[name] is not None: + setattr(self.model, name, disposed) + if self._original_forward is not None: + setattr(self.model, "forward", disposed) + setattr(self.model, "_arena_offload_disposed", True) + + def release(self) -> None: + if self.released or self._releasing: + return + self._releasing = True + self.closing = True + failures = [] + + if self.executor is not None and getattr( + self.executor, "active_executions", 0 + ): + self._releasing = False + self.closing = False + raise ArenaCleanupError( + (("active execution", RuntimeError("cannot_close_during_execution")),) + ) + + def attempt(label, operation): + try: + operation() + except BaseException as error: + failures.append((label, error)) + + from . import transfer + + try: + transfer_cleanup_failed = False + if self.owner_token is not None: + try: + transfer.drain_fetch_runtime(owner_token=self.owner_token) + except BaseException as error: + transfer_cleanup_failed = True + failures.append(("transfer tickets", error)) + if self.executor is not None: + attempt("immutable executor", self.executor.close) + for label, restore in reversed(self.fp8_restores): + attempt(label, restore) + self.fp8_restores.clear() + if self.residency is not None: + attempt("resident sidecars", self.residency.clear) + + for owner, name, value in reversed(self.published_attributes): + if getattr(owner, name, None) is value: + attempt( + f"published attribute {name}", + lambda owner=owner, name=name: delattr(owner, name), + ) + self.published_attributes.clear() + + if self.canonical_committed: + if self.arena is not None: + attempt( + "movement interception", + lambda: self.arena.unguard_whole_model_to(self.model), + ) + attempt("canonical arena", self.arena.release) + self.disposed = True + self._install_disposed_movement_guard() + elif self.canonical_build is not None: + attempt("canonical build", self.canonical_build.rollback) + + if self.owner_token is not None: + try: + transfer.release_fetch_runtime(self.owner_token) + except BaseException as error: + transfer_cleanup_failed = True + failures.append(("transfer runtime", error)) + if not transfer_cleanup_failed: + try: + release_process_owner(self.owner_token) + except BaseException as error: + failures.append(("process ownership", error)) + else: + self.owner_token = None + finally: + self.closing = False + self.released = self.owner_token is None + self._releasing = False + if self.runtime is not None: + self.runtime._closed = True + self.runtime._disposed = bool(self.canonical_committed) + self.canonical_modules = () + self.canonical_build = None + self.residency = None + self.executor = None + if failures: + raise ArenaCleanupError(failures) diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py new file mode 100644 index 0000000000..21c2ffebe1 --- /dev/null +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -0,0 +1,1217 @@ +"""The arena offload runtime facade. + +One object owns the arena, the residency state, the training plan, and the +immutable block executor. Model integrations and the shared trainer hold a +reference to it and nothing else -- no `_mm_*` field reads, no arena +construction, no residency manipulation. + +Cold planning, live policy, FP8 transforms, and teardown are arena-owned. The +legacy per-linear manager is deliberately outside this package. +""" + +from __future__ import annotations + +import contextlib +import time +from collections.abc import Sequence +from dataclasses import replace +from typing import Any + +from .. import allocator_cap +from ..canonical_arena import CanonicalArena +from ..residency import ResidencyPlan, ResidencyState +from ..vram_budget import apply_simulated_card +from .policy import ( + AGGRESSIVE_PROMOTION_MIN_CAPACITY, + ArenaResidencyController, + TrainingSignalWindow, +) +from .errors import ArenaCleanupError, ArenaSetupFatalError +from .fp8 import disable as disable_fp8 +from .fp8 import enable as enable_fp8 +from .fp8 import set_fp8_grad_input_enabled +from .planner import build_training_plan, resolve_margin_gib +from .resources import ArenaRuntimeResources + +RUNTIME_ATTR = "_arena_offload_runtime" + +GIB = 1024**3 +BOOTSTRAP_MARGIN_BYTES = GIB +BOOTSTRAP_MIN_STEP = 2 + + +class ArenaOffloadRuntime: + """Lifecycle + execution contexts for one arena-offloaded transformer.""" + + def __init__( + self, + model, + *, + device, + config, + arena, + residency, + executor, + training_plan, + smart_plan, + canonical_modules, + resources, + ) -> None: + self._model = model + self._device = device + self._config = config + self._arena = arena + self._residency = residency + self._executor = executor + self._training_plan = training_plan + self._smart_plan = smart_plan + self._canonical_modules = canonical_modules + self._resources = resources + self._closed = False + self._disposed = False + + # Set by training_step(); the residency controller (git-bug 0c577ef) + # reads these at the step boundary. + self._last_shape_key: tuple | None = None + self._last_step_num: int | None = None + self._successful_training_steps = 0 + self._signals = TrainingSignalWindow() + self._last_policy_error: str | None = None + self._last_failure_event: dict | None = None + self._policy = ArenaResidencyController() + self._last_training_cap_target_bytes: int | None = None + self._bootstrap_complete = False + self._bootstrap_min_free_bytes: int | None = None + self._bootstrap_budget_bytes = 0 + self._bootstrap_block_keys: tuple[str, ...] = () + self._training_fp8_restores = [] + self._training_fp8_canonical = 0 + self._training_fp8_singletons = 0 + self._sampling_fp8_canonical = 0 + self._sampling_fp8_singletons = 0 + self._permanent_placement = None + + # ------------------------------------------------------------------ + # construction + # ------------------------------------------------------------------ + + @classmethod + def _prepare( + cls, + transformer, + *, + device, + selection, + config, + ignore_modules: Sequence[Any] | None = None, + canonical_build=None, + ) -> ArenaOffloadRuntime: + existing = getattr(transformer, RUNTIME_ATTR, None) + if getattr(transformer, "_arena_offload_disposed", False): + raise RuntimeError("arena_offload_transformer_disposed") + if existing is not None: + if getattr(existing, "disposed", False): + raise RuntimeError("arena_offload_transformer_disposed") + raise RuntimeError("arena_offload_already_prepared") + + resources = getattr(canonical_build, "_arena_resources", None) + if resources is None: + resources = ArenaRuntimeResources(transformer, device) + resources.acquire_process_owner() + + try: + # Bind card simulation and allocator policy before any plan reads. + apply_simulated_card(config._simulated_vram_gib, device=device) + allocator_cap.apply_wddm_hard_allocator_cap( + device, config._policy.wddm_hard_gib, log_prefix="[ArenaOffload]" + ) + set_fp8_grad_input_enabled(config.fp8_backward) + + blocks = selection.blocks + block_keys = selection.block_keys + entries_by_block = { + key: list(selection.entries_by_block[key]) + for key in block_keys + } + if canonical_build is None: + arena = CanonicalArena() + canonical_build = arena.prepare(entries_by_block, model=transformer) + resources.adopt_canonical_build(canonical_build) + canonical_build.populate_from_model() + else: + resources.adopt_canonical_build(canonical_build) + if canonical_build.model is not transformer: + raise RuntimeError("arena_canonical_build_model_mismatch") + prepared_entries = { + key: tuple(module for _name, module in entries) + for key, entries in canonical_build.entries_by_block.items() + } + expected_entries = { + key: tuple(module for _name, module in entries) + for key, entries in entries_by_block.items() + } + if prepared_entries != expected_entries: + raise RuntimeError("arena_canonical_build_selection_mismatch") + arena = canonical_build.arena + + canonical_build.commit() + resources.mark_canonical_committed() + canonical_modules = tuple( + child for entries in entries_by_block.values() for _name, child in entries + ) + resources.canonical_modules = canonical_modules + + smart_plan = build_training_plan( + transformer, + arena, + canonical_modules, + device, + config, + block_keys=block_keys, + ) + residency = ResidencyState(arena, device) + resources.adopt_residency(residency) + training_plan = ResidencyPlan.from_smart_plan( + arena, smart_plan, phase="train" + ) + residency.reconcile(training_plan) + + policy = config._policy + executor_kwargs = dict( + depth=policy.prefetch_depth, + compile_blocks=config.compile_blocks, + compile_dynamic=config._compile_dynamic, + compile_dynamic_hints=config._compile_dynamic_hints, + protected_training_leaf_keys=smart_plan.get( + "protected_training_leaf_keys", () + ), + owner_token=resources.owner_token, + ) + from .dispatcher import prepare_block_dispatcher_runtime + + executor = prepare_block_dispatcher_runtime( + transformer, + residency, + selection=selection, + **executor_kwargs, + ) + resources.adopt_executor(executor) + + runtime = cls( + transformer, + device=device, + config=config, + arena=arena, + residency=residency, + executor=executor, + training_plan=training_plan, + smart_plan=smart_plan, + canonical_modules=canonical_modules, + resources=resources, + ) + resources.adopt_runtime(runtime) + resources.record_published_attribute( + transformer, RUNTIME_ATTR, runtime + ) + return runtime + except BaseException as error: + committed = resources.canonical_committed + try: + resources.release() + except ArenaCleanupError as cleanup_error: + try: + error.add_note(str(cleanup_error)) + except AttributeError: + pass + if committed: + raise ArenaSetupFatalError( + "arena setup failed after canonical commit" + ) from error + raise + + # ------------------------------------------------------------------ + # read-only state + # ------------------------------------------------------------------ + + @property + def model(self): + return self._model + + @property + def owns_compile(self) -> bool: + return True + + @property + def closed(self) -> bool: + return self._closed + + @property + def disposed(self) -> bool: + return self._disposed + + def place_permanent_modules(self, device, dtype=None) -> None: + """Move only noncanonical subtrees, preserving arena Parameter views.""" + self._require_open() + import torch + + target = (torch.device(device), dtype) + if getattr(self, "_permanent_placement", None) == target: + return + canonical = set(self._canonical_modules) + + def contains_canonical(module): + return any(child in canonical for child in module.modules()) + + def has_wrapped_parameter(module, *, recurse=True): + for parameter in module.parameters(recurse=recurse): + try: + names, _context = parameter.__tensor_flatten__() + except Exception: + continue + if names: + return True + return False + + def move_local_state(module): + local_dtype = ( + None + if dtype is None or has_wrapped_parameter(module, recurse=False) + else dtype + ) + + def convert(tensor): + if local_dtype is not None and ( + tensor.is_floating_point() or tensor.is_complex() + ): + return tensor.to(device=device, dtype=local_dtype) + return tensor.to(device=device) + + module._apply(convert, recurse=False) + + def move(module): + if module in canonical: + return + if not contains_canonical(module): + if dtype is None or has_wrapped_parameter(module): + module.to(device=device) + else: + module.to(device=device, dtype=dtype) + return + # A mixed parent can own direct permanent state (for example + # architecture pad tokens) in addition to canonical descendants. + # Move that local state before recursing into child subtrees. + move_local_state(module) + for child in module.children(): + move(child) + + try: + move(self._model) + self._permanent_placement = target + except BaseException as error: + self._fatal_setup_failure(error) + + @property + def device(self): + return self._device + + @property + def config(self): + return self._config + + @property + def block_count(self) -> int: + return len(self._arena.block_keys()) + + @property + def finalized(self) -> bool: + return bool(getattr(self._executor, "finalized", False)) + + # ------------------------------------------------------------------ + # lifecycle + # ------------------------------------------------------------------ + + def set_compile_dynamic_hints(self, hints) -> None: + """Install mark_dynamic hints on the block kernels (see ImmutableRuntime). + + The trainer derives sequence bounds from the datasets, which do not exist + when the runtime is prepared. Must be called before the first forward. + """ + self._require_open() + try: + self._executor.set_compile_dynamic_hints(hints) + self._config = replace( + self._config, + _compile_dynamic_hints=self._executor.compile_dynamic_hints, + ) + except BaseException as error: + self._fatal_setup_failure(error) + + def finalize(self, network=None): + """Build the permanent train/sample programs, then activate TRAIN. + + Must run AFTER the training network is applied so the dispatcher saves + the final installed block forwards. ``network`` is retained as a public + lifecycle argument, but adapter execution remains owned by those saved + model forwards rather than by arena-specific mappings. + """ + self._require_open() + try: + self._bind_training_cap() + del network + if self._config.fp8_forward: + canonical_ids = self._canonical_runtime_ids() + singleton_ids = self._singleton_runtime_ids() + self._training_fp8_restores = enable_fp8( + self._model, + include_ids=canonical_ids | singleton_ids, + live_ids=canonical_ids | singleton_ids, + training=True, + device=self._device, + ) + installed_ids = { + restore[4] for restore in self._training_fp8_restores + } + self._training_fp8_canonical = len(installed_ids & canonical_ids) + self._training_fp8_singletons = len(installed_ids & singleton_ids) + self._resources.record_fp8_restore( + "FP8 training restore", + lambda restores=self._training_fp8_restores: disable_fp8(restores), + ) + # Save block forwards only after the live-state FP8 transforms are + # installed, so compiled dispatcher kernels trace the selected + # execution policy rather than Quanto's materializing fallback. + self._executor.finalize_execution() + self._executor.activate(self._executor.TRAIN, self._training_plan) + return self + except BaseException as error: + self._fatal_setup_failure(error) + + def close(self) -> None: + """Release through the same owner used during preparation.""" + self._resources.release() + + def _fatal_setup_failure(self, error): + try: + self._resources.release() + except ArenaCleanupError as cleanup_error: + try: + error.add_note(str(cleanup_error)) + except AttributeError: + pass + raise ArenaSetupFatalError( + "arena setup failed after canonical commit" + ) from error + + def _require_open(self) -> None: + if self._closed: + if self._disposed: + raise RuntimeError("arena_offload_transformer_disposed") + raise RuntimeError("arena_offload_runtime_closed") + + # ------------------------------------------------------------------ + # execution contexts + # ------------------------------------------------------------------ + + @contextlib.contextmanager + def training_step(self, *, shape_key: tuple | None = None, step_num: int | None = None): + """The training phase boundary. Spans forward AND backward. + + Backward must be inside: checkpoint recomputation re-enters the block + runtime, so the source snapshot has to stay pinned for the whole step. + + This is the two-timescale residency controller's phase-boundary hook: + enter = plan/act, exit = observe. + """ + self._require_open() + self._last_shape_key = shape_key + self._last_step_num = step_num + try: + self._apply_training_policy() + except BaseException as error: + try: + self._handle_training_failure( + error, shape_key=shape_key, step_num=step_num + ) + except Exception as cleanup_error: + self._last_policy_error = ( + "training_policy_cleanup: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + raise + try: + import torch + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats(self._device) + except Exception: + pass + started_at = time.perf_counter() + succeeded = False + try: + with self._executor.execution(self._executor.TRAIN): + yield self + succeeded = True + except BaseException as error: + try: + self._handle_training_failure( + error, shape_key=shape_key, step_num=step_num + ) + except Exception as cleanup_error: + self._last_policy_error = ( + "training_failure_cleanup: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + raise + finally: + if succeeded: + try: + self._observe_training_step( + shape_key=shape_key, + step_num=step_num, + step_wall_ms=(time.perf_counter() - started_at) * 1000.0, + ) + self._last_policy_error = None + except Exception as error: + # Diagnostics must never mask a successful training step. + self._last_policy_error = f"{type(error).__name__}: {error}" + + def _handle_training_failure(self, error, *, shape_key, step_num): + """Clean arena-owned state after the executor has unwound.""" + import torch + + from . import transfer + from ..vram_budget import device_free_bytes + + abandoned = transfer.drain_fetch_runtime() + text = str(error) + allocation_failure = ( + isinstance(error, torch.cuda.OutOfMemoryError) + or "out of memory" in text.lower() + ) + rollback = None + if allocation_failure: + decision = self._policy.allocation_failure() + if decision.action == "rollback" and decision.block_key is not None: + rollback_keys = tuple(decision.block_keys or ()) + if rollback_keys: + self.transition_training_blocks( + rollback_keys, resident=False + ) + else: + self.transition_training_block( + decision.block_key, resident=False + ) + if decision.target_cap_bytes is not None: + allocator_cap.apply_wddm_hard_allocator_cap( + self._device, + self._config._policy.wddm_hard_gib, + target_cap_bytes=decision.target_cap_bytes, + log_prefix="[ArenaOffload]", + ) + self._last_training_cap_target_bytes = ( + decision.target_cap_bytes + ) + rollback = ( + list(decision.block_keys) + if decision.block_keys + else decision.block_key + ) + try: + stats = torch.cuda.memory_stats(self._device) + peak_allocated = int(torch.cuda.max_memory_allocated(self._device)) + peak_reserved = int(torch.cuda.max_memory_reserved(self._device)) + physical_free = int(device_free_bytes(self._device)) + except Exception: + stats = {} + peak_allocated = 0 + peak_reserved = 0 + physical_free = 0 + active_cap = allocator_cap.applied_cap_bytes(self._device) + policy = getattr(getattr(self, "_config", None), "_policy", None) + hard_gib = getattr(policy, "wddm_hard_gib", None) + hard_bytes = int(max(1.0, float(hard_gib or 1.0)) * GIB) + classification = "non_allocation_failure" + if allocation_failure: + classification = ( + "capped_allocator_rejection" + if active_cap is not None and physical_free > hard_bytes / 2 + else "cuda_allocation_failure_unknown" + ) + self._last_failure_event = { + "event": "training_allocation_failure", + "classification": classification, + "exception_type": type(error).__name__, + "exception": text, + "shape_key": shape_key, + "step_num": step_num, + "active_cap_bytes": active_cap, + "peak_allocated_bytes": peak_allocated, + "peak_reserved_bytes": peak_reserved, + "physical_free_bytes": physical_free, + "allocator_retry_delta": max( + 0, + int(stats.get("num_alloc_retries", 0) or 0) + - int( + ( + getattr(self._signals, "_allocator_previous", None) + or {} + ).get("num_alloc_retries", 0) + or 0 + ), + ), + "rollback_block": rollback, + "rejected_residency_bytes": ( + self._policy.last_rejected_residency_bytes + ), + "abandoned_fetch_tickets": int(abandoned), + } + + def _singleton_runtime_ids(self): + return set((self._smart_plan or {}).get("singleton_runtime_ids", ())) + + def _canonical_runtime_ids(self): + return {id(module) for module in self._canonical_modules} + + @contextlib.contextmanager + def sampling_session(self): + """Wrap a sampling run and restore TRAIN once at the end.""" + self._require_open() + sampling_restores = [] + if self._config.fp8_sampling: + canonical_ids = self._canonical_runtime_ids() + singleton_ids = self._singleton_runtime_ids() + sampling_restores = enable_fp8( + self._model, + include_ids=canonical_ids | singleton_ids, + live_ids=canonical_ids | singleton_ids, + training=False, + device=self._device, + ) + # Counts are eligibility diagnostics; training keeps its own + # persistent transforms underneath this temporary sampling layer. + installed_ids = {restore[4] for restore in sampling_restores} + self._sampling_fp8_canonical = len(installed_ids & canonical_ids) + self._sampling_fp8_singletons = len(installed_ids & singleton_ids) + try: + yield self + finally: + if sampling_restores: + disable_fp8(sampling_restores) + self._bind_training_cap() + self._executor.activate(self._executor.TRAIN, self._training_plan) + + @contextlib.contextmanager + def sampling_image(self, *, shape_key: tuple, cold_working_bytes: int): + """The sampling phase boundary for ONE image. + + Switches to the permanent SAMPLE program (forward-only, no + checkpointing) over the same arena. TRAIN is restored by the enclosing + `sampling_session()`. + """ + self._require_open() + policy = self._config._policy + + fixed_working_bytes = _fixed_working_bytes(policy.sampling_working_reserve_gib) + hard_gib = ( + 1.0 + if policy.sampling_wddm_hard_gib is None + else float(policy.sampling_wddm_hard_gib) + ) + allocator_cap.apply_wddm_hard_allocator_cap( + self._device, hard_gib, log_prefix="[ArenaOffload]" + ) + margin_gib = resolve_margin_gib( + self._device, + policy.sampling_wddm_margin_gib, + hard_gib=hard_gib, + ) + dequant_reserve = ( + 0 + if self._config.fp8_sampling + else int( + (self._smart_plan or {}).get( + "largest_singleton_bf16_dequant_bytes", 0 + ) + ) + ) + with self._executor.sampling( + shape_key=shape_key, + cold_working_bytes=int(cold_working_bytes), + fixed_working_bytes=fixed_working_bytes, + cold_floor_bytes=int(margin_gib * GIB) + dequant_reserve, + hot_floor_bytes=int((hard_gib + 0.25) * GIB) + dequant_reserve, + ): + yield self + + # ------------------------------------------------------------------ + def record_training_physical_free_min(self, free_bytes) -> None: + """Publish one successful step's physical high-water for bootstrap.""" + if ( + self._bootstrap_complete + or free_bytes is None + or self._successful_training_steps < BOOTSTRAP_MIN_STEP + ): + return + value = max(0, int(free_bytes)) + self._bootstrap_min_free_bytes = ( + value + if self._bootstrap_min_free_bytes is None + else min(self._bootstrap_min_free_bytes, value) + ) + + def _bootstrap_training_residency(self, active_cap_bytes) -> bool: + if ( + self._bootstrap_complete + or self._bootstrap_min_free_bytes is None + or int(self._successful_training_steps) < BOOTSTRAP_MIN_STEP + ): + return False + hard_gib = self._config._policy.wddm_hard_gib + hard_bytes = int( + (1.0 if hard_gib is None else max(1.0, float(hard_gib))) * GIB + ) + budget = max( + 0, + self._bootstrap_min_free_bytes + - hard_bytes + - BOOTSTRAP_MARGIN_BYTES, + ) + self._bootstrap_budget_bytes = budget + candidates = [] + protected = self._protected_training_blocks() + plan = getattr(self._residency, "plan", None) or self._training_plan + for order, block_key in enumerate(self._arena.block_keys()): + record = self._arena.block_record(block_key) + keys = tuple((block_key, name) for name in record.leaf_names) + if block_key in protected or any( + key in plan.resident_leaf_keys for key in keys + ): + continue + candidates.append( + (int(record.committed_bytes), order, str(block_key)) + ) + selected = [] + used = 0 + for block_bytes, _order, block_key in sorted(candidates): + if used + block_bytes > budget: + continue + selected.append(block_key) + used += block_bytes + if not selected: + self._bootstrap_complete = True + return False + resident_before = ( + int(self._residency.resident_bytes()) + + int((self._smart_plan or {}).get("singleton_resident_bytes", 0)) + ) + result = self.transition_training_blocks(selected, resident=True) + if not result.get("changed"): + return False + self._bootstrap_complete = True + self._bootstrap_block_keys = tuple(result["block_keys"]) + self._policy.begin_bootstrap_promotion( + self._bootstrap_block_keys, + used, + resident_before, + active_cap_bytes, + ) + return True + + def _protected_training_blocks(self): + return frozenset( + str(block) + for block, _leaf in (self._smart_plan or {}).get( + "protected_training_leaf_keys", () + ) + ) + + def _promotion_candidates(self): + plan = getattr(self._residency, "plan", None) or self._training_plan + protected = self._protected_training_blocks() + candidates = [] + for order, block_key in enumerate(self._arena.block_keys()): + record = self._arena.block_record(block_key) + keys = tuple((block_key, name) for name in record.leaf_names) + if block_key in protected or any( + key in plan.resident_leaf_keys for key in keys + ): + continue + candidates.append( + (int(record.committed_bytes), order, str(block_key)) + ) + return tuple( + {"block_key": block_key, "block_bytes": block_bytes} + for block_bytes, _order, block_key in sorted(candidates) + ) + + def _promotion_candidate(self): + candidates = self._promotion_candidates() + return candidates[0] if candidates else None + + def _aggressive_promotion_capacity(self, current_cap_bytes): + """Blocks that fit under both worst-shape safety budgets.""" + candidates = self._promotion_candidates() + allocator_slack = self._worst_shape_allocator_slack_bytes( + current_cap_bytes + ) + pad = int(self._policy.slack_pad_bytes) + used = 0 + capacity = 0 + for candidate in candidates: + used += int(candidate["block_bytes"]) + cumulative = {"block_bytes": used} + if self._worst_shape_candidate_margin_bytes(cumulative) < 0: + break + if allocator_slack <= used + pad: + break + capacity += 1 + return capacity + + def _demotion_candidate(self): + plan = getattr(self._residency, "plan", None) or self._training_plan + protected = self._protected_training_blocks() + candidates = [] + for order, block_key in enumerate(self._arena.block_keys()): + record = self._arena.block_record(block_key) + keys = tuple((block_key, name) for name in record.leaf_names) + if block_key in protected or not all( + key in plan.resident_leaf_keys for key in keys + ): + continue + actual = sum( + self._residency.resident_leaf_bytes(key) for key in keys + ) + candidates.append( + (actual or int(record.committed_bytes), -order, str(block_key)) + ) + if not candidates: + return None + block_bytes, _order, block_key = max(candidates) + return {"block_key": block_key, "block_bytes": block_bytes} + + def _worst_shape_candidate_margin_bytes(self, candidate): + if candidate is None: + return 0 + signal = self._signals.last_signal + peaks = self._signals.shape_peaks + if signal is None or not peaks: + return 0 + from .. import vram_budget + + total = int(vram_budget.device_total_bytes(self._device)) + worst_working = max( + int(peak.working_peak_bytes) + for peak in peaks.values() + if peak.steps > 0 + ) if any(peak.steps > 0 for peak in peaks.values()) else 0 + current_resident = ( + int(self._residency.resident_bytes()) + + int((self._smart_plan or {}).get("singleton_resident_bytes", 0)) + ) + current_ring = self._training_ring_bytes() + non_torch = max( + 0, + total + - int(signal.get("device_free_bytes", 0) or 0) + - int(signal.get("peak_reserved_bytes", 0) or 0), + ) + hard_gib = self._config._policy.wddm_hard_gib + hard_bytes = int( + (1.0 if hard_gib is None else max(1.0, float(hard_gib))) * GIB + ) + predicted_free = total - ( + worst_working + + current_resident + + current_ring + + non_torch + + int(candidate["block_bytes"]) + ) + return int(predicted_free - hard_bytes) + + def _worst_shape_allocator_slack_bytes(self, current_cap_bytes): + peaks = self._signals.shape_peaks + if not any(peak.steps > 0 for peak in peaks.values()): + return 0 + from .. import vram_budget + + worst_working = max( + int(peak.working_peak_bytes) + for peak in peaks.values() + if peak.steps > 0 + ) + current_resident = ( + int(self._residency.resident_bytes()) + + int((self._smart_plan or {}).get("singleton_resident_bytes", 0)) + ) + predicted_live = ( + worst_working + + current_resident + + self._training_ring_bytes() + ) + return vram_budget.allocator_allowance_bytes( + current_cap_bytes, predicted_live + ) + + def _apply_training_policy(self): + import torch + + if torch.device(self._device).type != "cuda" or not torch.cuda.is_available(): + return + candidate = self._promotion_candidate() + demote_candidate = self._demotion_candidate() + cliff_cap = allocator_cap.wddm_cliff_cap_bytes( + self._device, self._config._policy.wddm_hard_gib + ) + signal = self._signals.last_signal + current_cap = min( + cliff_cap, + int(self._last_training_cap_target_bytes or cliff_cap), + ) + if self._bootstrap_training_residency(current_cap): + return + aggressive_capacity = self._aggressive_promotion_capacity(current_cap) + decision = self._policy.step( + self._signals.last_signal, + candidate=candidate, + demote_candidate=demote_candidate, + cliff_cap_bytes=cliff_cap, + current_cap_bytes=current_cap, + worst_shape_free_bytes=self._worst_shape_candidate_margin_bytes( + candidate + ), + worst_shape_allocator_slack_bytes=( + self._worst_shape_allocator_slack_bytes(current_cap) + ), + aggressive_promotion_capacity=aggressive_capacity, + ) + if decision.action == "promote": + self.transition_training_block(decision.block_key, resident=True) + elif decision.action in ("demote", "rollback"): + rollback_keys = tuple(decision.block_keys or ()) + if rollback_keys: + self.transition_training_blocks( + rollback_keys, resident=False + ) + else: + self.transition_training_block( + decision.block_key, resident=False + ) + if ( + decision.action == "rollback" + and decision.target_cap_bytes is not None + ): + allocator_cap.apply_wddm_hard_allocator_cap( + self._device, + self._config._policy.wddm_hard_gib, + target_cap_bytes=decision.target_cap_bytes, + log_prefix="[ArenaOffload]", + ) + self._last_training_cap_target_bytes = ( + decision.target_cap_bytes + ) + elif decision.action == "raise_cap": + allocator_cap.apply_wddm_hard_allocator_cap( + self._device, + self._config._policy.wddm_hard_gib, + target_cap_bytes=decision.target_cap_bytes, + log_prefix="[ArenaOffload]", + ) + self._last_training_cap_target_bytes = decision.target_cap_bytes + + def transition_training_blocks(self, block_keys, *, resident: bool) -> dict: + """Apply one executor-owned multi-block transaction at a boundary.""" + self._require_open() + result = self._executor.transition_training_blocks( + tuple(block_keys), resident=bool(resident) + ) + if result.get("changed"): + self._training_plan = result["plan"] + return result + + def transition_training_block(self, block_key: str, *, resident: bool) -> dict: + """Apply one executor-owned whole-block transaction at a boundary.""" + self._require_open() + result = self._executor.transition_training_block( + str(block_key), resident=bool(resident) + ) + if result.get("changed"): + self._training_plan = result["plan"] + return result + + def _bind_training_cap(self) -> None: + allocator_cap.apply_wddm_hard_allocator_cap( + self._device, + self._config._policy.wddm_hard_gib, + log_prefix="[ArenaOffload]", + ) + + # diagnostics + # ------------------------------------------------------------------ + + def diagnostics(self) -> dict: + """One stable dict. Shared logging prints it; nobody reconstructs it.""" + active_plan = getattr(self._residency, "plan", None) or self._training_plan + canonical_resident = int(self._residency.resident_bytes()) + singleton_resident = int( + (self._smart_plan or {}).get("singleton_resident_bytes", 0) + ) + accounting = self._execution_accounting(active_plan) + selection = getattr(self._executor, "selection", None) + state_audit = getattr(selection, "accounting", None) + return { + "backend": "arena", + "blocks": self.block_count, + "finalized": self.finalized, + "resident_bytes": singleton_resident + canonical_resident, + "singleton_resident_bytes": singleton_resident, + "canonical_resident_bytes": canonical_resident, + "total_weight_resident_bytes": singleton_resident + canonical_resident, + "plan_fingerprint": getattr(active_plan, "fingerprint", None), + "checkpoint_owner": "model", + "accounting": accounting, + "state_audit": ( + None + if state_audit is None + else { + "managed_entries": int(state_audit.managed_entries), + "managed_bytes": int(state_audit.managed_bytes), + "trainable_entries": int(state_audit.trainable_entries), + "resident_entries": int(state_audit.resident_entries), + "resident_bytes": int(state_audit.resident_bytes), + } + ), + "prefetch_depth": int(getattr(self._executor, "depth", 0)), + "compile_blocks": bool(self._config.compile_blocks), + "compile_dynamic": bool(self._config._compile_dynamic), + "fp8_forward": bool(self._config.fp8_forward), + "fp8_backward": bool(self._config.fp8_backward), + "fp8_sampling": bool(self._config.fp8_sampling), + "training_fp8_canonical": getattr( + self, "_training_fp8_canonical", 0 + ), + "training_fp8_singletons": getattr( + self, "_training_fp8_singletons", 0 + ), + "sampling_fp8_canonical": getattr( + self, "_sampling_fp8_canonical", 0 + ), + "sampling_fp8_singletons": getattr( + self, "_sampling_fp8_singletons", 0 + ), + "largest_singleton_bf16_dequant_bytes": int( + (self._smart_plan or {}).get( + "largest_singleton_bf16_dequant_bytes", 0 + ) + ), + "training_cap_target_bytes": getattr( + self, "_last_training_cap_target_bytes", None + ), + "bootstrap_complete": self._bootstrap_complete, + "bootstrap_min_free_bytes": self._bootstrap_min_free_bytes, + "bootstrap_margin_bytes": BOOTSTRAP_MARGIN_BYTES, + "bootstrap_min_step": BOOTSTRAP_MIN_STEP, + "bootstrap_budget_bytes": self._bootstrap_budget_bytes, + "bootstrap_block_keys": self._bootstrap_block_keys, + "working_reserve_bytes": int( + (self._smart_plan or {}).get("working_reserve_bytes", 0) + ), + "all_resident_fit": bool( + (self._smart_plan or {}).get("all_resident_fit", False) + ), + "all_resident_working_reserve_bytes": int( + (self._smart_plan or {}).get( + "all_resident_working_reserve_bytes", 0 + ) + ), + "last_shape_key": self._last_shape_key, + "last_step_num": self._last_step_num, + "successful_training_steps": self._successful_training_steps, + "policy": { + **self._signals.diagnostics(), + "controller": self._policy.diagnostics(), + }, + "policy_error": self._last_policy_error, + "last_failure_event": self._last_failure_event, + } + + def _execution_accounting(self, plan) -> dict: + """Reconcile canonical payload, residency, and one execution's H2D plan.""" + from ..transfer_plan import build_transfer_plan + + resident_keys = frozenset(getattr(plan, "resident_leaf_keys", ())) + canonical_committed = 0 + canonical_payload = 0 + streamed_payload = 0 + planned_forward_bytes = 0 + planned_copies = 0 + resident_leaves = 0 + streamed_leaves = 0 + resident_blocks = 0 + streamed_blocks = 0 + partial_blocks = 0 + + for block_key in self._arena.block_keys(): + record = self._arena.block_record(block_key) + canonical_committed += int(record.committed_bytes) + all_leaves = tuple(record.leaf_names) + canonical_payload += sum( + int(tensor.nbytes) + for leaf in all_leaves + for tensor in record.leaf_spec(leaf).tensors + ) + streamed = tuple( + leaf + for leaf in all_leaves + if (block_key, leaf) not in resident_keys + ) + resident_count = len(all_leaves) - len(streamed) + resident_leaves += resident_count + streamed_leaves += len(streamed) + if not streamed: + resident_blocks += 1 + elif resident_count == 0: + streamed_blocks += 1 + else: + partial_blocks += 1 + if streamed: + transfer = build_transfer_plan(record, streamed) + streamed_payload += sum( + int(tensor.nbytes) + for leaf in streamed + for tensor in record.leaf_spec(leaf).tensors + ) + planned_forward_bytes += int(transfer.compact_nbytes) + planned_copies += int(transfer.num_ranges) + + canonical_resident_payload = int(self._residency.resident_bytes()) + protected = self._protected_training_blocks() + protected_resident = all( + all( + (block_key, leaf) in resident_keys + for leaf in self._arena.block_record(block_key).leaf_names + ) + for block_key in protected + ) + training_multiplier = 2 + return { + "phase": getattr(plan, "phase", None), + "canonical_committed_bytes": canonical_committed, + "canonical_payload_bytes": canonical_payload, + "canonical_padding_bytes": canonical_committed - canonical_payload, + "canonical_resident_payload_bytes": canonical_resident_payload, + "streamed_payload_bytes": streamed_payload, + "payload_reconciled": ( + canonical_payload + == canonical_resident_payload + streamed_payload + ), + "resident_blocks": resident_blocks, + "streamed_blocks": streamed_blocks, + "partially_resident_blocks": partial_blocks, + "resident_leaves": resident_leaves, + "streamed_leaves": streamed_leaves, + "mixed_residency": resident_leaves > 0 and streamed_leaves > 0, + "planned_forward_h2d_bytes": planned_forward_bytes, + "planned_forward_h2d_copies": planned_copies, + "planned_training_h2d_bytes": ( + planned_forward_bytes * training_multiplier + ), + "planned_training_h2d_copies": planned_copies * training_multiplier, + "protected_training_blocks": tuple(sorted(protected)), + "protected_training_blocks_resident": protected_resident, + } + + def _observe_training_step(self, *, shape_key, step_num, step_wall_ms) -> None: + """Collect the completed step's policy signals.""" + import torch + + from . import transfer + from ..vram_budget import device_free_bytes + + if torch.device(self._device).type != "cuda" or not torch.cuda.is_available(): + return + try: + stats = torch.cuda.memory_stats(self._device) + except Exception: + stats = {} + allocator = { + key: int(stats.get(key, 0) or 0) + for key in ("num_alloc_retries", "num_device_alloc", "num_device_free") + } + transfer_stats = ( + transfer.lifetime_fetch_stats() + if self._signals.transfer_snapshot_due + else None + ) + self._signals.observe( + shape_key=shape_key, + step_num=step_num, + allocator_counters=allocator, + peak_allocated_bytes=torch.cuda.max_memory_allocated(self._device), + peak_reserved_bytes=torch.cuda.max_memory_reserved(self._device), + device_free_bytes=device_free_bytes(self._device), + resident_bytes=( + self._residency.resident_bytes() + + int((self._smart_plan or {}).get("singleton_resident_bytes", 0)) + ), + ring_bytes=self._training_ring_bytes(), + compile_counters=_compile_counter_snapshot(torch), + transfer_counters=transfer_stats, + step_wall_ms=step_wall_ms, + ) + self._successful_training_steps += 1 + + def _training_ring_bytes(self) -> int: + from ..transfer_plan import build_transfer_plan + + largest = 0 + plan = getattr(self._residency, "plan", None) or self._training_plan + for block_key in self._arena.block_keys(): + record = self._arena.block_record(block_key) + streamed = tuple( + name + for name in record.leaf_names + if (block_key, name) not in plan.resident_leaf_keys + ) + if streamed: + transfer = build_transfer_plan(record, streamed) + largest = max(largest, transfer.compact_nbytes) + depth = max(1, int(getattr(self._executor, "depth", 1))) + return int(largest * depth) + + def report_foreign_vram_once(self, *, phase: str) -> None: + """Say so, once, if another tenant on the GPU is why we are streaming.""" + self._executor.report_foreign_vram_once( + self._residency.device, + phase=phase, + working_reserve_bytes=int( + (self._smart_plan or {}).get("working_reserve_bytes", 0) + ), + ) + +def _compile_counter_snapshot(torch_module): + try: + counters = torch_module._dynamo.utils.counters + except AttributeError: + return None + frames = int(counters["frames"].get("total", 0) or 0) + graphs = int(counters["stats"].get("unique_graphs", 0) or 0) + if frames == 0 and graphs == 0: + return None + return { + "frames": frames, + "graphs": graphs, + "graph_breaks": int(sum(counters["graph_break"].values())), + } + + +def _fixed_working_bytes(value) -> int | None: + """None when the sampling working reserve is auto (unset, negative, 'auto').""" + if value is None: + return None + try: + numeric = float(value) + except (TypeError, ValueError): + # Non-numeric (e.g. "auto") means auto-size, same as unset. + return None + if numeric < 0: + return None + return int(numeric * GIB) diff --git a/toolkit/memory_management/arena_offload/transfer.py b/toolkit/memory_management/arena_offload/transfer.py new file mode 100644 index 0000000000..eee81318cd --- /dev/null +++ b/toolkit/memory_management/arena_offload/transfer.py @@ -0,0 +1,816 @@ +"""Compact host-to-device transfer runtime for arena offload. + +Owns fetch plans at execution time, the reusable device ring, custom fetch +operators, statistics, and checkpoint/recompute ticket lifetime. +""" + +from __future__ import annotations + +import collections +import contextlib +import itertools +import threading +import time +from dataclasses import dataclass + +import torch +from toolkit.memory_management import pin_manager +from .ownership import validate_process_owner + +@dataclass +class _Ticket: + tid: int + device_buffer: torch.Tensor + ready_event: torch.cuda.Event | None + slot: int = -1 + # The exact key the slot was acquired under: torch.device("cuda") and + # buffer.device (cuda:0) do not compare equal as dict keys. + slot_device: torch.device | None = None + free_event: torch.cuda.Event | None = None + h2d_start: torch.cuda.Event | None = None + h2d_end: torch.cuda.Event | None = None + nbytes: int = 0 + copies: int = 0 + + +class _Slot: + """One reusable device buffer in the fetch ring. + + The buffer is allocated once and reused for every fetch that lands on this + slot. ``free_event`` is recorded on the COMPUTE stream by fetch_free, after + the last reader of the previous occupant; the transfer stream waits on it + device-side before overwriting the buffer. That ordering is what lets the + host submit ahead without ever blocking: the GPU enforces the recycle. + """ + + __slots__ = ("buffer", "free_event") + + def __init__(self) -> None: + self.buffer: torch.Tensor | None = None + self.free_event: torch.cuda.Event | None = None + + +_STATE_LOCK = threading.Lock() +_TICKETS: dict[int, _Ticket] = {} +_LIVE: collections.deque[int] = collections.deque() +_NEXT_ID = 0 +_DEPTH = 3 +_TRANSFER_STREAMS: dict[torch.device, torch.cuda.Stream] = {} +# Per-device ring of reusable device buffers, plus the indices currently +# available. A slot returns to _FREE_SLOTS when fetch_free SUBMITS (not when the +# GPU reaches it) -- the device-side wait on its free_event is what keeps the +# reuse correct, so the host never has to wait for compute to catch up. +_SLOTS: dict[torch.device, list[_Slot]] = {} +_FREE_SLOTS: dict[torch.device, collections.deque[int]] = {} +_STATS = { + "fetches": 0, + "bytes": 0, + "copies": 0, + "h2d_ms": 0.0, + "wait_ms": 0.0, + "depth_waits": 0, +} +_LIFETIME_STATS = dict(_STATS) +# (h2d_start, h2d_end) pairs awaiting timing. Drained only when the events have +# already completed, so accounting for a copy never blocks the host on it. +_PENDING_H2D: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] + +# Harness-only: restore the old behaviour of settling each copy's timing inside +# fetch_wait. Kept solely so a benchmark can A/B the cost of that host sync on +# the same build; production always drains lazily. +_BLOCKING_H2D_TIMING = False +_RUNTIME_OWNER_TOKEN = None + + +def set_h2d_timing_blocking(enabled: bool) -> None: + global _BLOCKING_H2D_TIMING + _BLOCKING_H2D_TIMING = bool(enabled) + + +def _drain_h2d(block: bool = False) -> None: + """Accumulate h2d_ms for finished copies. Host-blocking only if block=True. + + block=True is for the end-of-run report, where the in-flight copies are done + anyway; the hot path always calls this with block=False. + """ + if not _PENDING_H2D: + return + pending = [] + for h2d_start, h2d_end in _PENDING_H2D: + try: + if not block and not h2d_end.query(): + pending.append((h2d_start, h2d_end)) + continue + if block: + h2d_end.synchronize() + elapsed_ms = h2d_start.elapsed_time(h2d_end) + _STATS["h2d_ms"] += elapsed_ms + _LIFETIME_STATS["h2d_ms"] += elapsed_ms + except RuntimeError: + # Event never recorded (abandoned fetch, e.g. OOM unwind): drop it. + pass + _PENDING_H2D[:] = pending + + +def raise_dynamo_recompile_limit(min_limit: int = 128) -> None: + """Lift dynamo's per-code-object recompile cap for the in-graph trunks. + + Two legitimate recompile sources stack up on one code object: bucketed + training resolutions (dynamic=False -> one cache entry per distinct token + shape) and sampling-boundary rebuilds (fresh block-fn closures fail the + old entries' guards without evicting them). The default limit of 8 turned + that into FailOnRecompileLimitHit at the third boundary of a 200-step run + (~step 101). Each extra entry costs one ~3 min compile, not correctness; + the cap exists to flag accidental recompile storms, which the boundary + rebuild is not. + """ + config = torch._dynamo.config + for attribute in ("recompile_limit", "cache_size_limit"): + current = getattr(config, attribute, None) + if isinstance(current, int) and current < min_limit: + setattr(config, attribute, min_limit) + + +def configure_fetch_runtime(*, depth: int = 3, owner_token=None) -> None: + global _DEPTH, _NEXT_ID, _RUNTIME_OWNER_TOKEN + if owner_token is not None: + validate_process_owner(owner_token) + if torch.cuda.is_available() and _SLOTS: + # Slot buffers may still be in flight; settle before dropping them. + torch.cuda.synchronize() + _DEPTH = max(1, int(depth)) + with _STATE_LOCK: + _TICKETS.clear() + _LIVE.clear() + _PENDING_H2D.clear() + _SLOTS.clear() + _FREE_SLOTS.clear() + _NEXT_ID = 0 + _RUNTIME_OWNER_TOKEN = owner_token + + +def reset_fetch_stats() -> None: + _PENDING_H2D.clear() + for key in _STATS: + _STATS[key] = 0 + + +def fetch_stats(reset: bool = False) -> dict: + # Settle the copies still in flight so the reported h2d_ms covers every + # fetch, not just the ones that happened to finish before the last wait. + _drain_h2d(block=True) + stats = dict(_STATS) + if reset: + reset_fetch_stats() + return stats + + +def lifetime_fetch_stats() -> dict: + """Return monotonic fetch counters unaffected by report-window resets.""" + _drain_h2d(block=True) + return dict(_LIFETIME_STATS) + + +def fetch_performance_metrics(stats: dict, *, step_wall_ms=None) -> dict: + """Derive transfer-stream utilization from a settled reporting window. + + ``h2d_ms`` is CUDA-event time on the single serialized transfer stream. + Dividing its window total by the matching step-wall total estimates transfer + duty. ``wait_ms`` is deliberately excluded: it is host blocking around an + event wait and does not say whether the GPU compute stream was idle. + + H2D timing drains opportunistically, so callers should provide a multi-step + reporting window. Duty above 100% is retained and flagged rather than + clamped; it indicates accounting carried across a window boundary or a + mismatched denominator. + """ + h2d_ms = float((stats or {}).get("h2d_ms", 0.0) or 0.0) + byte_count = int((stats or {}).get("bytes", 0) or 0) + wall_ms = None if step_wall_ms is None else float(step_wall_ms) + duty_pct = None + if wall_ms is not None and wall_ms > 0.0: + duty_pct = 100.0 * h2d_ms / wall_ms + achieved_gbps = None + if h2d_ms > 0.0: + achieved_gbps = byte_count / (h2d_ms * 1_000_000.0) + return { + "step_wall_ms": wall_ms, + "h2d_duty_pct": duty_pct, + "h2d_duty_overflow": bool(duty_pct is not None and duty_pct > 100.0), + "achieved_gbps": achieved_gbps, + } + + +def fetch_report(reset: bool = False, *, step_wall_ms=None) -> str | None: + stats = fetch_stats(reset=reset) + if not stats["fetches"]: + return None + metrics = fetch_performance_metrics(stats, step_wall_ms=step_wall_ms) + gib = stats["bytes"] / 1024 ** 3 + duty = ( + "-" if metrics["h2d_duty_pct"] is None + else f"{metrics['h2d_duty_pct']:.1f}" + ) + gbps = ( + "-" if metrics["achieved_gbps"] is None + else f"{metrics['achieved_gbps']:.2f}" + ) + wall = ( + "-" if metrics["step_wall_ms"] is None + else f"{metrics['step_wall_ms']:.3f}" + ) + return ( + f"[InGraphStream] fetches={int(stats['fetches'])} " + f"copies={int(stats['copies'])} " + f"bytes={gib:.2f} GiB h2d_ms={stats['h2d_ms']:.3f} " + f"step_wall_ms={wall} h2d_duty_pct={duty} " + f"h2d_duty_overflow={int(metrics['h2d_duty_overflow'])} " + f"achieved_gbps={gbps} " + f"wait_ms={stats['wait_ms']:.3f} depth_waits={int(stats['depth_waits'])}" + ) + + +def _transfer_stream(device: torch.device): + stream = _TRANSFER_STREAMS.get(device) + if stream is None: + stream = torch.cuda.Stream(device=device) + _TRANSFER_STREAMS[device] = stream + return stream + + +def _ring_locked(device: torch.device) -> tuple[list[_Slot], collections.deque]: + slots = _SLOTS.get(device) + if slots is None or len(slots) != _DEPTH: + slots = [_Slot() for _ in range(_DEPTH)] + _SLOTS[device] = slots + _FREE_SLOTS[device] = collections.deque(range(_DEPTH)) + return slots, _FREE_SLOTS[device] + + +def _acquire_slot(device: torch.device, nbytes: int) -> tuple[int, _Slot]: + """Take a ring slot and make sure its buffer holds nbytes. + + Never blocks on GPU progress. An empty free list means the graph is holding + more than `depth` fetched buffers live at once -- the same condition the old + host-side reaper raised on, and still a bug rather than something to wait + out (waiting here would mean waiting on the compute stream, which only the + host can feed). + """ + with _STATE_LOCK: + slots, free = _ring_locked(device) + if not free: + raise RuntimeError( + "mm.fetch_start depth exceeded before fetch_free " + f"(ring depth {_DEPTH}); the graph holds more fetched buffers " + "live than the ring has slots" + ) + index = free.popleft() + slot = slots[index] + if slot.buffer is None or slot.buffer.numel() < nbytes: + # Growth happens during warmup, until every slot has seen the largest + # block. Settle the device first: the outgoing buffer may still be in + # flight, and dropping its last reference would hand the memory back to + # the caching allocator while a stream is still reading it. + if slot.buffer is not None: + torch.cuda.synchronize(device) + slot.buffer = torch.empty(nbytes, dtype=torch.uint8, device=device) + return index, slot + + +def _release_slot(device: torch.device, index: int) -> None: + with _STATE_LOCK: + _FREE_SLOTS.setdefault(device, collections.deque()).append(index) + + +def drain_fetch_runtime(*, owner_token=None) -> int: + """Abandon every outstanding fetch ticket (OOM-recovery path only). + + An OOM unwinds a forward between fetch_start and fetch_free, leaving + tickets whose free_event never records; the next fetch_start then blocks + on the depth limit and raises 'depth exceeded before fetch_free'. The + recovery path (mid-denoise demote / full streamed transition) calls this + AFTER the failed forward has fully unwound: nothing will consume the + in-flight device buffers anymore, so waiting out the transfer streams and + dropping the tickets is safe. Returns the number of tickets abandoned. + """ + token = _RUNTIME_OWNER_TOKEN if owner_token is None else owner_token + if token is not None: + validate_process_owner(token) + with _STATE_LOCK: + for stream in _TRANSFER_STREAMS.values(): + stream.synchronize() + abandoned = len(_LIVE) + _LIVE.clear() + _TICKETS.clear() + _PENDING_H2D.clear() + # The abandoned tickets never called fetch_free, so their slots were + # never returned. Rebuild the free list and clear the stale free_events + # (the recovery path has already unwound whatever would have read them). + for device, slots in _SLOTS.items(): + for slot in slots: + slot.free_event = None + _FREE_SLOTS[device] = collections.deque(range(len(slots))) + return abandoned + + +def release_fetch_runtime(owner_token) -> None: + """Release ring, streams, events, tickets, and reporting state for owner.""" + global _NEXT_ID, _RUNTIME_OWNER_TOKEN + validate_process_owner(owner_token) + drain_fetch_runtime(owner_token=owner_token) + with _STATE_LOCK: + _drain_h2d(block=True) + _TICKETS.clear() + _LIVE.clear() + _PENDING_H2D.clear() + _SLOTS.clear() + _FREE_SLOTS.clear() + _TRANSFER_STREAMS.clear() + _NEXT_ID = 0 + for key in _STATS: + _STATS[key] = 0 + _LIFETIME_STATS[key] = 0 + _RUNTIME_OWNER_TOKEN = None + + +def _fetch_start_impl(host_flat: torch.Tensor) -> torch.Tensor: + if _RUNTIME_OWNER_TOKEN is not None: + validate_process_owner(_RUNTIME_OWNER_TOKEN) + if host_flat.device.type != "cpu": + raise RuntimeError("mm.fetch_start expected a CPU host_flat tensor") + if torch.cuda.is_available() and not pin_manager.is_host_pinned(host_flat): + # is_host_pinned, not host_flat.is_pinned(): arena flats are pinned + # in place with cudaHostRegister, which torch's is_pinned() does not + # recognize (it only tracks its own caching-allocator pins). + raise RuntimeError("mm.fetch_start expected a pinned host_flat tensor") + device = torch.device("cuda") + stream = _transfer_stream(device) + nbytes = host_flat.numel() + index, slot = _acquire_slot(device, nbytes) + with _STATE_LOCK: + global _NEXT_ID + tid = _NEXT_ID + _NEXT_ID += 1 + _LIVE.append(tid) + h2d_start = torch.cuda.Event(enable_timing=True) + h2d_end = torch.cuda.Event(enable_timing=True) + ready = torch.cuda.Event() + device_buffer = slot.buffer[:nbytes] + with torch.cuda.stream(stream): + if slot.free_event is not None: + # Device-side recycle: the copy waits for the previous occupant's + # last reader, so the host does not have to. + stream.wait_event(slot.free_event) + h2d_start.record(stream) + device_buffer.copy_(host_flat, non_blocking=True) + ready.record(stream) + h2d_end.record(stream) + with _STATE_LOCK: + _TICKETS[tid] = _Ticket( + tid=tid, + device_buffer=device_buffer, + ready_event=ready, + slot=index, + slot_device=device, + h2d_start=h2d_start, + h2d_end=h2d_end, + nbytes=nbytes, + copies=1, + ) + _STATS["fetches"] += 1 + _STATS["bytes"] += int(host_flat.numel()) + _STATS["copies"] += 1 + _LIFETIME_STATS["fetches"] += 1 + _LIFETIME_STATS["bytes"] += int(host_flat.numel()) + _LIFETIME_STATS["copies"] += 1 + return torch.tensor([tid], dtype=torch.int64) + + +def _validated_transfer_ranges( + host_flat: torch.Tensor, + ranges: torch.Tensor, + compact_nbytes: int, +) -> list[tuple[int, int, int]]: + """Validate a static multi-range plan at the opaque runtime boundary. + + Ranges are tensor data rather than a closed-over Python plan so replacing + same-shaped canonical storage remains guard-stable. Destination spans must + form one dense compact flat; source spans must be ordered, non-overlapping + canonical bytes. + """ + if host_flat.device.type != "cpu" or host_flat.dtype != torch.uint8: + raise RuntimeError("mm.fetch_start_multi expected a CPU uint8 host_flat") + if not host_flat.is_contiguous(): + raise RuntimeError("mm.fetch_start_multi expected a contiguous host_flat") + if not pin_manager.is_arena_backed(host_flat): + raise RuntimeError( + "mm.fetch_start_multi expected a registered canonical arena source" + ) + if ranges.device.type != "cpu" or ranges.dtype != torch.int64: + raise RuntimeError("mm.fetch_start_multi expected CPU int64 ranges") + if ranges.ndim != 2 or ranges.shape[1] != 3 or ranges.shape[0] == 0: + raise RuntimeError("mm.fetch_start_multi expected non-empty Nx3 ranges") + compact_nbytes = int(compact_nbytes) + if compact_nbytes <= 0: + raise RuntimeError("mm.fetch_start_multi expected compact_nbytes > 0") + + rows = [tuple(int(value) for value in row) for row in ranges.tolist()] + previous_src_end = 0 + expected_dst = 0 + for index, (src_offset, dst_offset, nbytes) in enumerate(rows): + if src_offset < 0 or dst_offset < 0 or nbytes <= 0: + raise RuntimeError(f"mm.fetch_start_multi invalid range {index}") + if src_offset + nbytes > host_flat.numel(): + raise RuntimeError(f"mm.fetch_start_multi source range {index} out of bounds") + if index and src_offset < previous_src_end: + raise RuntimeError(f"mm.fetch_start_multi source range {index} overlaps") + if dst_offset != expected_dst: + raise RuntimeError( + f"mm.fetch_start_multi destination range {index} is not compact" + ) + previous_src_end = src_offset + nbytes + expected_dst = dst_offset + nbytes + if expected_dst != compact_nbytes: + raise RuntimeError("mm.fetch_start_multi ranges do not fill compact_nbytes") + return rows + + +def _fetch_start_multi_impl( + host_flat: torch.Tensor, + ranges: torch.Tensor, + compact_nbytes: int, +) -> torch.Tensor: + rows = _validated_transfer_ranges(host_flat, ranges, compact_nbytes) + device = torch.device("cuda") + stream = _transfer_stream(device) + compact_nbytes = int(compact_nbytes) + index, slot = _acquire_slot(device, compact_nbytes) + with _STATE_LOCK: + global _NEXT_ID + tid = _NEXT_ID + _NEXT_ID += 1 + _LIVE.append(tid) + + h2d_start = torch.cuda.Event(enable_timing=True) + h2d_end = torch.cuda.Event(enable_timing=True) + ready = torch.cuda.Event() + device_buffer = slot.buffer[:compact_nbytes] + with torch.cuda.stream(stream): + if slot.free_event is not None: + stream.wait_event(slot.free_event) + h2d_start.record(stream) + for src_offset, dst_offset, nbytes in rows: + device_buffer[dst_offset:dst_offset + nbytes].copy_( + host_flat[src_offset:src_offset + nbytes], non_blocking=True + ) + ready.record(stream) + h2d_end.record(stream) + + with _STATE_LOCK: + _TICKETS[tid] = _Ticket( + tid=tid, + device_buffer=device_buffer, + ready_event=ready, + slot=index, + slot_device=device, + h2d_start=h2d_start, + h2d_end=h2d_end, + nbytes=compact_nbytes, + copies=len(rows), + ) + _STATS["fetches"] += 1 + _STATS["bytes"] += compact_nbytes + _STATS["copies"] += len(rows) + _LIFETIME_STATS["fetches"] += 1 + _LIFETIME_STATS["bytes"] += compact_nbytes + _LIFETIME_STATS["copies"] += len(rows) + return torch.tensor([tid], dtype=torch.int64) + + +@torch.library.custom_op("mm::fetch_start", mutates_args=()) +def fetch_start(host_flat: torch.Tensor) -> torch.Tensor: + return _fetch_start_impl(host_flat) + + +@fetch_start.register_fake +def _(host_flat): + return torch.empty(1, dtype=torch.int64, device="cpu") + + +@torch.library.custom_op("mm::fetch_start_after", mutates_args=("guard",)) +def fetch_start_after(host_flat: torch.Tensor, guard: torch.Tensor) -> torch.Tensor: + return _fetch_start_impl(host_flat) + + +@fetch_start_after.register_fake +def _(host_flat, guard): + return torch.empty(1, dtype=torch.int64, device="cpu") + + +@torch.library.custom_op("mm::fetch_start_gated", mutates_args=()) +def fetch_start_gated(host_flat: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + """fetch_start with a REAL data dependency on `gate`. + + fetch_start_after's guard is a declared mutation; that bookkeeping does + not survive Inductor's scheduler/reinplacer, which may hoist backward + re-fetches above frees (ring overrun). Here the gate is an ordinary + input, so the dependency is genuine dataflow no scheduling stage can + drop. Emitted by the post-grad ordering pass (ingraph_stream_scheduling); + not intended for hand-written model code.""" + return _fetch_start_impl(host_flat) + + +@fetch_start_gated.register_fake +def _(host_flat, gate): + return torch.empty(1, dtype=torch.int64, device="cpu") + + +@torch.library.custom_op("mm::fetch_start_multi", mutates_args=()) +def fetch_start_multi( + host_flat: torch.Tensor, ranges: torch.Tensor, compact_nbytes: int +) -> torch.Tensor: + return _fetch_start_multi_impl(host_flat, ranges, compact_nbytes) + + +@fetch_start_multi.register_fake +def _(host_flat, ranges, compact_nbytes: int): + return torch.empty(1, dtype=torch.int64, device="cpu") + + +@torch.library.custom_op("mm::fetch_start_multi_after", mutates_args=("guard",)) +def fetch_start_multi_after( + host_flat: torch.Tensor, + ranges: torch.Tensor, + compact_nbytes: int, + guard: torch.Tensor, +) -> torch.Tensor: + return _fetch_start_multi_impl(host_flat, ranges, compact_nbytes) + + +@fetch_start_multi_after.register_fake +def _(host_flat, ranges, compact_nbytes: int, guard): + return torch.empty(1, dtype=torch.int64, device="cpu") + + +@torch.library.custom_op("mm::fetch_start_multi_gated", mutates_args=()) +def fetch_start_multi_gated( + host_flat: torch.Tensor, + ranges: torch.Tensor, + compact_nbytes: int, + gate: torch.Tensor, +) -> torch.Tensor: + return _fetch_start_multi_impl(host_flat, ranges, compact_nbytes) + + +@fetch_start_multi_gated.register_fake +def _(host_flat, ranges, compact_nbytes: int, gate): + return torch.empty(1, dtype=torch.int64, device="cpu") + + +@torch.library.custom_op("mm::fetch_wait", mutates_args=()) +def fetch_wait(token: torch.Tensor, nbytes: int) -> torch.Tensor: + tid = int(token[0].item()) + ticket = _TICKETS.get(tid) + if ticket is None: + raise RuntimeError(f"mm.fetch_wait got unknown ticket {tid}") + if int(nbytes) != ticket.nbytes: + raise RuntimeError( + f"mm.fetch_wait size mismatch for ticket {tid}: " + f"expected {ticket.nbytes}, got {int(nbytes)}" + ) + current = torch.cuda.current_stream() + start = time.perf_counter() + current.wait_event(ticket.ready_event) + elapsed_ms = (time.perf_counter() - start) * 1000.0 + _STATS["wait_ms"] += elapsed_ms + _LIFETIME_STATS["wait_ms"] += elapsed_ms + if ticket.h2d_start is not None and ticket.h2d_end is not None: + # Queue the pair for opportunistic draining; do NOT synchronize here. + # Blocking on h2d_end just to service a counter stalls the submit loop, + # and with the ring recycling device-side there is nothing to gain from + # the throttle it used to provide. + _PENDING_H2D.append((ticket.h2d_start, ticket.h2d_end)) + _drain_h2d(block=_BLOCKING_H2D_TIMING) + return ticket.device_buffer + + +@fetch_wait.register_fake +def _(token, nbytes: int): + return torch.empty(nbytes, dtype=torch.uint8, device="cuda") + + +def _fetch_free_impl(token: torch.Tensor) -> torch.Tensor: + tid = int(token[0].item()) + ticket = _TICKETS.get(tid) + if ticket is None: + raise RuntimeError(f"mm.fetch_free got unknown ticket {tid}") + free_event = torch.cuda.Event() + free_event.record(torch.cuda.current_stream()) + ticket.free_event = free_event + if ticket.slot >= 0 and ticket.slot_device is not None: + # The slot's next occupant waits on this event device-side before it + # overwrites the buffer, so the slot can go back into circulation now, + # at SUBMIT time, without the host waiting for the GPU to reach it. + _SLOTS[ticket.slot_device][ticket.slot].free_event = free_event + _release_slot(ticket.slot_device, ticket.slot) + with _STATE_LOCK: + _TICKETS.pop(tid, None) + if _LIVE and _LIVE[0] == tid: + _LIVE.popleft() + elif tid in _LIVE: + _LIVE.remove(tid) + return token.clone() + + +@torch.library.custom_op("mm::fetch_free", mutates_args=()) +def fetch_free(token: torch.Tensor) -> torch.Tensor: + return _fetch_free_impl(token) + + +@fetch_free.register_fake +def _(token): + return token.clone() + + +@torch.library.custom_op("mm::fetch_free_after", mutates_args=("guard",)) +def fetch_free_after(token: torch.Tensor, guard: torch.Tensor) -> torch.Tensor: + return _fetch_free_impl(token) + + +@fetch_free_after.register_fake +def _(token, guard): + return token.clone() + + +def _register_ordered_effects(): + """Pin the fetch ops to program order inside compiled graphs. + + Functionalized custom ops only carry data deps through their args; in the + AOT backward graph each checkpoint unit's re-fetch depends only on the + saved boundary activation (an immediately-available input), so Inductor + may hoist all re-fetches above the frees -- exceeding the ring depth and + deadlocking the host-side depth guard. Ordered effect tokens thread a + dependency chain through every fetch op, enforcing eager program order in + forward AND backward graphs (kernel-launch order only; stream overlap is + unaffected).""" + try: + from torch._higher_order_ops.effects import ( + _EffectType, + _register_effectful_op, + ) + + for op in ( + torch.ops.mm.fetch_start.default, + torch.ops.mm.fetch_start_after.default, + torch.ops.mm.fetch_start_multi.default, + torch.ops.mm.fetch_start_multi_after.default, + torch.ops.mm.fetch_wait.default, + torch.ops.mm.fetch_free.default, + torch.ops.mm.fetch_free_after.default, + ): + _register_effectful_op(op, _EffectType.ORDERED) + except Exception as error: # pragma: no cover - torch-version dependent + raise RuntimeError( + "in-graph streaming requires ordered-effect registration for its " + f"fetch ops (torch internal API changed?): {error!r}" + ) from error + + +# NOT registered at import time: in torch 2.12 ordered-effect tokens trip an +# internal token-erasure assertion inside the checkpoint HOP lowering +# (see tests/test_ingraph_training_ops.py, compiled xfail). Phase 4a S1 keeps +# this as the candidate ordering mechanism for the compiled trunk; call it +# explicitly once the HOP interaction is resolved (torch upgrade or flat-trunk +# design without the checkpoint HOP). + + +class _FreeOnBackwardFn(torch.autograd.Function): + """Anchor a ticket's free event to the consuming block's BACKWARD. + + Training-mode counterpart of `fetch_free_after`: under checkpoint + recompute, the block's backward (grad-input from the fetched weight + views) is the true last reader of the ticket's device buffer, so a + forward-side free lets the depth-K ring recycle the buffer under + backward kernels that are still reading it (silent corruption). + + Wrap the block INPUT, not its output: this node's backward runs last + in the block's backward (input side), i.e. after every weight-view + read, and `fetch_free_after`'s declared guard mutation on the incoming + grad keeps the free ordered after the kernels that produced it. + """ + + @staticmethod + def forward(ctx, x, token): + ctx.save_for_backward(token) + return x + + @staticmethod + def backward(ctx, grad_x): + (token,) = ctx.saved_tensors + if torch.compiler.is_compiling(): + # Declared guard mutation orders the free after the kernels that + # produced grad_x; functionalization makes it version-safe. + torch.ops.mm.fetch_free_after(token, grad_x) + else: + # Eager executes in program order -- and the guarded variant's + # version bump on grad_x would trip autograd's version checks. + torch.ops.mm.fetch_free(token) + return grad_x, None + + +def free_on_backward(x: torch.Tensor, token: torch.Tensor) -> torch.Tensor: + """Defer a ticket's free to the consuming block's backward (training). + + Two fetch generations exist under non-reentrant checkpoint: the + first-pass fetch (its views are dropped by the checkpoint hooks, so it + is safe to free after the block's forward) and the recompute fetch + (its views feed the real backward, so it must be freed after the + block's backward). The token passed here is saved via + ``save_for_backward`` -- checkpoint's saved-tensor machinery therefore + swaps it for the RECOMPUTE generation's token automatically, and this + node's backward frees exactly the ticket backward actually read. + + Canonical training block shape (see checkpoint_recompute_context): + + token = fetch_start_after(host, x) + flat = fetch_wait(token, nbytes) + ...views... + x = free_on_backward(x, token) + out = (x, views) + if not in_recompute(): + torch.ops.mm.fetch_free_after(token, out) # first-pass gen only + return out + """ + return _FreeOnBackwardFn.apply(x, token) + + +_IN_RECOMPUTE = threading.local() + + +def in_recompute() -> bool: + """True while a checkpoint recompute pass (via checkpoint_recompute_context) + is re-running the block fn.""" + return bool(getattr(_IN_RECOMPUTE, "value", False)) + + +class _RecomputeMarker: + def __enter__(self): + self._prev = getattr(_IN_RECOMPUTE, "value", False) + _IN_RECOMPUTE.value = True + return self + + def __exit__(self, exc_type, exc, tb): + _IN_RECOMPUTE.value = self._prev + return False + + +def checkpoint_recompute_context(): + """``context_fn`` for torch.utils.checkpoint: null forward context, and a + recompute context that flips in_recompute() so the block fn suppresses the + first-pass forward free during recompute (the recompute ticket is freed by + free_on_backward instead). EAGER ONLY -- compiled checkpoint requires + TorchDispatchMode contexts; use compiled_checkpoint_context there.""" + return contextlib.nullcontext(), _RecomputeMarker() + + +def _compiled_free_policy(ctx, op, *args, **kwargs): + from torch.utils.checkpoint import CheckpointPolicy + + if op in ( + torch.ops.mm.fetch_free_after.default, + torch.ops.mm.fetch_free.default, + ): + # Keep the forward-side free OUT of the backward replay: replayed, it + # would free the backward re-fetch's buffer before the grad kernels + # read it. free_on_backward's op is the backward-side free. + return CheckpointPolicy.MUST_SAVE + if op in ( + torch.ops.mm.fetch_start.default, + torch.ops.mm.fetch_start_after.default, + torch.ops.mm.fetch_start_multi.default, + torch.ops.mm.fetch_start_multi_after.default, + torch.ops.mm.fetch_wait.default, + ): + # The design's core invariant: fetched weights are NEVER saved for + # backward. PREFER_RECOMPUTE is advisory -- at Krea2 scale the + # partitioner chose to save all 28 fetched flats (12.25 GiB -> OOM). + return CheckpointPolicy.MUST_RECOMPUTE + # Everything else replays in backward (full-checkpoint mode). + return CheckpointPolicy.PREFER_RECOMPUTE + + +def compiled_checkpoint_context(): + """``context_fn`` for torch.utils.checkpoint under torch.compile.""" + from torch.utils.checkpoint import create_selective_checkpoint_contexts + + return create_selective_checkpoint_contexts(_compiled_free_policy) + + +# NOTE: there is deliberately NO helper that "picks the right checkpoint +# context automatically". The checkpoint HOP calls context_fn() OUTSIDE the +# compiling frame, so an is_compiling() check inside such a helper always +# reads False under compile and hands the HOP eager (non-TorchDispatchMode) +# contexts, failing its assertion. Select the context at trunk level instead: +# context_fn = (compiled_checkpoint_context +# if torch.compiler.is_compiling() +# else checkpoint_recompute_context) diff --git a/toolkit/memory_management/canonical_arena.py b/toolkit/memory_management/canonical_arena.py new file mode 100644 index 0000000000..43955924ff --- /dev/null +++ b/toolkit/memory_management/canonical_arena.py @@ -0,0 +1,276 @@ +"""Canonical host arena (Slice 1, tasks/open/IMMUTABLE_TRANSFER_ARENA_PLAN.md). + +One page-exclusive pinned host flat per block, immutable leaf metadata, and +a ONE-TIME repoint of frozen base Parameters into views over those flats +(Invariant 4). This is deliberately NOT the legacy ``pinned_arena.py``: +no generation counter, no ``is_current``/staleness oracle over live module +storage, no invalidate/restore/rebuild path, no borrowed-vs-owned pack +taxonomy. Per the plan's Decision section, promotion/demotion/sampling +transitions must never repoint a Parameter again once ``canonicalize()`` +has run -- that is the job of the residency sidecars (Slice 3), not this +module. + +Construction is destination-first and transactional: preparation allocates +the final host flats without mutating Parameters, population fills those flats, +and commit publishes all Parameter views atomically or restores the originals. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from toolkit.memory_management import pin_manager +from toolkit.memory_management.arena_offload.layout import BlockPack, LinearSpec, release_pack + + +ARENA_KIND = "weights" + + +class CanonicalArenaError(ValueError): + """A build/canonicalize/guard operation violated a canonical-arena + invariant (tasks/open/IMMUTABLE_TRANSFER_ARENA_PLAN.md).""" + + +def _entry_module_and_leaves(entry): + if len(entry) != 2 or entry[1] is None: + raise CanonicalArenaError( + "canonical_arena_requires_module_entries: blocks must be built " + "from (name, module) entries so Parameters can be repointed" + ) + name, module = entry + weight = module.weight + bias = getattr(module, "bias", None) + return name, module, weight, bias + + +def _assert_entries_frozen(entries) -> None: + """Invariant 4/Decision: the canonical arena is for FROZEN base weights + only; LoRA/adapters stay ordinary trainable Parameters outside it. Run + this over every block BEFORE any block is built, so a trainable leaf + anywhere in the batch fails closed without repointing a single + Parameter (no partial canonicalization from this particular cause).""" + for entry in entries: + name, _module, weight, bias = _entry_module_and_leaves(entry) + if getattr(weight, "requires_grad", False): + raise CanonicalArenaError(f"canonical_arena_trainable_leaf:{name}:weight") + if bias is not None and getattr(bias, "requires_grad", False): + raise CanonicalArenaError(f"canonical_arena_trainable_leaf:{name}:bias") + + +@dataclass(frozen=True) +class BlockRecord: + """One block's immutable canonical host representation. + + No residency state, no generation counter, no live-module currentness + test -- packs ARE views over arena records (plan Target Components #1). + """ + + block_key: str + pack: BlockPack + leaf_names: tuple[str, ...] + modules: tuple[torch.nn.Module, ...] + + def module_for_leaf(self, leaf_name: str) -> torch.nn.Module: + try: + return self.modules[self.leaf_names.index(leaf_name)] + except ValueError as error: + raise KeyError( + f"no leaf named {leaf_name!r} in block {self.block_key!r}" + ) from error + + @property + def host_flat(self) -> torch.Tensor: + return self.pack.host_flat + + @property + def committed_bytes(self) -> int: + return self.pack.required_pin_bytes + + def leaf_spec(self, leaf_name: str) -> LinearSpec: + for spec in self.pack.linears: + if spec.name == leaf_name: + return spec + raise KeyError(f"no leaf named {leaf_name!r} in block {self.block_key!r}") + + +@dataclass +class CanonicalArenaStats: + blocks: int = 0 + pinned_bytes: int = 0 + + +class CanonicalArena: + """Owns one process's canonical, page-exclusive pinned host blocks. + + Not thread-safe; ``canonicalize()``/``release()`` are expected to run on + the same thread that drives model attach/detach, exactly once each per + instance -- construct a new ``CanonicalArena`` rather than rebuilding. + """ + + def __init__(self) -> None: + self._blocks: dict[str, BlockRecord] = {} + self._canonicalized = False + + @property + def canonicalized(self) -> bool: + return self._canonicalized + + # -- canonicalize (Invariant 3 + 4) ------------------------------------ + + def canonicalize( + self, entries_by_block: dict, *, kind: str = ARENA_KIND + ) -> CanonicalArenaStats: + """Build every block's canonical flat and repoint its Parameters, + exactly once. ``entries_by_block`` maps ``block_key -> iterable of + (name, module)`` entries, all frozen (``requires_grad=False``). + + Caller sequencing responsibility (this method cannot see it): run + AFTER load/quantize/freeze and BEFORE LoRA attach, optimizer + construction, or compile (Invariant 4) -- once Parameters are + repointed here, nothing may replace them again for the life of the + arena. + + Fails loud and releases any blocks already built in this call if + any block cannot be admitted (trainable leaf, unsupported quant + wrapper, or pin budget exceeded) -- there is no silent pageable + fallback in this arena (that is an admission-policy decision for + the caller, Invariant 10, not a mechanism this class provides). + """ + if self._canonicalized: + raise CanonicalArenaError( + "canonical_arena_double_canonicalize: canonicalize() may " + "only run once per arena instance (Invariant 4) -- runtime " + "promotion/demotion/sampling transitions must never " + "repoint a Parameter again" + ) + # Validate every block's leaves are frozen BEFORE building any of + # them, so a trainable leaf anywhere fails closed with zero + # Parameters repointed (see _assert_entries_frozen). + normalized: dict[str, list] = {} + for block_key, raw_entries in entries_by_block.items(): + entries = list(raw_entries) + _assert_entries_frozen(entries) + normalized[block_key] = entries + + from toolkit.memory_management.arena_offload.construction import PreparedCanonicalBuild + + build = PreparedCanonicalBuild(self, normalized, kind=kind) + build.populate_from_model() + return build.commit() + + def prepare(self, entries_by_block: dict, *, model=None, kind: str = ARENA_KIND): + """Prepare final destinations without mutating model Parameters.""" + from toolkit.memory_management.arena_offload.construction import PreparedCanonicalBuild + + normalized = {key: list(entries) for key, entries in entries_by_block.items()} + for entries in normalized.values(): + _assert_entries_frozen(entries) + return PreparedCanonicalBuild(self, normalized, model=model, kind=kind) + + # -- whole-model .to() interception (Invariant 5) ---------------------- + + @staticmethod + def guard_whole_model_to(model: torch.nn.Module) -> None: + """Forbid whole-model ``.to()``/``.cuda()``/``.cpu()`` on a model + with canonicalized leaves: a model-wide move silently detaches + every Parameter from its arena flat (copy semantics on the full + move), the exact drift the legacy arena's ``restore_view`` existed + to repair after the fact. Idempotent. Callers that need to move + SOME parameters (LoRA, non-canonicalized submodules) must route + through a canonical-arena-aware helper instead of raw ``.to()``.""" + if getattr(model, "_mm_canonical_to_guarded", False): + return + original_to = model.to + + def _guarded_to(*args, **kwargs): + runtime = getattr(model, "_arena_offload_runtime", None) + placement = getattr(runtime, "_permanent_placement", None) + if placement is not None: + device, dtype, _non_blocking, _memory_format = ( + torch._C._nn._parse_to(*args, **kwargs) + ) + placed_device, placed_dtype = placement + same_device = ( + device is not None + and torch.device(device) == torch.device(placed_device) + ) + same_dtype = dtype is None or dtype == placed_dtype + if same_device and same_dtype: + return model + raise CanonicalArenaError( + "canonical_arena_whole_model_to: whole-model .to()/.cuda()/" + ".cpu() is forbidden once canonicalized leaves exist -- it " + "would silently detach every Parameter from its arena flat. " + "Route non-canonicalized regions through their own .to() " + "calls, or move canonicalized weights via the residency " + "sidecar path instead." + ) + + model.to = _guarded_to + model._mm_canonical_to_guarded = True + model._mm_canonical_to_original = original_to + + @staticmethod + def unguard_whole_model_to(model: torch.nn.Module) -> None: + original = getattr(model, "_mm_canonical_to_original", None) + if original is not None: + model.to = original + del model._mm_canonical_to_original + if hasattr(model, "_mm_canonical_to_guarded"): + del model._mm_canonical_to_guarded + + # -- introspection ------------------------------------------------------ + + def block_record(self, block_key: str) -> BlockRecord | None: + return self._blocks.get(block_key) + + def block_pack(self, block_key: str) -> BlockPack | None: + record = self._blocks.get(block_key) + return None if record is None else record.pack + + def block_keys(self) -> tuple[str, ...]: + return tuple(self._blocks.keys()) + + def committed_pinned_bytes(self) -> int: + return sum(record.committed_bytes for record in self._blocks.values()) + + def stats(self) -> CanonicalArenaStats: + return CanonicalArenaStats( + blocks=len(self._blocks), pinned_bytes=self.committed_pinned_bytes() + ) + + def immutable_signature(self) -> tuple: + """Return this arena's immutable host-storage commitment. + + Process-wide pin-ledger state is deliberately excluded: unrelated + consumers such as the bounce pool may grow or shrink while residency + sidecars change without mutating canonical host storage. + """ + return ( + self._canonicalized, + tuple( + ( + block_key, + record.host_flat.data_ptr(), + record.committed_bytes, + pin_manager.is_host_pinned(record.host_flat), + pin_manager.is_arena_backed(record.host_flat), + ) + for block_key, record in self._blocks.items() + ), + ) + + # -- explicit unload ------------------------------------------------ + + def release(self) -> None: + """Release every block's pin registration and every committed + byte (Invariant 1: pin_manager is the sole authority, every + registered byte is released explicitly). Safe to call on a + partially-built or already-released arena.""" + for record in self._blocks.values(): + pin_manager.unregister_arena_storage(record.pack.host_flat) + release_pack(record.pack) + self._blocks.clear() + self._canonicalized = False diff --git a/toolkit/memory_management/immutable_runtime.py b/toolkit/memory_management/immutable_runtime.py new file mode 100644 index 0000000000..d7e4e3138b --- /dev/null +++ b/toolkit/memory_management/immutable_runtime.py @@ -0,0 +1,880 @@ +"""Generic compile-neutral source state for immutable transformer runtimes. + +The source table owns only publication and execution exclusion. Canonical +storage remains owned by the arena, while ``ResidencyState`` continues to own +resident device sidecars. +""" + +from __future__ import annotations + +import hashlib +from contextlib import contextmanager +from dataclasses import dataclass + +import torch + +from toolkit.memory_management import vram_budget +from toolkit.memory_management.residency import ( + ResidencyDelta, + ResidencyPlan, + ResidencyState, +) +from toolkit.memory_management.transfer_plan import ( + BlockTransferPlan, + build_transfer_plan, +) + + +class ImmutableRuntimeError(RuntimeError): + pass + + +@dataclass(frozen=True) +class ImmutableBlockABI: + """Structural block information that remains stable for runtime life.""" + + block_key: str + leaf_names: tuple[str, ...] + leaf_layout: tuple + + +@dataclass(frozen=True) +class ImmutableBlockSourceSnapshot: + """Current source-selection metadata for one canonical block.""" + + block_key: str + leaf_names: tuple[str, ...] + resident_leaf_names: frozenset[str] + transfer: BlockTransferPlan | None + ranges: torch.Tensor | None + + def assemble_leaf_args( + self, + residency: ResidencyState, + compact_flat: torch.Tensor | None, + ) -> tuple: + record = residency.arena.block_record(self.block_key) + if record is None: + raise ImmutableRuntimeError(f"missing_canonical_block:{self.block_key}") + + args = [] + for leaf_name in self.leaf_names: + spec = record.leaf_spec(leaf_name) + if leaf_name in self.resident_leaf_names: + sidecar = residency.resident_leaf((self.block_key, leaf_name)) + if sidecar is None: + raise ImmutableRuntimeError(f"missing_resident_source:{self.block_key}.{leaf_name}") + args.append(sidecar.tensors) + continue + + if compact_flat is None or self.transfer is None: + raise ImmutableRuntimeError(f"missing_streamed_source:{self.block_key}.{leaf_name}") + + args.append( + tuple( + self.transfer.compact_leaf_view( + compact_flat, + leaf_name, + item.role, + ) + for item in spec.tensors + ) + ) + + return tuple(args) + + +def build_source_snapshot( + residency: ResidencyState, + plan: ResidencyPlan, + abi: ImmutableBlockABI, +) -> ImmutableBlockSourceSnapshot: + record = residency.arena.block_record(abi.block_key) + if record is None: + raise ImmutableRuntimeError(f"missing_canonical_block:{abi.block_key}") + + resident = plan.resident_in_block(abi.block_key) + unknown = resident - frozenset(abi.leaf_names) + if unknown: + leaf_name = sorted(unknown)[0] + raise ImmutableRuntimeError(f"unknown_residency_leaf:{abi.block_key}.{leaf_name}") + + streamed = tuple(leaf_name for leaf_name in abi.leaf_names if leaf_name not in resident) + transfer = build_transfer_plan(record, streamed) if streamed else None + ranges = None if transfer is None else transfer.ranges_tensor() + return ImmutableBlockSourceSnapshot( + block_key=abi.block_key, + leaf_names=abi.leaf_names, + resident_leaf_names=resident, + transfer=transfer, + ranges=ranges, + ) + + +@dataclass(frozen=True) +class ImmutableProgram: + """One permanent eager program for one execution mode.""" + + mode: str + fingerprint: str + trunk: object + + +def build_program_fingerprint( + mode: str, + block_abis, + *, + architecture_key: str, + depth: int, + checkpoint_mode: str, + adapter_shape=(), + has_multiplier: bool = False, +) -> str: + per_block = tuple( + ( + abi.block_key, + tuple(abi.leaf_names), + abi.leaf_layout, + ) + for abi in block_abis + ) + source = repr( + ( + "immutable-runtime-v1", + str(architecture_key), + str(mode), + per_block, + int(depth), + str(checkpoint_mode), + tuple(adapter_shape), + bool(has_multiplier), + ) + ) + return hashlib.sha1(source.encode("utf-8")).hexdigest()[:16] + + +class ImmutableRuntimeSourceTable: + """Atomically published per-block source snapshots.""" + + def __init__(self, residency: ResidencyState, block_abis) -> None: + self.residency = residency + self.block_abis = tuple(block_abis) + self._generation = 0 + self._active_executions = 0 + self._snapshots: tuple[ImmutableBlockSourceSnapshot, ...] | None = None + self._plan: ResidencyPlan | None = None + + if residency.plan.phase != "empty": + self._snapshots = self._build_snapshots(residency.plan) + self._plan = residency.plan + self._generation = 1 + + @property + def generation(self) -> int: + return self._generation + + @property + def plan(self) -> ResidencyPlan | None: + return self._plan + + @property + def active_executions(self) -> int: + return self._active_executions + + def _build_snapshots( + self, + plan: ResidencyPlan, + ) -> tuple[ImmutableBlockSourceSnapshot, ...]: + return tuple(build_source_snapshot(self.residency, plan, abi) for abi in self.block_abis) + + def begin_execution(self) -> int: + if self._snapshots is None: + raise ImmutableRuntimeError("no_residency_source_table") + self._active_executions += 1 + return self._generation + + def end_execution(self, generation: int) -> None: + if self._active_executions <= 0: + raise ImmutableRuntimeError("immutable_execution_not_active") + if int(generation) != self._generation: + raise ImmutableRuntimeError("immutable_execution_generation_mismatch") + self._active_executions -= 1 + + def source(self, block_index: int) -> ImmutableBlockSourceSnapshot: + snapshots = self._snapshots + if snapshots is None: + raise ImmutableRuntimeError("no_residency_source_table") + try: + return snapshots[int(block_index)] + except IndexError as error: + raise ImmutableRuntimeError(f"unknown_execution_block:{block_index}") from error + + def publish(self, plan: ResidencyPlan) -> ResidencyDelta: + if self._active_executions: + raise ImmutableRuntimeError("residency_transition_during_execution") + + snapshots = self._build_snapshots(plan) + if self.residency.plan.fingerprint == plan.fingerprint: + delta = ResidencyDelta((), (), self.residency.resident_bytes()) + else: + delta = self.residency.reconcile(plan) + + self._snapshots = snapshots + self._plan = plan + self._generation += 1 + return delta + + def clear(self) -> None: + if self._active_executions: + raise ImmutableRuntimeError("source_table_clear_during_execution") + self._snapshots = None + self._plan = None + + +def _leaf_layout(record) -> tuple: + layout = [] + for leaf_name in record.leaf_names: + spec = record.leaf_spec(leaf_name) + layout.append( + ( + leaf_name, + tuple( + (item.role, tuple(item.shape), str(item.dtype)) + for item in spec.tensors + ), + spec.execution_key, + ) + ) + return tuple(layout) + + +def build_block_abi( + residency: ResidencyState, + block_key: str, + expected: tuple[str, ...], +) -> ImmutableBlockABI: + record = residency.arena.block_record(block_key) + if record is None: + raise ImmutableRuntimeError(f"missing_canonical_block:{block_key}") + + if record.leaf_names != expected: + raise ImmutableRuntimeError( + f"canonical_leaf_order_mismatch:{block_key}:expected={expected}:actual={record.leaf_names}" + ) + + return ImmutableBlockABI( + block_key=block_key, + leaf_names=record.leaf_names, + leaf_layout=_leaf_layout(record), + ) + + +class ImmutableTransformerRuntime: + """Source publication and residency policy for the block dispatcher.""" + + TRAIN = "train" + SAMPLE = "sample" + + def __init__( + self, + model, + residency: ResidencyState, + *, + blocks, + block_keys, + entries_by_block, + depth: int = 3, + compile_blocks: bool = True, + compile_dynamic: bool | None = True, + compile_dynamic_hints: tuple[tuple[int, int | None, int | None], ...] = (), + protected_training_leaf_keys=(), + owner_token=None, + ) -> None: + self._sampling_working_bytes: dict[tuple, int] = {} + self._sampling_baseline = None + # External-VRAM check is reported once per phase, not per image/step. + self._foreign_vram_checked = False + self.model = model + self.residency = residency + self._blocks = tuple(blocks) + block_keys = tuple(str(key) for key in block_keys) + self.depth = max(1, int(depth)) + self.compile_blocks = bool(compile_blocks) + self.compile_dynamic = ( + None if compile_dynamic is None else bool(compile_dynamic) + ) + self.compile_dynamic_hints = tuple(compile_dynamic_hints or ()) + self.protected_training_leaf_keys = frozenset( + (str(block), str(leaf)) + for block, leaf in protected_training_leaf_keys + ) + self.owner_token = owner_token + self._hint_range_warned: set[tuple] = set() + self._arena_signature = self.residency.arena.immutable_signature() + + self._block_abis = tuple( + build_block_abi( + residency, + block_key, + tuple(name for name, _module in entries_by_block[block_key]), + ) + for block_key in block_keys + ) + if len(self._blocks) != len(self._block_abis): + raise ImmutableRuntimeError("dispatcher_block_count_mismatch") + self._sources = ImmutableRuntimeSourceTable(residency, self._block_abis) + self._block_kernels: dict[tuple[str, int], object] = {} + self._programs: dict[str, ImmutableProgram] = {} + self._finalized = False + self._finalization_signature = None + self._active_token = None + self._active_mode = None + self.stats = { + "residency_transitions": 0, + "source_generation": self._sources.generation, + } + self.sampling_fallback_plan = ResidencyPlan.build( + "sample_fallback", + (), + ) + + @property + def source_generation(self) -> int: + return self._sources.generation + + @property + def active_executions(self) -> int: + return self._sources.active_executions + + @property + def finalized(self) -> bool: + return self._finalized + + def source(self, block_index: int) -> ImmutableBlockSourceSnapshot: + """Current published source snapshot for one block.""" + return self._sources.source(block_index) + + def _require_finalized(self) -> None: + if not self._finalized: + raise ImmutableRuntimeError("immutable_runtime_not_finalized") + + + @contextmanager + def execution(self, mode: str): + self._require_finalized() + if mode not in (self.TRAIN, self.SAMPLE): + raise ImmutableRuntimeError(f"unknown_execution_mode:{mode}") + if self._active_token is not None: + raise ImmutableRuntimeError( + f"immutable_execution_already_active:{self._active_mode}" + ) + + token = object() + generation = self._sources.begin_execution() + self._active_token = token + self._active_mode = mode + try: + yield self + finally: + if self._active_token is not token: + raise ImmutableRuntimeError("immutable_execution_token_mismatch") + self._active_token = None + self._active_mode = None + self._sources.end_execution(generation) + + @contextmanager + def sampling(self, *, shape_key: tuple, **activate_kwargs): + succeeded = False + try: + self.activate_sampling_image( + shape_key=shape_key, + **activate_kwargs, + ) + with self.execution(self.SAMPLE): + yield self + succeeded = True + finally: + if succeeded: + self.finish_sampling_image(shape_key=shape_key) + else: + self._sampling_baseline = None + def _assert_arena_stable(self, where: str) -> None: + current = self.residency.arena.immutable_signature() + if current != self._arena_signature: + raise ImmutableRuntimeError( + f"arena_mutated_at_boundary:{where}: canonical host flats " + "or registrations changed across a phase boundary" + ) + + def set_compile_dynamic_hints(self, hints) -> None: + """Install mark_dynamic hints derived after the runtime was prepared. + + The trainer can only compute sequence bounds once the datasets exist, + which is long after `prepare_arena_offload`. Hints are read per call, so + installing them any time before the first compiled block call is enough. + A block kernel that already traced would have to re-specialize, so + refuse once anything is compiled rather than pay a silent recompile. + """ + hints = tuple(tuple(hint) for hint in (hints or ())) + if hints == self.compile_dynamic_hints: + return + if self._block_kernels: + raise RuntimeError( + "compile_dynamic_hints changed after block kernels were built; " + "set them before the first forward pass." + ) + self.compile_dynamic_hints = hints + self._hint_range_warned.clear() + + def _warn_hint_out_of_range(self, dim, size, lo, hi) -> None: + key = (dim, size) + if key in self._hint_range_warned: + return + self._hint_range_warned.add(key) + print( + f"[immutable] dim {dim} size {size} is outside the declared dynamic " + f"range [{lo}, {hi}]; compiling a dedicated shape for it. " + "Widen compile_dynamic_hints to avoid the extra compile." + ) + + def set_residency_plan(self, plan: ResidencyPlan) -> ResidencyDelta: + self._assert_arena_stable("pre_residency_publish") + delta = self._sources.publish(plan) + self._assert_arena_stable("post_residency_publish") + self.stats["residency_transitions"] += 1 + self.stats["source_generation"] = self._sources.generation + return delta + + def activate(self, mode: str, plan: ResidencyPlan) -> ImmutableProgram: + self._require_finalized() + if mode not in (self.TRAIN, self.SAMPLE): + raise ImmutableRuntimeError(f"unknown_execution_mode:{mode}") + self.set_residency_plan(plan) + return self.program(mode) + + def program(self, mode: str) -> ImmutableProgram: + self._require_finalized() + try: + return self._programs[mode] + except KeyError as error: + raise ImmutableRuntimeError(f"unknown_execution_mode:{mode}") from error + + def activate_sampling_fallback(self) -> ImmutableProgram: + self.set_residency_plan(self.sampling_fallback_plan) + return self.program(self.SAMPLE) + + def transition_training_blocks(self, block_keys, *, resident: bool) -> dict: + """Atomically add or remove complete training blocks in one plan.""" + current = self._sources.plan or self.residency.plan + if current.phase != self.TRAIN: + raise ImmutableRuntimeError( + f"training_block_transition_requires_train:{current.phase}" + ) + requested = tuple(dict.fromkeys(str(key) for key in block_keys)) + protected = self.protected_training_leaf_keys + next_keys = set(current.resident_leaf_keys) + changed = [] + for key in requested: + abi = next( + (item for item in self._block_abis if item.block_key == key), + None, + ) + if abi is None: + raise ImmutableRuntimeError(f"unknown_training_block:{key}") + leaf_keys = tuple((key, leaf) for leaf in abi.leaf_names) + present = tuple( + item for item in leaf_keys if item in current.resident_leaf_keys + ) + if present and len(present) != len(leaf_keys): + raise ImmutableRuntimeError( + f"partial_training_block_layout:{key}" + ) + if not resident and any(item in protected for item in leaf_keys): + raise ImmutableRuntimeError(f"protected_training_block:{key}") + if bool(present) == bool(resident): + continue + changed.append(key) + if resident: + next_keys.update(leaf_keys) + else: + next_keys.difference_update(leaf_keys) + if not changed: + return { + "changed": False, + "block_keys": requested, + "resident": bool(resident), + "plan": current, + } + next_plan = ResidencyPlan.build(self.TRAIN, next_keys) + delta = self.set_residency_plan(next_plan) + return { + "changed": True, + "block_keys": tuple(changed), + "resident": bool(resident), + "resident_bytes": self.residency.resident_bytes(), + "delta": delta, + "plan": next_plan, + } + + def transition_training_block(self, block_key: str, *, resident: bool) -> dict: + """Atomically add or remove one complete training block by stable key.""" + current = self._sources.plan or self.residency.plan + if current.phase != self.TRAIN: + raise ImmutableRuntimeError( + f"training_block_transition_requires_train:{current.phase}" + ) + key = str(block_key) + abi = next((item for item in self._block_abis if item.block_key == key), None) + if abi is None: + raise ImmutableRuntimeError(f"unknown_training_block:{key}") + leaf_keys = tuple((key, leaf) for leaf in abi.leaf_names) + present = tuple(item for item in leaf_keys if item in current.resident_leaf_keys) + if present and len(present) != len(leaf_keys): + raise ImmutableRuntimeError(f"partial_training_block_layout:{key}") + want_resident = bool(resident) + if bool(present) == want_resident: + return { + "changed": False, + "block_key": key, + "resident": want_resident, + "plan": current, + } + protected = self.protected_training_leaf_keys + if not want_resident and any(item in protected for item in leaf_keys): + raise ImmutableRuntimeError(f"protected_training_block:{key}") + + next_keys = set(current.resident_leaf_keys) + if want_resident: + next_keys.update(leaf_keys) + else: + next_keys.difference_update(leaf_keys) + next_plan = ResidencyPlan.build(self.TRAIN, next_keys) + delta = self.set_residency_plan(next_plan) + return { + "changed": True, + "block_key": key, + "resident": want_resident, + "resident_bytes": self.residency.resident_bytes(), + "delta": delta, + "plan": next_plan, + } + + def next_training_promotion_bytes(self) -> int: + current = self._sources.plan or self.residency.plan + if current.phase != self.TRAIN: + return 0 + protected = self.protected_training_leaf_keys + candidates = [] + for abi in self._block_abis: + keys = tuple((abi.block_key, leaf) for leaf in abi.leaf_names) + if any(key in current.resident_leaf_keys for key in keys): + continue + if any(key in protected for key in keys): + continue + record = self.residency.arena.block_record(abi.block_key) + candidates.append(int(record.committed_bytes)) + return min(candidates, default=0) + def increase_training_residency( + self, + available_growth_bytes: int, + *, + max_blocks: int = 1, + ) -> dict: + available = max(0, int(available_growth_bytes)) + current = self._sources.plan or self.residency.plan + if current.phase != self.TRAIN: + raise ImmutableRuntimeError( + f"training_residency_growth_requires_train:{current.phase}" + ) + + protected = self.protected_training_leaf_keys + candidates = [] + for order, abi in enumerate(self._block_abis): + keys = tuple((abi.block_key, leaf) for leaf in abi.leaf_names) + resident_count = sum( + key in current.resident_leaf_keys + for key in keys + ) + if resident_count: + # Initial and controller plans are whole-block. Preserve + # source-table support for partial plans without growing them. + continue + if any(key in protected for key in keys): + continue + record = self.residency.arena.block_record(abi.block_key) + candidates.append( + ( + int(record.committed_bytes), + order, + abi.block_key, + keys, + ) + ) + + candidates.sort(key=lambda item: (item[0], item[1])) + added = [] + predicted = 0 + limit = max(0, int(max_blocks)) + for nbytes, _order, _block_key, keys in candidates: + if limit and len(added) >= limit: + break + if predicted + nbytes > available: + continue + added.append((nbytes, keys)) + predicted += nbytes + + previous_plan = current + if added: + next_keys = set(current.resident_leaf_keys) + for _nbytes, keys in added: + next_keys.update(keys) + next_plan = ResidencyPlan.build(self.TRAIN, next_keys) + self.set_residency_plan(next_plan) + else: + next_plan = current + + added_keys = tuple( + sorted( + key + for _nbytes, keys in added + for key in keys + ) + ) + actual_growth = sum( + self.residency.resident_leaf_bytes(key) + for key in added_keys + ) + return { + "available_growth_bytes": available, + "predicted_growth_bytes": int(predicted), + "actual_growth_bytes": int(actual_growth), + "added_leaf_keys": added_keys, + "added_blocks": tuple(sorted({key[0] for key in added_keys})), + "previous_plan": previous_plan, + "plan": next_plan, + } + def reduce_training_residency(self, required_relief_bytes: int) -> dict: + requested = max(0, int(required_relief_bytes)) + current = self._sources.plan or self.residency.plan + if current.phase != self.TRAIN: + raise ImmutableRuntimeError(f"training_residency_reduction_requires_train:{current.phase}") + + protected = self.protected_training_leaf_keys + candidates = [] + for abi in self._block_abis: + keys = tuple((abi.block_key, leaf) for leaf in abi.leaf_names) + resident_keys = tuple(key for key in keys if key in current.resident_leaf_keys) + if not resident_keys or any(key in protected for key in keys): + continue + nbytes = sum(self.residency.resident_leaf_bytes(key) for key in resident_keys) + candidates.append((nbytes, abi.block_key, resident_keys)) + + candidates.sort(key=lambda item: (-item[0], item[1])) + removed = [] + relieved = 0 + for nbytes, _block_key, keys in candidates: + if relieved >= requested: + break + removed.extend(keys) + relieved += nbytes + + if removed: + next_keys = set(current.resident_leaf_keys) - set(removed) + next_plan = ResidencyPlan.build(self.TRAIN, next_keys) + self.set_residency_plan(next_plan) + else: + next_plan = current + + return { + "requested_relief_bytes": requested, + "relieved_bytes": int(relieved), + "removed_leaf_keys": tuple(sorted(removed)), + "removed_blocks": tuple(sorted({key[0] for key in removed})), + "remaining_resident_bytes": self.residency.resident_bytes(), + "plan": next_plan, + } + + def full_model_resident_bytes(self) -> int: + """Bytes to hold every canonical block resident (the 'want' figure).""" + arena = self.residency.arena + total = 0 + for block_key in arena.block_keys(): + record = arena.block_record(block_key) + if record is not None: + total += int(record.committed_bytes) + return total + + def report_foreign_vram_once(self, device, *, phase: str, working_reserve_bytes: int) -> None: + """Phase-start external-VRAM check for callers that know their reserve. + + ``have`` is the residency budget: what we already hold, plus whatever is + left on the card once the activation working set is reserved. + """ + if self._foreign_vram_checked: + return + self._foreign_vram_checked = True + if torch.device(device).type != "cuda": + return + try: + free_bytes = vram_budget.device_free_bytes(device) + reserved_bytes = int(torch.cuda.memory_reserved(device)) + allocated_bytes = int(torch.cuda.memory_allocated(device)) + except Exception: + return + reclaimable = max(0, reserved_bytes - allocated_bytes) + have = self.residency.resident_bytes() + max( + 0, free_bytes + reclaimable - max(0, int(working_reserve_bytes)) + ) + self._report_foreign_vram( + device, + phase=phase, + reserved_bytes=reserved_bytes, + have_bytes=have, + ) + + def _report_foreign_vram(self, device, *, phase: str, reserved_bytes, have_bytes) -> None: + """Say so when ANOTHER tenant on the GPU is what forces us to stream. + + Without this the failure is invisible: residency silently shrinks, the + step time multiplies, and nothing in the log points at the real cause + (a leftover job, a ComfyUI server, a game). Fires once per phase. + """ + try: + free_bytes, total_bytes = vram_budget.device_mem_info(device) + except Exception: + return + report = vram_budget.assess_foreign_vram( + total_bytes=total_bytes, + free_bytes=free_bytes, + torch_reserved_bytes=reserved_bytes, + want_bytes=self.full_model_resident_bytes(), + have_bytes=have_bytes, + ) + message = vram_budget.format_foreign_vram_warning(report, phase=phase) + if message: + print(message) + + def activate_sampling_image( + self, + *, + shape_key: tuple, + cold_working_bytes: int, + fixed_working_bytes: int | None, + cold_floor_bytes: int, + hot_floor_bytes: int, + measured_pad_bytes: int = 256 * 1024**2, + measured_floor_bytes: int = 512 * 1024**2, + ) -> ImmutableProgram: + device = self.residency.device + if device.type != "cuda": + self.set_residency_plan(self.sampling_fallback_plan) + return self.program(self.SAMPLE) + + learned = int(self._sampling_working_bytes.get(shape_key, 0)) + if fixed_working_bytes is not None: + working_bytes = max(0, int(fixed_working_bytes)) + floor_bytes = max(0, int(cold_floor_bytes)) + reserve_source = "fixed" + elif learned > 0: + working_bytes = max( + int(measured_floor_bytes), + learned + int(measured_pad_bytes), + ) + floor_bytes = max(0, int(hot_floor_bytes)) + reserve_source = "measured" + else: + working_bytes = max(0, int(cold_working_bytes)) + floor_bytes = max(0, int(cold_floor_bytes)) + reserve_source = "cold" + + # Physical free across ALL processes (NVML-backed). mem_get_info would + # over-report here whenever anything else is on the card -- a game, a + # ComfyUI server, an orphaned job -- and we would size residency against + # VRAM that does not exist, then page silently at ~4x the step time. + free_bytes = vram_budget.device_free_bytes(device) + allocated_bytes = torch.cuda.memory_allocated(device) + reserved_bytes = torch.cuda.memory_reserved(device) + reclaimable_cache = max(0, reserved_bytes - allocated_bytes) + current_sidecars = self.residency.resident_bytes() + resident_budget = max( + 0, + current_sidecars + int(free_bytes) + reclaimable_cache - working_bytes - floor_bytes, + ) + if not self._foreign_vram_checked: + self._foreign_vram_checked = True + self._report_foreign_vram( + device, + phase="sampling", + reserved_bytes=reserved_bytes, + have_bytes=resident_budget, + ) + + current_plan = self._sources.plan or self.residency.plan + plan = ResidencyPlan.fit_whole_blocks( + self.residency.arena, + resident_budget, + phase=self.SAMPLE, + prefer_resident_keys=current_plan.resident_leaf_keys, + ) + self.set_residency_plan(plan) + + torch.cuda.synchronize(device) + baseline_allocated = torch.cuda.memory_allocated(device) + baseline_reserved = torch.cuda.memory_reserved(device) + torch.cuda.reset_peak_memory_stats(device) + self._sampling_baseline = { + "shape_key": shape_key, + "allocated": baseline_allocated, + "reserved": baseline_reserved, + "working_bytes": working_bytes, + "floor_bytes": floor_bytes, + "source": reserve_source, + } + + print( + "[MemoryManager] immutable sampling layout: " + f"source={reserve_source} " + f"working={working_bytes / 1024**3:.2f} GiB " + f"floor={floor_bytes / 1024**3:.2f} GiB " + f"sidecars={self.residency.resident_bytes() / 1024**3:.2f} GiB " + f"device_free={vram_budget.device_free_bytes(device) / 1024**3:.2f} GiB " + f"plan={plan.fingerprint}" + ) + return self.program(self.SAMPLE) + + def finish_sampling_image(self, *, shape_key: tuple) -> int: + baseline = self._sampling_baseline + if baseline is None or baseline["shape_key"] != shape_key: + return 0 + + device = self.residency.device + torch.cuda.synchronize(device) + allocated_peak = torch.cuda.max_memory_allocated(device) + reserved_peak = torch.cuda.max_memory_reserved(device) + allocated_growth = max( + 0, + allocated_peak - int(baseline["allocated"]), + ) + reserved_growth = max( + 0, + reserved_peak - int(baseline["reserved"]), + ) + observed = max(allocated_growth, reserved_growth) + previous = int(self._sampling_working_bytes.get(shape_key, 0)) + self._sampling_working_bytes[shape_key] = max(previous, observed) + self._sampling_baseline = None + + print( + "[MemoryManager] immutable sampling measurement: " + f"shape={shape_key} " + f"observed_working={observed / 1024**3:.2f} GiB " + f"learned={self._sampling_working_bytes[shape_key] / 1024**3:.2f} GiB" + ) + return observed + + def close(self) -> None: + if self._sources.active_executions: + raise ImmutableRuntimeError("cannot_close_during_execution") + self._sources.clear() + self._programs.clear() + self._block_kernels.clear() diff --git a/toolkit/memory_management/residency.py b/toolkit/memory_management/residency.py new file mode 100644 index 0000000000..114ff6e91c --- /dev/null +++ b/toolkit/memory_management/residency.py @@ -0,0 +1,338 @@ +"""Manager-owned device residency sidecars for the immutable host arena. + +Residency changes never mutate a Parameter or the canonical host allocation. +They publish optional device tensors keyed by stable ``(block, leaf)`` keys; +execution adapters choose a sidecar or a compact fetched view from a static +transfer plan. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass + +import torch + +from toolkit.memory_management.canonical_arena import CanonicalArena +from toolkit.memory_management.arena_offload.layout import ( + + _flatten_leaves, + _rebuild_from_leaves, + leaf_view, +) + +LeafKey = tuple[str, str] + +def _interleave_priority(index: int, count: int) -> float: + if count <= 1: + return 0.0 + bits = (count - 1).bit_length() + value = index + reversed_bits = 0 + for _ in range(bits): + reversed_bits = (reversed_bits << 1) | (value & 1) + value >>= 1 + return reversed_bits / float(1 << bits) +class ResidencyError(RuntimeError): + """A residency plan or transition violated an immutable-arena invariant.""" + + +@dataclass(frozen=True) +class ResidencyPlan: + phase: str + resident_leaf_keys: frozenset[LeafKey] + fingerprint: str + @classmethod + def fit_whole_blocks( + cls, + arena: CanonicalArena, + resident_budget_bytes: int, + *, + phase: str, + prefer_resident_keys=(), + ) -> ResidencyPlan: + """Select complete canonical blocks within a sidecar-only byte budget. + + Existing fully resident blocks are preferred to avoid needless sidecar + churn when a sampling plan grows or shrinks between images. + """ + budget = max(0, int(resident_budget_bytes)) + preferred = frozenset( + (str(block), str(leaf)) + for block, leaf in prefer_resident_keys + ) + + blocks = [] + block_keys = arena.block_keys() + + for order, block_key in enumerate(block_keys): + record = arena.block_record(block_key) + leaf_keys = tuple( + (block_key, leaf_name) + for leaf_name in record.leaf_names + ) + already_full = all(key in preferred for key in leaf_keys) + + blocks.append( + { + "order": order, + "block_key": block_key, + # Slightly conservative because this includes page rounding. + "nbytes": int(record.committed_bytes), + "leaf_keys": leaf_keys, + "already_full": already_full, + } + ) + + blocks.sort( + key=lambda item: ( + not item["already_full"], + item["nbytes"], + _interleave_priority(item["order"], len(blocks)), + ) + ) + + selected = [] + used = 0 + + for item in blocks: + nbytes = item["nbytes"] + if used + nbytes > budget: + continue + selected.extend(item["leaf_keys"]) + used += nbytes + + return cls.build(phase, selected) + @classmethod + def build(cls, phase: str, resident_leaf_keys) -> ResidencyPlan: + keys = frozenset((str(block), str(leaf)) for block, leaf in resident_leaf_keys) + source = repr((str(phase), tuple(sorted(keys)))) + fingerprint = hashlib.sha1(source.encode("utf-8")).hexdigest()[:16] + return cls(str(phase), keys, fingerprint) + + @classmethod + def from_smart_plan( + cls, arena: CanonicalArena, smart_plan: dict, *, phase: str + ) -> ResidencyPlan: + """Adapt the existing planner's ``offload_ids`` decision to sidecars. + + This is the Slice 3 planner seam: priority and capacity remain owned by + ``MemoryManager.smart_training_plan``; only the mutation target changes. + """ + offload_ids = set(smart_plan.get("offload_ids", ())) + resident = [] + for block_key in arena.block_keys(): + block = arena.block_record(block_key) + block_is_fully_resident = all( + id(module) not in offload_ids + for module in block.modules + ) + if not block_is_fully_resident: + continue + resident.extend( + (block_key, leaf_name) + for leaf_name in block.leaf_names + ) + return cls.build(phase, resident) + + def resident_in_block(self, block_key: str) -> frozenset[str]: + return frozenset(leaf for block, leaf in self.resident_leaf_keys if block == block_key) + + +@dataclass(frozen=True) +class ResidentLeaf: + key: LeafKey + tensors: tuple[torch.Tensor, ...] + weight_leaf_count: int + weight_template: torch.Tensor + ready_event: torch.cuda.Event | None + nbytes: int + + @property + def weight(self): + weight_tensors = self.tensors[:self.weight_leaf_count] + if self.weight_leaf_count == 1: + return weight_tensors[0] + return _rebuild_from_leaves(self.weight_template, iter(weight_tensors)) + + @property + def bias(self): + if len(self.tensors) == self.weight_leaf_count: + return None + return self.tensors[self.weight_leaf_count] + + +@dataclass(frozen=True) +class ResidencyDelta: + promoted: tuple[LeafKey, ...] + demoted: tuple[LeafKey, ...] + resident_bytes: int + + +def _tensor_bytes(tensor: torch.Tensor | None) -> int: + if tensor is None: + return 0 + return sum(leaf.numel() * leaf.element_size() for leaf in _flatten_leaves(tensor)) + + +def _record_stream(tensor: torch.Tensor | None, stream) -> None: + if tensor is None: + return + for leaf in _flatten_leaves(tensor): + leaf.record_stream(stream) + + +class ResidencyState: + """Atomic per-Linear sidecar state over one immutable canonical arena.""" + + def __init__(self, arena: CanonicalArena, device) -> None: + if not arena.canonicalized: + raise ResidencyError("residency_requires_canonicalized_arena") + self.arena = arena + self.device = torch.device(device) + self._sidecars: dict[LeafKey, ResidentLeaf] = {} + self._plan = ResidencyPlan.build("empty", ()) + self._copy_stream = ( + torch.cuda.Stream(device=self.device) if self.device.type == "cuda" else None + ) + + @property + def plan(self) -> ResidencyPlan: + return self._plan + + def _all_keys(self) -> frozenset[LeafKey]: + return frozenset( + (block_key, leaf_name) + for block_key in self.arena.block_keys() + for leaf_name in self.arena.block_record(block_key).leaf_names + ) + + def _canonical_leaf(self, key: LeafKey): + block_key, leaf_name = key + block = self.arena.block_record(block_key) + if block is None: + raise ResidencyError(f"unknown_residency_block:{block_key}") + try: + spec = block.leaf_spec(leaf_name) + module = block.module_for_leaf(leaf_name) + except KeyError as error: + raise ResidencyError(f"unknown_residency_leaf:{block_key}.{leaf_name}") from error + return block, spec, module + + def _build_sidecar(self, key: LeafKey) -> ResidentLeaf: + block, spec, _module = self._canonical_leaf(key) + stream_context = ( + torch.cuda.stream(self._copy_stream) + if self._copy_stream is not None + else torch.no_grad() + ) + with torch.no_grad(), stream_context: + tensors = tuple( + leaf_view(block.host_flat, item).to( + self.device, non_blocking=self.device.type == "cuda" + ) + for item in spec.tensors + ) + event = None + if self._copy_stream is not None: + event = torch.cuda.Event() + event.record(self._copy_stream) + return ResidentLeaf( + key=key, + tensors=tensors, + weight_leaf_count=spec.weight_leaf_count, + weight_template=spec.weight_template, + ready_event=event, + nbytes=sum(_tensor_bytes(tensor) for tensor in tensors), + ) + + def reconcile(self, plan: ResidencyPlan) -> ResidencyDelta: + desired = plan.resident_leaf_keys + unknown = desired - self._all_keys() + if unknown: + block, leaf = sorted(unknown)[0] + raise ResidencyError(f"unknown_residency_leaf:{block}.{leaf}") + + before_arena = self.arena.immutable_signature() + current = set(self._sidecars) + additions = tuple(sorted(desired - current)) + removals = tuple(sorted(current - desired)) + pending: dict[LeafKey, ResidentLeaf] = {} + try: + for key in additions: + pending[key] = self._build_sidecar(key) + except Exception as error: + # Copies may already be queued on the private stream. Drain them + # before pending tensors are released, while leaving published + # sidecars and the active plan exactly unchanged. + if self._copy_stream is not None: + self._copy_stream.synchronize() + pending.clear() + if self.arena.immutable_signature() != before_arena: + raise ResidencyError( + "pin_ledger_changed_during_failed_promotion" + ) from error + raise + + next_sidecars = { + key: value for key, value in self._sidecars.items() if key not in removals + } + next_sidecars.update(pending) + self._sidecars = next_sidecars + self._plan = plan + if self.arena.immutable_signature() != before_arena: + raise ResidencyError("pin_ledger_changed_during_residency_transition") + return ResidencyDelta(additions, removals, self.resident_bytes()) + + def promote(self, key: LeafKey, *, phase: str | None = None) -> bool: + normalized = (str(key[0]), str(key[1])) + if normalized in self._sidecars: + return False + plan = ResidencyPlan.build( + phase or self._plan.phase, self._sidecars.keys() | {normalized} + ) + self.reconcile(plan) + return True + + def demote(self, key: LeafKey, *, phase: str | None = None) -> bool: + normalized = (str(key[0]), str(key[1])) + if normalized not in self._sidecars: + return False + plan = ResidencyPlan.build( + phase or self._plan.phase, set(self._sidecars) - {normalized} + ) + self.reconcile(plan) + return True + + def resident_leaf(self, key: LeafKey) -> ResidentLeaf | None: + sidecar = self._sidecars.get((str(key[0]), str(key[1]))) + if sidecar is None: + return None + if sidecar.ready_event is not None: + current = torch.cuda.current_stream(self.device) + current.wait_event(sidecar.ready_event) + for tensor in sidecar.tensors: + _record_stream(tensor, current) + return sidecar + + def resident_tensor(self, key: LeafKey) -> torch.Tensor | None: + sidecar = self.resident_leaf(key) + return None if sidecar is None else sidecar.weight + + def streamed_leaf_names(self, block_key: str) -> tuple[str, ...]: + block = self.arena.block_record(block_key) + if block is None: + raise ResidencyError(f"unknown_residency_block:{block_key}") + resident = self._plan.resident_in_block(block_key) + return tuple(name for name in block.leaf_names if name not in resident) + + def resident_leaf_bytes(self, key: LeafKey) -> int: + """Return published sidecar bytes without synchronizing its copy event.""" + sidecar = self._sidecars.get((str(key[0]), str(key[1]))) + return 0 if sidecar is None else int(sidecar.nbytes) + + def resident_bytes(self) -> int: + return sum(sidecar.nbytes for sidecar in self._sidecars.values()) + + def clear(self, *, phase: str = "clear") -> ResidencyDelta: + return self.reconcile(ResidencyPlan.build(phase, ())) diff --git a/toolkit/memory_management/runtime.py b/toolkit/memory_management/runtime.py new file mode 100644 index 0000000000..a28b65226a --- /dev/null +++ b/toolkit/memory_management/runtime.py @@ -0,0 +1,82 @@ +"""Generic memory-runtime discovery used by shared trainer code.""" + +from __future__ import annotations + + +RUNTIME_ATTR = "_arena_offload_runtime" + + +def unwrap_memory_model(model): + seen = set() + while model is not None and id(model) not in seen: + seen.add(id(model)) + original = getattr(model, "_orig_mod", None) + if original is not None and original is not model: + model = original + continue + if getattr(model, RUNTIME_ATTR, None) is not None: + return model + if hasattr(model, "_memory_manager"): + return model + inner = getattr(model, "module", None) + if inner is None or inner is model: + return model + model = inner + return model + + +def get_memory_runtime(model): + if model is None: + return None + return getattr(unwrap_memory_model(model), RUNTIME_ATTR, None) + + +def is_memory_managed(model) -> bool: + if model is None: + return False + inner = unwrap_memory_model(model) + return ( + get_memory_runtime(inner) is not None + or hasattr(inner, "_memory_manager") + or bool(getattr(inner, "_arena_offload_disposed", False)) + ) + + +def memory_runtime_owns_compile(model) -> bool: + runtime = get_memory_runtime(model) + return bool(runtime is not None and getattr(runtime, "owns_compile", True)) + + +def close_memory_runtime(model) -> None: + runtime = get_memory_runtime(model) + if runtime is not None: + runtime.close() + + +def close_memory_runtime_preparation(model_owner) -> None: + """Release a loader-scoped preparation that never published a runtime.""" + pending_models = getattr(model_owner, "_arena_pending_load_models", ()) + if hasattr(model_owner, "_arena_pending_load_models"): + delattr(model_owner, "_arena_pending_load_models") + if pending_models: + from .arena_offload.load_session import discard_pending_canonical_build + + for model in pending_models: + discard_pending_canonical_build(model) + operation = getattr(model_owner, "cleanup_memory_runtime_preparation", None) + if operation is not None: + operation() + + +def memory_sampling_step_trim(model) -> bool: + """Run the legacy backend's optional per-step sampling trim.""" + inner = unwrap_memory_model(model) + operation = getattr(inner, "_mm_sampling_step_trim", None) + return bool(operation is not None and operation()) + + +def memory_sampling_demote(model, *, reason) -> bool: + """Request sampling residency relief from the active legacy backend.""" + inner = unwrap_memory_model(model) + operation = getattr(inner, "_mm_sampling_demote", None) + return bool(operation is not None and operation(reason=reason)) diff --git a/toolkit/memory_management/transfer_plan.py b/toolkit/memory_management/transfer_plan.py new file mode 100644 index 0000000000..4d2674e9ed --- /dev/null +++ b/toolkit/memory_management/transfer_plan.py @@ -0,0 +1,216 @@ +"""Static multi-range transfer plans (Slice 2, +tasks/open/IMMUTABLE_TRANSFER_ARENA_PLAN.md). + +A ``BlockTransferPlan`` says which byte ranges of a canonical block's host +flat (``canonical_arena.BlockRecord``) need to move to the device for one +residency phase, coalesced and packed into a compact destination layout. +Pure data model + coalescing algorithm live here (no CUDA needed to build +or inspect a plan); the runtime that actually submits the copies is +``mm::fetch_start_multi`` in ``ingraph_stream``, reusing its existing +ticket/ring/backpressure machinery unchanged (Slice 2's "single-ticket +semantics"). + +Immutable and fingerprintable per Invariant 8: two plans built from the +same block + the same set of streamed leaf names always produce identical +ranges and the same fingerprint, independent of the actual tensor/storage +identity behind the block's host flat at build time. +""" + +from __future__ import annotations + +import hashlib +from dataclasses import dataclass +from itertools import pairwise +from types import MappingProxyType + +import torch + +from toolkit.memory_management.canonical_arena import BlockRecord +from toolkit.memory_management.arena_offload.layout import LEAF_ALIGN + + + +class TransferPlanError(ValueError): + """A transfer plan could not be built from a block's canonical layout.""" + + +@dataclass(frozen=True) +class LeafRange: + """One coalesced source(host) -> destination(compact device) copy span.""" + + src_offset: int + dst_offset: int + nbytes: int + + +@dataclass(frozen=True) +class CompactLeafSpec: + """A streamed leaf's location within the plan's compact destination + buffer -- same shape as ``ingraph_stream.LeafSpec`` but the offset is + in DESTINATION (compact device buffer) coordinates, not the arena + flat's.""" + + dst_offset: int + nbytes: int + dtype: torch.dtype + shape: tuple[int, ...] + role: str + + +@dataclass(frozen=True) +class BlockTransferPlan: + block_key: str + ranges: tuple[LeafRange, ...] + compact_nbytes: int + # leaf_name -> {declared tensor name: CompactLeafSpec} + leaf_specs: dict + streamed_leaf_names: tuple[str, ...] + fully_streamed: bool + fingerprint: str + + @property + def num_ranges(self) -> int: + return len(self.ranges) + + def compact_leaf_view(self, device_flat: torch.Tensor, leaf_name: str, role: str) -> torch.Tensor: + spec = self.leaf_specs[leaf_name][role] + return ( + device_flat[spec.dst_offset:spec.dst_offset + spec.nbytes] + .view(spec.dtype) + .reshape(spec.shape) + ) + + def ranges_tensor(self) -> torch.Tensor: + """(N, 3) int64 CPU tensor of [src_offset, dst_offset, nbytes] rows, + the exact argument shape the ``mm::fetch_start_multi`` custom op + expects (Slice 2's compile-visible extended fetch op).""" + if not self.ranges: + return torch.empty((0, 3), dtype=torch.int64) + return torch.tensor( + [[r.src_offset, r.dst_offset, r.nbytes] for r in self.ranges], + dtype=torch.int64, + ) + + +def _leaf_items(block: BlockRecord, streamed_leaf_names: frozenset): + """(src_offset, nbytes, leaf_name, role, LeafSpec) for every streamed + leaf's declared storage tensors, sorted by source offset.""" + items = [] + for spec in block.pack.linears: + if spec.name not in streamed_leaf_names: + continue + for leaf_spec in spec.tensors: + role = leaf_spec.role + items.append((leaf_spec.offset, leaf_spec.nbytes, spec.name, role, leaf_spec)) + items.sort(key=lambda item: item[0]) + for (off, n, *_), (noff, *_rest) in pairwise(items): + if noff < off + n: + raise TransferPlanError( + f"transfer_plan_overlap:{block.block_key}: leaf ranges overlap " + "-- canonical arena layout invariant violated" + ) + return items + + +def build_transfer_plan( + block: BlockRecord, + streamed_leaf_names, + *, + slack_bytes: int = LEAF_ALIGN - 1, +) -> BlockTransferPlan: + """Build a coalesced multi-range transfer plan for the STREAMED subset + of ``block``'s leaves. ``streamed_leaf_names`` is the residency + decision (owned by the planner/controllers, Invariant 10) -- this + function only turns "which leaves stream this phase" into "which byte + ranges to copy and where." + + Two source items coalesce into one range when the gap between them is + at most ``slack_bytes`` (default: one leaf's alignment padding, + ``LEAF_ALIGN - 1``) -- adjacent streamed leaves separated only by + padding merge into a single copy; a resident (non-streamed) leaf in + between breaks the run. Coalescing never changes a leaf's own + destination offset formula (range-relative), only how many discrete + copies the runtime submits. + """ + streamed = frozenset(streamed_leaf_names) + known = frozenset(spec.name for spec in block.pack.linears) + unknown = streamed - known + if unknown: + raise TransferPlanError( + f"transfer_plan_unknown_leaf:{block.block_key}:{sorted(unknown)[0]}" + ) + if slack_bytes < 0: + raise TransferPlanError("transfer_plan_slack_bytes must be non-negative") + items = _leaf_items(block, streamed) + if not items: + raise TransferPlanError(f"transfer_plan_empty:{block.block_key}: no streamed leaves") + + ranges: list[LeafRange] = [] + leaf_specs: dict[str, dict[str, CompactLeafSpec]] = {} + compact_total = 0 + + def flush(src_start: int, src_end: int, dst_start: int, pending: list) -> int: + length = src_end - src_start + ranges.append(LeafRange(src_offset=src_start, dst_offset=dst_start, nbytes=length)) + for off, n, name, role, leaf_spec in pending: + leaf_specs.setdefault(name, {})[role] = CompactLeafSpec( + dst_offset=dst_start + (off - src_start), + nbytes=n, + dtype=leaf_spec.dtype, + shape=leaf_spec.shape, + role=leaf_spec.role, + ) + return length + + src_start, first_len = items[0][0], items[0][1] + src_end = src_start + first_len + dst_start = 0 + pending = [items[0]] + for item in items[1:]: + off, n = item[0], item[1] + gap = off - src_end + if gap <= slack_bytes: + src_end = off + n + pending.append(item) + else: + compact_total += flush(src_start, src_end, dst_start, pending) + src_start, src_end, dst_start, pending = off, off + n, compact_total, [item] + compact_total += flush(src_start, src_end, dst_start, pending) + + all_leaf_names = tuple(spec.name for spec in block.pack.linears) + fully_streamed = streamed == frozenset(all_leaf_names) + frozen_leaf_specs = MappingProxyType( + { + name: MappingProxyType(dict(roles)) + for name, roles in leaf_specs.items() + } + ) + + fp_source = repr( + ( + block.block_key, + tuple(sorted(streamed)), + tuple((r.src_offset, r.dst_offset, r.nbytes) for r in ranges), + tuple( + ( + name, + tuple( + (role, spec.dst_offset, spec.nbytes, str(spec.dtype), spec.shape) + for role, spec in sorted(roles.items()) + ), + ) + for name, roles in sorted(frozen_leaf_specs.items()) + ), + ) + ) + fingerprint = hashlib.sha1(fp_source.encode("utf-8")).hexdigest()[:16] + + return BlockTransferPlan( + block_key=block.block_key, + ranges=tuple(ranges), + compact_nbytes=compact_total, + leaf_specs=frozen_leaf_specs, + streamed_leaf_names=tuple(sorted(streamed)), + fully_streamed=fully_streamed, + fingerprint=fingerprint, + ) diff --git a/toolkit/quantization/__init__.py b/toolkit/quantization/__init__.py new file mode 100644 index 0000000000..8eb944a32d --- /dev/null +++ b/toolkit/quantization/__init__.py @@ -0,0 +1 @@ +"""Quantized tensor execution policies independent of memory management.""" diff --git a/toolkit/quantization/fp8_linear.py b/toolkit/quantization/fp8_linear.py new file mode 100644 index 0000000000..0d1a72f55e --- /dev/null +++ b/toolkit/quantization/fp8_linear.py @@ -0,0 +1,752 @@ +"""Row-wise FP8 Linear execution independent of storage and residency. + +This module owns native qualification, activation quantization, scaled-matmul +execution, materializing fallback, and training grad-input policy. Callers +supply explicit qdata, row scale, and bias tensors; no arena or memory-manager +state is inspected here. +""" + +from __future__ import annotations + +import os +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from .fp8_transpose import column_major + + +FP8_STATS = { + "enabled": False, + "training_enabled": False, + "kernel_calls": 0, + "fallback_calls": 0, +} + +FP8_LINEAR_EXECUTION_KEY = "toolkit.quantization.fp8_linear" +DEFAULT_ACTIVATION_FP8_DTYPE = torch.float8_e4m3fn +NATIVE_SCALED_MM_FP8_DTYPE = torch.float8_e4m3fn + + +@dataclass(frozen=True) +class Fp8LinearSpec: + """Backend-neutral semantic description of an FP8 Linear weight.""" + + weight_dtype: torch.dtype + activation_dtype: torch.dtype + scale_granularity: str + scale_dtype: torch.dtype + has_zero_point: bool + weight_layout: str + execution_variant: str + + +@dataclass(frozen=True) +class Fp8LinearDeclaration: + spec: Fp8LinearSpec + qdata: torch.Tensor + scale: torch.Tensor + + @property + def tensors(self): + return (self.qdata, self.scale) + + +def fp8_execution_key(spec: Fp8LinearSpec) -> tuple: + return (FP8_LINEAR_EXECUTION_KEY, spec) + + +def fp8_spec_from_execution_key(execution_key) -> Fp8LinearSpec | None: + if ( + isinstance(execution_key, tuple) + and len(execution_key) == 2 + and execution_key[0] == FP8_LINEAR_EXECUTION_KEY + and isinstance(execution_key[1], Fp8LinearSpec) + ): + return execution_key[1] + return None + + +def _spec_for_payload(qdata, scale, granularity) -> Fp8LinearSpec: + return Fp8LinearSpec( + weight_dtype=qdata.dtype, + activation_dtype=DEFAULT_ACTIVATION_FP8_DTYPE, + scale_granularity=granularity, + scale_dtype=scale.dtype, + has_zero_point=False, + weight_layout="out_in", + execution_variant="scaled_mm_dynamic_activation", + ) + + +def _adapt_quanto_fp8(value) -> Fp8LinearDeclaration | None: + try: + from optimum.quanto.tensor.qbytes import QBytesTensor + except ImportError: + return None + if not isinstance(value, QBytesTensor): + return None + qtype = value.qtype + if not qtype.is_floating_point or qtype.bits != 8: + return None + axis = value.axis + if axis == 0: + granularity = "output_row" + elif axis is None: + granularity = "per_tensor" + elif axis == -1: + granularity = "input_column" + else: + return None + qdata, scale = value._data, value._scale + return Fp8LinearDeclaration( + _spec_for_payload(qdata, scale, granularity), + qdata, + scale, + ) + + +def _adapt_torchao_fp8(value) -> Fp8LinearDeclaration | None: + try: + from torchao.quantization import Float8Tensor + except ImportError: + return None + if not isinstance(value, Float8Tensor): + return None + qdata, scale = value.qdata, value.scale + block_size = tuple(value.block_size or ()) + if qdata.ndim == 2 and block_size == (1, qdata.shape[1]): + granularity = "output_row" + elif scale.numel() == 1 and block_size in ((), tuple(qdata.shape)): + granularity = "per_tensor" + else: + return None + return Fp8LinearDeclaration( + _spec_for_payload(qdata, scale, granularity), + qdata, + scale, + ) + + +def declare_fp8_linear(value) -> Fp8LinearDeclaration | None: + """Normalize supported backend wrappers without leaking them to callers.""" + value = value.data if isinstance(value, torch.nn.Parameter) else value + return _adapt_torchao_fp8(value) or _adapt_quanto_fp8(value) + +_FP8_GRAD_INPUT = os.environ.get("AI_TOOLKIT_FP8_GRAD_INPUT", "0").lower() not in ( + "0", "false", "no", "off", "", +) +_FP8_GRAD_VERIFIED = None +_REUSE_DEQUANT = os.environ.get("AI_TOOLKIT_REUSE_DEQUANT", "1").lower() not in ( + "0", "false", "no", "off", "", +) +_REUSE_VERIFIED = None + + +def set_fp8_grad_input_enabled(enabled: bool) -> None: + global _FP8_GRAD_INPUT, _FP8_GRAD_VERIFIED + if bool(enabled) and not _FP8_GRAD_INPUT: + _FP8_GRAD_VERIFIED = None + _FP8_GRAD_INPUT = bool(enabled) + + +def reference_dequantize_to(tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + try: + return tensor.dequantize(output_dtype=dtype) + except TypeError: + value = tensor.dequantize() + return value if value.dtype == dtype else value.to(dtype=dtype) + + +def _scale_view_shape(spec, qdata, scale): + if spec.scale_granularity == "output_row": + if scale.numel() != qdata.shape[0]: + raise ValueError("invalid_output_row_scale") + return [qdata.shape[0]] + [1] * (qdata.ndim - 1) + if spec.scale_granularity == "input_column": + if scale.numel() != qdata.shape[-1]: + raise ValueError("invalid_input_column_scale") + return [1] * (qdata.ndim - 1) + [qdata.shape[-1]] + if spec.scale_granularity == "per_tensor": + if scale.numel() != 1: + raise ValueError("invalid_per_tensor_scale") + return [1] * qdata.ndim + raise ValueError("unsupported_fp8_scale_granularity") + + +def materialize_fp8_weight(spec, qdata, scale, dtype): + if spec.has_zero_point or spec.weight_layout != "out_in": + raise ValueError("unsupported_fp8_materialization") + view_shape = _scale_view_shape(spec, qdata, scale) + return qdata.to(dtype) * scale.reshape(view_shape).to(dtype) + + +def dequantize_rowwise(qdata, scale, dtype): + spec = _spec_for_payload(qdata, scale, "output_row") + return materialize_fp8_weight(spec, qdata, scale, dtype) + + +def fast_dequantize_into(qweight, dest): + declaration = declare_fp8_linear(qweight) + if declaration is None or dest is None: + return None + spec = declaration.spec + qdata, scale = declaration.qdata, declaration.scale + if qdata.shape != dest.shape or spec.has_zero_point: + return None + try: + view_shape = _scale_view_shape(spec, qdata, scale) + except ValueError: + return None + dest.copy_(qdata) + dest.mul_(scale.reshape(view_shape).to(dest.dtype)) + return dest + + +def fast_dequantize(qweight, dtype): + global _REUSE_VERIFIED + if not _REUSE_DEQUANT or _REUSE_VERIFIED is False: + return None + if dtype not in (torch.bfloat16, torch.float16, torch.float32): + return None + declaration = declare_fp8_linear(qweight) + if declaration is None: + return None + qdata = declaration.qdata + dest = torch.empty(qdata.shape, dtype=dtype, device=qdata.device) + fast = fast_dequantize_into(qweight, dest) + if fast is None: + return None + if _REUSE_VERIFIED is None: + try: + reference = reference_dequantize_to(qweight, dtype) + ok = reference.shape == fast.shape and torch.allclose( + fast, reference, rtol=1e-2, atol=1e-2 + ) + except Exception: + ok = False + _REUSE_VERIFIED = bool(ok) + if not ok: + return None + return fast + + +def dequantize_to(tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + fast = fast_dequantize(tensor, dtype) + return fast if fast is not None else reference_dequantize_to(tensor, dtype) + + +def dequantize_into(qweight, dest): + global _REUSE_VERIFIED + if not _REUSE_DEQUANT or _REUSE_VERIFIED is False or dest is None: + return None + fast = fast_dequantize_into(qweight, dest) + if fast is None: + return None + if _REUSE_VERIFIED is None: + try: + reference = reference_dequantize_to(qweight, dest.dtype) + ok = reference.shape == fast.shape and torch.allclose( + fast, reference, rtol=1e-2, atol=1e-2 + ) + except Exception: + ok = False + _REUSE_VERIFIED = bool(ok) + if not ok: + return None + return fast + + +def native_variant_supported(spec) -> bool: + return bool( + spec.weight_dtype == NATIVE_SCALED_MM_FP8_DTYPE + and spec.activation_dtype == NATIVE_SCALED_MM_FP8_DTYPE + and spec.scale_granularity == "output_row" + and not spec.has_zero_point + and spec.weight_layout == "out_in" + and spec.execution_variant == "scaled_mm_dynamic_activation" + ) + + +def supports_native_scaled_mm(spec, qdata, scale, *, device=None) -> bool: + target = qdata.device if device is None else torch.device(device) + if not native_device_supported(target): + return False + return bool( + native_variant_supported(spec) + and qdata.dtype == spec.weight_dtype + and scale.dtype == spec.scale_dtype + and qdata.ndim == 2 + and scale.is_floating_point() + and scale.numel() == qdata.shape[0] + and qdata.shape[0] % 16 == 0 + and qdata.shape[1] % 16 == 0 + ) + + +def native_rowwise_qualifies(qdata, scale, *, device=None) -> bool: + spec = _spec_for_payload(qdata, scale, "output_row") + return supports_native_scaled_mm(spec, qdata, scale, device=device) + + +def native_device_supported(device) -> bool: + target = torch.device(device) + if target.type != "cuda" or not hasattr(torch, "_scaled_mm"): + return False + try: + return torch.cuda.get_device_capability(target) >= (8, 9) + except Exception: + return False + + +def weight_format_key(weight) -> str: + value = weight.data if isinstance(weight, torch.nn.Parameter) else weight + declaration = declare_fp8_linear(value) + if declaration is None: + return "other" + return ( + "rowwise_fp8" + if native_variant_supported(declaration.spec) + else "fp8" + ) + + +def fp8_sampling_qualifies(weight, *, device=None) -> bool: + declaration = declare_fp8_linear(weight) + return bool( + declaration is not None + and supports_native_scaled_mm( + declaration.spec, + declaration.qdata, + declaration.scale, + device=device, + ) + ) + + +def native_linear( + x, + qdata_t, + scale_row, + bias, + activation_dtype=DEFAULT_ACTIVATION_FP8_DTYPE, +): + shape = x.shape + x_2d = x.reshape(-1, shape[-1]) + info = torch.finfo(activation_dtype) + scale_x = torch.clamp( + x_2d.abs().amax().float() / info.max, + min=torch.finfo(torch.float32).tiny, + ) + x_fp8 = torch.clamp( + x_2d / scale_x.to(x_2d.dtype), min=info.min, max=info.max + ).to(activation_dtype) + one = torch.ones((), device=x.device, dtype=torch.float32) + out = torch._scaled_mm( + x_fp8, + qdata_t, + scale_a=scale_x, + scale_b=one, + out_dtype=x.dtype, + use_fast_accum=True, + ) + out = out * scale_row.reshape(1, -1).to(out.dtype) + if bias is not None: + out = out + bias.to(device=x.device, dtype=out.dtype) + return out.reshape(*shape[:-1], scale_row.shape[0]) + + +def materialized_linear(x, spec, qdata, scale, bias): + return F.linear(x, materialize_fp8_weight(spec, qdata, scale, x.dtype), bias) + + +def _grad_input_compute( + grad_out, + qdata, + scale, + target_dtype, + activation_dtype=DEFAULT_ACTIVATION_FP8_DTYPE, +): + try: + shape = grad_out.shape + grad = grad_out.reshape(-1, shape[-1]).to(torch.float32) + grad = grad * scale.reshape(1, -1).to(torch.float32) + info = torch.finfo(activation_dtype) + scale_grad = torch.clamp( + grad.abs().amax() / info.max, + min=torch.finfo(torch.float32).tiny, + ) + grad_fp8 = torch.clamp( + grad / scale_grad, min=info.min, max=info.max + ).to(activation_dtype) + one = torch.ones((), device=grad_out.device, dtype=torch.float32) + result = torch._scaled_mm( + grad_fp8, + column_major(qdata), + scale_a=scale_grad, + scale_b=one, + out_dtype=target_dtype, + use_fast_accum=True, + ) + return result.reshape(*shape[:-1], qdata.shape[1]) + except Exception: + return None + + +def grad_input_supported(qdata, scale, grad_out, spec=None) -> bool: + if not _FP8_GRAD_INPUT or _FP8_GRAD_VERIFIED is False: + return False + spec = spec or _spec_for_payload(qdata, scale, "output_row") + return bool( + supports_native_scaled_mm(spec, qdata, scale, device=grad_out.device) + and grad_out.dtype in (torch.bfloat16, torch.float16) + and grad_out.shape[-1] == qdata.shape[0] + ) + + +def grad_input_supported_weight(qweight, grad_out) -> bool: + declaration = declare_fp8_linear(qweight) + return bool( + declaration is not None + and grad_input_supported( + declaration.qdata, + declaration.scale, + grad_out, + declaration.spec, + ) + ) + + +def grad_input(grad_out, qweight, target_dtype): + declaration = declare_fp8_linear(qweight) + if declaration is None: + return None + spec = declaration.spec + qdata, scale = declaration.qdata, declaration.scale + global _FP8_GRAD_VERIFIED + out = ( + _grad_input_compute( + grad_out, + qdata, + scale, + target_dtype, + spec.activation_dtype, + ) + if grad_input_supported(qdata, scale, grad_out, spec) + else None + ) + if out is not None and _FP8_GRAD_VERIFIED is None: + try: + reference = grad_out.to(target_dtype) @ dequantize_to(qweight, target_dtype) + _FP8_GRAD_VERIFIED = bool( + torch.allclose(out, reference, rtol=2e-2, atol=2e-2) + ) + except Exception: + _FP8_GRAD_VERIFIED = False + if out is not None and _FP8_GRAD_VERIFIED: + return out + return grad_out.to(target_dtype) @ dequantize_to(qweight, target_dtype) + + +class _NativeTrainingFn(torch.autograd.Function): + @staticmethod + def forward(ctx, x, qdata_t, scale_row, bias, activation_dtype): + ctx.save_for_backward(qdata_t, scale_row) + ctx.input_dtype = x.dtype + ctx.fp8_grad_input = bool(_FP8_GRAD_INPUT) + ctx.activation_dtype = activation_dtype + return native_linear( + x, + qdata_t, + scale_row, + bias, + activation_dtype, + ) + + @staticmethod + def backward(ctx, grad_out): + qdata_t, scale_row = ctx.saved_tensors + qdata = qdata_t.t() + result = None + if ctx.fp8_grad_input: + result = _grad_input_compute( + grad_out, + qdata, + scale_row, + ctx.input_dtype, + ctx.activation_dtype, + ) + if result is None: + weight = dequantize_rowwise(qdata, scale_row, ctx.input_dtype) + result = grad_out.to(ctx.input_dtype) @ weight + return result.to(grad_out.dtype), None, None, None, None + + +def native_linear_training( + x, + qdata_t, + scale_row, + bias, + activation_dtype=DEFAULT_ACTIVATION_FP8_DTYPE, +): + return _NativeTrainingFn.apply( + x, + qdata_t, + scale_row, + bias, + activation_dtype, + ) + + +def fp8_linear_inference(x, weight, bias): + declaration = declare_fp8_linear(weight) + if ( + declaration is None + or x.dtype not in (torch.bfloat16, torch.float16) + or x.numel() == 0 + ): + FP8_STATS["fallback_calls"] += int( + FP8_STATS["enabled"] or FP8_STATS["training_enabled"] + ) + return None + spec = declaration.spec + qdata, scale = declaration.qdata, declaration.scale + if ( + not supports_native_scaled_mm(spec, qdata, scale, device=x.device) + or qdata.device != x.device + or scale.device != x.device + or (bias is not None and bias.device != x.device) + or x.shape[-1] != qdata.shape[1] + ): + FP8_STATS["fallback_calls"] += int( + FP8_STATS["enabled"] or FP8_STATS["training_enabled"] + ) + return None + try: + out = native_linear( + x, + qdata.t(), + scale, + bias, + spec.activation_dtype, + ) + except RuntimeError: + FP8_STATS["fallback_calls"] += int( + FP8_STATS["enabled"] or FP8_STATS["training_enabled"] + ) + return None + FP8_STATS["kernel_calls"] += int( + FP8_STATS["enabled"] or FP8_STATS["training_enabled"] + ) + return out + + +@dataclass(frozen=True) +class BoundFp8LinearOperation: + spec: Fp8LinearSpec + native: bool + bias_index: int | None = 2 + @property + def format_key(self): + return ( + "rowwise_fp8" + if native_variant_supported(self.spec) + else "fp8" + ) + + def explicit_tensors(self, weight, bias, scale): + values = [weight, scale] + if self.bias_index is not None: + values.append(bias) + return tuple(values) + + def _unpack(self, tensors): + qdata, scale = tensors[0], tensors[1] + bias = None if self.bias_index is None else tensors[self.bias_index] + return qdata, scale, bias + + def functional_components(self, tensors): + qdata, scale, bias = self._unpack(tensors) + return qdata, bias, scale + + def forward_sample(self, x, tensors): + qdata, scale, bias = self._unpack(tensors) + if self.native: + return native_linear( + x, + qdata.t(), + scale.reshape(-1), + bias, + self.spec.activation_dtype, + ) + return materialized_linear(x, self.spec, qdata, scale, bias) + + def forward_train(self, x, tensors): + qdata, scale, bias = self._unpack(tensors) + if self.native: + return native_linear_training( + x, + qdata.t(), + scale.reshape(-1), + bias, + self.spec.activation_dtype, + ) + return materialized_linear(x, self.spec, qdata, scale, bias) + + def forward_explicit(self, x, weight, bias, scale, *, training): + """Execute either the original FP8 tuple or a dense replacement.""" + if scale is None: + return BoundDenseLinearOperation()._forward(x, weight, bias) + tensors = self.explicit_tensors(weight, bias, scale) + forward = self.forward_train if training else self.forward_sample + return forward(x, tensors) + + def materialize(self, tensors, dtype=torch.bfloat16): + qdata, scale, _bias = self._unpack(tensors) + return materialize_fp8_weight(self.spec, qdata, scale, dtype) + + +def bind_fp8_linear( + spec, + qdata, + scale, + *, + device, + has_bias=True, +) -> BoundFp8LinearOperation: + return BoundFp8LinearOperation( + spec=spec, + native=supports_native_scaled_mm(spec, qdata, scale, device=device), + bias_index=2 if has_bias else None, + ) + + +def bind_rowwise_fp8(qdata, scale, *, device, has_bias=True) -> BoundFp8LinearOperation: + """Compatibility binder for an explicitly declared rowwise payload.""" + spec = _spec_for_payload(qdata, scale, "output_row") + return bind_fp8_linear( + spec, + qdata, + scale, + device=device, + has_bias=has_bias, + ) + + +@dataclass(frozen=True) +class BoundDenseLinearOperation: + format_key = "dense" + def _forward(self, x, weight, bias): + if weight.dtype != x.dtype and weight.dtype in ( + torch.float16, + torch.bfloat16, + torch.float32, + ): + weight = weight.to(dtype=x.dtype) + if bias is not None and bias.dtype != x.dtype: + bias = bias.to(dtype=x.dtype) + return F.linear(x, weight, bias) + + bias_index: int | None = 1 + + def explicit_tensors(self, weight, bias, _scale): + return (weight,) if self.bias_index is None else (weight, bias) + + def functional_components(self, tensors): + weight = tensors[0] + bias = None if self.bias_index is None else tensors[self.bias_index] + return weight, bias, None + + def forward_sample(self, x, tensors): + weight = tensors[0] + bias = None if self.bias_index is None else tensors[self.bias_index] + return self._forward(x, weight, bias) + + def forward_train(self, x, tensors): + weight = tensors[0] + bias = None if self.bias_index is None else tensors[self.bias_index] + return self._forward(x, weight, bias) + + def forward_explicit(self, x, weight, bias, scale, *, training): + del scale, training + return self._forward(x, weight, bias) + + def materialize(self, tensors, dtype=None): + weight = tensors[0] + return weight if dtype is None or weight.dtype == dtype else weight.to(dtype=dtype) + + +def bind_linear_operation(weight, bias=None, *, device): + value = weight.data if isinstance(weight, torch.nn.Parameter) else weight + declaration = declare_fp8_linear(value) + if declaration is None: + try: + value.__tensor_flatten__() + except Exception: + pass + else: + raise ValueError("unsupported_quantized_linear_operation") + return BoundDenseLinearOperation(bias_index=1 if bias is not None else None) + return bind_fp8_linear( + declaration.spec, + declaration.qdata, + declaration.scale, + device=device, + has_bias=bias is not None, + ) + + +def bind_parameter_operation(weight, bias=None, *, device): + """Bind an operation and snapshot its explicit ordered tensor tuple.""" + from .storage import linear_storage_binding + + binding = linear_storage_binding(weight, bias) + tensors = tuple(item.tensor for item in binding.tensors) + operation = bind_storage_operation( + tensors, + device=device, + weight_leaf_count=binding.weight_leaf_count, + execution_key=binding.execution_key, + ) + return operation, tensors + + +def bind_storage_operation( + tensors, + *, + device, + weight_leaf_count: int, + execution_key, +): + """Bind execution from an opaque tuple outside the memory manager.""" + tensors = tuple(tensors) + spec = fp8_spec_from_execution_key(execution_key) + if spec is not None: + if int(weight_leaf_count) != 2 or len(tensors) not in (2, 3): + raise ValueError("invalid_fp8_linear_storage") + return bind_fp8_linear( + spec, + tensors[0], + tensors[1], + device=device, + has_bias=len(tensors) > 2, + ) + dense_declaration = False + try: + declared_weight_leaves = tuple(execution_key[2]) + dense_declaration = ( + len(declared_weight_leaves) == 1 + and declared_weight_leaves[0][0] == "weight" + ) + except (IndexError, TypeError): + pass + if int(weight_leaf_count) == 1 and dense_declaration: + return BoundDenseLinearOperation(bias_index=1 if len(tensors) > 1 else None) + raise ValueError("unsupported_linear_storage_operation") + + +# Transitional names for callers migrating from manager_modules. +_fp8_linear_compiled = native_linear +_fp8_linear_training = native_linear_training +_fp8_grad_input_compute = _grad_input_compute diff --git a/toolkit/quantization/fp8_transpose.py b/toolkit/quantization/fp8_transpose.py new file mode 100644 index 0000000000..b74806a2ed --- /dev/null +++ b/toolkit/quantization/fp8_transpose.py @@ -0,0 +1,87 @@ +"""Tiled column-major conversion for one-byte FP8 matrix operands.""" + +from __future__ import annotations + +import torch + +try: + import triton + import triton.language as tl + + @triton.jit + def _fp8_transpose_kernel( + src_ptr, + dst_ptr, + m, + n, + stride_sm, + stride_sn, + stride_dm, + stride_dn, + BLOCK: tl.constexpr, + ): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + rm = pid_m * BLOCK + tl.arange(0, BLOCK) + rn = pid_n * BLOCK + tl.arange(0, BLOCK) + mask = (rm[:, None] < m) & (rn[None, :] < n) + tile = tl.load( + src_ptr + rm[:, None] * stride_sm + rn[None, :] * stride_sn, + mask=mask, + other=0, + ) + tl.store( + dst_ptr + rn[:, None] * stride_dm + rm[None, :] * stride_dn, + tl.trans(tile), + mask=tl.trans(mask), + ) + + _HAVE_TRITON = True +except ImportError: # pragma: no cover - triton-less installs only + _HAVE_TRITON = False + + +_BLOCK = 64 + + +def _tiled_supported(x: torch.Tensor) -> bool: + return ( + _HAVE_TRITON + and x.device.type == "cuda" + and x.ndim == 2 + and x.element_size() == 1 + and x.numel() > 0 + ) + + +@torch.library.custom_op("mm::transpose_contiguous_1byte", mutates_args=()) +def transpose_contiguous_1byte(x: torch.Tensor) -> torch.Tensor: + """Return the contiguous transpose of a two-dimensional one-byte tensor.""" + if not _tiled_supported(x): + return x.t().contiguous() + src = x.view(torch.uint8) + m, n = src.shape + dst = torch.empty((n, m), device=src.device, dtype=torch.uint8) + grid = (triton.cdiv(m, _BLOCK), triton.cdiv(n, _BLOCK)) + _fp8_transpose_kernel[grid]( + src, + dst, + m, + n, + src.stride(0), + src.stride(1), + dst.stride(0), + dst.stride(1), + BLOCK=_BLOCK, + ) + return dst.view(x.dtype) + + +@transpose_contiguous_1byte.register_fake +def _(x: torch.Tensor) -> torch.Tensor: + return torch.empty((x.shape[1], x.shape[0]), dtype=x.dtype, device=x.device) + + +def column_major(x: torch.Tensor) -> torch.Tensor: + """Return ``x`` as a column-major operand for ``torch._scaled_mm``.""" + return torch.ops.mm.transpose_contiguous_1byte(x).t() diff --git a/toolkit/quantization/storage.py b/toolkit/quantization/storage.py new file mode 100644 index 0000000000..f4402e2005 --- /dev/null +++ b/toolkit/quantization/storage.py @@ -0,0 +1,206 @@ +"""Opaque physical storage plus quantization-owned state substitution.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +import torch + + +@dataclass(frozen=True) +class TensorStorageBinding: + name: str + tensor: torch.Tensor + + +@dataclass(frozen=True) +class StateSubstitutionBinding: + """One module-state target reconstructed from declared physical tensors.""" + + name: str + tensor_indices: tuple[int, ...] + reconstruct: Callable[[tuple[torch.Tensor, ...]], torch.Tensor] + + +@dataclass(frozen=True) +class _TensorSubclassReconstruction: + cls: type + names: tuple[str, ...] + context: object + size: tuple[int, ...] + stride: tuple[int, ...] + + def __call__(self, tensors: tuple[torch.Tensor, ...]) -> torch.Tensor: + return self.cls.__tensor_unflatten__( + dict(zip(self.names, tensors, strict=True)), + self.context, + torch.Size(self.size), + self.stride, + ) + + +def _identity(tensors: tuple[torch.Tensor, ...]) -> torch.Tensor: + if len(tensors) != 1: + raise ValueError("storage_identity_requires_one_tensor") + return tensors[0] + + +@dataclass(frozen=True) +class LayerStorageBinding: + tensors: tuple[TensorStorageBinding, ...] + execution_key: tuple + weight_leaf_count: int + weight_template: torch.Tensor | None + substitutions: tuple[StateSubstitutionBinding, ...] + + +def _flatten_named(value, prefix): + try: + names, _context = value.__tensor_flatten__() + except Exception: + return [TensorStorageBinding(prefix, value)] + out = [] + for name in names: + child = getattr(value, name, None) + if child is not None: + child_prefix = f"{prefix}.{name}" if prefix else str(name) + out.extend(_flatten_named(child, child_prefix)) + return out + + +def named_tensor_storage(value, prefix: str = "") -> tuple[TensorStorageBinding, ...]: + """Return physical tensor leaves with stable dotted source names. + + State-dict consumers can pass the state key as ``prefix`` so tensor + subclasses such as ``weight._data`` / ``weight._scale`` are addressable + without retaining a second flattened value table. + """ + return tuple(_flatten_named(value, prefix)) + + +def linear_storage_binding(weight, bias=None) -> LayerStorageBinding: + """Describe physical storage without exposing its meaning to the mover.""" + weight_value = weight.data if isinstance(weight, torch.nn.Parameter) else weight + bias_value = bias.data if isinstance(bias, torch.nn.Parameter) else bias + from .fp8_linear import declare_fp8_linear, fp8_execution_key + + fp8 = declare_fp8_linear(weight_value) + if fp8 is None: + weight_tensors = _flatten_named(weight_value, "weight") + execution_key = ( + type(weight_value).__module__, + type(weight_value).__qualname__, + tuple( + (item.name, tuple(item.tensor.shape), str(item.tensor.dtype)) + for item in weight_tensors + ), + ) + weight_template = weight_value + else: + weight_tensors = [ + TensorStorageBinding("qdata", fp8.qdata), + TensorStorageBinding("scale", fp8.scale), + ] + execution_key = fp8_execution_key(fp8.spec) + weight_template = weight_value + bias_tensors = [] if bias_value is None else _flatten_named(bias_value, "bias") + tensors = tuple(weight_tensors + bias_tensors) + if len(weight_tensors) == 1: + weight_reconstruct = _identity + else: + names, context = weight_value.__tensor_flatten__() + weight_reconstruct = _TensorSubclassReconstruction( + cls=type(weight_value), + names=tuple(names), + context=context, + size=tuple(weight_value.shape), + stride=tuple(weight_value.stride()), + ) + substitutions = [ + StateSubstitutionBinding( + name="weight", + tensor_indices=tuple(range(len(weight_tensors))), + reconstruct=weight_reconstruct, + ) + ] + if bias_value is not None: + substitutions.append( + StateSubstitutionBinding( + name="bias", + tensor_indices=tuple(range(len(weight_tensors), len(tensors))), + reconstruct=_identity, + ) + ) + return LayerStorageBinding( + tensors=tensors, + execution_key=execution_key, + weight_leaf_count=len(weight_tensors), + weight_template=weight_template, + substitutions=tuple(substitutions), + ) + + +def module_storage_binding(module: torch.nn.Module) -> LayerStorageBinding: + """Declare storage for a supported Linear without materializing weights.""" + from toolkit.util.ostris_quant import OstrisLinear + + if not isinstance(module, OstrisLinear): + weight = module._parameters.get("weight") + if weight is None: + raise ValueError( + f"unsupported_managed_module:{type(module).__module__}." + f"{type(module).__qualname__}" + ) + return linear_storage_binding(weight, module._parameters.get("bias")) + + tensors = [] + substitutions = [] + for name, value in module._buffers.items(): + if value is None: + continue + # Ostris quantizers intentionally keep packed execution buffers out of + # state_dict and serialize them through their own cache format. Buffer + # persistence is a serialization policy, not an execution-storage + # policy: every live quantizer buffer is canonical arena state. + index = len(tensors) + tensors.append(TensorStorageBinding(f"buffer.{name}", value)) + substitutions.append( + StateSubstitutionBinding(name, (index,), _identity) + ) + bias = module._parameters.get("bias") + if bias is not None: + index = len(tensors) + tensors.append(TensorStorageBinding("bias", bias.data)) + substitutions.append( + StateSubstitutionBinding("bias", (index,), _identity) + ) + if not tensors: + raise ValueError("ostris_linear_has_no_declared_storage") + quantizer = module.ostris_quantizer + execution_key = ( + "ostris", + type(quantizer).__module__, + type(quantizer).__qualname__, + str(getattr(quantizer, "qtype", None)), + tuple( + (item.name, tuple(item.tensor.shape), str(item.tensor.dtype)) + for item in tensors + ), + ) + return LayerStorageBinding( + tensors=tuple(tensors), + execution_key=execution_key, + weight_leaf_count=0, + weight_template=None, + substitutions=tuple(substitutions), + ) + + +def temporary_materialization_bytes(value, dtype=torch.bfloat16) -> int: + """Return scratch bytes needed when wrapped storage must be materialized.""" + value = value.data if isinstance(value, torch.nn.Parameter) else value + binding = linear_storage_binding(value) + if binding.weight_leaf_count <= 1: + return 0 + return int(value.numel() * torch.empty((), dtype=dtype).element_size()) diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index a060a599e6..1aa603bdbb 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -1,16 +1,19 @@ from fnmatch import fnmatch from typing import List, Optional, Union, TYPE_CHECKING import torch +from torchao.quantization import Float8Tensor from optimum.quanto.quantize import _quantize_submodule from optimum.quanto.tensor import Optimizer, qtype, qtypes from torchao.quantization.quant_api import ( quantize_ as torchao_quantize_, + _is_linear as torchao_is_linear, Float8WeightOnlyConfig, - UIntXWeightOnlyConfig, + IntxWeightOnlyConfig, Int8WeightOnlyConfig ) from optimum.quanto import freeze +from optimum.quanto.tensor.qbytes import QBytesTensor from tqdm import tqdm from safetensors.torch import load_file from huggingface_hub import hf_hub_download @@ -27,6 +30,37 @@ if TYPE_CHECKING: from toolkit.models.base_model import BaseModel + +def tensor_subclass_leaves(value: torch.Tensor) -> list[torch.Tensor]: + try: + names, _context = value.__tensor_flatten__() + except Exception: + return [value] + leaves = [] + for name in names: + inner = getattr(value, name) + if inner is not None: + leaves.extend(tensor_subclass_leaves(inner)) + return leaves + + +def _tensor_subclass_to_meta(value: torch.Tensor) -> torch.Tensor: + try: + names, context = value.__tensor_flatten__() + except Exception: + return value.to(device="meta") + inner = { + name: ( + None + if getattr(value, name) is None + else _tensor_subclass_to_meta(getattr(value, name)) + ) + for name in names + } + return type(value).__tensor_unflatten__( + inner, context, value.size(), value.stride() + ) + # the quantize function in quanto had a bug where it was using exclude instead of include Q_MODULES = [ @@ -41,14 +75,13 @@ ] torchao_qtypes = { - # "int4": Int4WeightOnlyConfig(), - "uint2": UIntXWeightOnlyConfig(torch.uint2), - "uint3": UIntXWeightOnlyConfig(torch.uint3), - "uint4": UIntXWeightOnlyConfig(torch.uint4), - "uint5": UIntXWeightOnlyConfig(torch.uint5), - "uint6": UIntXWeightOnlyConfig(torch.uint6), - "uint7": UIntXWeightOnlyConfig(torch.uint7), - "uint8": UIntXWeightOnlyConfig(torch.uint8), + "uint2": IntxWeightOnlyConfig(torch.int2), + "uint3": IntxWeightOnlyConfig(torch.int3), + "uint4": IntxWeightOnlyConfig(torch.int4), + "uint5": IntxWeightOnlyConfig(torch.int5), + "uint6": IntxWeightOnlyConfig(torch.int6), + "uint7": IntxWeightOnlyConfig(torch.int7), + "uint8": Int8WeightOnlyConfig(), "int8": Int8WeightOnlyConfig(), "float8": Float8WeightOnlyConfig(), } @@ -169,6 +202,25 @@ def quantize( include = [include] if isinstance(include, str) else include if exclude is not None: exclude = [exclude] if isinstance(exclude, str) else exclude + + if isinstance(weights, aotype): + # TorchAO quantize_ already walks the entire module tree. Calling it for + # every item yielded by named_modules() quantizes children once through + # their parent and then attempts to quantize them again directly. + def filter_fn(module: torch.nn.Module, fqn: str) -> bool: + if not torchao_is_linear(module, fqn): + return False + if isinstance(module.weight, Float8Tensor): + return False + if include is not None and not any(fnmatch(fqn, pattern) for pattern in include): + return False + if exclude is not None and any(fnmatch(fqn, pattern) for pattern in exclude): + return False + return True + + torchao_quantize_(model, weights.config, filter_fn=filter_fn) + return + for name, m in model.named_modules(): if include is not None and not any( fnmatch(name, pattern) for pattern in include @@ -204,8 +256,6 @@ def quantize( if isinstance(weights, ostristype): if isinstance(m, torch.nn.Linear): convert_linear_to_ostris(m, weights.quantizer) - elif isinstance(weights, aotype): - torchao_quantize_(m, weights.config) else: _quantize_submodule( model, @@ -224,6 +274,181 @@ def quantize( # raise e +def quantize_module_at_path( + model: torch.nn.Module, + name: str, + *, + weights, +) -> torch.nn.Module: + """Quantize one already-materialized module and publish its replacement. + + Recursive Quanto quantization cannot replace the root module: its relative + name is empty, so the generated QLinear is assigned to an unusable empty + attribute while the original Linear's weight is cleared. Bounded loaders + know the real parent path and use this helper for root quantization units. + """ + module = model.get_submodule(name) + resolved = get_qtype(weights) + if isinstance(resolved, aotype): + quantize(module, weights=resolved) + elif isinstance(resolved, ostristype): + if isinstance(module, torch.nn.Linear): + # Shape-ineligible Ostris roots remain dense, matching recursive + # quantize() behavior for children (for example Krea's 12-wide + # text-fusion projector). + convert_linear_to_ostris(module, resolved.quantizer) + else: + quantize(module, weights=resolved) + else: + _quantize_submodule( + model, + name, + module, + weights=resolved, + ) + return model.get_submodule(name) + + +def assign_quantized_state_dict( + model: torch.nn.Module, + state_dict: dict, + weights, +) -> None: + """Assign cached quantized state, using a generic active arena session.""" + prepare_quantized_state_dict_model(model, state_dict, weights) + from toolkit.memory_management.arena_offload.load_session import ( + try_prepare_canonical_from_state_dict, + ) + + if try_prepare_canonical_from_state_dict(model, state_dict) is not None: + try: + assign_quantized_state_dict_subset(model, state_dict, weights) + return + except BaseException: + from toolkit.memory_management.arena_offload.load_session import ( + discard_pending_canonical_build, + ) + + discard_pending_canonical_build(model) + raise + missing, unexpected = model.load_state_dict(state_dict, strict=True, assign=True) + if missing or unexpected: + raise RuntimeError(f"missing={missing[:5]} unexpected={unexpected[:5]}") + model.requires_grad_(False) + + +def prepare_quantized_state_dict_model( + model: torch.nn.Module, + state_dict: dict, + weights, +) -> None: + """Reconstruct cached quantized wrappers with meta storage only.""" + resolved = get_qtype(weights) + quanto_data_suffix = ".weight._data" + quanto_prefixes = [ + key[: -len(quanto_data_suffix)] + for key in state_dict + if key.endswith(quanto_data_suffix) + ] + if quanto_prefixes: + if isinstance(resolved, (aotype, ostristype)): + raise ValueError("cached_quanto_state_qtype_mismatch") + quantize(model, weights=resolved) + modules = dict(model.named_modules()) + for prefix in quanto_prefixes: + module = modules[prefix] + data = state_dict[f"{prefix}.weight._data"] + scale = state_dict[f"{prefix}.weight._scale"] + template = module.weight + meta_data = torch.empty_like(data, device="meta") + meta_scale = torch.empty_like(scale, device="meta") + wrapper = QBytesTensor( + resolved, + 0, + template.size(), + template.stride(), + meta_data, + meta_scale, + requires_grad=False, + ) + module.weight = torch.nn.Parameter(wrapper, requires_grad=False) + bias = getattr(module, "bias", None) + if bias is not None: + bias.requires_grad_(False) + else: + modules = dict(model.named_modules()) + for key, value in state_dict.items(): + if not key.endswith(".weight") or len(tensor_subclass_leaves(value)) == 1: + continue + prefix = key[: -len(".weight")] + module = modules.get(prefix) + if module is None or not hasattr(module, "weight"): + continue + module.weight = torch.nn.Parameter( + _tensor_subclass_to_meta(value), requires_grad=False + ) + + +def assign_quantized_state_dict_subset( + model: torch.nn.Module, + state_dict: dict, + weights, + *, + excluded_keys=(), +) -> None: + """Assign cached values except leaves owned by a direct arena destination.""" + excluded = set(excluded_keys) + prepare_quantized_state_dict_model(model, state_dict, weights) + resolved = get_qtype(weights) + data_suffix = ".weight._data" + prefixes = { + key[: -len(data_suffix)] + for key in state_dict + if key.endswith(data_suffix) + } + handled = set() + modules = dict(model.named_modules()) + for prefix in prefixes: + data_key = f"{prefix}.weight._data" + scale_key = f"{prefix}.weight._scale" + if data_key in excluded and scale_key in excluded: + continue + if data_key in excluded or scale_key in excluded: + raise ValueError(f"partial_cached_quantized_weight:{prefix}") + module = modules[prefix] + template = module.weight + wrapper = QBytesTensor( + resolved, + 0, + template.size(), + template.stride(), + state_dict[data_key], + state_dict[scale_key], + requires_grad=False, + ) + module.weight = torch.nn.Parameter(wrapper, requires_grad=False) + handled.update((data_key, scale_key)) + + for key, value in state_dict.items(): + if key in excluded or key in handled: + continue + *parents, leaf = key.split(".") + target = model + for component in parents: + target = getattr(target, component) + if leaf in target._parameters: + old = target._parameters[leaf] + requires_grad = bool(old.requires_grad) if old is not None else False + target._parameters[leaf] = torch.nn.Parameter( + value, requires_grad=requires_grad + ) + elif leaf in target._buffers: + target._buffers[leaf] = value + else: + raise KeyError(f"unsupported_cached_state_leaf:{key}") + model.requires_grad_(False) + + def quantize_model( base_model: "BaseModel", model_to_quantize: torch.nn.Module, From f7131aefed9d9aefda5543ceb9ef37718dea17a6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Tue, 14 Jul 2026 21:56:17 +0200 Subject: [PATCH 03/20] Wire arena offload through shared training lifecycle --- jobs/BaseJob.py | 15 +- jobs/process/BaseProcess.py | 4 + jobs/process/BaseSDTrainProcess.py | 145 ++++++++++++- run.py | 5 +- tests/test_arena_lifecycle_contract.py | 286 +++++++++++++++++++++++++ tests/test_job_cleanup.py | 41 ++++ toolkit/config_modules.py | 39 ++++ toolkit/models/base_model.py | 42 +++- ui/src/app/jobs/new/SimpleJob.tsx | 30 ++- ui/src/app/jobs/new/utils.ts | 2 + ui/src/docs.tsx | 11 + ui/src/types.ts | 1 + 12 files changed, 592 insertions(+), 29 deletions(-) create mode 100644 tests/test_arena_lifecycle_contract.py create mode 100644 tests/test_job_cleanup.py diff --git a/jobs/BaseJob.py b/jobs/BaseJob.py index 5b339ebe6d..38001f453f 100644 --- a/jobs/BaseJob.py +++ b/jobs/BaseJob.py @@ -66,6 +66,15 @@ def load_processes(self, process_dict: dict): raise ValueError(f'config file is invalid. Unknown process type: {process["type"]}') def cleanup(self): - # if you implement this in child clas, - # be sure to call super().cleanup() LAST - del self + errors = [] + processes = list(getattr(self, "process", ())) + for process in reversed(processes): + try: + process.cleanup() + except Exception as error: + errors.append(f"{type(process).__name__}: {error}") + finally: + process.job = None + self.process = [] + if errors: + raise RuntimeError("job cleanup failed: " + "; ".join(errors)) diff --git a/jobs/process/BaseProcess.py b/jobs/process/BaseProcess.py index c58724c987..be95ded304 100644 --- a/jobs/process/BaseProcess.py +++ b/jobs/process/BaseProcess.py @@ -57,6 +57,10 @@ def run(self): # be sure to call super().run() first incase something is added here pass + def cleanup(self): + """Release resources owned by this process. Must be idempotent.""" + pass + def add_meta(self, additional_meta: OrderedDict): self.meta.update(additional_meta) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index beb36865cb..90d6f899eb 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -1,3 +1,4 @@ +import contextlib import copy import glob import inspect @@ -21,6 +22,11 @@ import torch.backends.cuda from huggingface_hub import HfApi, interpreter_login from toolkit.memory_management import MemoryManager +from toolkit.memory_management.runtime import ( + close_memory_runtime_preparation, + get_memory_runtime, + memory_runtime_owns_compile, +) from toolkit.basic import value_map from toolkit.clip_vision_adapter import ClipVisionAdapter @@ -196,6 +202,8 @@ def __init__(self, process_id: int, job, config: OrderedDict, custom_pipeline=No # to hold network if there is one self.network: Union[Network, None] = None + self._arena_runtime = None + self._cleanup_started = False self.adapter: Union[T2IAdapter, IPAdapter, ClipVisionAdapter, ReferenceAdapter, CustomAdapter, ControlNetModel, None] = None self.embedding: Union[Embedding, None] = None self.decorator: Union[Decorator, None] = None @@ -361,8 +369,16 @@ def sample(self, step=None, is_first=False): if self.adapter is not None and isinstance(self.adapter, CustomAdapter): self.adapter.is_sampling = True - # send to be generated - self.sd.generate_images(gen_img_config_list, sampler=sample_config.sampler) + arena_runtime = get_memory_runtime(getattr(self.sd, "unet", None)) + sampling_session = ( + arena_runtime.sampling_session() + if arena_runtime is not None + else contextlib.nullcontext() + ) + with sampling_session: + self.sd.generate_images( + gen_img_config_list, sampler=sample_config.sampler + ) if self.adapter is not None and isinstance(self.adapter, CustomAdapter): @@ -481,6 +497,45 @@ def clean_up_saves(self): def post_save_hook(self, save_path): # override in subclass pass + + def cleanup(self): + if self._cleanup_started: + return + self._cleanup_started = True + errors = [] + + def attempt(label, operation): + try: + operation() + except Exception as error: + errors.append(f"{label}: {type(error).__name__}: {error}") + + runtime = self._arena_runtime + sd = getattr(self, "sd", None) + if runtime is None and sd is not None: + runtime = get_memory_runtime(getattr(sd, "unet", None)) + if runtime is not None: + attempt("arena runtime", runtime.close) + self._arena_runtime = None + if sd is not None: + attempt( + "memory runtime preparation", + lambda: close_memory_runtime_preparation(sd), + ) + models = [getattr(sd, "unet", None)] + text_encoders = getattr(sd, "text_encoder", None) + if isinstance(text_encoders, (list, tuple)): + models.extend(text_encoders) + else: + models.append(text_encoders) + for model in models: + if model is not None: + attempt( + "legacy memory manager", + lambda model=model: MemoryManager.detach(model), + ) + if errors: + raise RuntimeError("; ".join(errors)) def done_hook(self): pass @@ -1579,6 +1634,20 @@ def run(self): ### HOOK ### self.hook_before_model_load() model_config_to_load = copy.deepcopy(self.model_config) + arena_requested = bool( + self.model_config.layer_offloading + and self.model_config.layer_offloading_smart + ) + if arena_requested: + if not self.train_config.gradient_checkpointing: + raise ValueError( + "arena offload training requires " + "train.gradient_checkpointing=true" + ) + # Models should load on CPU without selecting their legacy + # per-layer offloader. Generic arena attachment happens below. + model_config_to_load.layer_offloading = False + model_config_to_load.low_vram = True if self.is_fine_tuning or self.train_config.merge_network_on_save: # get the latest checkpoint @@ -1628,10 +1697,34 @@ def run(self): custom_pipeline=self.custom_pipeline, noise_scheduler=sampler, ) + self.sd.dataset_configs = self.dataset_configs + self.sd.train_config = self.train_config self.hook_after_sd_init_before_load() # run base sd process run - self.sd.load_model() + from toolkit.memory_management.arena_offload import model_load_arena_session + + with model_load_arena_session(self.sd, enabled=arena_requested): + self.sd.load_model() + + text_encoders = getattr(self.sd, "text_encoder", None) + if text_encoders is not None and not isinstance( + text_encoders, (list, tuple) + ): + text_encoders = (text_encoders,) + if ( + arena_requested + and text_encoders + and self.model_config.layer_offloading_text_encoder_percent > 0 + ): + for text_encoder in text_encoders: + MemoryManager.attach( + text_encoder, + self.device_torch, + offload_percent=( + self.model_config.layer_offloading_text_encoder_percent + ), + ) self.sd.add_after_sample_image_hook(self.sample_step_hook) @@ -1724,9 +1817,23 @@ def run(self): else: text_encoder.requires_grad_(False) text_encoder.eval() - unet.to(self.device_torch, dtype=dtype) unet.requires_grad_(False) unet.eval() + if arena_requested: + from toolkit.memory_management.arena_offload import ( + ArenaOffloadConfig, + prepare_arena_offload, + ) + + arena_runtime = prepare_arena_offload( + unet, + device=self.device_torch, + block_names=self.sd.get_transformer_block_names(), + config=ArenaOffloadConfig.from_model_config(self.model_config), + ) + arena_runtime.place_permanent_modules(self.device_torch, dtype) + else: + unet.to(self.device_torch, dtype=dtype) vae = vae.to(torch.device('cpu'), dtype=dtype) vae.requires_grad_(False) vae.eval() @@ -1981,6 +2088,11 @@ def run(self): self.setup_adapter() flush() + arena_runtime = get_memory_runtime(getattr(self.sd, "unet", None)) + if arena_runtime is not None: + arena_runtime.finalize(self.network) + self._arena_runtime = arena_runtime + ### HOOK ### params = self.hook_add_extra_train_params(params) self.params = params @@ -2120,6 +2232,9 @@ def run(self): compile_dynamic = getattr(self.model_config, 'compile_dynamic', True) compile_fullgraph = getattr(self.model_config, 'compile_fullgraph', False) block_compile = getattr(self.model_config, 'block_compile', False) + runtime_owns_block_compile = memory_runtime_owns_compile( + inner_unet_check + ) # quantized + offloaded unet is incompatible with fullgraph; force it off if is_unet_quantized and is_unet_offloaded and compile_fullgraph: @@ -2133,7 +2248,12 @@ def run(self): # ==================================================== # BLOCK COMPILE # ==================================================== - if block_compile: + if runtime_owns_block_compile and block_compile: + print_acc( + "Arena offload owns block compilation; " + "skipping trainer block compile." + ) + elif block_compile: BLOCK_LIST_ATTRS = self.sd.get_transformer_block_names() if BLOCK_LIST_ATTRS is None or len(BLOCK_LIST_ATTRS) == 0: @@ -2447,9 +2567,22 @@ def run(self): self.torch_profiler.start() did_oom = False loss_dict = None + arena_runtime = self._arena_runtime or get_memory_runtime( + getattr(self.sd, "unet", None) + ) + shape_key = MemoryManager.offload_shape_key_from_batch(batch_list) + execution_context = ( + arena_runtime.training_step( + shape_key=shape_key, + step_num=self.step_num, + ) + if arena_runtime is not None + else contextlib.nullcontext() + ) try: with self.accelerator.accumulate(self.modules_being_trained): - loss_dict = self.hook_train_loop(batch_list) + with execution_context: + loss_dict = self.hook_train_loop(batch_list) except torch.cuda.OutOfMemoryError: did_oom = True except RuntimeError as e: diff --git a/run.py b/run.py index 9238a83497..52ab3df98b 100644 --- a/run.py +++ b/run.py @@ -108,10 +108,10 @@ def main(): print_acc(f"Running {len(config_file_list)} job{'' if len(config_file_list) == 1 else 's'}") for config_file in config_file_list: + job = None try: job = get_job(config_file, args.name) job.run() - job.cleanup() jobs_completed += 1 except Exception as e: print_acc(f"Error running job: {e}") @@ -131,6 +131,9 @@ def main(): if not args.recover: print_end_message(jobs_completed, jobs_failed) raise e + finally: + if job is not None: + job.cleanup() if __name__ == '__main__': diff --git a/tests/test_arena_lifecycle_contract.py b/tests/test_arena_lifecycle_contract.py new file mode 100644 index 0000000000..9599f1170a --- /dev/null +++ b/tests/test_arena_lifecycle_contract.py @@ -0,0 +1,286 @@ +import ast +from pathlib import Path +from types import SimpleNamespace +from unittest import mock + +import pytest +import torch + +from toolkit.memory_management.arena_offload import ( + ArenaOffloadConfig, + ArenaSetupFatalError, + prepare_canonical_storage, + prepare_arena_offload, +) +from toolkit.memory_management.arena_offload.errors import ( + ArenaCleanupError, + is_fatal_arena_setup, + recover_allows_next_job, +) +from toolkit.memory_management.arena_offload.discovery import BlockDiscoveryError +from toolkit.memory_management.arena_offload.ownership import ( + acquire_process_owner, + active_process_owner, + release_process_owner, +) +from toolkit.memory_management.arena_offload.resources import ArenaRuntimeResources +from toolkit.memory_management.arena_offload.runtime import ArenaOffloadRuntime + +class _Block(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + + def forward(self, value): + return self.linear(value) + + +class _Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.blocks = torch.nn.ModuleList((_Block(), _Block())) + self.gradient_checkpointing = True + self._checkpoint_keep_last = 0 + + @property + def weight(self): + return self.blocks[0].linear.weight + + def forward(self, value): + for block in self.blocks: + value = block(value) + return value + + +def _frozen_linear(): + model = _Model() + model.requires_grad_(False) + return model + + +def test_process_owner_is_exclusive_sequential_and_stale_safe(): + first = acquire_process_owner("cpu") + try: + with pytest.raises(RuntimeError, match="arena_runtime_already_active"): + acquire_process_owner("cpu") + finally: + release_process_owner(first) + + second = acquire_process_owner("cpu") + try: + with pytest.raises(RuntimeError, match="owner_mismatch"): + release_process_owner(first) + assert active_process_owner() is second + finally: + release_process_owner(second) + + +def test_precommit_failure_preserves_original_classification_and_model(): + model = _frozen_linear() + original = model.weight + + with pytest.raises(BlockDiscoveryError, match="block_container_not_found"): + prepare_arena_offload( + model, + device="cpu", + block_names=("missing",), + config=ArenaOffloadConfig(enabled=True), + ) + + assert model.weight is original + assert active_process_owner() is None + assert not hasattr(model, "_arena_offload_runtime") + assert not hasattr(model, "_arena_offload_disposed") + + +def test_disabled_config_and_unsupported_architecture_fail_before_mutation(): + model = _frozen_linear() + original = model.weight + with pytest.raises(ValueError, match="arena_offload_not_enabled"): + prepare_arena_offload( + model, + device="cpu", + config=ArenaOffloadConfig(enabled=False), + ) + unsupported = torch.nn.Linear(4, 4) + unsupported.gradient_checkpointing = True + with pytest.raises(BlockDiscoveryError, match="no_repeated_block_container"): + prepare_arena_offload( + unsupported, + device="cpu", + config=ArenaOffloadConfig(enabled=True), + ) + assert model.weight is original + assert active_process_owner() is None + assert not hasattr(model, "_arena_offload_runtime") + assert not hasattr(model, "_arena_offload_disposed") + + +def test_direct_loader_rollback_releases_preparation_owner(): + model = _frozen_linear() + build = prepare_canonical_storage(model, block_names=("blocks",), device="cpu") + assert active_process_owner() is not None + build.rollback() + assert active_process_owner() is None + assert not hasattr(model, "_arena_offload_disposed") + + +def test_postcommit_failure_is_fatal_disposes_and_releases_owner(): + model = _frozen_linear() + original_error = RuntimeError("residency construction failed") + with mock.patch( + "toolkit.memory_management.arena_offload.runtime.build_training_plan", + side_effect=original_error, + ): + with pytest.raises(ArenaSetupFatalError) as caught: + prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=ArenaOffloadConfig(enabled=True), + ) + + wrapper = RuntimeError("wrapper") + wrapper.__cause__ = caught.value + assert caught.value.__cause__ is original_error + assert is_fatal_arena_setup(wrapper) + assert not recover_allows_next_job(wrapper, True) + assert recover_allows_next_job(RuntimeError("ordinary"), True) + assert active_process_owner() is None + assert model._arena_offload_disposed + with pytest.raises(RuntimeError, match="transformer_disposed"): + model(torch.randn(1, 4)) + with pytest.raises(RuntimeError, match="transformer_disposed"): + model.to("cpu") + + +@pytest.mark.parametrize( + "target", + ( + "toolkit.memory_management.arena_offload.runtime.ResidencyState", + "toolkit.memory_management.arena_offload.dispatcher.prepare_block_dispatcher_runtime", + ), +) +def test_postcommit_fault_boundaries_are_fatal_and_never_fall_back(target): + model = _frozen_linear() + failure = RuntimeError(f"injected:{target.rsplit('.', 1)[-1]}") + plan = { + "offload_ids": set(), + "protected_training_leaf_keys": frozenset(), + } + patches = [ + mock.patch( + "toolkit.memory_management.arena_offload.runtime.build_training_plan", + return_value=plan, + ), + mock.patch(target, side_effect=failure), + ] + with patches[0], patches[1]: + with pytest.raises(ArenaSetupFatalError) as caught: + prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=ArenaOffloadConfig(enabled=True), + ) + assert caught.value.__cause__ is failure + assert active_process_owner() is None + assert model._arena_offload_disposed + assert not hasattr(model, "_memory_manager") + + +def test_resource_release_continues_after_cleanup_error_and_is_idempotent(): + calls = [] + model = _frozen_linear() + resources = ArenaRuntimeResources(model, "cpu") + resources.acquire_process_owner() + resources.canonical_committed = True + resources.arena = SimpleNamespace( + unguard_whole_model_to=lambda _model: calls.append("unguard"), + release=lambda: calls.append("arena"), + ) + resources.residency = SimpleNamespace(clear=lambda: calls.append("residency")) + resources.executor = SimpleNamespace( + active_executions=0, + close=lambda: (_ for _ in ()).throw(RuntimeError("executor boom")), + ) + + with pytest.raises(ArenaCleanupError, match="executor boom"): + resources.release() + resources.release() + + assert calls == ["residency", "unguard", "arena"] + assert resources.released + assert resources.disposed + assert active_process_owner() is None + + +def test_transfer_cleanup_failure_retains_process_owner_until_retry(): + model = _frozen_linear() + resources = ArenaRuntimeResources(model, "cpu") + resources.acquire_process_owner() + token = resources.owner_token + + with mock.patch( + "toolkit.memory_management.arena_offload.transfer.release_fetch_runtime", + side_effect=RuntimeError("transfer cleanup failed"), + ): + with pytest.raises(ArenaCleanupError, match="transfer cleanup failed"): + resources.release() + + assert active_process_owner() is token + assert not resources.released + resources.release() + assert active_process_owner() is None + assert resources.released + + +def test_finalize_cap_failure_is_fatal_after_runtime_publication(): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._disposed = False + runtime._resources = SimpleNamespace(release=mock.Mock()) + failure = RuntimeError("allocator cap failed") + runtime._bind_training_cap = mock.Mock(side_effect=failure) + + with pytest.raises(ArenaSetupFatalError) as caught: + runtime.finalize() + + assert caught.value.__cause__ is failure + runtime._resources.release.assert_called_once_with() + + +def test_phase7_import_and_private_state_boundaries(): + root = Path(__file__).parents[1] + arena_sources = "\n".join( + path.read_text(encoding="utf-8") + for path in (root / "toolkit" / "memory_management" / "arena_offload").glob("*.py") + ) + assert "from ..manager import" not in arena_sources + assert "manager_modules" not in arena_sources + + for relative in ( + "jobs/process/BaseSDTrainProcess.py", + "extensions_built_in/sd_trainer/SDTrainer.py", + ): + source = (root / relative).read_text(encoding="utf-8") + tree = ast.parse(source) + assert not any( + isinstance(node, ast.Attribute) and node.attr.startswith("_mm_") + for node in ast.walk(tree) + ) + +def test_phase8_legacy_manager_has_no_arena_execution_bridge(): + root = Path(__file__).parents[1] + manager = ( + root / "toolkit" / "memory_management" / "manager.py" + ).read_text(encoding="utf-8") + for obsolete in ( + "attach_smart_training_immutable", + "smart_immutable", + "_mm_immutable_", + "_immutable_runtime", + "_mm_canonical_leaf", + "canonical_relief", + ): + assert obsolete not in manager diff --git a/tests/test_job_cleanup.py b/tests/test_job_cleanup.py new file mode 100644 index 0000000000..172fc4778d --- /dev/null +++ b/tests/test_job_cleanup.py @@ -0,0 +1,41 @@ +import pytest + +from jobs.BaseJob import BaseJob + + +class _CleanupProcess: + def __init__(self, name, calls, error=None): + self.name = name + self.calls = calls + self.error = error + self.job = object() + + def cleanup(self): + self.calls.append(self.name) + if self.error is not None: + raise self.error + + +def test_base_job_cleanup_releases_every_process_after_failure(): + calls = [] + job = BaseJob.__new__(BaseJob) + job.process = [ + _CleanupProcess("first", calls), + _CleanupProcess("second", calls, RuntimeError("boom")), + ] + processes = list(job.process) + + with pytest.raises(RuntimeError, match="_CleanupProcess: boom"): + job.cleanup() + + assert calls == ["second", "first"] + assert job.process == [] + assert all(process.job is None for process in processes) + + +def test_base_job_cleanup_is_safe_before_processes_are_loaded(): + job = BaseJob.__new__(BaseJob) + + job.cleanup() + + assert job.process == [] diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index ef5900226c..b7a65f5984 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -699,6 +699,45 @@ def __init__(self, **kwargs): # 0 is off and 1.0 is 100% of the layers self.layer_offloading_transformer_percent = kwargs.get("layer_offloading_transformer_percent", 1.0) self.layer_offloading_text_encoder_percent = kwargs.get("layer_offloading_text_encoder_percent", 1.0) + # Optional block-native arena backend. The existing per-layer backend + # remains the default when this is false. + self.layer_offloading_smart = kwargs.get("layer_offloading_smart", False) + self.layer_offloading_smart_working_reserve_gb = kwargs.get( + "layer_offloading_smart_working_reserve_gb", -1.0 + ) + self.layer_offloading_smart_wddm_margin_gb = kwargs.get( + "layer_offloading_smart_wddm_margin_gb", -1.0 + ) + self.layer_offloading_smart_wddm_hard_gb = kwargs.get( + "layer_offloading_smart_wddm_hard_gb", 1.0 + ) + self.layer_offloading_smart_sampling_working_reserve_gb = kwargs.get( + "layer_offloading_smart_sampling_working_reserve_gb", -1.0 + ) + self.layer_offloading_smart_sampling_wddm_margin_gb = kwargs.get( + "layer_offloading_smart_sampling_wddm_margin_gb", -1.0 + ) + self.layer_offloading_smart_sampling_wddm_hard_gb = kwargs.get( + "layer_offloading_smart_sampling_wddm_hard_gb", 1.0 + ) + self.layer_offloading_fp8_forward = kwargs.get( + "layer_offloading_fp8_forward", False + ) + self.layer_offloading_fp8_grad_input = kwargs.get( + "layer_offloading_fp8_grad_input", False + ) + self.layer_offloading_fp8_sampling = kwargs.get( + "layer_offloading_fp8_sampling", False + ) + self.layer_offloading_checkpoint_keep_last = kwargs.get( + "layer_offloading_checkpoint_keep_last", 0 + ) + self.layer_offloading_prefetch_depth = kwargs.get( + "layer_offloading_prefetch_depth", 3 + ) + self.layer_offloading_simulated_vram_gb = kwargs.get( + "layer_offloading_simulated_vram_gb", 0 + ) # can be used to load the extras like text encoder or vae from here # only setup for some models but will prevent having to download the te for diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 4caf3ae7a3..58d75e534e 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -1,3 +1,4 @@ +import contextlib import copy import gc import inspect @@ -657,14 +658,31 @@ def generate_images( unconditional_embeds = unconditional_embeds.to( self.device_torch, dtype=self.unet.dtype) - img = self.generate_single_image( - pipeline, - gen_config, - conditional_embeds, - unconditional_embeds, - generator, - extra, + from toolkit.memory_management.runtime import get_memory_runtime + + arena_runtime = get_memory_runtime(self.unet) + sampling_context = ( + arena_runtime.sampling_image( + shape_key=( + "sample", + int(gen_config.height), + int(gen_config.width), + bool(getattr(gen_config, "batch_cfg", False)), + ), + cold_working_bytes=3 * (1024 ** 3), + ) + if arena_runtime is not None + else contextlib.nullcontext() ) + with sampling_context: + img = self.generate_single_image( + pipeline, + gen_config, + conditional_embeds, + unconditional_embeds, + generator, + extra, + ) gen_config.save_image(img, i) gen_config.log_image(img, i) @@ -688,7 +706,15 @@ def generate_images( network.train() network.multiplier = start_multiplier - self.unet.to(self.device_torch, dtype=self.torch_dtype) + from toolkit.memory_management.runtime import get_memory_runtime + + arena_runtime = get_memory_runtime(self.unet) + if arena_runtime is not None: + arena_runtime.place_permanent_modules( + self.device_torch, self.torch_dtype + ) + else: + self.unet.to(self.device_torch, dtype=self.torch_dtype) if network.is_merged_in: network.merge_out(merge_multiplier) # self.tokenizer.to(original_device_dict['tokenizer']) diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index e5d51b75f1..e64e830cdb 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -360,18 +360,26 @@ export default function SimpleJob({ /> {jobConfig.config.process[0].model.layer_offloading && (
- - setJobConfig(value * 0.01, 'config.process[0].model.layer_offloading_transformer_percent') - } - min={0} - max={100} - step={1} + setJobConfig(value, 'config.process[0].model.layer_offloading_smart')} + docKey="model.layer_offloading_smart" /> + {!jobConfig.config.process[0].model.layer_offloading_smart && ( + + setJobConfig(value * 0.01, 'config.process[0].model.layer_offloading_transformer_percent') + } + min={0} + max={100} + step={1} + /> + )} ), }, + 'model.layer_offloading_smart': { + title: 'Automatic Arena Offloading', + description: ( + <> + Uses the generic block arena to plan transformer residency automatically. Frozen base weights stay in one + canonical pinned host representation and are streamed by execution block, while trainable adapter weights stay + on the normal training path. Gradient checkpointing is required. Turn this off to use percentage-based layer + offloading instead. + + ), + }, 'model.qie.match_target_res': { title: 'Match Target Res', description: ( diff --git a/ui/src/types.ts b/ui/src/types.ts index 07f070c5e9..073d7340a3 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -168,6 +168,7 @@ export interface ModelConfig { low_vram: boolean; model_kwargs: { [key: string]: any }; layer_offloading?: boolean; + layer_offloading_smart?: boolean; layer_offloading_transformer_percent?: number; layer_offloading_text_encoder_percent?: number; assistant_lora_path?: string; From 2eea80be0de362505f2da53a03fd86bf34dad8b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Tue, 14 Jul 2026 23:38:17 +0200 Subject: [PATCH 04/20] Complete arena lifecycle and compiled execution --- docs/ARENA_OFFLOAD_CONTRACT.md | 55 +++ extensions_built_in/sd_trainer/SDTrainer.py | 8 +- jobs/BaseJob.py | 6 +- jobs/process/BaseSDTrainProcess.py | 110 ++++-- run.py | 46 ++- tests/test_allocator_cap.py | 66 ++++ tests/test_arena_lifecycle_contract.py | 332 ++++++++++++++++-- tests/test_arena_offload_api.py | 14 +- tests/test_arena_offload_planner.py | 37 +- tests/test_arena_offload_policy.py | 60 +++- tests/test_generic_block_dispatcher.py | 35 +- tests/test_lora_compile_scalars.py | 132 +++++++ toolkit/compile_utils.py | 23 ++ toolkit/config_modules.py | 13 +- toolkit/lora_special.py | 2 +- toolkit/lycoris_special.py | 2 +- toolkit/memory_management/allocator_cap.py | 34 +- .../arena_offload/__init__.py | 2 + .../memory_management/arena_offload/api.py | 43 ++- .../arena_offload/dispatcher.py | 57 ++- .../arena_offload/planner.py | 104 ++++-- .../memory_management/arena_offload/policy.py | 46 ++- .../arena_offload/runtime.py | 152 ++++++-- .../arena_offload/transfer.py | 2 +- toolkit/memory_management/canonical_arena.py | 44 +-- toolkit/memory_management/manager.py | 2 +- toolkit/memory_management/vram_budget.py | 95 +---- toolkit/models/DoRA.py | 2 +- toolkit/models/base_model.py | 46 ++- toolkit/models/lokr.py | 4 +- toolkit/network_mixins.py | 19 +- toolkit/timer.py | 2 + 32 files changed, 1277 insertions(+), 318 deletions(-) create mode 100644 docs/ARENA_OFFLOAD_CONTRACT.md create mode 100644 tests/test_allocator_cap.py create mode 100644 tests/test_lora_compile_scalars.py create mode 100644 toolkit/compile_utils.py diff --git a/docs/ARENA_OFFLOAD_CONTRACT.md b/docs/ARENA_OFFLOAD_CONTRACT.md new file mode 100644 index 0000000000..250c5fd832 --- /dev/null +++ b/docs/ARENA_OFFLOAD_CONTRACT.md @@ -0,0 +1,55 @@ +# Arena offload capability contract + +Arena offload does not select support by model architecture or quantization +name. An explicitly selected model is accepted when its live module graph and +storage satisfy the following contracts; otherwise setup fails at the narrowest +known boundary with the unmet contract in the error. + +## Model contract + +- The transformer exposes one or more repeated `ModuleList` or `Sequential` + block containers, either by unambiguous discovery or by declarative container + paths from the model integration. +- Selected blocks keep their ordinary installed forwards. The dispatcher does + not reconstruct model dataflow or block arguments. +- A selected block receives at least one tensor leaf in its positional or + keyword argument structure and returns a tensor or a nested Python structure + containing a tensor leaf. +- Model-owned gradient checkpointing is enabled before canonical commit. + Intentionally uncheckpointed selected blocks remain resident. +- Canonical managed leaves are frozen base weights. Trainable adapters remain + ordinary state outside canonical storage. +- Every selected-block parameter and buffer is enumerable before commit. + Shared, parametrized, missing, or conflicting managed state fails before the + destructive boundary. + +## Quantization contract + +- Each managed linear exposes a `LayerStorageBinding` through + `module_storage_binding()`. +- The binding enumerates every physical tensor leaf without materializing a + dequantized weight, supplies stable execution metadata, and declares how live + parameter or buffer targets are reconstructed from those leaves. +- Construction and dispatcher finalization both verify that every substitution + target exists. Unknown tensor-subclass storage fails before canonical commit. +- Transfer and residency code treats declared leaves as opaque tensors. It does + not branch on qtype or quantization backend identity. + +## Maintainer validation matrix + +The upstream gate is the matrix, not an allowlist. Each selected production +model and quantization row must establish the relevant mechanisms below. + +| Gate | Required evidence | +| --- | --- | +| Model independence | At least two production transformer architectures use the same discovery, saved-forward dispatcher, checkpoint owner, and lifecycle APIs. | +| Quantization independence | Plain, TorchAO/Quanto tensor-subclass, and Ostris packed-buffer layouts pass declaration, canonical construction, resident/streamed execution, and teardown checks where available. | +| Training | Resident and streamed adapter training produce finite gradients; checkpoint backward performs the planned re-fetches. | +| Sampling transition | Train -> sample -> train preserves canonical storage and restores the training residency plan. | +| Structured ABI | Positional or keyword tensor inputs and tensor or nested tensor outputs preserve the original block result. | +| Sequential jobs | Arena ownership, legacy manager state, transfer runtime, device ring registry, and pin ledger return to baseline before the next job. | +| Save/resume | Supported quantized layouts dequantize/save and resume through their existing serialization path without arena-specific state. | + +Production runs remain maintainer- or user-launched. Focused tests should assert +mechanism and numerical parity; they must not replace this matrix with model or +qtype name checks. diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 6d19c649da..14aee62afb 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -40,6 +40,7 @@ from PIL import Image from torchvision.transforms import functional as TF from toolkit.basic import flush +from toolkit.memory_management.runtime import get_memory_runtime adapter_transforms = transforms.Compose([ @@ -242,7 +243,12 @@ def hook_before_train_loop(self): super().hook_before_train_loop() if self.is_caching_text_embeddings: # make sure model is on cpu for this part so we don't oom. - self.sd.unet.to('cpu') + arena_runtime = get_memory_runtime(self.sd.unet) + if arena_runtime is not None: + arena_runtime.park_residency_for_external_phase() + arena_runtime.place_permanent_modules('cpu') + else: + self.sd.unet.to('cpu') # cache unconditional embeds (blank prompt) with torch.no_grad(): diff --git a/jobs/BaseJob.py b/jobs/BaseJob.py index 38001f453f..016ed83318 100644 --- a/jobs/BaseJob.py +++ b/jobs/BaseJob.py @@ -73,8 +73,8 @@ def cleanup(self): process.cleanup() except Exception as error: errors.append(f"{type(process).__name__}: {error}") - finally: - process.job = None - self.process = [] if errors: raise RuntimeError("job cleanup failed: " + "; ".join(errors)) + for process in processes: + process.job = None + self.process = [] diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 90d6f899eb..be3fa49ce7 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -30,6 +30,7 @@ from toolkit.basic import value_map from toolkit.clip_vision_adapter import ClipVisionAdapter +from toolkit.compile_utils import configure_cuda_only_inductor from toolkit.custom_adapter import CustomAdapter from toolkit.data_loader import get_dataloader_from_datasets, trigger_dataloader_setup_epoch from toolkit.data_transfer_object.data_loader import FileItemDTO, DataLoaderBatchDTO @@ -203,7 +204,8 @@ def __init__(self, process_id: int, job, config: OrderedDict, custom_pipeline=No # to hold network if there is one self.network: Union[Network, None] = None self._arena_runtime = None - self._cleanup_started = False + self._cleanup_in_progress = False + self._cleanup_completed = False self.adapter: Union[T2IAdapter, IPAdapter, ClipVisionAdapter, ReferenceAdapter, CustomAdapter, ControlNetModel, None] = None self.embedding: Union[Embedding, None] = None self.decorator: Union[Decorator, None] = None @@ -499,9 +501,11 @@ def post_save_hook(self, save_path): pass def cleanup(self): - if self._cleanup_started: + if self._cleanup_completed: return - self._cleanup_started = True + if self._cleanup_in_progress: + raise RuntimeError("training process cleanup is already in progress") + self._cleanup_in_progress = True errors = [] def attempt(label, operation): @@ -510,32 +514,42 @@ def attempt(label, operation): except Exception as error: errors.append(f"{label}: {type(error).__name__}: {error}") - runtime = self._arena_runtime - sd = getattr(self, "sd", None) - if runtime is None and sd is not None: - runtime = get_memory_runtime(getattr(sd, "unet", None)) - if runtime is not None: - attempt("arena runtime", runtime.close) - self._arena_runtime = None - if sd is not None: - attempt( - "memory runtime preparation", - lambda: close_memory_runtime_preparation(sd), - ) - models = [getattr(sd, "unet", None)] - text_encoders = getattr(sd, "text_encoder", None) - if isinstance(text_encoders, (list, tuple)): - models.extend(text_encoders) - else: - models.append(text_encoders) - for model in models: - if model is not None: - attempt( - "legacy memory manager", - lambda model=model: MemoryManager.detach(model), + try: + runtime = self._arena_runtime + sd = getattr(self, "sd", None) + if runtime is None and sd is not None: + runtime = get_memory_runtime(getattr(sd, "unet", None)) + if runtime is not None: + try: + runtime.close() + except Exception as error: + errors.append( + f"arena runtime: {type(error).__name__}: {error}" ) - if errors: - raise RuntimeError("; ".join(errors)) + else: + self._arena_runtime = None + if sd is not None: + attempt( + "memory runtime preparation", + lambda: close_memory_runtime_preparation(sd), + ) + models = [getattr(sd, "unet", None)] + text_encoders = getattr(sd, "text_encoder", None) + if isinstance(text_encoders, (list, tuple)): + models.extend(text_encoders) + else: + models.append(text_encoders) + for model in models: + if model is not None: + attempt( + "legacy memory manager", + lambda model=model: MemoryManager.detach(model), + ) + if errors: + raise RuntimeError("; ".join(errors)) + self._cleanup_completed = True + finally: + self._cleanup_in_progress = False def done_hook(self): pass @@ -789,7 +803,13 @@ def prepare_accelerator(self): # # prepare all the models stuff for accelerator (hopefully we dont miss any) self.sd.vae = self.accelerator.prepare(self.sd.vae) if self.sd.unet is not None: - self.sd.unet = self.accelerator.prepare(self.sd.unet) + arena_runtime = get_memory_runtime(self.sd.unet) + # The arena runtime is the sole transformer placement and + # residency authority. Accelerate's model preparation starts + # with whole-model placement and may install its own wrappers, + # so do not give it a second ownership path. + if arena_runtime is None: + self.sd.unet = self.accelerator.prepare(self.sd.unet) # todo always tdo it? self.modules_being_trained.append(self.sd.unet) if self.sd.text_encoder is not None and self.train_config.train_text_encoder: @@ -1662,6 +1682,15 @@ def run(self): self.load_training_state_from_metadata(latest_save_path) ModelClass = get_model_class(self.model_config) + if arena_requested: + from toolkit.memory_management.arena_offload import ( + validate_arena_training_mode, + ) + + validate_arena_training_mode( + full_finetune=self.is_fine_tuning, + mutates_base_weights=self.train_config.merge_network_on_save, + ) # if the model class has get_train_scheduler static method if hasattr(ModelClass, 'get_train_scheduler'): sampler = ModelClass.get_train_scheduler() @@ -2204,7 +2233,10 @@ def run(self): compiled_refs = [] # (block_list, index, original_block) for rollback on failure try: inner_unet_check = unwrap_model(self.sd.unet) - is_unet_offloaded = hasattr(inner_unet_check, '_memory_manager') + is_unet_offloaded = ( + hasattr(inner_unet_check, '_memory_manager') + or get_memory_runtime(inner_unet_check) is not None + ) text_encoder = getattr(self.sd, "text_encoder", None) text_encoder_check = unwrap_model(text_encoder) if text_encoder is not None else None @@ -2220,7 +2252,7 @@ def run(self): user_set_cache_limit = cache_size_limit is not None if user_set_cache_limit: torch._dynamo.config.cache_size_limit = cache_size_limit - torch._dynamo.config.suppress_errors = False + configure_cuda_only_inductor() # torch 2.9 inductor bug: the new memory-coalescing tiling analysis # crashes on some dynamic-shape index expressions (sympy PowByNatural # "assert p >= 0", seen with Qwen Image). The analysis doesn't apply @@ -2248,10 +2280,10 @@ def run(self): # ==================================================== # BLOCK COMPILE # ==================================================== - if runtime_owns_block_compile and block_compile: + if runtime_owns_block_compile: print_acc( "Arena offload owns block compilation; " - "skipping trainer block compile." + "skipping trainer compile." ) elif block_compile: BLOCK_LIST_ATTRS = self.sd.get_transformer_block_names() @@ -2592,6 +2624,14 @@ def run(self): raise # not an OOM; surface real errors if did_oom: self.num_consecutive_oom += 1 + if arena_runtime is not None: + failure = arena_runtime.diagnostics().get( + 'last_failure_event' + ) + if failure is not None: + print_acc( + f"[ArenaOffload] training failure: {failure}" + ) if self.num_consecutive_oom > 3: raise RuntimeError("OOM during training step 3 times in a row, aborting training") optimizer.zero_grad(set_to_none=True) @@ -2790,6 +2830,10 @@ def run(self): repo_id=self.save_config.hf_repo_id, private=self.save_config.hf_private ) + # Deterministic teardown needs the model graph in order to close the + # arena and detach legacy managers (notably an offloaded text encoder). + # BaseJob.cleanup() becomes an idempotent no-op after this succeeds. + self.cleanup() del ( self.sd, unet, diff --git a/run.py b/run.py index 52ab3df98b..66fda7b581 100644 --- a/run.py +++ b/run.py @@ -40,6 +40,7 @@ from toolkit.job import get_job from toolkit.accelerator import get_accelerator from toolkit.print import print_acc, setup_log_to_file +from toolkit.memory_management.arena_offload.errors import recover_allows_next_job accelerator = get_accelerator() @@ -109,31 +110,58 @@ def main(): for config_file in config_file_list: job = None + failure = None + cleanup_failed = False try: job = get_job(config_file, args.name) job.run() - jobs_completed += 1 except Exception as e: + failure = e print_acc(f"Error running job: {e}") - jobs_failed += 1 try: job.process[0].on_error(e) except Exception as e2: print_acc(f"Error running on_error: {e2}") - if not args.recover: - print_end_message(jobs_completed, jobs_failed) - raise e + try: + e.add_note(f"on_error failed: {type(e2).__name__}: {e2}") + except AttributeError: + pass except KeyboardInterrupt as e: + failure = e try: job.process[0].on_error(e) except Exception as e2: print_acc(f"Error running on_error: {e2}") - if not args.recover: - print_end_message(jobs_completed, jobs_failed) - raise e + try: + e.add_note(f"on_error failed: {type(e2).__name__}: {e2}") + except AttributeError: + pass finally: if job is not None: - job.cleanup() + try: + job.cleanup() + except Exception as cleanup_error: + cleanup_failed = True + if failure is None: + failure = cleanup_error + print_acc(f"Error cleaning up job: {cleanup_error}") + else: + try: + failure.add_note( + "job cleanup failed: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + except AttributeError: + pass + + if failure is None: + jobs_completed += 1 + continue + + jobs_failed += 1 + if cleanup_failed or not recover_allows_next_job(failure, args.recover): + print_end_message(jobs_completed, jobs_failed) + raise failure if __name__ == '__main__': diff --git a/tests/test_allocator_cap.py b/tests/test_allocator_cap.py new file mode 100644 index 0000000000..8ddaab3e11 --- /dev/null +++ b/tests/test_allocator_cap.py @@ -0,0 +1,66 @@ +import pytest + +from toolkit.memory_management import allocator_cap + + +@pytest.fixture(autouse=True) +def _clear_applied_fractions(): + allocator_cap.APPLIED_FRACTIONS.clear() + yield + allocator_cap.APPLIED_FRACTIONS.clear() + + +def test_production_guard_removes_strict_cap_and_allows_spill(monkeypatch): + calls = [] + allocator_cap.APPLIED_FRACTIONS[0] = 0.5 + monkeypatch.setattr(allocator_cap.sys, "platform", "win32") + monkeypatch.setattr(allocator_cap.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(allocator_cap.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + allocator_cap.torch.cuda, + "set_per_process_memory_fraction", + lambda fraction, index: calls.append((fraction, index)), + ) + + result = allocator_cap.configure_wddm_allocator_guard( + "cuda", strict=False + ) + + assert result is None + assert calls == [(1.0, 0)] + assert 0 not in allocator_cap.APPLIED_FRACTIONS + + +def test_strict_development_guard_binds_allocator_cap(monkeypatch): + calls = [] + gib = 1024**3 + monkeypatch.setattr(allocator_cap.sys, "platform", "win32") + monkeypatch.setattr(allocator_cap.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(allocator_cap.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(allocator_cap.torch.cuda, "memory_reserved", lambda _i: 0) + monkeypatch.setattr( + allocator_cap.torch.cuda, + "set_per_process_memory_fraction", + lambda fraction, index: calls.append((fraction, index)), + ) + monkeypatch.setattr( + allocator_cap.vram_budget, "device_total_bytes", lambda _i: 12 * gib + ) + monkeypatch.setattr( + allocator_cap.vram_budget, + "real_device_total_bytes", + lambda _i: 12 * gib, + ) + monkeypatch.setattr( + allocator_cap.vram_budget, + "device_mem_info", + lambda _i: (12 * gib, 12 * gib), + ) + + result = allocator_cap.configure_wddm_allocator_guard( + "cuda", 1.0, strict=True + ) + + assert result == 11 / 12 + assert calls == [(11 / 12, 0)] + assert allocator_cap.APPLIED_FRACTIONS[0] == 11 / 12 diff --git a/tests/test_arena_lifecycle_contract.py b/tests/test_arena_lifecycle_contract.py index 9599f1170a..0ef89ee5df 100644 --- a/tests/test_arena_lifecycle_contract.py +++ b/tests/test_arena_lifecycle_contract.py @@ -25,6 +25,29 @@ ) from toolkit.memory_management.arena_offload.resources import ArenaRuntimeResources from toolkit.memory_management.arena_offload.runtime import ArenaOffloadRuntime +from toolkit.memory_management import pin_manager +from toolkit.memory_management.manager_modules import _DEVICE_STATE +from toolkit.memory_management.residency import ResidencyPlan +from toolkit.models.base_model import BaseModel +from jobs.process.BaseSDTrainProcess import BaseSDTrainProcess +from jobs.BaseJob import BaseJob + + +@pytest.fixture(autouse=True) +def _disable_cuda_host_registration(monkeypatch): + """Lifecycle ownership tests do not need a live CUDA pin registration.""" + def fake_commit(candidate, nbytes, kind, **_kwargs): + return SimpleNamespace( + tensor=candidate, + nbytes=int(nbytes), + kind=kind, + pinned=True, + mechanism="register", + ) + + monkeypatch.setattr(pin_manager, "pin_register_commit", fake_commit) + monkeypatch.setattr(pin_manager, "release", lambda _handle: None) + class _Block(torch.nn.Module): def __init__(self): @@ -125,14 +148,15 @@ def test_direct_loader_rollback_releases_preparation_owner(): assert not hasattr(model, "_arena_offload_disposed") -def test_postcommit_failure_is_fatal_disposes_and_releases_owner(): +def test_planning_failure_is_precommit_and_preserves_model(): model = _frozen_linear() + original = model.weight original_error = RuntimeError("residency construction failed") with mock.patch( "toolkit.memory_management.arena_offload.runtime.build_training_plan", side_effect=original_error, ): - with pytest.raises(ArenaSetupFatalError) as caught: + with pytest.raises(RuntimeError, match="residency construction failed") as caught: prepare_arena_offload( model, device="cpu", @@ -140,18 +164,10 @@ def test_postcommit_failure_is_fatal_disposes_and_releases_owner(): config=ArenaOffloadConfig(enabled=True), ) - wrapper = RuntimeError("wrapper") - wrapper.__cause__ = caught.value - assert caught.value.__cause__ is original_error - assert is_fatal_arena_setup(wrapper) - assert not recover_allows_next_job(wrapper, True) - assert recover_allows_next_job(RuntimeError("ordinary"), True) + assert caught.value is original_error assert active_process_owner() is None - assert model._arena_offload_disposed - with pytest.raises(RuntimeError, match="transformer_disposed"): - model(torch.randn(1, 4)) - with pytest.raises(RuntimeError, match="transformer_disposed"): - model.to("cpu") + assert model.weight is original + assert not hasattr(model, "_arena_offload_disposed") @pytest.mark.parametrize( @@ -167,6 +183,7 @@ def test_postcommit_fault_boundaries_are_fatal_and_never_fall_back(target): plan = { "offload_ids": set(), "protected_training_leaf_keys": frozenset(), + "fits": True, } patches = [ mock.patch( @@ -189,6 +206,16 @@ def test_postcommit_fault_boundaries_are_fatal_and_never_fall_back(target): assert not hasattr(model, "_memory_manager") +def test_fatal_setup_classification_blocks_recovery_through_wrappers(): + fatal = ArenaSetupFatalError("committed setup failed") + wrapper = RuntimeError("wrapper") + wrapper.__cause__ = fatal + + assert is_fatal_arena_setup(wrapper) + assert not recover_allows_next_job(wrapper, True) + assert recover_allows_next_job(RuntimeError("ordinary"), True) + + def test_resource_release_continues_after_cleanup_error_and_is_idempotent(): calls = [] model = _frozen_linear() @@ -250,14 +277,258 @@ def test_finalize_cap_failure_is_fatal_after_runtime_publication(): runtime._resources.release.assert_called_once_with() -def test_phase7_import_and_private_state_boundaries(): +def test_base_model_device_state_routes_arena_transformer_through_runtime(): + model = SimpleNamespace( + vae=mock.Mock(), + unet=mock.Mock(), + text_encoder=mock.Mock(), + adapter=None, + refiner_unet=None, + ) + runtime = mock.Mock() + state = { + "vae": {"training": False, "device": "cpu", "requires_grad": False}, + "unet": {"training": False, "device": "cpu", "requires_grad": False}, + "text_encoder": { + "training": False, + "device": "cpu", + "requires_grad": False, + }, + } + + with mock.patch( + "toolkit.memory_management.runtime.get_memory_runtime", + return_value=runtime, + ), mock.patch("toolkit.models.base_model.flush"): + BaseModel.set_device_state(model, state) + + runtime.place_permanent_modules.assert_called_once_with(torch.device("cpu")) + runtime.park_residency_for_external_phase.assert_called_once_with() + model.unet.to.assert_not_called() + model.unet.requires_grad_.assert_not_called() + + +def test_arena_external_phase_parks_and_restores_exact_training_plan(): + original = ResidencyPlan.build( + "train", (("blocks.0", "linear"), ("blocks.1", "linear")) + ) + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._disposed = False + runtime._device_state_parked_plan = None + runtime._residency = SimpleNamespace(plan=original) + runtime._executor = SimpleNamespace( + TRAIN="train", + finalized=True, + set_residency_plan=mock.Mock(), + ) + runtime._training_plan = original + + runtime.park_residency_for_external_phase() + parked = runtime._executor.set_residency_plan.call_args.args[0] + assert parked.phase == "train" + assert parked.resident_leaf_keys == frozenset() + + runtime.restore_residency_after_external_phase() + assert runtime._executor.set_residency_plan.call_args.args[0] is original + assert runtime._device_state_parked_plan is None + + +def test_base_model_restores_arena_only_after_text_encoder_offload(): + events = [] + text_encoder = mock.Mock() + text_encoder.to.side_effect = lambda device: events.append(("text", device)) + runtime = mock.Mock() + runtime.restore_residency_after_external_phase.side_effect = ( + lambda: events.append(("arena", "restore")) + ) + model = SimpleNamespace( + vae=mock.Mock(), + unet=mock.Mock(), + text_encoder=text_encoder, + adapter=None, + refiner_unet=None, + ) + state = { + "vae": {"training": False, "device": "cpu", "requires_grad": False}, + "unet": { + "training": True, + "device": torch.device("cuda"), + "requires_grad": False, + }, + "text_encoder": { + "training": False, + "device": "cpu", + "requires_grad": False, + }, + } + + with mock.patch( + "toolkit.memory_management.runtime.get_memory_runtime", + return_value=runtime, + ), mock.patch("toolkit.models.base_model.flush"): + BaseModel.set_device_state(model, state) + + assert events == [("text", "cpu"), ("arena", "restore")] + + +def test_base_model_text_cache_preset_activates_only_text_encoder(): + model = SimpleNamespace( + save_device_state=mock.Mock(), + set_device_state=mock.Mock(), + vae=mock.Mock(), + unet=mock.Mock(), + text_encoder=mock.Mock(), + adapter=None, + refiner_unet=None, + vae_device_torch=torch.device("cuda"), + device_torch=torch.device("cuda"), + te_device_torch=torch.device("cuda"), + ) + + BaseModel.set_device_state_preset(model, "cache_text_encoder") + + state = model.set_device_state.call_args.args[0] + assert state["vae"]["device"] == "cpu" + assert state["unet"]["device"] == "cpu" + assert state["text_encoder"]["device"] == torch.device("cuda") + + +def test_accelerator_preparation_skips_arena_managed_transformer(): + unet = mock.Mock() + accelerator = mock.Mock() + accelerator.prepare.side_effect = lambda value, **_kwargs: value + process = SimpleNamespace( + accelerator=accelerator, + sd=SimpleNamespace( + vae=mock.Mock(), + unet=unet, + text_encoder=None, + refiner_unet=None, + network=None, + ), + train_config=SimpleNamespace( + train_text_encoder=False, + train_refiner=False, + ), + modules_being_trained=[], + adapter=None, + optimizer=mock.Mock(), + lr_scheduler=None, + ) + + with mock.patch( + "jobs.process.BaseSDTrainProcess.get_memory_runtime", + return_value=mock.Mock(), + ): + BaseSDTrainProcess.prepare_accelerator(process) + + assert not any( + call.args and call.args[0] is unet + for call in accelerator.prepare.mock_calls + ) + assert process.modules_being_trained == [unet] + + +def test_process_cleanup_retains_runtime_after_failure_and_retries(): + runtime = mock.Mock() + runtime.close.side_effect = (RuntimeError("close failed"), None) + process = SimpleNamespace( + _cleanup_in_progress=False, + _cleanup_completed=False, + _arena_runtime=runtime, + sd=SimpleNamespace(unet=mock.Mock(), text_encoder=None), + ) + + with mock.patch( + "jobs.process.BaseSDTrainProcess.close_memory_runtime_preparation" + ), mock.patch("jobs.process.BaseSDTrainProcess.MemoryManager.detach"): + with pytest.raises(RuntimeError, match="close failed"): + BaseSDTrainProcess.cleanup(process) + assert process._arena_runtime is runtime + assert not process._cleanup_completed + + BaseSDTrainProcess.cleanup(process) + + assert process._arena_runtime is None + assert process._cleanup_completed + assert runtime.close.call_count == 2 + + +def test_sequential_success_cleanup_returns_global_owners_to_baseline(): + model = _frozen_linear() + resources = ArenaRuntimeResources(model, "cpu") + resources.acquire_process_owner() + runtime = SimpleNamespace(close=resources.release) + text_encoder = torch.nn.Linear(4, 4) + text_encoder._memory_manager = SimpleNamespace(unmanaged_modules=[]) + process = SimpleNamespace( + _cleanup_in_progress=False, + _cleanup_completed=False, + _arena_runtime=runtime, + sd=SimpleNamespace(unet=model, text_encoder=text_encoder), + ) + pin_baseline = pin_manager.total_pinned_bytes() + fake_cuda_device = torch.device("cuda") + _DEVICE_STATE[fake_cuda_device] = object() + + with mock.patch("torch.cuda.empty_cache"): + BaseSDTrainProcess.cleanup(process) + + assert active_process_owner() is None + assert process._arena_runtime is None + assert process._cleanup_completed + assert not hasattr(text_encoder, "_memory_manager") + assert fake_cuda_device not in _DEVICE_STATE + assert pin_manager.total_pinned_bytes() == pin_baseline + + +def test_base_job_retains_processes_until_retryable_cleanup_succeeds(): + process = SimpleNamespace(job=object(), cleanup=mock.Mock()) + process.cleanup.side_effect = (RuntimeError("retry me"), None) + job = BaseJob.__new__(BaseJob) + job.process = [process] + + with pytest.raises(RuntimeError, match="retry me"): + job.cleanup() + assert job.process == [process] + assert process.job is not None + + job.cleanup() + assert job.process == [] + assert process.job is None + + +def _imported_modules(path): + tree = ast.parse(path.read_text(encoding="utf-8")) + imported = set() + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported.update(alias.name for alias in node.names) + elif isinstance(node, ast.ImportFrom): + base = node.module or "" + imported.add(base) + imported.update( + f"{base}.{alias.name}" if base else alias.name + for alias in node.names + ) + return imported + + +def test_backend_import_and_shared_private_state_boundaries(): root = Path(__file__).parents[1] - arena_sources = "\n".join( - path.read_text(encoding="utf-8") - for path in (root / "toolkit" / "memory_management" / "arena_offload").glob("*.py") + arena_imports = set() + for path in (root / "toolkit" / "memory_management" / "arena_offload").glob( + "*.py" + ): + arena_imports.update(_imported_modules(path)) + assert not any( + name == "manager" + or name.endswith(".manager") + or name == "manager_modules" + or name.endswith(".manager_modules") + for name in arena_imports ) - assert "from ..manager import" not in arena_sources - assert "manager_modules" not in arena_sources for relative in ( "jobs/process/BaseSDTrainProcess.py", @@ -270,17 +541,16 @@ def test_phase7_import_and_private_state_boundaries(): for node in ast.walk(tree) ) -def test_phase8_legacy_manager_has_no_arena_execution_bridge(): + trainer_source = ( + root / "jobs" / "process" / "BaseSDTrainProcess.py" + ).read_text(encoding="utf-8") + assert "if runtime_owns_block_compile:" in trainer_source + assert "if runtime_owns_block_compile and block_compile:" not in trainer_source + + +def test_legacy_manager_does_not_import_arena_backend(): root = Path(__file__).parents[1] - manager = ( + imports = _imported_modules( root / "toolkit" / "memory_management" / "manager.py" - ).read_text(encoding="utf-8") - for obsolete in ( - "attach_smart_training_immutable", - "smart_immutable", - "_mm_immutable_", - "_immutable_runtime", - "_mm_canonical_leaf", - "canonical_relief", - ): - assert obsolete not in manager + ) + assert not any("arena_offload" in name for name in imports) diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py index 03a8c702cb..8b397c3959 100644 --- a/tests/test_arena_offload_api.py +++ b/tests/test_arena_offload_api.py @@ -19,6 +19,7 @@ is_arena_offloaded, is_memory_managed, memory_runtime_owns_compile, + validate_arena_training_mode, ) from toolkit.memory_management.arena_offload.api import RUNTIME_ATTR, unwrap from toolkit.memory_management.arena_offload.runtime import _fixed_working_bytes @@ -93,9 +94,18 @@ class _FakeModelConfig: layer_offloading_smart_sampling_working_reserve_gb = -1.0 layer_offloading_smart_sampling_wddm_margin_gb = -1.0 layer_offloading_smart_sampling_wddm_hard_gb = 1.0 + layer_offloading_strict_vram_cap = False class ArenaOffloadHelpersTest(unittest.TestCase): + def test_training_mode_requires_frozen_immutable_base_weights(self): + validate_arena_training_mode() + + with self.assertRaisesRegex(ValueError, "full-model fine-tuning"): + validate_arena_training_mode(full_finetune=True) + with self.assertRaisesRegex(ValueError, "merge_network_on_save"): + validate_arena_training_mode(mutates_base_weights=True) + def test_helpers_are_none_safe(self): self.assertIsNone(get_arena_runtime(None)) self.assertFalse(is_arena_offloaded(None)) @@ -169,6 +179,7 @@ def test_from_model_config_maps_the_public_surface(self): self.assertTrue(config.fp8_sampling) # compile_blocks is derived, not its own public knob. self.assertTrue(config.compile_blocks) + self.assertFalse(config.strict_vram_cap) self.assertEqual(config._policy.prefetch_depth, 3) self.assertEqual(config._policy.checkpoint_keep_last, 2) @@ -182,6 +193,7 @@ def test_public_surface_is_narrow(self): "fp8_backward", "fp8_sampling", "compile_blocks", + "strict_vram_cap", }, ) @@ -212,7 +224,7 @@ class Aliases: policy = ArenaOffloadConfig.from_model_config(Aliases())._policy self.assertEqual(policy.working_reserve_gib, 4.0) - self.assertEqual(policy.wddm_margin_gib, 1.5) + self.assertEqual(policy.physical_vram_headroom_gib, 1.5) self.assertEqual(policy.wddm_hard_gib, 0.75) def test_backward_without_fp8_forward_is_ignored_once(self): diff --git a/tests/test_arena_offload_planner.py b/tests/test_arena_offload_planner.py index 68a9dac11b..4233c7d5aa 100644 --- a/tests/test_arena_offload_planner.py +++ b/tests/test_arena_offload_planner.py @@ -1,7 +1,11 @@ from types import SimpleNamespace from unittest import mock -from toolkit.memory_management.arena_offload.planner import GIB, build_training_plan +from toolkit.memory_management.arena_offload.planner import ( + GIB, + build_training_plan, + impossible_training_plan_message, +) class _Arena: @@ -29,7 +33,7 @@ def _config(): _policy=SimpleNamespace( working_reserve_gib=-1, wddm_hard_gib=1.0, - wddm_margin_gib=1.0, + physical_vram_headroom_gib=1.0, checkpoint_keep_last=0, prefetch_depth=2, ) @@ -102,3 +106,32 @@ def test_explicit_working_reserve_controls_all_resident_fit(): assert plan["all_resident_fit"] assert plan["offloaded_layers"] == 0 assert plan["working_reserve_bytes"] == 2 * GIB + + +def test_minimum_layout_that_cannot_fit_has_actionable_admission_error(): + records = [_record("blocks.0", 2.0), _record("blocks.1", 2.0)] + with ( + mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(7 * GIB, 12 * GIB), + ), + mock.patch( + "toolkit.memory_management.arena_offload.planner._singleton_stats", + return_value=(1 * GIB, 0, set()), + ), + ): + plan = build_training_plan( + SimpleNamespace(), _Arena(records), (), "cuda", _config() + ) + + assert not plan["fits"] + assert plan["generic_resident_bytes"] == 0 + message = impossible_training_plan_message(plan) + for field in ( + "required_bytes=", + "available_bytes=", + "reserve_bytes=", + "ring_bytes=", + "singleton_bytes=", + ): + assert field in message diff --git a/tests/test_arena_offload_policy.py b/tests/test_arena_offload_policy.py index f425beb096..9e0addf24a 100644 --- a/tests/test_arena_offload_policy.py +++ b/tests/test_arena_offload_policy.py @@ -188,13 +188,13 @@ def test_worst_shape_allocator_slack_reconstructs_current_layout(): def test_aggressive_capacity_counts_exact_smallest_blocks_under_both_budgets(): runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) - runtime._policy = SimpleNamespace(slack_pad_bytes=5) + runtime._policy = SimpleNamespace(allocator_cache_headroom_bytes=5) runtime._promotion_candidates = lambda: tuple( {"block_key": f"blocks.{index}", "block_bytes": size} for index, size in enumerate((10, 20, 30, 40, 50)) ) runtime._worst_shape_allocator_slack_bytes = lambda _cap: 106 - runtime._worst_shape_candidate_margin_bytes = ( + runtime._worst_shape_candidate_physical_headroom_bytes = ( lambda candidate: 120 - candidate["block_bytes"] ) @@ -203,22 +203,27 @@ def test_aggressive_capacity_counts_exact_smallest_blocks_under_both_budgets(): assert runtime._aggressive_promotion_capacity(1000) == 4 -def test_training_cap_binding_uses_configured_phase_margin(monkeypatch): +def test_training_cap_binding_uses_configured_guard_mode(monkeypatch): calls = [] monkeypatch.setattr( "toolkit.memory_management.arena_offload.runtime." - "allocator_cap.apply_wddm_hard_allocator_cap", + "allocator_cap.configure_wddm_allocator_guard", lambda device, hard, **kwargs: calls.append((device, hard, kwargs)), ) runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) runtime._device = "cuda:1" runtime._config = SimpleNamespace( - _policy=SimpleNamespace(wddm_hard_gib=1.25) + strict_vram_cap=False, + _policy=SimpleNamespace(wddm_hard_gib=1.25), ) runtime._bind_training_cap() assert calls == [ - ("cuda:1", 1.25, {"log_prefix": "[ArenaOffload]"}) + ( + "cuda:1", + 1.25, + {"strict": False, "log_prefix": "[ArenaOffload]"}, + ) ] def test_shape_working_peak_excludes_residency(): window = TrainingSignalWindow() @@ -250,7 +255,9 @@ def test_transfer_benefit_gate_requires_valid_nonzero_streaming(): def test_controller_promotes_exact_candidate_then_rolls_it_back(): - controller = ArenaResidencyController(slack_pad_bytes=10) + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10 + ) candidate = {"block_key": "blocks.3", "block_bytes": 20} clean = { "allocator": {"alloc_retries_delta": 0}, @@ -318,7 +325,7 @@ def test_controller_promotes_exact_candidate_then_rolls_it_back(): worst_shape_allocator_slack_bytes=20, ) assert held.action != "promote" - assert held.reason == "allocator_headband" + assert held.reason == "allocator_cache_headroom" promoted_again = None for _ in range(4): @@ -350,7 +357,9 @@ def test_controller_cold_starts_one_whole_block_below(): def test_controller_promotes_each_step_with_four_block_headroom(): - controller = ArenaResidencyController(slack_pad_bytes=10) + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10 + ) controller.bootstrapped = True clean = { "allocator": { @@ -391,7 +400,9 @@ def test_controller_promotes_each_step_with_four_block_headroom(): def test_controller_four_block_fast_lane_keeps_safety_vetoes(): - controller = ArenaResidencyController(slack_pad_bytes=10) + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10 + ) controller.bootstrapped = True dirty = { "allocator": { @@ -417,7 +428,9 @@ def test_controller_four_block_fast_lane_keeps_safety_vetoes(): def test_controller_does_not_bypass_bootstrap_verification(): - controller = ArenaResidencyController(slack_pad_bytes=10) + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10 + ) controller.bootstrapped = True controller.begin_bootstrap_promotion( ("blocks.0", "blocks.1", "blocks.2", "blocks.3"), @@ -455,7 +468,9 @@ def test_controller_does_not_bypass_bootstrap_verification(): def test_controller_raises_cap_by_fixed_fsm_increment(): - controller = ArenaResidencyController(slack_pad_bytes=10) + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10 + ) controller.bootstrapped = True controller.state = type(controller.state)("stable", 2) signal = { @@ -554,7 +569,7 @@ def test_bf16_sampling_reserves_largest_singleton_dequant(monkeypatch): _policy=SimpleNamespace( sampling_working_reserve_gib="auto", sampling_wddm_hard_gib=1.0, - sampling_wddm_margin_gib=1.0, + sampling_physical_vram_headroom_gib=1.0, ), ) @@ -566,11 +581,12 @@ def sampling(**kwargs): runtime._executor = SimpleNamespace(sampling=sampling) monkeypatch.setattr( "toolkit.memory_management.arena_offload.runtime." - "allocator_cap.apply_wddm_hard_allocator_cap", + "allocator_cap.configure_wddm_allocator_guard", lambda *_args, **_kwargs: None, ) monkeypatch.setattr( - "toolkit.memory_management.arena_offload.runtime.resolve_margin_gib", + "toolkit.memory_management.arena_offload.runtime." + "resolve_physical_vram_headroom_gib", lambda *_args, **_kwargs: 1.0, ) @@ -601,8 +617,12 @@ def test_bootstrap_uses_min_physical_free_and_one_gib_margin(): runtime._bootstrap_block_keys = () runtime._last_step_num = 50_000 runtime._successful_training_steps = 1 + runtime._device = "cpu" runtime._config = SimpleNamespace( - _policy=SimpleNamespace(wddm_hard_gib=1.0) + _policy=SimpleNamespace( + wddm_hard_gib=1.0, + physical_vram_headroom_gib=1.0, + ) ) runtime._model = SimpleNamespace() runtime._arena = SimpleNamespace( @@ -671,8 +691,12 @@ def test_bootstrap_keeps_priority_over_four_block_fast_lane(): runtime._bootstrap_budget_bytes = 0 runtime._bootstrap_block_keys = () runtime._successful_training_steps = 2 + runtime._device = "cpu" runtime._config = SimpleNamespace( - _policy=SimpleNamespace(wddm_hard_gib=1.0) + _policy=SimpleNamespace( + wddm_hard_gib=1.0, + physical_vram_headroom_gib=1.0, + ) ) runtime._arena = SimpleNamespace( block_keys=lambda: tuple(records), @@ -725,7 +749,7 @@ def test_arena_allocation_failure_drains_and_rolls_back(monkeypatch): cap_calls = [] monkeypatch.setattr( "toolkit.memory_management.arena_offload.runtime." - "allocator_cap.apply_wddm_hard_allocator_cap", + "allocator_cap.configure_wddm_allocator_guard", lambda *args, **kwargs: cap_calls.append((args, kwargs)), ) diff --git a/tests/test_generic_block_dispatcher.py b/tests/test_generic_block_dispatcher.py index e23617a068..92f56c503c 100644 --- a/tests/test_generic_block_dispatcher.py +++ b/tests/test_generic_block_dispatcher.py @@ -13,6 +13,11 @@ prepare_arena_offload, ) from toolkit.memory_management.arena_offload.discovery import BlockDiscoveryError +from toolkit.memory_management.arena_offload.dispatcher import ( + _first_output_tensor, + _first_tensor_argument, + _replace_tensor_argument, +) from toolkit.memory_management.arena_offload.ownership import active_process_owner from toolkit.memory_management.residency import ResidencyPlan from toolkit.memory_management.runtime import get_memory_runtime @@ -104,7 +109,7 @@ def _fp8_runtime(model, device, *, forward, backward, compile_blocks): _policy=replace( config._policy, working_reserve_gib=0.0, - wddm_margin_gib=0.0, + physical_vram_headroom_gib=0.0, wddm_hard_gib=1.0, checkpoint_keep_last=1, ), @@ -189,6 +194,25 @@ def test_declared_container_discovery_accounts_all_block_state(): assert selection.accounting.managed_bytes > 0 +def test_dispatcher_abi_accepts_keyword_inputs_and_structured_outputs(): + hidden = torch.randn(2, 4, requires_grad=True) + mask = torch.ones(2, 4, dtype=torch.bool) + args = ("metadata", mask) + kwargs = {"inputs": {"hidden_states": hidden}, "mask": None} + + selected, location = _first_tensor_argument(args, kwargs) + assert selected is hidden + replacement = hidden + 1 + updated_args, updated_kwargs = _replace_tensor_argument( + args, kwargs, location, replacement + ) + assert updated_args == args + assert updated_kwargs["inputs"]["hidden_states"] is replacement + + output = ({"hidden_states": replacement}, (None, hidden)) + assert _first_output_tensor(output) is replacement + + def test_shared_managed_state_is_rejected_before_construction(): model = _frozen_transformer() shared = model.blocks[0].proj.weight @@ -222,7 +246,8 @@ def test_saved_installed_forward_checkpoint_backward_and_teardown(): "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", return_value=(8 * 1024**3, 12 * 1024**3), ), mock.patch( - "toolkit.memory_management.arena_offload.planner.vram_budget.auto_margin_gib", + "toolkit.memory_management.arena_offload.planner.vram_budget." + "auto_physical_vram_headroom_gib", return_value=1.0, ): config = ArenaOffloadConfig(enabled=True, compile_blocks=False) @@ -306,14 +331,16 @@ def test_cuda_streamed_compiled_train_sample_train(): _policy=replace( config._policy, working_reserve_gib=0.0, - wddm_margin_gib=0.0, + physical_vram_headroom_gib=0.0, wddm_hard_gib=1.0, checkpoint_keep_last=1, + prefetch_depth=1, ), ) with mock.patch( "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", - return_value=(1 * 1024**3, 12 * 1024**3), + # Force a mixed plan: two blocks fit, all three do not. + return_value=(700, 12 * 1024**3), ): runtime = prepare_arena_offload( model, diff --git a/tests/test_lora_compile_scalars.py b/tests/test_lora_compile_scalars.py new file mode 100644 index 0000000000..6673516d6c --- /dev/null +++ b/tests/test_lora_compile_scalars.py @@ -0,0 +1,132 @@ +import os + +import pytest +import torch + +from toolkit.compile_utils import configure_cuda_only_inductor +from toolkit.lora_special import LoRAModule +from toolkit.models.DoRA import DoRAModule +from toolkit.models.lokr import LokrModule + + +class _Network: + network_type = "lora" + + +def _linear(): + return torch.nn.Linear(8, 8, bias=False) + + +def test_cuda_only_inductor_keeps_compile_errors_visible(): + configure_cuda_only_inductor() + + assert torch._dynamo.config.suppress_errors is False + if os.name == "nt": + from torch._inductor import config as inductor_config + + assert inductor_config.cpp.vec_isa_ok is False + + +def test_cuda_compile_prefers_aot_eager_then_compile(monkeypatch): + stances = [] + monkeypatch.setattr( + torch.compiler, + "set_stance", + lambda stance: stances.append(stance), + ) + + configure_cuda_only_inductor() + + assert stances == ["aot_eager_then_compile"] + + +def test_lora_tensor_alpha_uses_device_owned_runtime_scale(): + module = LoRAModule( + "compile_scalar", + _linear(), + lora_dim=4, + alpha=torch.tensor(8, dtype=torch.bfloat16), + network=_Network(), + ) + + assert type(module.scale) is float + assert module.scale == 2.0 + assert module._runtime_scale.item() == 2.0 + assert "_runtime_scale" not in module.state_dict() + assert not hasattr(module, "scalar") + + +def test_dora_tensor_alpha_uses_device_owned_runtime_scale(): + module = DoRAModule( + "compile_scalar_dora", + _linear(), + lora_dim=4, + alpha=torch.tensor(8, dtype=torch.bfloat16), + network=_Network(), + ) + + assert type(module.scale) is float + assert module.scale == 2.0 + assert module._runtime_scale.item() == 2.0 + assert "_runtime_scale" not in module.state_dict() + assert not hasattr(module, "scalar") + + +def test_lokr_tensor_alpha_uses_device_owned_runtime_scale(): + module = LokrModule( + "compile_scalar_lokr", + _linear(), + lora_dim=2, + alpha=torch.tensor(4, dtype=torch.bfloat16), + network=_Network(), + ) + + assert type(module.scale) is float + assert module.scale == 1.0 + assert module._runtime_scale.item() == 1.0 + assert "_runtime_scale" not in module.state_dict() + assert module.get_weight().device.type == "cpu" + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_functional_call_compiles_lora_without_cpu_codegen(): + configure_cuda_only_inductor() + network = _Network() + network.is_lorm = False + network.is_active = True + network.is_merged_in = False + network._multiplier = 1.0 + network.torch_multiplier = torch.ones(1, device="cuda") + original = torch.nn.Linear( + 8, 8, bias=False, device="cuda", dtype=torch.bfloat16 + ) + module = LoRAModule( + "compile_scalar_cuda", + original, + lora_dim=4, + alpha=torch.tensor(8, dtype=torch.bfloat16), + network=network, + ).to("cuda") + module.org_forward = original.forward + state = dict(module.named_parameters()) + state.update(module.named_buffers()) + + def kernel(value): + return torch.func.functional_call( + module, + state, + (value,), + strict=False, + tie_weights=False, + ) + + compiled = torch.compile(kernel, fullgraph=True, dynamic=False) + value = torch.randn(2, 8, device="cuda", dtype=torch.bfloat16) + result = compiled(value) + result.square().mean().backward() + torch.cuda.synchronize() + + assert module._runtime_scale.device.type == "cuda" + assert result.device.type == "cuda" + assert module.lora_down.weight.grad is not None + assert module.lora_up.weight.grad is not None diff --git a/toolkit/compile_utils.py b/toolkit/compile_utils.py new file mode 100644 index 0000000000..7ef44c876f --- /dev/null +++ b/toolkit/compile_utils.py @@ -0,0 +1,23 @@ +import os + +import torch + + +def configure_cuda_only_inductor() -> None: + """Configure CUDA compilation without requiring a CPU toolchain.""" + torch._dynamo.config.suppress_errors = False + set_stance = getattr(torch.compiler, "set_stance", None) + if set_stance is not None: + try: + # The first AOT-eager pass preserves checkpointing's memory + # behavior and gives Dynamo real shape evidence before compiling. + set_stance("aot_eager_then_compile") + except (RuntimeError, ValueError): + # Older supported PyTorch builds do not expose this stance. + pass + if os.name == "nt": + # Inductor otherwise dry-compiles a CPU vector-ISA probe even when the + # requested graph is CUDA-only. This does not enable a CPU fallback. + from torch._inductor import config as inductor_config + + inductor_config.cpp.vec_isa_ok = False diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index b7a65f5984..b75d4930b8 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -705,8 +705,9 @@ def __init__(self, **kwargs): self.layer_offloading_smart_working_reserve_gb = kwargs.get( "layer_offloading_smart_working_reserve_gb", -1.0 ) - self.layer_offloading_smart_wddm_margin_gb = kwargs.get( - "layer_offloading_smart_wddm_margin_gb", -1.0 + self.layer_offloading_smart_physical_vram_headroom_gb = kwargs.get( + "layer_offloading_smart_physical_vram_headroom_gb", + kwargs.get("layer_offloading_smart_wddm_margin_gb", -1.0), ) self.layer_offloading_smart_wddm_hard_gb = kwargs.get( "layer_offloading_smart_wddm_hard_gb", 1.0 @@ -714,12 +715,16 @@ def __init__(self, **kwargs): self.layer_offloading_smart_sampling_working_reserve_gb = kwargs.get( "layer_offloading_smart_sampling_working_reserve_gb", -1.0 ) - self.layer_offloading_smart_sampling_wddm_margin_gb = kwargs.get( - "layer_offloading_smart_sampling_wddm_margin_gb", -1.0 + self.layer_offloading_smart_sampling_physical_vram_headroom_gb = kwargs.get( + "layer_offloading_smart_sampling_physical_vram_headroom_gb", + kwargs.get("layer_offloading_smart_sampling_wddm_margin_gb", -1.0), ) self.layer_offloading_smart_sampling_wddm_hard_gb = kwargs.get( "layer_offloading_smart_sampling_wddm_hard_gb", 1.0 ) + self.layer_offloading_strict_vram_cap = kwargs.get( + "layer_offloading_strict_vram_cap", False + ) self.layer_offloading_fp8_forward = kwargs.get( "layer_offloading_fp8_forward", False ) diff --git a/toolkit/lora_special.py b/toolkit/lora_special.py index bf8309db8b..86994dbc41 100644 --- a/toolkit/lora_special.py +++ b/toolkit/lora_special.py @@ -113,7 +113,7 @@ def __init__( if type(alpha) == torch.Tensor: alpha = float(alpha.detach().float().item()) alpha = self.lora_dim if alpha is None or alpha == 0 else alpha - self.scale = float(alpha) / self.lora_dim + self._set_runtime_scale(float(alpha) / self.lora_dim) self.register_buffer("alpha", torch.tensor(alpha)) # 定数として扱える # same as microsoft's diff --git a/toolkit/lycoris_special.py b/toolkit/lycoris_special.py index 8bafb6d925..1bce9dfeec 100644 --- a/toolkit/lycoris_special.py +++ b/toolkit/lycoris_special.py @@ -96,7 +96,7 @@ def __init__( if type(alpha) == torch.Tensor: alpha = float(alpha.detach().float().item()) alpha = lora_dim if alpha is None or alpha == 0 else alpha - self.scale = float(alpha) / self.lora_dim + self._set_runtime_scale(float(alpha) / self.lora_dim) self.register_buffer('alpha', torch.tensor(alpha)) # 定数として扱える # same as microsoft's diff --git a/toolkit/memory_management/allocator_cap.py b/toolkit/memory_management/allocator_cap.py index 47707e73bd..7eec779e5a 100644 --- a/toolkit/memory_management/allocator_cap.py +++ b/toolkit/memory_management/allocator_cap.py @@ -3,6 +3,7 @@ from __future__ import annotations import sys +import warnings import torch @@ -10,7 +11,6 @@ GIB = 1024 ** 3 APPLIED_FRACTIONS: dict[int, float] = {} -RELIEF_BYTES: dict[int, int] = {} @@ -52,18 +52,20 @@ def applied_cap_bytes(device) -> int | None: return int(float(fraction) * vram_budget.real_device_total_bytes(index)) -def apply_wddm_hard_allocator_cap( +def configure_wddm_allocator_guard( device, wddm_hard_gib=None, *, target_cap_bytes=None, + strict=False, log_prefix="[MemoryManager]", ): - """Bind torch's allocator below the WDDM dedicated-memory cliff. + """Configure the WDDM cliff guard for production or development. - Call only at a phase boundary. The governing capacity may be a simulated - smaller card, but torch's fraction is always converted against the physical - card total. + Production treats the cliff as a planning target and permits WDDM spill; + strict development mode binds torch's allocator below it so a breach raises + OOM. Call only at a phase boundary. A simulated governing capacity is + always converted against the physical card total. """ if sys.platform != "win32" or not torch.cuda.is_available(): return None @@ -71,6 +73,19 @@ def apply_wddm_hard_allocator_cap( if dev.type != "cuda": return None index = dev.index if dev.index is not None else torch.cuda.current_device() + if not strict: + previous = APPLIED_FRACTIONS.pop(index, None) + if previous is not None and previous < 1.0: + try: + torch.cuda.set_per_process_memory_fraction(1.0, index) + except Exception as error: + warnings.warn( + "could not remove the strict CUDA allocator cap; " + f"WDDM spill fallback may remain unavailable: {error}", + RuntimeWarning, + stacklevel=2, + ) + return None try: hard_gib = float(wddm_hard_gib) if wddm_hard_gib is not None else 1.0 except (TypeError, ValueError): @@ -91,9 +106,6 @@ def apply_wddm_hard_allocator_cap( fraction = max(0.1, min(cliff_fraction, target_fraction)) reclaimed = fraction < cliff_fraction - 1e-9 - relief_bytes = RELIEF_BYTES.get(index, 0) - if relief_bytes: - fraction = min(1.0, fraction + relief_bytes / float(total)) applied = fraction * total / float(real_total) previous = APPLIED_FRACTIONS.get(index) tolerance = (64 * 1024**2) / real_total @@ -108,12 +120,10 @@ def apply_wddm_hard_allocator_cap( if reclaimed else "cliff bound" ) - if relief_bytes: - source += f"; +{relief_bytes / GIB:.2f} GiB post-violation relief" if total != real_total: source += f"; SIMULATED {total / GIB:.2f} GiB card" print( - f"{log_prefix} WDDM hard allocator cap: " + f"{log_prefix} strict WDDM allocator cap: " f"{fraction * total / GIB:.2f}/{total / GIB:.2f} GiB " f"({source}; margin {hard_gib:.2f} GiB, " f"non_torch {non_torch / GIB:.2f} GiB; allocation beyond this " diff --git a/toolkit/memory_management/arena_offload/__init__.py b/toolkit/memory_management/arena_offload/__init__.py index 92e4e45666..a3002345f9 100644 --- a/toolkit/memory_management/arena_offload/__init__.py +++ b/toolkit/memory_management/arena_offload/__init__.py @@ -27,6 +27,7 @@ prepare_canonical_storage, prepare_canonical_storage_from_state_dict, prepare_arena_offload, + validate_arena_training_mode, ) from .runtime import ArenaOffloadRuntime from .dispatcher import DISPATCHER_GENERATION @@ -54,4 +55,5 @@ "prepare_canonical_storage", "prepare_canonical_storage_from_state_dict", "prepare_arena_offload", + "validate_arena_training_mode", ] diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index 3c2b028326..cae6057206 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -34,7 +34,8 @@ "layer_offloading_smart_working_reserve_gb": ( "layer_offloading_smart_headroom_gb", ), - "layer_offloading_smart_wddm_margin_gb": ( + "layer_offloading_smart_physical_vram_headroom_gb": ( + "layer_offloading_smart_wddm_margin_gb", "layer_offloading_smart_buffer_gb", ), "layer_offloading_smart_wddm_hard_gb": ( @@ -43,7 +44,8 @@ "layer_offloading_smart_sampling_working_reserve_gb": ( "layer_offloading_smart_sampling_headroom_gb", ), - "layer_offloading_smart_sampling_wddm_margin_gb": ( + "layer_offloading_smart_sampling_physical_vram_headroom_gb": ( + "layer_offloading_smart_sampling_wddm_margin_gb", "layer_offloading_smart_sampling_buffer_gb", ), "layer_offloading_smart_sampling_wddm_hard_gb": ( @@ -52,6 +54,27 @@ } +def validate_arena_training_mode( + *, full_finetune=False, mutates_base_weights=False +) -> None: + """Reject mutable-base configurations before model loading. + + Canonical arena leaves are immutable frozen base weights. Full-parameter + training or merge-in save workflows require a different storage + architecture, independent of the model or quantization integration. + """ + if full_finetune: + raise ValueError( + "arena offload requires frozen base transformer weights and does " + "not support full-model fine-tuning" + ) + if mutates_base_weights: + raise ValueError( + "arena offload requires immutable base transformer weights and " + "does not support merge_network_on_save" + ) + + def unwrap(model): """Peel Accelerate / DDP / torch.compile wrappers without importing them. @@ -67,13 +90,13 @@ class _ArenaPolicyOptions: """Internal policy inputs retained while fork job aliases are migrated.""" working_reserve_gib: float | None = None - wddm_margin_gib: float | None = None + physical_vram_headroom_gib: float | None = None wddm_hard_gib: float | None = None checkpoint_keep_last: int = 0 prefetch_depth: int = 3 sampling_working_reserve_gib: float | None = None - sampling_wddm_margin_gib: float | None = None + sampling_physical_vram_headroom_gib: float | None = None sampling_wddm_hard_gib: float | None = 1.0 @@ -90,6 +113,7 @@ class ArenaOffloadConfig: fp8_backward: bool = False fp8_sampling: bool = False compile_blocks: bool = False + strict_vram_cap: bool = False _compile_dynamic: bool | None = True _compile_dynamic_hints: tuple[tuple[int, int | None, int | None], ...] = () # Validation knob: pretend the card is this many GiB, so small-card @@ -167,6 +191,9 @@ def get(name: str, default: Any = None) -> Any: or get("compile_sample", False) or get("train_compile_blocks", False) ), + strict_vram_cap=bool( + get("layer_offloading_strict_vram_cap", False) + ), _compile_dynamic=( None if get("compile_dynamic", True) is None @@ -180,7 +207,9 @@ def get(name: str, default: Any = None) -> Any: ), _policy=_ArenaPolicyOptions( working_reserve_gib=working_reserve_gib, - wddm_margin_gib=get("layer_offloading_smart_wddm_margin_gb"), + physical_vram_headroom_gib=get( + "layer_offloading_smart_physical_vram_headroom_gb" + ), wddm_hard_gib=get("layer_offloading_smart_wddm_hard_gb"), checkpoint_keep_last=max( 0, int(get("layer_offloading_checkpoint_keep_last", 0) or 0) @@ -189,8 +218,8 @@ def get(name: str, default: Any = None) -> Any: sampling_working_reserve_gib=get( "layer_offloading_smart_sampling_working_reserve_gb" ), - sampling_wddm_margin_gib=get( - "layer_offloading_smart_sampling_wddm_margin_gb" + sampling_physical_vram_headroom_gib=get( + "layer_offloading_smart_sampling_physical_vram_headroom_gb" ), sampling_wddm_hard_gib=get( "layer_offloading_smart_sampling_wddm_hard_gb", 1.0 diff --git a/toolkit/memory_management/arena_offload/dispatcher.py b/toolkit/memory_management/arena_offload/dispatcher.py index 6cb66ab0c6..57c3d55676 100644 --- a/toolkit/memory_management/arena_offload/dispatcher.py +++ b/toolkit/memory_management/arena_offload/dispatcher.py @@ -6,6 +6,9 @@ import torch +from toolkit.compile_utils import configure_cuda_only_inductor +from torch.utils._pytree import tree_flatten, tree_unflatten + from toolkit.memory_management.immutable_runtime import ( ImmutableProgram, ImmutableRuntimeError, @@ -21,6 +24,44 @@ DISPATCHER_GENERATION = "generic-block-dispatcher-v1" +def _first_tensor_argument(args, kwargs): + """Locate the tensor leaf that best carries block execution lifetime.""" + leaves, spec = tree_flatten((args, kwargs)) + tensors = [ + (index, value) + for index, value in enumerate(leaves) + if isinstance(value, torch.Tensor) + ] + for index, value in tensors: + if value.requires_grad: + return value, (spec, index) + if tensors: + index, value = tensors[0] + return value, (spec, index) + raise ImmutableRuntimeError("unsupported_block_arguments:no_tensor_argument") + + +def _replace_tensor_argument(args, kwargs, location, value): + expected_spec, index = location + leaves, spec = tree_flatten((args, kwargs)) + if spec != expected_spec: + raise ImmutableRuntimeError("block_argument_structure_changed") + leaves[index] = value + return tree_unflatten(leaves, spec) + + +def _first_output_tensor(output): + """Return a lifetime guard without constraining the block output shape.""" + leaves, _spec = tree_flatten(output) + tensors = [value for value in leaves if isinstance(value, torch.Tensor)] + for value in tensors: + if value.requires_grad: + return value + if tensors: + return tensors[0] + raise ImmutableRuntimeError("unsupported_block_output:no_tensor_leaf") + + def _in_backward_graph_task() -> bool: """True for non-reentrant checkpoint replay inside autograd backward.""" try: @@ -165,6 +206,7 @@ def kernel(leaf_args, args, kwargs): ) if self.compile_blocks: + configure_cuda_only_inductor() kernel = torch.compile( kernel, mode="default", @@ -242,16 +284,17 @@ def dispatch(self, index, args, kwargs): raise ImmutableRuntimeError( "immutable_execution_mode_mismatch:active=sample:call=train" ) - if not args or not isinstance(args[0], torch.Tensor): + try: + first, first_location = _first_tensor_argument(args, kwargs) + except ImmutableRuntimeError as error: raise ImmutableRuntimeError( f"unsupported_block_arguments:{self._block_abis[index].block_key}" - ) + ) from error source = self._sources.source(index) transfer = source.transfer token = None compact_flat = None - first = args[0] training = self._active_mode == self.TRAIN if transfer is not None: if training and any( @@ -273,7 +316,9 @@ def dispatch(self, index, args, kwargs): compact_flat = torch.ops.mm.fetch_wait(token, nbytes) if training and torch.is_grad_enabled(): first = free_on_backward(first, token) - args = (first, *args[1:]) + args, kwargs = _replace_tensor_argument( + args, kwargs, first_location, first + ) leaf_args = source.assemble_leaf_args(self.residency, compact_flat) self._mark_dispatch_dynamic(first) @@ -288,7 +333,9 @@ def dispatch(self, index, args, kwargs): or not torch.is_grad_enabled() or not _in_backward_graph_task() ): - torch.ops.mm.fetch_free_after(token, output) + torch.ops.mm.fetch_free_after( + token, _first_output_tensor(output) + ) return output def close(self): diff --git a/toolkit/memory_management/arena_offload/planner.py b/toolkit/memory_management/arena_offload/planner.py index dfa28038fa..ecb5c8f66a 100644 --- a/toolkit/memory_management/arena_offload/planner.py +++ b/toolkit/memory_management/arena_offload/planner.py @@ -2,6 +2,8 @@ from __future__ import annotations +from types import SimpleNamespace + import torch from toolkit.quantization.storage import temporary_materialization_bytes @@ -12,15 +14,16 @@ GIB = 1024**3 DEFAULT_AUTO_WORKING_RESERVE_GIB = 5.0 # Full residency removes the transfer ring, but training still needs activation, -# adapter, dequantization, and allocator-fragmentation headroom. A production -# Orbit4 full-model smoke exhausted a 2 GiB reserve during checkpoint backward; -# 4 GiB is the narrowest evidence-backed automatic bound. Explicit policy values -# remain authoritative for workloads with measured tighter requirements. +# adapter, dequantization, and allocator-fragmentation headroom. Keep a +# conservative cold-start reserve on the target 12 GiB class of devices; +# explicit policy values remain authoritative for measured workloads. DEFAULT_ALL_RESIDENT_WORKING_RESERVE_GIB = 4.0 DEFAULT_RESIDENT_FLOOR_GIB = 2.0 -def resolve_margin_gib(device, value, *, hard_gib=0.0) -> float: +def resolve_physical_vram_headroom_gib( + device, value, *, hard_gib=0.0 +) -> float: try: margin = float(value) automatic = margin < 0 @@ -28,8 +31,11 @@ def resolve_margin_gib(device, value, *, hard_gib=0.0) -> float: automatic = value is None or str(value).strip().lower() == "auto" margin = -1.0 if automatic: - margin = vram_budget.auto_margin_gib(device) - return max(float(margin), float(hard_gib or 0.0)) + margin = max( + float(vram_budget.auto_physical_vram_headroom_gib(device)), + float(hard_gib or 0.0), + ) + return max(0.0, float(margin)) def training_pinned_keys_for_keep_last(model, keep_last, block_keys=None) -> set[str]: @@ -80,8 +86,51 @@ def _singleton_stats(model, canonical_modules) -> tuple[int, int, set[int]]: return total, largest_materialization, runtime_ids +def _planning_records(arena_or_build): + """Return the record surface needed by the cold planner. + + A prepared canonical build exposes final allocation sizes and module/leaf + membership before it publishes any Parameter views. Planning against that + surface keeps predictable admission failures on the rollback-safe side of + the canonical commit boundary. + """ + if hasattr(arena_or_build, "block_keys"): + return [ + arena_or_build.block_record(key) + for key in arena_or_build.block_keys() + if arena_or_build.block_record(key) is not None + ] + blocks = getattr(arena_or_build, "blocks", None) + if blocks is None: + raise TypeError("arena planner requires an arena or prepared build") + return [ + SimpleNamespace( + block_key=block.key, + committed_bytes=int(block.layout.nbytes), + modules=tuple(module for _name, module in block.entries), + leaf_names=tuple(name for name, _module in block.entries), + ) + for block in blocks + ] + + +def impossible_training_plan_message(plan) -> str: + """Describe a minimum-layout admission failure with actionable budgets.""" + return ( + "arena training minimum layout does not fit before canonical commit: " + f"required_bytes={int(plan['minimum_required_bytes'])}, " + f"available_bytes={int(plan['available_bytes'])}, " + f"reserve_bytes={int(plan['working_reserve_bytes'])}, " + f"ring_bytes={int(plan['ring_bytes'])}, " + f"singleton_bytes={int(plan['singleton_resident_bytes'])}, " + f"mandatory_resident_bytes={int(plan['pinned_resident_bytes'])}, " + "physical_vram_headroom_bytes=" + f"{int(plan['physical_vram_headroom_bytes'])}" + ) + + def build_training_plan( - model, arena, canonical_modules, device, config, *, block_keys=None + model, arena_or_build, canonical_modules, device, config, *, block_keys=None ) -> dict: """Choose an initial whole-block layout without the legacy manager.""" device = torch.device(device) @@ -98,19 +147,18 @@ def build_training_plan( working_value = DEFAULT_AUTO_WORKING_RESERVE_GIB hard_gib = float(policy.wddm_hard_gib or 1.0) - margin_gib = resolve_margin_gib( - device, policy.wddm_margin_gib, hard_gib=hard_gib + physical_headroom_gib = resolve_physical_vram_headroom_gib( + device, policy.physical_vram_headroom_gib, hard_gib=hard_gib ) free_bytes, total_bytes = vram_budget.device_mem_info(device) working_bytes = int(max(0.0, working_value) * GIB) - margin_bytes = int(margin_gib * GIB) + physical_headroom_bytes = int(physical_headroom_gib * GIB) hard_bytes = int(hard_gib * GIB) singleton_bytes, largest_singleton_dequant, runtime_ids = _singleton_stats( model, canonical_modules ) - records = [arena.block_record(key) for key in arena.block_keys()] - records = [record for record in records if record is not None] + records = _planning_records(arena_or_build) block_bytes = sum(record.committed_bytes for record in records) all_resident_working_bytes = working_bytes if automatic: @@ -120,7 +168,7 @@ def build_training_plan( ) all_resident_fit = ( singleton_bytes + block_bytes + all_resident_working_bytes - <= max(0, int(free_bytes) - margin_bytes) + <= max(0, int(free_bytes) - physical_headroom_bytes) ) if all_resident_fit: # A transfer ring and the generic 5 GiB cold-start reserve are both @@ -143,7 +191,9 @@ def build_training_plan( streamed = [record for record in records if record.block_key not in resident_keys] largest_stream = max((record.committed_bytes for record in streamed), default=0) ring_bytes = largest_stream * max(1, int(policy.prefetch_depth)) - usable = max(0, int(free_bytes) - margin_bytes - working_bytes) + usable = max( + 0, int(free_bytes) - physical_headroom_bytes - working_bytes + ) resident_budget = max(0, usable - singleton_bytes - ring_bytes) resident_bytes = sum( record.committed_bytes for record in records if record.block_key in resident_keys @@ -186,19 +236,25 @@ def build_training_plan( except Exception: system_reserve = max(0, int(total_bytes) - int(free_bytes)) + pinned_resident_bytes = sum( + record.committed_bytes + for record in records + if record.block_key in pinned_keys + ) + available_bytes = max(0, int(free_bytes) - physical_headroom_bytes) + minimum_required_bytes = ( + singleton_bytes + pinned_resident_bytes + ring_bytes + working_bytes + ) + return { "offload_ids": offload_ids, "offloaded_layers": len(offload_ids), "candidate_layers": sum(len(record.modules) for record in records), "model_bytes": singleton_bytes + block_bytes, "resident_bytes": singleton_bytes + resident_bytes, - "must_resident_bytes": singleton_bytes + sum( - record.committed_bytes for record in records if record.block_key in pinned_keys - ), + "must_resident_bytes": singleton_bytes + pinned_resident_bytes, "must_resident_layer_keys": set(), - "pinned_resident_bytes": sum( - record.committed_bytes for record in records if record.block_key in pinned_keys - ), + "pinned_resident_bytes": pinned_resident_bytes, "pinned_resident_keys": set(pinned_keys), "protected_training_leaf_keys": protected, "generic_resident_bytes": max(0, resident_bytes), @@ -206,12 +262,14 @@ def build_training_plan( "gpu_stream_need_bytes": ring_bytes, "gpu_stream_budget_bytes": ring_bytes, "working_reserve_bytes": working_bytes, - "wddm_margin_bytes": margin_bytes, + "physical_vram_headroom_bytes": physical_headroom_bytes, "wddm_hard_bytes": hard_bytes, "system_reserve_bytes": system_reserve, "usable_bytes": usable, "free_bytes": int(free_bytes), - "fits": singleton_bytes + resident_bytes + ring_bytes <= usable, + "fits": minimum_required_bytes <= available_bytes, + "minimum_required_bytes": minimum_required_bytes, + "available_bytes": available_bytes, "singleton_resident_bytes": singleton_bytes, "largest_singleton_bf16_dequant_bytes": largest_singleton_dequant, "singleton_runtime_ids": runtime_ids, diff --git a/toolkit/memory_management/arena_offload/policy.py b/toolkit/memory_management/arena_offload/policy.py index 616932bece..6cc682a8de 100644 --- a/toolkit/memory_management/arena_offload/policy.py +++ b/toolkit/memory_management/arena_offload/policy.py @@ -9,7 +9,7 @@ _ALLOC = ("num_alloc_retries", "num_device_alloc", "num_device_free") _COMPILE = ("frames", "graphs", "graph_breaks") -DEFAULT_SLACK_PAD_BYTES = 256 * 1024**2 +DEFAULT_ALLOCATOR_CACHE_HEADROOM_BYTES = 256 * 1024**2 AGGRESSIVE_PROMOTION_MIN_CAPACITY = 4 @@ -36,16 +36,24 @@ class PolicyDecision: class ArenaResidencyController: """Stateful wiring around the pure two-timescale residency FSM.""" - def __init__(self, *, slack_pad_bytes=DEFAULT_SLACK_PAD_BYTES): + def __init__( + self, + *, + allocator_cache_headroom_bytes=( + DEFAULT_ALLOCATOR_CACHE_HEADROOM_BYTES + ), + ): self.state = vram_budget.ResidencyFsmState() - self.slack_pad_bytes = max(0, int(slack_pad_bytes)) + self.allocator_cache_headroom_bytes = max( + 0, int(allocator_cache_headroom_bytes) + ) self.last_action = "hold" self.last_reason = "cold_start" self.last_promoted_key = None self.last_block_key = None self.last_block_bytes = 0 self.last_target_cap_bytes = None - self.last_worst_shape_margin_bytes = None + self.last_worst_shape_physical_headroom_bytes = None self.last_throughput_gate = None self.last_promote_gate = None self.last_cap_covers_promo = None @@ -87,7 +95,10 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, and throughput_ok and worst_ok and vram_budget.residency_promote_ok( - retries, allocator_slack, block_bytes, self.slack_pad_bytes + retries, + allocator_slack, + block_bytes, + self.allocator_cache_headroom_bytes, ) ) active_cap = ( @@ -97,7 +108,8 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, ) cap_covers = ( candidate is not None - and allocator_slack > block_bytes + self.slack_pad_bytes + and allocator_slack + > block_bytes + self.allocator_cache_headroom_bytes ) binding = retries > 0 or int(worst_shape_free_bytes) < 0 aggressive_capacity = max(0, int(aggressive_promotion_capacity or 0)) @@ -113,9 +125,12 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, and retries == 0 and device_frees == 0 and worst_ok - and allocator_slack > block_bytes + self.slack_pad_bytes + and allocator_slack + > block_bytes + self.allocator_cache_headroom_bytes + ) + self.last_worst_shape_physical_headroom_bytes = int( + worst_shape_free_bytes ) - self.last_worst_shape_margin_bytes = int(worst_shape_free_bytes) self.last_worst_shape_allocator_slack_bytes = allocator_slack self.last_throughput_gate = bool(throughput_ok) self.last_promote_gate = bool(promote_ok) @@ -125,7 +140,7 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, cap_raise_bytes = ( block_bytes if promote_ok and block_bytes > 0 - else self.slack_pad_bytes + else self.allocator_cache_headroom_bytes ) needed_cap = min( int(cliff_cap_bytes), @@ -204,9 +219,10 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, "throughput_gate" if candidate is not None and not throughput_ok else ( - "allocator_headband" + "allocator_cache_headroom" if candidate is not None - and allocator_slack <= block_bytes + self.slack_pad_bytes + and allocator_slack + <= block_bytes + self.allocator_cache_headroom_bytes else "fsm_hold" ) ) @@ -319,7 +335,9 @@ def diagnostics(self): "last_block_key": self.last_block_key, "last_block_bytes": self.last_block_bytes, "last_target_cap_bytes": self.last_target_cap_bytes, - "last_worst_shape_margin_bytes": self.last_worst_shape_margin_bytes, + "last_worst_shape_physical_headroom_bytes": ( + self.last_worst_shape_physical_headroom_bytes + ), "last_throughput_gate": self.last_throughput_gate, "last_promote_gate": self.last_promote_gate, "last_cap_covers_promo": self.last_cap_covers_promo, @@ -333,7 +351,9 @@ def diagnostics(self): "last_rejected_residency_bytes": ( self.last_rejected_residency_bytes ), - "slack_pad_bytes": self.slack_pad_bytes, + "allocator_cache_headroom_bytes": ( + self.allocator_cache_headroom_bytes + ), } diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index 21c2ffebe1..956991ca77 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -13,6 +13,7 @@ import contextlib import time +import warnings from collections.abc import Sequence from dataclasses import replace from typing import Any @@ -30,7 +31,11 @@ from .fp8 import disable as disable_fp8 from .fp8 import enable as enable_fp8 from .fp8 import set_fp8_grad_input_enabled -from .planner import build_training_plan, resolve_margin_gib +from .planner import ( + build_training_plan, + impossible_training_plan_message, + resolve_physical_vram_headroom_gib, +) from .resources import ArenaRuntimeResources RUNTIME_ATTR = "_arena_offload_runtime" @@ -90,6 +95,7 @@ def __init__( self._sampling_fp8_canonical = 0 self._sampling_fp8_singletons = 0 self._permanent_placement = None + self._device_state_parked_plan = None # ------------------------------------------------------------------ # construction @@ -122,8 +128,11 @@ def _prepare( try: # Bind card simulation and allocator policy before any plan reads. apply_simulated_card(config._simulated_vram_gib, device=device) - allocator_cap.apply_wddm_hard_allocator_cap( - device, config._policy.wddm_hard_gib, log_prefix="[ArenaOffload]" + allocator_cap.configure_wddm_allocator_guard( + device, + config._policy.wddm_hard_gib, + strict=config.strict_vram_cap, + log_prefix="[ArenaOffload]", ) set_fp8_grad_input_enabled(config.fp8_backward) @@ -154,21 +163,31 @@ def _prepare( raise RuntimeError("arena_canonical_build_selection_mismatch") arena = canonical_build.arena - canonical_build.commit() - resources.mark_canonical_committed() canonical_modules = tuple( child for entries in entries_by_block.values() for _name, child in entries ) - resources.canonical_modules = canonical_modules - smart_plan = build_training_plan( transformer, - arena, + canonical_build, canonical_modules, device, config, block_keys=block_keys, ) + if not smart_plan["fits"]: + message = impossible_training_plan_message(smart_plan) + if config.strict_vram_cap: + raise ValueError(message) + warnings.warn( + message + + "; continuing in production spill-permitted mode", + RuntimeWarning, + stacklevel=2, + ) + + canonical_build.commit() + resources.mark_canonical_committed() + resources.canonical_modules = canonical_modules residency = ResidencyState(arena, device) resources.adopt_residency(residency) training_plan = ResidencyPlan.from_smart_plan( @@ -330,6 +349,40 @@ def finalized(self) -> bool: # lifecycle # ------------------------------------------------------------------ + def park_residency_for_external_phase(self) -> None: + """Drop device sidecars while another large model component is active.""" + self._require_open() + if self._device_state_parked_plan is not None: + return + current = self._residency.plan + if current.phase != self._executor.TRAIN: + raise RuntimeError( + f"arena_external_phase_requires_train:{current.phase}" + ) + self._device_state_parked_plan = current + parked = ResidencyPlan.build(self._executor.TRAIN, ()) + try: + if self.finalized: + self._executor.set_residency_plan(parked) + else: + self._residency.reconcile(parked) + except BaseException: + self._device_state_parked_plan = None + raise + + def restore_residency_after_external_phase(self) -> None: + """Restore the exact training residency saved by the matching park.""" + self._require_open() + plan = self._device_state_parked_plan + if plan is None: + return + if self.finalized: + self._executor.set_residency_plan(plan) + else: + self._residency.reconcile(plan) + self._training_plan = plan + self._device_state_parked_plan = None + def set_compile_dynamic_hints(self, hints) -> None: """Install mark_dynamic hints on the block kernels (see ImmutableRuntime). @@ -501,10 +554,13 @@ def _handle_training_failure(self, error, *, shape_key, step_num): decision.block_key, resident=False ) if decision.target_cap_bytes is not None: - allocator_cap.apply_wddm_hard_allocator_cap( + allocator_cap.configure_wddm_allocator_guard( self._device, self._config._policy.wddm_hard_gib, target_cap_bytes=decision.target_cap_bytes, + strict=getattr( + self._config, "strict_vram_cap", False + ), log_prefix="[ArenaOffload]", ) self._last_training_cap_target_bytes = ( @@ -616,12 +672,15 @@ def sampling_image(self, *, shape_key: tuple, cold_working_bytes: int): if policy.sampling_wddm_hard_gib is None else float(policy.sampling_wddm_hard_gib) ) - allocator_cap.apply_wddm_hard_allocator_cap( - self._device, hard_gib, log_prefix="[ArenaOffload]" + allocator_cap.configure_wddm_allocator_guard( + self._device, + hard_gib, + strict=getattr(self._config, "strict_vram_cap", False), + log_prefix="[ArenaOffload]", ) - margin_gib = resolve_margin_gib( + physical_headroom_gib = resolve_physical_vram_headroom_gib( self._device, - policy.sampling_wddm_margin_gib, + policy.sampling_physical_vram_headroom_gib, hard_gib=hard_gib, ) dequant_reserve = ( @@ -637,7 +696,9 @@ def sampling_image(self, *, shape_key: tuple, cold_working_bytes: int): shape_key=shape_key, cold_working_bytes=int(cold_working_bytes), fixed_working_bytes=fixed_working_bytes, - cold_floor_bytes=int(margin_gib * GIB) + dequant_reserve, + cold_floor_bytes=( + int(physical_headroom_gib * GIB) + dequant_reserve + ), hot_floor_bytes=int((hard_gib + 0.25) * GIB) + dequant_reserve, ): yield self @@ -658,6 +719,16 @@ def record_training_physical_free_min(self, free_bytes) -> None: else min(self._bootstrap_min_free_bytes, value) ) + def _training_physical_vram_headroom_bytes(self) -> int: + policy = self._config._policy + hard_gib = float(getattr(policy, "wddm_hard_gib", None) or 1.0) + headroom_gib = resolve_physical_vram_headroom_gib( + self._device, + getattr(policy, "physical_vram_headroom_gib", -1.0), + hard_gib=hard_gib, + ) + return int(headroom_gib * GIB) + def _bootstrap_training_residency(self, active_cap_bytes) -> bool: if ( self._bootstrap_complete @@ -665,14 +736,10 @@ def _bootstrap_training_residency(self, active_cap_bytes) -> bool: or int(self._successful_training_steps) < BOOTSTRAP_MIN_STEP ): return False - hard_gib = self._config._policy.wddm_hard_gib - hard_bytes = int( - (1.0 if hard_gib is None else max(1.0, float(hard_gib))) * GIB - ) budget = max( 0, self._bootstrap_min_free_bytes - - hard_bytes + - self._training_physical_vram_headroom_bytes() - BOOTSTRAP_MARGIN_BYTES, ) self._bootstrap_budget_bytes = budget @@ -753,15 +820,22 @@ def _aggressive_promotion_capacity(self, current_cap_bytes): allocator_slack = self._worst_shape_allocator_slack_bytes( current_cap_bytes ) - pad = int(self._policy.slack_pad_bytes) + allocator_headroom = int( + self._policy.allocator_cache_headroom_bytes + ) used = 0 capacity = 0 for candidate in candidates: used += int(candidate["block_bytes"]) cumulative = {"block_bytes": used} - if self._worst_shape_candidate_margin_bytes(cumulative) < 0: + if ( + self._worst_shape_candidate_physical_headroom_bytes( + cumulative + ) + < 0 + ): break - if allocator_slack <= used + pad: + if allocator_slack <= used + allocator_headroom: break capacity += 1 return capacity @@ -788,7 +862,7 @@ def _demotion_candidate(self): block_bytes, _order, block_key = max(candidates) return {"block_key": block_key, "block_bytes": block_bytes} - def _worst_shape_candidate_margin_bytes(self, candidate): + def _worst_shape_candidate_physical_headroom_bytes(self, candidate): if candidate is None: return 0 signal = self._signals.last_signal @@ -814,10 +888,6 @@ def _worst_shape_candidate_margin_bytes(self, candidate): - int(signal.get("device_free_bytes", 0) or 0) - int(signal.get("peak_reserved_bytes", 0) or 0), ) - hard_gib = self._config._policy.wddm_hard_gib - hard_bytes = int( - (1.0 if hard_gib is None else max(1.0, float(hard_gib))) * GIB - ) predicted_free = total - ( worst_working + current_resident @@ -825,7 +895,10 @@ def _worst_shape_candidate_margin_bytes(self, candidate): + non_torch + int(candidate["block_bytes"]) ) - return int(predicted_free - hard_bytes) + return int( + predicted_free + - self._training_physical_vram_headroom_bytes() + ) def _worst_shape_allocator_slack_bytes(self, current_cap_bytes): peaks = self._signals.shape_peaks @@ -875,8 +948,10 @@ def _apply_training_policy(self): demote_candidate=demote_candidate, cliff_cap_bytes=cliff_cap, current_cap_bytes=current_cap, - worst_shape_free_bytes=self._worst_shape_candidate_margin_bytes( - candidate + worst_shape_free_bytes=( + self._worst_shape_candidate_physical_headroom_bytes( + candidate + ) ), worst_shape_allocator_slack_bytes=( self._worst_shape_allocator_slack_bytes(current_cap) @@ -899,20 +974,24 @@ def _apply_training_policy(self): decision.action == "rollback" and decision.target_cap_bytes is not None ): - allocator_cap.apply_wddm_hard_allocator_cap( + allocator_cap.configure_wddm_allocator_guard( self._device, self._config._policy.wddm_hard_gib, target_cap_bytes=decision.target_cap_bytes, + strict=getattr( + self._config, "strict_vram_cap", False + ), log_prefix="[ArenaOffload]", ) self._last_training_cap_target_bytes = ( decision.target_cap_bytes ) elif decision.action == "raise_cap": - allocator_cap.apply_wddm_hard_allocator_cap( + allocator_cap.configure_wddm_allocator_guard( self._device, self._config._policy.wddm_hard_gib, target_cap_bytes=decision.target_cap_bytes, + strict=getattr(self._config, "strict_vram_cap", False), log_prefix="[ArenaOffload]", ) self._last_training_cap_target_bytes = decision.target_cap_bytes @@ -938,9 +1017,10 @@ def transition_training_block(self, block_key: str, *, resident: bool) -> dict: return result def _bind_training_cap(self) -> None: - allocator_cap.apply_wddm_hard_allocator_cap( + allocator_cap.configure_wddm_allocator_guard( self._device, self._config._policy.wddm_hard_gib, + strict=getattr(self._config, "strict_vram_cap", False), log_prefix="[ArenaOffload]", ) @@ -982,6 +1062,9 @@ def diagnostics(self) -> dict: "prefetch_depth": int(getattr(self._executor, "depth", 0)), "compile_blocks": bool(self._config.compile_blocks), "compile_dynamic": bool(self._config._compile_dynamic), + "strict_vram_cap": bool( + getattr(self._config, "strict_vram_cap", False) + ), "fp8_forward": bool(self._config.fp8_forward), "fp8_backward": bool(self._config.fp8_backward), "fp8_sampling": bool(self._config.fp8_sampling), @@ -1008,6 +1091,9 @@ def diagnostics(self) -> dict: "bootstrap_complete": self._bootstrap_complete, "bootstrap_min_free_bytes": self._bootstrap_min_free_bytes, "bootstrap_margin_bytes": BOOTSTRAP_MARGIN_BYTES, + "training_physical_vram_headroom_bytes": ( + self._training_physical_vram_headroom_bytes() + ), "bootstrap_min_step": BOOTSTRAP_MIN_STEP, "bootstrap_budget_bytes": self._bootstrap_budget_bytes, "bootstrap_block_keys": self._bootstrap_block_keys, diff --git a/toolkit/memory_management/arena_offload/transfer.py b/toolkit/memory_management/arena_offload/transfer.py index eee81318cd..4248fbe072 100644 --- a/toolkit/memory_management/arena_offload/transfer.py +++ b/toolkit/memory_management/arena_offload/transfer.py @@ -678,7 +678,7 @@ def _register_ordered_effects(): # NOT registered at import time: in torch 2.12 ordered-effect tokens trip an # internal token-erasure assertion inside the checkpoint HOP lowering -# (see tests/test_ingraph_training_ops.py, compiled xfail). Phase 4a S1 keeps +# (see tests/test_ingraph_training_ops.py, compiled xfail). The current path keeps # this as the candidate ordering mechanism for the compiled trunk; call it # explicitly once the HOP interaction is resolved (torch upgrade or flat-trunk # design without the checkpoint HOP). diff --git a/toolkit/memory_management/canonical_arena.py b/toolkit/memory_management/canonical_arena.py index 43955924ff..ce2dd02dc9 100644 --- a/toolkit/memory_management/canonical_arena.py +++ b/toolkit/memory_management/canonical_arena.py @@ -1,13 +1,14 @@ -"""Canonical host arena (Slice 1, tasks/open/IMMUTABLE_TRANSFER_ARENA_PLAN.md). +"""Canonical host arena for immutable frozen base weights. One page-exclusive pinned host flat per block, immutable leaf metadata, and a ONE-TIME repoint of frozen base Parameters into views over those flats -(Invariant 4). This is deliberately NOT the legacy ``pinned_arena.py``: +(the canonical storage invariant). This is deliberately NOT the legacy +``pinned_arena.py``: no generation counter, no ``is_current``/staleness oracle over live module storage, no invalidate/restore/rebuild path, no borrowed-vs-owned pack -taxonomy. Per the plan's Decision section, promotion/demotion/sampling +taxonomy. Promotion, demotion, and sampling transitions must never repoint a Parameter again once ``canonicalize()`` -has run -- that is the job of the residency sidecars (Slice 3), not this +has run -- that is the job of the residency sidecars, not this module. Construction is destination-first and transactional: preparation allocates @@ -29,8 +30,7 @@ class CanonicalArenaError(ValueError): - """A build/canonicalize/guard operation violated a canonical-arena - invariant (tasks/open/IMMUTABLE_TRANSFER_ARENA_PLAN.md).""" + """A build, canonicalize, or guard operation violated an arena invariant.""" def _entry_module_and_leaves(entry): @@ -46,11 +46,12 @@ def _entry_module_and_leaves(entry): def _assert_entries_frozen(entries) -> None: - """Invariant 4/Decision: the canonical arena is for FROZEN base weights - only; LoRA/adapters stay ordinary trainable Parameters outside it. Run - this over every block BEFORE any block is built, so a trainable leaf - anywhere in the batch fails closed without repointing a single - Parameter (no partial canonicalization from this particular cause).""" + """Require frozen base weights before any Parameter is repointed. + + LoRA and adapter parameters remain ordinary trainable state outside the + arena. Checking the full batch first prevents partial canonicalization when + any managed leaf is trainable. + """ for entry in entries: name, _module, weight, bias = _entry_module_and_leaves(entry) if getattr(weight, "requires_grad", False): @@ -64,7 +65,7 @@ class BlockRecord: """One block's immutable canonical host representation. No residency state, no generation counter, no live-module currentness - test -- packs ARE views over arena records (plan Target Components #1). + test -- packs are views over arena records. """ block_key: str @@ -117,7 +118,7 @@ def __init__(self) -> None: def canonicalized(self) -> bool: return self._canonicalized - # -- canonicalize (Invariant 3 + 4) ------------------------------------ + # -- canonicalize ------------------------------------------------------- def canonicalize( self, entries_by_block: dict, *, kind: str = ARENA_KIND @@ -128,7 +129,7 @@ def canonicalize( Caller sequencing responsibility (this method cannot see it): run AFTER load/quantize/freeze and BEFORE LoRA attach, optimizer - construction, or compile (Invariant 4) -- once Parameters are + construction, or compile -- once Parameters are repointed here, nothing may replace them again for the life of the arena. @@ -136,12 +137,12 @@ def canonicalize( any block cannot be admitted (trainable leaf, unsupported quant wrapper, or pin budget exceeded) -- there is no silent pageable fallback in this arena (that is an admission-policy decision for - the caller, Invariant 10, not a mechanism this class provides). + the caller, not a mechanism this class provides). """ if self._canonicalized: raise CanonicalArenaError( "canonical_arena_double_canonicalize: canonicalize() may " - "only run once per arena instance (Invariant 4) -- runtime " + "only run once per arena instance -- runtime " "promotion/demotion/sampling transitions must never " "repoint a Parameter again" ) @@ -169,7 +170,7 @@ def prepare(self, entries_by_block: dict, *, model=None, kind: str = ARENA_KIND) _assert_entries_frozen(entries) return PreparedCanonicalBuild(self, normalized, model=model, kind=kind) - # -- whole-model .to() interception (Invariant 5) ---------------------- + # -- whole-model .to() interception ------------------------------------ @staticmethod def guard_whole_model_to(model: torch.nn.Module) -> None: @@ -265,10 +266,11 @@ def immutable_signature(self) -> tuple: # -- explicit unload ------------------------------------------------ def release(self) -> None: - """Release every block's pin registration and every committed - byte (Invariant 1: pin_manager is the sole authority, every - registered byte is released explicitly). Safe to call on a - partially-built or already-released arena.""" + """Release every block pin and every committed byte explicitly. + + ``pin_manager`` is the sole pin authority. Safe to call on a partially + built or already released arena. + """ for record in self._blocks.values(): pin_manager.unregister_arena_storage(record.pack.host_flat) release_pack(record.pack) diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index 3ea6aa9f13..b6017d29a4 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -109,7 +109,7 @@ def attach( offload_percent: float = 1.0, ignore_modules: list[torch.nn.Module] = [] ): - allocator_cap.apply_wddm_hard_allocator_cap(device) + allocator_cap.configure_wddm_allocator_guard(device, strict=False) if hasattr(module, "_memory_manager"): # already attached return diff --git a/toolkit/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py index b58f2e95a4..54a26d06b6 100644 --- a/toolkit/memory_management/vram_budget.py +++ b/toolkit/memory_management/vram_budget.py @@ -90,7 +90,7 @@ def reconcile_free_bytes(driver_free_bytes, physical_free_bytes) -> int: # same arithmetic reason it would be on the real small card. # # It does not shrink the physical card: the allocator cap is what actually makes -# an over-plan fail, and `_apply_wddm_hard_allocator_cap` converts the simulated +# an over-plan fail, and `configure_wddm_allocator_guard` converts the simulated # cap bytes back into a fraction of the REAL total before handing it to torch. _SIMULATED_CARD_BYTES: int | None = None @@ -386,77 +386,16 @@ def format(self) -> str: ) -@dataclass(frozen=True) -class WddmMargins: - """Dedicated-cliff margins for one phase (training attach / sampling start). - - Resolve ONCE at the phase boundary and pass by value; do not re-read env - vars mid-phase (they cannot change mid-run, and re-reads hide which value - actually governed a decision). - """ - - hard_gib: float - margin_gib: float - source: str # "config" | "env" | "auto" - - @property - def hard_bytes(self) -> int: - return int(self.hard_gib * GIB) - - @property - def margin_bytes(self) -> int: - return int(self.margin_gib * GIB) - - def format(self) -> str: - return ( - f"wddm_hard={self.hard_gib:.2f} GiB " - f"wddm_margin={self.margin_gib:.2f} GiB ({self.source})" - ) - - -def auto_margin_gib(device, pct: float = 0.10, floor_gib: float = 1.0) -> float: - """Auto planning margin: max(floor, pct * card size).""" +def auto_physical_vram_headroom_gib( + device, pct: float = 0.10, floor_gib: float = 1.0 +) -> float: + """Automatic dedicated-VRAM headroom: max(floor, card fraction).""" try: total_bytes = device_total_bytes(device) except Exception: total_bytes = 0 total_gib = max(0.0, float(total_bytes) / GIB) return max(float(floor_gib), float(pct) * total_gib) - - -def resolve_margins( - device, - margin_value, - hard_value, - *, - margin_env: str, - hard_env: str, -) -> WddmMargins: - """Resolve the phase's margins from config value > env > auto. - - ``margin_value`` / ``hard_value`` are the config-supplied values (``None`` - means "consult the env var"; a negative margin or "auto" means auto). - ``margin`` is clamped to at least ``hard``. - """ - hard_gib = float(_env(hard_env, "1.0")) if hard_value is None else float(hard_value) - raw = _env(margin_env, "-1.0") if margin_value is None else margin_value - source = "env" if margin_value is None else "config" - try: - margin_gib = float(raw) - auto = margin_gib < 0 - except (TypeError, ValueError): - auto = str(raw).strip().lower() == "auto" - margin_gib = -1.0 - if auto: - margin_gib = auto_margin_gib(device) - source = "auto" - return WddmMargins( - hard_gib=hard_gib, - margin_gib=max(margin_gib, hard_gib or 0.0), - source=source, - ) - - def cap_fraction(total_bytes, free_bytes, reserved_bytes, hard_gib) -> float: """Allocator-cap fraction so device_used stays <= total - hard (pure). @@ -702,11 +641,9 @@ def estimate_training_working_reserve_bytes( above. At low resolution that flat reserve is generous; at high resolution it is not enough, so the attach-time residency plan keeps too many blocks resident, leaves activations too little headroom, and the run discovers the - shortfall only via a cold-start WDDM-cap-violation storm -- each violation - widens the allocator cap by a fixed ``WDDM_CAP_RELIEF_BYTES`` (0.5 GiB), so - a large resolution jump can cost several wasted/skipped steps before the - cap finally catches up (observed: Krea2 LoKr at 1024x1024 skipped 5/5 fake - steps under the flat default, never reaching a real step). + shortfall only during execution. In strict development mode that becomes + an allocator-cap OOM; production permits WDDM spill, but the resulting + paging is still far slower than choosing the right cold layout. Linear-in-tokens model calibrated on Krea2 LoKr RTX 4070 smoke runs (2026-07-14, ``--block-stream-only`` so zero blocks are resident and @@ -831,7 +768,7 @@ def cap_bytes_for_live( def cap_can_host_promotion( live_bytes, block_bytes, - slack_pad_bytes, + allocator_cache_headroom_bytes, cliff_cap_bytes, *, gc_threshold=GC_THRESHOLD, @@ -839,7 +776,8 @@ def cap_can_host_promotion( """Can the cheap cap lever (tier 1) absorb one more resident block? (pure). Promoting a streamed block to resident raises live by ``block_bytes``. To - keep ``slack_pad_bytes`` of allowance afterward, the GC target must reach + keep ``allocator_cache_headroom_bytes`` of allowance afterward, the GC + target must reach ``live + block + slack``, i.e. the cap must reach ``(live + block + slack) / gc_threshold``. The cap lever can do this only if that target cap is still under the WDDM cliff bound; otherwise the cap is @@ -852,7 +790,7 @@ def cap_can_host_promotion( need_cap = ( float(max(0, int(live_bytes))) + float(max(0, int(block_bytes))) - + float(max(0, int(slack_pad_bytes))) + + float(max(0, int(allocator_cache_headroom_bytes))) ) / float(gc_threshold) return need_cap <= float(int(cliff_cap_bytes)) @@ -861,7 +799,7 @@ def residency_promote_ok( num_alloc_retries, allocator_slack_bytes, block_bytes, - slack_pad_bytes, + allocator_cache_headroom_bytes, ) -> bool: """Sampling climb gate: convert one streamed block to resident? (pure). @@ -871,14 +809,17 @@ def residency_promote_ok( * ``num_alloc_retries == 0`` over the window (nothing cap-binding), AND * worst-shape allocator slack (``0.95 * cap - predicted_live``) exceeds one block plus the pad, so the promotion still leaves - ``slack_pad`` of reusable-cache allowance. + ``allocator_cache_headroom`` of reusable-cache allowance. Both must hold: retries can be zero simply because residency is too low, so the worst-shape allocator-slack test is what proves there is room to spend. """ if int(num_alloc_retries or 0) > 0: return False - return float(allocator_slack_bytes or 0) > float(block_bytes) + float(slack_pad_bytes) + return ( + float(allocator_slack_bytes or 0) + > float(block_bytes) + float(allocator_cache_headroom_bytes) + ) # --- Hysteresis FSM (one transition per phase boundary) --------------------- diff --git a/toolkit/models/DoRA.py b/toolkit/models/DoRA.py index 1d402778fd..346ddda5a7 100644 --- a/toolkit/models/DoRA.py +++ b/toolkit/models/DoRA.py @@ -63,7 +63,7 @@ def __init__( if type(alpha) == torch.Tensor: alpha = float(alpha.detach().float().item()) alpha = self.lora_dim if alpha is None or alpha == 0 else alpha - self.scale = float(alpha) / self.lora_dim + self._set_runtime_scale(float(alpha) / self.lora_dim) # self.register_buffer("alpha", torch.tensor(alpha)) # 定数として扱える eng: treat as constant self.multiplier: Union[float, List[float]] = multiplier diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index 58d75e534e..d53c529211 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -936,13 +936,16 @@ def scale_model_input(model_input, timestep_tensor): f"Batch size of latents {latent_model_input.shape[0]} must be the same or half the batch size of timesteps {timestep.shape[0]}") # predict the noise residual - if self.unet.device != self.device_torch: - try: - self.unet.to(self.device_torch) - except Exception as e: - pass - if self.unet.dtype != self.torch_dtype: - self.unet = self.unet.to(dtype=self.torch_dtype) + from toolkit.memory_management.runtime import get_memory_runtime + + if get_memory_runtime(self.unet) is None: + if self.unet.device != self.device_torch: + try: + self.unet.to(self.device_torch) + except Exception: + pass + if self.unet.dtype != self.torch_dtype: + self.unet = self.unet.to(dtype=self.torch_dtype) # check if get_noise prediction has guidance_embedding_scale # if it does not, we dont pass it @@ -1485,11 +1488,23 @@ def set_device_state(self, state): self.unet.train() else: self.unet.eval() - self.unet.to(state['unet']['device']) - if state['unet']['requires_grad']: - self.unet.requires_grad_(True) + from toolkit.memory_management.runtime import get_memory_runtime + + arena_runtime = get_memory_runtime(self.unet) + if arena_runtime is not None: + # Canonical arena parameters are immutable host views. Only the + # noncanonical transformer state may follow device-state presets; + # residency remains owned by the arena runtime. + unet_device = torch.device(state['unet']['device']) + if unet_device.type == 'cpu': + arena_runtime.park_residency_for_external_phase() + arena_runtime.place_permanent_modules(unet_device) else: - self.unet.requires_grad_(False) + self.unet.to(state['unet']['device']) + if state['unet']['requires_grad']: + self.unet.requires_grad_(True) + else: + self.unet.requires_grad_(False) if isinstance(self.text_encoder, list): for i, encoder in enumerate(self.text_encoder): if isinstance(state['text_encoder'], list): @@ -1533,6 +1548,11 @@ def set_device_state(self, state): self.refiner_unet.train() else: self.refiner_unet.eval() + if arena_runtime is not None and unet_device.type != 'cpu': + # Restore only after inactive components (especially a large text + # encoder) have moved out, avoiding a transient model overlap. + arena_runtime.place_permanent_modules(unet_device) + arena_runtime.restore_residency_after_external_phase() flush() def set_device_state_preset(self, device_state_preset: DeviceStatePreset): @@ -1547,6 +1567,10 @@ def set_device_state_preset(self, device_state_preset: DeviceStatePreset): active_modules = ['vae'] if device_state_preset in ['cache_clip']: active_modules = ['clip'] + if device_state_preset in ['cache_text_encoder']: + active_modules = ['text_encoder'] + if device_state_preset in ['unload']: + active_modules = [] if device_state_preset in ['generate']: active_modules = ['vae', 'unet', 'text_encoder', 'adapter', 'refiner_unet'] diff --git a/toolkit/models/lokr.py b/toolkit/models/lokr.py index 8f11d6c396..158247a64d 100644 --- a/toolkit/models/lokr.py +++ b/toolkit/models/lokr.py @@ -212,7 +212,7 @@ def __init__( if self.use_w2 and self.use_w1: # use scale = 1 alpha = lora_dim - self.scale = float(alpha) / self.lora_dim + self._set_runtime_scale(float(alpha) / self.lora_dim) self.register_buffer('alpha', torch.tensor(alpha)) # treat as constant if self.use_w2: @@ -251,7 +251,7 @@ def get_weight(self, orig_weight=None): (self.lokr_w2 if self.use_w2 else make_weight_cp(self.lokr_t2, self.lokr_w2_a, self.lokr_w2_b) if self.cp else self.lokr_w2_a@self.lokr_w2_b), - self.scale + self._runtime_scale ) if orig_weight is not None: weight = weight.reshape(orig_weight.shape) diff --git a/toolkit/network_mixins.py b/toolkit/network_mixins.py index 2e3bfb3ff1..9650c60d93 100644 --- a/toolkit/network_mixins.py +++ b/toolkit/network_mixins.py @@ -158,7 +158,7 @@ def extract_weight( # set up alphas self.alpha = (self.alpha * 0) + down_weight.shape[0] - self.scale = float(self.alpha.detach().float().item()) / self.lora_dim + self._set_runtime_scale(float(self.alpha.detach().float().item()) / self.lora_dim) # assign them @@ -179,6 +179,19 @@ def __init__( self.is_checkpointing = False self._multiplier: Union[float, list, torch.Tensor] = None + def _set_runtime_scale(self: Module, value) -> None: + """Keep scale metadata as a float and compiled math on the module device.""" + self.scale = float(value) + runtime_scale = getattr(self, "_runtime_scale", None) + if runtime_scale is None: + self.register_buffer( + "_runtime_scale", + torch.tensor(self.scale, dtype=torch.float32), + persistent=False, + ) + else: + runtime_scale.fill_(self.scale) + def _call_forward(self: Module, x): # module dropout if self.module_dropout is not None and self.training: @@ -211,9 +224,9 @@ def _call_forward(self: Module, x): # scaling for rank dropout: treat as if the rank is changed # maskから計算することも考えられるが、augmentation的な効果を期待してrank_dropoutを用いる - scale = self.scale * (1.0 / (1.0 - self.rank_dropout)) # redundant for readability + scale = self._runtime_scale * (1.0 / (1.0 - self.rank_dropout)) # redundant for readability else: - scale = self.scale + scale = self._runtime_scale lx = self.lora_up(lx) diff --git a/toolkit/timer.py b/toolkit/timer.py index e849ba5faa..09779dc8e5 100644 --- a/toolkit/timer.py +++ b/toolkit/timer.py @@ -48,6 +48,8 @@ def print(self): timing_dict = {} # sort by longest at top for timer_name, timings in sorted(self.timers.items(), key=lambda x: sum(x[1]), reverse=True): + if not timings: + continue avg_time = sum(timings) / len(timings) if not is_ui: From ce40dfea98a4b60fe08059419af542b635b0a0c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Tue, 14 Jul 2026 23:41:36 +0200 Subject: [PATCH 05/20] Reject train-time text encoders with arena offload --- jobs/process/BaseSDTrainProcess.py | 5 +++++ tests/test_arena_offload_api.py | 4 ++++ toolkit/memory_management/arena_offload/api.py | 11 ++++++++++- 3 files changed, 19 insertions(+), 1 deletion(-) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index be3fa49ce7..79fe6169fd 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -1690,6 +1690,11 @@ def run(self): validate_arena_training_mode( full_finetune=self.is_fine_tuning, mutates_base_weights=self.train_config.merge_network_on_save, + train_text_encoder=self.train_config.train_text_encoder, + unload_text_encoder=( + self.train_config.unload_text_encoder + or self.is_caching_text_embeddings + ), ) # if the model class has get_train_scheduler static method if hasattr(ModelClass, 'get_train_scheduler'): diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py index 8b397c3959..19fa7f303a 100644 --- a/tests/test_arena_offload_api.py +++ b/tests/test_arena_offload_api.py @@ -105,6 +105,10 @@ def test_training_mode_requires_frozen_immutable_base_weights(self): validate_arena_training_mode(full_finetune=True) with self.assertRaisesRegex(ValueError, "merge_network_on_save"): validate_arena_training_mode(mutates_base_weights=True) + with self.assertRaisesRegex(ValueError, "text encoder during training"): + validate_arena_training_mode(train_text_encoder=True) + with self.assertRaisesRegex(ValueError, "text encoder during training"): + validate_arena_training_mode(unload_text_encoder=False) def test_helpers_are_none_safe(self): self.assertIsNone(get_arena_runtime(None)) diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index cae6057206..5fbe7a9833 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -55,7 +55,11 @@ def validate_arena_training_mode( - *, full_finetune=False, mutates_base_weights=False + *, + full_finetune=False, + mutates_base_weights=False, + train_text_encoder=False, + unload_text_encoder=True, ) -> None: """Reject mutable-base configurations before model loading. @@ -73,6 +77,11 @@ def validate_arena_training_mode( "arena offload requires immutable base transformer weights and " "does not support merge_network_on_save" ) + if train_text_encoder or not unload_text_encoder: + raise ValueError( + "arena offload does not support a text encoder during training; " + "cache text embeddings and unload the text encoder" + ) def unwrap(model): From d5b8271db209d27b42c0f87abbd9cc44f60f5e76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Wed, 15 Jul 2026 01:12:16 +0200 Subject: [PATCH 06/20] Fix arena offload lifecycle correctness --- docs/ARENA_OFFLOAD_CONTRACT.md | 26 ++- tests/test_allocator_cap.py | 38 ++++ tests/test_arena_lifecycle_contract.py | 164 +++++++++++++++++- tests/test_arena_load_session.py | 76 ++++++++ tests/test_arena_offload_api.py | 73 +++++++- tests/test_canonical_arena.py | 38 +++- tests/test_generic_block_dispatcher.py | 25 +++ tests/test_pin_manager.py | 35 ++++ toolkit/memory_management/allocator_cap.py | 32 ++++ .../memory_management/arena_offload/api.py | 9 +- .../arena_offload/construction.py | 44 ++++- .../arena_offload/discovery.py | 6 +- .../memory_management/arena_offload/fp8.py | 1 + .../arena_offload/load_session.py | 8 +- .../arena_offload/resources.py | 141 +++++++++++---- .../arena_offload/runtime.py | 47 +++++ toolkit/memory_management/canonical_arena.py | 83 +++++---- toolkit/memory_management/pin_manager.py | 11 +- toolkit/quantization/fp8_linear.py | 5 + 19 files changed, 765 insertions(+), 97 deletions(-) diff --git a/docs/ARENA_OFFLOAD_CONTRACT.md b/docs/ARENA_OFFLOAD_CONTRACT.md index 250c5fd832..feb8dc66f4 100644 --- a/docs/ARENA_OFFLOAD_CONTRACT.md +++ b/docs/ARENA_OFFLOAD_CONTRACT.md @@ -20,8 +20,9 @@ known boundary with the unmet contract in the error. - Canonical managed leaves are frozen base weights. Trainable adapters remain ordinary state outside canonical storage. - Every selected-block parameter and buffer is enumerable before commit. - Shared, parametrized, missing, or conflicting managed state fails before the - destructive boundary. + Managed leaves must not share one physical storage allocation, including + exact, overlapping, or disjoint views. Shared, parametrized, missing, or + conflicting managed state fails before the destructive boundary. ## Quantization contract @@ -35,6 +36,27 @@ known boundary with the unmet contract in the error. - Transfer and residency code treats declared leaves as opaque tensors. It does not branch on qtype or quantization backend identity. +## Loading and lifecycle contract + +- Direct checkpoint loading may fall back to ordinary `load_state_dict()` only + while the source mapping remains intact. Once a managed source entry has been + consumed, a later build failure aborts that model load instead of reusing the + partial mapping. +- Whole-model `.cpu()` and current-arena `.cuda()` or `.to(device)` requests are + interpreted by the runtime: permanent state follows the requested device and + training residency is parked or restored. Whole-model dtype conversion, + memory-format conversion, other CUDA devices, and arbitrary device + redistribution are unsupported. +- Cleanup is explicit and retryable. Arena-owned simulated-card policy, FP8 + grad-input policy, and Toolkit-tracked allocator fraction are restored before + process ownership is released to a sequential job. + +## Compilation contract + +Arena compilation follows Toolkit's supported model `compile` setting and owns +the one shared block dispatcher used by both training and sampling. Separate +train/sample arena compile policies and caches are not part of this contract. + ## Maintainer validation matrix The upstream gate is the matrix, not an allowlist. Each selected production diff --git a/tests/test_allocator_cap.py b/tests/test_allocator_cap.py index 8ddaab3e11..69130499b4 100644 --- a/tests/test_allocator_cap.py +++ b/tests/test_allocator_cap.py @@ -1,4 +1,5 @@ import pytest +from unittest import mock from toolkit.memory_management import allocator_cap @@ -64,3 +65,40 @@ def test_strict_development_guard_binds_allocator_cap(monkeypatch): assert result == 11 / 12 assert calls == [(11 / 12, 0)] assert allocator_cap.APPLIED_FRACTIONS[0] == 11 / 12 + + +def test_restore_reinstalls_previous_toolkit_fraction(monkeypatch): + calls = [] + monkeypatch.setattr(allocator_cap.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(allocator_cap.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + allocator_cap.torch.cuda, + "set_per_process_memory_fraction", + lambda fraction, index: calls.append((fraction, index)), + ) + allocator_cap.APPLIED_FRACTIONS[0] = 0.75 + + allocator_cap.restore_tracked_allocator_fraction("cuda", 0.5) + + assert calls == [(0.5, 0)] + assert allocator_cap.APPLIED_FRACTIONS[0] == 0.5 + + +def test_restore_removes_arena_fraction_only_after_cuda_succeeds(monkeypatch): + monkeypatch.setattr(allocator_cap.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(allocator_cap.torch.cuda, "current_device", lambda: 0) + setter = mock.Mock(side_effect=(RuntimeError("restore failed"), None)) + monkeypatch.setattr( + allocator_cap.torch.cuda, + "set_per_process_memory_fraction", + setter, + ) + allocator_cap.APPLIED_FRACTIONS[0] = 0.5 + + with pytest.raises(RuntimeError, match="restore failed"): + allocator_cap.restore_tracked_allocator_fraction("cuda", None) + assert allocator_cap.APPLIED_FRACTIONS[0] == 0.5 + + allocator_cap.restore_tracked_allocator_fraction("cuda", None) + assert setter.call_args_list == [mock.call(1.0, 0), mock.call(1.0, 0)] + assert 0 not in allocator_cap.APPLIED_FRACTIONS diff --git a/tests/test_arena_lifecycle_contract.py b/tests/test_arena_lifecycle_contract.py index 0ef89ee5df..9d6e562790 100644 --- a/tests/test_arena_lifecycle_contract.py +++ b/tests/test_arena_lifecycle_contract.py @@ -26,6 +26,11 @@ from toolkit.memory_management.arena_offload.resources import ArenaRuntimeResources from toolkit.memory_management.arena_offload.runtime import ArenaOffloadRuntime from toolkit.memory_management import pin_manager +from toolkit.memory_management import vram_budget +from toolkit.memory_management.arena_offload.fp8 import ( + fp8_grad_input_enabled, + set_fp8_grad_input_enabled, +) from toolkit.memory_management.manager_modules import _DEVICE_STATE from toolkit.memory_management.residency import ResidencyPlan from toolkit.models.base_model import BaseModel @@ -170,6 +175,88 @@ def test_planning_failure_is_precommit_and_preserves_model(): assert not hasattr(model, "_arena_offload_disposed") +def test_precommit_setup_failure_restores_process_policy(): + original_simulated = vram_budget.simulated_card_bytes() + original_fp8 = fp8_grad_input_enabled() + try: + vram_budget.set_simulated_card_bytes(1234) + set_fp8_grad_input_enabled(True) + config = ArenaOffloadConfig( + enabled=True, + fp8_backward=False, + _simulated_vram_gib=0.5, + ) + with mock.patch( + "toolkit.memory_management.arena_offload.runtime.apply_simulated_card", + side_effect=lambda _value, device=None: ( + vram_budget.set_simulated_card_bytes(5678) + ), + ), mock.patch( + "toolkit.memory_management.arena_offload.runtime.build_training_plan", + side_effect=RuntimeError("planning failed"), + ): + with pytest.raises(RuntimeError, match="planning failed"): + prepare_arena_offload( + _frozen_linear(), + device="cpu", + block_names=("blocks",), + config=config, + ) + + assert vram_budget.simulated_card_bytes() == 1234 + assert fp8_grad_input_enabled() + assert active_process_owner() is None + finally: + vram_budget.set_simulated_card_bytes(original_simulated) + set_fp8_grad_input_enabled(original_fp8) + + +def test_postcommit_setup_failure_restores_process_policy(): + original_simulated = vram_budget.simulated_card_bytes() + original_fp8 = fp8_grad_input_enabled() + model = _frozen_linear() + plan = { + "offload_ids": set(), + "protected_training_leaf_keys": frozenset(), + "fits": True, + } + try: + vram_budget.set_simulated_card_bytes(1234) + set_fp8_grad_input_enabled(True) + config = ArenaOffloadConfig( + enabled=True, + fp8_backward=False, + _simulated_vram_gib=0.5, + ) + with mock.patch( + "toolkit.memory_management.arena_offload.runtime.apply_simulated_card", + side_effect=lambda _value, device=None: ( + vram_budget.set_simulated_card_bytes(5678) + ), + ), mock.patch( + "toolkit.memory_management.arena_offload.runtime.build_training_plan", + return_value=plan, + ), mock.patch( + "toolkit.memory_management.arena_offload.runtime.ResidencyState", + side_effect=RuntimeError("postcommit failed"), + ): + with pytest.raises(ArenaSetupFatalError): + prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=config, + ) + + assert vram_budget.simulated_card_bytes() == 1234 + assert fp8_grad_input_enabled() + assert active_process_owner() is None + assert model._arena_offload_disposed + finally: + vram_budget.set_simulated_card_bytes(original_simulated) + set_fp8_grad_input_enabled(original_fp8) + + @pytest.mark.parametrize( "target", ( @@ -227,21 +314,94 @@ def test_resource_release_continues_after_cleanup_error_and_is_idempotent(): release=lambda: calls.append("arena"), ) resources.residency = SimpleNamespace(clear=lambda: calls.append("residency")) + executor_close = mock.Mock(side_effect=(RuntimeError("executor boom"), None)) resources.executor = SimpleNamespace( - active_executions=0, - close=lambda: (_ for _ in ()).throw(RuntimeError("executor boom")), + active_executions=0, close=executor_close ) with pytest.raises(ArenaCleanupError, match="executor boom"): resources.release() resources.release() + assert executor_close.call_count == 2 assert calls == ["residency", "unguard", "arena"] assert resources.released assert resources.disposed assert active_process_owner() is None +def test_arena_release_failure_retains_resource_for_retry(): + model = _frozen_linear() + resources = ArenaRuntimeResources(model, "cpu") + resources.acquire_process_owner() + token = resources.owner_token + arena = SimpleNamespace( + unguard_whole_model_to=mock.Mock(), + release=mock.Mock( + side_effect=(pin_manager.PinReleaseError("unregister failed"), None) + ), + ) + resources.canonical_committed = True + resources.arena = arena + + with pytest.raises(ArenaCleanupError, match="unregister failed"): + resources.release() + + assert resources.arena is arena + assert resources.owner_token is token + assert active_process_owner() is token + resources.release() + assert resources.arena is None + assert resources.released + assert active_process_owner() is None + + +def test_release_restores_arena_owned_process_globals(): + original_simulated = vram_budget.simulated_card_bytes() + original_fp8 = fp8_grad_input_enabled() + try: + vram_budget.set_simulated_card_bytes(1234) + set_fp8_grad_input_enabled(True) + resources = ArenaRuntimeResources(_frozen_linear(), "cpu") + resources.acquire_process_owner() + vram_budget.set_simulated_card_bytes(5678) + set_fp8_grad_input_enabled(False) + + resources.release() + + assert vram_budget.simulated_card_bytes() == 1234 + assert fp8_grad_input_enabled() + assert active_process_owner() is None + finally: + vram_budget.set_simulated_card_bytes(original_simulated) + set_fp8_grad_input_enabled(original_fp8) + + +def test_process_global_restore_failure_is_retryable(): + original_simulated = vram_budget.simulated_card_bytes() + try: + vram_budget.set_simulated_card_bytes(1234) + resources = ArenaRuntimeResources(_frozen_linear(), "cpu") + resources.acquire_process_owner() + token = resources.owner_token + vram_budget.set_simulated_card_bytes(5678) + + with mock.patch( + "toolkit.memory_management.vram_budget.set_simulated_card_bytes", + side_effect=RuntimeError("restore failed"), + ): + with pytest.raises(ArenaCleanupError, match="restore failed"): + resources.release() + + assert resources.owner_token is token + assert active_process_owner() is token + resources.release() + assert vram_budget.simulated_card_bytes() == 1234 + assert active_process_owner() is None + finally: + vram_budget.set_simulated_card_bytes(original_simulated) + + def test_transfer_cleanup_failure_retains_process_owner_until_retry(): model = _frozen_linear() resources = ArenaRuntimeResources(model, "cpu") diff --git a/tests/test_arena_load_session.py b/tests/test_arena_load_session.py index bb0c143927..497844bcca 100644 --- a/tests/test_arena_load_session.py +++ b/tests/test_arena_load_session.py @@ -1,5 +1,7 @@ from types import SimpleNamespace +from unittest import mock +import pytest import torch from toolkit.memory_management.arena_offload import model_load_arena_session @@ -7,6 +9,11 @@ PENDING_CANONICAL_BUILD_ATTR, claim_pending_canonical_build, ) +from toolkit.memory_management.arena_offload.construction import ( + CanonicalBuildError, + CanonicalStateConsumedError, + PreparedCanonicalBuild, +) from toolkit.memory_management.runtime import close_memory_runtime_preparation @@ -85,3 +92,72 @@ def test_trainable_target_falls_back_to_normal_assignment(): assert session.unsupported_reason is not None assert not hasattr(target, PENDING_CANONICAL_BUILD_ATTR) torch.testing.assert_close(target.blocks[0].linear.weight, expected) + + +def test_first_block_preparation_failure_keeps_mapping_and_falls_back(): + source = frozen_transformer() + state = {key: value.clone() for key, value in source.state_dict().items()} + original_keys = tuple(state) + target = frozen_transformer() + + with mock.patch.object( + PreparedCanonicalBuild, + "add_block", + side_effect=CanonicalBuildError("first block failed"), + ): + with model_load_arena_session(base_model()) as session: + incompatible = target.load_state_dict(state, strict=True, assign=True) + + assert not incompatible.missing_keys + assert not incompatible.unexpected_keys + assert tuple(state) == original_keys + assert session.unsupported_reason == "first block failed" + + +def test_failure_after_first_consumed_block_never_falls_back(): + source = frozen_transformer() + state = {key: value.clone() for key, value in source.state_dict().items()} + target = frozen_transformer() + original = PreparedCanonicalBuild.copy_state_entry + calls = 0 + + def fail_in_second_block(build, source_key, value): + nonlocal calls + calls += 1 + if calls == 3: + raise CanonicalBuildError("second block failed") + return original(build, source_key, value) + + with mock.patch.object( + PreparedCanonicalBuild, + "copy_state_entry", + fail_in_second_block, + ): + with model_load_arena_session(base_model()): + with pytest.raises(CanonicalStateConsumedError) as caught: + target.load_state_dict(state, strict=True, assign=True) + + assert isinstance(caught.value.__cause__, CanonicalBuildError) + assert "blocks.0.linear.weight" not in state + assert "blocks.1.linear.weight" in state + + +def test_final_population_failure_after_consumption_never_falls_back(): + source = frozen_transformer() + state = {key: value.clone() for key, value in source.state_dict().items()} + target = frozen_transformer() + original = PreparedCanonicalBuild._finish_population + + def fail_after_pin(build): + original(build) + raise CanonicalBuildError("final pin failed") + + with mock.patch.object( + PreparedCanonicalBuild, "_finish_population", fail_after_pin + ): + with model_load_arena_session(base_model()): + with pytest.raises(CanonicalStateConsumedError): + target.load_state_dict(state, strict=True, assign=True) + + assert "blocks.0.linear.weight" not in state + assert "blocks.1.linear.weight" not in state diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py index 19fa7f303a..1377408ec9 100644 --- a/tests/test_arena_offload_api.py +++ b/tests/test_arena_offload_api.py @@ -9,6 +9,7 @@ import ast from dataclasses import fields from pathlib import Path +from types import SimpleNamespace import unittest import torch @@ -81,7 +82,7 @@ class _FakeModelConfig: layer_offloading_fp8_forward = True layer_offloading_fp8_grad_input = True layer_offloading_fp8_sampling = True - compile = False + compile = True compile_sample = True train_compile_blocks = False layer_offloading_smart_working_reserve_gb = -1.0 @@ -172,6 +173,67 @@ def __init__(self): self.assertEqual(model.root_token.dtype, torch.float64) self.assertEqual(model.root_buffer.dtype, torch.float64) + def test_whole_model_move_parks_cpu_and_restores_arena_device(self): + model = object() + runtime = object.__new__(ArenaOffloadRuntime) + runtime._model = model + runtime._closed = False + runtime._disposed = False + runtime._device = torch.device("cuda:0") + runtime._executor = SimpleNamespace(active_executions=0) + runtime._device_state_parked_plan = None + runtime._permanent_placement = (torch.device("cuda:0"), torch.float32) + events = [] + + def place(device, dtype=None): + normalized = torch.device(device) + runtime._permanent_placement = (normalized, dtype) + events.append(("place", normalized, dtype)) + + runtime.place_permanent_modules = place + runtime.park_residency_for_external_phase = lambda: events.append( + ("park",) + ) + runtime.restore_residency_after_external_phase = lambda: events.append( + ("restore",) + ) + + self.assertIs(runtime.handle_whole_model_move("cpu"), model) + self.assertIs(runtime.handle_whole_model_move("cuda:0"), model) + self.assertEqual( + events, + [ + ("park",), + ("place", torch.device("cpu"), torch.float32), + ("place", torch.device("cuda:0"), torch.float32), + ("restore",), + ], + ) + + def test_whole_model_move_rejects_unsupported_intent_before_mutation(self): + runtime = object.__new__(ArenaOffloadRuntime) + runtime._model = object() + runtime._closed = False + runtime._disposed = False + runtime._device = torch.device("cuda:0") + runtime._executor = SimpleNamespace(active_executions=0) + runtime._device_state_parked_plan = None + runtime._permanent_placement = (torch.device("cuda:0"), torch.float32) + runtime.place_permanent_modules = unittest.mock.Mock() + + with self.assertRaisesRegex(RuntimeError, "dtype_change"): + runtime.handle_whole_model_move("cuda:0", dtype=torch.float64) + with self.assertRaisesRegex(RuntimeError, "cuda_device"): + runtime.handle_whole_model_move("cuda:1") + with self.assertRaisesRegex(RuntimeError, "memory_format"): + runtime.handle_whole_model_move( + "cuda:0", memory_format=torch.channels_last + ) + runtime._executor.active_executions = 1 + with self.assertRaisesRegex(RuntimeError, "during_execution"): + runtime.handle_whole_model_move("cpu") + runtime.place_permanent_modules.assert_not_called() + class ArenaOffloadConfigTest(unittest.TestCase): def test_from_model_config_maps_the_public_surface(self): @@ -220,6 +282,15 @@ def test_missing_attributes_fall_back_to_defaults(self): self.assertFalse(config.compile_blocks) self.assertEqual(config._policy.prefetch_depth, 3) + def test_dead_compile_aliases_do_not_enable_arena_compile(self): + class DeadAliases: + compile = False + compile_sample = True + train_compile_blocks = True + + config = ArenaOffloadConfig.from_model_config(DeadAliases()) + self.assertFalse(config.compile_blocks) + def test_compatibility_aliases_map_to_internal_policy(self): class Aliases: layer_offloading_smart_headroom_gb = 4.0 diff --git a/tests/test_canonical_arena.py b/tests/test_canonical_arena.py index 7ca914a7aa..a9489b437a 100644 --- a/tests/test_canonical_arena.py +++ b/tests/test_canonical_arena.py @@ -1,6 +1,7 @@ import io import unittest from types import SimpleNamespace +from unittest import mock import torch import torch.nn as nn @@ -126,8 +127,13 @@ def test_guarded_to_raises(self): try: arena.canonicalize({"blocks.0": [("lin", model[0])]}) CanonicalArena.guard_whole_model_to(model) - with self.assertRaises(CanonicalArenaError): - model.to(torch.device("cpu")) + for move in ( + lambda: model.to(torch.device("cpu")), + model.cpu, + model.cuda, + ): + with self.assertRaises(CanonicalArenaError): + move() finally: CanonicalArena.unguard_whole_model_to(model) arena.release() @@ -140,27 +146,41 @@ def test_guard_is_idempotent(self): self.assertIs(model.to, original) CanonicalArena.unguard_whole_model_to(model) - def test_guarded_to_allows_only_idempotent_runtime_placement(self): + def test_guarded_movement_routes_all_entry_points_to_runtime(self): model = nn.Sequential(_linear()) - model._arena_offload_runtime = SimpleNamespace( - _permanent_placement=(torch.device("cpu"), torch.float32) + runtime = SimpleNamespace( + device=torch.device("cuda:0"), + handle_whole_model_move=mock.Mock(return_value=model), ) + model._arena_offload_runtime = runtime CanonicalArena.guard_whole_model_to(model) try: self.assertIs(model.to(torch.device("cpu")), model) - self.assertIs( - model.to(device=torch.device("cpu"), dtype=torch.float32), model + self.assertIs(model.cpu(), model) + self.assertIs(model.cuda(), model) + self.assertEqual(runtime.handle_whole_model_move.call_count, 3) + self.assertEqual( + runtime.handle_whole_model_move.call_args_list[0].args, + (torch.device("cpu"),), + ) + self.assertEqual( + runtime.handle_whole_model_move.call_args_list[1].args, + ("cpu",), + ) + self.assertEqual( + runtime.handle_whole_model_move.call_args_list[2].args, + (torch.device("cuda:0"),), ) - with self.assertRaises(CanonicalArenaError): - model.to(device=torch.device("cpu"), dtype=torch.float64) finally: CanonicalArena.unguard_whole_model_to(model) del model._arena_offload_runtime def test_unguard_restores_normal_to(self): model = nn.Sequential(_linear()) + originals = (model.to, model.cuda, model.cpu) CanonicalArena.guard_whole_model_to(model) CanonicalArena.unguard_whole_model_to(model) + self.assertEqual((model.to, model.cuda, model.cpu), originals) # Ordinary .to() must work again (no canonicalized leaves here). model.to(torch.device("cpu")) diff --git a/tests/test_generic_block_dispatcher.py b/tests/test_generic_block_dispatcher.py index 92f56c503c..82c30db3d0 100644 --- a/tests/test_generic_block_dispatcher.py +++ b/tests/test_generic_block_dispatcher.py @@ -221,6 +221,31 @@ def test_shared_managed_state_is_rejected_before_construction(): discover_blocks(model, container_paths=("blocks",)) +@pytest.mark.parametrize("view_kind", ("exact", "partial", "disjoint")) +def test_same_block_shared_managed_storage_is_rejected(view_kind): + model = _frozen_transformer() + model.blocks[0].other = torch.nn.Linear(4, 4, bias=False) + storage = torch.randn(32) + if view_kind == "exact": + left = storage[:16] + right = storage[:16] + elif view_kind == "partial": + left = storage[:16] + right = storage[4:20] + else: + left = storage[:16] + right = storage[16:32] + model.blocks[0].proj.weight = torch.nn.Parameter( + left.view(4, 4), requires_grad=False + ) + model.blocks[0].other.weight = torch.nn.Parameter( + right.view(4, 4), requires_grad=False + ) + + with pytest.raises(BlockDiscoveryError, match="shared_managed_storage"): + discover_blocks(model, container_paths=("blocks",)) + + def test_checkpointing_rejection_precedes_canonical_commit(): model = _frozen_transformer() original = model.blocks[0].proj.weight diff --git a/tests/test_pin_manager.py b/tests/test_pin_manager.py index a0767c149d..b60bc54698 100644 --- a/tests/test_pin_manager.py +++ b/tests/test_pin_manager.py @@ -90,6 +90,41 @@ def test_release_handle_is_idempotent(self): pin_manager.release(handle) # second release must be a no-op self.assertEqual(pin_manager.total_pinned_bytes(), 0) + def test_failed_registered_release_retains_handle_and_ledger_for_retry(self): + tensor = mock.Mock() + handle = pin_manager.PinHandle( + tensor=tensor, + nbytes=64, + kind="weights", + pinned=True, + mechanism="register", + ) + pin_manager.register_pinned_bytes(64, "weights") + + with mock.patch.object( + pin_manager, "unpin_tensor_in_place", return_value=False + ): + with self.assertRaises(pin_manager.PinReleaseError): + pin_manager.release(handle) + + self.assertTrue(handle.pinned) + self.assertEqual(handle.nbytes, 64) + self.assertIs(handle.tensor, tensor) + self.assertEqual(pin_manager.total_pinned_bytes(), 64) + + def succeed(_tensor, kind): + pin_manager.release_pinned_bytes(64, kind) + return True + + with mock.patch.object( + pin_manager, "unpin_tensor_in_place", side_effect=succeed + ): + pin_manager.release(handle) + + self.assertFalse(handle.pinned) + self.assertEqual(handle.nbytes, 0) + self.assertEqual(pin_manager.total_pinned_bytes(), 0) + class PinConformanceTests(unittest.TestCase): """No direct pinning outside the pin manager (PIN_MANAGER_PLAN S2). diff --git a/toolkit/memory_management/allocator_cap.py b/toolkit/memory_management/allocator_cap.py index 7eec779e5a..3c6eefba52 100644 --- a/toolkit/memory_management/allocator_cap.py +++ b/toolkit/memory_management/allocator_cap.py @@ -13,6 +13,38 @@ APPLIED_FRACTIONS: dict[int, float] = {} +def _cuda_index(device) -> int | None: + if not torch.cuda.is_available(): + return None + dev = torch.device(device if device is not None else "cuda") + if dev.type != "cuda": + return None + return dev.index if dev.index is not None else torch.cuda.current_device() + + +def tracked_allocator_fraction(device) -> float | None: + """Return only allocator policy previously installed by Toolkit.""" + index = _cuda_index(device) + return None if index is None else APPLIED_FRACTIONS.get(index) + + +def restore_tracked_allocator_fraction(device, previous: float | None) -> None: + """Restore a Toolkit-owned allocator setting captured before arena setup.""" + index = _cuda_index(device) + if index is None: + return + current = APPLIED_FRACTIONS.get(index) + if previous is None: + if current is not None and current < 1.0: + torch.cuda.set_per_process_memory_fraction(1.0, index) + APPLIED_FRACTIONS.pop(index, None) + return + previous = float(previous) + if current is None or abs(current - previous) > 1e-12: + torch.cuda.set_per_process_memory_fraction(previous, index) + APPLIED_FRACTIONS[index] = previous + + def wddm_cliff_cap_bytes(device, wddm_hard_gib=None) -> int: diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index 5fbe7a9833..7ec578bcd5 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -195,11 +195,10 @@ def get(name: str, default: Any = None) -> Any: and requested_backward, fp8_sampling=fp8_weights and requested_sampling, - compile_blocks=bool( - get("compile", False) - or get("compile_sample", False) - or get("train_compile_blocks", False) - ), + # Arena execution has one shared block dispatcher for training + # and sampling, so Toolkit's supported model compile setting owns + # both phases. + compile_blocks=bool(get("compile", False)), strict_vram_cap=bool( get("layer_offloading_strict_vram_cap", False) ), diff --git a/toolkit/memory_management/arena_offload/construction.py b/toolkit/memory_management/arena_offload/construction.py index 1b5995e3b9..deeb533d98 100644 --- a/toolkit/memory_management/arena_offload/construction.py +++ b/toolkit/memory_management/arena_offload/construction.py @@ -25,6 +25,10 @@ class CanonicalBuildError(RuntimeError): pass +class CanonicalStateConsumedError(CanonicalBuildError): + """A direct-load source mapping was mutated before construction failed.""" + + class CanonicalStateInferenceError(CanonicalBuildError): pass @@ -353,8 +357,24 @@ def populate_from_state_dict_consuming( self._finish_population() return tuple(consumed) - except Exception: - self.rollback() + except Exception as error: + cleanup_error = None + try: + self.rollback() + except BaseException as rollback_error: + cleanup_error = rollback_error + if consumed: + consumed_error = CanonicalStateConsumedError( + "canonical_state_consumed_before_build_failure" + ) + if cleanup_error is not None: + consumed_error.add_note( + "canonical rollback also failed: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + raise consumed_error from error + if cleanup_error is not None: + raise cleanup_error from error raise def copy_state_entry(self, source_key: str, value) -> bool: @@ -577,14 +597,26 @@ def rollback(self) -> None: module._buffers[target] = value if self.model is not None: self.arena.unguard_whole_model_to(self.model) + first_error = None for block in self.blocks: + try: + pin_manager.release(block.handle) + except BaseException as error: + if first_error is None: + first_error = error + continue if block.pack is not None: pin_manager.unregister_arena_storage(block.flat) - pin_manager.release(block.handle) + self.arena._blocks.pop(block.key, None) block.handle = None - self.arena._blocks.clear() - self.arena._canonicalized = False + self.arena._canonicalized = bool(self.arena._blocks) self._populated = False resources = getattr(self, "_arena_resources", None) - if resources is not None and not resources.canonical_committed: + if ( + first_error is None + and resources is not None + and not resources.canonical_committed + ): resources.release() + if first_error is not None: + raise first_error diff --git a/toolkit/memory_management/arena_offload/discovery.py b/toolkit/memory_management/arena_offload/discovery.py index e5f1f4daf6..53697cb8d7 100644 --- a/toolkit/memory_management/arena_offload/discovery.py +++ b/toolkit/memory_management/arena_offload/discovery.py @@ -123,9 +123,7 @@ def _tensor_identity(tensor): try: return ( tensor.device.type, - tensor.untyped_storage().data_ptr(), - tensor.storage_offset(), - tensor.numel(), + tensor.untyped_storage()._cdata, ) except Exception: return ("object", id(tensor)) @@ -185,7 +183,7 @@ def _audit_state(blocks, block_keys, entries_by_block) -> BlockStateAccounting: for tensor in binding.tensors: identity = _tensor_identity(tensor.tensor) previous = managed_storage.get(identity) - if previous is not None and previous[0] != block_key: + if previous is not None: raise BlockDiscoveryError( f"shared_managed_storage:{previous[1]}:{block_key}.{module_path}" ) diff --git a/toolkit/memory_management/arena_offload/fp8.py b/toolkit/memory_management/arena_offload/fp8.py index 391061cdf0..0f605e702b 100644 --- a/toolkit/memory_management/arena_offload/fp8.py +++ b/toolkit/memory_management/arena_offload/fp8.py @@ -7,6 +7,7 @@ from toolkit.quantization.fp8_linear import ( bind_parameter_operation, declare_fp8_linear, + fp8_grad_input_enabled, set_fp8_grad_input_enabled, ) diff --git a/toolkit/memory_management/arena_offload/load_session.py b/toolkit/memory_management/arena_offload/load_session.py index 52f96f68f8..0b94e1086e 100644 --- a/toolkit/memory_management/arena_offload/load_session.py +++ b/toolkit/memory_management/arena_offload/load_session.py @@ -135,7 +135,11 @@ def try_prepare_canonical_from_state_dict(model, state_dict): if session is None: return None from .api import prepare_canonical_storage_from_state_dict - from .construction import CanonicalBuildError, CanonicalStateInferenceError + from .construction import ( + CanonicalBuildError, + CanonicalStateConsumedError, + CanonicalStateInferenceError, + ) from .discovery import BlockDiscoveryError from ..canonical_arena import CanonicalArenaError @@ -146,6 +150,8 @@ def try_prepare_canonical_from_state_dict(model, state_dict): block_names=session.block_names, device=session.device, ) + except CanonicalStateConsumedError: + raise except ( BlockDiscoveryError, CanonicalArenaError, diff --git a/toolkit/memory_management/arena_offload/resources.py b/toolkit/memory_management/arena_offload/resources.py index 0f41cd1be8..7fee4fea95 100644 --- a/toolkit/memory_management/arena_offload/resources.py +++ b/toolkit/memory_management/arena_offload/resources.py @@ -26,6 +26,15 @@ def __init__(self, model, device) -> None: self.disposed = False self.released = False self._releasing = False + self._movement_restored = False + self._transfer_runtime_released = False + self._process_state_captured = False + self._restore_fp8_policy = False + self._restore_simulated_card = False + self._restore_allocator_fraction = False + self._previous_fp8_grad_input = None + self._previous_simulated_card = None + self._previous_allocator_fraction = None self._original_movement = { name: getattr(model, name, None) for name in ("to", "cuda", "cpu") } @@ -34,6 +43,23 @@ def __init__(self, model, device) -> None: def acquire_process_owner(self) -> None: if self.owner_token is None: self.owner_token = acquire_process_owner(self.device) + try: + from .. import allocator_cap, vram_budget + from .fp8 import fp8_grad_input_enabled + + self._previous_fp8_grad_input = fp8_grad_input_enabled() + self._previous_simulated_card = vram_budget.simulated_card_bytes() + self._previous_allocator_fraction = ( + allocator_cap.tracked_allocator_fraction(self.device) + ) + except BaseException: + release_process_owner(self.owner_token) + self.owner_token = None + raise + self._process_state_captured = True + self._restore_fp8_policy = True + self._restore_simulated_card = True + self._restore_allocator_fraction = True def adopt_canonical_build(self, build) -> None: self.canonical_build = build @@ -93,63 +119,122 @@ def release(self) -> None: (("active execution", RuntimeError("cannot_close_during_execution")),) ) - def attempt(label, operation): + def attempt(label, operation, completed=None): try: operation() except BaseException as error: failures.append((label, error)) + return False + if completed is not None: + completed() + return True from . import transfer try: - transfer_cleanup_failed = False - if self.owner_token is not None: + if self.owner_token is not None and not self._transfer_runtime_released: try: transfer.drain_fetch_runtime(owner_token=self.owner_token) except BaseException as error: - transfer_cleanup_failed = True failures.append(("transfer tickets", error)) if self.executor is not None: - attempt("immutable executor", self.executor.close) + attempt( + "immutable executor", + self.executor.close, + lambda: setattr(self, "executor", None), + ) + retained_fp8_restores = [] for label, restore in reversed(self.fp8_restores): - attempt(label, restore) - self.fp8_restores.clear() + if not attempt(label, restore): + retained_fp8_restores.append((label, restore)) + self.fp8_restores = list(reversed(retained_fp8_restores)) if self.residency is not None: - attempt("resident sidecars", self.residency.clear) + attempt( + "resident sidecars", + self.residency.clear, + lambda: setattr(self, "residency", None), + ) + retained_attributes = [] for owner, name, value in reversed(self.published_attributes): if getattr(owner, name, None) is value: - attempt( + if not attempt( f"published attribute {name}", lambda owner=owner, name=name: delattr(owner, name), - ) - self.published_attributes.clear() + ): + retained_attributes.append((owner, name, value)) + self.published_attributes = list(reversed(retained_attributes)) if self.canonical_committed: if self.arena is not None: + if not self._movement_restored: + attempt( + "movement interception", + lambda: self.arena.unguard_whole_model_to(self.model), + lambda: setattr(self, "_movement_restored", True), + ) + arena = self.arena attempt( - "movement interception", - lambda: self.arena.unguard_whole_model_to(self.model), + "canonical arena", + arena.release, + lambda: setattr(self, "arena", None), ) - attempt("canonical arena", self.arena.release) self.disposed = True self._install_disposed_movement_guard() elif self.canonical_build is not None: - attempt("canonical build", self.canonical_build.rollback) - - if self.owner_token is not None: + build = self.canonical_build + attempt( + "canonical build", + build.rollback, + lambda: setattr(self, "canonical_build", None), + ) + + if self.owner_token is not None and not self._transfer_runtime_released: try: transfer.release_fetch_runtime(self.owner_token) except BaseException as error: - transfer_cleanup_failed = True failures.append(("transfer runtime", error)) - if not transfer_cleanup_failed: - try: - release_process_owner(self.owner_token) - except BaseException as error: - failures.append(("process ownership", error)) - else: - self.owner_token = None + else: + self._transfer_runtime_released = True + + if self._restore_fp8_policy: + from .fp8 import set_fp8_grad_input_enabled + + attempt( + "FP8 grad-input policy", + lambda: set_fp8_grad_input_enabled( + self._previous_fp8_grad_input + ), + lambda: setattr(self, "_restore_fp8_policy", False), + ) + if self._restore_simulated_card: + from ..vram_budget import set_simulated_card_bytes + + attempt( + "simulated card policy", + lambda: set_simulated_card_bytes( + self._previous_simulated_card + ), + lambda: setattr(self, "_restore_simulated_card", False), + ) + if self._restore_allocator_fraction: + from ..allocator_cap import restore_tracked_allocator_fraction + + attempt( + "allocator fraction policy", + lambda: restore_tracked_allocator_fraction( + self.device, self._previous_allocator_fraction + ), + lambda: setattr(self, "_restore_allocator_fraction", False), + ) + + if self.owner_token is not None and not failures: + try: + release_process_owner(self.owner_token) + except BaseException as error: + failures.append(("process ownership", error)) + else: + self.owner_token = None finally: self.closing = False self.released = self.owner_token is None @@ -157,9 +242,7 @@ def attempt(label, operation): if self.runtime is not None: self.runtime._closed = True self.runtime._disposed = bool(self.canonical_committed) - self.canonical_modules = () - self.canonical_build = None - self.residency = None - self.executor = None + if self.released: + self.canonical_modules = () if failures: raise ArenaCleanupError(failures) diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index 956991ca77..a267e74462 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -36,6 +36,7 @@ impossible_training_plan_message, resolve_physical_vram_headroom_gib, ) +from .ownership import normalize_device from .resources import ArenaRuntimeResources RUNTIME_ATTR = "_arena_offload_runtime" @@ -329,6 +330,52 @@ def move(module): except BaseException as error: self._fatal_setup_failure(error) + def handle_whole_model_move( + self, device, *, dtype=None, memory_format=None + ): + """Interpret safe whole-model device intent without moving arena leaves.""" + self._require_open() + if getattr(self._executor, "active_executions", 0): + raise RuntimeError("arena_whole_model_move_during_execution") + if memory_format is not None: + raise RuntimeError("arena_whole_model_memory_format_unsupported") + + placement = self._permanent_placement + if placement is None: + raise RuntimeError("arena_whole_model_move_before_placement") + placed_device, placed_dtype = placement + placed_device = normalize_device(placed_device) + if dtype is not None and dtype != placed_dtype: + raise RuntimeError("arena_whole_model_dtype_change_unsupported") + + target_device = placed_device if device is None else normalize_device(device) + arena_device = normalize_device(self._device) + if target_device.type not in ("cpu", "cuda"): + raise RuntimeError( + f"arena_whole_model_device_unsupported:{target_device.type}" + ) + if target_device.type == "cuda" and target_device != arena_device: + raise RuntimeError( + f"arena_whole_model_cuda_device_unsupported:{target_device}" + ) + + if target_device == placed_device: + if ( + target_device == arena_device + and self._device_state_parked_plan is not None + ): + self.restore_residency_after_external_phase() + return self._model + + if target_device.type == "cpu": + self.park_residency_for_external_phase() + self.place_permanent_modules(target_device, placed_dtype) + return self._model + + self.place_permanent_modules(arena_device, placed_dtype) + self.restore_residency_after_external_phase() + return self._model + @property def device(self): return self._device diff --git a/toolkit/memory_management/canonical_arena.py b/toolkit/memory_management/canonical_arena.py index ce2dd02dc9..f1cbd24211 100644 --- a/toolkit/memory_management/canonical_arena.py +++ b/toolkit/memory_management/canonical_arena.py @@ -174,51 +174,52 @@ def prepare(self, entries_by_block: dict, *, model=None, kind: str = ARENA_KIND) @staticmethod def guard_whole_model_to(model: torch.nn.Module) -> None: - """Forbid whole-model ``.to()``/``.cuda()``/``.cpu()`` on a model - with canonicalized leaves: a model-wide move silently detaches - every Parameter from its arena flat (copy semantics on the full - move), the exact drift the legacy arena's ``restore_view`` existed - to repair after the fact. Idempotent. Callers that need to move - SOME parameters (LoRA, non-canonicalized submodules) must route - through a canonical-arena-aware helper instead of raw ``.to()``.""" + """Route whole-model movement through the published arena runtime.""" if getattr(model, "_mm_canonical_to_guarded", False): return - original_to = model.to - def _guarded_to(*args, **kwargs): + originals = { + name: getattr(model, name) for name in ("to", "cuda", "cpu") + } + + def runtime_authority(): runtime = getattr(model, "_arena_offload_runtime", None) - placement = getattr(runtime, "_permanent_placement", None) - if placement is not None: - device, dtype, _non_blocking, _memory_format = ( - torch._C._nn._parse_to(*args, **kwargs) - ) - placed_device, placed_dtype = placement - same_device = ( - device is not None - and torch.device(device) == torch.device(placed_device) + if runtime is None: + raise CanonicalArenaError( + "canonical_arena_whole_model_move_before_runtime" ) - same_dtype = dtype is None or dtype == placed_dtype - if same_device and same_dtype: - return model - raise CanonicalArenaError( - "canonical_arena_whole_model_to: whole-model .to()/.cuda()/" - ".cpu() is forbidden once canonicalized leaves exist -- it " - "would silently detach every Parameter from its arena flat. " - "Route non-canonicalized regions through their own .to() " - "calls, or move canonicalized weights via the residency " - "sidecar path instead." + return runtime + + def _guarded_to(*args, **kwargs): + device, dtype, _non_blocking, memory_format = ( + torch._C._nn._parse_to(*args, **kwargs) + ) + return runtime_authority().handle_whole_model_move( + device, dtype=dtype, memory_format=memory_format + ) + + def _guarded_cuda(device=None): + runtime = runtime_authority() + return runtime.handle_whole_model_move( + runtime.device if device is None else device ) + def _guarded_cpu(): + return runtime_authority().handle_whole_model_move("cpu") + model.to = _guarded_to + model.cuda = _guarded_cuda + model.cpu = _guarded_cpu model._mm_canonical_to_guarded = True - model._mm_canonical_to_original = original_to + model._mm_canonical_movement_originals = originals @staticmethod def unguard_whole_model_to(model: torch.nn.Module) -> None: - original = getattr(model, "_mm_canonical_to_original", None) - if original is not None: - model.to = original - del model._mm_canonical_to_original + originals = getattr(model, "_mm_canonical_movement_originals", None) + if originals is not None: + for name, original in originals.items(): + setattr(model, name, original) + del model._mm_canonical_movement_originals if hasattr(model, "_mm_canonical_to_guarded"): del model._mm_canonical_to_guarded @@ -271,8 +272,16 @@ def release(self) -> None: ``pin_manager`` is the sole pin authority. Safe to call on a partially built or already released arena. """ - for record in self._blocks.values(): + first_error = None + for block_key, record in tuple(self._blocks.items()): + try: + release_pack(record.pack) + except BaseException as error: + if first_error is None: + first_error = error + continue pin_manager.unregister_arena_storage(record.pack.host_flat) - release_pack(record.pack) - self._blocks.clear() - self._canonicalized = False + self._blocks.pop(block_key, None) + self._canonicalized = bool(self._blocks) + if first_error is not None: + raise first_error diff --git a/toolkit/memory_management/pin_manager.py b/toolkit/memory_management/pin_manager.py index 23242d8a21..c5fb000c6d 100644 --- a/toolkit/memory_management/pin_manager.py +++ b/toolkit/memory_management/pin_manager.py @@ -28,6 +28,10 @@ class PinBudgetExceeded(RuntimeError): """Raised when a must-pin allocation cannot fit the current host-pin budget.""" +class PinReleaseError(RuntimeError): + """Raised when a registered host allocation could not be unpinned.""" + + @dataclass class PinHandle: tensor: torch.Tensor @@ -466,7 +470,12 @@ def release(handle: PinHandle) -> None: if getattr(handle, "mechanism", "alloc") == "register": # unpin_tensor_in_place does the ledger release itself (and the # cudaHostUnregister returns DXGI budget immediately). - unpin_tensor_in_place(handle.tensor, getattr(handle, "kind", None)) + if not unpin_tensor_in_place( + handle.tensor, getattr(handle, "kind", None) + ): + raise PinReleaseError( + f"cudaHostUnregister failed for {getattr(handle, 'kind', 'unknown')}" + ) else: release_pinned_bytes(int(handle.nbytes), getattr(handle, "kind", "unknown")) handle.pinned = False diff --git a/toolkit/quantization/fp8_linear.py b/toolkit/quantization/fp8_linear.py index 0d1a72f55e..7f04b0f72d 100644 --- a/toolkit/quantization/fp8_linear.py +++ b/toolkit/quantization/fp8_linear.py @@ -151,6 +151,11 @@ def set_fp8_grad_input_enabled(enabled: bool) -> None: _FP8_GRAD_INPUT = bool(enabled) +def fp8_grad_input_enabled() -> bool: + """Return the Toolkit-owned process policy for FP8 input gradients.""" + return bool(_FP8_GRAD_INPUT) + + def reference_dequantize_to(tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: try: return tensor.dequantize(output_dtype=dtype) From 3f3bf4b9cd8196ebbb91d369186216ce80d00795 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Wed, 15 Jul 2026 14:53:47 +0200 Subject: [PATCH 07/20] Fix checkpoint replay and compile policy --- tests/test_generic_block_dispatcher.py | 17 ++++++++-- toolkit/compile_utils.py | 9 ----- .../arena_offload/dispatcher.py | 34 +++++++++++++------ 3 files changed, 38 insertions(+), 22 deletions(-) diff --git a/tests/test_generic_block_dispatcher.py b/tests/test_generic_block_dispatcher.py index 82c30db3d0..796f1cef30 100644 --- a/tests/test_generic_block_dispatcher.py +++ b/tests/test_generic_block_dispatcher.py @@ -400,16 +400,19 @@ def installed_forward(self, value, _saved=saved): assert any(streamed) assert not streamed[-1] - def train_once(step): + def train_once(step, *, input_requires_grad=True): transfer_before = transfer.lifetime_fetch_stats()["bytes"] planned = runtime.diagnostics()["accounting"][ "planned_training_h2d_bytes" ] - value = torch.randn(2, 4, device=device, requires_grad=True) + value = torch.randn( + 2, 4, device=device, requires_grad=input_requires_grad + ) with runtime.training_step(shape_key=(2, 4), step_num=step): output = model(value) output.square().mean().backward() - assert value.grad is not None + if input_requires_grad: + assert value.grad is not None assert all(parameter.grad is not None for parameter in adapters) for parameter in adapters: parameter.grad = None @@ -429,8 +432,16 @@ def train_once(step): assert torch.isfinite(sampled).all() runtime._executor.activate(runtime._executor.TRAIN, runtime._training_plan) second = train_once(2) + # ZImage enters some checkpointed block families through frozen setup + # layers, so their first streamed block has no gradient-bearing input. + # Two consecutive steps prove recompute returns those ring slots instead + # of relying on a backward hook that autograd cannot schedule. + third = train_once(3, input_requires_grad=False) + fourth = train_once(4, input_requires_grad=False) assert torch.isfinite(first).all() assert torch.isfinite(second).all() + assert torch.isfinite(third).all() + assert torch.isfinite(fourth).all() close_arena_offload(model) assert active_process_owner() is None diff --git a/toolkit/compile_utils.py b/toolkit/compile_utils.py index 7ef44c876f..e83d5cfa6d 100644 --- a/toolkit/compile_utils.py +++ b/toolkit/compile_utils.py @@ -6,15 +6,6 @@ def configure_cuda_only_inductor() -> None: """Configure CUDA compilation without requiring a CPU toolchain.""" torch._dynamo.config.suppress_errors = False - set_stance = getattr(torch.compiler, "set_stance", None) - if set_stance is not None: - try: - # The first AOT-eager pass preserves checkpointing's memory - # behavior and gives Dynamo real shape evidence before compiling. - set_stance("aot_eager_then_compile") - except (RuntimeError, ValueError): - # Older supported PyTorch builds do not expose this stance. - pass if os.name == "nt": # Inductor otherwise dry-compiles a CPU vector-ISA probe even when the # requested graph is CUDA-only. This does not enable a CPU fallback. diff --git a/toolkit/memory_management/arena_offload/dispatcher.py b/toolkit/memory_management/arena_offload/dispatcher.py index 57c3d55676..fb2d7ece5c 100644 --- a/toolkit/memory_management/arena_offload/dispatcher.py +++ b/toolkit/memory_management/arena_offload/dispatcher.py @@ -296,6 +296,7 @@ def dispatch(self, index, args, kwargs): token = None compact_flat = None training = self._active_mode == self.TRAIN + release_on_backward = False if transfer is not None: if training and any( (source.block_key, leaf) in self.protected_training_leaf_keys @@ -314,7 +315,10 @@ def dispatch(self, index, args, kwargs): guard, ) compact_flat = torch.ops.mm.fetch_wait(token, nbytes) - if training and torch.is_grad_enabled(): + release_on_backward = ( + training and torch.is_grad_enabled() and first.requires_grad + ) + if release_on_backward: first = free_on_backward(first, token) args, kwargs = _replace_tensor_argument( args, kwargs, first_location, first @@ -322,17 +326,27 @@ def dispatch(self, index, args, kwargs): leaf_args = source.assemble_leaf_args(self.residency, compact_flat) self._mark_dispatch_dynamic(first) - output = self._get_dispatch_kernel(index)(leaf_args, args, kwargs) + try: + output = self._get_dispatch_kernel(index)(leaf_args, args, kwargs) + except BaseException: + # Non-reentrant checkpoint replay raises its private early-stop + # control-flow exception as soon as it has regenerated every + # tensor backward requested. That can unwind the block before the + # normal post-forward release below. With no gradient-bearing + # input, frozen streamed state is not a backward dependency and + # there is no free_on_backward node, so release on that unwind. + if token is not None and not release_on_backward: + torch.ops.mm.fetch_free(token) + raise if token is not None: # The first checkpoint pass discards its fetched views, so return - # that slot after forward. Replay runs inside an autograd graph - # task; its token is instead released by free_on_backward after - # the compiled block backward has consumed the substituted state. - if ( - not training - or not torch.is_grad_enabled() - or not _in_backward_graph_task() - ): + # that slot after forward. A replay with a gradient-bearing input + # instead releases through free_on_backward after the compiled + # block backward consumes the substituted state. When no input + # requires gradients there is no backward node to run that hook, + # and the frozen streamed state is not a backward dependency, so + # the replay must also release at forward completion. + if not release_on_backward or not _in_backward_graph_task(): torch.ops.mm.fetch_free_after( token, _first_output_tensor(output) ) From a1e6284e9dc2b18b1777dce709694e7756344870 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Wed, 15 Jul 2026 15:37:02 +0200 Subject: [PATCH 08/20] Make arena pinning residency-aware --- tests/test_canonical_arena.py | 28 ++++++ tests/test_generic_block_dispatcher.py | 16 +++- tests/test_residency.py | 91 +++++++++++++++++++ tests/test_transfer_runtime.py | 2 +- .../memory_management/arena_offload/api.py | 6 +- .../arena_offload/construction.py | 24 ++++- .../memory_management/arena_offload/layout.py | 1 + .../arena_offload/runtime.py | 55 +++++++---- .../arena_offload/transfer.py | 6 +- toolkit/memory_management/canonical_arena.py | 67 ++++++++++++-- .../memory_management/immutable_runtime.py | 88 +++++++++++++++++- toolkit/memory_management/pin_manager.py | 20 ++-- toolkit/memory_management/residency.py | 70 +++++++++++++- 13 files changed, 427 insertions(+), 47 deletions(-) diff --git a/tests/test_canonical_arena.py b/tests/test_canonical_arena.py index a9489b437a..d634a41170 100644 --- a/tests/test_canonical_arena.py +++ b/tests/test_canonical_arena.py @@ -110,6 +110,34 @@ def test_release_returns_pin_ledger_bytes(self): self.assertEqual(arena.committed_pinned_bytes(), 0) self.assertEqual(arena.block_keys(), ()) + def test_pageable_build_can_register_and_unregister_same_storage(self): + layer = _linear(in_f=64, out_f=64, bias=True) + expected = layer.weight.detach().clone() + arena = CanonicalArena() + build = arena.prepare( + {"blocks.0": [("lin", layer)]}, + pin_on_finish=False, + ) + try: + build.populate_from_model() + stats = build.commit() + record = arena.block_record("blocks.0") + pointer = record.host_flat.data_ptr() + self.assertEqual(stats.pinned_bytes, 0) + self.assertFalse(record.pack.pinned) + self.assertEqual(arena.committed_pinned_bytes(), 0) + + if torch.cuda.is_available(): + self.assertTrue(arena.pin_block("blocks.0", required=True)) + self.assertTrue(record.pack.pinned) + self.assertEqual(record.host_flat.data_ptr(), pointer) + self.assertTrue(arena.unpin_block("blocks.0")) + self.assertFalse(record.pack.pinned) + self.assertEqual(record.host_flat.data_ptr(), pointer) + torch.testing.assert_close(layer.weight, expected) + finally: + arena.release() + def test_unknown_block_lookup_returns_none(self): arena = CanonicalArena() try: diff --git a/tests/test_generic_block_dispatcher.py b/tests/test_generic_block_dispatcher.py index 796f1cef30..8e79a89d41 100644 --- a/tests/test_generic_block_dispatcher.py +++ b/tests/test_generic_block_dispatcher.py @@ -384,7 +384,8 @@ def installed_forward(self, value, _saved=saved): block.forward = MethodType(installed_forward, block) adapters.append(block.adapter_gain) runtime.finalize() - accounting = runtime.diagnostics()["accounting"] + diagnostics = runtime.diagnostics() + accounting = diagnostics["accounting"] assert accounting["payload_reconciled"] assert accounting["mixed_residency"] assert accounting["resident_blocks"] >= 1 @@ -399,6 +400,13 @@ def installed_forward(self, value, _saved=saved): ] assert any(streamed) assert not streamed[-1] + streamed_keys = { + f"blocks.{index}" for index, is_streamed in enumerate(streamed) + if is_streamed + } + pinned_keys = set(diagnostics["canonical_pinned_block_keys"]) + assert streamed_keys <= pinned_keys + assert len(pinned_keys - streamed_keys) <= 2 def train_once(step, *, input_requires_grad=True): transfer_before = transfer.lifetime_fetch_stats()["bytes"] @@ -422,6 +430,9 @@ def train_once(step, *, input_requires_grad=True): first = train_once(1) sample_plan = ResidencyPlan.build("sample", ()) runtime._executor.activate(runtime._executor.SAMPLE, sample_plan) + assert runtime._arena.pinned_block_keys() == frozenset( + f"blocks.{index}" for index in range(runtime.block_count) + ) sample_accounting = runtime.diagnostics()["accounting"] sample_transfer_before = transfer.lifetime_fetch_stats()["bytes"] with torch.no_grad(), runtime._executor.execution(runtime._executor.SAMPLE): @@ -431,6 +442,9 @@ def train_once(step, *, input_requires_grad=True): ) assert torch.isfinite(sampled).all() runtime._executor.activate(runtime._executor.TRAIN, runtime._training_plan) + restored_pins = runtime._arena.pinned_block_keys() + assert streamed_keys <= restored_pins + assert len(restored_pins - streamed_keys) <= 2 second = train_once(2) # ZImage enters some checkpointed block families through frozen setup # layers, so their first streamed block has no gradient-bearing input. diff --git a/tests/test_residency.py b/tests/test_residency.py index 13d2b9f963..2d9fa21e15 100644 --- a/tests/test_residency.py +++ b/tests/test_residency.py @@ -10,6 +10,8 @@ ResidencyError, ResidencyPlan, ResidencyState, + ordered_demotion_block_keys, + pin_requirements_for_plan, ) @@ -43,6 +45,95 @@ def test_phase_plan_and_existing_planner_seam(arena_layers): assert plan.fingerprint != ResidencyPlan.build("sample", plan.resident_leaf_keys).fingerprint +def test_pin_plan_keeps_streamed_blocks_and_two_exact_demotion_candidates(): + layers_by_block = { + "blocks.0": {"a": _linear(1)}, + "blocks.1": {"a": torch.nn.Linear(16, 16, bias=False)}, + "blocks.2": {"a": torch.nn.Linear(12, 12, bias=False)}, + "blocks.3": {"a": _linear(4)}, + } + for layers in layers_by_block.values(): + for layer in layers.values(): + layer.requires_grad_(False) + arena = CanonicalArena() + arena.canonicalize( + {key: list(layers.items()) for key, layers in layers_by_block.items()} + ) + try: + state = ResidencyState(arena, "cpu") + resident = { + (block_key, leaf_name) + for block_key, layers in layers_by_block.items() + if block_key != "blocks.3" + for leaf_name in layers + } + plan = ResidencyPlan.build("train", resident) + state.reconcile(plan) + + ordered = ordered_demotion_block_keys(arena, state, plan) + required, reserve = pin_requirements_for_plan( + arena, state, plan, reserve_blocks=2 + ) + + assert ordered[:3] == ("blocks.1", "blocks.2", "blocks.0") + assert required == frozenset({"blocks.3"}) + assert reserve == ("blocks.1", "blocks.2") + finally: + arena.release() + + +def test_protected_resident_block_is_not_a_demotion_pin_candidate(arena_layers): + arena, layers = arena_layers + state = ResidencyState(arena, "cpu") + plan = ResidencyPlan.build( + "train", (("blocks.0", name) for name in layers) + ) + state.reconcile(plan) + required, reserve = pin_requirements_for_plan( + arena, + state, + plan, + protected_leaf_keys=(("blocks.0", "a"),), + ) + assert required == frozenset() + assert reserve == () + + +def test_required_repin_failure_precedes_residency_demotion( + arena_layers, monkeypatch +): + arena, layers = arena_layers + state = ResidencyState(arena, "cpu") + current = ResidencyPlan.build( + "train", (("blocks.0", name) for name in layers) + ) + state.reconcile(current) + block = SimpleNamespace(entries=tuple(layers.items())) + runtime = ImmutableTransformerRuntime( + SimpleNamespace(blocks=(block,)), + state, + blocks=(block,), + block_keys=("blocks.0",), + entries_by_block={"blocks.0": tuple(layers.items())}, + compile_blocks=False, + ) + state.device = torch.device("cuda") + generation = runtime.source_generation + monkeypatch.setattr(arena, "pinned_block_keys", lambda: frozenset()) + + def refuse(*_args, **_kwargs): + raise RuntimeError("synthetic pin refusal") + + monkeypatch.setattr(arena, "pin_block", refuse) + target = ResidencyPlan.build("train", ()) + with pytest.raises(RuntimeError, match="synthetic pin refusal"): + runtime.set_residency_plan(target) + + assert state.plan is current + assert runtime.source_generation == generation + assert state.resident_bytes() > 0 + + def test_runtime_training_transitions_are_whole_block(arena_layers): arena, layers = arena_layers block = SimpleNamespace(entries=tuple(layers.items())) diff --git a/tests/test_transfer_runtime.py b/tests/test_transfer_runtime.py index b35821530f..c6e717d612 100644 --- a/tests/test_transfer_runtime.py +++ b/tests/test_transfer_runtime.py @@ -121,7 +121,7 @@ def test_invalid_ranges_fail_closed(arena_block, ranges, compact_nbytes, match): def test_pageable_non_arena_source_and_unknown_ticket_fail_closed(): host = torch.empty(64, dtype=torch.uint8) ranges = torch.tensor([[0, 0, 64]], dtype=torch.int64) - with pytest.raises(RuntimeError, match="registered canonical arena"): + with pytest.raises(RuntimeError, match="canonical arena"): torch.ops.mm.fetch_start_multi(host, ranges, 64) with pytest.raises(RuntimeError, match="unknown ticket"): torch.ops.mm.fetch_wait(torch.tensor([987654], dtype=torch.int64), 64) diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index 7ec578bcd5..5607f399a2 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -266,7 +266,11 @@ def prepare_canonical_storage( } ) arena = CanonicalArena() - build = arena.prepare(entries, model=transformer) + build = arena.prepare( + entries, + model=transformer, + pin_on_finish=False, + ) if resources is not None: resources.adopt_canonical_build(build) return build diff --git a/toolkit/memory_management/arena_offload/construction.py b/toolkit/memory_management/arena_offload/construction.py index deeb533d98..a492c79539 100644 --- a/toolkit/memory_management/arena_offload/construction.py +++ b/toolkit/memory_management/arena_offload/construction.py @@ -214,10 +214,19 @@ class _PreparedBlock: class PreparedCanonicalBuild: """A prepared arena build whose model publication is atomic.""" - def __init__(self, arena, entries_by_block, *, model=None, kind="weights"): + def __init__( + self, + arena, + entries_by_block, + *, + model=None, + kind="weights", + pin_on_finish: bool = True, + ): self.arena = arena self.model = model self.kind = kind + self.pin_on_finish = bool(pin_on_finish) self.blocks = [] self.destinations = {} self.entries_by_block = {} @@ -497,6 +506,8 @@ def model_source_leaves(self, *, block_key: str | None = None): def _finish_population(self) -> None: for block in self.blocks: + if not self.pin_on_finish: + continue handle = pin_manager.pin_register_commit( block.flat, block.layout.nbytes, self.kind, required=False ) @@ -569,7 +580,7 @@ def commit(self): block.flat, tuple(specs), block.layout.nbytes, - True, + bool(block.handle and block.handle.pinned), pin_handle=block.handle, ) pack.view_maker = make_block_view_maker(pack) @@ -582,7 +593,14 @@ def commit(self): if self.model is not None: self.arena.guard_whole_model_to(self.model) self._committed = True - return CanonicalArenaStats(len(self.blocks), sum(b.layout.nbytes for b in self.blocks)) + return CanonicalArenaStats( + len(self.blocks), + sum( + b.layout.nbytes + for b in self.blocks + if b.handle is not None and b.handle.pinned + ), + ) except Exception: self.rollback() raise diff --git a/toolkit/memory_management/arena_offload/layout.py b/toolkit/memory_management/arena_offload/layout.py index 8d8ac77474..7e50fb602d 100644 --- a/toolkit/memory_management/arena_offload/layout.py +++ b/toolkit/memory_management/arena_offload/layout.py @@ -142,6 +142,7 @@ def release_pack(pack: "BlockPack | None") -> None: return pin_manager.release(pack.pin_handle) pack.pin_handle = None + pack.pinned = False def pack_block_host( diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index a267e74462..a97d02e9c8 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -20,7 +20,11 @@ from .. import allocator_cap from ..canonical_arena import CanonicalArena -from ..residency import ResidencyPlan, ResidencyState +from ..residency import ( + ResidencyPlan, + ResidencyState, + ordered_demotion_block_keys, +) from ..vram_budget import apply_simulated_card from .policy import ( AGGRESSIVE_PROMOTION_MIN_CAPACITY, @@ -145,7 +149,11 @@ def _prepare( } if canonical_build is None: arena = CanonicalArena() - canonical_build = arena.prepare(entries_by_block, model=transformer) + canonical_build = arena.prepare( + entries_by_block, + model=transformer, + pin_on_finish=False, + ) resources.adopt_canonical_build(canonical_build) canonical_build.populate_from_model() else: @@ -215,6 +223,7 @@ def _prepare( selection=selection, **executor_kwargs, ) + executor.reconcile_pin_policy(training_plan) resources.adopt_executor(executor) runtime = cls( @@ -889,24 +898,22 @@ def _aggressive_promotion_capacity(self, current_cap_bytes): def _demotion_candidate(self): plan = getattr(self._residency, "plan", None) or self._training_plan - protected = self._protected_training_blocks() - candidates = [] - for order, block_key in enumerate(self._arena.block_keys()): - record = self._arena.block_record(block_key) - keys = tuple((block_key, name) for name in record.leaf_names) - if block_key in protected or not all( - key in plan.resident_leaf_keys for key in keys - ): - continue - actual = sum( - self._residency.resident_leaf_bytes(key) for key in keys - ) - candidates.append( - (actual or int(record.committed_bytes), -order, str(block_key)) - ) - if not candidates: + ordered = ordered_demotion_block_keys( + self._arena, + self._residency, + plan, + protected_leaf_keys=(self._smart_plan or {}).get( + "protected_training_leaf_keys", () + ), + ) + if not ordered: return None - block_bytes, _order, block_key = max(candidates) + block_key = ordered[0] + record = self._arena.block_record(block_key) + keys = tuple((block_key, name) for name in record.leaf_names) + block_bytes = sum( + self._residency.resident_leaf_bytes(key) for key in keys + ) or int(record.committed_bytes) return {"block_key": block_key, "block_bytes": block_bytes} def _worst_shape_candidate_physical_headroom_bytes(self, candidate): @@ -1082,6 +1089,7 @@ def diagnostics(self) -> dict: (self._smart_plan or {}).get("singleton_resident_bytes", 0) ) accounting = self._execution_accounting(active_plan) + pinned_block_keys = self._arena.pinned_block_keys() selection = getattr(self._executor, "selection", None) state_audit = getattr(selection, "accounting", None) return { @@ -1091,6 +1099,15 @@ def diagnostics(self) -> dict: "resident_bytes": singleton_resident + canonical_resident, "singleton_resident_bytes": singleton_resident, "canonical_resident_bytes": canonical_resident, + "canonical_pinned_bytes": self._arena.committed_pinned_bytes(), + "canonical_pinned_blocks": len(pinned_block_keys), + "canonical_pinned_block_keys": tuple(sorted(pinned_block_keys)), + "demotion_pin_reserve_blocks": int( + getattr(self._executor, "pin_reserve_blocks", 0) + ), + "pin_trim_failures": tuple( + getattr(self._executor, "last_pin_trim_failures", ()) + ), "total_weight_resident_bytes": singleton_resident + canonical_resident, "plan_fingerprint": getattr(active_plan, "fingerprint", None), "checkpoint_owner": "model", diff --git a/toolkit/memory_management/arena_offload/transfer.py b/toolkit/memory_management/arena_offload/transfer.py index 4248fbe072..2dfff9df78 100644 --- a/toolkit/memory_management/arena_offload/transfer.py +++ b/toolkit/memory_management/arena_offload/transfer.py @@ -404,7 +404,11 @@ def _validated_transfer_ranges( raise RuntimeError("mm.fetch_start_multi expected a contiguous host_flat") if not pin_manager.is_arena_backed(host_flat): raise RuntimeError( - "mm.fetch_start_multi expected a registered canonical arena source" + "mm.fetch_start_multi expected a canonical arena source" + ) + if torch.cuda.is_available() and not pin_manager.is_host_pinned(host_flat): + raise RuntimeError( + "mm.fetch_start_multi expected the streamed canonical block to be pinned" ) if ranges.device.type != "cpu" or ranges.dtype != torch.int64: raise RuntimeError("mm.fetch_start_multi expected CPU int64 ranges") diff --git a/toolkit/memory_management/canonical_arena.py b/toolkit/memory_management/canonical_arena.py index f1cbd24211..8fabbb70df 100644 --- a/toolkit/memory_management/canonical_arena.py +++ b/toolkit/memory_management/canonical_arena.py @@ -161,14 +161,27 @@ def canonicalize( build.populate_from_model() return build.commit() - def prepare(self, entries_by_block: dict, *, model=None, kind: str = ARENA_KIND): + def prepare( + self, + entries_by_block: dict, + *, + model=None, + kind: str = ARENA_KIND, + pin_on_finish: bool = True, + ): """Prepare final destinations without mutating model Parameters.""" from toolkit.memory_management.arena_offload.construction import PreparedCanonicalBuild normalized = {key: list(entries) for key, entries in entries_by_block.items()} for entries in normalized.values(): _assert_entries_frozen(entries) - return PreparedCanonicalBuild(self, normalized, model=model, kind=kind) + return PreparedCanonicalBuild( + self, + normalized, + model=model, + kind=kind, + pin_on_finish=pin_on_finish, + ) # -- whole-model .to() interception ------------------------------------ @@ -236,7 +249,50 @@ def block_keys(self) -> tuple[str, ...]: return tuple(self._blocks.keys()) def committed_pinned_bytes(self) -> int: - return sum(record.committed_bytes for record in self._blocks.values()) + return sum( + record.committed_bytes + for record in self._blocks.values() + if record.pack.pinned + ) + + def pinned_block_keys(self) -> frozenset[str]: + return frozenset( + block_key + for block_key, record in self._blocks.items() + if record.pack.pinned + ) + + def pin_block(self, block_key: str, *, required: bool = True, device=None) -> bool: + """Register one populated canonical flat without replacing its storage.""" + record = self._blocks.get(str(block_key)) + if record is None: + raise CanonicalArenaError(f"unknown_canonical_block:{block_key}") + if record.pack.pinned: + return False + handle = pin_manager.pin_register_commit( + record.host_flat, + record.committed_bytes, + ARENA_KIND, + device=device, + required=required, + ) + if not handle.pinned: + return False + record.pack.pin_handle = handle + record.pack.pinned = True + return True + + def unpin_block(self, block_key: str) -> bool: + """Unregister one canonical flat while retaining its populated bytes.""" + record = self._blocks.get(str(block_key)) + if record is None: + raise CanonicalArenaError(f"unknown_canonical_block:{block_key}") + if not record.pack.pinned: + return False + pin_manager.release(record.pack.pin_handle) + record.pack.pin_handle = None + record.pack.pinned = False + return True def stats(self) -> CanonicalArenaStats: return CanonicalArenaStats( @@ -247,8 +303,8 @@ def immutable_signature(self) -> tuple: """Return this arena's immutable host-storage commitment. Process-wide pin-ledger state is deliberately excluded: unrelated - consumers such as the bounce pool may grow or shrink while residency - sidecars change without mutating canonical host storage. + consumers and canonical registration policy may grow or shrink while + residency sidecars change without mutating canonical host storage. """ return ( self._canonicalized, @@ -257,7 +313,6 @@ def immutable_signature(self) -> tuple: block_key, record.host_flat.data_ptr(), record.committed_bytes, - pin_manager.is_host_pinned(record.host_flat), pin_manager.is_arena_backed(record.host_flat), ) for block_key, record in self._blocks.items() diff --git a/toolkit/memory_management/immutable_runtime.py b/toolkit/memory_management/immutable_runtime.py index d7e4e3138b..afe4e4b12b 100644 --- a/toolkit/memory_management/immutable_runtime.py +++ b/toolkit/memory_management/immutable_runtime.py @@ -15,9 +15,11 @@ from toolkit.memory_management import vram_budget from toolkit.memory_management.residency import ( + DEFAULT_DEMOTION_PIN_RESERVE_BLOCKS, ResidencyDelta, ResidencyPlan, ResidencyState, + pin_requirements_for_plan, ) from toolkit.memory_management.transfer_plan import ( BlockTransferPlan, @@ -311,6 +313,8 @@ def __init__( self.owner_token = owner_token self._hint_range_warned: set[tuple] = set() self._arena_signature = self.residency.arena.immutable_signature() + self.pin_reserve_blocks = DEFAULT_DEMOTION_PIN_RESERVE_BLOCKS + self.last_pin_trim_failures: tuple[str, ...] = () self._block_abis = tuple( build_block_abi( @@ -403,9 +407,77 @@ def _assert_arena_stable(self, where: str) -> None: if current != self._arena_signature: raise ImmutableRuntimeError( f"arena_mutated_at_boundary:{where}: canonical host flats " - "or registrations changed across a phase boundary" + "changed across a phase boundary" ) + def _prepare_plan_pins(self, plan: ResidencyPlan): + """Pin target stream sources before any device sidecar is removed.""" + arena = self.residency.arena + previous = arena.pinned_block_keys() + if self.residency.device.type != "cuda": + return previous, set(), set() + required, reserve = pin_requirements_for_plan( + arena, + self.residency, + plan, + protected_leaf_keys=self.protected_training_leaf_keys, + reserve_blocks=self.pin_reserve_blocks, + ) + newly_pinned = set() + keep = set(required) + try: + for block_key in arena.block_keys(): + if block_key not in required or block_key in previous: + continue + arena.pin_block( + block_key, + required=True, + device=self.residency.device, + ) + newly_pinned.add(block_key) + for block_key in reserve: + if block_key in arena.pinned_block_keys(): + keep.add(block_key) + continue + if arena.pin_block( + block_key, + required=False, + device=self.residency.device, + ): + newly_pinned.add(block_key) + keep.add(block_key) + except BaseException: + for block_key in tuple(newly_pinned): + try: + arena.unpin_block(block_key) + except BaseException: + pass + raise + return previous, newly_pinned, keep + + def _trim_plan_pins(self, keep) -> None: + """Best-effort release of resident pins after promotion copies settle.""" + arena = self.residency.arena + if self.residency.device.type != "cuda": + return + self.residency.synchronize_copies() + failures = [] + for block_key in arena.block_keys(): + if block_key in keep or block_key not in arena.pinned_block_keys(): + continue + try: + arena.unpin_block(block_key) + except BaseException as error: + failures.append( + f"{block_key}:{type(error).__name__}:{error}" + ) + self.last_pin_trim_failures = tuple(failures) + + def reconcile_pin_policy(self, plan: ResidencyPlan) -> None: + """Converge registration to streamed blocks plus two known demotions.""" + _previous, _newly_pinned, keep = self._prepare_plan_pins(plan) + self._trim_plan_pins(keep) + def set_compile_dynamic_hints(self, hints) -> None: """Install mark_dynamic hints derived after the runtime was prepared. @@ -439,7 +511,19 @@ def _warn_hint_out_of_range(self, dim, size, lo, hi) -> None: def set_residency_plan(self, plan: ResidencyPlan) -> ResidencyDelta: self._assert_arena_stable("pre_residency_publish") - delta = self._sources.publish(plan) + previous, newly_pinned, keep = self._prepare_plan_pins(plan) + try: + delta = self._sources.publish(plan) + except BaseException: + for block_key in tuple(newly_pinned): + if block_key in previous: + continue + try: + self.residency.arena.unpin_block(block_key) + except BaseException: + pass + raise + self._trim_plan_pins(keep) self._assert_arena_stable("post_residency_publish") self.stats["residency_transitions"] += 1 self.stats["source_generation"] = self._sources.generation diff --git a/toolkit/memory_management/pin_manager.py b/toolkit/memory_management/pin_manager.py index c5fb000c6d..ee09528572 100644 --- a/toolkit/memory_management/pin_manager.py +++ b/toolkit/memory_management/pin_manager.py @@ -400,16 +400,12 @@ def unpin_tensor_in_place(t: torch.Tensor, kind: Optional[str] = None) -> bool: return True -# Storage-base data_ptrs of pinned arena flats. A streamed leaf is a VIEW into -# one of these flats at an offset, so its OWN data_ptr misses the exact-ptr -# _REGISTERED_HOST_PINS table (which keys on the registered flat ptr, not the -# view). But every such view shares the flat's untyped storage, whose base ptr -# is recorded here. The eager/pre-compile streaming pinned-bypass -# (manager_modules._profile_is_pinned / bounce_pool._is_pinned) consults this so -# register-pinned arena views are recognized as pinned and skip bounce staging. -# O(1) storage-base lookup, NOT the rejected per-forward range scan. Refcounted -# so a rebuild that recycles the same storage base ptr (release old flat, alloc -# new) never leaves a transient gap. +# Storage-base data_ptrs owned by canonical arena flats. Registration is now a +# residency policy: streamed blocks and the next known demotion candidates are +# pinned, while other fully resident blocks remain pageable. This registry says +# only "canonical source storage," not "currently pinned"; callers that require +# direct async H2D must also consult is_host_pinned. Refcounting keeps ownership +# correct when an allocator later recycles the same storage base pointer. _ARENA_BACKED_STORAGE_LOCK = threading.Lock() _ARENA_BACKED_STORAGE_PTRS: dict[int, int] = {} @@ -427,7 +423,7 @@ def _storage_base_ptr(t: torch.Tensor) -> Optional[int]: def register_arena_storage(t: torch.Tensor) -> None: - """Mark a pinned arena flat's storage so views into it read as pinned.""" + """Mark a canonical arena flat's storage independently of pinnedness.""" ptr = _storage_base_ptr(t) if ptr is None: return @@ -450,7 +446,7 @@ def unregister_arena_storage(t: torch.Tensor) -> None: def is_arena_backed(t: torch.Tensor) -> bool: - """True if this CPU tensor is a view into a pinned arena flat.""" + """True if this CPU tensor is a view into canonical arena storage.""" ptr = _storage_base_ptr(t) if ptr is None: return False diff --git a/toolkit/memory_management/residency.py b/toolkit/memory_management/residency.py index 114ff6e91c..1280edd5fc 100644 --- a/toolkit/memory_management/residency.py +++ b/toolkit/memory_management/residency.py @@ -13,6 +13,7 @@ import torch +from toolkit.memory_management import pin_manager from toolkit.memory_management.canonical_arena import CanonicalArena from toolkit.memory_management.arena_offload.layout import ( @@ -22,6 +23,7 @@ ) LeafKey = tuple[str, str] +DEFAULT_DEMOTION_PIN_RESERVE_BLOCKS = 2 def _interleave_priority(index: int, count: int) -> float: if count <= 1: @@ -221,6 +223,10 @@ def _canonical_leaf(self, key: LeafKey): def _build_sidecar(self, key: LeafKey) -> ResidentLeaf: block, spec, _module = self._canonical_leaf(key) + non_blocking = ( + self.device.type == "cuda" + and pin_manager.is_host_pinned(block.host_flat) + ) stream_context = ( torch.cuda.stream(self._copy_stream) if self._copy_stream is not None @@ -229,7 +235,7 @@ def _build_sidecar(self, key: LeafKey) -> ResidentLeaf: with torch.no_grad(), stream_context: tensors = tuple( leaf_view(block.host_flat, item).to( - self.device, non_blocking=self.device.type == "cuda" + self.device, non_blocking=non_blocking ) for item in spec.tensors ) @@ -334,5 +340,67 @@ def resident_leaf_bytes(self, key: LeafKey) -> int: def resident_bytes(self) -> int: return sum(sidecar.nbytes for sidecar in self._sidecars.values()) + def synchronize_copies(self) -> None: + """Settle queued promotions before their host sources are unpinned.""" + if self._copy_stream is not None: + self._copy_stream.synchronize() + def clear(self, *, phase: str = "clear") -> ResidencyDelta: return self.reconcile(ResidencyPlan.build(phase, ())) + + +def ordered_demotion_block_keys( + arena: CanonicalArena, + residency: ResidencyState, + plan: ResidencyPlan, + *, + protected_leaf_keys=(), +) -> tuple[str, ...]: + """Fully resident blocks in the controller's deterministic demotion order.""" + protected = frozenset( + (str(block), str(leaf)) for block, leaf in protected_leaf_keys + ) + candidates = [] + for order, block_key in enumerate(arena.block_keys()): + record = arena.block_record(block_key) + leaf_keys = tuple((block_key, name) for name in record.leaf_names) + if any(key in protected for key in leaf_keys) or not all( + key in plan.resident_leaf_keys for key in leaf_keys + ): + continue + payload_bytes = sum( + item.nbytes + for leaf_name in record.leaf_names + for item in record.leaf_spec(leaf_name).tensors + ) + candidates.append( + (-payload_bytes, order, str(block_key)) + ) + return tuple(block_key for _bytes, _order, block_key in sorted(candidates)) + + +def pin_requirements_for_plan( + arena: CanonicalArena, + residency: ResidencyState, + plan: ResidencyPlan, + *, + protected_leaf_keys=(), + reserve_blocks: int = DEFAULT_DEMOTION_PIN_RESERVE_BLOCKS, +) -> tuple[frozenset[str], tuple[str, ...]]: + """Return required streamed pins and optional known demotion reserves.""" + streamed = set() + for block_key in arena.block_keys(): + record = arena.block_record(block_key) + if any( + (block_key, leaf_name) not in plan.resident_leaf_keys + for leaf_name in record.leaf_names + ): + streamed.add(str(block_key)) + ordered = ordered_demotion_block_keys( + arena, + residency, + plan, + protected_leaf_keys=protected_leaf_keys, + ) + reserve = ordered[:max(0, int(reserve_blocks))] + return frozenset(streamed), reserve From 287fb44ec6dfa7b359b4273b710afc0cae147460 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Wed, 15 Jul 2026 18:50:07 +0200 Subject: [PATCH 09/20] Restore arena WDDM allocator control --- jobs/process/BaseSDTrainProcess.py | 40 ++++- tests/test_allocator_cap.py | 57 ++++++- tests/test_arena_offload_api.py | 29 ++++ tests/test_arena_offload_policy.py | 48 ++++++ toolkit/memory_management/allocator_cap.py | 64 +++++--- .../arena_offload/__init__.py | 2 + .../memory_management/arena_offload/api.py | 39 +++++ .../memory_management/arena_offload/policy.py | 14 ++ .../arena_offload/runtime.py | 147 ++++++++++++++++-- 9 files changed, 400 insertions(+), 40 deletions(-) diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 79fe6169fd..780b88adfd 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -1856,14 +1856,31 @@ def run(self): if arena_requested: from toolkit.memory_management.arena_offload import ( ArenaOffloadConfig, + estimate_training_working_reserve_hint_bytes, prepare_arena_offload, ) + training_reserve_hint = ( + estimate_training_working_reserve_hint_bytes( + self.dataset_configs, + batch_size=self.train_config.batch_size, + ) + ) + if training_reserve_hint is not None: + print_acc( + "[ArenaOffload] configured-shape training reserve: " + f"{training_reserve_hint / 1024**3:.2f} GiB" + ) arena_runtime = prepare_arena_offload( unet, device=self.device_torch, block_names=self.sd.get_transformer_block_names(), - config=ArenaOffloadConfig.from_model_config(self.model_config), + config=ArenaOffloadConfig.from_model_config( + self.model_config, + training_working_reserve_hint_bytes=( + training_reserve_hint + ), + ), ) arena_runtime.place_permanent_modules(self.device_torch, dtype) else: @@ -2628,24 +2645,37 @@ def run(self): else: raise # not an OOM; surface real errors if did_oom: - self.num_consecutive_oom += 1 + recoverable_oom = False if arena_runtime is not None: failure = arena_runtime.diagnostics().get( 'last_failure_event' ) if failure is not None: + recoverable_oom = bool( + failure.get('recoverable', False) + ) print_acc( f"[ArenaOffload] training failure: {failure}" ) - if self.num_consecutive_oom > 3: - raise RuntimeError("OOM during training step 3 times in a row, aborting training") + if recoverable_oom: + self.num_consecutive_oom = 0 + else: + self.num_consecutive_oom += 1 + if self.num_consecutive_oom > 3: + raise RuntimeError("OOM during training step 3 times in a row, aborting training") optimizer.zero_grad(set_to_none=True) flush() torch.cuda.ipc_collect() # skip this step and keep going print_acc("") print_acc("################################################") - print_acc(f"# OOM during training step, skipping batch {self.num_consecutive_oom}/3 #") + if recoverable_oom: + print_acc( + "# Allocator guard recovered; skipping this batch " + "and continuing #" + ) + else: + print_acc(f"# OOM during training step, skipping batch {self.num_consecutive_oom}/3 #") print_acc("################################################") print_acc("") else: diff --git a/tests/test_allocator_cap.py b/tests/test_allocator_cap.py index 69130499b4..1923a411d0 100644 --- a/tests/test_allocator_cap.py +++ b/tests/test_allocator_cap.py @@ -7,29 +7,45 @@ @pytest.fixture(autouse=True) def _clear_applied_fractions(): allocator_cap.APPLIED_FRACTIONS.clear() + allocator_cap.RELIEF_BYTES.clear() yield allocator_cap.APPLIED_FRACTIONS.clear() + allocator_cap.RELIEF_BYTES.clear() -def test_production_guard_removes_strict_cap_and_allows_spill(monkeypatch): +def test_production_guard_binds_the_same_allocator_cap(monkeypatch): calls = [] - allocator_cap.APPLIED_FRACTIONS[0] = 0.5 + gib = 1024**3 monkeypatch.setattr(allocator_cap.sys, "platform", "win32") monkeypatch.setattr(allocator_cap.torch.cuda, "is_available", lambda: True) monkeypatch.setattr(allocator_cap.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr(allocator_cap.torch.cuda, "memory_reserved", lambda _i: 0) monkeypatch.setattr( allocator_cap.torch.cuda, "set_per_process_memory_fraction", lambda fraction, index: calls.append((fraction, index)), ) + monkeypatch.setattr( + allocator_cap.vram_budget, "device_total_bytes", lambda _i: 12 * gib + ) + monkeypatch.setattr( + allocator_cap.vram_budget, + "real_device_total_bytes", + lambda _i: 12 * gib, + ) + monkeypatch.setattr( + allocator_cap.vram_budget, + "device_mem_info", + lambda _i: (12 * gib, 12 * gib), + ) result = allocator_cap.configure_wddm_allocator_guard( - "cuda", strict=False + "cuda", 1.0, strict=False ) - assert result is None - assert calls == [(1.0, 0)] - assert 0 not in allocator_cap.APPLIED_FRACTIONS + assert result == 11 / 12 + assert calls == [(11 / 12, 0)] + assert allocator_cap.APPLIED_FRACTIONS[0] == 11 / 12 def test_strict_development_guard_binds_allocator_cap(monkeypatch): @@ -67,6 +83,35 @@ def test_strict_development_guard_binds_allocator_cap(monkeypatch): assert allocator_cap.APPLIED_FRACTIONS[0] == 11 / 12 +def test_only_non_strict_failure_policy_can_relax_a_rejected_cap(monkeypatch): + calls = [] + gib = 1024**3 + allocator_cap.APPLIED_FRACTIONS[0] = 10 / 12 + monkeypatch.setattr(allocator_cap.sys, "platform", "win32") + monkeypatch.setattr(allocator_cap.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(allocator_cap.torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + allocator_cap.vram_budget, + "real_device_total_bytes", + lambda _i: 12 * gib, + ) + monkeypatch.setattr( + allocator_cap.torch.cuda, + "set_per_process_memory_fraction", + lambda fraction, index: calls.append((fraction, index)), + ) + + assert allocator_cap.relieve_wddm_allocator_guard_after_oom( + "cuda", strict=True + ) is False + assert allocator_cap.relieve_wddm_allocator_guard_after_oom( + "cuda", strict=False + ) is True + + assert calls == [((10.5 / 12), 0)] + assert allocator_cap.RELIEF_BYTES[0] == allocator_cap.CAP_RELIEF_BYTES + + def test_restore_reinstalls_previous_toolkit_fraction(monkeypatch): calls = [] monkeypatch.setattr(allocator_cap.torch.cuda, "is_available", lambda: True) diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py index 1377408ec9..a082a4b433 100644 --- a/tests/test_arena_offload_api.py +++ b/tests/test_arena_offload_api.py @@ -16,6 +16,7 @@ from toolkit.memory_management.arena_offload import ( ArenaOffloadConfig, + estimate_training_working_reserve_hint_bytes, get_arena_runtime, is_arena_offloaded, is_memory_managed, @@ -236,6 +237,34 @@ def test_whole_model_move_rejects_unsupported_intent_before_mutation(self): class ArenaOffloadConfigTest(unittest.TestCase): + def test_auto_training_reserve_uses_largest_configured_resolution(self): + datasets = [ + SimpleNamespace(resolution=256), + SimpleNamespace(resolution=512), + SimpleNamespace(resolution=1024), + ] + + hint = estimate_training_working_reserve_hint_bytes(datasets) + config = ArenaOffloadConfig.from_model_config( + _FakeModelConfig(), + training_working_reserve_hint_bytes=hint, + ) + + self.assertIsNotNone(hint) + self.assertAlmostEqual(config._policy.working_reserve_gib, hint / GIB) + self.assertGreater(config._policy.working_reserve_gib, 5.0) + + def test_explicit_training_reserve_overrides_configured_shape_hint(self): + class Manual(_FakeModelConfig): + layer_offloading_smart_working_reserve_gb = 6.0 + + config = ArenaOffloadConfig.from_model_config( + Manual(), + training_working_reserve_hint_bytes=10 * GIB, + ) + + self.assertEqual(config._policy.working_reserve_gib, 6.0) + def test_from_model_config_maps_the_public_surface(self): config = ArenaOffloadConfig.from_model_config(_FakeModelConfig()) diff --git a/tests/test_arena_offload_policy.py b/tests/test_arena_offload_policy.py index 9e0addf24a..672848a70a 100644 --- a/tests/test_arena_offload_policy.py +++ b/tests/test_arena_offload_policy.py @@ -160,6 +160,53 @@ def test_runtime_preserves_shape_peaks_after_layout_change(): assert signals.shape_peaks +def test_pressure_relief_reclaims_cache_then_demotes_enough_blocks(monkeypatch): + snapshots = iter( + [ + { + "device_free_bytes": 100, + "predicted_peak_free_bytes": 200, + "deficit_bytes": 300, + }, + { + "device_free_bytes": 180, + "predicted_peak_free_bytes": 240, + "deficit_bytes": 260, + }, + { + "device_free_bytes": 500, + "predicted_peak_free_bytes": 540, + "deficit_bytes": 0, + }, + ] + ) + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._signals = SimpleNamespace(last_signal={"step_num": 8}) + runtime._policy = ArenaResidencyController() + runtime._last_training_pressure_relief = None + runtime._training_pressure_snapshot = lambda: next(snapshots) + runtime._demotion_candidates = lambda: ( + {"block_key": "blocks.1", "block_bytes": 150}, + {"block_key": "blocks.2", "block_bytes": 150}, + {"block_key": "blocks.3", "block_bytes": 150}, + ) + transitions = [] + runtime.transition_training_blocks = lambda keys, resident: ( + transitions.append((tuple(keys), resident)) + or {"changed": True} + ) + empty_cache_calls = [] + monkeypatch.setattr( + "torch.cuda.empty_cache", lambda: empty_cache_calls.append(True) + ) + + assert runtime._relieve_training_physical_pressure() is True + assert transitions == [(('blocks.1', 'blocks.2'), False)] + assert len(empty_cache_calls) == 2 + assert runtime._policy.last_reason == "physical_pressure_relief" + assert runtime._last_training_pressure_relief["demoted_bytes"] == 300 + + def test_worst_shape_allocator_slack_reconstructs_current_layout(): signals = TrainingSignalWindow() observe( @@ -763,6 +810,7 @@ def test_arena_allocation_failure_drains_and_rolls_back(monkeypatch): assert runtime._policy.last_safe_residency_bytes == 200 assert runtime._policy.last_rejected_residency_bytes == 220 assert runtime._last_failure_event["rollback_block"] == ["blocks.7"] + assert runtime._last_failure_event["recoverable"] is True assert runtime._last_failure_event["abandoned_fetch_tickets"] == 2 assert cap_calls[0][1]["target_cap_bytes"] == 1000 diff --git a/toolkit/memory_management/allocator_cap.py b/toolkit/memory_management/allocator_cap.py index 3c6eefba52..fb2c9f7dfa 100644 --- a/toolkit/memory_management/allocator_cap.py +++ b/toolkit/memory_management/allocator_cap.py @@ -3,7 +3,6 @@ from __future__ import annotations import sys -import warnings import torch @@ -11,6 +10,8 @@ GIB = 1024 ** 3 APPLIED_FRACTIONS: dict[int, float] = {} +RELIEF_BYTES: dict[int, int] = {} +CAP_RELIEF_BYTES = int(0.5 * GIB) def _cuda_index(device) -> int | None: @@ -38,6 +39,7 @@ def restore_tracked_allocator_fraction(device, previous: float | None) -> None: if current is not None and current < 1.0: torch.cuda.set_per_process_memory_fraction(1.0, index) APPLIED_FRACTIONS.pop(index, None) + RELIEF_BYTES.pop(index, None) return previous = float(previous) if current is None or abs(current - previous) > 1e-12: @@ -92,12 +94,12 @@ def configure_wddm_allocator_guard( strict=False, log_prefix="[MemoryManager]", ): - """Configure the WDDM cliff guard for production or development. + """Bind the allocator below the WDDM cliff in every guard mode. - Production treats the cliff as a planning target and permits WDDM spill; - strict development mode binds torch's allocator below it so a breach raises - OOM. Call only at a phase boundary. A simulated governing capacity is - always converted against the physical card total. + ``strict`` controls how the caller handles a cap rejection; it must not + disable the cap or the FSM's allocator steering. Call only at a phase + boundary. A simulated governing capacity is always converted against the + physical card total. """ if sys.platform != "win32" or not torch.cuda.is_available(): return None @@ -105,19 +107,6 @@ def configure_wddm_allocator_guard( if dev.type != "cuda": return None index = dev.index if dev.index is not None else torch.cuda.current_device() - if not strict: - previous = APPLIED_FRACTIONS.pop(index, None) - if previous is not None and previous < 1.0: - try: - torch.cuda.set_per_process_memory_fraction(1.0, index) - except Exception as error: - warnings.warn( - "could not remove the strict CUDA allocator cap; " - f"WDDM spill fallback may remain unavailable: {error}", - RuntimeWarning, - stacklevel=2, - ) - return None try: hard_gib = float(wddm_hard_gib) if wddm_hard_gib is not None else 1.0 except (TypeError, ValueError): @@ -138,6 +127,10 @@ def configure_wddm_allocator_guard( fraction = max(0.1, min(cliff_fraction, target_fraction)) reclaimed = fraction < cliff_fraction - 1e-9 + relief_bytes = RELIEF_BYTES.get(index, 0) + if relief_bytes: + fraction = min(1.0, fraction + relief_bytes / float(total)) + applied = fraction * total / float(real_total) previous = APPLIED_FRACTIONS.get(index) tolerance = (64 * 1024**2) / real_total @@ -154,11 +147,42 @@ def configure_wddm_allocator_guard( ) if total != real_total: source += f"; SIMULATED {total / GIB:.2f} GiB card" + if relief_bytes: + source += f"; +{relief_bytes / GIB:.2f} GiB recovery relief" print( - f"{log_prefix} strict WDDM allocator cap: " + f"{log_prefix} WDDM allocator cap: " f"{fraction * total / GIB:.2f}/{total / GIB:.2f} GiB " f"({source}; margin {hard_gib:.2f} GiB, " f"non_torch {non_torch / GIB:.2f} GiB; allocation beyond this " "recycles cache or raises OOM instead of silently paging)" ) return applied + + +def relieve_wddm_allocator_guard_after_oom( + device, *, strict=False, context="training step", log_prefix="[MemoryManager]" +) -> bool: + """Widen a capped allocator only when non-strict recovery has no layout relief.""" + if strict or sys.platform != "win32" or not torch.cuda.is_available(): + return False + dev = torch.device(device if device is not None else "cuda") + if dev.type != "cuda": + return False + index = dev.index if dev.index is not None else torch.cuda.current_device() + applied = APPLIED_FRACTIONS.get(index) + if applied is None or applied >= 1.0: + return False + + real_total = vram_budget.real_device_total_bytes(index) + relief = RELIEF_BYTES.get(index, 0) + CAP_RELIEF_BYTES + widened = min(1.0, applied + CAP_RELIEF_BYTES / float(real_total)) + torch.cuda.set_per_process_memory_fraction(widened, index) + APPLIED_FRACTIONS[index] = widened + RELIEF_BYTES[index] = relief + print( + f"{log_prefix} allocator cap rejected {context}: no resident layout " + f"relief remained, widening {applied * real_total / GIB:.2f}->" + f"{widened * real_total / GIB:.2f} GiB so the non-strict job can " + "continue" + ) + return True diff --git a/toolkit/memory_management/arena_offload/__init__.py b/toolkit/memory_management/arena_offload/__init__.py index a3002345f9..433f3bed4f 100644 --- a/toolkit/memory_management/arena_offload/__init__.py +++ b/toolkit/memory_management/arena_offload/__init__.py @@ -20,6 +20,7 @@ from .api import ( ArenaOffloadConfig, close_arena_offload, + estimate_training_working_reserve_hint_bytes, get_arena_runtime, is_arena_offloaded, is_memory_managed, @@ -48,6 +49,7 @@ "get_arena_runtime", "get_memory_runtime", "discover_blocks", + "estimate_training_working_reserve_hint_bytes", "is_arena_offloaded", "is_memory_managed", "memory_runtime_owns_compile", diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index 5607f399a2..686e97b4f9 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -54,6 +54,45 @@ } +def estimate_training_working_reserve_hint_bytes( + dataset_configs, *, batch_size: int = 1 +) -> int | None: + """Estimate the worst configured image-training working set. + + Dataset ``resolution`` is an area target: a 1024 bucket is approximately + 1024**2 pixels regardless of aspect ratio. Diffusion transformers normally + see one token per 16x16 image pixels after VAE and patch compression. + Planning from the largest configured bucket prevents earlier low-resolution + steps from licensing residency that the later bucket cannot support. + """ + max_image_tokens = 0 + for dataset in dataset_configs or (): + resolution = getattr(dataset, "resolution", 0) + if isinstance(resolution, Sequence) and not isinstance( + resolution, (str, bytes) + ): + candidates = resolution + else: + candidates = (resolution,) + for candidate in candidates: + try: + side = max(0, int(candidate)) + except (TypeError, ValueError): + continue + image_tokens = (side * side + 255) // 256 + max_image_tokens = max(max_image_tokens, image_tokens) + if max_image_tokens <= 0: + return None + + from ..vram_budget import estimate_training_working_reserve_bytes + + batch = max(1, int(batch_size or 1)) + return estimate_training_working_reserve_bytes( + max_image_tokens * batch, + text_tokens=512 * batch, + ) + + def validate_arena_training_mode( *, full_finetune=False, diff --git a/toolkit/memory_management/arena_offload/policy.py b/toolkit/memory_management/arena_offload/policy.py index 6cc682a8de..7c315a19a2 100644 --- a/toolkit/memory_management/arena_offload/policy.py +++ b/toolkit/memory_management/arena_offload/policy.py @@ -300,6 +300,20 @@ def _begin_promotion( def allocation_failure(self): return self.reject_pending_promotion("promotion_allocation_failure") + def record_physical_pressure_relief(self, block_keys, block_bytes): + """Put the controller in cooldown after emergency cache/layout relief.""" + keys = tuple(str(key) for key in block_keys) + self.pending_promotion = None + self.last_promoted_key = None + self.state = vram_budget.ResidencyFsmState( + vram_budget.FSM_COOLDOWN, 0 + ) + self.last_action = "demote" if keys else "hold" + self.last_reason = "physical_pressure_relief" + self.last_block_key = keys[-1] if keys else None + self.last_block_bytes = int(block_bytes) + self.last_target_cap_bytes = None + def _decision(self, action, candidate, *, reason): self.last_action = action self.last_reason = reason diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index a97d02e9c8..cd32e12120 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -90,6 +90,7 @@ def __init__( self._last_failure_event: dict | None = None self._policy = ArenaResidencyController() self._last_training_cap_target_bytes: int | None = None + self._last_training_pressure_relief: dict | None = None self._bootstrap_complete = False self._bootstrap_min_free_bytes: int | None = None self._bootstrap_budget_bytes = 0 @@ -597,6 +598,7 @@ def _handle_training_failure(self, error, *, shape_key, step_num): or "out of memory" in text.lower() ) rollback = None + recoverable = False if allocation_failure: decision = self._policy.allocation_failure() if decision.action == "rollback" and decision.block_key is not None: @@ -627,6 +629,35 @@ def _handle_training_failure(self, error, *, shape_key, step_num): if decision.block_keys else decision.block_key ) + torch.cuda.empty_cache() + recoverable = not bool( + getattr(self._config, "strict_vram_cap", False) + ) + elif not bool( + getattr(self._config, "strict_vram_cap", False) + ): + candidates = self._demotion_candidates() + if candidates: + candidate = candidates[0] + self.transition_training_block( + candidate["block_key"], resident=False + ) + rollback = candidate["block_key"] + self._policy.record_physical_pressure_relief( + (candidate["block_key"],), + candidate["block_bytes"], + ) + torch.cuda.empty_cache() + recoverable = True + else: + recoverable = ( + allocator_cap.relieve_wddm_allocator_guard_after_oom( + self._device, + strict=False, + context="training step", + log_prefix="[ArenaOffload]", + ) + ) try: stats = torch.cuda.memory_stats(self._device) peak_allocated = int(torch.cuda.max_memory_allocated(self._device)) @@ -671,6 +702,7 @@ def _handle_training_failure(self, error, *, shape_key, step_num): ), ), "rollback_block": rollback, + "recoverable": bool(recoverable), "rejected_residency_bytes": ( self._policy.last_rejected_residency_bytes ), @@ -897,6 +929,10 @@ def _aggressive_promotion_capacity(self, current_cap_bytes): return capacity def _demotion_candidate(self): + candidates = self._demotion_candidates() + return candidates[0] if candidates else None + + def _demotion_candidates(self): plan = getattr(self._residency, "plan", None) or self._training_plan ordered = ordered_demotion_block_keys( self._arena, @@ -906,15 +942,103 @@ def _demotion_candidate(self): "protected_training_leaf_keys", () ), ) - if not ordered: - return None - block_key = ordered[0] - record = self._arena.block_record(block_key) - keys = tuple((block_key, name) for name in record.leaf_names) - block_bytes = sum( - self._residency.resident_leaf_bytes(key) for key in keys - ) or int(record.committed_bytes) - return {"block_key": block_key, "block_bytes": block_bytes} + candidates = [] + for block_key in ordered: + record = self._arena.block_record(block_key) + keys = tuple((block_key, name) for name in record.leaf_names) + block_bytes = sum( + self._residency.resident_leaf_bytes(key) for key in keys + ) or int(record.committed_bytes) + candidates.append( + {"block_key": block_key, "block_bytes": block_bytes} + ) + return tuple(candidates) + + def _training_pressure_snapshot(self): + import torch + + from .. import vram_budget + + total = int(vram_budget.device_total_bytes(self._device)) + free = int(vram_budget.device_free_bytes(self._device)) + reserved = int(torch.cuda.memory_reserved(self._device)) + signal = self._signals.last_signal or {} + peak_allocated = int(signal.get("peak_allocated_bytes", 0) or 0) + for peak in self._signals.shape_peaks.values(): + if peak.steps > 0: + peak_allocated = max( + peak_allocated, int(peak.peak_allocated_bytes) + ) + non_torch = max(0, total - free - reserved) + predicted_free = total - peak_allocated - non_torch + headroom = self._training_physical_vram_headroom_bytes() + governing_free = min(free, predicted_free) + return { + "total_bytes": total, + "device_free_bytes": free, + "torch_reserved_bytes": reserved, + "peak_allocated_bytes": peak_allocated, + "non_torch_bytes": non_torch, + "predicted_peak_free_bytes": predicted_free, + "headroom_bytes": headroom, + "deficit_bytes": max(0, headroom - governing_free), + } + + def _relieve_training_physical_pressure(self) -> bool: + """Reclaim cache, then demote enough residents to preserve the floor.""" + import torch + + if self._signals.last_signal is None: + return False + before = self._training_pressure_snapshot() + if before["deficit_bytes"] <= 0: + return False + + torch.cuda.empty_cache() + after_cache = self._training_pressure_snapshot() + selected = [] + selected_bytes = 0 + remaining = int(after_cache["deficit_bytes"]) + if remaining > 0: + for candidate in self._demotion_candidates(): + selected.append(candidate["block_key"]) + selected_bytes += int(candidate["block_bytes"]) + if selected_bytes >= remaining: + break + if selected: + result = self.transition_training_blocks( + selected, resident=False + ) + if not result.get("changed"): + selected = [] + selected_bytes = 0 + torch.cuda.empty_cache() + + after = self._training_pressure_snapshot() + self._policy.record_physical_pressure_relief( + selected, selected_bytes + ) + self._last_training_pressure_relief = { + "before": before, + "after_cache": after_cache, + "after": after, + "demoted_block_keys": tuple(selected), + "demoted_bytes": selected_bytes, + } + print( + "[ArenaOffload] WDDM pressure relief: " + f"cache_reclaimed=" + f"{max(0, after_cache['device_free_bytes'] - before['device_free_bytes']) / GIB:.2f} GiB, " + f"demoted_blocks={len(selected)}, " + f"demoted={selected_bytes / GIB:.2f} GiB, " + f"predicted_free=" + f"{before['predicted_peak_free_bytes'] / GIB:.2f}->" + f"{after['predicted_peak_free_bytes'] / GIB:.2f} GiB, " + f"device_free=" + f"{before['device_free_bytes'] / GIB:.2f}->" + f"{after['device_free_bytes'] / GIB:.2f} GiB" + ) + return True def _worst_shape_candidate_physical_headroom_bytes(self, candidate): if candidate is None: @@ -983,6 +1107,8 @@ def _apply_training_policy(self): if torch.device(self._device).type != "cuda" or not torch.cuda.is_available(): return + if self._relieve_training_physical_pressure(): + return candidate = self._promotion_candidate() demote_candidate = self._demotion_candidate() cliff_cap = allocator_cap.wddm_cliff_cap_bytes( @@ -1152,6 +1278,9 @@ def diagnostics(self) -> dict: "training_cap_target_bytes": getattr( self, "_last_training_cap_target_bytes", None ), + "last_training_pressure_relief": getattr( + self, "_last_training_pressure_relief", None + ), "bootstrap_complete": self._bootstrap_complete, "bootstrap_min_free_bytes": self._bootstrap_min_free_bytes, "bootstrap_margin_bytes": BOOTSTRAP_MARGIN_BYTES, From 3fa1435bcb408b7218c4fe555f652a5ec9891182 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Wed, 15 Jul 2026 18:55:44 +0200 Subject: [PATCH 10/20] Rebind arena allocator cap per step --- tests/test_arena_offload_policy.py | 24 ++++++++++++++++++- .../arena_offload/runtime.py | 2 ++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/test_arena_offload_policy.py b/tests/test_arena_offload_policy.py index 672848a70a..e82885a578 100644 --- a/tests/test_arena_offload_policy.py +++ b/tests/test_arena_offload_policy.py @@ -263,15 +263,37 @@ def test_training_cap_binding_uses_configured_guard_mode(monkeypatch): strict_vram_cap=False, _policy=SimpleNamespace(wddm_hard_gib=1.25), ) + runtime._last_training_cap_target_bytes = 900 runtime._bind_training_cap() assert calls == [ ( "cuda:1", 1.25, - {"strict": False, "log_prefix": "[ArenaOffload]"}, + { + "target_cap_bytes": 900, + "strict": False, + "log_prefix": "[ArenaOffload]", + }, ) ] + + +def test_training_policy_rebinds_cap_before_pressure_relief(monkeypatch): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._device = "cuda" + events = [] + runtime._bind_training_cap = lambda: events.append("bind") + runtime._relieve_training_physical_pressure = lambda: ( + events.append("relieve") or True + ) + monkeypatch.setattr("torch.cuda.is_available", lambda: True) + + runtime._apply_training_policy() + + assert events == ["bind", "relieve"] + + def test_shape_working_peak_excludes_residency(): window = TrainingSignalWindow() observe(window, peak_allocated_bytes=900, resident_bytes=200) diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index cd32e12120..e261460de9 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -1107,6 +1107,7 @@ def _apply_training_policy(self): if torch.device(self._device).type != "cuda" or not torch.cuda.is_available(): return + self._bind_training_cap() if self._relieve_training_physical_pressure(): return candidate = self._promotion_candidate() @@ -1200,6 +1201,7 @@ def _bind_training_cap(self) -> None: allocator_cap.configure_wddm_allocator_guard( self._device, self._config._policy.wddm_hard_gib, + target_cap_bytes=self._last_training_cap_target_bytes, strict=getattr(self._config, "strict_vram_cap", False), log_prefix="[ArenaOffload]", ) From 6b215dc106a27c38ebcb38c4afcffdef20293721 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Wed, 15 Jul 2026 19:03:50 +0200 Subject: [PATCH 11/20] Preserve allocator GC feedback --- tests/test_arena_offload_policy.py | 9 ++---- tests/test_residency_two_timescale.py | 12 ++++++++ .../memory_management/arena_offload/policy.py | 2 +- .../arena_offload/runtime.py | 29 +++++++------------ toolkit/memory_management/vram_budget.py | 6 ++-- 5 files changed, 30 insertions(+), 28 deletions(-) diff --git a/tests/test_arena_offload_policy.py b/tests/test_arena_offload_policy.py index e82885a578..d052160e85 100644 --- a/tests/test_arena_offload_policy.py +++ b/tests/test_arena_offload_policy.py @@ -160,7 +160,7 @@ def test_runtime_preserves_shape_peaks_after_layout_change(): assert signals.shape_peaks -def test_pressure_relief_reclaims_cache_then_demotes_enough_blocks(monkeypatch): +def test_pressure_relief_demotes_enough_live_blocks_without_emptying_cache(monkeypatch): snapshots = iter( [ { @@ -168,11 +168,6 @@ def test_pressure_relief_reclaims_cache_then_demotes_enough_blocks(monkeypatch): "predicted_peak_free_bytes": 200, "deficit_bytes": 300, }, - { - "device_free_bytes": 180, - "predicted_peak_free_bytes": 240, - "deficit_bytes": 260, - }, { "device_free_bytes": 500, "predicted_peak_free_bytes": 540, @@ -202,7 +197,7 @@ def test_pressure_relief_reclaims_cache_then_demotes_enough_blocks(monkeypatch): assert runtime._relieve_training_physical_pressure() is True assert transitions == [(('blocks.1', 'blocks.2'), False)] - assert len(empty_cache_calls) == 2 + assert empty_cache_calls == [] assert runtime._policy.last_reason == "physical_pressure_relief" assert runtime._last_training_pressure_relief["demoted_bytes"] == 300 diff --git a/tests/test_residency_two_timescale.py b/tests/test_residency_two_timescale.py index 1cdb644f00..381ee2e7de 100644 --- a/tests/test_residency_two_timescale.py +++ b/tests/test_residency_two_timescale.py @@ -78,6 +78,18 @@ def test_cold_settles_to_stable_after_k_clean(): assert s.name == vb.FSM_STABLE +def test_cold_pressure_raises_cap_when_it_can_relieve(): + s = vb.ResidencyFsmState() + s, a = drive(s, {"binding": True, "cap_can_relieve": True}) + assert s.name == vb.FSM_CAP_VERIFY and a == vb.ACT_RAISE_CAP + + +def test_cold_pressure_demotes_when_cap_is_pinned(): + s = vb.ResidencyFsmState() + s, a = drive(s, {"binding": True, "cap_can_relieve": False}) + assert s.name == vb.FSM_COLD and a == vb.ACT_DEMOTE + + def _stable(): s = vb.ResidencyFsmState() for _ in range(2): diff --git a/toolkit/memory_management/arena_offload/policy.py b/toolkit/memory_management/arena_offload/policy.py index 7c315a19a2..8fa66bcdb9 100644 --- a/toolkit/memory_management/arena_offload/policy.py +++ b/toolkit/memory_management/arena_offload/policy.py @@ -301,7 +301,7 @@ def allocation_failure(self): return self.reject_pending_promotion("promotion_allocation_failure") def record_physical_pressure_relief(self, block_keys, block_bytes): - """Put the controller in cooldown after emergency cache/layout relief.""" + """Put the controller in cooldown after emergency layout relief.""" keys = tuple(str(key) for key in block_keys) self.pending_promotion = None self.last_promoted_key = None diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index e261460de9..1d457cbbb3 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -972,7 +972,6 @@ def _training_pressure_snapshot(self): non_torch = max(0, total - free - reserved) predicted_free = total - peak_allocated - non_torch headroom = self._training_physical_vram_headroom_bytes() - governing_free = min(free, predicted_free) return { "total_bytes": total, "device_free_bytes": free, @@ -981,30 +980,28 @@ def _training_pressure_snapshot(self): "non_torch_bytes": non_torch, "predicted_peak_free_bytes": predicted_free, "headroom_bytes": headroom, - "deficit_bytes": max(0, headroom - governing_free), + # The allocator cap owns idle-cache reclamation. This guard asks + # the different question: would live memory alone breach the + # physical floor even after allocator GC did everything it could? + "deficit_bytes": max(0, headroom - predicted_free), } def _relieve_training_physical_pressure(self) -> bool: - """Reclaim cache, then demote enough residents to preserve the floor.""" - import torch - + """Demote live residency when allocator GC cannot preserve the floor.""" if self._signals.last_signal is None: return False before = self._training_pressure_snapshot() if before["deficit_bytes"] <= 0: return False - torch.cuda.empty_cache() - after_cache = self._training_pressure_snapshot() selected = [] selected_bytes = 0 - remaining = int(after_cache["deficit_bytes"]) - if remaining > 0: - for candidate in self._demotion_candidates(): - selected.append(candidate["block_key"]) - selected_bytes += int(candidate["block_bytes"]) - if selected_bytes >= remaining: - break + remaining = int(before["deficit_bytes"]) + for candidate in self._demotion_candidates(): + selected.append(candidate["block_key"]) + selected_bytes += int(candidate["block_bytes"]) + if selected_bytes >= remaining: + break if selected: result = self.transition_training_blocks( selected, resident=False @@ -1012,7 +1009,6 @@ def _relieve_training_physical_pressure(self) -> bool: if not result.get("changed"): selected = [] selected_bytes = 0 - torch.cuda.empty_cache() after = self._training_pressure_snapshot() self._policy.record_physical_pressure_relief( @@ -1020,15 +1016,12 @@ def _relieve_training_physical_pressure(self) -> bool: ) self._last_training_pressure_relief = { "before": before, - "after_cache": after_cache, "after": after, "demoted_block_keys": tuple(selected), "demoted_bytes": selected_bytes, } print( "[ArenaOffload] WDDM pressure relief: " - f"cache_reclaimed=" - f"{max(0, after_cache['device_free_bytes'] - before['device_free_bytes']) / GIB:.2f} GiB, " f"demoted_blocks={len(selected)}, " f"demoted={selected_bytes / GIB:.2f} GiB, " f"predicted_free=" diff --git a/toolkit/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py index 54a26d06b6..9b71ae7826 100644 --- a/toolkit/memory_management/vram_budget.py +++ b/toolkit/memory_management/vram_budget.py @@ -894,8 +894,10 @@ def demote(): return enter(FSM_COLD) if name == FSM_COLD: - if invalid or binding: - return ResidencyFsmState(FSM_COLD, 0 if invalid else w), ACT_HOLD + if invalid: + return ResidencyFsmState(FSM_COLD, 0), ACT_HOLD + if binding: + return enter(FSM_CAP_VERIFY, ACT_RAISE_CAP) if cap_relieve else demote() return enter(FSM_STABLE) if w >= k_clean else stay() if name == FSM_STABLE: From 28aa9ade4aafdff159178c76fc3c25a83c5364ef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Wed, 15 Jul 2026 21:55:28 +0200 Subject: [PATCH 12/20] Stabilize arena calibration and compiled sampling --- jobs/process/BaseSDTrainProcess.py | 17 +- requirements_base.txt | 2 +- tests/test_arena_cap_calibrator.py | 158 ++++ tests/test_arena_offload_api.py | 35 + tests/test_arena_offload_policy.py | 130 +++- tests/test_arena_sampling_cap.py | 169 +++++ tests/test_compile_utils.py | 49 ++ tests/test_generic_block_dispatcher.py | 123 ++++ tests/test_residency_two_timescale.py | 13 + tests/test_torchao_compat.py | 27 + toolkit/compile_utils.py | 17 + toolkit/config_modules.py | 6 + toolkit/memory_management/allocator_cap.py | 3 +- .../memory_management/arena_offload/api.py | 28 +- .../arena_offload/cap_calibrator.py | 485 +++++++++++++ .../arena_offload/dispatcher.py | 56 +- .../memory_management/arena_offload/policy.py | 138 +++- .../arena_offload/runtime.py | 682 +++++++++++++++++- .../memory_management/immutable_runtime.py | 46 +- toolkit/memory_management/residency.py | 9 + toolkit/memory_management/vram_budget.py | 18 + toolkit/models/base_model.py | 9 +- toolkit/quantization/fp8_linear.py | 11 +- toolkit/quantization/torchao_compat.py | 96 +++ toolkit/util/quantize.py | 26 +- ui/src/app/jobs/new/SimpleJob.tsx | 15 + ui/src/app/jobs/new/utils.ts | 2 + ui/src/types.ts | 2 + 28 files changed, 2261 insertions(+), 111 deletions(-) create mode 100644 tests/test_arena_cap_calibrator.py create mode 100644 tests/test_arena_sampling_cap.py create mode 100644 tests/test_compile_utils.py create mode 100644 tests/test_torchao_compat.py create mode 100644 toolkit/memory_management/arena_offload/cap_calibrator.py create mode 100644 toolkit/quantization/torchao_compat.py diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 780b88adfd..7581366064 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -30,7 +30,10 @@ from toolkit.basic import value_map from toolkit.clip_vision_adapter import ClipVisionAdapter -from toolkit.compile_utils import configure_cuda_only_inductor +from toolkit.compile_utils import ( + configure_cuda_only_inductor, + configure_quantized_compile_tuning, +) from toolkit.custom_adapter import CustomAdapter from toolkit.data_loader import get_dataloader_from_datasets, trigger_dataloader_setup_epoch from toolkit.data_transfer_object.data_loader import FileItemDTO, DataLoaderBatchDTO @@ -373,7 +376,9 @@ def sample(self, step=None, is_first=False): arena_runtime = get_memory_runtime(getattr(self.sd, "unet", None)) sampling_session = ( - arena_runtime.sampling_session() + arena_runtime.sampling_session( + gen_configs=gen_img_config_list, + ) if arena_runtime is not None else contextlib.nullcontext() ) @@ -1741,6 +1746,14 @@ def run(self): with model_load_arena_session(self.sd, enabled=arena_requested): self.sd.load_model() + coordinate_descent = configure_quantized_compile_tuning(self.model_config) + if coordinate_descent is not None: + state = "enabled" if coordinate_descent else "disabled" + print_acc( + "Quantized compile coordinate-descent tuning explicitly " + f"{state} by job config." + ) + text_encoders = getattr(self.sd, "text_encoder", None) if text_encoders is not None and not isinstance( text_encoders, (list, tuple) diff --git a/requirements_base.txt b/requirements_base.txt index 0c0472491d..9ff8fc4516 100644 --- a/requirements_base.txt +++ b/requirements_base.txt @@ -1,4 +1,4 @@ -torchao==0.17.0 +torchao>=0.10.0,<0.18.0 safetensors git+https://github.com/huggingface/diffusers.git@c943837899b16cbae2f619b8dd4f7bb6f07dd81a #pip install git+https://github.com/huggingface/diffusers.git@refs/pull/13432/head diff --git a/tests/test_arena_cap_calibrator.py b/tests/test_arena_cap_calibrator.py new file mode 100644 index 0000000000..aa43438c93 --- /dev/null +++ b/tests/test_arena_cap_calibrator.py @@ -0,0 +1,158 @@ +from types import SimpleNamespace + +from toolkit.memory_management.arena_offload.cap_calibrator import ( + CAP_PROBE_SETTLE, + CAP_PROBE_VERIFY, + CAP_RESTORE_VERIFY, + CAP_SET, + CAP_SETTLED, + TrainingCapCalibrator, +) + + +def peak(working, steps=2): + return SimpleNamespace(working_peak_bytes=working, steps=steps) + + +def signal(shape, *, retries=0, frees=0, compile_invalid=False): + return { + "shape_key": shape, + "compile_invalid": compile_invalid, + "allocator": { + "alloc_retries_delta": retries, + "free_count_delta": frees, + }, + } + + +def drive(calibrator, previous, upcoming, peaks, current): + return calibrator.step( + previous, + upcoming_shape_key=upcoming, + shape_peaks=peaks, + resident_bytes=100, + ring_bytes=0, + cliff_cap_bytes=1000, + current_cap_bytes=current, + ) + + +def test_calibrator_verifies_every_bucket_and_restores_last_clean_cap(): + calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + peaks = {("a",): peak(300), ("b",): peak(400)} + + decision = drive(calibrator, None, ("a",), peaks, 1000) + assert decision.action == CAP_SET + assert decision.target_cap_bytes == 700 + assert calibrator.state == CAP_PROBE_SETTLE + + decision = drive(calibrator, signal(("a",)), ("a",), peaks, 700) + assert decision.action != CAP_SET + assert calibrator.state == CAP_PROBE_VERIFY + + decision = drive(calibrator, signal(("a",)), ("b",), peaks, 700) + assert decision.reason == "awaiting_probe_buckets" + decision = drive(calibrator, signal(("b",)), ("a",), peaks, 700) + assert decision.target_cap_bytes == 600 + + drive(calibrator, signal(("a",)), ("a",), peaks, 600) + drive(calibrator, signal(("a",)), ("b",), peaks, 600) + decision = drive(calibrator, signal(("b",)), ("a",), peaks, 600) + assert decision.target_cap_bytes == 500 + + decision = drive( + calibrator, signal(("a",), frees=1), ("a",), peaks, 500 + ) + assert decision.action == CAP_SET + assert decision.target_cap_bytes == 600 + assert decision.reason == "probe_allocator_gc" + assert calibrator.state == CAP_RESTORE_VERIFY + + decision = drive(calibrator, signal(("a",)), ("a",), peaks, 600) + assert decision.hold_residency is False + assert calibrator.state == CAP_SETTLED + assert calibrator.settled_cap_bytes == 600 + assert calibrator.learned_cache_pad_bytes == 70 + profiles = calibrator.diagnostics()["buckets"] + assert {row["shape_key"] for row in profiles} == {("a",), ("b",)} + assert any(row["dirty_at_probe_cap"] == 500 for row in profiles) + + +def test_first_gc_during_probe_restores_last_clean_cap(): + calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + peaks = {("a",): peak(400)} + drive(calibrator, None, ("a",), peaks, 1000) + + decision = drive( + calibrator, signal(("a",), frees=1), ("a",), peaks, 700 + ) + + assert decision.action == CAP_SET + assert decision.target_cap_bytes == 1000 + assert decision.reason == "probe_allocator_gc" + assert calibrator.last_dirty_cap_bytes == 700 + assert calibrator.state == CAP_RESTORE_VERIFY + + +def test_unseen_bucket_restores_cliff_and_invalidates_settlement(): + calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + calibrator.state = CAP_SETTLED + calibrator.settled_cap_bytes = 600 + calibrator.learned_cache_pad_bytes = 70 + calibrator.bucket_profiles[("a",)] = SimpleNamespace( + working_peak_bytes=400, + valid_observations=2, + seen_at_probe_cap=600, + dirty_at_probe_cap=None, + ) + + decision = drive(calibrator, signal(("a",)), ("new",), {}, 600) + assert decision.action == CAP_SET + assert decision.target_cap_bytes == 1000 + assert decision.hold_residency is True + + +def test_compile_invalid_restores_cliff_and_discards_bucket_profiles(): + calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + peaks = {("a",): peak(400)} + drive(calibrator, None, ("a",), peaks, 1000) + + decision = drive( + calibrator, + signal(("a",), compile_invalid=True), + ("a",), + {}, + 700, + ) + assert decision.action == CAP_SET + assert decision.target_cap_bytes == 1000 + assert calibrator.bucket_profiles == {} + + +def test_probe_oom_restores_before_residency_recovery(): + calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + peaks = {("a",): peak(400)} + drive(calibrator, None, ("a",), peaks, 1000) + + decision = calibrator.allocation_failure( + cliff_cap_bytes=1000, current_cap_bytes=700 + ) + assert decision.action == CAP_SET + assert decision.target_cap_bytes == 1000 + assert calibrator.state == CAP_RESTORE_VERIFY + + +def test_probe_oom_adds_two_notches_when_last_clean_is_too_close(): + calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + calibrator.state = CAP_PROBE_SETTLE + calibrator.probe_cap_bytes = 500 + calibrator.last_clean_cap_bytes = 600 + + decision = calibrator.allocation_failure( + cliff_cap_bytes=1000, current_cap_bytes=500 + ) + + assert decision.action == CAP_SET + assert decision.target_cap_bytes == 700 + assert calibrator.last_clean_cap_bytes == 700 + assert calibrator.state == CAP_RESTORE_VERIFY diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py index a082a4b433..4a8d4b2077 100644 --- a/tests/test_arena_offload_api.py +++ b/tests/test_arena_offload_api.py @@ -89,6 +89,7 @@ class _FakeModelConfig: layer_offloading_smart_working_reserve_gb = -1.0 layer_offloading_smart_wddm_margin_gb = None layer_offloading_smart_wddm_hard_gb = 1.0 + layer_offloading_smart_cap_calibration = True layer_offloading_wddm_spill_reserve_pct = 0.10 layer_offloading_block_stream_only = False layer_offloading_checkpoint_keep_last = 2 @@ -277,6 +278,7 @@ def test_from_model_config_maps_the_public_surface(self): self.assertFalse(config.strict_vram_cap) self.assertEqual(config._policy.prefetch_depth, 3) self.assertEqual(config._policy.checkpoint_keep_last, 2) + self.assertTrue(config._policy.cap_calibration) def test_public_surface_is_narrow(self): public = {field.name for field in fields(ArenaOffloadConfig) if not field.name.startswith("_")} @@ -305,11 +307,44 @@ class NoQuant(_FakeModelConfig): self.assertFalse(config.fp8_sampling) self.assertTrue(config.enabled) + def test_old_torchao_disables_only_torchao_arena_fp8(self): + class TorchAOFloat8(_FakeModelConfig): + qtype = "float8" + + with unittest.mock.patch( + "toolkit.memory_management.arena_offload.api." + "torchao_arena_fp8_supported", + return_value=False, + ), unittest.mock.patch( + "toolkit.memory_management.arena_offload.api.TORCHAO_VERSION", + "0.10.0", + ): + with self.assertWarnsRegex(RuntimeWarning, "requires_0.17.0"): + config = ArenaOffloadConfig.from_model_config(TorchAOFloat8()) + + self.assertTrue(config.enabled) + self.assertFalse(config.fp8_forward) + self.assertFalse(config.fp8_backward) + self.assertFalse(config.fp8_sampling) + + def test_quanto_fp8_does_not_require_new_torchao_tensor_format(self): + with unittest.mock.patch( + "toolkit.memory_management.arena_offload.api." + "torchao_arena_fp8_supported", + return_value=False, + ): + config = ArenaOffloadConfig.from_model_config(_FakeModelConfig()) + + self.assertTrue(config.fp8_forward) + self.assertTrue(config.fp8_backward) + self.assertTrue(config.fp8_sampling) + def test_missing_attributes_fall_back_to_defaults(self): config = ArenaOffloadConfig.from_model_config(object()) self.assertFalse(config.enabled) self.assertFalse(config.compile_blocks) self.assertEqual(config._policy.prefetch_depth, 3) + self.assertFalse(config._policy.cap_calibration) def test_dead_compile_aliases_do_not_enable_arena_compile(self): class DeadAliases: diff --git a/tests/test_arena_offload_policy.py b/tests/test_arena_offload_policy.py index d052160e85..35d94d11b2 100644 --- a/tests/test_arena_offload_policy.py +++ b/tests/test_arena_offload_policy.py @@ -4,6 +4,10 @@ import pytest +from toolkit.memory_management.arena_offload.cap_calibrator import ( + CAP_SET, + CapCalibrationDecision, +) from toolkit.memory_management.arena_offload.policy import ( ArenaResidencyController, TrainingSignalWindow, @@ -289,6 +293,71 @@ def test_training_policy_rebinds_cap_before_pressure_relief(monkeypatch): assert events == ["bind", "relieve"] +def test_training_policy_applies_calibration_before_residency(monkeypatch): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._device = "cuda" + runtime._config = SimpleNamespace( + strict_vram_cap=False, + _policy=SimpleNamespace(wddm_hard_gib=1.0), + ) + runtime._last_training_cap_target_bytes = None + runtime._signals = SimpleNamespace(last_signal=None, shape_peaks={}) + runtime._residency = SimpleNamespace(resident_bytes=lambda: 100) + runtime._smart_plan = {"singleton_resident_bytes": 0} + runtime._bind_training_cap = lambda: None + runtime._relieve_training_physical_pressure = lambda: False + runtime._promotion_candidate = lambda: None + runtime._demotion_candidate = lambda: None + runtime._demotion_candidates = lambda: () + runtime._training_ring_bytes = lambda: 50 + seen = {} + + class Calibrator: + enabled = True + learned_cache_pad_bytes = None + + def step(self, signal, **kwargs): + seen.update(kwargs) + return CapCalibrationDecision( + action=CAP_SET, + target_cap_bytes=700, + reason="initial_probe", + hold_residency=True, + ) + + def diagnostics(self): + return { + "state": "probe_settle", + "predicted_initial_cap_bytes": 700, + "last_clean_cap_bytes": 1000, + "last_dirty_cap_bytes": None, + "learned_cache_pad_bytes": None, + "buckets": [], + } + + runtime._cap_calibrator = Calibrator() + cap_calls = [] + monkeypatch.setattr("torch.cuda.is_available", lambda: True) + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime.allocator_cap." + "wddm_cliff_cap_bytes", + lambda *_args, **_kwargs: 1000, + ) + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime.allocator_cap." + "configure_wddm_allocator_guard", + lambda *args, **kwargs: cap_calls.append((args, kwargs)), + ) + + runtime._apply_training_policy(shape_key=(768, 768)) + + assert seen["upcoming_shape_key"] == (768, 768) + assert seen["resident_bytes"] == 100 + assert seen["ring_bytes"] == 50 + assert cap_calls[0][1]["target_cap_bytes"] == 700 + assert runtime._last_training_cap_target_bytes == 700 + + def test_shape_working_peak_excludes_residency(): window = TrainingSignalWindow() observe(window, peak_allocated_bytes=900, resident_bytes=200) @@ -556,6 +625,63 @@ def test_controller_raises_cap_by_fixed_fsm_increment(): assert decision.target_cap_bytes == 710 +def test_controller_exactly_prefunds_promotion_after_cap_calibration(): + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10 + ) + controller.bootstrapped = True + controller.state = type(controller.state)("stable", 2) + signal = { + "allocator": { + "alloc_retries_delta": 0, + "free_count_delta": 0, + }, + "compile_invalid": False, + "transfer": {"bytes": 100, "h2d_ms": 10.0}, + } + decision = controller.step( + signal, + candidate={"block_key": "blocks.1", "block_bytes": 100}, + demote_candidate=None, + cliff_cap_bytes=1000, + current_cap_bytes=600, + worst_shape_free_bytes=100, + worst_shape_allocator_slack_bytes=70, + worst_shape_live_bytes=500, + learned_cache_pad_bytes=50, + ) + assert decision.action == "raise_cap" + assert decision.target_cap_bytes == int(650 / 0.95) + + +def test_controller_holds_when_exact_promotion_cap_exceeds_cliff(): + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10 + ) + controller.bootstrapped = True + controller.state = type(controller.state)("stable", 2) + decision = controller.step( + { + "allocator": { + "alloc_retries_delta": 0, + "free_count_delta": 0, + }, + "compile_invalid": False, + "transfer": {"bytes": 100, "h2d_ms": 10.0}, + }, + candidate={"block_key": "blocks.1", "block_bytes": 100}, + demote_candidate=None, + cliff_cap_bytes=650, + current_cap_bytes=600, + worst_shape_free_bytes=100, + worst_shape_allocator_slack_bytes=70, + worst_shape_live_bytes=500, + learned_cache_pad_bytes=50, + ) + assert decision.action == "hold" + assert decision.reason == "promotion_exceeds_cliff" + + def test_arena_sampling_binds_fp8_to_canonical_and_singletons(monkeypatch): model = SimpleNamespace() canonical = SimpleNamespace() @@ -839,7 +965,7 @@ def test_failed_training_step_preserves_original_error_if_cleanup_fails(): runtime._last_step_num = None runtime._device = "cpu" runtime._last_policy_error = None - runtime._apply_training_policy = lambda: None + runtime._apply_training_policy = lambda **_kwargs: None runtime._handle_training_failure = ( lambda *_args, **_kwargs: (_ for _ in ()).throw( RuntimeError("cleanup failed") @@ -862,7 +988,7 @@ def test_failed_training_step_does_not_publish_partial_peak(monkeypatch): runtime._last_shape_key = None runtime._last_step_num = None runtime._device = "cpu" - runtime._apply_training_policy = lambda: None + runtime._apply_training_policy = lambda **_kwargs: None runtime._handle_training_failure = lambda *_args, **_kwargs: None runtime._executor = SimpleNamespace( TRAIN="train", diff --git a/tests/test_arena_sampling_cap.py b/tests/test_arena_sampling_cap.py new file mode 100644 index 0000000000..6a9976e424 --- /dev/null +++ b/tests/test_arena_sampling_cap.py @@ -0,0 +1,169 @@ +from types import SimpleNamespace + +import torch + +from toolkit.memory_management.arena_offload.cap_calibrator import ( + CAP_SET, + TrainingCapCalibrator, +) +from toolkit.memory_management.arena_offload.policy import ( + ShapePeak, + TrainingSignalWindow, +) +from toolkit.memory_management.arena_offload.runtime import ( + ArenaOffloadRuntime, + _SamplingCapProfile, + _sampling_cold_working_bytes, + _sampling_config_shape_key, +) +from toolkit.memory_management.immutable_runtime import ImmutableTransformerRuntime + + +def config(**overrides): + values = { + "width": 1024, + "height": 1024, + "num_frames": 1, + "guidance_scale": 4.5, + "batch_cfg": False, + "ctrl_idx": None, + "ctrl_img": None, + "ctrl_img_1": None, + "ctrl_img_2": None, + "ctrl_img_3": None, + "extra_values": [], + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_sampling_config_key_is_generic_and_shape_conservative(): + base = config() + assert _sampling_config_shape_key(base) == _sampling_config_shape_key(config()) + assert _sampling_config_shape_key(base) != _sampling_config_shape_key( + config(guidance_scale=1.0) + ) + assert _sampling_config_shape_key(base) != _sampling_config_shape_key( + config(ctrl_img="reference.png", ctrl_img_1="reference.png") + ) + assert _sampling_cold_working_bytes( + config(ctrl_img="reference.png", ctrl_img_1="reference.png") + ) > _sampling_cold_working_bytes(base) + + +def test_only_repeated_or_previously_seen_profiles_enable_calibration(): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._cap_calibration_enabled = True + runtime._sampling_cap_profiles = {} + key = _sampling_config_shape_key(config()) + + assert runtime._sampling_profile(key, 1) is None + repeated = runtime._sampling_profile(key, 2) + assert repeated is not None + assert repeated.calibrator.enabled + assert runtime._sampling_profile(key, 1) is repeated + + +def test_sampling_session_preflights_configs_and_restores_cap_before_decode(): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + runtime._closed = False + runtime._config = SimpleNamespace(fp8_sampling=False) + runtime._canonical_modules = () + runtime._smart_plan = {} + runtime._training_plan = object() + runtime._executor = SimpleNamespace( + TRAIN="train", activate=lambda *_args: None + ) + runtime._bind_training_cap = lambda: None + restored = [] + runtime._bind_sampling_cap = lambda target=None, **_kwargs: restored.append(target) + class Vae(torch.nn.Module): + def decode(self, value): + return ("decoded", value) + + owner = SimpleNamespace(vae=Vae()) + first = config() + second = config() + + with runtime.sampling_session( + gen_configs=[first, second], generation_owner=owner + ): + key = _sampling_config_shape_key(first) + assert runtime._sampling_session_occurrences[key] == 2 + assert owner.vae.decode("latent") == ("decoded", "latent") + + assert restored == [None] + assert "decode" not in owner.vae.__dict__ + + +def test_sampling_forward_begin_uses_completed_generic_fsm(monkeypatch): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + key = _sampling_config_shape_key(config()) + profile = _SamplingCapProfile( + calibrator=TrainingCapCalibrator(enabled=True, notch_bytes=100), + signals=TrainingSignalWindow(), + ) + profile.signals._shape_peaks[key] = ShapePeak( + steps=2, working_peak_bytes=400 + ) + runtime._active_sampling_cap_profile = profile + runtime._active_sampling_shape_key = key + runtime._device = "cuda" + runtime._residency = SimpleNamespace(resident_bytes=lambda: 100) + runtime._smart_plan = {} + runtime._active_sampling_resident_bytes = 100 + runtime._active_sampling_ring_bytes = 0 + runtime._training_ring_bytes = lambda: 0 + runtime._sampling_cliff_cap_bytes = lambda: 1000 + targets = [] + counter_snapshots = [] + runtime._bind_sampling_cap = lambda target, **kwargs: targets.append( + (target, kwargs) + ) + runtime._sampling_allocator_counters = lambda: { + "num_alloc_retries": 0, + "num_device_alloc": 12, + "num_device_free": 5, + } + original_prime = profile.signals.prime_counters + + def record_prime(**kwargs): + counter_snapshots.append(kwargs) + original_prime(**kwargs) + + profile.signals.prime_counters = record_prime + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime.allocator_cap.applied_cap_bytes", + lambda _device: 1000, + ) + monkeypatch.setattr( + "torch.cuda.reset_peak_memory_stats", lambda _device: None + ) + + runtime._sampling_forward_begin() + + assert profile.calibrator.last_action == CAP_SET + assert targets == [(700, {"reclaim": True})] + assert counter_snapshots == [ + { + "allocator_counters": { + "num_alloc_retries": 0, + "num_device_alloc": 12, + "num_device_free": 5, + } + } + ] + + +def test_per_forward_resets_preserve_the_image_wide_sampling_peak(): + executor = ImmutableTransformerRuntime.__new__(ImmutableTransformerRuntime) + executor._sampling_baseline = { + "external_peak_allocated": 100, + "external_peak_reserved": 120, + } + + executor.record_sampling_peak(allocated_bytes=300, reserved_bytes=350) + executor.record_sampling_peak(allocated_bytes=250, reserved_bytes=320) + + assert executor._sampling_baseline["external_peak_allocated"] == 300 + assert executor._sampling_baseline["external_peak_reserved"] == 350 diff --git a/tests/test_compile_utils.py b/tests/test_compile_utils.py new file mode 100644 index 0000000000..da6cd6e950 --- /dev/null +++ b/tests/test_compile_utils.py @@ -0,0 +1,49 @@ +from types import SimpleNamespace + +import torch + +from toolkit.compile_utils import configure_quantized_compile_tuning + + +def test_quantized_compile_tuning_preserves_torchao_default_when_unset(): + config = SimpleNamespace(compile=True, quantize=True) + original_tuning = torch._inductor.config.coordinate_descent_tuning + original_directions = ( + torch._inductor.config.coordinate_descent_check_all_directions + ) + try: + result = configure_quantized_compile_tuning(config) + assert result is None + assert torch._inductor.config.coordinate_descent_tuning == original_tuning + assert ( + torch._inductor.config.coordinate_descent_check_all_directions + == original_directions + ) + finally: + torch._inductor.config.coordinate_descent_tuning = original_tuning + torch._inductor.config.coordinate_descent_check_all_directions = ( + original_directions + ) + + +def test_quantized_compile_tuning_can_disable_torchao_search(): + config = SimpleNamespace( + compile=True, + quantize=True, + compile_coordinate_descent=False, + ) + original_tuning = torch._inductor.config.coordinate_descent_tuning + original_directions = ( + torch._inductor.config.coordinate_descent_check_all_directions + ) + try: + assert configure_quantized_compile_tuning(config) is False + assert torch._inductor.config.coordinate_descent_tuning is False + assert ( + torch._inductor.config.coordinate_descent_check_all_directions is False + ) + finally: + torch._inductor.config.coordinate_descent_tuning = original_tuning + torch._inductor.config.coordinate_descent_check_all_directions = ( + original_directions + ) diff --git a/tests/test_generic_block_dispatcher.py b/tests/test_generic_block_dispatcher.py index 8e79a89d41..065ee7cd8d 100644 --- a/tests/test_generic_block_dispatcher.py +++ b/tests/test_generic_block_dispatcher.py @@ -261,6 +261,129 @@ def test_checkpointing_rejection_precedes_canonical_commit(): assert not hasattr(model, "_arena_offload_runtime") +def test_sampling_callbacks_bracket_one_complete_transformer_forward(): + model = _frozen_transformer() + model.enable_gradient_checkpointing(keep_last=1) + with mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(8 * 1024**3, 12 * 1024**3), + ), mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget." + "auto_physical_vram_headroom_gib", + return_value=1.0, + ): + config = ArenaOffloadConfig(enabled=True, compile_blocks=False) + config = replace( + config, + _policy=replace(config._policy, checkpoint_keep_last=1), + ) + runtime = prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=config, + ) + try: + runtime.finalize() + events = [] + runtime._executor.set_sampling_forward_callbacks( + lambda: events.append("begin"), + lambda: events.append("end"), + ) + runtime._executor.activate( + runtime._executor.SAMPLE, + ResidencyPlan.build(runtime._executor.SAMPLE, ()), + ) + promotions = [] + runtime._executor.set_residency_promotion_callback( + lambda nbytes, plan: promotions.append((nbytes, plan.phase)) + ) + resident = [ + (block_key, leaf_name) + for block_key in runtime._arena.block_keys() + for leaf_name in runtime._arena.block_record(block_key).leaf_names + ] + runtime._executor.activate( + runtime._executor.SAMPLE, + ResidencyPlan.build(runtime._executor.SAMPLE, resident), + ) + with torch.no_grad(), runtime._executor.execution( + runtime._executor.SAMPLE + ): + model(torch.randn(2, 4)) + assert events == ["begin", "end"] + assert promotions + assert promotions[-1][0] > 0 + assert promotions[-1][1] == runtime._executor.SAMPLE + finally: + close_arena_offload(model) + + +def test_sampling_dispatch_retries_one_failed_compiled_block_in_eager_wrapper(): + model = _frozen_transformer() + model.enable_gradient_checkpointing(keep_last=1) + with mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(8 * 1024**3, 12 * 1024**3), + ), mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget." + "auto_physical_vram_headroom_gib", + return_value=1.0, + ): + config = ArenaOffloadConfig(enabled=True, compile_blocks=False) + config = replace( + config, + _policy=replace(config._policy, checkpoint_keep_last=1), + ) + runtime = prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=config, + ) + try: + runtime.finalize() + executor = runtime._executor + original_get_kernel = executor._get_dispatch_kernel + attempts = [] + + def flaky_get_kernel(index): + kernel = original_get_kernel(index) + + def flaky(*args, **kwargs): + attempts.append(index) + if len(attempts) == 1: + raise torch.OutOfMemoryError("synthetic capped OOM") + return kernel(*args, **kwargs) + + return flaky + + executor._get_dispatch_kernel = flaky_get_kernel + recoveries = [] + executor.set_sampling_forward_callbacks( + allocation_failure=lambda error: recoveries.append(error) or True + ) + resident = [ + (block_key, leaf_name) + for block_key in runtime._arena.block_keys() + for leaf_name in runtime._arena.block_record(block_key).leaf_names + ] + executor.activate( + executor.SAMPLE, + ResidencyPlan.build(executor.SAMPLE, resident), + ) + + with torch.no_grad(), executor.execution(executor.SAMPLE): + output = model(torch.randn(2, 4)) + + assert output.shape == (2, 4) + assert len(recoveries) == 1 + assert isinstance(recoveries[0], torch.OutOfMemoryError) + assert attempts[:2] == [0, 0] + finally: + close_arena_offload(model) + + def test_saved_installed_forward_checkpoint_backward_and_teardown(): torch.manual_seed(17) model = _frozen_transformer() diff --git a/tests/test_residency_two_timescale.py b/tests/test_residency_two_timescale.py index 381ee2e7de..98d117c2cc 100644 --- a/tests/test_residency_two_timescale.py +++ b/tests/test_residency_two_timescale.py @@ -36,6 +36,19 @@ def test_cap_for_live_is_allowance_inverse(): assert vb.allocator_allowance_bytes(cap, gib(6.77)) == pytest.approx(gib(0.21), abs=gib(0.01)) +def test_promotion_cap_growth_preserves_allocator_allowance(): + current = gib(6.5) + promoted = gib(0.5) + target = vb.cap_bytes_preserving_allowance_after_promotion( + current, promoted, gib(9.5) + ) + + before = vb.allocator_allowance_bytes(current, gib(5.5)) + after = vb.allocator_allowance_bytes(target, gib(6.0)) + assert target > current + promoted + assert after >= before + + def test_cap_for_live_clamped_to_cliff(): cap = vb.cap_bytes_for_live(gib(11.0), gib(2.0), cliff_cap_bytes=gib(9.85)) assert cap == gib(9.85) diff --git a/tests/test_torchao_compat.py b/tests/test_torchao_compat.py new file mode 100644 index 0000000000..f0914316b4 --- /dev/null +++ b/tests/test_torchao_compat.py @@ -0,0 +1,27 @@ +from toolkit.quantization.torchao_compat import ( + _release_tuple, + intx_weight_only_config, + torchao_arena_fp8_supported, +) + + +def test_release_tuple_ignores_local_and_prerelease_suffixes(): + assert _release_tuple("0.17.0+cu132") == (0, 17, 0) + assert _release_tuple("0.17.0.dev20260715") == (0, 17, 0) + assert _release_tuple("unknown") == () + + +def test_arena_fp8_requires_tested_version_and_tensor_format(): + assert not torchao_arena_fp8_supported( + "0.10.0", float8_tensor_available=True + ) + assert not torchao_arena_fp8_supported( + "0.17.0", float8_tensor_available=False + ) + assert torchao_arena_fp8_supported( + "0.17.0", float8_tensor_available=True + ) + + +def test_current_intx_config_factory_is_available(): + assert intx_weight_only_config(4) is not None diff --git a/toolkit/compile_utils.py b/toolkit/compile_utils.py index e83d5cfa6d..37b1d1a9e3 100644 --- a/toolkit/compile_utils.py +++ b/toolkit/compile_utils.py @@ -12,3 +12,20 @@ def configure_cuda_only_inductor() -> None: from torch._inductor import config as inductor_config inductor_config.cpp.vec_isa_ok = False + + +def configure_quantized_compile_tuning(model_config) -> bool | None: + """Apply an explicit coordinate-descent policy after TorchAO quantization.""" + if not getattr(model_config, "compile", False): + return None + if not getattr(model_config, "quantize", False): + return None + + requested = getattr(model_config, "compile_coordinate_descent", None) + if requested is None: + return None + + enabled = bool(requested) + torch._inductor.config.coordinate_descent_tuning = enabled + torch._inductor.config.coordinate_descent_check_all_directions = enabled + return enabled diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index b75d4930b8..8e4f7449e1 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -712,6 +712,9 @@ def __init__(self, **kwargs): self.layer_offloading_smart_wddm_hard_gb = kwargs.get( "layer_offloading_smart_wddm_hard_gb", 1.0 ) + self.layer_offloading_smart_cap_calibration = kwargs.get( + "layer_offloading_smart_cap_calibration", False + ) self.layer_offloading_smart_sampling_working_reserve_gb = kwargs.get( "layer_offloading_smart_sampling_working_reserve_gb", -1.0 ) @@ -765,6 +768,9 @@ def __init__(self, **kwargs): self.compile_mode = kwargs.get("compile_mode", "default") self.compile_fullgraph = kwargs.get("compile_fullgraph", False) self.compile_dynamic = kwargs.get("compile_dynamic", True) + self.compile_coordinate_descent = kwargs.get( + "compile_coordinate_descent", None + ) self.cache_size_limit = kwargs.get("cache_size_limit", None) # kwargs to pass to the model diff --git a/toolkit/memory_management/allocator_cap.py b/toolkit/memory_management/allocator_cap.py index fb2c9f7dfa..b27517f052 100644 --- a/toolkit/memory_management/allocator_cap.py +++ b/toolkit/memory_management/allocator_cap.py @@ -93,6 +93,7 @@ def configure_wddm_allocator_guard( target_cap_bytes=None, strict=False, log_prefix="[MemoryManager]", + force=False, ): """Bind the allocator below the WDDM cliff in every guard mode. @@ -134,7 +135,7 @@ def configure_wddm_allocator_guard( applied = fraction * total / float(real_total) previous = APPLIED_FRACTIONS.get(index) tolerance = (64 * 1024**2) / real_total - if previous is not None and abs(previous - applied) < tolerance: + if not force and previous is not None and abs(previous - applied) < tolerance: return previous torch.cuda.set_per_process_memory_fraction(applied, index) diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index 686e97b4f9..ef2f8d9a4c 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -18,6 +18,12 @@ from typing import Any import warnings +from toolkit.quantization.torchao_compat import ( + TORCHAO_ARENA_FP8_MIN_VERSION, + TORCHAO_VERSION, + torchao_arena_fp8_supported, +) + from ..runtime import ( RUNTIME_ATTR, close_memory_runtime, @@ -140,6 +146,7 @@ class _ArenaPolicyOptions: working_reserve_gib: float | None = None physical_vram_headroom_gib: float | None = None wddm_hard_gib: float | None = None + cap_calibration: bool = False checkpoint_keep_last: int = 0 prefetch_depth: int = 3 @@ -203,6 +210,11 @@ def get(name: str, default: Any = None) -> Any: requested_backward = bool(get("layer_offloading_fp8_grad_input", False)) requested_sampling = bool(get("layer_offloading_fp8_sampling", False)) ignored = [] + torchao_fp8_unavailable = bool( + fp8_weights + and get("qtype") == "float8" + and not torchao_arena_fp8_supported() + ) if not fp8_weights: ignored.extend( name @@ -215,6 +227,13 @@ def get(name: str, default: Any = None) -> Any: ) elif requested_backward and not requested_forward: ignored.append("fp8_backward_without_fp8_forward") + if torchao_fp8_unavailable and any( + (requested_forward, requested_backward, requested_sampling) + ): + ignored.append( + "torchao_fp8_requires_" + f"{TORCHAO_ARENA_FP8_MIN_VERSION}_installed_{TORCHAO_VERSION}" + ) if ignored: warnings.warn( "arena offload ignored irrelevant FP8 options: " @@ -228,11 +247,15 @@ def get(name: str, default: Any = None) -> Any: get("layer_offloading", False) and get("layer_offloading_smart", False) ), - fp8_forward=fp8_weights and requested_forward, + fp8_forward=( + fp8_weights and not torchao_fp8_unavailable and requested_forward + ), fp8_backward=fp8_weights + and not torchao_fp8_unavailable and requested_forward and requested_backward, fp8_sampling=fp8_weights + and not torchao_fp8_unavailable and requested_sampling, # Arena execution has one shared block dispatcher for training # and sampling, so Toolkit's supported model compile setting owns @@ -258,6 +281,9 @@ def get(name: str, default: Any = None) -> Any: "layer_offloading_smart_physical_vram_headroom_gb" ), wddm_hard_gib=get("layer_offloading_smart_wddm_hard_gb"), + cap_calibration=bool( + get("layer_offloading_smart_cap_calibration", False) + ), checkpoint_keep_last=max( 0, int(get("layer_offloading_checkpoint_keep_last", 0) or 0) ), diff --git a/toolkit/memory_management/arena_offload/cap_calibrator.py b/toolkit/memory_management/arena_offload/cap_calibrator.py new file mode 100644 index 0000000000..7990f171aa --- /dev/null +++ b/toolkit/memory_management/arena_offload/cap_calibrator.py @@ -0,0 +1,485 @@ +"""Bounded allocator-cap calibration for arena training. + +The PyTorch allocator cap is process-global, so bucket measurements feed one +worst-bucket cap. The calibrator moves that cap only during a bounded startup +or invalidation probe; it never switches caps on every bucket transition. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from .. import vram_budget + + +CAP_WARMUP = "warmup" +CAP_PROBE_SETTLE = "probe_settle" +CAP_PROBE_VERIFY = "probe_verify" +CAP_SETTLED = "settled" +CAP_RESTORE_VERIFY = "restore_verify" + +CAP_HOLD = "hold" +CAP_SET = "set_cap" + +DEFAULT_CAP_NOTCH_BYTES = 256 * 1024**2 + + +@dataclass +class BucketCapProfile: + working_peak_bytes: int = 0 + valid_observations: int = 0 + seen_at_probe_cap: int | None = None + dirty_at_probe_cap: int | None = None + + +@dataclass(frozen=True) +class CapCalibrationDecision: + action: str = CAP_HOLD + target_cap_bytes: int | None = None + reason: str = "" + hold_residency: bool = True + + +def _ceil_to_notch(value: int, notch_bytes: int) -> int: + value = max(0, int(value)) + notch = max(1, int(notch_bytes)) + return ((value + notch - 1) // notch) * notch + + +class TrainingCapCalibrator: + """Learn one safe allocator cap from per-bucket working peaks. + + A retry, allocator GC, OOM, or compile invalidation rejects a lowered cap. + Runtime-owned cache trims must be excluded from the allocator counters + before the probe window begins. + Every bucket known when the probe begins must then produce a clean step. + """ + + def __init__( + self, + *, + enabled=False, + notch_bytes=DEFAULT_CAP_NOTCH_BYTES, + monitor_settled=False, + ): + self.enabled = bool(enabled) + self.monitor_settled = bool(monitor_settled) + self.notch_bytes = max(1, int(notch_bytes)) + self.state = CAP_WARMUP + self.bucket_profiles: dict[tuple, BucketCapProfile] = {} + self.probe_cap_bytes: int | None = None + self.last_clean_cap_bytes: int | None = None + self.settled_cap_bytes: int | None = None + self.learned_cache_pad_bytes: int | None = None + self.predicted_initial_cap_bytes: int | None = None + self.last_dirty_cap_bytes: int | None = None + self._verify_pending: set[tuple] = set() + self.last_action = CAP_HOLD + self.last_reason = "disabled" if not self.enabled else "warmup" + + @property + def active(self) -> bool: + return self.enabled and self.state != CAP_SETTLED + + def step( + self, + signal, + *, + upcoming_shape_key, + shape_peaks, + resident_bytes, + ring_bytes, + cliff_cap_bytes, + current_cap_bytes, + ) -> CapCalibrationDecision: + if not self.enabled: + return self._decision(CAP_HOLD, reason="disabled", hold=False) + + cliff = max(0, int(cliff_cap_bytes)) + current = max(0, int(current_cap_bytes)) + upcoming = _shape_key(upcoming_shape_key) + self._sync_profiles(shape_peaks) + + if signal is not None and bool(signal.get("compile_invalid")): + self._reset_for_invalid_measurements() + return self._set_or_hold( + cliff, current, reason="compile_invalid", hold=True + ) + + decision = self._consume_previous_signal( + signal, + shape_peaks=shape_peaks, + resident_bytes=resident_bytes, + ring_bytes=ring_bytes, + cliff_cap_bytes=cliff, + current_cap_bytes=current, + ) + if decision is not None: + current = ( + current + if decision.target_cap_bytes is None + else int(decision.target_cap_bytes) + ) + + valid = self._valid_bucket_keys() + if upcoming not in valid: + self._restart_warmup(reason="unseen_bucket") + return self._set_or_hold( + cliff, current, reason="unseen_bucket", hold=True + ) + + if decision is not None: + return decision + + if self.state == CAP_WARMUP: + if not valid: + return self._decision(CAP_HOLD, reason="awaiting_bucket", hold=True) + allocator = (signal or {}).get("allocator") or {} + if int(allocator.get("alloc_retries_delta", 0) or 0) > 0 or int( + allocator.get("free_count_delta", 0) or 0 + ) > 0: + return self._decision( + CAP_HOLD, reason="awaiting_clean_warmup", hold=True + ) + return self._begin_initial_probe( + resident_bytes=resident_bytes, + ring_bytes=ring_bytes, + cliff_cap_bytes=cliff, + current_cap_bytes=current, + ) + + return self._decision( + CAP_HOLD, + reason="settled" if self.state == CAP_SETTLED else self.state, + hold=self.state != CAP_SETTLED, + ) + + def allocation_failure(self, *, cliff_cap_bytes, current_cap_bytes): + """Reject an active probe before residency/OOM recovery is attempted.""" + if not self.active or self.state == CAP_WARMUP: + return None + dirty_cap = self.probe_cap_bytes or int(current_cap_bytes) + self._mark_dirty(None, dirty_cap) + recovery_floor = min( + int(cliff_cap_bytes), + int(dirty_cap) + 2 * self.notch_bytes, + ) + return self._begin_restore( + cliff_cap_bytes=int(cliff_cap_bytes), + current_cap_bytes=int(current_cap_bytes), + reason="probe_allocation_failure", + minimum_cap_bytes=recovery_floor, + ) + + def invalidate_for_residency_growth(self): + """Keep working peaks but relearn the cap for a larger resident layout.""" + if not self.enabled: + return + self.last_clean_cap_bytes = None + self.settled_cap_bytes = None + self.learned_cache_pad_bytes = None + self.predicted_initial_cap_bytes = None + self.last_dirty_cap_bytes = None + self._restart_warmup(reason="residency_growth") + + def _consume_previous_signal( + self, + signal, + *, + shape_peaks, + resident_bytes, + ring_bytes, + cliff_cap_bytes, + current_cap_bytes, + ): + if signal is None: + return None + allocator = signal.get("allocator") or {} + retries = int(allocator.get("alloc_retries_delta", 0) or 0) + frees = int(allocator.get("free_count_delta", 0) or 0) + shape_key = _shape_key(signal.get("shape_key")) + + if self.state == CAP_SETTLED and self.monitor_settled: + if retries > 0 or frees > 0: + dirty_cap = int(self.settled_cap_bytes or current_cap_bytes) + self._mark_dirty(shape_key, dirty_cap) + widened = min( + int(cliff_cap_bytes), dirty_cap + self.notch_bytes + ) + if widened > dirty_cap: + self.last_clean_cap_bytes = widened + self.settled_cap_bytes = None + self.state = CAP_RESTORE_VERIFY + return self._set_or_hold( + widened, + current_cap_bytes, + reason=( + "settled_allocator_retry" + if retries > 0 + else "settled_allocator_gc" + ), + hold=True, + ) + return None + + if self.state == CAP_PROBE_SETTLE: + if retries > 0 or frees > 0: + self._mark_dirty(shape_key, self.probe_cap_bytes) + return self._begin_restore( + cliff_cap_bytes=cliff_cap_bytes, + current_cap_bytes=current_cap_bytes, + reason=( + "probe_allocator_retry" + if retries > 0 + else "probe_allocator_gc" + ), + ) + self.state = CAP_PROBE_VERIFY + self._verify_pending = set(self._valid_bucket_keys()) + return self._decision(CAP_HOLD, reason="probe_settled", hold=True) + + if self.state == CAP_PROBE_VERIFY: + if retries > 0 or frees > 0: + self._mark_dirty(shape_key, self.probe_cap_bytes) + return self._begin_restore( + cliff_cap_bytes=cliff_cap_bytes, + current_cap_bytes=current_cap_bytes, + reason=( + "probe_allocator_retry" + if retries > 0 + else "probe_allocator_gc" + ), + ) + self._verify_pending.discard(shape_key) + profile = self.bucket_profiles.get(shape_key) + if profile is not None: + profile.seen_at_probe_cap = self.probe_cap_bytes + if self._verify_pending: + return self._decision( + CAP_HOLD, reason="awaiting_probe_buckets", hold=True + ) + return self._advance_clean_probe( + shape_peaks=shape_peaks, + resident_bytes=resident_bytes, + ring_bytes=ring_bytes, + cliff_cap_bytes=cliff_cap_bytes, + current_cap_bytes=current_cap_bytes, + ) + + if self.state == CAP_RESTORE_VERIFY: + if retries > 0 or frees > 0: + self._restart_warmup(reason="restore_not_clean") + return self._set_or_hold( + cliff_cap_bytes, + current_cap_bytes, + reason="restore_not_clean", + hold=True, + ) + self.state = CAP_SETTLED + self.settled_cap_bytes = self.last_clean_cap_bytes + return self._decision(CAP_HOLD, reason="calibration_settled", hold=False) + + return None + + def _begin_initial_probe( + self, *, resident_bytes, ring_bytes, cliff_cap_bytes, current_cap_bytes + ): + worst_live = self._worst_live_bytes(resident_bytes, ring_bytes) + predicted = _ceil_to_notch( + int(worst_live / vram_budget.GC_THRESHOLD), self.notch_bytes + ) + self.notch_bytes + candidate = min(int(cliff_cap_bytes), predicted) + self.predicted_initial_cap_bytes = candidate + self.last_clean_cap_bytes = int(cliff_cap_bytes) + if candidate >= int(cliff_cap_bytes): + self.state = CAP_SETTLED + self.settled_cap_bytes = int(cliff_cap_bytes) + self.learned_cache_pad_bytes = max( + 0, + vram_budget.allocator_allowance_bytes( + self.settled_cap_bytes, worst_live + ), + ) + return self._decision( + CAP_HOLD, reason="no_reclaimable_notch", hold=False + ) + return self._begin_probe(candidate, current_cap_bytes, "initial_probe") + + def _advance_clean_probe( + self, + *, + shape_peaks, + resident_bytes, + ring_bytes, + cliff_cap_bytes, + current_cap_bytes, + ): + clean_cap = int(self.probe_cap_bytes or current_cap_bytes) + self.last_clean_cap_bytes = clean_cap + worst_live = self._worst_live_bytes(resident_bytes, ring_bytes) + self.learned_cache_pad_bytes = max( + 0, vram_budget.allocator_allowance_bytes(clean_cap, worst_live) + ) + next_cap = clean_cap - self.notch_bytes + minimum = _ceil_to_notch( + int(worst_live / vram_budget.GC_THRESHOLD), self.notch_bytes + ) - self.notch_bytes + next_cap = max(self.notch_bytes, minimum, next_cap) + if next_cap >= clean_cap: + self.state = CAP_SETTLED + self.settled_cap_bytes = clean_cap + return self._decision(CAP_HOLD, reason="calibration_floor", hold=False) + return self._begin_probe(next_cap, current_cap_bytes, "lower_probe") + + def _begin_probe(self, target_cap_bytes, current_cap_bytes, reason): + self.state = CAP_PROBE_SETTLE + self.probe_cap_bytes = int(target_cap_bytes) + self._verify_pending.clear() + return self._set_or_hold( + self.probe_cap_bytes, + int(current_cap_bytes), + reason=reason, + hold=True, + ) + + def _begin_restore( + self, + *, + cliff_cap_bytes, + current_cap_bytes, + reason, + minimum_cap_bytes=0, + ): + restore = min( + int(cliff_cap_bytes), + max( + int(self.last_clean_cap_bytes or cliff_cap_bytes), + int(minimum_cap_bytes), + ), + ) + self.last_clean_cap_bytes = restore + self.state = CAP_RESTORE_VERIFY + self.probe_cap_bytes = None + self._verify_pending.clear() + return self._set_or_hold( + restore, int(current_cap_bytes), reason=reason, hold=True + ) + + def _sync_profiles(self, shape_peaks): + for key, peak in (shape_peaks or {}).items(): + normalized = _shape_key(key) + steps = int(getattr(peak, "steps", 0) or 0) + if steps <= 0: + continue + profile = self.bucket_profiles.setdefault( + normalized, BucketCapProfile() + ) + profile.working_peak_bytes = max( + profile.working_peak_bytes, + int(getattr(peak, "working_peak_bytes", 0) or 0), + ) + profile.valid_observations = max(profile.valid_observations, steps) + + def _valid_bucket_keys(self): + return { + key + for key, profile in self.bucket_profiles.items() + if profile.valid_observations > 0 + } + + def _worst_live_bytes(self, resident_bytes, ring_bytes): + worst_working = max( + ( + profile.working_peak_bytes + for profile in self.bucket_profiles.values() + if profile.valid_observations > 0 + ), + default=0, + ) + return ( + int(worst_working) + + max(0, int(resident_bytes)) + + max(0, int(ring_bytes)) + ) + + def _mark_dirty(self, shape_key, cap_bytes): + self.last_dirty_cap_bytes = ( + None if cap_bytes is None else int(cap_bytes) + ) + if shape_key is None: + return + profile = self.bucket_profiles.get(_shape_key(shape_key)) + if profile is not None: + profile.dirty_at_probe_cap = ( + None if cap_bytes is None else int(cap_bytes) + ) + + def _reset_for_invalid_measurements(self): + self.bucket_profiles.clear() + self.last_clean_cap_bytes = None + self.settled_cap_bytes = None + self.learned_cache_pad_bytes = None + self.predicted_initial_cap_bytes = None + self.last_dirty_cap_bytes = None + self._restart_warmup(reason="compile_invalid") + + def _restart_warmup(self, *, reason): + self.state = CAP_WARMUP + self.probe_cap_bytes = None + self.settled_cap_bytes = None + self._verify_pending.clear() + self.last_reason = reason + + def _set_or_hold(self, target, current, *, reason, hold): + if int(target) == int(current): + return self._decision(CAP_HOLD, reason=reason, hold=hold) + return self._decision(CAP_SET, int(target), reason=reason, hold=hold) + + def _decision(self, action, target=None, *, reason, hold): + self.last_action = action + self.last_reason = reason + return CapCalibrationDecision( + action=action, + target_cap_bytes=target, + reason=reason, + hold_residency=bool(hold), + ) + + def diagnostics(self): + return { + "enabled": self.enabled, + "monitor_settled": self.monitor_settled, + "state": self.state, + "notch_bytes": self.notch_bytes, + "probe_cap_bytes": self.probe_cap_bytes, + "last_clean_cap_bytes": self.last_clean_cap_bytes, + "settled_cap_bytes": self.settled_cap_bytes, + "learned_cache_pad_bytes": self.learned_cache_pad_bytes, + "predicted_initial_cap_bytes": self.predicted_initial_cap_bytes, + "last_dirty_cap_bytes": self.last_dirty_cap_bytes, + "verify_pending": tuple(sorted(self._verify_pending, key=repr)), + "last_action": self.last_action, + "last_reason": self.last_reason, + "buckets": [ + { + "shape_key": key, + "working_peak_bytes": profile.working_peak_bytes, + "valid_observations": profile.valid_observations, + "seen_at_probe_cap": profile.seen_at_probe_cap, + "dirty_at_probe_cap": profile.dirty_at_probe_cap, + } + for key, profile in sorted( + self.bucket_profiles.items(), key=lambda item: repr(item[0]) + ) + ], + } + + +def _shape_key(value): + if value is None: + return ("unknown",) + if isinstance(value, tuple): + return value + if isinstance(value, list): + return tuple(value) + return (value,) diff --git a/toolkit/memory_management/arena_offload/dispatcher.py b/toolkit/memory_management/arena_offload/dispatcher.py index fb2d7ece5c..a94ad2e35e 100644 --- a/toolkit/memory_management/arena_offload/dispatcher.py +++ b/toolkit/memory_management/arena_offload/dispatcher.py @@ -134,6 +134,17 @@ def __init__( self._dispatchers = () self._saved_forwards = () self._replacements = () + self._sampling_forward_begin = None + self._sampling_forward_end = None + self._sampling_allocation_failure = None + + def set_sampling_forward_callbacks( + self, begin=None, end=None, allocation_failure=None + ): + """Install eager callbacks around one complete sampled transformer pass.""" + self._sampling_forward_begin = begin + self._sampling_forward_end = end + self._sampling_allocation_failure = allocation_failure def _replacement_plan(self, index, invoker): abi = self._block_abis[index] @@ -284,6 +295,9 @@ def dispatch(self, index, args, kwargs): raise ImmutableRuntimeError( "immutable_execution_mode_mismatch:active=sample:call=train" ) + sampling = self._active_mode == self.SAMPLE + if sampling and index == 0 and self._sampling_forward_begin is not None: + self._sampling_forward_begin() try: first, first_location = _first_tensor_argument(args, kwargs) except ImmutableRuntimeError as error: @@ -326,18 +340,27 @@ def dispatch(self, index, args, kwargs): leaf_args = source.assemble_leaf_args(self.residency, compact_flat) self._mark_dispatch_dynamic(first) - try: - output = self._get_dispatch_kernel(index)(leaf_args, args, kwargs) - except BaseException: - # Non-reentrant checkpoint replay raises its private early-stop - # control-flow exception as soon as it has regenerated every - # tensor backward requested. That can unwind the block before the - # normal post-forward release below. With no gradient-bearing - # input, frozen streamed state is not a backward dependency and - # there is no free_on_backward node, so release on that unwind. - if token is not None and not release_on_backward: - torch.ops.mm.fetch_free(token) - raise + while True: + try: + output = self._get_dispatch_kernel(index)(leaf_args, args, kwargs) + break + except BaseException as error: + recover = ( + sampling + and self._sampling_allocation_failure is not None + and self._sampling_allocation_failure(error) + ) + if recover: + continue + # Non-reentrant checkpoint replay raises its private early-stop + # control-flow exception as soon as it has regenerated every + # tensor backward requested. That can unwind the block before the + # normal post-forward release below. With no gradient-bearing + # input, frozen streamed state is not a backward dependency and + # there is no free_on_backward node, so release on that unwind. + if token is not None and not release_on_backward: + torch.ops.mm.fetch_free(token) + raise if token is not None: # The first checkpoint pass discards its fetched views, so return # that slot after forward. A replay with a gradient-bearing input @@ -350,6 +373,12 @@ def dispatch(self, index, args, kwargs): torch.ops.mm.fetch_free_after( token, _first_output_tensor(output) ) + if ( + sampling + and index == len(self._blocks) - 1 + and self._sampling_forward_end is not None + ): + self._sampling_forward_end() return output def close(self): @@ -366,6 +395,9 @@ def close(self): self._saved_forwards = () self._invokers = () self._replacements = () + self._sampling_forward_begin = None + self._sampling_forward_end = None + self._sampling_allocation_failure = None super().close() diff --git a/toolkit/memory_management/arena_offload/policy.py b/toolkit/memory_management/arena_offload/policy.py index 8fa66bcdb9..e63ce20770 100644 --- a/toolkit/memory_management/arena_offload/policy.py +++ b/toolkit/memory_management/arena_offload/policy.py @@ -60,6 +60,8 @@ def __init__( self.last_worst_shape_allocator_slack_bytes = None self.last_aggressive_capacity = 0 self.last_aggressive_gate = False + self.last_needed_promotion_cap_bytes = None + self.last_cache_pad_bytes = self.allocator_cache_headroom_bytes self.pending_promotion = None self.last_safe_residency_bytes = None self.last_rejected_residency_bytes = None @@ -67,7 +69,8 @@ def __init__( def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, worst_shape_free_bytes, worst_shape_allocator_slack_bytes=None, - current_cap_bytes=None, aggressive_promotion_capacity=0): + current_cap_bytes=None, aggressive_promotion_capacity=0, + worst_shape_live_bytes=None, learned_cache_pad_bytes=None): if not self.bootstrapped: self.bootstrapped = True if demote_candidate is not None: @@ -88,30 +91,90 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, if worst_shape_allocator_slack_bytes is None else worst_shape_allocator_slack_bytes ) + exact_funding = ( + learned_cache_pad_bytes is not None + and worst_shape_live_bytes is not None + ) + cache_pad = ( + self.allocator_cache_headroom_bytes + if learned_cache_pad_bytes is None + else max(0, int(learned_cache_pad_bytes)) + ) throughput_ok = transfer_benefits_from_residency(signal.get("transfer")) worst_ok = candidate is not None and int(worst_shape_free_bytes) >= 0 - promote_ok = ( - candidate is not None - and throughput_ok - and worst_ok - and vram_budget.residency_promote_ok( - retries, - allocator_slack, - block_bytes, - self.allocator_cache_headroom_bytes, + if exact_funding: + promotion_cap_possible = ( + candidate is not None + and vram_budget.cap_can_host_promotion( + int(worst_shape_live_bytes), + block_bytes, + cache_pad, + int(cliff_cap_bytes), + ) + ) + promote_ok = ( + candidate is not None + and promotion_cap_possible + and throughput_ok + and worst_ok + and retries == 0 + and device_frees == 0 + ) + else: + promotion_cap_possible = True + promote_ok = ( + candidate is not None + and throughput_ok + and worst_ok + and vram_budget.residency_promote_ok( + retries, + allocator_slack, + block_bytes, + cache_pad, + ) ) - ) active_cap = ( int(cliff_cap_bytes) if current_cap_bytes is None else int(current_cap_bytes) ) - cap_covers = ( - candidate is not None - and allocator_slack - > block_bytes + self.allocator_cache_headroom_bytes - ) - binding = retries > 0 or int(worst_shape_free_bytes) < 0 + if exact_funding: + needed_promotion_cap = vram_budget.cap_bytes_for_live( + int(worst_shape_live_bytes) + block_bytes, + cache_pad, + int(cliff_cap_bytes), + ) + cap_covers = ( + candidate is not None + and active_cap >= needed_promotion_cap + ) + binding = retries > 0 + pressure_needed_cap = vram_budget.cap_bytes_for_live( + int(worst_shape_live_bytes), + cache_pad, + int(cliff_cap_bytes), + ) + needed_cap = ( + needed_promotion_cap + if promote_ok and not cap_covers + else pressure_needed_cap + ) + else: + needed_promotion_cap = None + cap_covers = ( + candidate is not None + and allocator_slack > block_bytes + cache_pad + ) + binding = retries > 0 or int(worst_shape_free_bytes) < 0 + cap_raise_bytes = ( + block_bytes + if promote_ok and block_bytes > 0 + else cache_pad + ) + needed_cap = min( + int(cliff_cap_bytes), + active_cap + max(0, int(cap_raise_bytes)), + ) aggressive_capacity = max(0, int(aggressive_promotion_capacity or 0)) bootstrap_pending = bool( self.pending_promotion is not None @@ -125,8 +188,7 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, and retries == 0 and device_frees == 0 and worst_ok - and allocator_slack - > block_bytes + self.allocator_cache_headroom_bytes + and allocator_slack > block_bytes + cache_pad ) self.last_worst_shape_physical_headroom_bytes = int( worst_shape_free_bytes @@ -137,15 +199,8 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, self.last_cap_covers_promo = bool(cap_covers) self.last_aggressive_capacity = aggressive_capacity self.last_aggressive_gate = bool(aggressive_ok) - cap_raise_bytes = ( - block_bytes - if promote_ok and block_bytes > 0 - else self.allocator_cache_headroom_bytes - ) - needed_cap = min( - int(cliff_cap_bytes), - active_cap + max(0, int(cap_raise_bytes)), - ) + self.last_needed_promotion_cap_bytes = needed_promotion_cap + self.last_cache_pad_bytes = cache_pad if ( self.pending_promotion is not None and (retries > 0 or device_frees > 0) @@ -216,15 +271,21 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, "worst_shape_veto" if candidate is not None and not worst_ok else ( + "promotion_exceeds_cliff" + if candidate is not None + and exact_funding + and not promotion_cap_possible + else ( "throughput_gate" if candidate is not None and not throughput_ok else ( "allocator_cache_headroom" if candidate is not None - and allocator_slack - <= block_bytes + self.allocator_cache_headroom_bytes + and not exact_funding + and allocator_slack <= block_bytes + cache_pad else "fsm_hold" ) + ) ) ) return self._hold(reason, candidate=candidate) @@ -360,6 +421,10 @@ def diagnostics(self): ), "last_aggressive_capacity": self.last_aggressive_capacity, "last_aggressive_gate": self.last_aggressive_gate, + "last_needed_promotion_cap_bytes": ( + self.last_needed_promotion_cap_bytes + ), + "last_cache_pad_bytes": self.last_cache_pad_bytes, "pending_promotion": self.pending_promotion, "last_safe_residency_bytes": self.last_safe_residency_bytes, "last_rejected_residency_bytes": ( @@ -416,6 +481,19 @@ def last_signal(self): def transfer_snapshot_due(self): return self._transfer_steps + 1 >= self.transfer_window_steps + def prime_counters(self, *, allocator_counters=None, compile_counters=None): + """Start a new observation phase without attributing old activity.""" + if allocator_counters is not None: + self._allocator_previous = { + key: int((allocator_counters or {}).get(key, 0) or 0) + for key in _ALLOC + } + if compile_counters is not None: + self._compile_previous = { + key: int((compile_counters or {}).get(key, 0) or 0) + for key in _COMPILE + } + def invalidate_shape_peaks(self): self._shape_peaks.clear() self._last_signal = None diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index 1d457cbbb3..1799324e29 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -12,10 +12,12 @@ from __future__ import annotations import contextlib +import sys import time import warnings +from collections import Counter from collections.abc import Sequence -from dataclasses import replace +from dataclasses import dataclass, replace from typing import Any from .. import allocator_cap @@ -31,6 +33,7 @@ ArenaResidencyController, TrainingSignalWindow, ) +from .cap_calibrator import CAP_SET, TrainingCapCalibrator from .errors import ArenaCleanupError, ArenaSetupFatalError from .fp8 import disable as disable_fp8 from .fp8 import enable as enable_fp8 @@ -50,6 +53,14 @@ BOOTSTRAP_MIN_STEP = 2 +@dataclass +class _SamplingCapProfile: + calibrator: TrainingCapCalibrator + signals: TrainingSignalWindow + occurrences: int = 0 + forward_count: int = 0 + + class ArenaOffloadRuntime: """Lifecycle + execution contexts for one arena-offloaded transformer.""" @@ -89,6 +100,20 @@ def __init__( self._last_policy_error: str | None = None self._last_failure_event: dict | None = None self._policy = ArenaResidencyController() + cap_calibration_requested = bool(config._policy.cap_calibration) + cap_calibration_enabled = ( + cap_calibration_requested and sys.platform == "win32" + ) + if cap_calibration_requested and not cap_calibration_enabled: + warnings.warn( + "arena allocator-cap calibration is currently Windows-only", + RuntimeWarning, + stacklevel=2, + ) + self._cap_calibrator = TrainingCapCalibrator( + enabled=cap_calibration_enabled + ) + self._cap_calibration_enabled = cap_calibration_enabled self._last_training_cap_target_bytes: int | None = None self._last_training_pressure_relief: dict | None = None self._bootstrap_complete = False @@ -100,6 +125,28 @@ def __init__( self._training_fp8_singletons = 0 self._sampling_fp8_canonical = 0 self._sampling_fp8_singletons = 0 + self._sampling_cap_profiles: dict[tuple, _SamplingCapProfile] = {} + self._sampling_session_occurrences = Counter() + self._active_sampling_shape_key: tuple | None = None + self._active_sampling_cap_profile: _SamplingCapProfile | None = None + self._active_sampling_resident_bytes = 0 + self._active_sampling_ring_bytes = 0 + self._sampling_forward_started_at: float | None = None + self._sampling_oom_retry_used = False + callback_setter = getattr( + self._executor, "set_sampling_forward_callbacks", None + ) + if callback_setter is not None: + callback_setter( + self._sampling_forward_begin, + self._sampling_forward_end, + self._sampling_allocation_failure, + ) + promotion_setter = getattr( + self._executor, "set_residency_promotion_callback", None + ) + if promotion_setter is not None: + promotion_setter(self._reserve_allocator_for_promotion) self._permanent_placement = None self._device_state_parked_plan = None @@ -536,7 +583,7 @@ def training_step(self, *, shape_key: tuple | None = None, step_num: int | None self._last_shape_key = shape_key self._last_step_num = step_num try: - self._apply_training_policy() + self._apply_training_policy(shape_key=shape_key) except BaseException as error: try: self._handle_training_failure( @@ -600,8 +647,44 @@ def _handle_training_failure(self, error, *, shape_key, step_num): rollback = None recoverable = False if allocation_failure: - decision = self._policy.allocation_failure() - if decision.action == "rollback" and decision.block_key is not None: + cap_decision = None + cap_calibrator = getattr(self, "_cap_calibrator", None) + if cap_calibrator is not None and cap_calibrator.active: + cliff_cap = allocator_cap.wddm_cliff_cap_bytes( + self._device, self._config._policy.wddm_hard_gib + ) + current_cap = int( + self._last_training_cap_target_bytes or cliff_cap + ) + cap_decision = cap_calibrator.allocation_failure( + cliff_cap_bytes=cliff_cap, + current_cap_bytes=current_cap, + ) + if cap_decision is not None: + if cap_decision.target_cap_bytes is not None: + allocator_cap.configure_wddm_allocator_guard( + self._device, + self._config._policy.wddm_hard_gib, + target_cap_bytes=cap_decision.target_cap_bytes, + strict=getattr( + self._config, "strict_vram_cap", False + ), + log_prefix="[ArenaOffload]", + ) + self._last_training_cap_target_bytes = ( + cap_decision.target_cap_bytes + ) + rollback = "allocator_cap_probe" + recoverable = not bool( + getattr(self._config, "strict_vram_cap", False) + ) + else: + decision = self._policy.allocation_failure() + if ( + cap_decision is None + and decision.action == "rollback" + and decision.block_key is not None + ): rollback_keys = tuple(decision.block_keys or ()) if rollback_keys: self.transition_training_blocks( @@ -633,7 +716,7 @@ def _handle_training_failure(self, error, *, shape_key, step_num): recoverable = not bool( getattr(self._config, "strict_vram_cap", False) ) - elif not bool( + elif cap_decision is None and not bool( getattr(self._config, "strict_vram_cap", False) ): candidates = self._demotion_candidates() @@ -715,10 +798,318 @@ def _singleton_runtime_ids(self): def _canonical_runtime_ids(self): return {id(module) for module in self._canonical_modules} + def _sampling_hard_gib(self) -> float: + value = self._config._policy.sampling_wddm_hard_gib + return 1.0 if value is None else float(value) + + def _sampling_cliff_cap_bytes(self) -> int: + import torch + + if ( + torch.device(self._device).type != "cuda" + or not torch.cuda.is_available() + ): + return 0 + return allocator_cap.wddm_cliff_cap_bytes( + self._device, self._sampling_hard_gib() + ) + + def _reserve_allocator_for_promotion(self, promotion_bytes, plan) -> None: + """Grow the cap before resident sidecars consume calibrated allowance.""" + from .. import vram_budget + + phase = str(getattr(plan, "phase", "")) + sampling = phase.startswith("sample") + hard_gib = ( + self._sampling_hard_gib() + if sampling + else self._config._policy.wddm_hard_gib + ) + cliff = ( + self._sampling_cliff_cap_bytes() + if sampling + else allocator_cap.wddm_cliff_cap_bytes(self._device, hard_gib) + ) + current = int(allocator_cap.applied_cap_bytes(self._device) or cliff) + target = vram_budget.cap_bytes_preserving_allowance_after_promotion( + current, + int(promotion_bytes), + cliff, + ) + if target <= current: + return + allocator_cap.configure_wddm_allocator_guard( + self._device, + hard_gib, + target_cap_bytes=target, + strict=getattr(self._config, "strict_vram_cap", False), + log_prefix="[ArenaOffload]", + force=True, + ) + applied = int(allocator_cap.applied_cap_bytes(self._device) or target) + if sampling: + profile = getattr(self, "_active_sampling_cap_profile", None) + if profile is not None: + profile.calibrator.invalidate_for_residency_growth() + else: + self._last_training_cap_target_bytes = applied + self._cap_calibrator.invalidate_for_residency_growth() + print( + "[ArenaOffload] allocator cap matched residency promotion: " + f"resident=+{int(promotion_bytes) / GIB:.2f} GiB " + f"cap={current / GIB:.2f}->{applied / GIB:.2f} GiB" + ) + + def _bind_sampling_cap(self, target_cap_bytes=None, *, reclaim=False) -> int: + """Bind one sampling cap, optionally forcing one bounded cache settle.""" + import torch + + cliff = self._sampling_cliff_cap_bytes() + if cliff <= 0: + return 0 + target = cliff if target_cap_bytes is None else min( + cliff, max(0, int(target_cap_bytes)) + ) + before = allocator_cap.applied_cap_bytes(self._device) + allocator_cap.configure_wddm_allocator_guard( + self._device, + self._sampling_hard_gib(), + target_cap_bytes=target, + strict=getattr(self._config, "strict_vram_cap", False), + log_prefix="[ArenaOffload]", + ) + applied = int(allocator_cap.applied_cap_bytes(self._device) or target) + if ( + reclaim + and before is not None + and applied < int(before) - 64 * 1024**2 + and torch.cuda.is_available() + ): + # Same-shape cache hits bypass both the cap and gc_threshold. One + # explicit phase-boundary trim guarantees that the next forward is + # the settlement window instead of silently reusing the old cache. + torch.cuda.empty_cache() + return applied + + def _sampling_allocator_counters(self): + import torch + + try: + stats = torch.cuda.memory_stats(self._device) + except Exception: + stats = {} + return { + key: int(stats.get(key, 0) or 0) + for key in ( + "num_alloc_retries", + "num_device_alloc", + "num_device_free", + ) + } + + def _sampling_profile(self, shape_key, occurrences): + profiles = getattr(self, "_sampling_cap_profiles", None) + if profiles is None: + profiles = self._sampling_cap_profiles = {} + profile = profiles.get(shape_key) + eligible = bool( + getattr(self, "_cap_calibration_enabled", False) + and (int(occurrences) > 1 or profile is not None) + ) + if profile is None and eligible: + profile = _SamplingCapProfile( + calibrator=TrainingCapCalibrator( + enabled=True, monitor_settled=True + ), + signals=TrainingSignalWindow(), + ) + profiles[shape_key] = profile + if profile is not None: + profile.occurrences = max(profile.occurrences, int(occurrences)) + return profile + + def _sampling_profile_cap(self, profile, cliff_cap_bytes): + if profile is None: + return int(cliff_cap_bytes) + calibrator = profile.calibrator + return int( + calibrator.probe_cap_bytes + or calibrator.settled_cap_bytes + or calibrator.last_clean_cap_bytes + or cliff_cap_bytes + ) + + def _sampling_forward_begin(self): + """Act on the previous pass, then start one transformer measurement.""" + import torch + + profile = getattr(self, "_active_sampling_cap_profile", None) + shape_key = getattr(self, "_active_sampling_shape_key", None) + if profile is None or shape_key is None: + return + self._sampling_oom_retry_used = False + cliff = self._sampling_cliff_cap_bytes() + current = int(allocator_cap.applied_cap_bytes(self._device) or cliff) + decision = profile.calibrator.step( + profile.signals.last_signal, + upcoming_shape_key=shape_key, + shape_peaks=profile.signals.shape_peaks, + resident_bytes=int( + getattr(self, "_active_sampling_resident_bytes", 0) + ), + ring_bytes=int(getattr(self, "_active_sampling_ring_bytes", 0)), + cliff_cap_bytes=cliff, + current_cap_bytes=current, + ) + if decision.action == CAP_SET: + lowering = ( + decision.target_cap_bytes is not None + and int(decision.target_cap_bytes) < current + ) + self._bind_sampling_cap( + decision.target_cap_bytes, + reclaim=lowering, + ) + if lowering: + # The phase-boundary trim uses the same num_device_free + # counter as allocator GC. Exclude our own trim so the first + # free observed during the forward is unambiguously pressure. + profile.signals.prime_counters( + allocator_counters=self._sampling_allocator_counters() + ) + print( + "[ArenaOffload] sampling cap calibration: " + f"shape={shape_key} state={profile.calibrator.state} " + f"reason={decision.reason} target={decision.target_cap_bytes}" + ) + torch.cuda.reset_peak_memory_stats(self._device) + self._sampling_forward_started_at = time.perf_counter() + + def _sampling_allocation_failure(self, error) -> bool: + """Widen a rejected sampling probe and allow one block retry.""" + if getattr(self, "_sampling_oom_retry_used", False): + return False + recovery = self._widen_sampling_cap_after_oom(error) + if recovery is None: + return False + current, applied = recovery + self._sampling_oom_retry_used = True + print( + "[ArenaOffload] sampling allocator OOM: " + f"widening {current / GIB:.2f}->{applied / GIB:.2f} GiB " + "and retrying the transformer block once" + ) + return True + + def _widen_sampling_cap_after_oom(self, error): + """Return (old, new) after rejecting one capped sampling allocation.""" + import torch + + if not isinstance(error, torch.cuda.OutOfMemoryError): + return None + profile = getattr(self, "_active_sampling_cap_profile", None) + if profile is None: + return None + cliff = self._sampling_cliff_cap_bytes() + current = int(allocator_cap.applied_cap_bytes(self._device) or cliff) + decision = profile.calibrator.allocation_failure( + cliff_cap_bytes=cliff, + current_cap_bytes=current, + ) + if decision is None or decision.target_cap_bytes is None: + return None + target = int(decision.target_cap_bytes) + if target <= current: + return None + applied = self._bind_sampling_cap(target) + if applied <= current: + return None + # Attribute neither the rejected allocation retry nor its allocator + # cleanup to the widened-cap verification pass. + profile.signals.prime_counters( + allocator_counters=self._sampling_allocator_counters() + ) + return current, applied + + def _sampling_forward_end(self): + """Publish one completed transformer pass to the sampling FSM.""" + import torch + + from ..vram_budget import device_free_bytes + + profile = getattr(self, "_active_sampling_cap_profile", None) + shape_key = getattr(self, "_active_sampling_shape_key", None) + if profile is None or shape_key is None: + return + started = self._sampling_forward_started_at + peak_allocated = torch.cuda.max_memory_allocated(self._device) + peak_reserved = torch.cuda.max_memory_reserved(self._device) + record_peak = getattr(self._executor, "record_sampling_peak", None) + if record_peak is not None: + record_peak( + allocated_bytes=peak_allocated, + reserved_bytes=peak_reserved, + ) + profile.signals.observe( + shape_key=shape_key, + step_num=profile.forward_count, + allocator_counters=self._sampling_allocator_counters(), + peak_allocated_bytes=peak_allocated, + peak_reserved_bytes=peak_reserved, + device_free_bytes=device_free_bytes(self._device), + resident_bytes=int( + getattr(self, "_active_sampling_resident_bytes", 0) + ), + ring_bytes=int(getattr(self, "_active_sampling_ring_bytes", 0)), + compile_counters=_compile_counter_snapshot(torch), + transfer_counters=None, + step_wall_ms=( + 0.0 + if started is None + else (time.perf_counter() - started) * 1000.0 + ), + ) + profile.forward_count += 1 + self._sampling_forward_started_at = None + + @contextlib.contextmanager + def _sampling_decode_cap_guard(self, generation_owner): + """Restore the decode-safe cliff at a generic VAE decode boundary.""" + vae = getattr(generation_owner, "vae", None) + original_decode = getattr(vae, "decode", None) + if vae is None or not callable(original_decode): + yield + return + instance_dict = getattr(vae, "__dict__", {}) + had_instance_decode = "decode" in instance_dict + original_instance_decode = instance_dict.get("decode") + + runtime = self + + def guarded_decode(*args, **kwargs): + runtime._bind_sampling_cap(None) + return original_decode(*args, **kwargs) + + try: + setattr(vae, "decode", guarded_decode) + except (AttributeError, RuntimeError, TypeError) as error: + raise RuntimeError("sampling_decode_cap_guard_unavailable") from error + try: + yield + finally: + if getattr(vae, "decode", None) is guarded_decode: + if had_instance_decode: + setattr(vae, "decode", original_instance_decode) + else: + delattr(vae, "decode") + @contextlib.contextmanager - def sampling_session(self): + def sampling_session(self, *, gen_configs=(), generation_owner=None): """Wrap a sampling run and restore TRAIN once at the end.""" self._require_open() + self._sampling_session_occurrences = Counter( + _sampling_config_shape_key(config) for config in (gen_configs or ()) + ) sampling_restores = [] if self._config.fp8_sampling: canonical_ids = self._canonical_runtime_ids() @@ -736,15 +1127,30 @@ def sampling_session(self): self._sampling_fp8_canonical = len(installed_ids & canonical_ids) self._sampling_fp8_singletons = len(installed_ids & singleton_ids) try: - yield self + with self._sampling_decode_cap_guard(generation_owner): + yield self finally: + self._active_sampling_shape_key = None + self._active_sampling_cap_profile = None + self._active_sampling_resident_bytes = 0 + self._active_sampling_ring_bytes = 0 + self._sampling_forward_started_at = None + self._sampling_oom_retry_used = False + self._sampling_session_occurrences = Counter() if sampling_restores: disable_fp8(sampling_restores) self._bind_training_cap() self._executor.activate(self._executor.TRAIN, self._training_plan) @contextlib.contextmanager - def sampling_image(self, *, shape_key: tuple, cold_working_bytes: int): + def sampling_image( + self, + *, + gen_config=None, + generation_owner=None, + shape_key: tuple | None = None, + cold_working_bytes: int | None = None, + ): """The sampling phase boundary for ONE image. Switches to the permanent SAMPLE program (forward-only, no @@ -754,18 +1160,44 @@ def sampling_image(self, *, shape_key: tuple, cold_working_bytes: int): self._require_open() policy = self._config._policy - fixed_working_bytes = _fixed_working_bytes(policy.sampling_working_reserve_gib) - hard_gib = ( - 1.0 - if policy.sampling_wddm_hard_gib is None - else float(policy.sampling_wddm_hard_gib) + if gen_config is not None: + shape_key = _sampling_config_shape_key(gen_config) + if cold_working_bytes is None: + cold_working_bytes = _sampling_cold_working_bytes( + gen_config, fp8_native=bool(self._config.fp8_sampling) + ) + if shape_key is None: + raise ValueError("sampling_shape_key_required") + if cold_working_bytes is None: + raise ValueError("sampling_cold_working_bytes_required") + + occurrences = int( + getattr(self, "_sampling_session_occurrences", {}).get(shape_key, 1) ) - allocator_cap.configure_wddm_allocator_guard( - self._device, - hard_gib, - strict=getattr(self._config, "strict_vram_cap", False), - log_prefix="[ArenaOffload]", + profile = self._sampling_profile(shape_key, occurrences) + decode = getattr(getattr(generation_owner, "vae", None), "decode", None) + if profile is not None and not callable(decode): + # A lowered transformer cap must never leak into an unknown decode + # path. Keep measuring the ordinary image reserve, but calibrate + # only when the generic VAE boundary can restore the cliff. + profile = None + cliff_cap = self._sampling_cliff_cap_bytes() + target_cap = self._sampling_profile_cap(profile, cliff_cap) + active_cap = self._bind_sampling_cap( + target_cap, reclaim=target_cap < cliff_cap ) + if profile is not None: + import torch + + profile.signals.prime_counters( + allocator_counters=self._sampling_allocator_counters(), + compile_counters=_compile_counter_snapshot(torch), + ) + self._active_sampling_shape_key = shape_key + self._active_sampling_cap_profile = profile + + fixed_working_bytes = _fixed_working_bytes(policy.sampling_working_reserve_gib) + hard_gib = self._sampling_hard_gib() physical_headroom_gib = resolve_physical_vram_headroom_gib( self._device, policy.sampling_physical_vram_headroom_gib, @@ -780,16 +1212,65 @@ def sampling_image(self, *, shape_key: tuple, cold_working_bytes: int): ) ) ) - with self._executor.sampling( - shape_key=shape_key, - cold_working_bytes=int(cold_working_bytes), - fixed_working_bytes=fixed_working_bytes, - cold_floor_bytes=( - int(physical_headroom_gib * GIB) + dequant_reserve - ), - hot_floor_bytes=int((hard_gib + 0.25) * GIB) + dequant_reserve, - ): - yield self + setup_retry_used = False + try: + with self._sampling_decode_cap_guard(generation_owner): + while True: + yielded = False + try: + with self._executor.sampling( + shape_key=shape_key, + cold_working_bytes=int(cold_working_bytes), + fixed_working_bytes=fixed_working_bytes, + cold_floor_bytes=( + int(physical_headroom_gib * GIB) + dequant_reserve + ), + hot_floor_bytes=( + int((hard_gib + 0.25) * GIB) + dequant_reserve + ), + # Keep planning against the rejected cap so the + # recovery margin is not spent on more residents. + allocator_cap_bytes=active_cap, + allocator_hard_bytes=int(hard_gib * GIB), + ): + residency = getattr(self, "_residency", None) + self._active_sampling_resident_bytes = ( + 0 + if residency is None + else int(residency.resident_bytes()) + ) + int( + (self._smart_plan or {}).get( + "singleton_resident_bytes", 0 + ) + ) + self._active_sampling_ring_bytes = ( + self._training_ring_bytes() + if hasattr(self, "_arena") and residency is not None + else 0 + ) + yielded = True + yield self + break + except BaseException as error: + if yielded or setup_retry_used: + raise + recovery = self._widen_sampling_cap_after_oom(error) + if recovery is None: + raise + current, applied = recovery + setup_retry_used = True + print( + "[ArenaOffload] sampling setup OOM: " + f"widening {current / GIB:.2f}->{applied / GIB:.2f} GiB " + "and retrying residency setup once" + ) + finally: + self._active_sampling_shape_key = None + self._active_sampling_cap_profile = None + self._active_sampling_resident_bytes = 0 + self._active_sampling_ring_bytes = 0 + self._sampling_forward_started_at = None + self._sampling_oom_retry_used = False # ------------------------------------------------------------------ def record_training_physical_free_min(self, free_bytes) -> None: @@ -1072,11 +1553,20 @@ def _worst_shape_candidate_physical_headroom_bytes(self, candidate): ) def _worst_shape_allocator_slack_bytes(self, current_cap_bytes): - peaks = self._signals.shape_peaks - if not any(peak.steps > 0 for peak in peaks.values()): + predicted_live = self._worst_shape_live_bytes() + if predicted_live is None: return 0 from .. import vram_budget + return vram_budget.allocator_allowance_bytes( + current_cap_bytes, predicted_live + ) + + def _worst_shape_live_bytes(self): + peaks = self._signals.shape_peaks + if not any(peak.steps > 0 for peak in peaks.values()): + return None + worst_working = max( int(peak.working_peak_bytes) for peak in peaks.values() @@ -1086,16 +1576,13 @@ def _worst_shape_allocator_slack_bytes(self, current_cap_bytes): int(self._residency.resident_bytes()) + int((self._smart_plan or {}).get("singleton_resident_bytes", 0)) ) - predicted_live = ( + return ( worst_working + current_resident + self._training_ring_bytes() ) - return vram_budget.allocator_allowance_bytes( - current_cap_bytes, predicted_live - ) - def _apply_training_policy(self): + def _apply_training_policy(self, *, shape_key=None): import torch if torch.device(self._device).type != "cuda" or not torch.cuda.is_available(): @@ -1113,7 +1600,56 @@ def _apply_training_policy(self): cliff_cap, int(self._last_training_cap_target_bytes or cliff_cap), ) - if self._bootstrap_training_residency(current_cap): + cap_decision = self._cap_calibrator.step( + signal, + upcoming_shape_key=shape_key, + shape_peaks=self._signals.shape_peaks, + resident_bytes=( + int(self._residency.resident_bytes()) + + int((self._smart_plan or {}).get("singleton_resident_bytes", 0)) + ), + ring_bytes=self._training_ring_bytes(), + cliff_cap_bytes=cliff_cap, + current_cap_bytes=current_cap, + ) + if self._cap_calibrator.enabled and ( + cap_decision.action == CAP_SET + or cap_decision.reason == "calibration_settled" + ): + cap_diag = self._cap_calibrator.diagnostics() + transfer_diag = (signal or {}).get("transfer") or {} + print( + "[ArenaOffload] cap calibration: " + f"state={cap_diag['state']} reason={cap_decision.reason} " + f"target={cap_decision.target_cap_bytes} " + f"predicted={cap_diag['predicted_initial_cap_bytes']} " + f"clean={cap_diag['last_clean_cap_bytes']} " + f"dirty={cap_diag['last_dirty_cap_bytes']} " + f"cache_pad={cap_diag['learned_cache_pad_bytes']} " + f"buckets={len(cap_diag['buckets'])} " + f"resident_blocks={len(self._demotion_candidates())} " + f"h2d_bytes={int(transfer_diag.get('bytes', 0) or 0)} " + f"h2d_ms={float(transfer_diag.get('h2d_ms', 0.0) or 0.0):.3f} " + f"step_ms={float(transfer_diag.get('step_wall_ms', 0.0) or 0.0):.3f}" + ) + if cap_decision.action == CAP_SET: + allocator_cap.configure_wddm_allocator_guard( + self._device, + self._config._policy.wddm_hard_gib, + target_cap_bytes=cap_decision.target_cap_bytes, + strict=getattr(self._config, "strict_vram_cap", False), + log_prefix="[ArenaOffload]", + ) + self._last_training_cap_target_bytes = ( + cap_decision.target_cap_bytes + ) + current_cap = int(cap_decision.target_cap_bytes) + if cap_decision.hold_residency: + return + if ( + not self._cap_calibrator.enabled + and self._bootstrap_training_residency(current_cap) + ): return aggressive_capacity = self._aggressive_promotion_capacity(current_cap) decision = self._policy.step( @@ -1130,6 +1666,10 @@ def _apply_training_policy(self): worst_shape_allocator_slack_bytes=( self._worst_shape_allocator_slack_bytes(current_cap) ), + worst_shape_live_bytes=self._worst_shape_live_bytes(), + learned_cache_pad_bytes=( + self._cap_calibrator.learned_cache_pad_bytes + ), aggressive_promotion_capacity=aggressive_capacity, ) if decision.action == "promote": @@ -1265,6 +1805,18 @@ def diagnostics(self) -> dict: "sampling_fp8_singletons": getattr( self, "_sampling_fp8_singletons", 0 ), + "sampling_cap_profiles": [ + { + "shape_key": shape_key, + "occurrences": profile.occurrences, + "forward_count": profile.forward_count, + **profile.calibrator.diagnostics(), + } + for shape_key, profile in sorted( + getattr(self, "_sampling_cap_profiles", {}).items(), + key=lambda item: repr(item[0]), + ) + ], "largest_singleton_bf16_dequant_bytes": int( (self._smart_plan or {}).get( "largest_singleton_bf16_dequant_bytes", 0 @@ -1301,6 +1853,7 @@ def diagnostics(self) -> dict: "successful_training_steps": self._successful_training_steps, "policy": { **self._signals.diagnostics(), + "cap_calibration": self._cap_calibrator.diagnostics(), "controller": self._policy.diagnostics(), }, "policy_error": self._last_policy_error, @@ -1477,6 +2030,63 @@ def _compile_counter_snapshot(torch_module): } +def _config_value(config, name, default=None): + if isinstance(config, dict): + return config.get(name, default) + return getattr(config, name, default) + + +def _sampling_config_shape_key(config) -> tuple: + """Conservative, model-agnostic sampling key from job configuration.""" + controls = tuple( + None if value is None else str(value) + for value in ( + _config_value(config, "ctrl_img"), + _config_value(config, "ctrl_img_1"), + _config_value(config, "ctrl_img_2"), + _config_value(config, "ctrl_img_3"), + ) + ) + extra_values = _config_value(config, "extra_values", ()) or () + return ( + "sample", + int(_config_value(config, "height", 0) or 0), + int(_config_value(config, "width", 0) or 0), + int(_config_value(config, "num_frames", 1) or 1), + float(_config_value(config, "guidance_scale", 0.0) or 0.0), + bool(_config_value(config, "batch_cfg", False)), + _config_value(config, "ctrl_idx"), + controls, + len(extra_values), + ) + + +def _sampling_cold_working_bytes(config, *, fp8_native=True) -> int: + """Build the existing cold reserve estimate directly from job config.""" + from .. import vram_budget + + width = max(1, int(_config_value(config, "width", 1) or 1)) + height = max(1, int(_config_value(config, "height", 1) or 1)) + frames = max(1, int(_config_value(config, "num_frames", 1) or 1)) + image_tokens = ((width + 15) // 16) * ((height + 15) // 16) * frames + references = { + str(value) + for value in ( + _config_value(config, "ctrl_img"), + _config_value(config, "ctrl_img_1"), + _config_value(config, "ctrl_img_2"), + _config_value(config, "ctrl_img_3"), + ) + if value is not None + } + image_tokens *= 1 + len(references) + return vram_budget.estimate_sampling_working_reserve_bytes( + image_tokens, + batch_cfg=bool(_config_value(config, "batch_cfg", False)), + fp8_native=bool(fp8_native), + ) + + def _fixed_working_bytes(value) -> int | None: """None when the sampling working reserve is auto (unset, negative, 'auto').""" if value is None: diff --git a/toolkit/memory_management/immutable_runtime.py b/toolkit/memory_management/immutable_runtime.py index afe4e4b12b..e26aef8404 100644 --- a/toolkit/memory_management/immutable_runtime.py +++ b/toolkit/memory_management/immutable_runtime.py @@ -333,6 +333,7 @@ def __init__( self._finalization_signature = None self._active_token = None self._active_mode = None + self._residency_promotion_callback = None self.stats = { "residency_transitions": 0, "source_generation": self._sources.generation, @@ -358,6 +359,10 @@ def source(self, block_index: int) -> ImmutableBlockSourceSnapshot: """Current published source snapshot for one block.""" return self._sources.source(block_index) + def set_residency_promotion_callback(self, callback=None) -> None: + """Install a pre-allocation hook for exact sidecar promotion bytes.""" + self._residency_promotion_callback = callback + def _require_finalized(self) -> None: if not self._finalized: raise ImmutableRuntimeError("immutable_runtime_not_finalized") @@ -511,6 +516,9 @@ def _warn_hint_out_of_range(self, dim, size, lo, hi) -> None: def set_residency_plan(self, plan: ResidencyPlan) -> ResidencyDelta: self._assert_arena_stable("pre_residency_publish") + promotion_bytes = self.residency.planned_addition_bytes(plan) + if promotion_bytes and self._residency_promotion_callback is not None: + self._residency_promotion_callback(promotion_bytes, plan) previous, newly_pinned, keep = self._prepare_plan_pins(plan) try: delta = self._sources.publish(plan) @@ -846,6 +854,8 @@ def activate_sampling_image( fixed_working_bytes: int | None, cold_floor_bytes: int, hot_floor_bytes: int, + allocator_cap_bytes: int | None = None, + allocator_hard_bytes: int = 0, measured_pad_bytes: int = 256 * 1024**2, measured_floor_bytes: int = 512 * 1024**2, ) -> ImmutableProgram: @@ -879,6 +889,16 @@ def activate_sampling_image( allocated_bytes = torch.cuda.memory_allocated(device) reserved_bytes = torch.cuda.memory_reserved(device) reclaimable_cache = max(0, reserved_bytes - allocated_bytes) + if allocator_cap_bytes is not None: + total_bytes = vram_budget.device_total_bytes(device) + allocator_free = vram_budget.sampling_allocator_budget_free_bytes( + total_bytes, + allocated_bytes, + float(allocator_cap_bytes) / float(max(1, total_bytes)), + int(allocator_hard_bytes), + ) + if allocator_free is not None: + free_bytes = max(int(free_bytes), int(allocator_free)) current_sidecars = self.residency.resident_bytes() resident_budget = max( 0, @@ -910,6 +930,8 @@ def activate_sampling_image( "shape_key": shape_key, "allocated": baseline_allocated, "reserved": baseline_reserved, + "external_peak_allocated": baseline_allocated, + "external_peak_reserved": baseline_reserved, "working_bytes": working_bytes, "floor_bytes": floor_bytes, "source": reserve_source, @@ -926,6 +948,20 @@ def activate_sampling_image( ) return self.program(self.SAMPLE) + def record_sampling_peak(self, *, allocated_bytes: int, reserved_bytes: int): + """Preserve pass peaks when a sampling controller resets CUDA stats.""" + baseline = self._sampling_baseline + if baseline is None: + return + baseline["external_peak_allocated"] = max( + int(baseline.get("external_peak_allocated", 0)), + int(allocated_bytes), + ) + baseline["external_peak_reserved"] = max( + int(baseline.get("external_peak_reserved", 0)), + int(reserved_bytes), + ) + def finish_sampling_image(self, *, shape_key: tuple) -> int: baseline = self._sampling_baseline if baseline is None or baseline["shape_key"] != shape_key: @@ -933,8 +969,14 @@ def finish_sampling_image(self, *, shape_key: tuple) -> int: device = self.residency.device torch.cuda.synchronize(device) - allocated_peak = torch.cuda.max_memory_allocated(device) - reserved_peak = torch.cuda.max_memory_reserved(device) + allocated_peak = max( + torch.cuda.max_memory_allocated(device), + int(baseline.get("external_peak_allocated", 0)), + ) + reserved_peak = max( + torch.cuda.max_memory_reserved(device), + int(baseline.get("external_peak_reserved", 0)), + ) allocated_growth = max( 0, allocated_peak - int(baseline["allocated"]), diff --git a/toolkit/memory_management/residency.py b/toolkit/memory_management/residency.py index 1280edd5fc..b88cdcde35 100644 --- a/toolkit/memory_management/residency.py +++ b/toolkit/memory_management/residency.py @@ -340,6 +340,15 @@ def resident_leaf_bytes(self, key: LeafKey) -> int: def resident_bytes(self) -> int: return sum(sidecar.nbytes for sidecar in self._sidecars.values()) + def planned_addition_bytes(self, plan: ResidencyPlan) -> int: + """Exact sidecar payload allocated before an atomic plan is published.""" + additions = plan.resident_leaf_keys - set(self._sidecars) + total = 0 + for key in additions: + _block, spec, _module = self._canonical_leaf(key) + total += sum(int(item.nbytes) for item in spec.tensors) + return total + def synchronize_copies(self) -> None: """Settle queued promotions before their host sources are unpinned.""" if self._copy_stream is not None: diff --git a/toolkit/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py index 9b71ae7826..e61d8684a6 100644 --- a/toolkit/memory_management/vram_budget.py +++ b/toolkit/memory_management/vram_budget.py @@ -47,6 +47,7 @@ from __future__ import annotations +import math import os from dataclasses import dataclass from typing import Optional @@ -765,6 +766,23 @@ def cap_bytes_for_live( return int(want) +def cap_bytes_preserving_allowance_after_promotion( + cap_bytes, + promotion_bytes, + cliff_cap_bytes, + *, + gc_threshold=GC_THRESHOLD, +) -> int: + """Raise a cap enough that resident growth does not consume GC allowance.""" + current = max(0, int(cap_bytes)) + promoted = max(0, int(promotion_bytes)) + cliff = max(0, int(cliff_cap_bytes)) + if promoted <= 0 or current >= cliff: + return min(current, cliff) + growth = math.ceil(float(promoted) / float(gc_threshold)) + return min(cliff, current + int(growth)) + + def cap_can_host_promotion( live_bytes, block_bytes, diff --git a/toolkit/models/base_model.py b/toolkit/models/base_model.py index d53c529211..b69de56d58 100644 --- a/toolkit/models/base_model.py +++ b/toolkit/models/base_model.py @@ -663,13 +663,8 @@ def generate_images( arena_runtime = get_memory_runtime(self.unet) sampling_context = ( arena_runtime.sampling_image( - shape_key=( - "sample", - int(gen_config.height), - int(gen_config.width), - bool(getattr(gen_config, "batch_cfg", False)), - ), - cold_working_bytes=3 * (1024 ** 3), + gen_config=gen_config, + generation_owner=self, ) if arena_runtime is not None else contextlib.nullcontext() diff --git a/toolkit/quantization/fp8_linear.py b/toolkit/quantization/fp8_linear.py index 7f04b0f72d..2006de07fe 100644 --- a/toolkit/quantization/fp8_linear.py +++ b/toolkit/quantization/fp8_linear.py @@ -108,11 +108,14 @@ def _adapt_quanto_fp8(value) -> Fp8LinearDeclaration | None: def _adapt_torchao_fp8(value) -> Fp8LinearDeclaration | None: - try: - from torchao.quantization import Float8Tensor - except ImportError: + from .torchao_compat import ( + torchao_arena_fp8_supported, + torchao_is_float8_tensor, + ) + + if not torchao_arena_fp8_supported(): return None - if not isinstance(value, Float8Tensor): + if not torchao_is_float8_tensor(value): return None qdata, scale = value.qdata, value.scale block_size = tuple(value.block_size or ()) diff --git a/toolkit/quantization/torchao_compat.py b/toolkit/quantization/torchao_compat.py new file mode 100644 index 0000000000..60368208f4 --- /dev/null +++ b/toolkit/quantization/torchao_compat.py @@ -0,0 +1,96 @@ +"""TorchAO API compatibility and optional arena-FP8 capability checks.""" + +from __future__ import annotations + +from importlib.metadata import PackageNotFoundError, version +import re + +import torch + +from torchao.quantization.quant_api import ( + Float8WeightOnlyConfig, + Int8WeightOnlyConfig, + quantize_ as torchao_quantize_, +) + +try: + from torchao.quantization.quant_api import _is_linear as torchao_is_linear +except ImportError: + def torchao_is_linear(module: torch.nn.Module, _fqn: str) -> bool: + return isinstance(module, torch.nn.Linear) + +try: + from torchao.quantization import Float8Tensor as _Float8Tensor +except ImportError: + _Float8Tensor = None + +try: + from torchao.quantization.quant_api import IntxWeightOnlyConfig as _IntxConfig +except ImportError: + _IntxConfig = None + +try: + from torchao.quantization.quant_api import UIntXWeightOnlyConfig as _UIntXConfig +except ImportError: + _UIntXConfig = None + + +TORCHAO_ARENA_FP8_MIN_VERSION = "0.17.0" +_TORCHAO_ARENA_FP8_MIN_RELEASE = (0, 17, 0) + + +def _release_tuple(value: str) -> tuple[int, ...]: + match = re.match(r"\s*(\d+(?:\.\d+)*)", str(value)) + if match is None: + return () + return tuple(int(component) for component in match.group(1).split(".")) + + +try: + TORCHAO_VERSION = version("torchao") +except PackageNotFoundError: + TORCHAO_VERSION = "unknown" + + +def torchao_arena_fp8_supported( + installed_version: str | None = None, + *, + float8_tensor_available: bool | None = None, +) -> bool: + """Whether the tested TorchAO Float8Tensor arena adapter is available.""" + release = _release_tuple( + TORCHAO_VERSION if installed_version is None else installed_version + ) + has_tensor = ( + _Float8Tensor is not None + if float8_tensor_available is None + else bool(float8_tensor_available) + ) + return bool(has_tensor and release >= _TORCHAO_ARENA_FP8_MIN_RELEASE) + + +def torchao_is_float8_tensor(value) -> bool: + return bool(_Float8Tensor is not None and isinstance(value, _Float8Tensor)) + + +def intx_weight_only_config(bits: int): + """Build the low-bit config using either the current or TorchAO 0.10 API.""" + bits = int(bits) + if _IntxConfig is not None: + return _IntxConfig(getattr(torch, f"int{bits}")) + if _UIntXConfig is not None: + return _UIntXConfig(getattr(torch, f"uint{bits}")) + raise RuntimeError("installed TorchAO has no supported IntX weight-only config") + + +__all__ = [ + "Float8WeightOnlyConfig", + "Int8WeightOnlyConfig", + "TORCHAO_ARENA_FP8_MIN_VERSION", + "TORCHAO_VERSION", + "intx_weight_only_config", + "torchao_arena_fp8_supported", + "torchao_is_float8_tensor", + "torchao_is_linear", + "torchao_quantize_", +] diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index 1aa603bdbb..6a4e03d715 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -1,16 +1,16 @@ from fnmatch import fnmatch from typing import List, Optional, Union, TYPE_CHECKING import torch -from torchao.quantization import Float8Tensor from optimum.quanto.quantize import _quantize_submodule from optimum.quanto.tensor import Optimizer, qtype, qtypes -from torchao.quantization.quant_api import ( - quantize_ as torchao_quantize_, - _is_linear as torchao_is_linear, +from toolkit.quantization.torchao_compat import ( Float8WeightOnlyConfig, - IntxWeightOnlyConfig, - Int8WeightOnlyConfig + Int8WeightOnlyConfig, + intx_weight_only_config, + torchao_is_float8_tensor, + torchao_is_linear, + torchao_quantize_, ) from optimum.quanto import freeze from optimum.quanto.tensor.qbytes import QBytesTensor @@ -75,12 +75,12 @@ def _tensor_subclass_to_meta(value: torch.Tensor) -> torch.Tensor: ] torchao_qtypes = { - "uint2": IntxWeightOnlyConfig(torch.int2), - "uint3": IntxWeightOnlyConfig(torch.int3), - "uint4": IntxWeightOnlyConfig(torch.int4), - "uint5": IntxWeightOnlyConfig(torch.int5), - "uint6": IntxWeightOnlyConfig(torch.int6), - "uint7": IntxWeightOnlyConfig(torch.int7), + "uint2": intx_weight_only_config(2), + "uint3": intx_weight_only_config(3), + "uint4": intx_weight_only_config(4), + "uint5": intx_weight_only_config(5), + "uint6": intx_weight_only_config(6), + "uint7": intx_weight_only_config(7), "uint8": Int8WeightOnlyConfig(), "int8": Int8WeightOnlyConfig(), "float8": Float8WeightOnlyConfig(), @@ -210,7 +210,7 @@ def quantize( def filter_fn(module: torch.nn.Module, fqn: str) -> bool: if not torchao_is_linear(module, fqn): return False - if isinstance(module.weight, Float8Tensor): + if torchao_is_float8_tensor(module.weight): return False if include is not None and not any(fnmatch(fqn, pattern) for pattern in include): return False diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index e64e830cdb..3cc44f8448 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -366,6 +366,21 @@ export default function SimpleJob({ onChange={value => setJobConfig(value, 'config.process[0].model.layer_offloading_smart')} docKey="model.layer_offloading_smart" /> + {jobConfig.config.process[0].model.layer_offloading_smart && ( + + setJobConfig( + value, + 'config.process[0].model.layer_offloading_smart_cap_calibration', + ) + } + docKey="model.layer_offloading_smart_cap_calibration" + /> + )} {!jobConfig.config.process[0].model.layer_offloading_smart && ( Date: Thu, 16 Jul 2026 14:03:41 +0200 Subject: [PATCH 13/20] Remove retired Arena predecessor layout --- docs/ARENA_OFFLOAD_CONTRACT.md | 9 + tests/test_pin_manager.py | 5 +- .../arena_offload/construction.py | 2 - .../memory_management/arena_offload/layout.py | 587 +----------------- .../arena_offload/transfer.py | 33 - toolkit/memory_management/canonical_arena.py | 8 +- toolkit/memory_management/pin_manager.py | 21 +- toolkit/memory_management/residency.py | 15 +- toolkit/memory_management/transfer_plan.py | 9 +- toolkit/memory_management/vram_budget.py | 16 +- 10 files changed, 44 insertions(+), 661 deletions(-) diff --git a/docs/ARENA_OFFLOAD_CONTRACT.md b/docs/ARENA_OFFLOAD_CONTRACT.md index feb8dc66f4..4651d476cf 100644 --- a/docs/ARENA_OFFLOAD_CONTRACT.md +++ b/docs/ARENA_OFFLOAD_CONTRACT.md @@ -42,6 +42,11 @@ known boundary with the unmet contract in the error. while the source mapping remains intact. Once a managed source entry has been consumed, a later build failure aborts that model load instead of reusing the partial mapping. +- Dispatcher finalization occurs after the adapter or network is installed, so + the immutable execution description captures the model's final forwards and + trainable leaves exactly once. +- Residency/source publication and teardown reject changes while an execution + generation is active. - Whole-model `.cpu()` and current-arena `.cuda()` or `.to(device)` requests are interpreted by the runtime: permanent state follows the requested device and training residency is parked or restored. Whole-model dtype conversion, @@ -57,6 +62,10 @@ Arena compilation follows Toolkit's supported model `compile` setting and owns the one shared block dispatcher used by both training and sampling. Separate train/sample arena compile policies and caches are not part of this contract. +Transfer slots are reusable only after the compute stream's final reader has +been ordered past them. For training this includes checkpoint recomputation and +backward, not only the original forward call. + ## Maintainer validation matrix The upstream gate is the matrix, not an allowlist. Each selected production diff --git a/tests/test_pin_manager.py b/tests/test_pin_manager.py index b60bc54698..f2c3327bd0 100644 --- a/tests/test_pin_manager.py +++ b/tests/test_pin_manager.py @@ -136,11 +136,12 @@ class PinConformanceTests(unittest.TestCase): SCOPED_FILES = ( "toolkit/async_save.py", "toolkit/memory_management/bounce_pool.py", - "toolkit/memory_management/ingraph_stream.py", + "toolkit/memory_management/canonical_arena.py", "toolkit/memory_management/manager.py", "toolkit/memory_management/manager_modules.py", "toolkit/memory_management/checkpoint_autotuner.py", - "toolkit/memory_management/pinned_arena.py", + "toolkit/memory_management/arena_offload/construction.py", + "toolkit/memory_management/arena_offload/transfer.py", ) PATTERN = re.compile(r"pin_memory\s*=\s*True|\.pin_memory\(\)") diff --git a/toolkit/memory_management/arena_offload/construction.py b/toolkit/memory_management/arena_offload/construction.py index a492c79539..cc1056f90d 100644 --- a/toolkit/memory_management/arena_offload/construction.py +++ b/toolkit/memory_management/arena_offload/construction.py @@ -15,7 +15,6 @@ inspect_block, LayerStorageView, linear_views, - make_block_view_maker, substitution_views, typed_view, ) @@ -583,7 +582,6 @@ def commit(self): bool(block.handle and block.handle.pinned), pin_handle=block.handle, ) - pack.view_maker = make_block_view_maker(pack) block.pack = pack names = tuple(name for name, _ in block.entries) modules = tuple(module for _, module in block.entries) diff --git a/toolkit/memory_management/arena_offload/layout.py b/toolkit/memory_management/arena_offload/layout.py index 7e50fb602d..effa0ad480 100644 --- a/toolkit/memory_management/arena_offload/layout.py +++ b/toolkit/memory_management/arena_offload/layout.py @@ -11,7 +11,7 @@ import torch -from toolkit.quantization.storage import linear_storage_binding, module_storage_binding +from toolkit.quantization.storage import module_storage_binding from toolkit.memory_management import pin_manager LEAF_ALIGN = 256 @@ -48,15 +48,7 @@ class BlockPack: linears: tuple[LinearSpec, ...] required_pin_bytes: int pinned: bool - view_maker: object | None = None - # Ownership of ``host_flat``'s pin grant. ``pin_handle`` is the PinHandle - # returned by pin_manager.pin_alloc for this pack's OWN flat allocation - # (None when the pack didn't allocate -- e.g. it borrows another owner's - # storage). ``owns_flat`` gates release_pack: a borrowed pack (arena-backed, - # borrowed_from_arena=True) must never release someone else's handle. pin_handle: object | None = None - owns_flat: bool = True - borrowed_from_arena: bool = False @dataclass(frozen=True) @@ -72,584 +64,15 @@ class LayerStorageView: tensors: tuple[torch.Tensor, ...] -def layer_storage_views(pack: BlockPack) -> tuple[LayerStorageView, ...]: - """Expose a block's immutable execution declarations in leaf order.""" - return tuple( - LayerStorageView( - spec=spec, - tensors=tuple(typed_view(pack.host_flat, leaf) for leaf in spec.tensors), - ) - for spec in pack.linears - ) - - -def _rebuild_from_leaves(src, leaves_iter): - try: - names, ctx = src.__tensor_flatten__() - except Exception: - return next(leaves_iter) - moved = {} - for name in names: - inner = getattr(src, name, None) - moved[name] = None if inner is None else _rebuild_from_leaves(inner, leaves_iter) - return type(src).__tensor_unflatten__(moved, ctx, src.size(), src.stride()) - - -def _aligned_offsets(leaves: Iterable[torch.Tensor], align: int = LEAF_ALIGN): - offsets = [] - total = 0 - for leaf in leaves: - total = (total + align - 1) // align * align - offsets.append(total) - total += leaf.numel() * leaf.element_size() - return offsets, total - - -def _empty_host_flat( - nbytes: int, - *, - pin: bool = True, - kind: str = "ingraph_pack", - pin_mechanism: str = "alloc", -) -> tuple[torch.Tensor, bool, object | None]: - if not pin: - return torch.empty(nbytes, dtype=torch.uint8), False, None - if pin_mechanism == "register": - # I1: prepare the page-aligned buffer WITHOUT registering it yet -- - # pack_block_host copies the leaves into it (ordinary pageable - # memcpy) before the caller commits the cudaHostRegister pin. See - # pin_register_prepare's docstring for why population-before-pin is - # faster than registering a virgin buffer. - candidate, padded = pin_manager.pin_register_prepare(nbytes) - return candidate, False, ("register_pending", padded, kind) - else: - handle = pin_manager.pin_alloc( - nbytes, - kind, - required=False, - mode="sampling", - ) - return handle.tensor, bool(handle.pinned), handle - - def release_pack(pack: "BlockPack | None") -> None: - """Release a pack's own pin grant. - - A borrowed pack (``owns_flat=False``, e.g. arena-backed) must never - release someone else's handle -- the owner (the arena) is responsible for - its own flat's lifetime.""" - if pack is None or not pack.owns_flat: + """Release a canonical block pack's pin grant.""" + if pack is None: return pin_manager.release(pack.pin_handle) pack.pin_handle = None pack.pinned = False -def pack_block_host( - block_key: str, - linears, - *, - repoint: bool = True, - pin: bool = True, - kind: str = "ingraph_pack", - pin_mechanism: str = "alloc", -) -> BlockPack: - """Pack declared Linear storage tuples into one aligned host buffer.""" - normalized = [] - leaves = [] - for entry in linears: - if len(entry) == 2: - name, module = entry - weight = module.weight - bias = getattr(module, "bias", None) - else: - name, weight, bias = entry - module = None - binding = linear_storage_binding(weight, bias) - tensors = tuple(item.tensor for item in binding.tensors) - normalized.append((name, module, weight, bias, binding, tensors)) - leaves.extend(tensors) - - offsets, total = _aligned_offsets(leaves) - host, pinned, pin_handle = _empty_host_flat( - total, pin=pin, kind=kind, pin_mechanism=pin_mechanism - ) - register_pending = isinstance(pin_handle, tuple) and pin_handle[:1] == ( - "register_pending", - ) - try: - for leaf, offset in zip(leaves, offsets): - nbytes = leaf.numel() * leaf.element_size() - host[offset:offset + nbytes].view(leaf.dtype).reshape(leaf.shape).copy_(leaf) - - if register_pending: - _, _padded, register_kind = pin_handle - pin_handle = pin_manager.pin_register_commit( - host, total, register_kind, required=False - ) - host = pin_handle.tensor - pinned = bool(pin_handle.pinned) - - cursor = 0 - specs = [] - for name, module, weight, bias, binding, tensors in normalized: - tensor_specs = [] - views = [] - for declared, leaf in zip(binding.tensors, tensors): - offset = offsets[cursor] - nbytes = leaf.numel() * leaf.element_size() - tensor_specs.append( - LeafSpec( - offset=offset, - nbytes=nbytes, - dtype=leaf.dtype, - shape=tuple(leaf.shape), - role=declared.name, - ) - ) - views.append( - host[offset:offset + nbytes].view(leaf.dtype).reshape(leaf.shape) - ) - cursor += 1 - - if repoint and module is not None: - weight_views = views[:binding.weight_leaf_count] - weight_view = ( - weight_views[0] - if binding.weight_leaf_count == 1 - else _rebuild_from_leaves(binding.weight_template, iter(weight_views)) - ) - module.weight = torch.nn.Parameter( - weight_view, - requires_grad=getattr(weight, "requires_grad", False), - ) - if bias is not None: - module.bias = torch.nn.Parameter( - views[binding.weight_leaf_count], - requires_grad=getattr(bias, "requires_grad", False), - ) - - specs.append( - LinearSpec( - name=name, - tensors=tuple(tensor_specs), - execution_key=binding.execution_key, - weight_leaf_count=binding.weight_leaf_count, - weight_template=binding.weight_template, - weight_requires_grad=getattr(weight, "requires_grad", False), - bias_requires_grad=( - getattr(bias, "requires_grad", False) - if bias is not None - else False - ), - substitutions=binding.substitutions, - ) - ) - except Exception: - pin_manager.release(pin_handle) - raise - - pack = BlockPack( - block_key=block_key, - host_flat=host, - linears=tuple(specs), - required_pin_bytes=int(total), - pinned=bool(pinned), - pin_handle=pin_handle, - owns_flat=True, - borrowed_from_arena=False, - ) - pack.view_maker = make_block_view_maker(pack) - return pack - - -class ArenaBorrowError(ValueError): - """A block's live params don't actually live in the flat they were - expected to borrow from -- caller must fall back to an owned pack.""" - - -def pack_block_host_from_flat(block_key: str, linears, flat: torch.Tensor) -> "BlockPack": - """Describe declared storage tuples already resident in an owned flat.""" - flat_storage = flat.untyped_storage() - flat_ptr = flat.data_ptr() - flat_end = flat_ptr + flat.numel() * flat.element_size() - - def offset_of(leaf: torch.Tensor) -> int: - if leaf.untyped_storage().data_ptr() != flat_storage.data_ptr(): - raise ArenaBorrowError(f"arena_layout_mismatch:{block_key}:not_in_flat") - offset = leaf.data_ptr() - flat_ptr - nbytes = leaf.numel() * leaf.element_size() - if offset < 0 or offset + nbytes > flat_end - flat_ptr: - raise ArenaBorrowError(f"arena_layout_mismatch:{block_key}:out_of_range") - return offset - - specs = [] - for entry in linears: - if len(entry) == 2: - name, module = entry - weight = module.weight - bias = getattr(module, "bias", None) - else: - name, weight, bias = entry - binding = linear_storage_binding(weight, bias) - tensor_specs = [] - for declared in binding.tensors: - leaf = declared.tensor - tensor_specs.append( - LeafSpec( - offset=offset_of(leaf), - nbytes=leaf.numel() * leaf.element_size(), - dtype=leaf.dtype, - shape=tuple(leaf.shape), - role=declared.name, - ) - ) - specs.append( - LinearSpec( - name=name, - tensors=tuple(tensor_specs), - execution_key=binding.execution_key, - weight_leaf_count=binding.weight_leaf_count, - weight_template=binding.weight_template, - weight_requires_grad=getattr(weight, "requires_grad", False), - bias_requires_grad=( - getattr(bias, "requires_grad", False) - if bias is not None - else False - ), - substitutions=binding.substitutions, - ) - ) - - pack = BlockPack( - block_key=block_key, - host_flat=flat, - linears=tuple(specs), - required_pin_bytes=int(flat.numel() * flat.element_size()), - pinned=bool(pin_manager.is_host_pinned(flat)), - pin_handle=None, - owns_flat=False, - borrowed_from_arena=True, - ) - pack.view_maker = make_block_view_maker(pack) - return pack - - -class IngraphPackError(RuntimeError): - """A block pack could not be built or borrowed. ``reasons`` carries the - stable fail-closed tokens callers surface as ``_ingraph_unavailable_reasons`` - (``non_pinned_pack``, ``unsupported_quant_wrapper``, ``wrapper_pack_missing``, - ``arena_borrow_required``).""" - - def __init__(self, reasons, message: str = ""): - self.reasons = tuple(dict.fromkeys(reasons)) - super().__init__(message or ",".join(self.reasons)) - - -@dataclass -class PackBuildResult: - # Keyed by the model's STABLE block_key string, never a positional index -- - # the caller maps its own indices back locally. - packs: "dict[str, BlockPack]" - borrowed: int - owned: int - pageable: int # always 0 on success (a pageable pack raises non_pinned_pack) - reasons: tuple = () - - -def build_or_borrow_block_packs( - arena, - entries_by_block: dict, - *, - repoint: bool = False, - pin_mechanism: str = "register", - allow_owned_fallback: bool = True, -) -> PackBuildResult: - """Borrow each block's pack from the pinned arena, else build an owned one. - - The single place the in-graph pack policy lives, shared by every model's - ``enable_ingraph_sampling`` / ``enable_ingraph_training`` glue (see the - "in-graph arena protocol" in ``pinned_arena``). Nothing here knows about any - particular model: ``entries_by_block`` maps a stable ``block_key`` to that - block's ``(name, module)`` linear entries, and ``arena`` is duck-typed (any - object exposing ``try_borrow_pack(block_key, entries)``), so this module - never imports ``pinned_arena`` -- which imports it. - - Policy, centralized so callers cannot re-implement it inconsistently: - - * Borrow when the arena holds a current, pinned flat for the block: zero - alloc, zero copy, no second pin of the same bytes. - * Otherwise build an owned pack, but only if ``allow_owned_fallback``. - Under strict pinned-arena validation the caller passes False so a silent - fall back to owned packs cannot make a run "pass" without proving a - single borrow. - * Every streamed pack must be pinned; a pageable one fails the whole set - closed (``non_pinned_pack``) -- strict in-graph is all-or-nothing. - * On any failure, release ONLY packs we own. ``release_pack`` no-ops on a - borrowed pack (``owns_flat=False``), so the arena's flats are never freed - out from under it. - """ - packs: "dict[str, BlockPack]" = {} - borrowed = 0 - owned = 0 - try: - for block_key, raw_entries in entries_by_block.items(): - entries = list(raw_entries) - pack = arena.try_borrow_pack(block_key, entries) if arena is not None else None - if pack is not None: - borrowed += 1 - else: - if not allow_owned_fallback: - raise IngraphPackError( - ("arena_borrow_required",), - f"arena_borrow_required: block {block_key!r} is not " - "borrowable from the pinned arena", - ) - try: - pack = pack_block_host( - block_key, - entries, - repoint=repoint, - pin_mechanism=pin_mechanism, - ) - except ValueError as error: - message = str(error) - reason = ( - "wrapper_pack_missing" - if "wrapper packing" in message - else "unsupported_quant_wrapper" - ) - raise IngraphPackError( - (reason,), f"{reason} ({message})" - ) from error - owned += 1 - packs[block_key] = pack - if any(not pack.pinned for pack in packs.values()): - raise IngraphPackError(("non_pinned_pack",)) - except BaseException: - for pack in packs.values(): - release_pack(pack) - raise - return PackBuildResult(packs=packs, borrowed=borrowed, owned=owned, pageable=0) - - -def is_streamed_module(module) -> bool: - """The memory manager's marker for "this Linear's weights live on the host - and are fetched per call". - - Read it BEFORE stripping compile contaminants -- the strip deletes the - attribute, after which every leaf looks resident. - """ - return hasattr(module, "_layer_memory_manager") - - -def resident_linear_tensors(module) -> tuple: - """Return a Linear's declared opaque storage tuple in stable order.""" - binding = linear_storage_binding(module.weight, getattr(module, "bias", None)) - return tuple(item.tensor for item in binding.tensors) - - -@dataclass(frozen=True) -class BlockLeafPlan: - """Where each of a block's Linear leaves gets its weights this phase. - - The pack is a transfer-coalescing device (one H2D for N leaves), NOT a - residency decision. The memory planner splits residency per-Linear, so a - block is routinely part streamed / part resident. ``sources`` records, in - the caller's canonical leaf order, whether each leaf reads from the fetched - flat (``(True, i)`` -> ``streamed_views[i]``) or straight off its resident - Parameter (``(False, i)`` -> ``resident_args[i]``). Both are trace-time - constants, so the compiled block specializes on its residency pattern. - - ``pack is None`` means every leaf is resident: no flat, no fetch, no token. - """ - - block_key: str - pack: "BlockPack | None" - sources: tuple - resident_args: tuple - borrowed_from_arena: bool = False - - @property - def streams(self) -> bool: - return self.pack is not None - - -def assemble_leaf_args(plan: BlockLeafPlan, streamed_views: tuple = ()) -> tuple: - """Interleave fetched views and resident Parameters back into the block's - canonical leaf order. Pure Python over trace-time constants.""" - return tuple( - streamed_views[index] if from_pack else plan.resident_args[index] - for from_pack, index in plan.sources - ) - - -@dataclass -class BlockPlanResult: - plans: "dict[str, BlockLeafPlan]" - borrowed: int - owned: int - fully_resident: int - streamed_leaves: int - resident_leaves: int - reasons: tuple = () - - -def build_block_leaf_plans( - arena, - entries_by_block: dict, - *, - is_streamed=is_streamed_module, - repoint: bool = False, - pin_mechanism: str = "register", - allow_owned_fallback: bool = True, -) -> BlockPlanResult: - """Plan every block's leaves, packing only the ones the manager streams. - - ``entries_by_block`` maps a stable ``block_key`` to that block's FULL - ``(name, module)`` leaf list in canonical order. This splits each block by - ``is_streamed``, builds/borrows a pack over the streamed subset only, and - reads the resident leaves straight off their Parameters. - - Packing only the streamed subset is what lets the trunk coexist with the - planner's per-Linear residency: a partially-resident block yields a smaller - flat (so a smaller prefetch ring) and skips the fetch entirely for leaves - already on the device. Asking the arena for leaves it never offloaded is - what produced ``borrow refused: stale_modules=3/8``. - """ - streamed_by_block: dict = {} - for block_key, raw_entries in entries_by_block.items(): - streamed = [(name, module) for name, module in raw_entries if is_streamed(module)] - if streamed: - streamed_by_block[block_key] = streamed - - result = build_or_borrow_block_packs( - arena, - streamed_by_block, - repoint=repoint, - pin_mechanism=pin_mechanism, - allow_owned_fallback=allow_owned_fallback, - ) - try: - plans: "dict[str, BlockLeafPlan]" = {} - streamed_leaves = 0 - resident_leaves = 0 - for block_key, raw_entries in entries_by_block.items(): - pack = result.packs.get(block_key) - stream_index = { - name: index - for index, (name, _) in enumerate(streamed_by_block.get(block_key, ())) - } - sources = [] - resident_args = [] - for name, module in raw_entries: - index = stream_index.get(name) - if index is not None: - sources.append((True, index)) - streamed_leaves += 1 - continue - try: - tensors = resident_linear_tensors(module) - except ValueError as error: - raise IngraphPackError( - ("unsupported_quant_wrapper",), - f"unsupported_quant_wrapper ({block_key}.{name}: {error})", - ) from error - sources.append((False, len(resident_args))) - resident_args.append(tensors) - resident_leaves += 1 - plans[block_key] = BlockLeafPlan( - block_key=block_key, - pack=pack, - sources=tuple(sources), - resident_args=tuple(resident_args), - borrowed_from_arena=bool(pack is not None and pack.borrowed_from_arena), - ) - except BaseException: - for pack in result.packs.values(): - release_pack(pack) - raise - return BlockPlanResult( - plans=plans, - borrowed=result.borrowed, - owned=result.owned, - fully_resident=sum(1 for plan in plans.values() if not plan.streams), - streamed_leaves=streamed_leaves, - resident_leaves=resident_leaves, - ) - - -def _flat_view( - flat: torch.Tensor, - offset: int, - nbytes: int, - dtype: torch.dtype, - shape: tuple[int, ...], -) -> torch.Tensor: - return flat[offset:offset + nbytes].view(dtype).reshape(shape) - - -def _flat_clone_view( - flat: torch.Tensor, - offset: int, - nbytes: int, - dtype: torch.dtype, - shape: tuple[int, ...], -) -> torch.Tensor: - return flat[offset:offset + nbytes].clone().view(dtype).reshape(shape) - -def leaf_view(flat: torch.Tensor, spec: LeafSpec) -> torch.Tensor: - return _flat_view(flat, spec.offset, spec.nbytes, spec.dtype, spec.shape) - - -def block_storage_views( - flat: torch.Tensor, - pack: BlockPack, -) -> dict[str, LayerStorageView]: - """Return opaque ordered tensor views without binding execution.""" - out = {} - for spec in pack.linears: - tensors = tuple(leaf_view(flat, item) for item in spec.tensors) - out[spec.name] = LayerStorageView( - spec=spec, - tensors=tensors, - ) - return out - - -def make_block_view_maker(pack: BlockPack): - """Return a flat-buffer view maker that yields only tensor tuples.""" - entries = [] - for spec in pack.linears: - entries.append(tuple( - (item.offset, item.nbytes, item.dtype, item.shape) - for item in spec.tensors - )) - entries = tuple(entries) - - def view_maker(flat: torch.Tensor, _entries=entries): - out = [] - for tensor_entries in _entries: - tensors = [] - for index, item in enumerate(tensor_entries): - view = _flat_view(flat, item[0], item[1], item[2], item[3]) - tensors.append(view if index == 0 else view.clone()) - out.append(tuple(tensors)) - return tuple(out) - - return view_maker - - -def block_tensor_views(flat: torch.Tensor, pack: BlockPack) -> tuple: - maker = pack.view_maker - if maker is None: - maker = make_block_view_maker(pack) - pack.view_maker = maker - return maker(flat) - - - - @dataclass(frozen=True) class LeafDescriptor: role: str @@ -694,10 +117,6 @@ def flatten_leaves(value: torch.Tensor) -> list[torch.Tensor]: return leaves -# Compatibility for legacy manager/residency imports during extraction. -_flatten_leaves = flatten_leaves - - def rebuild_from_leaves(template: torch.Tensor, leaves: Iterable[torch.Tensor]): iterator = iter(leaves) diff --git a/toolkit/memory_management/arena_offload/transfer.py b/toolkit/memory_management/arena_offload/transfer.py index 2dfff9df78..dbaf076fc3 100644 --- a/toolkit/memory_management/arena_offload/transfer.py +++ b/toolkit/memory_management/arena_offload/transfer.py @@ -511,24 +511,6 @@ def _(host_flat, guard): return torch.empty(1, dtype=torch.int64, device="cpu") -@torch.library.custom_op("mm::fetch_start_gated", mutates_args=()) -def fetch_start_gated(host_flat: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - """fetch_start with a REAL data dependency on `gate`. - - fetch_start_after's guard is a declared mutation; that bookkeeping does - not survive Inductor's scheduler/reinplacer, which may hoist backward - re-fetches above frees (ring overrun). Here the gate is an ordinary - input, so the dependency is genuine dataflow no scheduling stage can - drop. Emitted by the post-grad ordering pass (ingraph_stream_scheduling); - not intended for hand-written model code.""" - return _fetch_start_impl(host_flat) - - -@fetch_start_gated.register_fake -def _(host_flat, gate): - return torch.empty(1, dtype=torch.int64, device="cpu") - - @torch.library.custom_op("mm::fetch_start_multi", mutates_args=()) def fetch_start_multi( host_flat: torch.Tensor, ranges: torch.Tensor, compact_nbytes: int @@ -556,21 +538,6 @@ def _(host_flat, ranges, compact_nbytes: int, guard): return torch.empty(1, dtype=torch.int64, device="cpu") -@torch.library.custom_op("mm::fetch_start_multi_gated", mutates_args=()) -def fetch_start_multi_gated( - host_flat: torch.Tensor, - ranges: torch.Tensor, - compact_nbytes: int, - gate: torch.Tensor, -) -> torch.Tensor: - return _fetch_start_multi_impl(host_flat, ranges, compact_nbytes) - - -@fetch_start_multi_gated.register_fake -def _(host_flat, ranges, compact_nbytes: int, gate): - return torch.empty(1, dtype=torch.int64, device="cpu") - - @torch.library.custom_op("mm::fetch_wait", mutates_args=()) def fetch_wait(token: torch.Tensor, nbytes: int) -> torch.Tensor: tid = int(token[0].item()) diff --git a/toolkit/memory_management/canonical_arena.py b/toolkit/memory_management/canonical_arena.py index 8fabbb70df..6870a7c987 100644 --- a/toolkit/memory_management/canonical_arena.py +++ b/toolkit/memory_management/canonical_arena.py @@ -2,11 +2,9 @@ One page-exclusive pinned host flat per block, immutable leaf metadata, and a ONE-TIME repoint of frozen base Parameters into views over those flats -(the canonical storage invariant). This is deliberately NOT the legacy -``pinned_arena.py``: -no generation counter, no ``is_current``/staleness oracle over live module -storage, no invalidate/restore/rebuild path, no borrowed-vs-owned pack -taxonomy. Promotion, demotion, and sampling +(the canonical storage invariant). There is no generation counter or +``is_current``/staleness oracle over live module storage, and no +invalidate/restore/rebuild path. Promotion, demotion, and sampling transitions must never repoint a Parameter again once ``canonicalize()`` has run -- that is the job of the residency sidecars, not this module. diff --git a/toolkit/memory_management/pin_manager.py b/toolkit/memory_management/pin_manager.py index ee09528572..4352427912 100644 --- a/toolkit/memory_management/pin_manager.py +++ b/toolkit/memory_management/pin_manager.py @@ -55,7 +55,7 @@ class PinHandle: # allocation strategy): they may reclaim the torch host cache during # reconcile, but must never shrink evictable higher-priority consumers # (the bounce pool) to make room for themselves. -_WEIGHT_TIER_KINDS = ("weights", "ingraph_pack") +_WEIGHT_TIER_KINDS = ("weights",) _SPILL_RESERVE_FLOOR_GIB_OVERRIDE: Optional[float] = None _SPILL_RESERVE_PCT_OVERRIDE: Optional[float] = None @@ -174,16 +174,11 @@ def dxgi_spill_reserve_bytes(budget_bytes: Optional[int] = None) -> int: def _spill_reserve_for_kind(kind: str, budget_bytes: Optional[int]) -> int: - # Weight-tier pins are one-shot STATIC commitments sized at attach/enable - # and fail-closed (strict mode raises rather than degrade): the per-tensor - # weight pins, the ingraph packs, and the pinned arena (kind="weights"). + # Weight-tier pins are one-shot static commitments sized at attach/enable: + # per-tensor legacy weight pins and canonical Arena storage (kind="weights"). # The pct-based reserve exists as slack for the *dynamic* streaming # consumer (the bounce pool), which grows at runtime -- a static commitment - # does not need it and keeps only the floor. The all-28 ingraph proof ran - # at 14.07/15.13 GiB committed, which the pct reserve would have refused; - # the pinned arena feeding the same all-streamed trunk (Phase 3 Slice B) is - # the identical commitment and must get the same floor, or a full-model - # training arena loses its last block to the pct reserve -> non_pinned_pack. + # does not need it and keeps only the floor. if kind in _WEIGHT_TIER_KINDS: return int(_spill_reserve_floor_gib() * GIB) return dxgi_spill_reserve_bytes(budget_bytes) @@ -365,9 +360,9 @@ def is_host_pinned(t: torch.Tensor) -> bool: caching host allocator; memory pinned in place with cudaHostRegister (``pin_tensor_in_place`` / ``pin_register`` -- the weight/arena tier) reports ``is_pinned() == False`` even though CUDA treats it as pinned for - transfer purposes. Consult the registration table too so consumers that - gate on "is this flat pinned" (e.g. borrowed ingraph packs) see registered - arena flats as pinned rather than falsely rejecting them as pageable. + transfer purposes. Consult the registration table too so canonical Arena + consumers recognize registered flats rather than falsely treating them as + pageable. """ if not isinstance(t, torch.Tensor): return False @@ -631,7 +626,7 @@ def pin_register( :func:`pin_register_commit` for callers with no data to populate before pinning (e.g. tests). Callers that populate a leaf-carrying flat should call the two steps directly with the copy in between (see - ``ingraph_stream.pack_block_host``). + canonical Arena construction). """ nbytes = int(nbytes) kind = str(kind or "unknown") diff --git a/toolkit/memory_management/residency.py b/toolkit/memory_management/residency.py index b88cdcde35..e57272aae3 100644 --- a/toolkit/memory_management/residency.py +++ b/toolkit/memory_management/residency.py @@ -16,10 +16,9 @@ from toolkit.memory_management import pin_manager from toolkit.memory_management.canonical_arena import CanonicalArena from toolkit.memory_management.arena_offload.layout import ( - - _flatten_leaves, - _rebuild_from_leaves, - leaf_view, + flatten_leaves, + rebuild_from_leaves, + typed_view, ) LeafKey = tuple[str, str] @@ -155,7 +154,7 @@ def weight(self): weight_tensors = self.tensors[:self.weight_leaf_count] if self.weight_leaf_count == 1: return weight_tensors[0] - return _rebuild_from_leaves(self.weight_template, iter(weight_tensors)) + return rebuild_from_leaves(self.weight_template, weight_tensors) @property def bias(self): @@ -174,13 +173,13 @@ class ResidencyDelta: def _tensor_bytes(tensor: torch.Tensor | None) -> int: if tensor is None: return 0 - return sum(leaf.numel() * leaf.element_size() for leaf in _flatten_leaves(tensor)) + return sum(leaf.numel() * leaf.element_size() for leaf in flatten_leaves(tensor)) def _record_stream(tensor: torch.Tensor | None, stream) -> None: if tensor is None: return - for leaf in _flatten_leaves(tensor): + for leaf in flatten_leaves(tensor): leaf.record_stream(stream) @@ -234,7 +233,7 @@ def _build_sidecar(self, key: LeafKey) -> ResidentLeaf: ) with torch.no_grad(), stream_context: tensors = tuple( - leaf_view(block.host_flat, item).to( + typed_view(block.host_flat, item).to( self.device, non_blocking=non_blocking ) for item in spec.tensors diff --git a/toolkit/memory_management/transfer_plan.py b/toolkit/memory_management/transfer_plan.py index 4d2674e9ed..8a10b0586b 100644 --- a/toolkit/memory_management/transfer_plan.py +++ b/toolkit/memory_management/transfer_plan.py @@ -5,10 +5,9 @@ flat (``canonical_arena.BlockRecord``) need to move to the device for one residency phase, coalesced and packed into a compact destination layout. Pure data model + coalescing algorithm live here (no CUDA needed to build -or inspect a plan); the runtime that actually submits the copies is -``mm::fetch_start_multi`` in ``ingraph_stream``, reusing its existing -ticket/ring/backpressure machinery unchanged (Slice 2's "single-ticket -semantics"). +or inspect a plan); the Arena transfer runtime submits the copies through +``mm::fetch_start_multi`` while preserving single-ticket ring/backpressure +semantics. Immutable and fingerprintable per Invariant 8: two plans built from the same block + the same set of streamed leaf names always produce identical @@ -46,7 +45,7 @@ class LeafRange: @dataclass(frozen=True) class CompactLeafSpec: """A streamed leaf's location within the plan's compact destination - buffer -- same shape as ``ingraph_stream.LeafSpec`` but the offset is + buffer -- the offset is in DESTINATION (compact device buffer) coordinates, not the arena flat's.""" diff --git a/toolkit/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py index e61d8684a6..fe1a16c5a7 100644 --- a/toolkit/memory_management/vram_budget.py +++ b/toolkit/memory_management/vram_budget.py @@ -589,8 +589,7 @@ def estimate_sampling_working_reserve_bytes( Used before any measured peak exists, so a high-resolution first sample plans enough streaming up front instead of discovering the working set - via mid-denoise OOM demotions (every demote invalidates compiled state - and, under strict ingraph, changes the pack set). + through mid-denoise OOM demotions. Linear-in-tokens model calibrated on Krea2 RTX 4070 smoke runs (2026-07-07/08, fp8 + cutlass attention, sequential CFG, partial @@ -612,10 +611,9 @@ def estimate_sampling_working_reserve_bytes( variation, and ``headroom_bytes`` (flat +1 GiB) deliberately overestimates: streaming one extra block costs a little bandwidth, while underestimating costs a - mid-denoise demote -- which invalidates compiled state, mutates the - strict-ingraph pack set, and (observed at 2000px) can cascade into a - full streamed transition. The learned per-run reserve replaces this - estimate after the first measured sample. + mid-denoise demote, which can cascade into a full streamed transition. The + learned per-run reserve replaces this estimate after the first measured + sample. """ tokens = max(0, int(image_tokens)) + max(0, int(text_tokens)) token_bytes = int(tokens * per_token_bytes * (2.5 if batch_cfg else 1.0)) @@ -646,9 +644,9 @@ def estimate_training_working_reserve_bytes( an allocator-cap OOM; production permits WDDM spill, but the resulting paging is still far slower than choosing the right cold layout. - Linear-in-tokens model calibrated on Krea2 LoKr RTX 4070 smoke runs - (2026-07-14, ``--block-stream-only`` so zero blocks are resident and - ``torch_max_allocated`` is purely the forward+backward+optimizer + Linear-in-tokens model calibrated on fully streamed Krea2 LoKr RTX 4070 + smoke runs (2026-07-14), where zero blocks were resident and + ``torch_max_allocated`` isolated the forward+backward+optimizer footprint, uncontaminated by the residency split this estimate feeds): 512x512 -> 1024 tokens, torch_max_allocated ~= 5.76 GiB From 285c4c0c61467e9ed76a8b4f0fade1d2f575362c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Thu, 16 Jul 2026 14:12:34 +0200 Subject: [PATCH 14/20] Remove unpublished Arena configuration and compile tuning --- jobs/process/BaseSDTrainProcess.py | 13 +---- tests/test_arena_offload_api.py | 15 +----- tests/test_compile_utils.py | 49 ------------------- toolkit/compile_utils.py | 17 ------- toolkit/config_modules.py | 9 +--- .../memory_management/arena_offload/api.py | 29 +---------- ui/src/types.ts | 1 - 7 files changed, 6 insertions(+), 127 deletions(-) delete mode 100644 tests/test_compile_utils.py diff --git a/jobs/process/BaseSDTrainProcess.py b/jobs/process/BaseSDTrainProcess.py index 7581366064..622a95d265 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -30,10 +30,7 @@ from toolkit.basic import value_map from toolkit.clip_vision_adapter import ClipVisionAdapter -from toolkit.compile_utils import ( - configure_cuda_only_inductor, - configure_quantized_compile_tuning, -) +from toolkit.compile_utils import configure_cuda_only_inductor from toolkit.custom_adapter import CustomAdapter from toolkit.data_loader import get_dataloader_from_datasets, trigger_dataloader_setup_epoch from toolkit.data_transfer_object.data_loader import FileItemDTO, DataLoaderBatchDTO @@ -1746,14 +1743,6 @@ def run(self): with model_load_arena_session(self.sd, enabled=arena_requested): self.sd.load_model() - coordinate_descent = configure_quantized_compile_tuning(self.model_config) - if coordinate_descent is not None: - state = "enabled" if coordinate_descent else "disabled" - print_acc( - "Quantized compile coordinate-descent tuning explicitly " - f"{state} by job config." - ) - text_encoders = getattr(self.sd, "text_encoder", None) if text_encoders is not None and not isinstance( text_encoders, (list, tuple) diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py index 4a8d4b2077..81e831c55b 100644 --- a/tests/test_arena_offload_api.py +++ b/tests/test_arena_offload_api.py @@ -87,7 +87,7 @@ class _FakeModelConfig: compile_sample = True train_compile_blocks = False layer_offloading_smart_working_reserve_gb = -1.0 - layer_offloading_smart_wddm_margin_gb = None + layer_offloading_smart_physical_vram_headroom_gb = None layer_offloading_smart_wddm_hard_gb = 1.0 layer_offloading_smart_cap_calibration = True layer_offloading_wddm_spill_reserve_pct = 0.10 @@ -95,7 +95,7 @@ class _FakeModelConfig: layer_offloading_checkpoint_keep_last = 2 layer_offloading_prefetch_depth = 3 layer_offloading_smart_sampling_working_reserve_gb = -1.0 - layer_offloading_smart_sampling_wddm_margin_gb = -1.0 + layer_offloading_smart_sampling_physical_vram_headroom_gb = -1.0 layer_offloading_smart_sampling_wddm_hard_gb = 1.0 layer_offloading_strict_vram_cap = False @@ -355,17 +355,6 @@ class DeadAliases: config = ArenaOffloadConfig.from_model_config(DeadAliases()) self.assertFalse(config.compile_blocks) - def test_compatibility_aliases_map_to_internal_policy(self): - class Aliases: - layer_offloading_smart_headroom_gb = 4.0 - layer_offloading_smart_buffer_gb = 1.5 - layer_offloading_smart_hard_buffer_gb = 0.75 - - policy = ArenaOffloadConfig.from_model_config(Aliases())._policy - self.assertEqual(policy.working_reserve_gib, 4.0) - self.assertEqual(policy.physical_vram_headroom_gib, 1.5) - self.assertEqual(policy.wddm_hard_gib, 0.75) - def test_backward_without_fp8_forward_is_ignored_once(self): class Invalid: quantize = True diff --git a/tests/test_compile_utils.py b/tests/test_compile_utils.py deleted file mode 100644 index da6cd6e950..0000000000 --- a/tests/test_compile_utils.py +++ /dev/null @@ -1,49 +0,0 @@ -from types import SimpleNamespace - -import torch - -from toolkit.compile_utils import configure_quantized_compile_tuning - - -def test_quantized_compile_tuning_preserves_torchao_default_when_unset(): - config = SimpleNamespace(compile=True, quantize=True) - original_tuning = torch._inductor.config.coordinate_descent_tuning - original_directions = ( - torch._inductor.config.coordinate_descent_check_all_directions - ) - try: - result = configure_quantized_compile_tuning(config) - assert result is None - assert torch._inductor.config.coordinate_descent_tuning == original_tuning - assert ( - torch._inductor.config.coordinate_descent_check_all_directions - == original_directions - ) - finally: - torch._inductor.config.coordinate_descent_tuning = original_tuning - torch._inductor.config.coordinate_descent_check_all_directions = ( - original_directions - ) - - -def test_quantized_compile_tuning_can_disable_torchao_search(): - config = SimpleNamespace( - compile=True, - quantize=True, - compile_coordinate_descent=False, - ) - original_tuning = torch._inductor.config.coordinate_descent_tuning - original_directions = ( - torch._inductor.config.coordinate_descent_check_all_directions - ) - try: - assert configure_quantized_compile_tuning(config) is False - assert torch._inductor.config.coordinate_descent_tuning is False - assert ( - torch._inductor.config.coordinate_descent_check_all_directions is False - ) - finally: - torch._inductor.config.coordinate_descent_tuning = original_tuning - torch._inductor.config.coordinate_descent_check_all_directions = ( - original_directions - ) diff --git a/toolkit/compile_utils.py b/toolkit/compile_utils.py index 37b1d1a9e3..e83d5cfa6d 100644 --- a/toolkit/compile_utils.py +++ b/toolkit/compile_utils.py @@ -12,20 +12,3 @@ def configure_cuda_only_inductor() -> None: from torch._inductor import config as inductor_config inductor_config.cpp.vec_isa_ok = False - - -def configure_quantized_compile_tuning(model_config) -> bool | None: - """Apply an explicit coordinate-descent policy after TorchAO quantization.""" - if not getattr(model_config, "compile", False): - return None - if not getattr(model_config, "quantize", False): - return None - - requested = getattr(model_config, "compile_coordinate_descent", None) - if requested is None: - return None - - enabled = bool(requested) - torch._inductor.config.coordinate_descent_tuning = enabled - torch._inductor.config.coordinate_descent_check_all_directions = enabled - return enabled diff --git a/toolkit/config_modules.py b/toolkit/config_modules.py index 8e4f7449e1..7cbe06c5f7 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -706,8 +706,7 @@ def __init__(self, **kwargs): "layer_offloading_smart_working_reserve_gb", -1.0 ) self.layer_offloading_smart_physical_vram_headroom_gb = kwargs.get( - "layer_offloading_smart_physical_vram_headroom_gb", - kwargs.get("layer_offloading_smart_wddm_margin_gb", -1.0), + "layer_offloading_smart_physical_vram_headroom_gb", -1.0 ) self.layer_offloading_smart_wddm_hard_gb = kwargs.get( "layer_offloading_smart_wddm_hard_gb", 1.0 @@ -719,8 +718,7 @@ def __init__(self, **kwargs): "layer_offloading_smart_sampling_working_reserve_gb", -1.0 ) self.layer_offloading_smart_sampling_physical_vram_headroom_gb = kwargs.get( - "layer_offloading_smart_sampling_physical_vram_headroom_gb", - kwargs.get("layer_offloading_smart_sampling_wddm_margin_gb", -1.0), + "layer_offloading_smart_sampling_physical_vram_headroom_gb", -1.0 ) self.layer_offloading_smart_sampling_wddm_hard_gb = kwargs.get( "layer_offloading_smart_sampling_wddm_hard_gb", 1.0 @@ -768,9 +766,6 @@ def __init__(self, **kwargs): self.compile_mode = kwargs.get("compile_mode", "default") self.compile_fullgraph = kwargs.get("compile_fullgraph", False) self.compile_dynamic = kwargs.get("compile_dynamic", True) - self.compile_coordinate_descent = kwargs.get( - "compile_coordinate_descent", None - ) self.cache_size_limit = kwargs.get("cache_size_limit", None) # kwargs to pass to the model diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index ef2f8d9a4c..404898f3d1 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -36,28 +36,6 @@ GIB = 1024**3 _FP8_QTYPES = ("qfloat8", "float8") -_COMPATIBILITY_ALIASES = { - "layer_offloading_smart_working_reserve_gb": ( - "layer_offloading_smart_headroom_gb", - ), - "layer_offloading_smart_physical_vram_headroom_gb": ( - "layer_offloading_smart_wddm_margin_gb", - "layer_offloading_smart_buffer_gb", - ), - "layer_offloading_smart_wddm_hard_gb": ( - "layer_offloading_smart_hard_buffer_gb", - ), - "layer_offloading_smart_sampling_working_reserve_gb": ( - "layer_offloading_smart_sampling_headroom_gb", - ), - "layer_offloading_smart_sampling_physical_vram_headroom_gb": ( - "layer_offloading_smart_sampling_wddm_margin_gb", - "layer_offloading_smart_sampling_buffer_gb", - ), - "layer_offloading_smart_sampling_wddm_hard_gb": ( - "layer_offloading_smart_sampling_hard_buffer_gb", - ), -} def estimate_training_working_reserve_hint_bytes( @@ -184,12 +162,7 @@ def from_model_config( cls, model_config, *, training_working_reserve_hint_bytes: int | None = None ) -> ArenaOffloadConfig: def get(name: str, default: Any = None) -> Any: - if hasattr(model_config, name): - return getattr(model_config, name) - for alias in _COMPATIBILITY_ALIASES.get(name, ()): - if hasattr(model_config, alias): - return getattr(model_config, alias) - return default + return getattr(model_config, name, default) raw_working_reserve_gib = get("layer_offloading_smart_working_reserve_gb") working_reserve_gib = raw_working_reserve_gib diff --git a/ui/src/types.ts b/ui/src/types.ts index 648d2206dd..9cb0fef7ac 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -179,7 +179,6 @@ export interface ModelConfig { compile_mode?: 'default' | 'max-autotune' | 'fastest'; compile_fullgraph?: boolean; compile_dynamic?: boolean; - compile_coordinate_descent?: boolean; cache_size_limit?: number; } From a3bcca426b0f19835f2a253baaca097165a6a658 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Thu, 16 Jul 2026 14:20:34 +0200 Subject: [PATCH 15/20] Pin tested TorchAO and preserve 0.10 compatibility --- docs/ARENA_OFFLOAD_CONTRACT.md | 14 ++++++++++++++ requirements_base.txt | 2 +- tests/test_torchao_compat.py | 16 ++++++++++++++++ 3 files changed, 31 insertions(+), 1 deletion(-) diff --git a/docs/ARENA_OFFLOAD_CONTRACT.md b/docs/ARENA_OFFLOAD_CONTRACT.md index 4651d476cf..919b987d28 100644 --- a/docs/ARENA_OFFLOAD_CONTRACT.md +++ b/docs/ARENA_OFFLOAD_CONTRACT.md @@ -36,6 +36,20 @@ known boundary with the unmet contract in the error. - Transfer and residency code treats declared leaves as opaque tensors. It does not branch on qtype or quantization backend identity. +## TorchAO dependency policy + +Fresh installations use the exact tested `torchao==0.17.0` release. This is the +minimum submitted TorchAO endpoint whose `Float8Tensor` representation is +supported by Arena's native FP8 forward and grad-input adapter. + +A code-only pull does not update an existing environment, so the compatibility +layer remains importable with upstream's previous `torchao==0.10.0` install. In +that stale environment, Toolkit configuration, the legacy offloader, dense +Arena storage, and compatible Quanto, OstrisLinear, and TorchAO operations stay +available. Only Arena's TorchAO-native FP8 execution is disabled, with a warning +that reports the installed version and the required `0.17.0` version. Versions +between `0.10.0` and `0.17.0` are not claimed as tested endpoints. + ## Loading and lifecycle contract - Direct checkpoint loading may fall back to ordinary `load_state_dict()` only diff --git a/requirements_base.txt b/requirements_base.txt index 9ff8fc4516..0c0472491d 100644 --- a/requirements_base.txt +++ b/requirements_base.txt @@ -1,4 +1,4 @@ -torchao>=0.10.0,<0.18.0 +torchao==0.17.0 safetensors git+https://github.com/huggingface/diffusers.git@c943837899b16cbae2f619b8dd4f7bb6f07dd81a #pip install git+https://github.com/huggingface/diffusers.git@refs/pull/13432/head diff --git a/tests/test_torchao_compat.py b/tests/test_torchao_compat.py index f0914316b4..47a5d70ac2 100644 --- a/tests/test_torchao_compat.py +++ b/tests/test_torchao_compat.py @@ -1,3 +1,6 @@ +import torch + +import toolkit.quantization.torchao_compat as torchao_compat from toolkit.quantization.torchao_compat import ( _release_tuple, intx_weight_only_config, @@ -25,3 +28,16 @@ def test_arena_fp8_requires_tested_version_and_tensor_format(): def test_current_intx_config_factory_is_available(): assert intx_weight_only_config(4) is not None + + +def test_torchao_010_uintx_config_factory_remains_supported(monkeypatch): + class LegacyUIntXConfig: + def __init__(self, dtype): + self.dtype = dtype + + monkeypatch.setattr(torchao_compat, "_IntxConfig", None) + monkeypatch.setattr(torchao_compat, "_UIntXConfig", LegacyUIntXConfig) + + config = intx_weight_only_config(4) + assert isinstance(config, LegacyUIntXConfig) + assert config.dtype is torch.uint4 From 5ad04abb8faebd2e7c210a6d4a67b6d63d60cd5f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Thu, 16 Jul 2026 14:23:26 +0200 Subject: [PATCH 16/20] Narrow Arena PR surface --- docs/ARENA_OFFLOAD_CONTRACT.md | 28 ++++++++ jobs/BaseJob.py | 4 +- tests/test_arena_canonical_transaction.py | 2 +- tests/test_arena_lifecycle_contract.py | 2 +- tests/test_arena_offload_api.py | 64 +++++++++++++++--- tests/test_canonical_arena.py | 4 +- tests/test_generic_block_dispatcher.py | 41 ++++++++++-- tests/test_lora_compile_scalars.py | 13 ---- tests/test_transfer_runtime.py | 2 +- .../arena_offload/__init__.py | 66 ++++++++++--------- toolkit/memory_management/residency.py | 6 +- toolkit/memory_management/transfer_plan.py | 20 +++--- toolkit/memory_management/vram_budget.py | 2 +- toolkit/timer.py | 2 - 14 files changed, 174 insertions(+), 82 deletions(-) diff --git a/docs/ARENA_OFFLOAD_CONTRACT.md b/docs/ARENA_OFFLOAD_CONTRACT.md index 919b987d28..352aa0be51 100644 --- a/docs/ARENA_OFFLOAD_CONTRACT.md +++ b/docs/ARENA_OFFLOAD_CONTRACT.md @@ -5,6 +5,34 @@ name. An explicitly selected model is accepted when its live module graph and storage satisfy the following contracts; otherwise setup fails at the narrowest known boundary with the unmet contract in the error. +## Configuration contract + +Arena is selected only when both `layer_offloading` and +`layer_offloading_smart` are true. Its public reserve controls are the canonical +`layer_offloading_smart_working_reserve_gb`, +`layer_offloading_smart_physical_vram_headroom_gb`, and +`layer_offloading_smart_wddm_hard_gb` names, plus their corresponding +`layer_offloading_smart_sampling_*` names. Negative reserve values request +automatic planning. + +The remaining public controls are +`layer_offloading_smart_cap_calibration`, `layer_offloading_strict_vram_cap`, +`layer_offloading_fp8_forward`, `layer_offloading_fp8_grad_input`, +`layer_offloading_fp8_sampling`, `layer_offloading_checkpoint_keep_last`, +`layer_offloading_prefetch_depth`, and the validation-only +`layer_offloading_simulated_vram_gb`. Development aliases are not accepted. + +## Integration surface + +The package facade exposes configuration and lifecycle entry points only: +`ArenaOffloadConfig`, setup and close operations, Arena runtime lookup and state +inspection, training-mode validation, direct canonical-loading operations, and +the model-load session. Generic memory-runtime lookup remains in +`toolkit.memory_management.runtime`. Discovery, dispatcher, layout, runtime +classes, implementation exceptions, and dispatcher constants are internal and +must be imported from their defining modules only by Arena implementation code +and focused tests. + ## Model contract - The transformer exposes one or more repeated `ModuleList` or `Sequential` diff --git a/jobs/BaseJob.py b/jobs/BaseJob.py index 016ed83318..ac99158e01 100644 --- a/jobs/BaseJob.py +++ b/jobs/BaseJob.py @@ -73,8 +73,8 @@ def cleanup(self): process.cleanup() except Exception as error: errors.append(f"{type(process).__name__}: {error}") - if errors: - raise RuntimeError("job cleanup failed: " + "; ".join(errors)) for process in processes: process.job = None self.process = [] + if errors: + raise RuntimeError("job cleanup failed: " + "; ".join(errors)) diff --git a/tests/test_arena_canonical_transaction.py b/tests/test_arena_canonical_transaction.py index 53252abffe..b3215127cc 100644 --- a/tests/test_arena_canonical_transaction.py +++ b/tests/test_arena_canonical_transaction.py @@ -212,7 +212,7 @@ def test_unsupported_layout_mid_stack_leaves_originals_untouched(self): layers = [frozen_linear(), frozen_linear()] originals = [layer.weight for layer in layers] arena = CanonicalArena() - from toolkit.memory_management.arena_offload import construction + import toolkit.memory_management.arena_offload.construction as construction real_inspect = construction.inspect_block calls = 0 diff --git a/tests/test_arena_lifecycle_contract.py b/tests/test_arena_lifecycle_contract.py index 9d6e562790..0b1570d187 100644 --- a/tests/test_arena_lifecycle_contract.py +++ b/tests/test_arena_lifecycle_contract.py @@ -8,12 +8,12 @@ from toolkit.memory_management.arena_offload import ( ArenaOffloadConfig, - ArenaSetupFatalError, prepare_canonical_storage, prepare_arena_offload, ) from toolkit.memory_management.arena_offload.errors import ( ArenaCleanupError, + ArenaSetupFatalError, is_fatal_arena_setup, recover_allows_next_job, ) diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py index 81e831c55b..3a2a2a6c47 100644 --- a/tests/test_arena_offload_api.py +++ b/tests/test_arena_offload_api.py @@ -1,35 +1,81 @@ -"""Facade-level coverage for `toolkit.memory_management.arena_offload`. - -These test the seam, not the machine: the helpers shared code now relies on -(`get_arena_runtime`, `is_memory_managed`, `memory_runtime_owns_compile`) and -the config mapping. Building a real arena needs CUDA and a real model; lifecycle -is covered by the arena contract tests and the Krea2 train smoke. -""" +"""Public Arena facade and generic memory-runtime seam coverage.""" import ast from dataclasses import fields from pathlib import Path +import subprocess +import sys from types import SimpleNamespace import unittest import torch +import toolkit.memory_management.arena_offload as arena_offload from toolkit.memory_management.arena_offload import ( ArenaOffloadConfig, estimate_training_working_reserve_hint_bytes, get_arena_runtime, is_arena_offloaded, - is_memory_managed, - memory_runtime_owns_compile, validate_arena_training_mode, ) from toolkit.memory_management.arena_offload.api import RUNTIME_ATTR, unwrap from toolkit.memory_management.arena_offload.runtime import _fixed_working_bytes from toolkit.memory_management.arena_offload.runtime import ArenaOffloadRuntime +from toolkit.memory_management.runtime import ( + is_memory_managed, + memory_runtime_owns_compile, +) GIB = 1024**3 +def test_canonical_arena_import_does_not_depend_on_facade_import_order(): + result = subprocess.run( + [ + sys.executable, + "-c", + ( + "import toolkit.memory_management.canonical_arena; " + "from toolkit.memory_management.arena_offload import " + "ArenaOffloadConfig; assert ArenaOffloadConfig" + ), + ], + cwd=Path(__file__).parents[1], + capture_output=True, + text=True, + check=False, + ) + assert result.returncode == 0, result.stderr + + +def test_public_facade_is_limited_to_integration_entry_points(): + assert arena_offload.__all__ == [ + "ArenaOffloadConfig", + "close_arena_offload", + "get_arena_runtime", + "estimate_training_working_reserve_hint_bytes", + "is_arena_offloaded", + "model_load_arena_session", + "prepare_canonical_storage", + "prepare_canonical_storage_from_state_dict", + "prepare_arena_offload", + "validate_arena_training_mode", + ] + for implementation_name in ( + "ArenaCleanupError", + "ArenaOffloadRuntime", + "ArenaSetupFatalError", + "BlockDiscoveryError", + "DISPATCHER_GENERATION", + "close_memory_runtime", + "discover_blocks", + "get_memory_runtime", + "is_memory_managed", + "memory_runtime_owns_compile", + ): + assert not hasattr(arena_offload, implementation_name) + + def test_arena_runtime_excludes_legacy_training_policy_calls(): source_path = ( Path(__file__).parents[1] / "jobs" / "process" / "BaseSDTrainProcess.py" diff --git a/tests/test_canonical_arena.py b/tests/test_canonical_arena.py index d634a41170..4dd4cbb9e7 100644 --- a/tests/test_canonical_arena.py +++ b/tests/test_canonical_arena.py @@ -85,8 +85,8 @@ def test_state_dict_round_trip_after_canonicalize(self): loaded = torch.load(buffer, weights_only=True) for key, value in expected.items(): self.assertTrue(torch.equal(value, loaded[key]), key) - # load_state_dict must copy IN PLACE, preserving the arena view - # (Parameter identity/storage unchanged -- Invariant 4). + # load_state_dict must copy in place so Parameter identity and the + # canonical arena storage view remain unchanged. flat_ptr = arena.block_pack("blocks.0").host_flat.untyped_storage().data_ptr() layer.load_state_dict(loaded) self.assertEqual(layer.weight.untyped_storage().data_ptr(), flat_ptr) diff --git a/tests/test_generic_block_dispatcher.py b/tests/test_generic_block_dispatcher.py index 065ee7cd8d..5494907b2d 100644 --- a/tests/test_generic_block_dispatcher.py +++ b/tests/test_generic_block_dispatcher.py @@ -9,10 +9,12 @@ from toolkit.memory_management.arena_offload import ( ArenaOffloadConfig, close_arena_offload, - discover_blocks, prepare_arena_offload, ) -from toolkit.memory_management.arena_offload.discovery import BlockDiscoveryError +from toolkit.memory_management.arena_offload.discovery import ( + BlockDiscoveryError, + discover_blocks, +) from toolkit.memory_management.arena_offload.dispatcher import ( _first_output_tensor, _first_tensor_argument, @@ -59,7 +61,7 @@ def _frozen_transformer(): return model -def _fp8_transformer(device, count=3, width=32): +def _fp8_transformer(device, count=3, width=32, *, qtype="float8"): from optimum.quanto import freeze from toolkit.util.quantize import get_qtype, quantize @@ -91,8 +93,10 @@ def forward(self, value): return value model = Fp8Transformer().to(device=device, dtype=torch.bfloat16) - quantize(model, weights=get_qtype("float8")) - freeze(model) + quantize(model, weights=get_qtype(qtype)) + if qtype in ("float8", "qfloat8"): + freeze(model) + model.requires_grad_(False) return model @@ -463,7 +467,7 @@ def installed_forward(self, value, _saved=saved): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_cuda_streamed_compiled_train_sample_train(): - from toolkit.memory_management.arena_offload import transfer + import toolkit.memory_management.arena_offload.transfer as transfer torch.manual_seed(23) device = torch.device("cuda") @@ -583,6 +587,31 @@ def train_once(step, *, input_requires_grad=True): assert active_process_owner() is None +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize( + ("qtype", "native_fp8"), + (("qfloat8", True), ("orbit4", False)), +) +def test_cuda_retained_quantization_backends_stream_through_arena( + qtype, native_fp8 +): + device = torch.device("cuda") + model = _fp8_transformer(device, qtype=qtype) + runtime = _fp8_runtime( + model, + device, + forward=native_fp8, + backward=native_fp8, + compile_blocks=False, + ) + try: + runtime.finalize() + output = _fp8_train_once(model, runtime, device, 1) + assert torch.isfinite(output).all() + finally: + close_arena_offload(model) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") def test_cuda_fp8_gates_select_distinct_canonical_arena_paths(): from toolkit.quantization import fp8_linear diff --git a/tests/test_lora_compile_scalars.py b/tests/test_lora_compile_scalars.py index 6673516d6c..aa261c1e74 100644 --- a/tests/test_lora_compile_scalars.py +++ b/tests/test_lora_compile_scalars.py @@ -27,19 +27,6 @@ def test_cuda_only_inductor_keeps_compile_errors_visible(): assert inductor_config.cpp.vec_isa_ok is False -def test_cuda_compile_prefers_aot_eager_then_compile(monkeypatch): - stances = [] - monkeypatch.setattr( - torch.compiler, - "set_stance", - lambda stance: stances.append(stance), - ) - - configure_cuda_only_inductor() - - assert stances == ["aot_eager_then_compile"] - - def test_lora_tensor_alpha_uses_device_owned_runtime_scale(): module = LoRAModule( "compile_scalar", diff --git a/tests/test_transfer_runtime.py b/tests/test_transfer_runtime.py index c6e717d612..7834b91f95 100644 --- a/tests/test_transfer_runtime.py +++ b/tests/test_transfer_runtime.py @@ -2,7 +2,7 @@ import torch import torch.nn.functional as F -from toolkit.memory_management.arena_offload import transfer as ingraph_stream +import toolkit.memory_management.arena_offload.transfer as ingraph_stream from toolkit.memory_management.canonical_arena import CanonicalArena from toolkit.memory_management.transfer_plan import build_transfer_plan diff --git a/toolkit/memory_management/arena_offload/__init__.py b/toolkit/memory_management/arena_offload/__init__.py index 433f3bed4f..cc728cff0e 100644 --- a/toolkit/memory_management/arena_offload/__init__.py +++ b/toolkit/memory_management/arena_offload/__init__.py @@ -17,45 +17,51 @@ the legacy manager remains a separate backend. """ -from .api import ( - ArenaOffloadConfig, - close_arena_offload, - estimate_training_working_reserve_hint_bytes, - get_arena_runtime, - is_arena_offloaded, - is_memory_managed, - memory_runtime_owns_compile, - prepare_canonical_storage, - prepare_canonical_storage_from_state_dict, - prepare_arena_offload, - validate_arena_training_mode, -) -from .runtime import ArenaOffloadRuntime -from .dispatcher import DISPATCHER_GENERATION -from .discovery import BlockDiscoveryError, discover_blocks -from .errors import ArenaCleanupError, ArenaSetupFatalError -from .load_session import model_load_arena_session -from ..runtime import close_memory_runtime, get_memory_runtime - __all__ = [ "ArenaOffloadConfig", - "ArenaCleanupError", - "ArenaOffloadRuntime", - "ArenaSetupFatalError", - "BlockDiscoveryError", - "DISPATCHER_GENERATION", "close_arena_offload", - "close_memory_runtime", "get_arena_runtime", - "get_memory_runtime", - "discover_blocks", "estimate_training_working_reserve_hint_bytes", "is_arena_offloaded", - "is_memory_managed", - "memory_runtime_owns_compile", "model_load_arena_session", "prepare_canonical_storage", "prepare_canonical_storage_from_state_dict", "prepare_arena_offload", "validate_arena_training_mode", ] + +_PUBLIC_ENTRY_POINTS = { + "ArenaOffloadConfig": (".api", "ArenaOffloadConfig"), + "close_arena_offload": (".api", "close_arena_offload"), + "get_arena_runtime": (".api", "get_arena_runtime"), + "estimate_training_working_reserve_hint_bytes": ( + ".api", + "estimate_training_working_reserve_hint_bytes", + ), + "is_arena_offloaded": (".api", "is_arena_offloaded"), + "model_load_arena_session": (".load_session", "model_load_arena_session"), + "prepare_canonical_storage": (".api", "prepare_canonical_storage"), + "prepare_canonical_storage_from_state_dict": ( + ".api", + "prepare_canonical_storage_from_state_dict", + ), + "prepare_arena_offload": (".api", "prepare_arena_offload"), + "validate_arena_training_mode": (".api", "validate_arena_training_mode"), +} + + +def __getattr__(name): + entry_point = _PUBLIC_ENTRY_POINTS.get(name) + if entry_point is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + from importlib import import_module + + module_name, attribute_name = entry_point + value = getattr(import_module(module_name, __name__), attribute_name) + globals()[name] = value + return value + + +def __dir__(): + return sorted(set(globals()) | set(__all__)) diff --git a/toolkit/memory_management/residency.py b/toolkit/memory_management/residency.py index e57272aae3..0420071b6e 100644 --- a/toolkit/memory_management/residency.py +++ b/toolkit/memory_management/residency.py @@ -115,10 +115,10 @@ def build(cls, phase: str, resident_leaf_keys) -> ResidencyPlan: def from_smart_plan( cls, arena: CanonicalArena, smart_plan: dict, *, phase: str ) -> ResidencyPlan: - """Adapt the existing planner's ``offload_ids`` decision to sidecars. + """Adapt the Arena planner's ``offload_ids`` decision to sidecars. - This is the Slice 3 planner seam: priority and capacity remain owned by - ``MemoryManager.smart_training_plan``; only the mutation target changes. + Priority and capacity remain planner-owned; this method only translates + the decision into immutable canonical leaf residency. """ offload_ids = set(smart_plan.get("offload_ids", ())) resident = [] diff --git a/toolkit/memory_management/transfer_plan.py b/toolkit/memory_management/transfer_plan.py index 8a10b0586b..07bd076f22 100644 --- a/toolkit/memory_management/transfer_plan.py +++ b/toolkit/memory_management/transfer_plan.py @@ -1,5 +1,4 @@ -"""Static multi-range transfer plans (Slice 2, -tasks/open/IMMUTABLE_TRANSFER_ARENA_PLAN.md). +"""Static multi-range transfer plans. A ``BlockTransferPlan`` says which byte ranges of a canonical block's host flat (``canonical_arena.BlockRecord``) need to move to the device for one @@ -9,10 +8,10 @@ ``mm::fetch_start_multi`` while preserving single-ticket ring/backpressure semantics. -Immutable and fingerprintable per Invariant 8: two plans built from the -same block + the same set of streamed leaf names always produce identical -ranges and the same fingerprint, independent of the actual tensor/storage -identity behind the block's host flat at build time. +Plans are immutable and fingerprintable: two plans built from the same block +and the same set of streamed leaf names always produce identical ranges and +the same fingerprint, independent of the tensor or storage identity behind the +block's host flat at build time. """ from __future__ import annotations @@ -81,8 +80,7 @@ def compact_leaf_view(self, device_flat: torch.Tensor, leaf_name: str, role: str def ranges_tensor(self) -> torch.Tensor: """(N, 3) int64 CPU tensor of [src_offset, dst_offset, nbytes] rows, - the exact argument shape the ``mm::fetch_start_multi`` custom op - expects (Slice 2's compile-visible extended fetch op).""" + matching the ``mm::fetch_start_multi`` custom op's argument shape.""" if not self.ranges: return torch.empty((0, 3), dtype=torch.int64) return torch.tensor( @@ -119,9 +117,9 @@ def build_transfer_plan( ) -> BlockTransferPlan: """Build a coalesced multi-range transfer plan for the STREAMED subset of ``block``'s leaves. ``streamed_leaf_names`` is the residency - decision (owned by the planner/controllers, Invariant 10) -- this - function only turns "which leaves stream this phase" into "which byte - ranges to copy and where." + decision owned by the planner and controllers; this function only turns + "which leaves stream this phase" into "which byte ranges to copy and + where." Two source items coalesce into one range when the gap between them is at most ``slack_bytes`` (default: one leaf's alignment padding, diff --git a/toolkit/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py index fe1a16c5a7..022672944a 100644 --- a/toolkit/memory_management/vram_budget.py +++ b/toolkit/memory_management/vram_budget.py @@ -711,7 +711,7 @@ def training_guard_pressure(dxgi: dict, physical: dict) -> dict: # --------------------------------------------------------------------------- -# Two-timescale residency control (see tasks/done/RESIDENCY_TWO_TIMESCALE_PLAN.md) +# Two-timescale residency control # # Allowance lives in *target-space* (0.95*cap - live); the allocator cap is set # in *cap-space*. The two differ by the gc_threshold factor: a cap raise of ``d`` diff --git a/toolkit/timer.py b/toolkit/timer.py index 09779dc8e5..e849ba5faa 100644 --- a/toolkit/timer.py +++ b/toolkit/timer.py @@ -48,8 +48,6 @@ def print(self): timing_dict = {} # sort by longest at top for timer_name, timings in sorted(self.timers.items(), key=lambda x: sum(x[1]), reverse=True): - if not timings: - continue avg_time = sum(timings) / len(timings) if not is_ui: From ed5545d7576bd37b09e15e41039e38352a3deb1a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Thu, 16 Jul 2026 15:33:21 +0200 Subject: [PATCH 17/20] Remove orphaned Arena support paths --- jobs/BaseJob.py | 12 +- tests/test_arena_lifecycle_contract.py | 2 +- tests/test_arena_offload_api.py | 41 +-- tests/test_arena_offload_policy.py | 163 ----------- tests/test_dxgi_meminfo.py | 35 --- tests/test_pin_manager.py | 57 +--- tests/test_residency.py | 28 -- .../memory_management/arena_offload/api.py | 3 - .../memory_management/arena_offload/fp8.py | 2 - .../memory_management/arena_offload/policy.py | 26 -- .../arena_offload/resources.py | 8 +- .../arena_offload/runtime.py | 94 +------ .../arena_offload/transfer.py | 173 +----------- toolkit/memory_management/dxgi_meminfo.py | 114 +------- .../memory_management/immutable_runtime.py | 139 --------- toolkit/memory_management/manager.py | 2 +- toolkit/memory_management/pin_manager.py | 215 +------------- toolkit/memory_management/runtime.py | 14 - toolkit/memory_management/vram_budget.py | 143 +--------- toolkit/quantization/fp8_linear.py | 266 +----------------- toolkit/util/quantize.py | 208 -------------- 21 files changed, 58 insertions(+), 1687 deletions(-) diff --git a/jobs/BaseJob.py b/jobs/BaseJob.py index ac99158e01..0693928b34 100644 --- a/jobs/BaseJob.py +++ b/jobs/BaseJob.py @@ -31,9 +31,9 @@ def get_conf(self, key, default=None, required=False): def run(self): print("") - print(f"#############################################") + print("#############################################") print(f"# Running job: {self.name}") - print(f"#############################################") + print("#############################################") print("") # implement in child class # be sure to call super().run() first @@ -67,14 +67,16 @@ def load_processes(self, process_dict: dict): def cleanup(self): errors = [] + failed = [] processes = list(getattr(self, "process", ())) for process in reversed(processes): try: process.cleanup() except Exception as error: errors.append(f"{type(process).__name__}: {error}") - for process in processes: - process.job = None - self.process = [] + failed.append(process) + else: + process.job = None + self.process = list(reversed(failed)) if errors: raise RuntimeError("job cleanup failed: " + "; ".join(errors)) diff --git a/tests/test_arena_lifecycle_contract.py b/tests/test_arena_lifecycle_contract.py index 0b1570d187..344f82064e 100644 --- a/tests/test_arena_lifecycle_contract.py +++ b/tests/test_arena_lifecycle_contract.py @@ -27,7 +27,7 @@ from toolkit.memory_management.arena_offload.runtime import ArenaOffloadRuntime from toolkit.memory_management import pin_manager from toolkit.memory_management import vram_budget -from toolkit.memory_management.arena_offload.fp8 import ( +from toolkit.quantization.fp8_linear import ( fp8_grad_input_enabled, set_fp8_grad_input_enabled, ) diff --git a/tests/test_arena_offload_api.py b/tests/test_arena_offload_api.py index 3a2a2a6c47..33c74343b2 100644 --- a/tests/test_arena_offload_api.py +++ b/tests/test_arena_offload_api.py @@ -1,6 +1,5 @@ """Public Arena facade and generic memory-runtime seam coverage.""" -import ast from dataclasses import fields from pathlib import Path import subprocess @@ -18,10 +17,11 @@ is_arena_offloaded, validate_arena_training_mode, ) -from toolkit.memory_management.arena_offload.api import RUNTIME_ATTR, unwrap +from toolkit.memory_management.arena_offload.api import unwrap from toolkit.memory_management.arena_offload.runtime import _fixed_working_bytes from toolkit.memory_management.arena_offload.runtime import ArenaOffloadRuntime from toolkit.memory_management.runtime import ( + RUNTIME_ATTR, is_memory_managed, memory_runtime_owns_compile, ) @@ -76,43 +76,6 @@ def test_public_facade_is_limited_to_integration_entry_points(): assert not hasattr(arena_offload, implementation_name) -def test_arena_runtime_excludes_legacy_training_policy_calls(): - source_path = ( - Path(__file__).parents[1] / "jobs" / "process" / "BaseSDTrainProcess.py" - ) - tree = ast.parse(source_path.read_text(encoding="utf-8")) - parents = {} - for parent in ast.walk(tree): - for child in ast.iter_child_nodes(parent): - parents[child] = parent - - guarded_calls = [] - for node in ast.walk(tree): - if not isinstance(node, ast.Call): - continue - func = node.func - name = func.attr if isinstance(func, ast.Attribute) else None - if name not in { - "prepare_training_memory_for_shape", - "auto_tune_training_memory", - }: - continue - ancestor = parents.get(node) - guarded = False - while ancestor is not None: - if isinstance(ancestor, ast.If): - condition = ast.unparse(ancestor.test) - if "arena_runtime is None" in condition: - guarded = True - break - ancestor = parents.get(ancestor) - guarded_calls.append((name, guarded)) - - # Upstream's legacy backend has no fork-local autotune calls. If those - # calls are added later, they must be explicitly excluded for arena runs. - assert all(guarded for _name, guarded in guarded_calls) - - class _Wrapper(torch.nn.Module): """Stands in for Accelerate/DDP, which expose the real model at `.module`.""" diff --git a/tests/test_arena_offload_policy.py b/tests/test_arena_offload_policy.py index 35d94d11b2..14266ba9a8 100644 --- a/tests/test_arena_offload_policy.py +++ b/tests/test_arena_offload_policy.py @@ -560,46 +560,6 @@ def test_controller_four_block_fast_lane_keeps_safety_vetoes(): assert controller.diagnostics()["last_aggressive_gate"] is False -def test_controller_does_not_bypass_bootstrap_verification(): - controller = ArenaResidencyController( - allocator_cache_headroom_bytes=10 - ) - controller.bootstrapped = True - controller.begin_bootstrap_promotion( - ("blocks.0", "blocks.1", "blocks.2", "blocks.3"), - 80, - 200, - 1000, - ) - clean = { - "allocator": { - "alloc_retries_delta": 0, - "free_count_delta": 0, - }, - "resident_bytes": 280, - "compile_invalid": False, - "transfer": None, - } - decision = controller.step( - clean, - candidate={"block_key": "blocks.4", "block_bytes": 20}, - demote_candidate=None, - cliff_cap_bytes=1000, - current_cap_bytes=1000, - worst_shape_free_bytes=100, - worst_shape_allocator_slack_bytes=200, - aggressive_promotion_capacity=4, - ) - assert decision.action == "hold" - assert controller.pending_promotion["block_keys"] == ( - "blocks.0", - "blocks.1", - "blocks.2", - "blocks.3", - ) - assert controller.diagnostics()["last_aggressive_gate"] is False - - def test_controller_raises_cap_by_fixed_fsm_increment(): controller = ArenaResidencyController( allocator_cache_headroom_bytes=10 @@ -789,129 +749,6 @@ def sampling(**kwargs): assert captured["hot_floor_bytes"] == int(1.25 * gib) + dequant -def test_bootstrap_uses_min_physical_free_and_one_gib_margin(): - gib = 1024 ** 3 - block_bytes = 200 * 1024 ** 2 - records = { - f"blocks.{index}": SimpleNamespace( - committed_bytes=block_bytes, - leaf_names=("linear",), - ) - for index in range(3) - } - plan = SimpleNamespace(resident_leaf_keys=frozenset()) - runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) - runtime._bootstrap_complete = False - runtime._bootstrap_min_free_bytes = 2 * gib + 450 * 1024 ** 2 - runtime._bootstrap_budget_bytes = 0 - runtime._bootstrap_block_keys = () - runtime._last_step_num = 50_000 - runtime._successful_training_steps = 1 - runtime._device = "cpu" - runtime._config = SimpleNamespace( - _policy=SimpleNamespace( - wddm_hard_gib=1.0, - physical_vram_headroom_gib=1.0, - ) - ) - runtime._model = SimpleNamespace() - runtime._arena = SimpleNamespace( - block_keys=lambda: tuple(records), - block_record=lambda key: records[key], - ) - runtime._residency = SimpleNamespace( - plan=plan, - resident_bytes=lambda: 0, - ) - runtime._training_plan = plan - runtime._smart_plan = {"singleton_resident_bytes": 100} - runtime._policy = ArenaResidencyController() - transitions = [] - runtime.transition_training_blocks = lambda keys, resident: ( - transitions.append((tuple(keys), resident)) - or { - "changed": True, - "block_keys": tuple(keys), - "plan": object(), - } - ) - - assert runtime._bootstrap_training_residency(10 * gib) is False - assert runtime._bootstrap_complete is False - - runtime._successful_training_steps = 2 - assert runtime._bootstrap_training_residency(10 * gib) is True - assert runtime._bootstrap_budget_bytes == 450 * 1024 ** 2 - assert transitions == [(("blocks.0", "blocks.1"), True)] - assert runtime._policy.pending_promotion["block_keys"] == ( - "blocks.0", - "blocks.1", - ) - assert runtime._policy.pending_promotion["resident_bytes_before"] == 100 - - -def test_bootstrap_ignores_first_runtime_sample_after_checkpoint_resume(): - runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) - runtime._bootstrap_complete = False - runtime._bootstrap_min_free_bytes = None - runtime._successful_training_steps = 1 - - runtime.record_training_physical_free_min(123) - assert runtime._bootstrap_min_free_bytes is None - - runtime._successful_training_steps = 2 - runtime.record_training_physical_free_min(456) - assert runtime._bootstrap_min_free_bytes == 456 - - -def test_bootstrap_keeps_priority_over_four_block_fast_lane(): - gib = 1024 ** 3 - block_bytes = 100 * 1024 ** 2 - records = { - f"blocks.{index}": SimpleNamespace( - committed_bytes=block_bytes, - leaf_names=("linear",), - ) - for index in range(5) - } - plan = SimpleNamespace(resident_leaf_keys=frozenset()) - runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) - runtime._bootstrap_complete = False - runtime._bootstrap_min_free_bytes = 3 * gib - runtime._bootstrap_budget_bytes = 0 - runtime._bootstrap_block_keys = () - runtime._successful_training_steps = 2 - runtime._device = "cpu" - runtime._config = SimpleNamespace( - _policy=SimpleNamespace( - wddm_hard_gib=1.0, - physical_vram_headroom_gib=1.0, - ) - ) - runtime._arena = SimpleNamespace( - block_keys=lambda: tuple(records), - block_record=lambda key: records[key], - ) - runtime._residency = SimpleNamespace( - plan=plan, - resident_bytes=lambda: 0, - ) - runtime._training_plan = plan - runtime._smart_plan = {"singleton_resident_bytes": 0} - runtime._policy = ArenaResidencyController() - transitions = [] - runtime.transition_training_blocks = lambda keys, resident: ( - transitions.append((tuple(keys), resident)) - or {"changed": True, "block_keys": tuple(keys), "plan": object()} - ) - - assert runtime._bootstrap_training_residency(10 * gib) is True - assert transitions == [ - (("blocks.0", "blocks.1", "blocks.2", "blocks.3", "blocks.4"), True) - ] - assert runtime._bootstrap_complete is True - - def test_arena_allocation_failure_drains_and_rolls_back(monkeypatch): runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) runtime._device = "cpu" diff --git a/tests/test_dxgi_meminfo.py b/tests/test_dxgi_meminfo.py index a9033f2cbc..4cbd82c945 100644 --- a/tests/test_dxgi_meminfo.py +++ b/tests/test_dxgi_meminfo.py @@ -114,13 +114,6 @@ def test_control_disable_keeps_legacy_proxy_for_control(self): class SpillReserveMarginTests(unittest.TestCase): - def setUp(self): - def _clear(): - pin_manager._SPILL_RESERVE_FLOOR_GIB_OVERRIDE = None - pin_manager._SPILL_RESERVE_PCT_OVERRIDE = None - _clear() - self.addCleanup(_clear) - def test_margin_is_pct_of_budget_when_pct_dominates(self): with mock.patch.dict( os.environ, @@ -154,33 +147,5 @@ def test_margin_falls_back_to_floor_without_budget(self): self.assertEqual(pin_manager.dxgi_spill_reserve_bytes(None), 2 * GIB) self.assertEqual(pin_manager.dxgi_spill_reserve_bytes(0), 2 * GIB) - def test_config_override_supersedes_env(self): - with mock.patch.dict( - os.environ, - {"AI_TOOLKIT_WDDM_SPILL_RESERVE_FLOOR_GIB": "2.0", - "AI_TOOLKIT_WDDM_SPILL_RESERVE_PCT": "0.20"}, - clear=False, - ): - pin_manager.set_spill_reserve_policy(floor_gib=3.0, pct=0.10) - # config floor 3 GiB dominates 0.10 * 16 = 1.6 GiB. - self.assertEqual(pin_manager.dxgi_spill_reserve_bytes(16 * GIB), 3 * GIB) - - -class SafeForControlTests(unittest.TestCase): - def test_confident_auto_detect_is_safe(self): - self.assertTrue(dxgi_meminfo.safe_for_control("single_nvidia")) - self.assertTrue(dxgi_meminfo.safe_for_control("luid")) - - def test_manual_override_is_safe_but_flagged_manual(self): - self.assertTrue(dxgi_meminfo.safe_for_control("env_override")) - self.assertTrue(dxgi_meminfo.is_manual_control("env_override")) - self.assertFalse(dxgi_meminfo.is_manual_control("single_nvidia")) - - def test_fallback_methods_are_not_safe(self): - for method in ("sole_hardware_adapter", "global_conservative", "", None): - self.assertFalse(dxgi_meminfo.safe_for_control(method)) - self.assertFalse(dxgi_meminfo.is_manual_control(method)) - - if __name__ == "__main__": unittest.main() diff --git a/tests/test_pin_manager.py b/tests/test_pin_manager.py index f2c3327bd0..4c59789581 100644 --- a/tests/test_pin_manager.py +++ b/tests/test_pin_manager.py @@ -1,3 +1,4 @@ +import os import re import unittest from pathlib import Path @@ -13,11 +14,9 @@ class PinManagerTests(unittest.TestCase): def setUp(self): pin_manager.reset_for_tests() - pin_manager.set_host_cache_reserve_bytes(None) def tearDown(self): pin_manager.reset_for_tests() - pin_manager.set_host_cache_reserve_bytes(None) def test_ledger_tracks_kinds(self): pin_manager.register_pinned_bytes(100, "weights") @@ -31,35 +30,16 @@ def test_ledger_tracks_kinds(self): ) def test_reserve_reduces_available_grant(self): - pin_manager.set_host_cache_reserve_bytes(2 * GIB) - with mock.patch.object(pin_manager, "pinned_bytes_headroom", return_value=8 * GIB): - self.assertEqual(pin_manager.available_for_pin(mode="training"), 6 * GIB) - self.assertTrue(pin_manager.can_pin(6 * GIB, mode="training")) - self.assertFalse(pin_manager.can_pin(7 * GIB, mode="training")) - - def test_plan_full_pin_disables_bounce(self): - pin_manager.set_host_cache_reserve_bytes(1 * GIB) - with mock.patch.object(pin_manager, "pinned_bytes_headroom", return_value=10 * GIB): - plan = pin_manager.plan_budgets( - offloaded_weight_bytes=8 * GIB, - requested_bounce_bytes=4 * GIB, - mode="training", - ) - self.assertEqual(plan["strategy"], "full_pin_no_bounce") - self.assertEqual(plan["weight_budget_bytes"], 8 * GIB) - self.assertEqual(plan["bounce_budget_bytes"], 0) - - def test_plan_partial_prioritizes_bounce_then_weights(self): - pin_manager.set_host_cache_reserve_bytes(1 * GIB) - with mock.patch.object(pin_manager, "pinned_bytes_headroom", return_value=10 * GIB): - plan = pin_manager.plan_budgets( - offloaded_weight_bytes=12 * GIB, - requested_bounce_bytes=4 * GIB, - mode="training", - ) - self.assertEqual(plan["strategy"], "partial_bounce_first") - self.assertEqual(plan["bounce_budget_bytes"], 4 * GIB) - self.assertEqual(plan["weight_budget_bytes"], 5 * GIB) + with mock.patch.dict( + os.environ, + {"AI_TOOLKIT_PIN_HOST_CACHE_RESERVE_GIB": "2.0"}, + ): + with mock.patch.object( + pin_manager, "pinned_bytes_headroom", return_value=8 * GIB + ): + self.assertEqual( + pin_manager.available_for_pin(mode="training"), 6 * GIB + ) def test_release_clamps_within_kind_only(self): @@ -69,17 +49,6 @@ def test_release_clamps_within_kind_only(self): pin_manager.release_pinned_bytes(50, "bounce") self.assertEqual(pin_manager.pinned_bytes_by_kind(), {"weights": 100}) - def test_weight_tier_reconcile_cannot_shrink_evictables(self): - # Priority guard: weights are the lowest pin tier, so a weights-tier - # shortfall may empty the host cache but must never evict the bounce - # pool to make room for itself. - calls = [] - pin_manager.register_evictable(lambda need: calls.append(need) or 0) - pin_manager.reconcile(1 * GIB, allow_shrink=False) - self.assertEqual(calls, []) - pin_manager.reconcile(1 * GIB, allow_shrink=True) - self.assertEqual(calls, [1 * GIB]) - def test_release_handle_is_idempotent(self): pin_manager.register_pinned_bytes(64, "save_stager") handle = pin_manager.PinHandle( @@ -127,7 +96,7 @@ def succeed(_tensor, kind): class PinConformanceTests(unittest.TestCase): - """No direct pinning outside the pin manager (PIN_MANAGER_PLAN S2). + """No direct pinning outside the pin manager. Every page-lock in the memory subsystem must route through pin_manager so the ledger stays authoritative. Scope is the offload subsystem + async @@ -135,11 +104,9 @@ class PinConformanceTests(unittest.TestCase): SCOPED_FILES = ( "toolkit/async_save.py", - "toolkit/memory_management/bounce_pool.py", "toolkit/memory_management/canonical_arena.py", "toolkit/memory_management/manager.py", "toolkit/memory_management/manager_modules.py", - "toolkit/memory_management/checkpoint_autotuner.py", "toolkit/memory_management/arena_offload/construction.py", "toolkit/memory_management/arena_offload/transfer.py", ) diff --git a/tests/test_residency.py b/tests/test_residency.py index 2d9fa21e15..c966460b6b 100644 --- a/tests/test_residency.py +++ b/tests/test_residency.py @@ -134,34 +134,6 @@ def refuse(*_args, **_kwargs): assert state.resident_bytes() > 0 -def test_runtime_training_transitions_are_whole_block(arena_layers): - arena, layers = arena_layers - block = SimpleNamespace(entries=tuple(layers.items())) - model = SimpleNamespace(blocks=(block,)) - state = ResidencyState(arena, "cpu") - state.reconcile(ResidencyPlan.build("train", ())) - runtime = ImmutableTransformerRuntime( - model, - state, - blocks=model.blocks, - block_keys=("blocks.0",), - entries_by_block={"blocks.0": tuple(layers.items())}, - compile_blocks=False, - ) - - growth = runtime.increase_training_residency( - arena.block_record("blocks.0").committed_bytes, - ) - expected = frozenset( - ("blocks.0", leaf_name) for leaf_name in layers - ) - assert state.plan.resident_leaf_keys == expected - assert growth["added_blocks"] == ("blocks.0",) - - relief = runtime.reduce_training_residency(1) - assert relief["removed_blocks"] == ("blocks.0",) - assert state.plan.resident_leaf_keys == frozenset() - def test_exact_training_block_transaction_uses_stable_key(arena_layers): arena, layers = arena_layers block = SimpleNamespace(entries=tuple(layers.items())) diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index 404898f3d1..928d9439ca 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -25,11 +25,8 @@ ) from ..runtime import ( - RUNTIME_ATTR, close_memory_runtime, get_memory_runtime, - is_memory_managed, - memory_runtime_owns_compile, unwrap_memory_model, ) from .runtime import ArenaOffloadRuntime diff --git a/toolkit/memory_management/arena_offload/fp8.py b/toolkit/memory_management/arena_offload/fp8.py index 0f605e702b..1ebd9ab228 100644 --- a/toolkit/memory_management/arena_offload/fp8.py +++ b/toolkit/memory_management/arena_offload/fp8.py @@ -7,8 +7,6 @@ from toolkit.quantization.fp8_linear import ( bind_parameter_operation, declare_fp8_linear, - fp8_grad_input_enabled, - set_fp8_grad_input_enabled, ) diff --git a/toolkit/memory_management/arena_offload/policy.py b/toolkit/memory_management/arena_offload/policy.py index e63ce20770..dfa1dadc09 100644 --- a/toolkit/memory_management/arena_offload/policy.py +++ b/toolkit/memory_management/arena_offload/policy.py @@ -176,14 +176,9 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, active_cap + max(0, int(cap_raise_bytes)), ) aggressive_capacity = max(0, int(aggressive_promotion_capacity or 0)) - bootstrap_pending = bool( - self.pending_promotion is not None - and self.pending_promotion.get("block_keys") - ) aggressive_ok = ( candidate is not None and aggressive_capacity >= AGGRESSIVE_PROMOTION_MIN_CAPACITY - and not bootstrap_pending and not bool(signal.get("compile_invalid")) and retries == 0 and device_frees == 0 @@ -321,27 +316,6 @@ def reject_pending_promotion(self, reason): ), ) - def begin_bootstrap_promotion( - self, block_keys, block_bytes, resident_bytes_before, active_cap_bytes - ): - keys = tuple(str(key) for key in block_keys) - self.last_promoted_key = keys[-1] if keys else None - self.pending_promotion = { - "block_key": self.last_promoted_key, - "block_keys": keys, - "block_bytes": int(block_bytes), - "resident_bytes_before": int(resident_bytes_before), - "resident_bytes_after": int(resident_bytes_before) + int(block_bytes), - "previous_cap_target_bytes": int(active_cap_bytes), - } - self.state = vram_budget.ResidencyFsmState( - vram_budget.FSM_PROMOTION_VERIFY, 0 - ) - self.last_action = "promote" - self.last_reason = "bootstrap_physical_free" - self.last_block_key = self.last_promoted_key - self.last_block_bytes = int(block_bytes) - def _begin_promotion( self, candidate, *, resident_bytes_before, active_cap_bytes ): diff --git a/toolkit/memory_management/arena_offload/resources.py b/toolkit/memory_management/arena_offload/resources.py index 7fee4fea95..2f9bad1036 100644 --- a/toolkit/memory_management/arena_offload/resources.py +++ b/toolkit/memory_management/arena_offload/resources.py @@ -45,7 +45,9 @@ def acquire_process_owner(self) -> None: self.owner_token = acquire_process_owner(self.device) try: from .. import allocator_cap, vram_budget - from .fp8 import fp8_grad_input_enabled + from toolkit.quantization.fp8_linear import ( + fp8_grad_input_enabled, + ) self._previous_fp8_grad_input = fp8_grad_input_enabled() self._previous_simulated_card = vram_budget.simulated_card_bytes() @@ -198,7 +200,9 @@ def attempt(label, operation, completed=None): self._transfer_runtime_released = True if self._restore_fp8_policy: - from .fp8 import set_fp8_grad_input_enabled + from toolkit.quantization.fp8_linear import ( + set_fp8_grad_input_enabled, + ) attempt( "FP8 grad-input policy", diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index 1799324e29..1b6f6d1afd 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -29,7 +29,6 @@ ) from ..vram_budget import apply_simulated_card from .policy import ( - AGGRESSIVE_PROMOTION_MIN_CAPACITY, ArenaResidencyController, TrainingSignalWindow, ) @@ -37,7 +36,7 @@ from .errors import ArenaCleanupError, ArenaSetupFatalError from .fp8 import disable as disable_fp8 from .fp8 import enable as enable_fp8 -from .fp8 import set_fp8_grad_input_enabled +from toolkit.quantization.fp8_linear import set_fp8_grad_input_enabled from .planner import ( build_training_plan, impossible_training_plan_message, @@ -49,8 +48,6 @@ RUNTIME_ATTR = "_arena_offload_runtime" GIB = 1024**3 -BOOTSTRAP_MARGIN_BYTES = GIB -BOOTSTRAP_MIN_STEP = 2 @dataclass @@ -91,8 +88,8 @@ def __init__( self._closed = False self._disposed = False - # Set by training_step(); the residency controller (git-bug 0c577ef) - # reads these at the step boundary. + # Set by training_step() and read by the residency controller at the + # step boundary. self._last_shape_key: tuple | None = None self._last_step_num: int | None = None self._successful_training_steps = 0 @@ -116,10 +113,6 @@ def __init__( self._cap_calibration_enabled = cap_calibration_enabled self._last_training_cap_target_bytes: int | None = None self._last_training_pressure_relief: dict | None = None - self._bootstrap_complete = False - self._bootstrap_min_free_bytes: int | None = None - self._bootstrap_budget_bytes = 0 - self._bootstrap_block_keys: tuple[str, ...] = () self._training_fp8_restores = [] self._training_fp8_canonical = 0 self._training_fp8_singletons = 0 @@ -189,7 +182,6 @@ def _prepare( ) set_fp8_grad_input_enabled(config.fp8_backward) - blocks = selection.blocks block_keys = selection.block_keys entries_by_block = { key: list(selection.entries_by_block[key]) @@ -1273,21 +1265,6 @@ def sampling_image( self._sampling_oom_retry_used = False # ------------------------------------------------------------------ - def record_training_physical_free_min(self, free_bytes) -> None: - """Publish one successful step's physical high-water for bootstrap.""" - if ( - self._bootstrap_complete - or free_bytes is None - or self._successful_training_steps < BOOTSTRAP_MIN_STEP - ): - return - value = max(0, int(free_bytes)) - self._bootstrap_min_free_bytes = ( - value - if self._bootstrap_min_free_bytes is None - else min(self._bootstrap_min_free_bytes, value) - ) - def _training_physical_vram_headroom_bytes(self) -> int: policy = self._config._policy hard_gib = float(getattr(policy, "wddm_hard_gib", None) or 1.0) @@ -1298,60 +1275,6 @@ def _training_physical_vram_headroom_bytes(self) -> int: ) return int(headroom_gib * GIB) - def _bootstrap_training_residency(self, active_cap_bytes) -> bool: - if ( - self._bootstrap_complete - or self._bootstrap_min_free_bytes is None - or int(self._successful_training_steps) < BOOTSTRAP_MIN_STEP - ): - return False - budget = max( - 0, - self._bootstrap_min_free_bytes - - self._training_physical_vram_headroom_bytes() - - BOOTSTRAP_MARGIN_BYTES, - ) - self._bootstrap_budget_bytes = budget - candidates = [] - protected = self._protected_training_blocks() - plan = getattr(self._residency, "plan", None) or self._training_plan - for order, block_key in enumerate(self._arena.block_keys()): - record = self._arena.block_record(block_key) - keys = tuple((block_key, name) for name in record.leaf_names) - if block_key in protected or any( - key in plan.resident_leaf_keys for key in keys - ): - continue - candidates.append( - (int(record.committed_bytes), order, str(block_key)) - ) - selected = [] - used = 0 - for block_bytes, _order, block_key in sorted(candidates): - if used + block_bytes > budget: - continue - selected.append(block_key) - used += block_bytes - if not selected: - self._bootstrap_complete = True - return False - resident_before = ( - int(self._residency.resident_bytes()) - + int((self._smart_plan or {}).get("singleton_resident_bytes", 0)) - ) - result = self.transition_training_blocks(selected, resident=True) - if not result.get("changed"): - return False - self._bootstrap_complete = True - self._bootstrap_block_keys = tuple(result["block_keys"]) - self._policy.begin_bootstrap_promotion( - self._bootstrap_block_keys, - used, - resident_before, - active_cap_bytes, - ) - return True - def _protected_training_blocks(self): return frozenset( str(block) @@ -1646,11 +1569,6 @@ def _apply_training_policy(self, *, shape_key=None): current_cap = int(cap_decision.target_cap_bytes) if cap_decision.hold_residency: return - if ( - not self._cap_calibrator.enabled - and self._bootstrap_training_residency(current_cap) - ): - return aggressive_capacity = self._aggressive_promotion_capacity(current_cap) decision = self._policy.step( self._signals.last_signal, @@ -1828,15 +1746,9 @@ def diagnostics(self) -> dict: "last_training_pressure_relief": getattr( self, "_last_training_pressure_relief", None ), - "bootstrap_complete": self._bootstrap_complete, - "bootstrap_min_free_bytes": self._bootstrap_min_free_bytes, - "bootstrap_margin_bytes": BOOTSTRAP_MARGIN_BYTES, "training_physical_vram_headroom_bytes": ( self._training_physical_vram_headroom_bytes() ), - "bootstrap_min_step": BOOTSTRAP_MIN_STEP, - "bootstrap_budget_bytes": self._bootstrap_budget_bytes, - "bootstrap_block_keys": self._bootstrap_block_keys, "working_reserve_bytes": int( (self._smart_plan or {}).get("working_reserve_bytes", 0) ), diff --git a/toolkit/memory_management/arena_offload/transfer.py b/toolkit/memory_management/arena_offload/transfer.py index dbaf076fc3..54ab32fd0f 100644 --- a/toolkit/memory_management/arena_offload/transfer.py +++ b/toolkit/memory_management/arena_offload/transfer.py @@ -7,8 +7,6 @@ from __future__ import annotations import collections -import contextlib -import itertools import threading import time from dataclasses import dataclass @@ -75,18 +73,9 @@ def __init__(self) -> None: # already completed, so accounting for a copy never blocks the host on it. _PENDING_H2D: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] -# Harness-only: restore the old behaviour of settling each copy's timing inside -# fetch_wait. Kept solely so a benchmark can A/B the cost of that host sync on -# the same build; production always drains lazily. -_BLOCKING_H2D_TIMING = False _RUNTIME_OWNER_TOKEN = None -def set_h2d_timing_blocking(enabled: bool) -> None: - global _BLOCKING_H2D_TIMING - _BLOCKING_H2D_TIMING = bool(enabled) - - def _drain_h2d(block: bool = False) -> None: """Accumulate h2d_ms for finished copies. Host-blocking only if block=True. @@ -112,25 +101,6 @@ def _drain_h2d(block: bool = False) -> None: _PENDING_H2D[:] = pending -def raise_dynamo_recompile_limit(min_limit: int = 128) -> None: - """Lift dynamo's per-code-object recompile cap for the in-graph trunks. - - Two legitimate recompile sources stack up on one code object: bucketed - training resolutions (dynamic=False -> one cache entry per distinct token - shape) and sampling-boundary rebuilds (fresh block-fn closures fail the - old entries' guards without evicting them). The default limit of 8 turned - that into FailOnRecompileLimitHit at the third boundary of a 200-step run - (~step 101). Each extra entry costs one ~3 min compile, not correctness; - the cap exists to flag accidental recompile storms, which the boundary - rebuild is not. - """ - config = torch._dynamo.config - for attribute in ("recompile_limit", "cache_size_limit"): - current = getattr(config, attribute, None) - if isinstance(current, int) and current < min_limit: - setattr(config, attribute, min_limit) - - def configure_fetch_runtime(*, depth: int = 3, owner_token=None) -> None: global _DEPTH, _NEXT_ID, _RUNTIME_OWNER_TOKEN if owner_token is not None: @@ -171,65 +141,6 @@ def lifetime_fetch_stats() -> dict: return dict(_LIFETIME_STATS) -def fetch_performance_metrics(stats: dict, *, step_wall_ms=None) -> dict: - """Derive transfer-stream utilization from a settled reporting window. - - ``h2d_ms`` is CUDA-event time on the single serialized transfer stream. - Dividing its window total by the matching step-wall total estimates transfer - duty. ``wait_ms`` is deliberately excluded: it is host blocking around an - event wait and does not say whether the GPU compute stream was idle. - - H2D timing drains opportunistically, so callers should provide a multi-step - reporting window. Duty above 100% is retained and flagged rather than - clamped; it indicates accounting carried across a window boundary or a - mismatched denominator. - """ - h2d_ms = float((stats or {}).get("h2d_ms", 0.0) or 0.0) - byte_count = int((stats or {}).get("bytes", 0) or 0) - wall_ms = None if step_wall_ms is None else float(step_wall_ms) - duty_pct = None - if wall_ms is not None and wall_ms > 0.0: - duty_pct = 100.0 * h2d_ms / wall_ms - achieved_gbps = None - if h2d_ms > 0.0: - achieved_gbps = byte_count / (h2d_ms * 1_000_000.0) - return { - "step_wall_ms": wall_ms, - "h2d_duty_pct": duty_pct, - "h2d_duty_overflow": bool(duty_pct is not None and duty_pct > 100.0), - "achieved_gbps": achieved_gbps, - } - - -def fetch_report(reset: bool = False, *, step_wall_ms=None) -> str | None: - stats = fetch_stats(reset=reset) - if not stats["fetches"]: - return None - metrics = fetch_performance_metrics(stats, step_wall_ms=step_wall_ms) - gib = stats["bytes"] / 1024 ** 3 - duty = ( - "-" if metrics["h2d_duty_pct"] is None - else f"{metrics['h2d_duty_pct']:.1f}" - ) - gbps = ( - "-" if metrics["achieved_gbps"] is None - else f"{metrics['achieved_gbps']:.2f}" - ) - wall = ( - "-" if metrics["step_wall_ms"] is None - else f"{metrics['step_wall_ms']:.3f}" - ) - return ( - f"[InGraphStream] fetches={int(stats['fetches'])} " - f"copies={int(stats['copies'])} " - f"bytes={gib:.2f} GiB h2d_ms={stats['h2d_ms']:.3f} " - f"step_wall_ms={wall} h2d_duty_pct={duty} " - f"h2d_duty_overflow={int(metrics['h2d_duty_overflow'])} " - f"achieved_gbps={gbps} " - f"wait_ms={stats['wait_ms']:.3f} depth_waits={int(stats['depth_waits'])}" - ) - - def _transfer_stream(device: torch.device): stream = _TRANSFER_STREAMS.get(device) if stream is None: @@ -561,7 +472,7 @@ def fetch_wait(token: torch.Tensor, nbytes: int) -> torch.Tensor: # and with the ring recycling device-side there is nothing to gain from # the throttle it used to provide. _PENDING_H2D.append((ticket.h2d_start, ticket.h2d_end)) - _drain_h2d(block=_BLOCKING_H2D_TIMING) + _drain_h2d(block=False) return ticket.device_buffer @@ -701,87 +612,5 @@ def free_on_backward(x: torch.Tensor, token: torch.Tensor) -> torch.Tensor: swaps it for the RECOMPUTE generation's token automatically, and this node's backward frees exactly the ticket backward actually read. - Canonical training block shape (see checkpoint_recompute_context): - - token = fetch_start_after(host, x) - flat = fetch_wait(token, nbytes) - ...views... - x = free_on_backward(x, token) - out = (x, views) - if not in_recompute(): - torch.ops.mm.fetch_free_after(token, out) # first-pass gen only - return out """ return _FreeOnBackwardFn.apply(x, token) - - -_IN_RECOMPUTE = threading.local() - - -def in_recompute() -> bool: - """True while a checkpoint recompute pass (via checkpoint_recompute_context) - is re-running the block fn.""" - return bool(getattr(_IN_RECOMPUTE, "value", False)) - - -class _RecomputeMarker: - def __enter__(self): - self._prev = getattr(_IN_RECOMPUTE, "value", False) - _IN_RECOMPUTE.value = True - return self - - def __exit__(self, exc_type, exc, tb): - _IN_RECOMPUTE.value = self._prev - return False - - -def checkpoint_recompute_context(): - """``context_fn`` for torch.utils.checkpoint: null forward context, and a - recompute context that flips in_recompute() so the block fn suppresses the - first-pass forward free during recompute (the recompute ticket is freed by - free_on_backward instead). EAGER ONLY -- compiled checkpoint requires - TorchDispatchMode contexts; use compiled_checkpoint_context there.""" - return contextlib.nullcontext(), _RecomputeMarker() - - -def _compiled_free_policy(ctx, op, *args, **kwargs): - from torch.utils.checkpoint import CheckpointPolicy - - if op in ( - torch.ops.mm.fetch_free_after.default, - torch.ops.mm.fetch_free.default, - ): - # Keep the forward-side free OUT of the backward replay: replayed, it - # would free the backward re-fetch's buffer before the grad kernels - # read it. free_on_backward's op is the backward-side free. - return CheckpointPolicy.MUST_SAVE - if op in ( - torch.ops.mm.fetch_start.default, - torch.ops.mm.fetch_start_after.default, - torch.ops.mm.fetch_start_multi.default, - torch.ops.mm.fetch_start_multi_after.default, - torch.ops.mm.fetch_wait.default, - ): - # The design's core invariant: fetched weights are NEVER saved for - # backward. PREFER_RECOMPUTE is advisory -- at Krea2 scale the - # partitioner chose to save all 28 fetched flats (12.25 GiB -> OOM). - return CheckpointPolicy.MUST_RECOMPUTE - # Everything else replays in backward (full-checkpoint mode). - return CheckpointPolicy.PREFER_RECOMPUTE - - -def compiled_checkpoint_context(): - """``context_fn`` for torch.utils.checkpoint under torch.compile.""" - from torch.utils.checkpoint import create_selective_checkpoint_contexts - - return create_selective_checkpoint_contexts(_compiled_free_policy) - - -# NOTE: there is deliberately NO helper that "picks the right checkpoint -# context automatically". The checkpoint HOP calls context_fn() OUTSIDE the -# compiling frame, so an is_compiling() check inside such a helper always -# reads False under compile and hands the HOP eager (non-TorchDispatchMode) -# contexts, failing its assertion. Select the context at trunk level instead: -# context_fn = (compiled_checkpoint_context -# if torch.compiler.is_compiling() -# else checkpoint_recompute_context) diff --git a/toolkit/memory_management/dxgi_meminfo.py b/toolkit/memory_management/dxgi_meminfo.py index 802fb14e4c..de7d90a76a 100644 --- a/toolkit/memory_management/dxgi_meminfo.py +++ b/toolkit/memory_management/dxgi_meminfo.py @@ -82,45 +82,6 @@ class DxgiMemoryInfo(NamedTuple): current_reservation_bytes: int -# Which adapter-match methods are trustworthy enough to *drive control* -# (aggressive pin-for-speed), versus only telemetry. Auto-detected confident -# matches (single_nvidia, and the Stage-B luid match) are safe. An explicit -# env override is honored for control but is a human-forced adapter, so it is -# logged distinctly as "manual" -- a misconfiguration must be visible, not -# silently trusted as if auto-verified. Everything else (sole_hardware_adapter, -# global_conservative, or no adapter at all) stays conservative. -_CONTROL_SAFE_METHODS = frozenset({"single_nvidia", "luid"}) -_CONTROL_MANUAL_METHODS = frozenset({"env_override"}) - - -def safe_for_control(match_method: Optional[str]) -> bool: - """Is a reading from this match method trustworthy enough to drive control? - - Pure classifier (unit-testable, no ctypes). True for confidently - auto-detected adapters and for an explicit manual override; False for - ambiguous/fallback selections and when the sensor is unavailable. - """ - if not match_method: - return False - return match_method in _CONTROL_SAFE_METHODS or match_method in _CONTROL_MANUAL_METHODS - - -def is_manual_control(match_method: Optional[str]) -> bool: - """True when control is driven off a human-forced (env-override) adapter.""" - return bool(match_method) and match_method in _CONTROL_MANUAL_METHODS - - -class DxgiAdapterInfo(NamedTuple): - index: int - description: str - vendor_id: int - device_id: int - luid: str - match_method: str - safe_for_control: bool - manual_control: bool - - class DxgiAdapterRecord(NamedTuple): index: int description: str @@ -136,7 +97,6 @@ class DxgiAdapterRecord(NamedTuple): class _AdapterSelection(NamedTuple): ptr: ctypes.c_void_p - info: DxgiAdapterInfo _selection_lock = threading.Lock() @@ -247,23 +207,6 @@ def _enum_adapters1(factory) -> list[tuple[ctypes.c_void_p, DxgiAdapterRecord]]: return adapters -def enumerate_adapters() -> list[DxgiAdapterRecord]: - """Return all DXGI adapters for diagnostics/probes, or [] if unavailable.""" - if os.name != "nt": - return [] - try: - factory = _create_dxgi_factory1() - if factory is None: - return [] - return [record for _adapter, record in _enum_adapters1(factory)] - except Exception as exc: - _log_once( - "enumerate_failed", - f"[DXGI] adapter enumeration unavailable: {exc}", - ) - return [] - - def _query_interface_adapter3(adapter) -> ctypes.c_void_p: adapter3 = ctypes.c_void_p() # IUnknown::QueryInterface slot 0. @@ -333,35 +276,22 @@ def _select_hardware_adapter3() -> Optional[_AdapterSelection]: ) _log_once( "ambiguous_adapter", - "[DXGI] adapter selection is ambiguous in Stage A; using legacy pin proxy", + "[DXGI] adapter selection is ambiguous; using legacy pin proxy", ) return None def _selection_from_record(adapter, record: DxgiAdapterRecord, match_method: str): adapter3 = _query_interface_adapter3(adapter) - manual = is_manual_control(match_method) - if manual: + if match_method == "env_override": # A human forced the adapter via AI_TOOLKIT_WDDM_DXGI_ADAPTER_INDEX. It - # still drives control, but log it distinctly so a misconfiguration is - # visible rather than lumped in with confident auto-detection. + # is worth logging distinctly so a misconfiguration is visible. _log_once( - "manual_control_override", + "manual_adapter_override", f"[DXGI] manual adapter override in effect (index={record.index} " - f"{record.description!r}); driving control as 'manual' -- verify " - "this is the training GPU", + f"{record.description!r}); verify this is the training GPU", ) - info = DxgiAdapterInfo( - index=record.index, - description=record.description, - vendor_id=record.vendor_id, - device_id=record.device_id, - luid=record.luid, - match_method=match_method, - safe_for_control=safe_for_control(match_method), - manual_control=manual, - ) - return _AdapterSelection(adapter3, info) + return _AdapterSelection(adapter3) def _selected_adapter() -> Optional[_AdapterSelection]: @@ -386,32 +316,13 @@ def _selected_adapter() -> Optional[_AdapterSelection]: return _selection -def selected_adapter_info() -> Optional[DxgiAdapterInfo]: - selection = _selected_adapter() - return selection.info if selection is not None else None - - -def control_is_eligible(cuda_device_index: int = 0) -> bool: - """True when the resolved adapter is trustworthy enough to drive control. - - Consumed by the pin-for-speed policy: aggressive full-pin engages only when - this is True (confident auto-detect or an explicit manual override). When - DXGI is unavailable/ambiguous this is False and the policy stays - conservative. ``cuda_device_index`` is accepted for the Stage-B per-device - match; Stage A resolves the single cached adapter. - """ - del cuda_device_index # Stage A: single cached adapter. - info = selected_adapter_info() - return bool(info is not None and info.safe_for_control) - - def query_video_memory_info( cuda_device_index: int = 0, segment_group: int = _DXGI_MEMORY_SEGMENT_GROUP_NON_LOCAL, min_interval_s: float = 0.5, ) -> Optional[DxgiMemoryInfo]: """Query LOCAL or NON_LOCAL DXGI video-memory info for the selected adapter.""" - del cuda_device_index # Stage A uses one cached adapter; Stage B matches per CUDA device. + del cuda_device_index # The current implementation selects one adapter. if os.name != "nt": return None selection = _selected_adapter() @@ -471,17 +382,6 @@ def query_non_local_video_memory_info( ) -def query_local_video_memory_info( - cuda_device_index: int = 0, - min_interval_s: float = 0.5, -) -> Optional[DxgiMemoryInfo]: - return query_video_memory_info( - cuda_device_index=cuda_device_index, - segment_group=_DXGI_MEMORY_SEGMENT_GROUP_LOCAL, - min_interval_s=min_interval_s, - ) - - def compute_non_local_headroom_bytes( budget_bytes: int, current_usage_bytes: int, diff --git a/toolkit/memory_management/immutable_runtime.py b/toolkit/memory_management/immutable_runtime.py index e26aef8404..c74cb49e9d 100644 --- a/toolkit/memory_management/immutable_runtime.py +++ b/toolkit/memory_management/immutable_runtime.py @@ -551,10 +551,6 @@ def program(self, mode: str) -> ImmutableProgram: except KeyError as error: raise ImmutableRuntimeError(f"unknown_execution_mode:{mode}") from error - def activate_sampling_fallback(self) -> ImmutableProgram: - self.set_residency_plan(self.sampling_fallback_plan) - return self.program(self.SAMPLE) - def transition_training_blocks(self, block_keys, *, resident: bool) -> dict: """Atomically add or remove complete training blocks in one plan.""" current = self._sources.plan or self.residency.plan @@ -651,141 +647,6 @@ def transition_training_block(self, block_key: str, *, resident: bool) -> dict: "plan": next_plan, } - def next_training_promotion_bytes(self) -> int: - current = self._sources.plan or self.residency.plan - if current.phase != self.TRAIN: - return 0 - protected = self.protected_training_leaf_keys - candidates = [] - for abi in self._block_abis: - keys = tuple((abi.block_key, leaf) for leaf in abi.leaf_names) - if any(key in current.resident_leaf_keys for key in keys): - continue - if any(key in protected for key in keys): - continue - record = self.residency.arena.block_record(abi.block_key) - candidates.append(int(record.committed_bytes)) - return min(candidates, default=0) - def increase_training_residency( - self, - available_growth_bytes: int, - *, - max_blocks: int = 1, - ) -> dict: - available = max(0, int(available_growth_bytes)) - current = self._sources.plan or self.residency.plan - if current.phase != self.TRAIN: - raise ImmutableRuntimeError( - f"training_residency_growth_requires_train:{current.phase}" - ) - - protected = self.protected_training_leaf_keys - candidates = [] - for order, abi in enumerate(self._block_abis): - keys = tuple((abi.block_key, leaf) for leaf in abi.leaf_names) - resident_count = sum( - key in current.resident_leaf_keys - for key in keys - ) - if resident_count: - # Initial and controller plans are whole-block. Preserve - # source-table support for partial plans without growing them. - continue - if any(key in protected for key in keys): - continue - record = self.residency.arena.block_record(abi.block_key) - candidates.append( - ( - int(record.committed_bytes), - order, - abi.block_key, - keys, - ) - ) - - candidates.sort(key=lambda item: (item[0], item[1])) - added = [] - predicted = 0 - limit = max(0, int(max_blocks)) - for nbytes, _order, _block_key, keys in candidates: - if limit and len(added) >= limit: - break - if predicted + nbytes > available: - continue - added.append((nbytes, keys)) - predicted += nbytes - - previous_plan = current - if added: - next_keys = set(current.resident_leaf_keys) - for _nbytes, keys in added: - next_keys.update(keys) - next_plan = ResidencyPlan.build(self.TRAIN, next_keys) - self.set_residency_plan(next_plan) - else: - next_plan = current - - added_keys = tuple( - sorted( - key - for _nbytes, keys in added - for key in keys - ) - ) - actual_growth = sum( - self.residency.resident_leaf_bytes(key) - for key in added_keys - ) - return { - "available_growth_bytes": available, - "predicted_growth_bytes": int(predicted), - "actual_growth_bytes": int(actual_growth), - "added_leaf_keys": added_keys, - "added_blocks": tuple(sorted({key[0] for key in added_keys})), - "previous_plan": previous_plan, - "plan": next_plan, - } - def reduce_training_residency(self, required_relief_bytes: int) -> dict: - requested = max(0, int(required_relief_bytes)) - current = self._sources.plan or self.residency.plan - if current.phase != self.TRAIN: - raise ImmutableRuntimeError(f"training_residency_reduction_requires_train:{current.phase}") - - protected = self.protected_training_leaf_keys - candidates = [] - for abi in self._block_abis: - keys = tuple((abi.block_key, leaf) for leaf in abi.leaf_names) - resident_keys = tuple(key for key in keys if key in current.resident_leaf_keys) - if not resident_keys or any(key in protected for key in keys): - continue - nbytes = sum(self.residency.resident_leaf_bytes(key) for key in resident_keys) - candidates.append((nbytes, abi.block_key, resident_keys)) - - candidates.sort(key=lambda item: (-item[0], item[1])) - removed = [] - relieved = 0 - for nbytes, _block_key, keys in candidates: - if relieved >= requested: - break - removed.extend(keys) - relieved += nbytes - - if removed: - next_keys = set(current.resident_leaf_keys) - set(removed) - next_plan = ResidencyPlan.build(self.TRAIN, next_keys) - self.set_residency_plan(next_plan) - else: - next_plan = current - - return { - "requested_relief_bytes": requested, - "relieved_bytes": int(relieved), - "removed_leaf_keys": tuple(sorted(removed)), - "removed_blocks": tuple(sorted({key[0] for key in removed})), - "remaining_resident_bytes": self.residency.resident_bytes(), - "plan": next_plan, - } - def full_model_resident_bytes(self) -> int: """Bytes to hold every canonical block resident (the 'want' figure).""" arena = self.residency.arena diff --git a/toolkit/memory_management/manager.py b/toolkit/memory_management/manager.py index b6017d29a4..fcc9bfa9e3 100644 --- a/toolkit/memory_management/manager.py +++ b/toolkit/memory_management/manager.py @@ -282,5 +282,5 @@ def detach(cls, module: torch.nn.Module): for key in keys_to_delete: del _DEVICE_STATE[key] - pin_manager.reconcile(allow_shrink=False) + pin_manager.reconcile() torch.cuda.empty_cache() diff --git a/toolkit/memory_management/pin_manager.py b/toolkit/memory_management/pin_manager.py index 4352427912..7c5f1b1ed8 100644 --- a/toolkit/memory_management/pin_manager.py +++ b/toolkit/memory_management/pin_manager.py @@ -10,7 +10,7 @@ import os import threading from dataclasses import dataclass -from typing import Callable, Optional +from typing import Optional import torch @@ -49,18 +49,11 @@ class PinHandle: _LOCK = threading.RLock() _LEDGER: dict[str, int] = {} -_EVICTABLES: list[Callable[[int], int]] = [] -# Weight-tier consumers are the LOWEST pin priority (PIN_MANAGER_PLAN -# allocation strategy): they may reclaim the torch host cache during -# reconcile, but must never shrink evictable higher-priority consumers -# (the bounce pool) to make room for themselves. +# Weight-tier consumers use only the fixed spill-reserve floor because their +# static commitments do not need the percentage reserve held for dynamic pins. _WEIGHT_TIER_KINDS = ("weights",) -_SPILL_RESERVE_FLOOR_GIB_OVERRIDE: Optional[float] = None -_SPILL_RESERVE_PCT_OVERRIDE: Optional[float] = None -_HOST_CACHE_RESERVE_BYTES_OVERRIDE: Optional[int] = None - dxgi_meminfo = None @@ -125,22 +118,9 @@ def release_pinned_bytes(n: int, kind: str = "unknown") -> None: def reset_for_tests() -> None: with _LOCK: _LEDGER.clear() - _EVICTABLES.clear() - - -def set_spill_reserve_policy( - floor_gib: Optional[float] = None, pct: Optional[float] = None -) -> None: - global _SPILL_RESERVE_FLOOR_GIB_OVERRIDE, _SPILL_RESERVE_PCT_OVERRIDE - if floor_gib is not None: - _SPILL_RESERVE_FLOOR_GIB_OVERRIDE = max(0.0, float(floor_gib)) - if pct is not None: - _SPILL_RESERVE_PCT_OVERRIDE = max(0.0, float(pct)) def _spill_reserve_floor_gib() -> float: - if _SPILL_RESERVE_FLOOR_GIB_OVERRIDE is not None: - return _SPILL_RESERVE_FLOOR_GIB_OVERRIDE for name, default in ( ("AI_TOOLKIT_WDDM_SPILL_RESERVE_FLOOR_GIB", None), ("AI_TOOLKIT_WDDM_SPILL_RESERVE_GIB", "1.0"), @@ -158,8 +138,6 @@ def _spill_reserve_floor_gib() -> float: def _spill_reserve_pct() -> float: - if _SPILL_RESERVE_PCT_OVERRIDE is not None: - return _SPILL_RESERVE_PCT_OVERRIDE try: return max(0.0, float(os.environ.get("AI_TOOLKIT_WDDM_SPILL_RESERVE_PCT", "0.10"))) except (TypeError, ValueError): @@ -239,14 +217,7 @@ def pinned_bytes_headroom( return max(0, int(total * fraction) - total_pinned_bytes()) -def set_host_cache_reserve_bytes(nbytes: Optional[int]) -> None: - global _HOST_CACHE_RESERVE_BYTES_OVERRIDE - _HOST_CACHE_RESERVE_BYTES_OVERRIDE = None if nbytes is None else max(0, int(nbytes)) - - def host_cache_reserve_bytes(mode: Optional[str] = None) -> int: - if _HOST_CACHE_RESERVE_BYTES_OVERRIDE is not None: - return _HOST_CACHE_RESERVE_BYTES_OVERRIDE if mode == "sampling": default = "0.0" else: @@ -282,40 +253,9 @@ def _empty_host_pin_cache() -> None: pass -def register_evictable(shrink: Callable[[int], int]) -> None: - with _LOCK: - if shrink not in _EVICTABLES: - _EVICTABLES.append(shrink) - - -def unregister_evictable(shrink: Callable[[int], int]) -> None: - with _LOCK: - if shrink in _EVICTABLES: - _EVICTABLES.remove(shrink) - - -def reconcile(required_bytes: int = 0, *, device=None, allow_shrink: bool = True) -> int: - """Escalation before failing a pin request: empty the torch host-pin cache, - then (for non-weight-tier requests) ask evictable consumers to shrink. - - ``allow_shrink=False`` is the priority guard: weight-tier requests may not - evict the bounce pool -- eviction runs in reverse priority order, and - weights are already the lowest tier.""" +def reconcile() -> None: + """Return cached host-pin allocations before retrying a pin request.""" _empty_host_pin_cache() - if not allow_shrink: - return 0 - freed = 0 - need = max(0, int(required_bytes or 0)) - with _LOCK: - evictables = list(_EVICTABLES) - for shrink in evictables: - try: - freed += max(0, int(shrink(max(0, need - freed)))) - except Exception: - pass - if need and freed >= need: - break - return freed _REGISTERED_HOST_PIN_LOCK = threading.Lock() @@ -331,11 +271,7 @@ def pin_tensor_in_place(t: torch.Tensor, kind: str = "weights", *, device=None) return False available = available_for_pin(kind=kind, nbytes=size, device=device) if available is not None and size > available: - reconcile( - size - available, - device=device, - allow_shrink=kind not in _WEIGHT_TIER_KINDS, - ) + reconcile() available = available_for_pin(kind=kind, nbytes=size, device=device) if available is not None and size > available: return False @@ -358,7 +294,8 @@ def is_host_pinned(t: torch.Tensor) -> bool: torch's ``is_pinned()`` only recognizes buffers allocated by its own caching host allocator; memory pinned in place with cudaHostRegister - (``pin_tensor_in_place`` / ``pin_register`` -- the weight/arena tier) + (``pin_tensor_in_place`` / register-mechanism pinning -- the + weight/arena tier) reports ``is_pinned() == False`` even though CUDA treats it as pinned for transfer purposes. Consult the registration table too so canonical Arena consumers recognize registered flats rather than falsely treating them as @@ -490,12 +427,11 @@ def pin_alloc( tensor = torch.empty(nbytes, dtype=torch.uint8) return PinHandle(tensor=tensor, nbytes=nbytes, kind=kind, pinned=False) - allow_shrink = kind not in _WEIGHT_TIER_KINDS available = available_for_pin( kind=kind, nbytes=nbytes, device=device, reserve_bytes=reserve_bytes, mode=mode ) if available is not None and nbytes > available: - reconcile(nbytes - available, device=device, allow_shrink=allow_shrink) + reconcile() available = available_for_pin( kind=kind, nbytes=nbytes, device=device, reserve_bytes=reserve_bytes, mode=mode ) @@ -512,7 +448,7 @@ def pin_alloc( try: tensor = torch.empty(nbytes, dtype=torch.uint8, pin_memory=True) except RuntimeError: - reconcile(nbytes, device=device, allow_shrink=allow_shrink) + reconcile() try: tensor = torch.empty(nbytes, dtype=torch.uint8, pin_memory=True) except RuntimeError as error: @@ -534,16 +470,16 @@ def pin_register_prepare(nbytes: int) -> tuple[torch.Tensor, int]: """Allocate the page-aligned pageable buffer a register-mechanism pin will need, WITHOUT pinning it yet. - Split out of ``pin_register`` so callers who need to populate the buffer - (e.g. copying leaf tensors into a block flat) can do so on ordinary + Split into prepare/commit steps so callers can populate the buffer (for + example, copying leaf tensors into a block flat) on ordinary pageable memory -- a plain memcpy that faults in pages at normal RAM bandwidth -- before ``cudaHostRegister`` runs. Registering a virgin, never-touched buffer forces the OS to commit+pin every page during the syscall itself, which is measurably slower than registering pages that - are already resident (I1, ~1-1.5s per full arena build). + are already resident, saving roughly 1-1.5 seconds per full arena build. Returns ``(candidate, padded_nbytes)``; ``candidate`` is an untouched - pageable view, exactly the layout ``pin_register`` used to build inline. + pageable view with the layout required by ``pin_register_commit``. """ nbytes = int(nbytes) if nbytes <= 0: @@ -605,111 +541,6 @@ def pin_register_commit( mechanism="register") -def pin_register( - nbytes: int, - kind: str, - *, - device=None, - required: bool = False, -) -> PinHandle: - """Exact-size host buffer pinned with cudaHostRegister. - - Unlike pin_alloc, this never touches torch's caching host allocator, so - the DXGI shared-budget cost is exactly ``nbytes`` (the caching allocator - rounds up to power-of-two buckets: observed live, 8.86 GiB of pin_alloc - flats committed 12.70 GiB of DXGI usage -- ~40% invisible overhead) and - release returns the budget immediately. Intended for large long-lived - buffers (the pinned weight arena); small/churny consumers should keep - using pin_alloc. - - Convenience wrapper over :func:`pin_register_prepare` + - :func:`pin_register_commit` for callers with no data to populate before - pinning (e.g. tests). Callers that populate a leaf-carrying flat should - call the two steps directly with the copy in between (see - canonical Arena construction). - """ - nbytes = int(nbytes) - kind = str(kind or "unknown") - if nbytes <= 0 or not torch.cuda.is_available(): - tensor = torch.empty(max(0, nbytes), dtype=torch.uint8) - return PinHandle(tensor=tensor, nbytes=max(0, nbytes) if nbytes > 0 else 0, - kind=kind, pinned=False, mechanism="register") - candidate, _padded = pin_register_prepare(nbytes) - return pin_register_commit(candidate, nbytes, kind, device=device, required=required) - - -def pin_empty(shape, dtype, kind: str, *, device=None, required: bool = False): - element_size = torch.empty((), dtype=dtype).element_size() - n = 1 - for dim in tuple(shape): - n *= int(dim) - handle = pin_alloc(n * element_size, kind, device=device, required=required) - return handle.tensor.view(dtype).reshape(tuple(shape)), handle.pinned - - -def can_pin( - nbytes: int, - *, - kind: str = "unknown", - device=None, - reserve_bytes: int = 0, - mode: Optional[str] = None, -) -> bool: - available = available_for_pin( - kind=kind, - nbytes=nbytes, - device=device, - reserve_bytes=reserve_bytes, - mode=mode, - ) - return available is None or int(nbytes) <= available - - - -def plan_budgets( - *, - offloaded_weight_bytes: int, - requested_bounce_bytes: int, - device=None, - mode: str = "training", -) -> dict: - """Return pin budgets using the fixed priority policy from PIN_MANAGER_PLAN.""" - weights = max(0, int(offloaded_weight_bytes or 0)) - requested_bounce = max(0, int(requested_bounce_bytes or 0)) - reserve = host_cache_reserve_bytes(mode) - headroom = pinned_bytes_headroom(_cuda_device_index(device)) - if headroom is None: - # No authoritative probe: keep the old requested shape, but still report - # the reserve so diagnostics show the implicit consumer exists. - return { - "mode": mode, - "strategy": "fallback_no_probe", - "headroom_bytes": None, - "reserve_bytes": reserve, - "bounce_budget_bytes": requested_bounce, - "weight_budget_bytes": weights, - } - usable = max(0, int(headroom) - reserve) - if weights and weights <= usable: - return { - "mode": mode, - "strategy": "full_pin_no_bounce", - "headroom_bytes": int(headroom), - "reserve_bytes": reserve, - "bounce_budget_bytes": 0, - "weight_budget_bytes": weights, - } - bounce_budget = min(requested_bounce, usable) - weight_budget = max(0, usable - bounce_budget) - return { - "mode": mode, - "strategy": "partial_bounce_first", - "headroom_bytes": int(headroom), - "reserve_bytes": reserve, - "bounce_budget_bytes": bounce_budget, - "weight_budget_bytes": min(weights, weight_budget), - } - def snapshot(device=None, mode: Optional[str] = None) -> dict: headroom = pinned_bytes_headroom(_cuda_device_index(device)) reserve = host_cache_reserve_bytes(mode) @@ -724,24 +555,6 @@ def snapshot(device=None, mode: Optional[str] = None) -> dict: } -def format_snapshot(device=None, mode: Optional[str] = None) -> str: - snap = snapshot(device=device, mode=mode) - by_kind = " ".join( - f"{key}={value / GIB:.2f}GiB" - for key, value in sorted(snap["pinned_by_kind"].items()) - ) - headroom = snap["headroom_bytes"] - avail = snap["available_bytes"] - return ( - "pin ledger: " - f"total={snap['pinned_total_bytes'] / GIB:.2f}GiB " - f"headroom={'n/a' if headroom is None else f'{headroom / GIB:.2f}GiB'} " - f"reserve={snap['host_cache_reserve_bytes'] / GIB:.2f}GiB " - f"free={'n/a' if avail is None else f'{avail / GIB:.2f}GiB'} " - f"{by_kind}" - ).strip() - - def _budget_message(kind: str, nbytes: int, available: Optional[int], *, device=None) -> str: snap = snapshot(device=device) return ( diff --git a/toolkit/memory_management/runtime.py b/toolkit/memory_management/runtime.py index a28b65226a..d00ef19f8c 100644 --- a/toolkit/memory_management/runtime.py +++ b/toolkit/memory_management/runtime.py @@ -66,17 +66,3 @@ def close_memory_runtime_preparation(model_owner) -> None: operation = getattr(model_owner, "cleanup_memory_runtime_preparation", None) if operation is not None: operation() - - -def memory_sampling_step_trim(model) -> bool: - """Run the legacy backend's optional per-step sampling trim.""" - inner = unwrap_memory_model(model) - operation = getattr(inner, "_mm_sampling_step_trim", None) - return bool(operation is not None and operation()) - - -def memory_sampling_demote(model, *, reason) -> bool: - """Request sampling residency relief from the active legacy backend.""" - inner = unwrap_memory_model(model) - operation = getattr(inner, "_mm_sampling_demote", None) - return bool(operation is not None and operation(reason=reason)) diff --git a/toolkit/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py index 022672944a..3676514aef 100644 --- a/toolkit/memory_management/vram_budget.py +++ b/toolkit/memory_management/vram_budget.py @@ -12,8 +12,8 @@ >= hard). * **Shared / DXGI NON_LOCAL budget** (NOT this module): pinned host memory commits against it and exhausting it is a hard cudaErrorMemoryAllocation. - That reserve lives in ``pin_manager`` / ``bounce_pool`` - (``dxgi_spill_reserve_bytes``); do not conflate the two. + That reserve lives in ``pin_manager`` (``dxgi_spill_reserve_bytes``); do + not conflate the two. The free signal (do not regress this) ------------------------------------- @@ -48,7 +48,6 @@ from __future__ import annotations import math -import os from dataclasses import dataclass from typing import Optional @@ -59,11 +58,6 @@ GIB = 1024 ** 3 -def _env(name: str, default: str) -> str: - value = os.environ.get(name) - return default if value is None or value == "" else value - - def reconcile_free_bytes(driver_free_bytes, physical_free_bytes) -> int: """Pick the governing device-free value (pure; CPU-testable). @@ -448,34 +442,6 @@ def sampling_allocator_budget_free_bytes( ) -def sampling_guard_predicted_peak_free(total_b, free_b, reserved_b, peak_reserved_b): - """Predicted free VRAM at the next forward's peak (pure). - - ``non_torch = (total - free) - reserved`` plus the worst forward's reserved - high-water is what the next peak will occupy; the prediction is ``total`` - minus that. It shrinks one-for-one as external use grows -- which is the - cohabitation guard's trigger. Forward-only sampling never OOMs at the cliff - (it pages silently), so the guard watches this instead of an exception. - """ - other_b = max(0, (total_b - free_b) - reserved_b) - return total_b - (peak_reserved_b + other_b) - - -def training_cliff_predicted_peak_free_gib( - total_gib, device_free_gib, torch_reserved_gib, peak_allocated_gib -): - """Driver free expected when the next step rebuilds its live peak (pure). - - ``empty_cache`` can make step-end free look healthy by dropping idle cached - blocks, but the next forward/backward will recreate the live peak. Keep - non-allocator residents (``non_torch``) from the current snapshot and ask - whether peak allocated memory itself clears the WDDM hard floor. - """ - device_used_gib = max(0.0, total_gib - device_free_gib) - non_torch_gib = max(0.0, device_used_gib - torch_reserved_gib) - return total_gib - (max(0.0, peak_allocated_gib) + non_torch_gib) - - def training_promotion_worst_shape_free_gib( *, resident_gib, @@ -513,66 +479,6 @@ def training_promotion_worst_shape_free_gib( return float(total_gib) - predicted_used -def training_eager_promote_blocks( - *, - resident_gib, - block_gib, - ring_gib, - worst_working_reserve_gib, - other_gib, - total_gib, - promote_floor_gib, - max_blocks, -): - """How many equal-sized blocks may be promoted at once while keeping the - predicted worst-shape free margin at or above ``promote_floor_gib`` (pure). - - This is the eager-fill counterpart of the one-block-at-a-time climb: a roomy - card leaves GiBs idle if residency only ever grows one block per cadence - window. The prediction is the same conservative worst-measured-resolution - model as ``training_promotion_worst_shape_free_gib`` -- the blocks are assumed - to add their full size and the ring is assumed not to shrink -- so the floor is - what the run actually keeps free on its tightest measured shape. Returns 0 when - not even one block fits, which the caller reports as a worst-shape veto. - """ - block = float(block_gib) - limit = int(max_blocks) - if block <= 0.0 or limit <= 0: - return 0 - free_now = training_promotion_worst_shape_free_gib( - resident_gib=resident_gib, - added_block_gib=0.0, - ring_gib=ring_gib, - worst_working_reserve_gib=worst_working_reserve_gib, - other_gib=other_gib, - total_gib=total_gib, - ) - room = free_now - float(promote_floor_gib) - if room < block: - return 0 - return min(limit, int(room // block)) - - -def sampling_step_should_trim(free_before_b, trim_margin_b) -> bool: - """Whether realized device-free warrants a per-step cache trim (pure). - - WDDM pages on the committed footprint silently, so the trigger is realized - free, not an allocated-side or peak signal. Trim (empty_cache) is cheap and - non-destructive, so the bar is just "free has dropped into the margin." - """ - return free_before_b < trim_margin_b - - -def sampling_step_should_demote(free_after_b, hard_floor_b) -> bool: - """Whether to escalate to a block demote after a trim (pure). - - Only when trimming left free still under the hard floor -- i.e. there was - no idle cache to reclaim, so the pressure is real (external) and the only - relief is giving back resident weights. - """ - return free_after_b < hard_floor_b - - def estimate_sampling_working_reserve_bytes( image_tokens: int, text_tokens: int = 512, @@ -665,51 +571,6 @@ def estimate_training_working_reserve_bytes( return int(estimate * float(safety)) + int(headroom_bytes) -def sampling_overshoot_margin_bytes( - overshoot_gib: float = 0.86, - safety_gib: float = 0.375, - hard_bytes: int = 0, -) -> int: - """Auto sampling margin in the allocator-cap era (pure, CPU-testable). - - With the reclaim allocator cap guarding the WDDM cliff (a capped allocation - recycles cache or raises a loud OOM, it never silently pages), the sampling - margin's only remaining job is to cover the caching allocator's - reserved-over-allocated overshoot. Measured on Krea2 fp8 512px that overshoot - is ~0.86 GiB and rock-steady (std ~0.05 GiB across 8 seeds), so the auto - margin is that measured overshoot plus one safety block -- NOT the old - ``0.10 * card`` cushion, which was sized for a chaotic allocator that kept - jumping over the limit and no longer misbehaves. Narrowing it hands the - difference straight to resident weights (fewer streamed blocks). Floored at - the hard margin so it can never drop below the WDDM device-free floor. - """ - return int( - max(float(overshoot_gib) + float(safety_gib), float(max(0, int(hard_bytes))) / GIB) - * GIB - ) - - -def training_guard_pressure(dxgi: dict, physical: dict) -> dict: - """Combine the DXGI LOCAL and physical cliff signals (pure). - - Pressure if EITHER signal predicts the next step's peak crosses its floor. - The DXGI LOCAL budget is a per-process OS grant and its usage counter - excludes other processes, so it can bless a layout the physical - (mem_get_info) view already knows will overfill the card -- and vice versa - when the OS shrinks the budget early. The merged dict keeps the DXGI - fields at the top level (``source`` compatibility) and carries the - physical signal under ``physical_*``. - """ - merged = dict(dxgi) - merged["pressure"] = bool(dxgi.get("pressure")) or bool(physical.get("pressure")) - merged["physical_predicted_peak_free_gib"] = physical.get("predicted_peak_free_gib") - merged["physical_target_free_gib"] = physical.get("target_free_gib") - merged["pressure_sources"] = [ - src["source"] for src in (dxgi, physical) if src.get("pressure") - ] - return merged - - # --------------------------------------------------------------------------- # Two-timescale residency control # diff --git a/toolkit/quantization/fp8_linear.py b/toolkit/quantization/fp8_linear.py index 2006de07fe..3a042a16e0 100644 --- a/toolkit/quantization/fp8_linear.py +++ b/toolkit/quantization/fp8_linear.py @@ -8,7 +8,6 @@ from __future__ import annotations -import os from dataclasses import dataclass import torch @@ -17,13 +16,6 @@ from .fp8_transpose import column_major -FP8_STATS = { - "enabled": False, - "training_enabled": False, - "kernel_calls": 0, - "fallback_calls": 0, -} - FP8_LINEAR_EXECUTION_KEY = "toolkit.quantization.fp8_linear" DEFAULT_ACTIVATION_FP8_DTYPE = torch.float8_e4m3fn NATIVE_SCALED_MM_FP8_DTYPE = torch.float8_e4m3fn @@ -137,20 +129,11 @@ def declare_fp8_linear(value) -> Fp8LinearDeclaration | None: value = value.data if isinstance(value, torch.nn.Parameter) else value return _adapt_torchao_fp8(value) or _adapt_quanto_fp8(value) -_FP8_GRAD_INPUT = os.environ.get("AI_TOOLKIT_FP8_GRAD_INPUT", "0").lower() not in ( - "0", "false", "no", "off", "", -) -_FP8_GRAD_VERIFIED = None -_REUSE_DEQUANT = os.environ.get("AI_TOOLKIT_REUSE_DEQUANT", "1").lower() not in ( - "0", "false", "no", "off", "", -) -_REUSE_VERIFIED = None +_FP8_GRAD_INPUT = False def set_fp8_grad_input_enabled(enabled: bool) -> None: - global _FP8_GRAD_INPUT, _FP8_GRAD_VERIFIED - if bool(enabled) and not _FP8_GRAD_INPUT: - _FP8_GRAD_VERIFIED = None + global _FP8_GRAD_INPUT _FP8_GRAD_INPUT = bool(enabled) @@ -159,14 +142,6 @@ def fp8_grad_input_enabled() -> bool: return bool(_FP8_GRAD_INPUT) -def reference_dequantize_to(tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - try: - return tensor.dequantize(output_dtype=dtype) - except TypeError: - value = tensor.dequantize() - return value if value.dtype == dtype else value.to(dtype=dtype) - - def _scale_view_shape(spec, qdata, scale): if spec.scale_granularity == "output_row": if scale.numel() != qdata.shape[0]: @@ -195,77 +170,6 @@ def dequantize_rowwise(qdata, scale, dtype): return materialize_fp8_weight(spec, qdata, scale, dtype) -def fast_dequantize_into(qweight, dest): - declaration = declare_fp8_linear(qweight) - if declaration is None or dest is None: - return None - spec = declaration.spec - qdata, scale = declaration.qdata, declaration.scale - if qdata.shape != dest.shape or spec.has_zero_point: - return None - try: - view_shape = _scale_view_shape(spec, qdata, scale) - except ValueError: - return None - dest.copy_(qdata) - dest.mul_(scale.reshape(view_shape).to(dest.dtype)) - return dest - - -def fast_dequantize(qweight, dtype): - global _REUSE_VERIFIED - if not _REUSE_DEQUANT or _REUSE_VERIFIED is False: - return None - if dtype not in (torch.bfloat16, torch.float16, torch.float32): - return None - declaration = declare_fp8_linear(qweight) - if declaration is None: - return None - qdata = declaration.qdata - dest = torch.empty(qdata.shape, dtype=dtype, device=qdata.device) - fast = fast_dequantize_into(qweight, dest) - if fast is None: - return None - if _REUSE_VERIFIED is None: - try: - reference = reference_dequantize_to(qweight, dtype) - ok = reference.shape == fast.shape and torch.allclose( - fast, reference, rtol=1e-2, atol=1e-2 - ) - except Exception: - ok = False - _REUSE_VERIFIED = bool(ok) - if not ok: - return None - return fast - - -def dequantize_to(tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: - fast = fast_dequantize(tensor, dtype) - return fast if fast is not None else reference_dequantize_to(tensor, dtype) - - -def dequantize_into(qweight, dest): - global _REUSE_VERIFIED - if not _REUSE_DEQUANT or _REUSE_VERIFIED is False or dest is None: - return None - fast = fast_dequantize_into(qweight, dest) - if fast is None: - return None - if _REUSE_VERIFIED is None: - try: - reference = reference_dequantize_to(qweight, dest.dtype) - ok = reference.shape == fast.shape and torch.allclose( - fast, reference, rtol=1e-2, atol=1e-2 - ) - except Exception: - ok = False - _REUSE_VERIFIED = bool(ok) - if not ok: - return None - return fast - - def native_variant_supported(spec) -> bool: return bool( spec.weight_dtype == NATIVE_SCALED_MM_FP8_DTYPE @@ -293,11 +197,6 @@ def supports_native_scaled_mm(spec, qdata, scale, *, device=None) -> bool: ) -def native_rowwise_qualifies(qdata, scale, *, device=None) -> bool: - spec = _spec_for_payload(qdata, scale, "output_row") - return supports_native_scaled_mm(spec, qdata, scale, device=device) - - def native_device_supported(device) -> bool: target = torch.device(device) if target.type != "cuda" or not hasattr(torch, "_scaled_mm"): @@ -308,31 +207,6 @@ def native_device_supported(device) -> bool: return False -def weight_format_key(weight) -> str: - value = weight.data if isinstance(weight, torch.nn.Parameter) else weight - declaration = declare_fp8_linear(value) - if declaration is None: - return "other" - return ( - "rowwise_fp8" - if native_variant_supported(declaration.spec) - else "fp8" - ) - - -def fp8_sampling_qualifies(weight, *, device=None) -> bool: - declaration = declare_fp8_linear(weight) - return bool( - declaration is not None - and supports_native_scaled_mm( - declaration.spec, - declaration.qdata, - declaration.scale, - device=device, - ) - ) - - def native_linear( x, qdata_t, @@ -402,61 +276,6 @@ def _grad_input_compute( return None -def grad_input_supported(qdata, scale, grad_out, spec=None) -> bool: - if not _FP8_GRAD_INPUT or _FP8_GRAD_VERIFIED is False: - return False - spec = spec or _spec_for_payload(qdata, scale, "output_row") - return bool( - supports_native_scaled_mm(spec, qdata, scale, device=grad_out.device) - and grad_out.dtype in (torch.bfloat16, torch.float16) - and grad_out.shape[-1] == qdata.shape[0] - ) - - -def grad_input_supported_weight(qweight, grad_out) -> bool: - declaration = declare_fp8_linear(qweight) - return bool( - declaration is not None - and grad_input_supported( - declaration.qdata, - declaration.scale, - grad_out, - declaration.spec, - ) - ) - - -def grad_input(grad_out, qweight, target_dtype): - declaration = declare_fp8_linear(qweight) - if declaration is None: - return None - spec = declaration.spec - qdata, scale = declaration.qdata, declaration.scale - global _FP8_GRAD_VERIFIED - out = ( - _grad_input_compute( - grad_out, - qdata, - scale, - target_dtype, - spec.activation_dtype, - ) - if grad_input_supported(qdata, scale, grad_out, spec) - else None - ) - if out is not None and _FP8_GRAD_VERIFIED is None: - try: - reference = grad_out.to(target_dtype) @ dequantize_to(qweight, target_dtype) - _FP8_GRAD_VERIFIED = bool( - torch.allclose(out, reference, rtol=2e-2, atol=2e-2) - ) - except Exception: - _FP8_GRAD_VERIFIED = False - if out is not None and _FP8_GRAD_VERIFIED: - return out - return grad_out.to(target_dtype) @ dequantize_to(qweight, target_dtype) - - class _NativeTrainingFn(torch.autograd.Function): @staticmethod def forward(ctx, x, qdata_t, scale_row, bias, activation_dtype): @@ -507,49 +326,6 @@ def native_linear_training( ) -def fp8_linear_inference(x, weight, bias): - declaration = declare_fp8_linear(weight) - if ( - declaration is None - or x.dtype not in (torch.bfloat16, torch.float16) - or x.numel() == 0 - ): - FP8_STATS["fallback_calls"] += int( - FP8_STATS["enabled"] or FP8_STATS["training_enabled"] - ) - return None - spec = declaration.spec - qdata, scale = declaration.qdata, declaration.scale - if ( - not supports_native_scaled_mm(spec, qdata, scale, device=x.device) - or qdata.device != x.device - or scale.device != x.device - or (bias is not None and bias.device != x.device) - or x.shape[-1] != qdata.shape[1] - ): - FP8_STATS["fallback_calls"] += int( - FP8_STATS["enabled"] or FP8_STATS["training_enabled"] - ) - return None - try: - out = native_linear( - x, - qdata.t(), - scale, - bias, - spec.activation_dtype, - ) - except RuntimeError: - FP8_STATS["fallback_calls"] += int( - FP8_STATS["enabled"] or FP8_STATS["training_enabled"] - ) - return None - FP8_STATS["kernel_calls"] += int( - FP8_STATS["enabled"] or FP8_STATS["training_enabled"] - ) - return out - - @dataclass(frozen=True) class BoundFp8LinearOperation: spec: Fp8LinearSpec @@ -630,18 +406,6 @@ def bind_fp8_linear( ) -def bind_rowwise_fp8(qdata, scale, *, device, has_bias=True) -> BoundFp8LinearOperation: - """Compatibility binder for an explicitly declared rowwise payload.""" - spec = _spec_for_payload(qdata, scale, "output_row") - return bind_fp8_linear( - spec, - qdata, - scale, - device=device, - has_bias=has_bias, - ) - - @dataclass(frozen=True) class BoundDenseLinearOperation: format_key = "dense" @@ -685,26 +449,6 @@ def materialize(self, tensors, dtype=None): return weight if dtype is None or weight.dtype == dtype else weight.to(dtype=dtype) -def bind_linear_operation(weight, bias=None, *, device): - value = weight.data if isinstance(weight, torch.nn.Parameter) else weight - declaration = declare_fp8_linear(value) - if declaration is None: - try: - value.__tensor_flatten__() - except Exception: - pass - else: - raise ValueError("unsupported_quantized_linear_operation") - return BoundDenseLinearOperation(bias_index=1 if bias is not None else None) - return bind_fp8_linear( - declaration.spec, - declaration.qdata, - declaration.scale, - device=device, - has_bias=bias is not None, - ) - - def bind_parameter_operation(weight, bias=None, *, device): """Bind an operation and snapshot its explicit ordered tensor tuple.""" from .storage import linear_storage_binding @@ -752,9 +496,3 @@ def bind_storage_operation( if int(weight_leaf_count) == 1 and dense_declaration: return BoundDenseLinearOperation(bias_index=1 if len(tensors) > 1 else None) raise ValueError("unsupported_linear_storage_operation") - - -# Transitional names for callers migrating from manager_modules. -_fp8_linear_compiled = native_linear -_fp8_linear_training = native_linear_training -_fp8_grad_input_compute = _grad_input_compute diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index 6a4e03d715..b67da99b51 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -13,7 +13,6 @@ torchao_quantize_, ) from optimum.quanto import freeze -from optimum.quanto.tensor.qbytes import QBytesTensor from tqdm import tqdm from safetensors.torch import load_file from huggingface_hub import hf_hub_download @@ -30,37 +29,6 @@ if TYPE_CHECKING: from toolkit.models.base_model import BaseModel - -def tensor_subclass_leaves(value: torch.Tensor) -> list[torch.Tensor]: - try: - names, _context = value.__tensor_flatten__() - except Exception: - return [value] - leaves = [] - for name in names: - inner = getattr(value, name) - if inner is not None: - leaves.extend(tensor_subclass_leaves(inner)) - return leaves - - -def _tensor_subclass_to_meta(value: torch.Tensor) -> torch.Tensor: - try: - names, context = value.__tensor_flatten__() - except Exception: - return value.to(device="meta") - inner = { - name: ( - None - if getattr(value, name) is None - else _tensor_subclass_to_meta(getattr(value, name)) - ) - for name in names - } - return type(value).__tensor_unflatten__( - inner, context, value.size(), value.stride() - ) - # the quantize function in quanto had a bug where it was using exclude instead of include Q_MODULES = [ @@ -274,181 +242,6 @@ def filter_fn(module: torch.nn.Module, fqn: str) -> bool: # raise e -def quantize_module_at_path( - model: torch.nn.Module, - name: str, - *, - weights, -) -> torch.nn.Module: - """Quantize one already-materialized module and publish its replacement. - - Recursive Quanto quantization cannot replace the root module: its relative - name is empty, so the generated QLinear is assigned to an unusable empty - attribute while the original Linear's weight is cleared. Bounded loaders - know the real parent path and use this helper for root quantization units. - """ - module = model.get_submodule(name) - resolved = get_qtype(weights) - if isinstance(resolved, aotype): - quantize(module, weights=resolved) - elif isinstance(resolved, ostristype): - if isinstance(module, torch.nn.Linear): - # Shape-ineligible Ostris roots remain dense, matching recursive - # quantize() behavior for children (for example Krea's 12-wide - # text-fusion projector). - convert_linear_to_ostris(module, resolved.quantizer) - else: - quantize(module, weights=resolved) - else: - _quantize_submodule( - model, - name, - module, - weights=resolved, - ) - return model.get_submodule(name) - - -def assign_quantized_state_dict( - model: torch.nn.Module, - state_dict: dict, - weights, -) -> None: - """Assign cached quantized state, using a generic active arena session.""" - prepare_quantized_state_dict_model(model, state_dict, weights) - from toolkit.memory_management.arena_offload.load_session import ( - try_prepare_canonical_from_state_dict, - ) - - if try_prepare_canonical_from_state_dict(model, state_dict) is not None: - try: - assign_quantized_state_dict_subset(model, state_dict, weights) - return - except BaseException: - from toolkit.memory_management.arena_offload.load_session import ( - discard_pending_canonical_build, - ) - - discard_pending_canonical_build(model) - raise - missing, unexpected = model.load_state_dict(state_dict, strict=True, assign=True) - if missing or unexpected: - raise RuntimeError(f"missing={missing[:5]} unexpected={unexpected[:5]}") - model.requires_grad_(False) - - -def prepare_quantized_state_dict_model( - model: torch.nn.Module, - state_dict: dict, - weights, -) -> None: - """Reconstruct cached quantized wrappers with meta storage only.""" - resolved = get_qtype(weights) - quanto_data_suffix = ".weight._data" - quanto_prefixes = [ - key[: -len(quanto_data_suffix)] - for key in state_dict - if key.endswith(quanto_data_suffix) - ] - if quanto_prefixes: - if isinstance(resolved, (aotype, ostristype)): - raise ValueError("cached_quanto_state_qtype_mismatch") - quantize(model, weights=resolved) - modules = dict(model.named_modules()) - for prefix in quanto_prefixes: - module = modules[prefix] - data = state_dict[f"{prefix}.weight._data"] - scale = state_dict[f"{prefix}.weight._scale"] - template = module.weight - meta_data = torch.empty_like(data, device="meta") - meta_scale = torch.empty_like(scale, device="meta") - wrapper = QBytesTensor( - resolved, - 0, - template.size(), - template.stride(), - meta_data, - meta_scale, - requires_grad=False, - ) - module.weight = torch.nn.Parameter(wrapper, requires_grad=False) - bias = getattr(module, "bias", None) - if bias is not None: - bias.requires_grad_(False) - else: - modules = dict(model.named_modules()) - for key, value in state_dict.items(): - if not key.endswith(".weight") or len(tensor_subclass_leaves(value)) == 1: - continue - prefix = key[: -len(".weight")] - module = modules.get(prefix) - if module is None or not hasattr(module, "weight"): - continue - module.weight = torch.nn.Parameter( - _tensor_subclass_to_meta(value), requires_grad=False - ) - - -def assign_quantized_state_dict_subset( - model: torch.nn.Module, - state_dict: dict, - weights, - *, - excluded_keys=(), -) -> None: - """Assign cached values except leaves owned by a direct arena destination.""" - excluded = set(excluded_keys) - prepare_quantized_state_dict_model(model, state_dict, weights) - resolved = get_qtype(weights) - data_suffix = ".weight._data" - prefixes = { - key[: -len(data_suffix)] - for key in state_dict - if key.endswith(data_suffix) - } - handled = set() - modules = dict(model.named_modules()) - for prefix in prefixes: - data_key = f"{prefix}.weight._data" - scale_key = f"{prefix}.weight._scale" - if data_key in excluded and scale_key in excluded: - continue - if data_key in excluded or scale_key in excluded: - raise ValueError(f"partial_cached_quantized_weight:{prefix}") - module = modules[prefix] - template = module.weight - wrapper = QBytesTensor( - resolved, - 0, - template.size(), - template.stride(), - state_dict[data_key], - state_dict[scale_key], - requires_grad=False, - ) - module.weight = torch.nn.Parameter(wrapper, requires_grad=False) - handled.update((data_key, scale_key)) - - for key, value in state_dict.items(): - if key in excluded or key in handled: - continue - *parents, leaf = key.split(".") - target = model - for component in parents: - target = getattr(target, component) - if leaf in target._parameters: - old = target._parameters[leaf] - requires_grad = bool(old.requires_grad) if old is not None else False - target._parameters[leaf] = torch.nn.Parameter( - value, requires_grad=requires_grad - ) - elif leaf in target._buffers: - target._buffers[leaf] = value - else: - raise KeyError(f"unsupported_cached_state_leaf:{key}") - model.requires_grad_(False) - - def quantize_model( base_model: "BaseModel", model_to_quantize: torch.nn.Module, @@ -503,7 +296,6 @@ def quantize_model( "transformer_only": False, } first_key = list(lora_state_dict.keys())[0] - first_weight = lora_state_dict[first_key] # if it starts with lycoris and includes lokr if first_key.startswith("lycoris") and any( "lokr" in key for key in lora_state_dict.keys() From 8ca15af86c9f694d9d829254f81f247e990b9f24 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Thu, 16 Jul 2026 15:40:57 +0200 Subject: [PATCH 18/20] Enable Arena fullgraph MegaCache --- tests/test_arena_compile_cache.py | 95 ++++++++++ tests/test_generic_block_dispatcher.py | 97 +++++++++++ .../memory_management/arena_offload/api.py | 4 + .../arena_offload/compile_cache.py | 162 ++++++++++++++++++ .../arena_offload/dispatcher.py | 46 ++++- .../arena_offload/runtime.py | 23 +++ 6 files changed, 424 insertions(+), 3 deletions(-) create mode 100644 tests/test_arena_compile_cache.py create mode 100644 toolkit/memory_management/arena_offload/compile_cache.py diff --git a/tests/test_arena_compile_cache.py b/tests/test_arena_compile_cache.py new file mode 100644 index 0000000000..dd63841753 --- /dev/null +++ b/tests/test_arena_compile_cache.py @@ -0,0 +1,95 @@ +from types import SimpleNamespace +from unittest import mock + +import torch + +from toolkit.memory_management.arena_offload.compile_cache import ( + ArenaCompileCacheSession, + arena_compile_cache_key, +) + + +class _Transformer(torch.nn.Module): + pass + + +def _config(**overrides): + values = { + "compile_blocks": True, + "fp8_forward": True, + "fp8_backward": False, + "fp8_sampling": True, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def _executor(**overrides): + values = { + "compile_fullgraph": True, + "compile_dynamic": False, + "compile_dynamic_hints": ((1, 64, 4096),), + "_programs": { + "train": SimpleNamespace(fingerprint="train-abi"), + "sample": SimpleNamespace(fingerprint="sample-abi"), + }, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def test_cache_key_is_stable_and_covers_compile_policy(): + model = _Transformer() + config = _config() + executor = _executor() + + first = arena_compile_cache_key(model, config, executor) + second = arena_compile_cache_key(model, config, executor) + changed = arena_compile_cache_key( + model, config, _executor(compile_fullgraph=False) + ) + + assert first == second + assert first != changed + assert len(first) == 64 + + +def test_cache_session_saves_atomically_and_loads(tmp_path): + session = ArenaCompileCacheSession.for_runtime( + _Transformer(), _config(), _executor(), cache_root=tmp_path + ) + info = SimpleNamespace(artifacts={}) + + with mock.patch.object( + torch.compiler, + "save_cache_artifacts", + return_value=(b"arena-artifacts", info), + ) as save_artifacts: + assert session.save(force=True) + + assert session.path.read_bytes() == b"arena-artifacts" + assert not list(tmp_path.glob("*.tmp")) + save_artifacts.assert_called_once_with() + + restored = ArenaCompileCacheSession(session.path, session.key) + with mock.patch.object( + torch.compiler, "load_cache_artifacts", return_value=info + ) as load_artifacts: + assert restored.load() + + load_artifacts.assert_called_once_with(b"arena-artifacts") + assert restored.diagnostics()["byte_count"] == len(b"arena-artifacts") + + +def test_cache_io_failure_is_non_fatal(tmp_path): + session = ArenaCompileCacheSession.for_runtime( + _Transformer(), _config(), _executor(), cache_root=tmp_path + ) + with mock.patch.object( + torch.compiler, + "save_cache_artifacts", + side_effect=OSError("synthetic write failure"), + ), mock.patch("toolkit.memory_management.arena_offload.compile_cache.warnings.warn"): + assert not session.save(force=True) + + assert "synthetic write failure" in session.diagnostics()["error"] diff --git a/tests/test_generic_block_dispatcher.py b/tests/test_generic_block_dispatcher.py index 5494907b2d..3cbe5a9935 100644 --- a/tests/test_generic_block_dispatcher.py +++ b/tests/test_generic_block_dispatcher.py @@ -18,6 +18,7 @@ from toolkit.memory_management.arena_offload.dispatcher import ( _first_output_tensor, _first_tensor_argument, + _is_dynamo_compile_failure, _replace_tensor_argument, ) from toolkit.memory_management.arena_offload.ownership import active_process_owner @@ -388,6 +389,98 @@ def flaky(*args, **kwargs): close_arena_offload(model) +def test_model_config_propagates_fullgraph_only_when_compile_is_enabled(): + enabled = ArenaOffloadConfig.from_model_config( + type( + "Config", + (), + { + "layer_offloading": True, + "layer_offloading_smart": True, + "compile": True, + "compile_fullgraph": True, + }, + )() + ) + disabled = ArenaOffloadConfig.from_model_config( + type( + "Config", + (), + { + "layer_offloading": True, + "layer_offloading_smart": True, + "compile": False, + "compile_fullgraph": True, + }, + )() + ) + + assert enabled._compile_fullgraph is True + assert disabled._compile_fullgraph is False + + +def test_strict_compile_failure_classifier_includes_recompile_limit(): + error = torch._dynamo.exc.FailOnRecompileLimitHit( + "synthetic recompile limit" + ) + assert _is_dynamo_compile_failure(error) + assert not _is_dynamo_compile_failure(RuntimeError("ordinary block error")) + + +def test_later_strict_recompile_failure_keeps_block_identity(): + model = _frozen_transformer() + model.enable_gradient_checkpointing() + with mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget.device_mem_info", + return_value=(8 * 1024**3, 12 * 1024**3), + ), mock.patch( + "toolkit.memory_management.arena_offload.planner.vram_budget." + "auto_physical_vram_headroom_gib", + return_value=1.0, + ): + runtime = prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=ArenaOffloadConfig(enabled=True, compile_blocks=False), + ) + try: + runtime.finalize() + executor = runtime._executor + executor.compile_fullgraph = True + executor._strict_kernels_executed.add(0) + + def failed_recompile(_index): + def fail(*_args, **_kwargs): + raise torch._dynamo.exc.FailOnRecompileLimitHit( + "synthetic later recompile" + ) + + return fail + + executor._get_dispatch_kernel = failed_recompile + resident = [ + (block_key, leaf_name) + for block_key in runtime._arena.block_keys() + for leaf_name in runtime._arena.block_record(block_key).leaf_names + ] + executor.activate( + executor.SAMPLE, + ResidencyPlan.build(executor.SAMPLE, resident), + ) + with torch.no_grad(), executor.execution(executor.SAMPLE), pytest.raises( + RuntimeError, + match="fullgraph_block_compile_failed:blocks.0", + ) as raised: + model(torch.randn(2, 4)) + assert isinstance( + raised.value.__cause__, + torch._dynamo.exc.FailOnRecompileLimitHit, + ) + finally: + close_arena_offload(model) + + def test_saved_installed_forward_checkpoint_backward_and_teardown(): torch.manual_seed(17) model = _frozen_transformer() @@ -477,6 +570,7 @@ def test_cuda_streamed_compiled_train_sample_train(): enabled=True, compile_blocks=True, _compile_dynamic=False, + _compile_fullgraph=True, ) config = replace( config, @@ -512,6 +606,9 @@ def installed_forward(self, value, _saved=saved): adapters.append(block.adapter_gain) runtime.finalize() diagnostics = runtime.diagnostics() + assert diagnostics["compile_fullgraph"] is True + assert diagnostics["compile_cache"]["enabled"] is True + assert diagnostics["compile_cache"]["load_attempted"] is True accounting = diagnostics["accounting"] assert accounting["payload_reconciled"] assert accounting["mixed_residency"] diff --git a/toolkit/memory_management/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py index 928d9439ca..b65a6375cb 100644 --- a/toolkit/memory_management/arena_offload/api.py +++ b/toolkit/memory_management/arena_offload/api.py @@ -145,6 +145,7 @@ class ArenaOffloadConfig: compile_blocks: bool = False strict_vram_cap: bool = False _compile_dynamic: bool | None = True + _compile_fullgraph: bool = False _compile_dynamic_hints: tuple[tuple[int, int | None, int | None], ...] = () # Validation knob: pretend the card is this many GiB, so small-card # behaviour (deeper streaming, tighter caps, a residency plan that cannot @@ -239,6 +240,9 @@ def get(name: str, default: Any = None) -> Any: if get("compile_dynamic", True) is None else bool(get("compile_dynamic", True)) ), + _compile_fullgraph=bool( + get("compile", False) and get("compile_fullgraph", False) + ), _compile_dynamic_hints=tuple( tuple(hint) for hint in (get("compile_dynamic_hints", ()) or ()) ), diff --git a/toolkit/memory_management/arena_offload/compile_cache.py b/toolkit/memory_management/arena_offload/compile_cache.py new file mode 100644 index 0000000000..dfe4b2740d --- /dev/null +++ b/toolkit/memory_management/arena_offload/compile_cache.py @@ -0,0 +1,162 @@ +"""Persistent torch.compile artifacts for the Arena block dispatcher.""" + +from __future__ import annotations + +import hashlib +import json +import os +import warnings +from pathlib import Path + +import torch + + +_CACHE_SCHEMA = "aitk-arena-megacache-v1" + + +def _dynamo_frame_count() -> int: + try: + return int(torch._dynamo.utils.counters["frames"].get("total", 0)) + except Exception: + return 0 + + +def _default_cache_root() -> Path: + from torch._inductor.runtime.runtime_utils import cache_dir + + return Path(cache_dir()) / "aitk_arena_megacache" + + +def arena_compile_cache_key(model, config, executor) -> str: + """Hash the coarse identity of one cumulative Arena compile cache. + + Torch still validates every artifact's graph and guards. This outer key + separates compiler/runtime policies and immutable dispatcher ABIs while + allowing train, sample, shape, and residency variants to accumulate. + """ + from .dispatcher import DISPATCHER_GENERATION + + model_type = type(model) + programs = getattr(executor, "_programs", {}) + identity = { + "schema": _CACHE_SCHEMA, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "model_type": f"{model_type.__module__}.{model_type.__qualname__}", + "dispatcher": DISPATCHER_GENERATION, + "programs": sorted( + (str(mode), str(program.fingerprint)) + for mode, program in programs.items() + ), + "compile": { + "fullgraph": bool(getattr(executor, "compile_fullgraph", False)), + "dynamic": getattr(executor, "compile_dynamic", True), + "dynamic_hints": tuple( + getattr(executor, "compile_dynamic_hints", ()) or () + ), + }, + "fp8": { + "forward": bool(getattr(config, "fp8_forward", False)), + "backward": bool(getattr(config, "fp8_backward", False)), + "sampling": bool(getattr(config, "fp8_sampling", False)), + }, + } + encoded = json.dumps( + identity, sort_keys=True, separators=(",", ":"), default=str + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +class ArenaCompileCacheSession: + """Best-effort load/save around the Arena's lazy compile lifecycle.""" + + def __init__(self, path: Path | None, key: str | None): + self.path = path + self.key = key + self.enabled = path is not None + self.load_attempted = False + self.loaded = False + self.saved = False + self.error: str | None = None + self.byte_count = 0 + self._last_saved_frames = _dynamo_frame_count() + + @classmethod + def for_runtime(cls, model, config, executor, *, cache_root=None): + supported = all( + callable(getattr(torch.compiler, name, None)) + for name in ("load_cache_artifacts", "save_cache_artifacts") + ) + if not bool(getattr(config, "compile_blocks", False)) or not supported: + return cls(None, None) + key = arena_compile_cache_key(model, config, executor) + root = Path(cache_root) if cache_root is not None else _default_cache_root() + return cls(root / f"{key}.torchcompile", key) + + def _warn(self, operation: str, error: BaseException) -> None: + self.error = f"{type(error).__name__}: {error}" + warnings.warn( + f"Arena MegaCache {operation} failed; continuing without it: " + f"{self.error}", + RuntimeWarning, + stacklevel=2, + ) + + def load(self) -> bool: + if not self.enabled or self.load_attempted: + return False + self.load_attempted = True + if not self.path.is_file(): + return False + try: + artifacts = self.path.read_bytes() + info = torch.compiler.load_cache_artifacts(artifacts) + except Exception as error: + self._warn("load", error) + return False + self.loaded = info is not None + self.byte_count = len(artifacts) if self.loaded else 0 + self._last_saved_frames = _dynamo_frame_count() + return self.loaded + + def save(self, *, force: bool = False) -> bool: + if not self.enabled: + return False + frames = _dynamo_frame_count() + if not force and frames <= self._last_saved_frames: + return False + temporary = self.path.with_name(f"{self.path.name}.{os.getpid()}.tmp") + try: + result = torch.compiler.save_cache_artifacts() + if result is None: + self._last_saved_frames = frames + return False + artifacts, _info = result + self.path.parent.mkdir(parents=True, exist_ok=True) + temporary.write_bytes(artifacts) + os.replace(temporary, self.path) + except Exception as error: + try: + temporary.unlink(missing_ok=True) + except OSError: + pass + self._last_saved_frames = frames + self._warn("save", error) + return False + self.saved = True + self.error = None + self.byte_count = len(artifacts) + self._last_saved_frames = frames + return True + + def diagnostics(self) -> dict: + return { + "enabled": self.enabled, + "key": self.key, + "path": None if self.path is None else str(self.path), + "load_attempted": self.load_attempted, + "loaded": self.loaded, + "saved": self.saved, + "byte_count": self.byte_count, + "error": self.error, + } diff --git a/toolkit/memory_management/arena_offload/dispatcher.py b/toolkit/memory_management/arena_offload/dispatcher.py index a94ad2e35e..25e72b75fc 100644 --- a/toolkit/memory_management/arena_offload/dispatcher.py +++ b/toolkit/memory_management/arena_offload/dispatcher.py @@ -21,7 +21,7 @@ ) -DISPATCHER_GENERATION = "generic-block-dispatcher-v1" +DISPATCHER_GENERATION = "generic-block-dispatcher-v2-fullgraph" def _first_tensor_argument(args, kwargs): @@ -70,6 +70,24 @@ def _in_backward_graph_task() -> bool: return False +def _is_dynamo_compile_failure(error: BaseException) -> bool: + """Recognize strict-capture failures, including later recompiles.""" + dynamo_errors = getattr(torch._dynamo, "exc", None) + if dynamo_errors is None: + return False + error_types = tuple( + error_type + for name in ( + "BackendCompilerFailed", + "FailOnRecompileLimitHit", + "RecompileError", + "Unsupported", + ) + if isinstance((error_type := getattr(dynamo_errors, name, None)), type) + ) + return bool(error_types) and isinstance(error, error_types) + + class OriginalBlockInvoker(torch.nn.Module): """Own a selected block while calling its preserved installed forward.""" @@ -112,11 +130,14 @@ def __init__( depth=3, compile_blocks=True, compile_dynamic=True, + compile_fullgraph=False, compile_dynamic_hints=(), protected_training_leaf_keys=(), owner_token=None, ): self.selection = selection + self.compile_fullgraph = bool(compile_blocks and compile_fullgraph) + self._strict_kernels_executed = set() super().__init__( model, residency, @@ -221,7 +242,7 @@ def kernel(leaf_args, args, kwargs): kernel = torch.compile( kernel, mode="default", - fullgraph=False, + fullgraph=self.compile_fullgraph, dynamic=self.compile_dynamic, ) self._block_kernels[key] = kernel @@ -343,6 +364,8 @@ def dispatch(self, index, args, kwargs): while True: try: output = self._get_dispatch_kernel(index)(leaf_args, args, kwargs) + if self.compile_fullgraph: + self._strict_kernels_executed.add(int(index)) break except BaseException as error: recover = ( @@ -358,8 +381,22 @@ def dispatch(self, index, args, kwargs): # normal post-forward release below. With no gradient-bearing # input, frozen streamed state is not a backward dependency and # there is no free_on_backward node, so release on that unwind. - if token is not None and not release_on_backward: + strict_compile_failed = ( + self.compile_fullgraph + and ( + int(index) not in self._strict_kernels_executed + or _is_dynamo_compile_failure(error) + ) + ) + if token is not None and ( + not release_on_backward or strict_compile_failed + ): torch.ops.mm.fetch_free(token) + if strict_compile_failed: + raise ImmutableRuntimeError( + "fullgraph_block_compile_failed:" + f"{self._block_abis[index].block_key}" + ) from error raise if token is not None: # The first checkpoint pass discards its fetched views, so return @@ -398,6 +435,7 @@ def close(self): self._sampling_forward_begin = None self._sampling_forward_end = None self._sampling_allocation_failure = None + self._strict_kernels_executed.clear() super().close() @@ -409,6 +447,7 @@ def prepare_block_dispatcher_runtime( depth=3, compile_blocks=True, compile_dynamic=True, + compile_fullgraph=False, compile_dynamic_hints=(), protected_training_leaf_keys=(), owner_token=None, @@ -420,6 +459,7 @@ def prepare_block_dispatcher_runtime( depth=depth, compile_blocks=compile_blocks, compile_dynamic=compile_dynamic, + compile_fullgraph=compile_fullgraph, compile_dynamic_hints=compile_dynamic_hints, protected_training_leaf_keys=protected_training_leaf_keys, owner_token=owner_token, diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index 1b6f6d1afd..24ef1ecf10 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -33,6 +33,7 @@ TrainingSignalWindow, ) from .cap_calibrator import CAP_SET, TrainingCapCalibrator +from .compile_cache import ArenaCompileCacheSession from .errors import ArenaCleanupError, ArenaSetupFatalError from .fp8 import disable as disable_fp8 from .fp8 import enable as enable_fp8 @@ -85,6 +86,7 @@ def __init__( self._smart_plan = smart_plan self._canonical_modules = canonical_modules self._resources = resources + self._compile_cache = None self._closed = False self._disposed = False @@ -249,6 +251,7 @@ def _prepare( depth=policy.prefetch_depth, compile_blocks=config.compile_blocks, compile_dynamic=config._compile_dynamic, + compile_fullgraph=config._compile_fullgraph, compile_dynamic_hints=config._compile_dynamic_hints, protected_training_leaf_keys=smart_plan.get( "protected_training_leaf_keys", () @@ -530,6 +533,10 @@ def finalize(self, network=None): # installed, so compiled dispatcher kernels trace the selected # execution policy rather than Quanto's materializing fallback. self._executor.finalize_execution() + self._compile_cache = ArenaCompileCacheSession.for_runtime( + self._model, self._config, self._executor + ) + self._compile_cache.load() self._executor.activate(self._executor.TRAIN, self._training_plan) return self except BaseException as error: @@ -537,8 +544,16 @@ def finalize(self, network=None): def close(self) -> None: """Release through the same owner used during preparation.""" + self._save_compile_cache(force=True) self._resources.release() + def _save_compile_cache(self, *, force: bool = False) -> bool: + cache = getattr(self, "_compile_cache", None) + executor = getattr(self, "_executor", None) + if cache is None or getattr(executor, "active_executions", 0): + return False + return cache.save(force=force) + def _fatal_setup_failure(self, error): try: self._resources.release() @@ -622,6 +637,7 @@ def training_step(self, *, shape_key: tuple | None = None, step_num: int | None except Exception as error: # Diagnostics must never mask a successful training step. self._last_policy_error = f"{type(error).__name__}: {error}" + self._save_compile_cache() def _handle_training_failure(self, error, *, shape_key, step_num): """Clean arena-owned state after the executor has unwound.""" @@ -1242,6 +1258,7 @@ def sampling_image( ) yielded = True yield self + self._save_compile_cache() break except BaseException as error: if yielded or setup_retry_used: @@ -1705,6 +1722,12 @@ def diagnostics(self) -> dict: "prefetch_depth": int(getattr(self._executor, "depth", 0)), "compile_blocks": bool(self._config.compile_blocks), "compile_dynamic": bool(self._config._compile_dynamic), + "compile_fullgraph": bool(self._config._compile_fullgraph), + "compile_cache": ( + None + if self._compile_cache is None + else self._compile_cache.diagnostics() + ), "strict_vram_cap": bool( getattr(self._config, "strict_vram_cap", False) ), From 23d6a54d5706c05660e4ab383f47d28983b814fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Thu, 16 Jul 2026 15:59:39 +0200 Subject: [PATCH 19/20] Use runtime allocator GC threshold in Arena --- tests/test_allocator_gc_threshold.py | 75 +++++++++++++ tests/test_arena_cap_calibrator.py | 32 +++++- tests/test_arena_offload_policy.py | 9 +- tests/test_arena_sampling_cap.py | 4 +- tests/test_residency_two_timescale.py | 49 +++++++-- .../arena_offload/cap_calibrator.py | 29 ++++- .../memory_management/arena_offload/policy.py | 15 +++ .../arena_offload/runtime.py | 29 ++++- toolkit/memory_management/vram_budget.py | 103 +++++++++++++++--- 9 files changed, 302 insertions(+), 43 deletions(-) create mode 100644 tests/test_allocator_gc_threshold.py diff --git a/tests/test_allocator_gc_threshold.py b/tests/test_allocator_gc_threshold.py new file mode 100644 index 0000000000..cb8db1fe8b --- /dev/null +++ b/tests/test_allocator_gc_threshold.py @@ -0,0 +1,75 @@ +import pytest + +from toolkit.memory_management import vram_budget + + +def test_runtime_allocator_settings_prefers_torch_runtime(monkeypatch): + monkeypatch.setenv( + "PYTORCH_ALLOC_CONF", "garbage_collection_threshold:0.61" + ) + monkeypatch.setattr( + vram_budget.torch._C, + "_accelerator_getAllocatorSettings", + lambda: "garbage_collection_threshold:0.79", + ) + + assert vram_budget._runtime_allocator_settings() == ( + "garbage_collection_threshold:0.79" + ) + + +def test_runtime_allocator_settings_falls_back_for_older_torch(monkeypatch): + monkeypatch.setenv( + "PYTORCH_CUDA_ALLOC_CONF", "garbage_collection_threshold:0.67" + ) + monkeypatch.delenv("PYTORCH_ALLOC_CONF", raising=False) + monkeypatch.setattr( + vram_budget.torch._C, + "_accelerator_getAllocatorSettings", + None, + ) + + assert vram_budget._runtime_allocator_settings() == ( + "garbage_collection_threshold:0.67" + ) + + +def test_allocator_gc_threshold_reads_live_torch_settings(monkeypatch): + monkeypatch.setattr( + vram_budget, + "_runtime_allocator_settings", + lambda: ( + "expandable_segments:True," + "garbage_collection_threshold:0.73" + ), + ) + + assert vram_budget.allocator_gc_threshold() == pytest.approx(0.73) + assert vram_budget.allocator_allowance_bytes(1_000, 600) == 130 + assert vram_budget.cap_bytes_for_live(600, 130, 2_000) == 1_000 + assert vram_budget.sampling_allocator_budget_free_bytes( + 1_000, 600, 1.0, 0 + ) == 130 + + +@pytest.mark.parametrize("settings", ["", "backend:cudaMallocAsync"]) +def test_allocator_without_native_gc_uses_cap_as_effective_target( + monkeypatch, settings +): + monkeypatch.setattr( + vram_budget, "_runtime_allocator_settings", lambda: settings + ) + + assert vram_budget.allocator_gc_threshold() == 1.0 + assert vram_budget.allocator_allowance_bytes(1_000, 600) == 400 + + +def test_invalid_fallback_allocator_threshold_fails_closed(monkeypatch): + monkeypatch.setattr( + vram_budget, + "_runtime_allocator_settings", + lambda: "garbage_collection_threshold:not-a-number", + ) + + with pytest.raises(ValueError, match="invalid runtime"): + vram_budget.allocator_gc_threshold() diff --git a/tests/test_arena_cap_calibrator.py b/tests/test_arena_cap_calibrator.py index aa43438c93..1682b9b992 100644 --- a/tests/test_arena_cap_calibrator.py +++ b/tests/test_arena_cap_calibrator.py @@ -1,5 +1,6 @@ from types import SimpleNamespace +from toolkit.memory_management import vram_budget from toolkit.memory_management.arena_offload.cap_calibrator import ( CAP_PROBE_SETTLE, CAP_PROBE_VERIFY, @@ -9,6 +10,12 @@ TrainingCapCalibrator, ) +GC_THRESHOLD = 0.95 + + +def make_calibrator(**kwargs): + return TrainingCapCalibrator(gc_threshold=GC_THRESHOLD, **kwargs) + def peak(working, steps=2): return SimpleNamespace(working_peak_bytes=working, steps=steps) @@ -38,7 +45,7 @@ def drive(calibrator, previous, upcoming, peaks, current): def test_calibrator_verifies_every_bucket_and_restores_last_clean_cap(): - calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + calibrator = make_calibrator(enabled=True, notch_bytes=100) peaks = {("a",): peak(300), ("b",): peak(400)} decision = drive(calibrator, None, ("a",), peaks, 1000) @@ -79,7 +86,7 @@ def test_calibrator_verifies_every_bucket_and_restores_last_clean_cap(): def test_first_gc_during_probe_restores_last_clean_cap(): - calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + calibrator = make_calibrator(enabled=True, notch_bytes=100) peaks = {("a",): peak(400)} drive(calibrator, None, ("a",), peaks, 1000) @@ -95,7 +102,7 @@ def test_first_gc_during_probe_restores_last_clean_cap(): def test_unseen_bucket_restores_cliff_and_invalidates_settlement(): - calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + calibrator = make_calibrator(enabled=True, notch_bytes=100) calibrator.state = CAP_SETTLED calibrator.settled_cap_bytes = 600 calibrator.learned_cache_pad_bytes = 70 @@ -113,7 +120,7 @@ def test_unseen_bucket_restores_cliff_and_invalidates_settlement(): def test_compile_invalid_restores_cliff_and_discards_bucket_profiles(): - calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + calibrator = make_calibrator(enabled=True, notch_bytes=100) peaks = {("a",): peak(400)} drive(calibrator, None, ("a",), peaks, 1000) @@ -130,7 +137,7 @@ def test_compile_invalid_restores_cliff_and_discards_bucket_profiles(): def test_probe_oom_restores_before_residency_recovery(): - calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + calibrator = make_calibrator(enabled=True, notch_bytes=100) peaks = {("a",): peak(400)} drive(calibrator, None, ("a",), peaks, 1000) @@ -143,7 +150,7 @@ def test_probe_oom_restores_before_residency_recovery(): def test_probe_oom_adds_two_notches_when_last_clean_is_too_close(): - calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + calibrator = make_calibrator(enabled=True, notch_bytes=100) calibrator.state = CAP_PROBE_SETTLE calibrator.probe_cap_bytes = 500 calibrator.last_clean_cap_bytes = 600 @@ -156,3 +163,16 @@ def test_probe_oom_adds_two_notches_when_last_clean_is_too_close(): assert decision.target_cap_bytes == 700 assert calibrator.last_clean_cap_bytes == 700 assert calibrator.state == CAP_RESTORE_VERIFY + + +def test_calibrator_uses_runtime_allocator_gc_threshold(monkeypatch): + monkeypatch.setattr( + vram_budget, "allocator_gc_threshold", lambda: 0.8 + ) + calibrator = TrainingCapCalibrator(enabled=True, notch_bytes=100) + peaks = {("a",): peak(400)} + + decision = drive(calibrator, None, ("a",), peaks, 1000) + + assert calibrator.gc_threshold == 0.8 + assert decision.target_cap_bytes == 800 diff --git a/tests/test_arena_offload_policy.py b/tests/test_arena_offload_policy.py index 14266ba9a8..2e0edf1dfa 100644 --- a/tests/test_arena_offload_policy.py +++ b/tests/test_arena_offload_policy.py @@ -227,6 +227,7 @@ def test_worst_shape_allocator_slack_reconstructs_current_layout(): runtime._residency = SimpleNamespace(resident_bytes=lambda: 120) runtime._smart_plan = {"singleton_resident_bytes": 30} runtime._training_ring_bytes = lambda: 50 + runtime._allocator_gc_threshold = 0.95 # working=600, current layout=150 resident + 50 ring => live=800. assert runtime._worst_shape_allocator_slack_bytes(1000) == 150 @@ -587,7 +588,8 @@ def test_controller_raises_cap_by_fixed_fsm_increment(): def test_controller_exactly_prefunds_promotion_after_cap_calibration(): controller = ArenaResidencyController( - allocator_cache_headroom_bytes=10 + allocator_cache_headroom_bytes=10, + gc_threshold=0.8, ) controller.bootstrapped = True controller.state = type(controller.state)("stable", 2) @@ -611,12 +613,13 @@ def test_controller_exactly_prefunds_promotion_after_cap_calibration(): learned_cache_pad_bytes=50, ) assert decision.action == "raise_cap" - assert decision.target_cap_bytes == int(650 / 0.95) + assert decision.target_cap_bytes == int(650 / 0.8) def test_controller_holds_when_exact_promotion_cap_exceeds_cliff(): controller = ArenaResidencyController( - allocator_cache_headroom_bytes=10 + allocator_cache_headroom_bytes=10, + gc_threshold=0.95, ) controller.bootstrapped = True controller.state = type(controller.state)("stable", 2) diff --git a/tests/test_arena_sampling_cap.py b/tests/test_arena_sampling_cap.py index 6a9976e424..eebc30d3c4 100644 --- a/tests/test_arena_sampling_cap.py +++ b/tests/test_arena_sampling_cap.py @@ -100,7 +100,9 @@ def test_sampling_forward_begin_uses_completed_generic_fsm(monkeypatch): runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) key = _sampling_config_shape_key(config()) profile = _SamplingCapProfile( - calibrator=TrainingCapCalibrator(enabled=True, notch_bytes=100), + calibrator=TrainingCapCalibrator( + enabled=True, notch_bytes=100, gc_threshold=0.95 + ), signals=TrainingSignalWindow(), ) profile.signals._shape_peaks[key] = ShapePeak( diff --git a/tests/test_residency_two_timescale.py b/tests/test_residency_two_timescale.py index 98d117c2cc..5988f32170 100644 --- a/tests/test_residency_two_timescale.py +++ b/tests/test_residency_two_timescale.py @@ -11,6 +11,7 @@ from toolkit.memory_management import vram_budget as vb GIB = vb.GIB +GC_THRESHOLD = 0.95 def gib(x): @@ -22,35 +23,55 @@ def gib(x): def test_allowance_matches_measured_knee(): # 0.95*7.35 - 6.77 = +0.21 GiB (clean, the floor cap); 0.95*7.10 - 6.77 = # -0.02 GiB (dirty). The model validated to within one 0.25 notch. - clean = vb.allocator_allowance_bytes(gib(7.35), gib(6.77)) - dirty = vb.allocator_allowance_bytes(gib(7.10), gib(6.77)) + clean = vb.allocator_allowance_bytes( + gib(7.35), gib(6.77), gc_threshold=GC_THRESHOLD + ) + dirty = vb.allocator_allowance_bytes( + gib(7.10), gib(6.77), gc_threshold=GC_THRESHOLD + ) assert clean == pytest.approx(gib(0.2125), abs=gib(0.01)) assert dirty < 0 def test_cap_for_live_is_allowance_inverse(): # To host 6.77 live + 0.21 cache budget the cap must be ~7.35 (the floor). - cap = vb.cap_bytes_for_live(gib(6.77), gib(0.21), cliff_cap_bytes=gib(9.85)) + cap = vb.cap_bytes_for_live( + gib(6.77), + gib(0.21), + cliff_cap_bytes=gib(9.85), + gc_threshold=GC_THRESHOLD, + ) assert cap == pytest.approx(gib(7.35), abs=gib(0.02)) # Round-trips: that cap yields ~the requested cache budget back. - assert vb.allocator_allowance_bytes(cap, gib(6.77)) == pytest.approx(gib(0.21), abs=gib(0.01)) + assert vb.allocator_allowance_bytes( + cap, gib(6.77), gc_threshold=GC_THRESHOLD + ) == pytest.approx(gib(0.21), abs=gib(0.01)) def test_promotion_cap_growth_preserves_allocator_allowance(): current = gib(6.5) promoted = gib(0.5) target = vb.cap_bytes_preserving_allowance_after_promotion( - current, promoted, gib(9.5) + current, promoted, gib(9.5), gc_threshold=GC_THRESHOLD ) - before = vb.allocator_allowance_bytes(current, gib(5.5)) - after = vb.allocator_allowance_bytes(target, gib(6.0)) + before = vb.allocator_allowance_bytes( + current, gib(5.5), gc_threshold=GC_THRESHOLD + ) + after = vb.allocator_allowance_bytes( + target, gib(6.0), gc_threshold=GC_THRESHOLD + ) assert target > current + promoted assert after >= before def test_cap_for_live_clamped_to_cliff(): - cap = vb.cap_bytes_for_live(gib(11.0), gib(2.0), cliff_cap_bytes=gib(9.85)) + cap = vb.cap_bytes_for_live( + gib(11.0), + gib(2.0), + cliff_cap_bytes=gib(9.85), + gc_threshold=GC_THRESHOLD, + ) assert cap == gib(9.85) @@ -58,12 +79,18 @@ def test_promotion_precheck_uses_the_0p95_divisor(): # need_cap = (6.77 + 0.375 + 0.21) / 0.95 = 7.742 GiB. live, block, slack = gib(6.77), gib(0.375), gib(0.21) # Sampling: cliff ~9.85 has room -> cap lever can fund the block. - assert vb.cap_can_host_promotion(live, block, slack, gib(9.85)) is True + assert vb.cap_can_host_promotion( + live, block, slack, gib(9.85), gc_threshold=GC_THRESHOLD + ) is True # Training-like: cap pinned at the 7.35 floor/cliff -> must demote instead. - assert vb.cap_can_host_promotion(live, block, slack, gib(7.35)) is False + assert vb.cap_can_host_promotion( + live, block, slack, gib(7.35), gc_threshold=GC_THRESHOLD + ) is False # The naive (no /0.95) test would wrongly pass at cliff = 7.36 # (6.77+0.375+0.21 = 7.355 < 7.36); the real need_cap 7.742 rejects it. - assert vb.cap_can_host_promotion(live, block, slack, gib(7.36)) is False + assert vb.cap_can_host_promotion( + live, block, slack, gib(7.36), gc_threshold=GC_THRESHOLD + ) is False def test_promote_gate_needs_zero_retries_and_a_block_of_slack(): diff --git a/toolkit/memory_management/arena_offload/cap_calibrator.py b/toolkit/memory_management/arena_offload/cap_calibrator.py index 7990f171aa..3b3685faa0 100644 --- a/toolkit/memory_management/arena_offload/cap_calibrator.py +++ b/toolkit/memory_management/arena_offload/cap_calibrator.py @@ -61,10 +61,21 @@ def __init__( enabled=False, notch_bytes=DEFAULT_CAP_NOTCH_BYTES, monitor_settled=False, + gc_threshold=None, ): self.enabled = bool(enabled) self.monitor_settled = bool(monitor_settled) self.notch_bytes = max(1, int(notch_bytes)) + self.gc_threshold = ( + vram_budget.allocator_gc_threshold() + if gc_threshold is None + else float(gc_threshold) + ) + if not 0.0 < self.gc_threshold <= 1.0: + raise ValueError( + "effective allocator GC threshold must be in (0, 1], " + f"got {self.gc_threshold}" + ) self.state = CAP_WARMUP self.bucket_profiles: dict[tuple, BucketCapProfile] = {} self.probe_cap_bytes: int | None = None @@ -286,7 +297,8 @@ def _begin_initial_probe( ): worst_live = self._worst_live_bytes(resident_bytes, ring_bytes) predicted = _ceil_to_notch( - int(worst_live / vram_budget.GC_THRESHOLD), self.notch_bytes + int(worst_live / self.gc_threshold), + self.notch_bytes, ) + self.notch_bytes candidate = min(int(cliff_cap_bytes), predicted) self.predicted_initial_cap_bytes = candidate @@ -297,7 +309,9 @@ def _begin_initial_probe( self.learned_cache_pad_bytes = max( 0, vram_budget.allocator_allowance_bytes( - self.settled_cap_bytes, worst_live + self.settled_cap_bytes, + worst_live, + gc_threshold=self.gc_threshold, ), ) return self._decision( @@ -318,11 +332,17 @@ def _advance_clean_probe( self.last_clean_cap_bytes = clean_cap worst_live = self._worst_live_bytes(resident_bytes, ring_bytes) self.learned_cache_pad_bytes = max( - 0, vram_budget.allocator_allowance_bytes(clean_cap, worst_live) + 0, + vram_budget.allocator_allowance_bytes( + clean_cap, + worst_live, + gc_threshold=self.gc_threshold, + ), ) next_cap = clean_cap - self.notch_bytes minimum = _ceil_to_notch( - int(worst_live / vram_budget.GC_THRESHOLD), self.notch_bytes + int(worst_live / self.gc_threshold), + self.notch_bytes, ) - self.notch_bytes next_cap = max(self.notch_bytes, minimum, next_cap) if next_cap >= clean_cap: @@ -449,6 +469,7 @@ def diagnostics(self): return { "enabled": self.enabled, "monitor_settled": self.monitor_settled, + "gc_threshold": self.gc_threshold, "state": self.state, "notch_bytes": self.notch_bytes, "probe_cap_bytes": self.probe_cap_bytes, diff --git a/toolkit/memory_management/arena_offload/policy.py b/toolkit/memory_management/arena_offload/policy.py index dfa1dadc09..0975de43af 100644 --- a/toolkit/memory_management/arena_offload/policy.py +++ b/toolkit/memory_management/arena_offload/policy.py @@ -42,8 +42,19 @@ def __init__( allocator_cache_headroom_bytes=( DEFAULT_ALLOCATOR_CACHE_HEADROOM_BYTES ), + gc_threshold=None, ): self.state = vram_budget.ResidencyFsmState() + self.gc_threshold = ( + vram_budget.allocator_gc_threshold() + if gc_threshold is None + else float(gc_threshold) + ) + if not 0.0 < self.gc_threshold <= 1.0: + raise ValueError( + "effective allocator GC threshold must be in (0, 1], " + f"got {self.gc_threshold}" + ) self.allocator_cache_headroom_bytes = max( 0, int(allocator_cache_headroom_bytes) ) @@ -110,6 +121,7 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, block_bytes, cache_pad, int(cliff_cap_bytes), + gc_threshold=self.gc_threshold, ) ) promote_ok = ( @@ -143,6 +155,7 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, int(worst_shape_live_bytes) + block_bytes, cache_pad, int(cliff_cap_bytes), + gc_threshold=self.gc_threshold, ) cap_covers = ( candidate is not None @@ -153,6 +166,7 @@ def step(self, signal, *, candidate, demote_candidate, cliff_cap_bytes, int(worst_shape_live_bytes), cache_pad, int(cliff_cap_bytes), + gc_threshold=self.gc_threshold, ) needed_cap = ( needed_promotion_cap @@ -407,6 +421,7 @@ def diagnostics(self): "allocator_cache_headroom_bytes": ( self.allocator_cache_headroom_bytes ), + "gc_threshold": self.gc_threshold, } diff --git a/toolkit/memory_management/arena_offload/runtime.py b/toolkit/memory_management/arena_offload/runtime.py index 24ef1ecf10..61d4201dd6 100644 --- a/toolkit/memory_management/arena_offload/runtime.py +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -27,7 +27,7 @@ ResidencyState, ordered_demotion_block_keys, ) -from ..vram_budget import apply_simulated_card +from ..vram_budget import allocator_gc_threshold, apply_simulated_card from .policy import ( ArenaResidencyController, TrainingSignalWindow, @@ -98,7 +98,10 @@ def __init__( self._signals = TrainingSignalWindow() self._last_policy_error: str | None = None self._last_failure_event: dict | None = None - self._policy = ArenaResidencyController() + self._allocator_gc_threshold = allocator_gc_threshold() + self._policy = ArenaResidencyController( + gc_threshold=self._allocator_gc_threshold + ) cap_calibration_requested = bool(config._policy.cap_calibration) cap_calibration_enabled = ( cap_calibration_requested and sys.platform == "win32" @@ -110,7 +113,8 @@ def __init__( stacklevel=2, ) self._cap_calibrator = TrainingCapCalibrator( - enabled=cap_calibration_enabled + enabled=cap_calibration_enabled, + gc_threshold=self._allocator_gc_threshold, ) self._cap_calibration_enabled = cap_calibration_enabled self._last_training_cap_target_bytes: int | None = None @@ -843,6 +847,7 @@ def _reserve_allocator_for_promotion(self, promotion_bytes, plan) -> None: current, int(promotion_bytes), cliff, + gc_threshold=self._allocator_gc_threshold, ) if target <= current: return @@ -925,9 +930,14 @@ def _sampling_profile(self, shape_key, occurrences): and (int(occurrences) > 1 or profile is not None) ) if profile is None and eligible: + gc_threshold = getattr(self, "_allocator_gc_threshold", None) + if gc_threshold is None: + gc_threshold = allocator_gc_threshold() profile = _SamplingCapProfile( calibrator=TrainingCapCalibrator( - enabled=True, monitor_settled=True + enabled=True, + monitor_settled=True, + gc_threshold=gc_threshold, ), signals=TrainingSignalWindow(), ) @@ -1498,8 +1508,13 @@ def _worst_shape_allocator_slack_bytes(self, current_cap_bytes): return 0 from .. import vram_budget + gc_threshold = getattr(self, "_allocator_gc_threshold", None) + if gc_threshold is None: + gc_threshold = allocator_gc_threshold() return vram_budget.allocator_allowance_bytes( - current_cap_bytes, predicted_live + current_cap_bytes, + predicted_live, + gc_threshold=gc_threshold, ) def _worst_shape_live_bytes(self): @@ -1688,6 +1703,9 @@ def diagnostics(self) -> dict: pinned_block_keys = self._arena.pinned_block_keys() selection = getattr(self._executor, "selection", None) state_audit = getattr(selection, "accounting", None) + gc_threshold = getattr(self, "_allocator_gc_threshold", None) + if gc_threshold is None: + gc_threshold = allocator_gc_threshold() return { "backend": "arena", "blocks": self.block_count, @@ -1731,6 +1749,7 @@ def diagnostics(self) -> dict: "strict_vram_cap": bool( getattr(self._config, "strict_vram_cap", False) ), + "allocator_gc_threshold": gc_threshold, "fp8_forward": bool(self._config.fp8_forward), "fp8_backward": bool(self._config.fp8_backward), "fp8_sampling": bool(self._config.fp8_sampling), diff --git a/toolkit/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py index 3676514aef..1b637a5f6e 100644 --- a/toolkit/memory_management/vram_budget.py +++ b/toolkit/memory_management/vram_budget.py @@ -41,13 +41,15 @@ * ``hard`` -- device-free floor the card must keep (WDDM spill guard). * ``margin`` -- planning headroom subtracted from budgets; ``>= hard``. -Everything here is pure (CPU-testable) except ``DeviceSnapshot.capture`` and -``device_free_bytes``. +Everything here is pure (CPU-testable) except the runtime allocator-setting +query, ``DeviceSnapshot.capture``, and ``device_free_bytes``. """ from __future__ import annotations import math +import os +import re from dataclasses import dataclass from typing import Optional @@ -58,6 +60,78 @@ GIB = 1024 ** 3 +def _runtime_allocator_settings() -> str: + """Return the allocator configuration Torch actually initialized.""" + getter = getattr(torch._C, "_accelerator_getAllocatorSettings", None) + if getter is not None: + try: + return str(getter() or "") + except (AttributeError, RuntimeError): + pass + + # Older Torch builds do not expose the runtime settings. Match Torch's + # preferred/legacy environment-variable precedence as a compatibility + # fallback. In supported builds the runtime getter above remains the source + # of truth, so later environment mutations cannot lie about active policy. + for key in ("PYTORCH_ALLOC_CONF", "PYTORCH_CUDA_ALLOC_CONF"): + if key in os.environ: + return os.environ[key] + return "" + + +def _allocator_setting(settings: str, name: str) -> str | None: + matches = re.findall( + rf"(?:^|,)\s*{re.escape(name)}\s*:\s*([^,\s]+)", + str(settings), + ) + return matches[-1] if matches else None + + +def allocator_gc_threshold() -> float: + """Return the effective GC target configured in Torch's live allocator. + + Torch's native allocator disables proactive garbage collection when the + option is absent. For cap/allowance arithmetic that means allocations may + grow to the cap itself, represented here by an effective threshold of 1.0. + The option is also ignored by ``cudaMallocAsync``, with the same result for + this model. + """ + settings = _runtime_allocator_settings() + backend = _allocator_setting(settings, "backend") + if backend is not None and backend.lower() == "cudamallocasync": + return 1.0 + + configured = _allocator_setting(settings, "garbage_collection_threshold") + if configured is None: + return 1.0 + try: + threshold = float(configured) + except (TypeError, ValueError) as error: + raise ValueError( + "invalid runtime garbage_collection_threshold: " + f"{configured!r}" + ) from error + if not 0.0 < threshold < 1.0: + raise ValueError( + "runtime garbage_collection_threshold must be between 0 and 1, " + f"got {threshold}" + ) + return threshold + + +def _resolve_gc_threshold(gc_threshold) -> float: + threshold = ( + allocator_gc_threshold() + if gc_threshold is None + else float(gc_threshold) + ) + if not 0.0 < threshold <= 1.0: + raise ValueError( + f"effective allocator GC threshold must be in (0, 1], got {threshold}" + ) + return threshold + + def reconcile_free_bytes(driver_free_bytes, physical_free_bytes) -> int: """Pick the governing device-free value (pure; CPU-testable). @@ -413,7 +487,7 @@ def sampling_allocator_budget_free_bytes( cap_fraction, hard_bytes, *, - gc_threshold=0.95, + gc_threshold=None, ): """Allocated-side equivalent of driver-free for the sampling planner (pure). @@ -434,6 +508,7 @@ def sampling_allocator_budget_free_bytes( """ if cap_fraction is None: return None + gc_threshold = _resolve_gc_threshold(gc_threshold) cap_bytes = float(cap_fraction) * float(max(1, int(total_bytes))) return int( float(gc_threshold) * cap_bytes @@ -574,7 +649,8 @@ def estimate_training_working_reserve_bytes( # --------------------------------------------------------------------------- # Two-timescale residency control # -# Allowance lives in *target-space* (0.95*cap - live); the allocator cap is set +# Allowance lives in *target-space* (gc_threshold*cap - live); the allocator +# cap is set # in *cap-space*. The two differ by the gc_threshold factor: a cap raise of ``d`` # only adds ``gc_threshold * d`` of GC target / allowance. Every conversion below # carries the ``/ gc_threshold`` so no call site open-codes it (that missing @@ -586,10 +662,7 @@ def estimate_training_working_reserve_bytes( # counters that lag its move by one window (the FSM's verify phases absorb this). # --------------------------------------------------------------------------- -GC_THRESHOLD = 0.95 - - -def allocator_allowance_bytes(cap_bytes, live_bytes, *, gc_threshold=GC_THRESHOLD) -> int: +def allocator_allowance_bytes(cap_bytes, live_bytes, *, gc_threshold=None) -> int: """Idle-cache allowance under the cap: ``gc_threshold*cap - live`` (pure). The caching allocator sweeps idle segments when reserved would cross the GC @@ -600,6 +673,7 @@ def allocator_allowance_bytes(cap_bytes, live_bytes, *, gc_threshold=GC_THRESHOL every sweep dumps all cache and every reuse re-mallocs (self-sustaining thrash), so callers must keep this positive at the live peak. """ + gc_threshold = _resolve_gc_threshold(gc_threshold) return int(float(gc_threshold) * float(max(0, int(cap_bytes))) - float(max(0, int(live_bytes)))) @@ -609,7 +683,7 @@ def cap_bytes_for_live( cliff_cap_bytes, *, floor_cap_bytes=0, - gc_threshold=GC_THRESHOLD, + gc_threshold=None, ) -> int: """Cap that hosts ``planned_live`` plus an idle-cache budget (pure). @@ -619,6 +693,7 @@ def cap_bytes_for_live( Clamped to the WDDM cliff bound above (never license silent paging; see :func:`cap_fraction`) and an optional floor below. """ + gc_threshold = _resolve_gc_threshold(gc_threshold) want = (float(max(0, int(planned_live_bytes))) + float(max(0, int(cache_budget_bytes)))) / float(gc_threshold) want = min(want, float(int(cliff_cap_bytes))) want = max(want, float(max(0, int(floor_cap_bytes)))) @@ -630,7 +705,7 @@ def cap_bytes_preserving_allowance_after_promotion( promotion_bytes, cliff_cap_bytes, *, - gc_threshold=GC_THRESHOLD, + gc_threshold=None, ) -> int: """Raise a cap enough that resident growth does not consume GC allowance.""" current = max(0, int(cap_bytes)) @@ -638,6 +713,7 @@ def cap_bytes_preserving_allowance_after_promotion( cliff = max(0, int(cliff_cap_bytes)) if promoted <= 0 or current >= cliff: return min(current, cliff) + gc_threshold = _resolve_gc_threshold(gc_threshold) growth = math.ceil(float(promoted) / float(gc_threshold)) return min(cliff, current + int(growth)) @@ -648,7 +724,7 @@ def cap_can_host_promotion( allocator_cache_headroom_bytes, cliff_cap_bytes, *, - gc_threshold=GC_THRESHOLD, + gc_threshold=None, ) -> bool: """Can the cheap cap lever (tier 1) absorb one more resident block? (pure). @@ -662,8 +738,9 @@ def cap_can_host_promotion( expensive resident demote (tier 2). The ``/ gc_threshold`` is load-bearing: a naive ``cliff - cap >= block`` test - under-reserves by the 0.95 factor. + under-reserves whenever allocator GC is configured below the cap. """ + gc_threshold = _resolve_gc_threshold(gc_threshold) need_cap = ( float(max(0, int(live_bytes))) + float(max(0, int(block_bytes))) @@ -684,7 +761,7 @@ def residency_promote_ok( the telemetry proves the room is really there -- * ``num_alloc_retries == 0`` over the window (nothing cap-binding), AND - * worst-shape allocator slack (``0.95 * cap - predicted_live``) + * worst-shape allocator slack (``gc_threshold * cap - predicted_live``) exceeds one block plus the pad, so the promotion still leaves ``allocator_cache_headroom`` of reusable-cache allowance. From 5f1f8ace9d24998b843228399d3fc4b905f60e86 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ryd=C3=A9n=20Johan?= Date: Thu, 16 Jul 2026 16:05:26 +0200 Subject: [PATCH 20/20] Ignore canceled timer measurements --- tests/test_timer.py | 35 +++++++++++++++++++++++++++++++++++ toolkit/timer.py | 6 ++++-- 2 files changed, 39 insertions(+), 2 deletions(-) create mode 100644 tests/test_timer.py diff --git a/tests/test_timer.py b/tests/test_timer.py new file mode 100644 index 0000000000..7334a280ea --- /dev/null +++ b/tests/test_timer.py @@ -0,0 +1,35 @@ +import pytest + +import toolkit.timer as timer_module +from toolkit.timer import Timer + + +def test_timer_print_ignores_canceled_empty_measurements(capsys): + timer = Timer("recovered step") + reports = [] + timer.add_after_print_hook(reports.append) + + with pytest.raises(RuntimeError, match="recoverable failure"): + with timer("failed operation"): + raise RuntimeError("recoverable failure") + + timer.print() + + assert reports == [{}] + assert "failed operation" not in capsys.readouterr().out + + +def test_timer_print_still_reports_completed_measurements(monkeypatch, capsys): + clock = iter((10.0, 12.5)) + monkeypatch.setattr(timer_module.time, "time", lambda: next(clock)) + timer = Timer("completed step") + reports = [] + timer.add_after_print_hook(reports.append) + + with timer("completed operation"): + pass + + timer.print() + + assert reports == [{"completed operation": 2.5}] + assert "2.5000s avg - completed operation" in capsys.readouterr().out diff --git a/toolkit/timer.py b/toolkit/timer.py index e849ba5faa..61920b057d 100644 --- a/toolkit/timer.py +++ b/toolkit/timer.py @@ -1,11 +1,11 @@ -import time from collections import OrderedDict, deque -import sys import os +import time # check if is ui process will have IS_AI_TOOLKIT_UI in env is_ui = os.environ.get("IS_AI_TOOLKIT_UI", "0") == "1" + class Timer: def __init__(self, name='Timer', max_buffer=10): self.name = name @@ -48,6 +48,8 @@ def print(self): timing_dict = {} # sort by longest at top for timer_name, timings in sorted(self.timers.items(), key=lambda x: sum(x[1]), reverse=True): + if not timings: + continue avg_time = sum(timings) / len(timings) if not is_ui: