diff --git a/docs/ARENA_OFFLOAD_CONTRACT.md b/docs/ARENA_OFFLOAD_CONTRACT.md new file mode 100644 index 000000000..352aa0be5 --- /dev/null +++ b/docs/ARENA_OFFLOAD_CONTRACT.md @@ -0,0 +1,128 @@ +# 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. + +## 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` + 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. + 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 + +- 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. + +## 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 + 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, + 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. + +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 +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 6d19c649d..14aee62af 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 5b339ebe6..0693928b3 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 @@ -66,6 +66,17 @@ 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 = [] + 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}") + failed.append(process) + else: + process.job = None + self.process = list(reversed(failed)) + if errors: + raise RuntimeError("job cleanup failed: " + "; ".join(errors)) diff --git a/jobs/process/BaseProcess.py b/jobs/process/BaseProcess.py index c58724c98..be95ded30 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 beb36865c..622a95d26 100644 --- a/jobs/process/BaseSDTrainProcess.py +++ b/jobs/process/BaseSDTrainProcess.py @@ -1,3 +1,4 @@ +import contextlib import copy import glob import inspect @@ -21,9 +22,15 @@ 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 +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 @@ -196,6 +203,9 @@ 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_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 @@ -361,8 +371,18 @@ 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( + gen_configs=gen_img_config_list, + ) + 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 +501,57 @@ def clean_up_saves(self): def post_save_hook(self, save_path): # override in subclass pass + + def cleanup(self): + if self._cleanup_completed: + return + if self._cleanup_in_progress: + raise RuntimeError("training process cleanup is already in progress") + self._cleanup_in_progress = True + errors = [] + + def attempt(label, operation): + try: + operation() + except Exception as error: + errors.append(f"{label}: {type(error).__name__}: {error}") + + 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}" + ) + 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 @@ -734,7 +805,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: @@ -1579,6 +1656,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 @@ -1593,6 +1684,20 @@ 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, + 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'): sampler = ModelClass.get_train_scheduler() @@ -1628,10 +1733,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 +1853,40 @@ 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, + 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, + training_working_reserve_hint_bytes=( + training_reserve_hint + ), + ), + ) + 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 +2141,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 @@ -2092,7 +2257,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 @@ -2108,7 +2276,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 @@ -2120,6 +2288,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 +2304,12 @@ def run(self): # ==================================================== # BLOCK COMPILE # ==================================================== - if block_compile: + if runtime_owns_block_compile: + print_acc( + "Arena offload owns block compilation; " + "skipping trainer 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 +2623,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: @@ -2458,16 +2647,37 @@ def run(self): else: raise # not an OOM; surface real errors if did_oom: - self.num_consecutive_oom += 1 - if self.num_consecutive_oom > 3: - raise RuntimeError("OOM during training step 3 times in a row, aborting training") + 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 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: @@ -2657,6 +2867,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/requirements_base.txt b/requirements_base.txt index afbff335f..0c0472491 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/run.py b/run.py index 9238a8349..66fda7b58 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() @@ -108,29 +109,59 @@ 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 + failure = None + cleanup_failed = False try: job = get_job(config_file, args.name) job.run() - job.cleanup() - 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: + 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 000000000..1923a411d --- /dev/null +++ b/tests/test_allocator_cap.py @@ -0,0 +1,149 @@ +import pytest +from unittest import mock + +from toolkit.memory_management import allocator_cap + + +@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_binds_the_same_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=False + ) + + 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): + 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 + + +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) + 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_allocator_gc_threshold.py b/tests/test_allocator_gc_threshold.py new file mode 100644 index 000000000..cb8db1fe8 --- /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_canonical_transaction.py b/tests/test_arena_canonical_transaction.py new file mode 100644 index 000000000..b3215127c --- /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() + import toolkit.memory_management.arena_offload.construction as 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_cap_calibrator.py b/tests/test_arena_cap_calibrator.py new file mode 100644 index 000000000..1682b9b99 --- /dev/null +++ b/tests/test_arena_cap_calibrator.py @@ -0,0 +1,178 @@ +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, + CAP_RESTORE_VERIFY, + CAP_SET, + CAP_SETTLED, + 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) + + +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 = make_calibrator(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 = make_calibrator(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 = make_calibrator(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 = make_calibrator(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 = make_calibrator(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 = make_calibrator(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 + + +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_compile_cache.py b/tests/test_arena_compile_cache.py new file mode 100644 index 000000000..dd6384175 --- /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_arena_lifecycle_contract.py b/tests/test_arena_lifecycle_contract.py new file mode 100644 index 000000000..344f82064 --- /dev/null +++ b/tests/test_arena_lifecycle_contract.py @@ -0,0 +1,716 @@ +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, + prepare_canonical_storage, + prepare_arena_offload, +) +from toolkit.memory_management.arena_offload.errors import ( + ArenaCleanupError, + ArenaSetupFatalError, + 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 +from toolkit.memory_management import pin_manager +from toolkit.memory_management import vram_budget +from toolkit.quantization.fp8_linear 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 +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): + 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_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(RuntimeError, match="residency construction failed") as caught: + prepare_arena_offload( + model, + device="cpu", + block_names=("blocks",), + config=ArenaOffloadConfig(enabled=True), + ) + + assert caught.value is original_error + assert active_process_owner() is None + assert model.weight is original + 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", + ( + "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(), + "fits": True, + } + 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_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() + 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")) + executor_close = mock.Mock(side_effect=(RuntimeError("executor boom"), None)) + resources.executor = SimpleNamespace( + 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") + 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_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_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 + ) + + 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) + ) + + 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] + imports = _imported_modules( + root / "toolkit" / "memory_management" / "manager.py" + ) + assert not any("arena_offload" in name for name in imports) diff --git a/tests/test_arena_load_session.py b/tests/test_arena_load_session.py new file mode 100644 index 000000000..497844bcc --- /dev/null +++ b/tests/test_arena_load_session.py @@ -0,0 +1,163 @@ +from types import SimpleNamespace +from unittest import mock + +import pytest +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.arena_offload.construction import ( + CanonicalBuildError, + CanonicalStateConsumedError, + PreparedCanonicalBuild, +) +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) + + +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 new file mode 100644 index 000000000..33c74343b --- /dev/null +++ b/tests/test_arena_offload_api.py @@ -0,0 +1,390 @@ +"""Public Arena facade and generic memory-runtime seam coverage.""" + +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, + validate_arena_training_mode, +) +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, +) + +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) + + +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 = True + compile_sample = True + train_compile_blocks = False + layer_offloading_smart_working_reserve_gb = -1.0 + 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 + 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_physical_vram_headroom_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) + 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)) + 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) + + 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_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()) + + 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.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("_")} + self.assertEqual( + public, + { + "enabled", + "fp8_forward", + "fp8_backward", + "fp8_sampling", + "compile_blocks", + "strict_vram_cap", + }, + ) + + 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_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: + compile = False + compile_sample = True + train_compile_blocks = True + + config = ArenaOffloadConfig.from_model_config(DeadAliases()) + self.assertFalse(config.compile_blocks) + + 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 000000000..4233c7d5a --- /dev/null +++ b/tests/test_arena_offload_planner.py @@ -0,0 +1,137 @@ +from types import SimpleNamespace +from unittest import mock + +from toolkit.memory_management.arena_offload.planner import ( + GIB, + build_training_plan, + impossible_training_plan_message, +) + + +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, + physical_vram_headroom_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 + + +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 new file mode 100644 index 000000000..2e0edf1df --- /dev/null +++ b/tests/test_arena_offload_policy.py @@ -0,0 +1,844 @@ +import contextlib +import json +from types import SimpleNamespace + +import pytest + +from toolkit.memory_management.arena_offload.cap_calibrator import ( + CAP_SET, + CapCalibrationDecision, +) +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_pressure_relief_demotes_enough_live_blocks_without_emptying_cache(monkeypatch): + snapshots = iter( + [ + { + "device_free_bytes": 100, + "predicted_peak_free_bytes": 200, + "deficit_bytes": 300, + }, + { + "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 empty_cache_calls == [] + 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( + 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 + 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 + + +def test_aggressive_capacity_counts_exact_smallest_blocks_under_both_budgets(): + runtime = ArenaOffloadRuntime.__new__(ArenaOffloadRuntime) + 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_physical_headroom_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_guard_mode(monkeypatch): + calls = [] + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime." + "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( + 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, + { + "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_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) + 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( + allocator_cache_headroom_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_cache_headroom" + + 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( + allocator_cache_headroom_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( + allocator_cache_headroom_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_raises_cap_by_fixed_fsm_increment(): + controller = ArenaResidencyController( + allocator_cache_headroom_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_controller_exactly_prefunds_promotion_after_cap_calibration(): + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10, + gc_threshold=0.8, + ) + 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.8) + + +def test_controller_holds_when_exact_promotion_cap_exceeds_cliff(): + controller = ArenaResidencyController( + allocator_cache_headroom_bytes=10, + gc_threshold=0.95, + ) + 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() + 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_physical_vram_headroom_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.configure_wddm_allocator_guard", + lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + "toolkit.memory_management.arena_offload.runtime." + "resolve_physical_vram_headroom_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_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.configure_wddm_allocator_guard", + 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["recoverable"] is True + 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 **_kwargs: 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 **_kwargs: 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_arena_sampling_cap.py b/tests/test_arena_sampling_cap.py new file mode 100644 index 000000000..eebc30d3c --- /dev/null +++ b/tests/test_arena_sampling_cap.py @@ -0,0 +1,171 @@ +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, gc_threshold=0.95 + ), + 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_canonical_arena.py b/tests/test_canonical_arena.py new file mode 100644 index 000000000..4dd4cbb9e --- /dev/null +++ b/tests/test_canonical_arena.py @@ -0,0 +1,221 @@ +import io +import unittest +from types import SimpleNamespace +from unittest import mock + +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 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) + 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_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: + 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) + 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() + + 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_movement_routes_all_entry_points_to_runtime(self): + model = nn.Sequential(_linear()) + 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.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"),), + ) + 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")) + + 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_dxgi_meminfo.py b/tests/test_dxgi_meminfo.py new file mode 100644 index 000000000..4cbd82c94 --- /dev/null +++ b/tests/test_dxgi_meminfo.py @@ -0,0 +1,151 @@ +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 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) + +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 000000000..3cbe5a993 --- /dev/null +++ b/tests/test_generic_block_dispatcher.py @@ -0,0 +1,808 @@ +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, + prepare_arena_offload, +) +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, + _is_dynamo_compile_failure, + _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 + + +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, *, qtype="float8"): + 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(qtype)) + if qtype in ("float8", "qfloat8"): + freeze(model) + model.requires_grad_(False) + 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, + physical_vram_headroom_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_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 + model.blocks[1].proj.weight = shared + with pytest.raises(BlockDiscoveryError, match="shared_managed"): + 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 + 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_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_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() + 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_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, + ) + 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(): + import toolkit.memory_management.arena_offload.transfer as 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, + _compile_fullgraph=True, + ) + config = replace( + config, + _policy=replace( + config._policy, + working_reserve_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", + # Force a mixed plan: two blocks fit, all three do not. + return_value=(700, 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() + 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"] + 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] + 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"] + planned = runtime.diagnostics()["accounting"][ + "planned_training_h2d_bytes" + ] + 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() + 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 + 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) + 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): + 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) + 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. + # 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 + + +@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 + + 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_job_cleanup.py b/tests/test_job_cleanup.py new file mode 100644 index 000000000..172fc4778 --- /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/tests/test_lora_compile_scalars.py b/tests/test_lora_compile_scalars.py new file mode 100644 index 000000000..aa261c1e7 --- /dev/null +++ b/tests/test_lora_compile_scalars.py @@ -0,0 +1,119 @@ +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_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/tests/test_nvml_free_signal.py b/tests/test_nvml_free_signal.py new file mode 100644 index 000000000..46686c5a4 --- /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 000000000..4c5978958 --- /dev/null +++ b/tests/test_pin_manager.py @@ -0,0 +1,137 @@ +import os +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() + + def tearDown(self): + pin_manager.reset_for_tests() + + 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): + 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): + # 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_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) + + 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. + + 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/canonical_arena.py", + "toolkit/memory_management/manager.py", + "toolkit/memory_management/manager_modules.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\(\)") + + 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/tests/test_residency.py b/tests/test_residency.py new file mode 100644 index 000000000..c966460b6 --- /dev/null +++ b/tests/test_residency.py @@ -0,0 +1,270 @@ +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, + ordered_demotion_block_keys, + pin_requirements_for_plan, +) + + + + +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_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_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 000000000..5988f3217 --- /dev/null +++ b/tests/test_residency_two_timescale.py @@ -0,0 +1,270 @@ +"""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 +GC_THRESHOLD = 0.95 + + +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), 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), + 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), 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), gc_threshold=GC_THRESHOLD + ) + + 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), + gc_threshold=GC_THRESHOLD, + ) + 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), 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), 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), gc_threshold=GC_THRESHOLD + ) 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 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): + 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_timer.py b/tests/test_timer.py new file mode 100644 index 000000000..7334a280e --- /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/tests/test_torchao_compat.py b/tests/test_torchao_compat.py new file mode 100644 index 000000000..47a5d70ac --- /dev/null +++ b/tests/test_torchao_compat.py @@ -0,0 +1,43 @@ +import torch + +import toolkit.quantization.torchao_compat as torchao_compat +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 + + +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 diff --git a/tests/test_transfer_plan.py b/tests/test_transfer_plan.py new file mode 100644 index 000000000..5d7358019 --- /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 000000000..7834b91f9 --- /dev/null +++ b/tests/test_transfer_runtime.py @@ -0,0 +1,165 @@ +import pytest +import torch +import torch.nn.functional as F + +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 + +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="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/compile_utils.py b/toolkit/compile_utils.py new file mode 100644 index 000000000..e83d5cfa6 --- /dev/null +++ b/toolkit/compile_utils.py @@ -0,0 +1,14 @@ +import os + +import torch + + +def configure_cuda_only_inductor() -> None: + """Configure CUDA compilation without requiring a CPU toolchain.""" + torch._dynamo.config.suppress_errors = False + 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 ef5900226..7cbe06c5f 100644 --- a/toolkit/config_modules.py +++ b/toolkit/config_modules.py @@ -699,6 +699,51 @@ 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_physical_vram_headroom_gb = kwargs.get( + "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 + ) + 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 + ) + self.layer_offloading_smart_sampling_physical_vram_headroom_gb = kwargs.get( + "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 + ) + 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 + ) + 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/lora_special.py b/toolkit/lora_special.py index bf8309db8..86994dbc4 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 8bafb6d92..1bce9dfee 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 new file mode 100644 index 000000000..b27517f05 --- /dev/null +++ b/toolkit/memory_management/allocator_cap.py @@ -0,0 +1,189 @@ +"""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] = {} +CAP_RELIEF_BYTES = int(0.5 * GIB) + + +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) + RELIEF_BYTES.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: + """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 configure_wddm_allocator_guard( + device, + wddm_hard_gib=None, + *, + target_cap_bytes=None, + strict=False, + log_prefix="[MemoryManager]", + force=False, +): + """Bind the allocator below the WDDM cliff in every guard mode. + + ``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 + 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 not force and 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 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} 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 new file mode 100644 index 000000000..cc728cff0 --- /dev/null +++ b/toolkit/memory_management/arena_offload/__init__.py @@ -0,0 +1,67 @@ +"""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. +""" + +__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", +] + +_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/arena_offload/api.py b/toolkit/memory_management/arena_offload/api.py new file mode 100644 index 000000000..b65a6375c --- /dev/null +++ b/toolkit/memory_management/arena_offload/api.py @@ -0,0 +1,438 @@ +"""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 toolkit.quantization.torchao_compat import ( + TORCHAO_ARENA_FP8_MIN_VERSION, + TORCHAO_VERSION, + torchao_arena_fp8_supported, +) + +from ..runtime import ( + close_memory_runtime, + get_memory_runtime, + unwrap_memory_model, +) +from .runtime import ArenaOffloadRuntime + +GIB = 1024**3 +_FP8_QTYPES = ("qfloat8", "float8") + + +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, + mutates_base_weights=False, + train_text_encoder=False, + unload_text_encoder=True, +) -> 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" + ) + 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): + """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 + 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 + + sampling_working_reserve_gib: float | None = None + sampling_physical_vram_headroom_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 + 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 + # 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: + return getattr(model_config, name, 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 = [] + torchao_fp8_unavailable = bool( + fp8_weights + and get("qtype") == "float8" + and not torchao_arena_fp8_supported() + ) + 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 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: " + + ", ".join(ignored), + RuntimeWarning, + stacklevel=2, + ) + + return cls( + enabled=bool( + get("layer_offloading", False) + and get("layer_offloading_smart", False) + ), + 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 + # both phases. + compile_blocks=bool(get("compile", False)), + strict_vram_cap=bool( + get("layer_offloading_strict_vram_cap", False) + ), + _compile_dynamic=( + None + 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 ()) + ), + _simulated_vram_gib=( + float(get("layer_offloading_simulated_vram_gb") or 0.0) or None + ), + _policy=_ArenaPolicyOptions( + working_reserve_gib=working_reserve_gib, + physical_vram_headroom_gib=get( + "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) + ), + prefetch_depth=int(get("layer_offloading_prefetch_depth", 3) or 3), + sampling_working_reserve_gib=get( + "layer_offloading_smart_sampling_working_reserve_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 + ), + ), + ) + + + +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, + pin_on_finish=False, + ) + 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/cap_calibrator.py b/toolkit/memory_management/arena_offload/cap_calibrator.py new file mode 100644 index 000000000..3b3685faa --- /dev/null +++ b/toolkit/memory_management/arena_offload/cap_calibrator.py @@ -0,0 +1,506 @@ +"""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, + 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 + 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 / self.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, + gc_threshold=self.gc_threshold, + ), + ) + 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, + gc_threshold=self.gc_threshold, + ), + ) + next_cap = clean_cap - self.notch_bytes + minimum = _ceil_to_notch( + 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: + 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, + "gc_threshold": self.gc_threshold, + "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/compile_cache.py b/toolkit/memory_management/arena_offload/compile_cache.py new file mode 100644 index 000000000..dfe4b2740 --- /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/construction.py b/toolkit/memory_management/arena_offload/construction.py new file mode 100644 index 000000000..cc1056f90 --- /dev/null +++ b/toolkit/memory_management/arena_offload/construction.py @@ -0,0 +1,638 @@ +"""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, + substitution_views, + typed_view, +) + + +class CanonicalBuildError(RuntimeError): + pass + + +class CanonicalStateConsumedError(CanonicalBuildError): + """A direct-load source mapping was mutated before construction failed.""" + + +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", + 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 = {} + 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 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: + """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: + if not self.pin_on_finish: + continue + 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, + bool(block.handle and block.handle.pinned), + pin_handle=block.handle, + ) + 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 + if b.handle is not None and b.handle.pinned + ), + ) + 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) + 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) + self.arena._blocks.pop(block.key, None) + block.handle = None + self.arena._canonicalized = bool(self.arena._blocks) + self._populated = False + resources = getattr(self, "_arena_resources", None) + 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 new file mode 100644 index 000000000..53697cb8d --- /dev/null +++ b/toolkit/memory_management/arena_offload/discovery.py @@ -0,0 +1,263 @@ +"""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()._cdata, + ) + 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: + 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 000000000..25e72b75f --- /dev/null +++ b/toolkit/memory_management/arena_offload/dispatcher.py @@ -0,0 +1,466 @@ +"""Generic saved-forward dispatcher over canonical arena storage.""" + +from __future__ import annotations + +from dataclasses import dataclass + +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, + ImmutableTransformerRuntime, + build_program_fingerprint, +) +from toolkit.memory_management.arena_offload.transfer import ( + configure_fetch_runtime, + free_on_backward, +) + + +DISPATCHER_GENERATION = "generic-block-dispatcher-v2-fullgraph" + + +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: + return int(torch._C._current_graph_task_id()) >= 0 + except (AttributeError, RuntimeError): + 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.""" + + 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_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, + 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 = () + 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] + 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: + configure_cuda_only_inductor() + kernel = torch.compile( + kernel, + mode="default", + fullgraph=self.compile_fullgraph, + 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" + ) + 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: + 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 + 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 + 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) + 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 + ) + + leaf_args = source.assemble_leaf_args(self.residency, compact_flat) + self._mark_dispatch_dynamic(first) + 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 = ( + 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. + 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 + # 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) + ) + 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): + 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 = () + self._sampling_forward_begin = None + self._sampling_forward_end = None + self._sampling_allocation_failure = None + self._strict_kernels_executed.clear() + super().close() + + +def prepare_block_dispatcher_runtime( + transformer, + residency, + *, + selection, + depth=3, + compile_blocks=True, + compile_dynamic=True, + compile_fullgraph=False, + 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_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/errors.py b/toolkit/memory_management/arena_offload/errors.py new file mode 100644 index 000000000..84da9d292 --- /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 000000000..1ebd9ab22 --- /dev/null +++ b/toolkit/memory_management/arena_offload/fp8.py @@ -0,0 +1,112 @@ +"""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, +) + + +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 000000000..effa0ad48 --- /dev/null +++ b/toolkit/memory_management/arena_offload/layout.py @@ -0,0 +1,189 @@ +"""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 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 + pin_handle: object | None = None + + +@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 release_pack(pack: "BlockPack | None") -> None: + """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 + + +@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 + + +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 000000000..0b94e1086 --- /dev/null +++ b/toolkit/memory_management/arena_offload/load_session.py @@ -0,0 +1,186 @@ +"""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, + CanonicalStateConsumedError, + 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 CanonicalStateConsumedError: + raise + 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 000000000..33f81af6f --- /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 000000000..ecb5c8f66 --- /dev/null +++ b/toolkit/memory_management/arena_offload/planner.py @@ -0,0 +1,279 @@ +"""Cold-start, whole-block residency planning owned by arena offload.""" + +from __future__ import annotations + +from types import SimpleNamespace + +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. 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_physical_vram_headroom_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 = 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]: + 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 _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_or_build, 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) + 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) + 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 = _planning_records(arena_or_build) + 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) - physical_headroom_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) - 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 + ) + 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)) + + 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 + pinned_resident_bytes, + "must_resident_layer_keys": set(), + "pinned_resident_bytes": pinned_resident_bytes, + "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, + "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": 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, + "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 000000000..0975de43a --- /dev/null +++ b/toolkit/memory_management/arena_offload/policy.py @@ -0,0 +1,630 @@ +"""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_ALLOCATOR_CACHE_HEADROOM_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, + *, + 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) + ) + 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_physical_headroom_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.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 + 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, + worst_shape_live_bytes=None, learned_cache_pad_bytes=None): + 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 + ) + 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 + 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), + gc_threshold=self.gc_threshold, + ) + ) + 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) + ) + 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), + gc_threshold=self.gc_threshold, + ) + 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), + gc_threshold=self.gc_threshold, + ) + 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)) + aggressive_ok = ( + candidate is not None + and aggressive_capacity >= AGGRESSIVE_PROMOTION_MIN_CAPACITY + and not bool(signal.get("compile_invalid")) + and retries == 0 + and device_frees == 0 + and worst_ok + and allocator_slack > block_bytes + cache_pad + ) + self.last_worst_shape_physical_headroom_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) + 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) + ): + 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 ( + "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 not exact_funding + and allocator_slack <= block_bytes + cache_pad + 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_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 record_physical_pressure_relief(self, block_keys, block_bytes): + """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 + 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 + 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_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, + "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, + "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": ( + self.last_rejected_residency_bytes + ), + "allocator_cache_headroom_bytes": ( + self.allocator_cache_headroom_bytes + ), + "gc_threshold": self.gc_threshold, + } + + +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 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 + + 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 000000000..2f9bad103 --- /dev/null +++ b/toolkit/memory_management/arena_offload/resources.py @@ -0,0 +1,252 @@ +"""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._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") + } + 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) + try: + from .. import allocator_cap, vram_budget + 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() + 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 + 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, 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: + 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: + failures.append(("transfer tickets", error)) + if self.executor is not None: + attempt( + "immutable executor", + self.executor.close, + lambda: setattr(self, "executor", None), + ) + retained_fp8_restores = [] + for label, restore in reversed(self.fp8_restores): + 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, + lambda: setattr(self, "residency", None), + ) + + retained_attributes = [] + for owner, name, value in reversed(self.published_attributes): + if getattr(owner, name, None) is value: + if not attempt( + f"published attribute {name}", + lambda owner=owner, name=name: delattr(owner, name), + ): + 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( + "canonical arena", + arena.release, + lambda: setattr(self, "arena", None), + ) + self.disposed = True + self._install_disposed_movement_guard() + elif self.canonical_build 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: + failures.append(("transfer runtime", error)) + else: + self._transfer_runtime_released = True + + if self._restore_fp8_policy: + from toolkit.quantization.fp8_linear 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 + self._releasing = False + if self.runtime is not None: + self.runtime._closed = True + self.runtime._disposed = bool(self.canonical_committed) + 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 new file mode 100644 index 000000000..61d4201dd --- /dev/null +++ b/toolkit/memory_management/arena_offload/runtime.py @@ -0,0 +1,2055 @@ +"""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 sys +import time +import warnings +from collections import Counter +from collections.abc import Sequence +from dataclasses import dataclass, replace +from typing import Any + +from .. import allocator_cap +from ..canonical_arena import CanonicalArena +from ..residency import ( + ResidencyPlan, + ResidencyState, + ordered_demotion_block_keys, +) +from ..vram_budget import allocator_gc_threshold, apply_simulated_card +from .policy import ( + ArenaResidencyController, + 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 +from toolkit.quantization.fp8_linear import set_fp8_grad_input_enabled +from .planner import ( + build_training_plan, + impossible_training_plan_message, + resolve_physical_vram_headroom_gib, +) +from .ownership import normalize_device +from .resources import ArenaRuntimeResources + +RUNTIME_ATTR = "_arena_offload_runtime" + +GIB = 1024**3 + + +@dataclass +class _SamplingCapProfile: + calibrator: TrainingCapCalibrator + signals: TrainingSignalWindow + occurrences: int = 0 + forward_count: int = 0 + + +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._compile_cache = None + self._closed = False + self._disposed = False + + # 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 + self._signals = TrainingSignalWindow() + self._last_policy_error: str | None = None + self._last_failure_event: dict | None = None + 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" + ) + 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, + gc_threshold=self._allocator_gc_threshold, + ) + 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._training_fp8_restores = [] + self._training_fp8_canonical = 0 + 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 + + # ------------------------------------------------------------------ + # 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.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) + + 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, + pin_on_finish=False, + ) + 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_modules = tuple( + child for entries in entries_by_block.values() for _name, child in entries + ) + smart_plan = build_training_plan( + transformer, + 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( + 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_fullgraph=config._compile_fullgraph, + 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, + ) + executor.reconcile_pin_policy(training_plan) + 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) + + 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 + + @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 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). + + 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._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: + self._fatal_setup_failure(error) + + 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() + 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(shape_key=shape_key) + 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}" + self._save_compile_cache() + + 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 + recoverable = False + if allocation_failure: + 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( + rollback_keys, resident=False + ) + else: + self.transition_training_block( + decision.block_key, resident=False + ) + if decision.target_cap_bytes is not None: + 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 + ) + rollback = ( + list(decision.block_keys) + if decision.block_keys + else decision.block_key + ) + torch.cuda.empty_cache() + recoverable = not bool( + getattr(self._config, "strict_vram_cap", False) + ) + elif cap_decision is None and 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)) + 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, + "recoverable": bool(recoverable), + "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} + + 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, + gc_threshold=self._allocator_gc_threshold, + ) + 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: + 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, + gc_threshold=gc_threshold, + ), + 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, *, 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() + 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: + 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, + *, + 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 + checkpointing) over the same arena. TRAIN is restored by the enclosing + `sampling_session()`. + """ + self._require_open() + policy = self._config._policy + + 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) + ) + 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, + 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 + ) + ) + ) + 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 + self._save_compile_cache() + 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 _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 _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 + ) + 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_physical_headroom_bytes( + cumulative + ) + < 0 + ): + break + if allocator_slack <= used + allocator_headroom: + break + capacity += 1 + 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, + self._residency, + plan, + protected_leaf_keys=(self._smart_plan or {}).get( + "protected_training_leaf_keys", () + ), + ) + 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() + 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, + # 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: + """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 + + selected = [] + selected_bytes = 0 + 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 + ) + if not result.get("changed"): + selected = [] + selected_bytes = 0 + + after = self._training_pressure_snapshot() + self._policy.record_physical_pressure_relief( + selected, selected_bytes + ) + self._last_training_pressure_relief = { + "before": before, + "after": after, + "demoted_block_keys": tuple(selected), + "demoted_bytes": selected_bytes, + } + print( + "[ArenaOffload] WDDM pressure relief: " + 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: + 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), + ) + predicted_free = total - ( + worst_working + + current_resident + + current_ring + + non_torch + + int(candidate["block_bytes"]) + ) + return int( + predicted_free + - self._training_physical_vram_headroom_bytes() + ) + + def _worst_shape_allocator_slack_bytes(self, current_cap_bytes): + predicted_live = self._worst_shape_live_bytes() + if predicted_live is None: + 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, + gc_threshold=gc_threshold, + ) + + 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() + if peak.steps > 0 + ) + current_resident = ( + int(self._residency.resident_bytes()) + + int((self._smart_plan or {}).get("singleton_resident_bytes", 0)) + ) + return ( + worst_working + + current_resident + + self._training_ring_bytes() + ) + + def _apply_training_policy(self, *, shape_key=None): + import torch + + 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() + 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), + ) + 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 + 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_physical_headroom_bytes( + candidate + ) + ), + 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": + 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.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.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 + + 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.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]", + ) + + # 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) + 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, + "finalized": self.finalized, + "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", + "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), + "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) + ), + "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), + "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 + ), + "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 + ) + ), + "training_cap_target_bytes": getattr( + self, "_last_training_cap_target_bytes", None + ), + "last_training_pressure_relief": getattr( + self, "_last_training_pressure_relief", None + ), + "training_physical_vram_headroom_bytes": ( + self._training_physical_vram_headroom_bytes() + ), + "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(), + "cap_calibration": self._cap_calibrator.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 _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: + 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 000000000..54ab32fd0 --- /dev/null +++ b/toolkit/memory_management/arena_offload/transfer.py @@ -0,0 +1,616 @@ +"""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 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]] = [] + +_RUNTIME_OWNER_TOKEN = None + + +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 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 _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 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") + 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_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_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=False) + 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). 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). + + +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. + + """ + return _FreeOnBackwardFn.apply(x, token) diff --git a/toolkit/memory_management/canonical_arena.py b/toolkit/memory_management/canonical_arena.py new file mode 100644 index 000000000..6870a7c98 --- /dev/null +++ b/toolkit/memory_management/canonical_arena.py @@ -0,0 +1,340 @@ +"""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 +(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. + +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, or guard operation violated an arena invariant.""" + + +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: + """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): + 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. + """ + + 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 ------------------------------------------------------- + + 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 -- 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, not a mechanism this class provides). + """ + if self._canonicalized: + raise CanonicalArenaError( + "canonical_arena_double_canonicalize: canonicalize() may " + "only run once per arena instance -- 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, + 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, + pin_on_finish=pin_on_finish, + ) + + # -- whole-model .to() interception ------------------------------------ + + @staticmethod + def guard_whole_model_to(model: torch.nn.Module) -> None: + """Route whole-model movement through the published arena runtime.""" + if getattr(model, "_mm_canonical_to_guarded", False): + return + + originals = { + name: getattr(model, name) for name in ("to", "cuda", "cpu") + } + + def runtime_authority(): + runtime = getattr(model, "_arena_offload_runtime", None) + if runtime is None: + raise CanonicalArenaError( + "canonical_arena_whole_model_move_before_runtime" + ) + 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_movement_originals = originals + + @staticmethod + def unguard_whole_model_to(model: torch.nn.Module) -> None: + 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 + + # -- 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() + 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( + 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 and canonical registration policy 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_arena_backed(record.host_flat), + ) + for block_key, record in self._blocks.items() + ), + ) + + # -- explicit unload ------------------------------------------------ + + def release(self) -> None: + """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. + """ + 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) + 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/dxgi_meminfo.py b/toolkit/memory_management/dxgi_meminfo.py new file mode 100644 index 000000000..de7d90a76 --- /dev/null +++ b/toolkit/memory_management/dxgi_meminfo.py @@ -0,0 +1,393 @@ +"""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 + + +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 + + +_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 _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; using legacy pin proxy", + ) + return None + + +def _selection_from_record(adapter, record: DxgiAdapterRecord, match_method: str): + adapter3 = _query_interface_adapter3(adapter) + if match_method == "env_override": + # A human forced the adapter via AI_TOOLKIT_WDDM_DXGI_ADAPTER_INDEX. It + # is worth logging distinctly so a misconfiguration is visible. + _log_once( + "manual_adapter_override", + f"[DXGI] manual adapter override in effect (index={record.index} " + f"{record.description!r}); verify this is the training GPU", + ) + return _AdapterSelection(adapter3) + + +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 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 # The current implementation selects one adapter. + 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 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/immutable_runtime.py b/toolkit/memory_management/immutable_runtime.py new file mode 100644 index 000000000..c74cb49e9 --- /dev/null +++ b/toolkit/memory_management/immutable_runtime.py @@ -0,0 +1,867 @@ +"""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 ( + DEFAULT_DEMOTION_PIN_RESERVE_BLOCKS, + ResidencyDelta, + ResidencyPlan, + ResidencyState, + pin_requirements_for_plan, +) +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.pin_reserve_blocks = DEFAULT_DEMOTION_PIN_RESERVE_BLOCKS + self.last_pin_trim_failures: tuple[str, ...] = () + + 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._residency_promotion_callback = 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 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") + + + @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 " + "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. + + 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") + 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) + 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 + 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 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 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, + 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: + 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) + 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, + 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, + "external_peak_allocated": baseline_allocated, + "external_peak_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 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: + return 0 + + device = self.residency.device + torch.cuda.synchronize(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"]), + ) + 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/manager.py b/toolkit/memory_management/manager.py index c048e940d..fcc9bfa9e 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.configure_wddm_allocator_guard(device, strict=False) 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() torch.cuda.empty_cache() diff --git a/toolkit/memory_management/manager_modules.py b/toolkit/memory_management/manager_modules.py index 44f3b2888..5217e7b1b 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 000000000..550cafff6 --- /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 000000000..7c5f1b1ed --- /dev/null +++ b/toolkit/memory_management/pin_manager.py @@ -0,0 +1,565 @@ +"""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 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.""" + + +class PinReleaseError(RuntimeError): + """Raised when a registered host allocation could not be unpinned.""" + + +@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] = {} + +# 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",) + +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() + + +def _spill_reserve_floor_gib() -> float: + 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: + 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: + # 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. + 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 host_cache_reserve_bytes(mode: Optional[str] = None) -> int: + 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 reconcile() -> None: + """Return cached host-pin allocations before retrying a pin request.""" + _empty_host_pin_cache() + + +_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() + 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`` / 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 + 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 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] = {} + + +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 canonical arena flat's storage independently of pinnedness.""" + 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 canonical arena storage.""" + 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). + 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 + 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) + + 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() + 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() + 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 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, saving roughly 1-1.5 seconds per full arena build. + + Returns ``(candidate, padded_nbytes)``; ``candidate`` is an untouched + pageable view with the layout required by ``pin_register_commit``. + """ + 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 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 _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/residency.py b/toolkit/memory_management/residency.py new file mode 100644 index 000000000..0420071b6 --- /dev/null +++ b/toolkit/memory_management/residency.py @@ -0,0 +1,414 @@ +"""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 import pin_manager +from toolkit.memory_management.canonical_arena import CanonicalArena +from toolkit.memory_management.arena_offload.layout import ( + flatten_leaves, + rebuild_from_leaves, + typed_view, +) + +LeafKey = tuple[str, str] +DEFAULT_DEMOTION_PIN_RESERVE_BLOCKS = 2 + +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 Arena planner's ``offload_ids`` decision to sidecars. + + 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 = [] + 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, 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) + 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 + else torch.no_grad() + ) + with torch.no_grad(), stream_context: + tensors = tuple( + typed_view(block.host_flat, item).to( + self.device, non_blocking=non_blocking + ) + 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 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: + 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 diff --git a/toolkit/memory_management/runtime.py b/toolkit/memory_management/runtime.py new file mode 100644 index 000000000..d00ef19f8 --- /dev/null +++ b/toolkit/memory_management/runtime.py @@ -0,0 +1,68 @@ +"""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() diff --git a/toolkit/memory_management/transfer_plan.py b/toolkit/memory_management/transfer_plan.py new file mode 100644 index 000000000..07bd076f2 --- /dev/null +++ b/toolkit/memory_management/transfer_plan.py @@ -0,0 +1,213 @@ +"""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 +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 Arena transfer runtime submits the copies through +``mm::fetch_start_multi`` while preserving single-ticket ring/backpressure +semantics. + +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 + +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 -- 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, + 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( + [[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 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, + ``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/memory_management/vram_budget.py b/toolkit/memory_management/vram_budget.py new file mode 100644 index 000000000..1b637a5f6 --- /dev/null +++ b/toolkit/memory_management/vram_budget.py @@ -0,0 +1,885 @@ +"""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`` (``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 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 + +import torch + +from . import nvml_meminfo + +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). + + ``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 `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 + + +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" + ) + + +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 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=None, +): + """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 + 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 + - float(max(0, int(allocated_bytes))) + + float(max(0, int(hard_bytes))) + ) + + +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 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 + 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 + 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 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 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 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 + 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) + + +# --------------------------------------------------------------------------- +# Two-timescale residency control +# +# 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 +# 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). +# --------------------------------------------------------------------------- + +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 + 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. + """ + gc_threshold = _resolve_gc_threshold(gc_threshold) + 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=None, +) -> 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. + """ + 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)))) + return int(want) + + +def cap_bytes_preserving_allowance_after_promotion( + cap_bytes, + promotion_bytes, + cliff_cap_bytes, + *, + gc_threshold=None, +) -> 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) + gc_threshold = _resolve_gc_threshold(gc_threshold) + growth = math.ceil(float(promoted) / float(gc_threshold)) + return min(cliff, current + int(growth)) + + +def cap_can_host_promotion( + live_bytes, + block_bytes, + allocator_cache_headroom_bytes, + cliff_cap_bytes, + *, + gc_threshold=None, +) -> 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 ``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 + 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 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))) + + float(max(0, int(allocator_cache_headroom_bytes))) + ) / float(gc_threshold) + return need_cap <= float(int(cliff_cap_bytes)) + + +def residency_promote_ok( + num_alloc_retries, + allocator_slack_bytes, + block_bytes, + allocator_cache_headroom_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 (``gc_threshold * cap - predicted_live``) + exceeds one block plus the pad, so the promotion still leaves + ``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(allocator_cache_headroom_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: + 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: + 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) diff --git a/toolkit/models/DoRA.py b/toolkit/models/DoRA.py index 1d402778f..346ddda5a 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 4caf3ae7a..b69de56d5 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,26 @@ 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( + gen_config=gen_config, + generation_owner=self, + ) + 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 +701,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']) @@ -910,13 +931,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 @@ -1459,11 +1483,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): @@ -1507,6 +1543,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): @@ -1521,6 +1562,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 8f11d6c39..158247a64 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 2e3bfb3ff..9650c60d9 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/quantization/__init__.py b/toolkit/quantization/__init__.py new file mode 100644 index 000000000..8eb944a32 --- /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 000000000..3a042a16e --- /dev/null +++ b/toolkit/quantization/fp8_linear.py @@ -0,0 +1,498 @@ +"""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 + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +from .fp8_transpose import column_major + + +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: + from .torchao_compat import ( + torchao_arena_fp8_supported, + torchao_is_float8_tensor, + ) + + if not torchao_arena_fp8_supported(): + return None + if not torchao_is_float8_tensor(value): + 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 = False + + +def set_fp8_grad_input_enabled(enabled: bool) -> None: + global _FP8_GRAD_INPUT + _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 _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 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_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 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 + + +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, + ) + + +@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, + ) + + +@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_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") diff --git a/toolkit/quantization/fp8_transpose.py b/toolkit/quantization/fp8_transpose.py new file mode 100644 index 000000000..b74806a2e --- /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 000000000..f4402e200 --- /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/quantization/torchao_compat.py b/toolkit/quantization/torchao_compat.py new file mode 100644 index 000000000..60368208f --- /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/timer.py b/toolkit/timer.py index e849ba5fa..61920b057 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: diff --git a/toolkit/util/quantize.py b/toolkit/util/quantize.py index a060a599e..b67da99b5 100644 --- a/toolkit/util/quantize.py +++ b/toolkit/util/quantize.py @@ -4,11 +4,13 @@ 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_, +from toolkit.quantization.torchao_compat import ( Float8WeightOnlyConfig, - UIntXWeightOnlyConfig, - Int8WeightOnlyConfig + Int8WeightOnlyConfig, + intx_weight_only_config, + torchao_is_float8_tensor, + torchao_is_linear, + torchao_quantize_, ) from optimum.quanto import freeze from tqdm import tqdm @@ -41,14 +43,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": 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(), } @@ -169,6 +170,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 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 + 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 +224,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, @@ -278,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() diff --git a/ui/src/app/jobs/new/SimpleJob.tsx b/ui/src/app/jobs/new/SimpleJob.tsx index e5d51b75f..3cc44f844 100644 --- a/ui/src/app/jobs/new/SimpleJob.tsx +++ b/ui/src/app/jobs/new/SimpleJob.tsx @@ -360,18 +360,41 @@ 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, + 'config.process[0].model.layer_offloading_smart_cap_calibration', + ) + } + docKey="model.layer_offloading_smart_cap_calibration" + /> + )} + {!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 07f070c5e..9cb0fef7a 100644 --- a/ui/src/types.ts +++ b/ui/src/types.ts @@ -168,6 +168,8 @@ export interface ModelConfig { low_vram: boolean; model_kwargs: { [key: string]: any }; layer_offloading?: boolean; + layer_offloading_smart?: boolean; + layer_offloading_smart_cap_calibration?: boolean; layer_offloading_transformer_percent?: number; layer_offloading_text_encoder_percent?: number; assistant_lora_path?: string;