Retain pinned CPU backups across resume cycles - #86
Conversation
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
There was a problem hiding this comment.
nit: given retain_cpu_backup looks like sth similar to other flags such as enable-cpu-backup, enable-disk-backup, etc, maybe we can make the configuration code mimic them
There was a problem hiding this comment.
(see the comment below about this ai review)
This review was performed by an AI coding agent (Codex), not a human reviewer. I found four blocking regressions/design issues; details are inline. I reviewed head 098b0e9 against its merge base and also checked integration with current master (39093d6). The PR is currently conflicting with master, but the comments below are semantic issues in the PR itself rather than mechanical merge conflicts. I did not find an additional hot-path performance regression beyond the documented opt-in pinned-RAM tradeoff.
| return nullptr; | ||
| return metadata.cpu_backup == nullptr | ||
| ? nullptr | ||
| : (uint8_t*) metadata.cpu_backup + offset; |
There was a problem hiding this comment.
[P1] This ACTIVE-state return is compiled for ROCm too, even though the new policy is advertised as CUDA-only. The ROCm 6.x resume path restores from metadata.cpu_backup but does not free or null it, so after this change get_cpu_backup() starts returning a retained/stale snapshot for ACTIVE ROCm allocations while retain_cpu_backup still reports False and enabling it is rejected. ROCm 7.x behaves differently because it follows this file's non-retaining resume path. Please gate the new ACTIVE visibility to CUDA or first normalize the ROCm backup lifetime.
| @retain_cpu_backup.setter | ||
| def retain_cpu_backup(self, value: bool): | ||
| self._ensure_initialized() | ||
| self._impl._binary_wrapper.cdll.tms_set_retain_cpu_backup(value) |
There was a problem hiding this comment.
[P2] Treating this as an immediate runtime call makes configuration order observable in a surprising way: torch_memory_saver.retain_cpu_backup = True; torch_memory_saver.hook_mode = "torch" initializes the default preload implementation on the first statement, deletes _impl_ctor_kwargs, and makes the second statement fail. Merely reading the new getter has the same side effect. Please put this policy into the pending initialization config (alongside hook_mode) and apply it when the implementation is created, so inspecting or setting one config value does not silently choose the allocator backend.
| atexit.register(self._cleanup_cuda_at_exit) | ||
|
|
||
| def _cleanup_cuda_at_exit(self): | ||
| self._binary_wrapper.cdll.tms_set_interesting_region(True) |
There was a problem hiding this comment.
[P2] This permanently changes the current thread's allocator policy. Python runs atexit handlers in LIFO order, so an application handler registered before lazy TMS initialization will run after this cleanup; any CUDA work in that later handler can now be intercepted even if the original interesting_region value was false, and it happens after the one-shot pool/backup cleanup has already run. Please save the original value and restore it in finally, or avoid mutating allocator policy as part of shutdown.
|
(I will make a final manual review after the human suggestion above is resolved; the agent's auto review is for reference, not a strict requirement) |
2 participants
Summary
Add an opt-in CUDA policy that retains pinned CPU backups across resume cycles:
torch_memory_saver.retain_cpu_backup = TrueTMS_RETAIN_CPU_BACKUP=1at process startupA later pause refreshes the retained storage in place. The existing backup query exposes the retained snapshot after resume, and CUDA process cleanup releases any remaining backups after memory pools are destroyed. ROCm remains unchanged and rejects enabling the policy.
This avoids repeated pinned-host allocation for large, interruptible model-weight workloads. It deliberately trades pinned RAM equal to backed-up allocations for faster repeated resume.
Validation
Against current
master(6d5bce4):torchandpreloadhook modes on H100 hardwarecudaFreeHoston resume deterministically fails the retained active-backup assertiongit diff --checkpassesProduction validation on an 8x H100 model server: