File size: 2,820 Bytes
9118991 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Exact saved-tensor placement with a bounded GPU retention budget."""
import torch
def storage_key(tensor):
return (str(tensor.device), tensor.untyped_storage().data_ptr())
class BudgetedSavedTensors(torch.autograd.graph.saved_tensors_hooks):
"""Keep existing backbone storage and a bounded set of additional storages.
Count whole unique storages, not views. Stored tensors are detached aliases;
no casting, quantization or recomputation changes backward values.
"""
def __init__(self, resident_tensors=(), keep_bytes=0):
if keep_bytes < 0:
raise ValueError('saved-tensor GPU budget must be nonnegative')
resident = {storage_key(t) for t in resident_tensors}
kept = set()
self.retained_bytes = 0
self.resident_hits = 0
self.cpu_copies = 0
cpu = torch.autograd.graph.save_on_cpu(pin_memory=True, device_type='cuda')
def pack(tensor):
key = storage_key(tensor)
size = tensor.untyped_storage().nbytes()
if key in resident:
self.resident_hits += 1
return ('resident', tensor.detach())
if key in kept or self.retained_bytes + size <= keep_bytes:
if key not in kept:
kept.add(key)
self.retained_bytes += size
return ('resident', tensor.detach())
self.cpu_copies += 1
return ('cpu', cpu.pack_hook(tensor))
def unpack(packed):
kind, value = packed
return value if kind == 'resident' else cpu.unpack_hook(value)
super().__init__(pack, unpack)
class AsyncSavedTensors(torch.autograd.graph.saved_tensors_hooks):
"""Pinned D2H without a host synchronization at every saved tensor.
An event explicitly orders each H2D restore after its D2H copy, including
when autograd executes that restore on a different CUDA stream.
"""
def __init__(self):
def pack(tensor):
if tensor.device.type != 'cuda':
return (tensor.device, tensor.detach(), None)
packed = torch.empty(tensor.size(), dtype=tensor.dtype,
layout=tensor.layout, pin_memory=True)
stream = torch.cuda.current_stream(tensor.device)
packed.copy_(tensor, non_blocking=True)
tensor.record_stream(stream)
ready = torch.cuda.Event()
ready.record(stream)
return (tensor.device, packed, ready)
def unpack(saved):
device, tensor, ready = saved
if ready is None:
return tensor
torch.cuda.current_stream(device).wait_event(ready)
return tensor.to(device, non_blocking=True)
super().__init__(pack, unpack)
|