Spaces:
Running on Zero
Running on Zero
| """ | |
| This code was heavily inspired by the work of Lodestone-Rock, pretty much all credit goes | |
| to them. The original code can be found here: | |
| https://github.com/lodestone-rock/RamTorch/blob/main/ramtorch/modules/linear.py | |
| I simply modified it to work with a memory management model and with AI Toolkit's models | |
| """ | |
| import os | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from typing import TYPE_CHECKING, Optional, Tuple | |
| from torch.overrides import has_torch_function_unary # (ADD) torchao detection | |
| if TYPE_CHECKING: | |
| from .manager import MemoryManager | |
| # --- Per-device global state registry --- | |
| _DEVICE_STATE = {} | |
| # How many layers deep to prefetch weights. The old ping-pong used 2 slots, which | |
| # only lets one transfer overlap one compute (1-deep). A deeper ring lets Python | |
| # enqueue several layers ahead so the H2D stream stays saturated instead of | |
| # stalling on a per-layer sync. Override with AI_TOOLKIT_OFFLOAD_DEPTH. | |
| PIPELINE_DEPTH = int(os.environ.get("AI_TOOLKIT_OFFLOAD_DEPTH", "4")) | |
| def _get_device_state(device: torch.device): | |
| """Get or initialize per-device state.""" | |
| if isinstance(device, str): | |
| device = torch.device(device) | |
| # CPU path needs no CUDA state | |
| if device.type != "cuda": | |
| if device not in _DEVICE_STATE: | |
| _DEVICE_STATE[device] = {} | |
| return _DEVICE_STATE[device] | |
| if device not in _DEVICE_STATE: | |
| d = max(2, PIPELINE_DEPTH) | |
| with torch.cuda.device(device): | |
| _DEVICE_STATE[device] = { | |
| "depth": d, | |
| # streams | |
| "transfer_stream": torch.cuda.Stream(device=device), | |
| "transfer_grad_stream": torch.cuda.Stream(device=device), | |
| # forward weight ring: slot_ready = H2D done, slot_free = compute | |
| # that consumed the slot done (so it can be overwritten). | |
| "w_buffers": [None] * d, | |
| "b_buffers": [None] * d, | |
| "fwd_slot_ready": [torch.cuda.Event() for _ in range(d)], | |
| "fwd_slot_free": [torch.cuda.Event() for _ in range(d)], | |
| "forward_clk": 0, | |
| # backward weight ring (re-fetch for grad-input). | |
| "w_bwd_buffers": [None] * d, | |
| "bwd_slot_ready": [torch.cuda.Event() for _ in range(d)], | |
| "bwd_slot_free": [torch.cuda.Event() for _ in range(d)], | |
| "backward_clk": 0, | |
| # backward grad-staging ring (device-side grads -> CPU). | |
| "w_grad_buffers": [None] * d, | |
| "b_grad_buffers": [None] * d, | |
| "grad_compute_done": [torch.cuda.Event() for _ in range(d)], | |
| "grad_xfer_done": [torch.cuda.Event() for _ in range(d)], | |
| } | |
| return _DEVICE_STATE[device] | |
| # ---- ring-buffer staging helpers ----------------------------------------- | |
| # | |
| # Each transfer waits only on the event for the *specific slot* it is about to | |
| # overwrite (the compute that used that slot D layers ago), not on a single | |
| # global "compute started" event. With D slots that prior compute is long done, | |
| # so the transfer stream never actually stalls and stays D layers ahead of | |
| # compute. This is the deeper-pipeline + relaxed-dependency change in one. | |
| def _stage_forward_weight(state, device, materialize, weight_cpu, bias_cpu): | |
| """H2D the next forward weight (+bias) into its ring slot; return (idx, w, b). | |
| Caller runs compute, then calls _release_forward_slot(state, idx).""" | |
| d = state["depth"] | |
| idx = state["forward_clk"] | |
| state["forward_clk"] = (idx + 1) % d | |
| ts = state["transfer_stream"] | |
| with torch.cuda.stream(ts): | |
| ts.wait_event(state["fwd_slot_free"][idx]) | |
| state["w_buffers"][idx] = materialize(weight_cpu, device) | |
| state["b_buffers"][idx] = ( | |
| bias_cpu.to(device, non_blocking=True) if bias_cpu is not None else None | |
| ) | |
| state["fwd_slot_ready"][idx].record() | |
| torch.cuda.current_stream().wait_event(state["fwd_slot_ready"][idx]) | |
| return idx, state["w_buffers"][idx], state["b_buffers"][idx] | |
| def _release_forward_slot(state, idx): | |
| # Slot is reusable once the compute stream finishes the op that read it. | |
| state["fwd_slot_free"][idx].record() | |
| def _stage_backward_weight(state, device, materialize, weight_cpu): | |
| """H2D the next backward weight into its ring slot; return (idx, w). | |
| Caller runs grad-input compute, then _release_backward_weight_slot.""" | |
| d = state["depth"] | |
| idx = state["backward_clk"] | |
| state["backward_clk"] = (idx + 1) % d | |
| ts = state["transfer_stream"] | |
| with torch.cuda.stream(ts): | |
| ts.wait_event(state["bwd_slot_free"][idx]) | |
| state["w_bwd_buffers"][idx] = materialize(weight_cpu) | |
| state["bwd_slot_ready"][idx].record() | |
| torch.cuda.current_stream().wait_event(state["bwd_slot_ready"][idx]) | |
| return idx, state["w_bwd_buffers"][idx] | |
| def _release_backward_weight_slot(state, idx): | |
| state["bwd_slot_free"][idx].record() | |
| def _stage_grads_to_cpu(state, idx, grad_w_gpu, grad_b_gpu, weight_cpu, bias_cpu): | |
| """Copy freshly-computed device grads (in staging slot idx) to CPU on the | |
| grad stream, overlapping the next H2D. Returns (grad_w_cpu, grad_b_cpu). | |
| The returned tensors are pinned-memory destinations of an ASYNC copy: their | |
| contents are undefined until the grad stream reaches grad_xfer_done. GPU | |
| consumers are ordered by that event; host consumers must join it first — | |
| the optimizer/clip path does so via sync_grad_transfers(). The one host | |
| read we can't defer is grad accumulation: when the param already holds a | |
| .grad, AccumulateGrad does `grad += returned` on the engine thread the | |
| moment backward returns, so block here until the copy has landed.""" | |
| gs = state["transfer_grad_stream"] | |
| state["grad_compute_done"][idx].record() # on the compute stream | |
| grad_w_cpu = grad_b_cpu = None | |
| with torch.cuda.stream(gs): | |
| gs.wait_event(state["grad_compute_done"][idx]) | |
| if grad_w_gpu is not None: | |
| grad_w_cpu = grad_w_gpu.to("cpu", non_blocking=True) | |
| if grad_b_gpu is not None: | |
| grad_b_cpu = grad_b_gpu.to("cpu", non_blocking=True) | |
| state["grad_xfer_done"][idx].record() | |
| if (grad_w_cpu is not None and weight_cpu.grad is not None) or ( | |
| grad_b_cpu is not None and bias_cpu.grad is not None | |
| ): | |
| state["grad_xfer_done"][idx].synchronize() | |
| return grad_w_cpu, grad_b_cpu | |
| def sync_grad_transfers(): | |
| """Host-join every device's grad D2H stream. | |
| Staged weight/bias grads of memory-managed layers are async copies into | |
| pinned CPU tensors; nothing else orders those copies against the host. | |
| Call this after backward and before anything on the CPU reads .grad of a | |
| memory-managed parameter (grad clipping, optimizer step). No-op when no | |
| offloading is active.""" | |
| for state in _DEVICE_STATE.values(): | |
| stream = state.get("transfer_grad_stream") | |
| if stream is not None: | |
| stream.synchronize() | |
| # (ADD) detect torchao wrapper tensors | |
| def _is_ao_quantized_tensor(t: Optional[torch.Tensor]) -> bool: | |
| if t is None: | |
| return False | |
| try: | |
| if has_torch_function_unary(t): | |
| return t.__class__.__module__.startswith("torchao.") | |
| except Exception: | |
| pass | |
| for attr in ( | |
| "_scale", | |
| "_scales", | |
| "_zero_point", | |
| "_zp", | |
| "_block_size", | |
| "_group_size", | |
| "_pack_dim", | |
| ): | |
| if hasattr(t, attr): | |
| return True | |
| return False | |
| def _is_quantized_tensor(t: Optional[torch.Tensor]) -> bool: | |
| if t is None: | |
| return False | |
| # torch quantized tensors | |
| try: | |
| if torch.is_quantized(t): # type: ignore[attr-defined] | |
| return True | |
| except Exception: | |
| pass | |
| # (ADD) torchao quantized wrappers | |
| if _is_ao_quantized_tensor(t): | |
| return True | |
| # packed/int formats (weight-only) | |
| return not t.dtype.is_floating_point | |
| def _pin_inner_tensors(t: torch.Tensor) -> None: | |
| """Pin the leaf storage of a tensor-subclass (e.g. torchao float8) in place. | |
| Quantized wrappers can't be pin_memory()'d directly, but they expose their | |
| real data as inner tensors via __tensor_flatten__. Pinning those lets the | |
| per-layer H2D bounce run async and overlap with compute instead of blocking. | |
| """ | |
| try: | |
| names, _ = t.__tensor_flatten__() | |
| except Exception: | |
| return | |
| for name in names: | |
| inner = getattr(t, name, None) | |
| if inner is 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 | |
| def _ensure_cpu_pinned(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]: | |
| if t is None: | |
| return None | |
| if t.device.type != "cpu": | |
| try: | |
| t = t.to("cpu", copy=True) | |
| except Exception: | |
| t = t.to("cpu") | |
| # Quantized wrappers can't be pin_memory()'d directly, but pinning their | |
| # inner storage gives the same async-transfer benefit. | |
| if _is_quantized_tensor(t): | |
| if torch.cuda.is_available(): | |
| _pin_inner_tensors(t) | |
| return t | |
| if torch.cuda.is_available(): | |
| try: | |
| t = t.pin_memory() | |
| except RuntimeError: | |
| pass | |
| return t | |
| def _storage_device(t: torch.Tensor) -> torch.device: | |
| """Device of a tensor's actual storage. A wrapper subclass (torchao) can report one | |
| device on the outside while its inner tensors sit elsewhere (a `.data =` move updates | |
| only the outer metadata), so walk __tensor_flatten__ down to the leaves.""" | |
| try: | |
| names, _ = t.__tensor_flatten__() | |
| except Exception: | |
| return t.device | |
| for name in names: | |
| inner = getattr(t, name, None) | |
| if isinstance(inner, torch.Tensor): | |
| return _storage_device(inner) | |
| return t.device | |
| def _move_own_param(module: nn.Module, name: str, param: nn.Parameter, device) -> None: | |
| """Move one of ``module``'s direct parameters to ``device``. Plain tensors go through | |
| `.data =`; quantized tensor subclasses (torchao AffineQuantizedTensor) get the whole | |
| Parameter replaced -- `.data =` on them moves only the outer metadata, leaving the | |
| inner storage behind (same reason _move_params_to_cpu_and_pin replaces them).""" | |
| with torch.no_grad(): | |
| moved = param.data.to(device) | |
| if _is_quantized_tensor(param.data): | |
| setattr(module, name, nn.Parameter(moved, requires_grad=param.requires_grad)) | |
| else: | |
| param.data = moved | |
| def _move_params_to_cpu_and_pin(module: nn.Module): | |
| """Force parameters to CPU (+pinned) so we can 'bounce' them per forward/backward.""" | |
| with torch.no_grad(): | |
| for name in ("weight", "bias"): | |
| param = getattr(module, name, None) | |
| if not isinstance(param, nn.Parameter): | |
| continue | |
| cpu_data = _ensure_cpu_pinned(param.data).detach() | |
| if _is_quantized_tensor(param.data): | |
| # Tensor-subclass weights (e.g. torchao float8 AffineQuantizedTensor) | |
| # ignore `param.data = ...`: the wrapper reports CPU but its inner | |
| # storage stays on the GPU, so the weight never actually offloads. | |
| # Replace the whole Parameter so the device move sticks. | |
| setattr( | |
| module, | |
| name, | |
| nn.Parameter(cpu_data, requires_grad=param.requires_grad), | |
| ) | |
| else: | |
| param.data = cpu_data | |
| # ========================== | |
| # Autograd functions (CUDA) | |
| # ========================== | |
| class _BouncingLinearFn(torch.autograd.Function): | |
| def forward(ctx, x, weight_cpu, bias_cpu, device: torch.device): | |
| # choose compute dtype to match activations | |
| target_dtype = ( | |
| x.dtype | |
| if x.dtype in (torch.bfloat16, torch.float16, torch.float32) | |
| else torch.bfloat16 | |
| ) | |
| # GPU-side dequant/cast for quantized; float path unchanged | |
| def _materialize_linear_weight(cpu_w, dev): | |
| if _is_quantized_tensor(cpu_w): | |
| # move quantized wrapper to GPU -> dequantize on GPU -> cast on GPU | |
| w_q_gpu = cpu_w.to(dev, non_blocking=True) | |
| try: | |
| w_fp_gpu = w_q_gpu.dequantize() | |
| except Exception: | |
| w_fp_gpu = w_q_gpu.to(dtype=torch.float32, non_blocking=True) | |
| if w_fp_gpu.dtype != target_dtype: | |
| w_fp_gpu = w_fp_gpu.to(target_dtype, non_blocking=True) | |
| return w_fp_gpu | |
| # float path (preserve original behavior: NO dtype cast) | |
| w_gpu = cpu_w.to(dev, non_blocking=True) | |
| return w_gpu | |
| if device.type != "cuda": | |
| out = F.linear( | |
| x.to("cpu"), | |
| _materialize_linear_weight(weight_cpu, torch.device("cpu")), | |
| bias_cpu, | |
| ) | |
| ctx.save_for_backward(x.to("cpu"), weight_cpu, bias_cpu) | |
| ctx.device = torch.device("cpu") | |
| return out.to(x.device) | |
| if x.device != device: | |
| # a pipeline that derives its execution device from the module's | |
| # (fully offloaded -> "cpu") parameters can feed cpu activations; | |
| # compute happens on the staged device and the output stays there | |
| x = x.to(device, non_blocking=True) | |
| state = _get_device_state(device) | |
| # the guard makes current_stream() (used by the staging helpers' event | |
| # waits/records) resolve to the process device; without it they hit | |
| # device 0's streams when training on another gpu and nothing orders | |
| # the H2D against the compute | |
| with torch.cuda.device(device): | |
| idx, w_gpu, b_gpu = _stage_forward_weight( | |
| state, device, _materialize_linear_weight, weight_cpu, bias_cpu | |
| ) | |
| out = F.linear(x, w_gpu, b_gpu) | |
| _release_forward_slot(state, idx) | |
| ctx.save_for_backward(x, weight_cpu, bias_cpu) | |
| ctx.device = device | |
| ctx.target_dtype = target_dtype | |
| return out | |
| def backward(ctx, grad_out): | |
| x, weight_cpu, bias_cpu = ctx.saved_tensors | |
| device = ctx.device | |
| target_dtype = getattr(ctx, "target_dtype", grad_out.dtype) | |
| if device.type != "cuda": | |
| go_cpu = grad_out.to("cpu") | |
| x_cpu = x.to("cpu") | |
| w_mat = ( | |
| weight_cpu.dequantize() | |
| if _is_quantized_tensor(weight_cpu) | |
| else weight_cpu | |
| ) | |
| if w_mat.dtype != target_dtype and target_dtype in ( | |
| torch.bfloat16, | |
| torch.float16, | |
| torch.float32, | |
| ): | |
| w_mat = w_mat.to(target_dtype) | |
| grad_input = go_cpu @ w_mat | |
| grad_weight = ( | |
| go_cpu.flatten(0, -2).T @ x_cpu.flatten(0, -2) | |
| if getattr(weight_cpu, "requires_grad", False) | |
| and weight_cpu.dtype.is_floating_point | |
| else None | |
| ) | |
| grad_bias = ( | |
| go_cpu.sum(dim=tuple(range(go_cpu.ndim - 1))) | |
| if (bias_cpu is not None and getattr(bias_cpu, "requires_grad", False)) | |
| else None | |
| ) | |
| return grad_input.to(grad_out.device), grad_weight, grad_bias, None | |
| state = _get_device_state(device) | |
| # GPU-side dequant/cast for quantized; float path unchanged | |
| def _materialize_for_bwd(cpu_w): | |
| if _is_quantized_tensor(cpu_w): | |
| w_q_gpu = cpu_w.to(device, non_blocking=True) | |
| try: | |
| w_fp_gpu = w_q_gpu.dequantize() | |
| except Exception: | |
| w_fp_gpu = w_q_gpu.to(dtype=torch.float32, non_blocking=True) | |
| if w_fp_gpu.dtype != target_dtype: | |
| w_fp_gpu = w_fp_gpu.to(target_dtype, non_blocking=True) | |
| return w_fp_gpu | |
| # float path (preserve original behavior: NO dtype cast) | |
| w = cpu_w.to(device, non_blocking=True) | |
| return w | |
| idx, w_bwd = _stage_backward_weight( | |
| state, device, _materialize_for_bwd, weight_cpu | |
| ) | |
| # grad wrt input (GPU) | |
| grad_input = grad_out.to(dtype=target_dtype) @ w_bwd | |
| _release_backward_weight_slot(state, idx) | |
| # compute grads if float masters exist (frozen/quantized bases skip this) | |
| grad_weight = None | |
| grad_bias = None | |
| need_w = ( | |
| getattr(weight_cpu, "requires_grad", False) | |
| and weight_cpu.dtype.is_floating_point | |
| ) | |
| need_b = bias_cpu is not None and getattr(bias_cpu, "requires_grad", False) | |
| if need_w or need_b: | |
| # ensure the prior grad D2H using this staging slot finished | |
| torch.cuda.current_stream().wait_event(state["grad_xfer_done"][idx]) | |
| w_grad_gpu = b_grad_gpu = None | |
| if need_w: | |
| w_grad_gpu = grad_out.flatten(0, -2).T @ x.flatten(0, -2) | |
| state["w_grad_buffers"][idx] = w_grad_gpu | |
| if need_b: | |
| b_grad_gpu = grad_out.sum(dim=tuple(range(grad_out.ndim - 1))) | |
| state["b_grad_buffers"][idx] = b_grad_gpu | |
| grad_weight, grad_bias = _stage_grads_to_cpu( | |
| state, idx, w_grad_gpu, b_grad_gpu, weight_cpu, bias_cpu | |
| ) | |
| return grad_input.to(dtype=grad_out.dtype), grad_weight, grad_bias, None | |
| class _BouncingConv2dFn(torch.autograd.Function): | |
| def forward( | |
| ctx, | |
| x, | |
| weight_cpu, | |
| bias_cpu, | |
| device: torch.device, | |
| stride: Tuple[int, int], | |
| padding: Tuple[int, int], | |
| dilation: Tuple[int, int], | |
| groups: int, | |
| ): | |
| target_dtype = ( | |
| x.dtype | |
| if x.dtype in (torch.bfloat16, torch.float16, torch.float32) | |
| else torch.bfloat16 | |
| ) | |
| # GPU-side dequant/cast for quantized; float path unchanged | |
| def _materialize_conv_weight(cpu_w, dev): | |
| if _is_quantized_tensor(cpu_w): | |
| w_q_gpu = cpu_w.to(dev, non_blocking=True) | |
| try: | |
| w_fp_gpu = w_q_gpu.dequantize() | |
| except Exception: | |
| w_fp_gpu = w_q_gpu.to(dtype=torch.float32, non_blocking=True) | |
| if w_fp_gpu.dtype != target_dtype: | |
| w_fp_gpu = w_fp_gpu.to(target_dtype, non_blocking=True) | |
| return w_fp_gpu | |
| # float path (preserve original behavior: NO dtype cast) | |
| w_gpu = cpu_w.to(dev, non_blocking=True) | |
| return w_gpu | |
| if device.type != "cuda": | |
| out = F.conv2d( | |
| x.to("cpu"), | |
| _materialize_conv_weight(weight_cpu, torch.device("cpu")), | |
| bias_cpu, | |
| stride, | |
| padding, | |
| dilation, | |
| groups, | |
| ) | |
| ctx.save_for_backward(x.to("cpu"), weight_cpu, bias_cpu) | |
| ctx.meta = ("cpu", stride, padding, dilation, groups, target_dtype) | |
| return out.to(x.device) | |
| if x.device != device: | |
| # cpu activations from a fully-offloaded pipeline: see | |
| # _BouncingLinearFn.forward | |
| x = x.to(device, non_blocking=True) | |
| state = _get_device_state(device) | |
| # device guard: see _BouncingLinearFn.forward | |
| with torch.cuda.device(device): | |
| idx, w_gpu, b_gpu = _stage_forward_weight( | |
| state, device, _materialize_conv_weight, weight_cpu, bias_cpu | |
| ) | |
| out = F.conv2d(x, w_gpu, b_gpu, stride, padding, dilation, groups) | |
| _release_forward_slot(state, idx) | |
| ctx.save_for_backward(x, weight_cpu, bias_cpu) | |
| ctx.meta = (device, stride, padding, dilation, groups, target_dtype) | |
| return out | |
| def backward(ctx, grad_out): | |
| x, weight_cpu, bias_cpu = ctx.saved_tensors | |
| device, stride, padding, dilation, groups, target_dtype = ctx.meta | |
| if ( | |
| isinstance(device, torch.device) and device.type != "cuda" | |
| ) or device == "cpu": | |
| go = grad_out.to("cpu") | |
| x_cpu = x.to("cpu") | |
| w_cpu = ( | |
| weight_cpu.dequantize() | |
| if _is_quantized_tensor(weight_cpu) | |
| else weight_cpu | |
| ) | |
| if w_cpu.dtype != target_dtype and target_dtype in ( | |
| torch.bfloat16, | |
| torch.float16, | |
| torch.float32, | |
| ): | |
| w_cpu = w_cpu.to(target_dtype) | |
| from torch.nn.grad import conv2d_input, conv2d_weight # type: ignore | |
| grad_input = conv2d_input( | |
| x_cpu.shape, | |
| w_cpu, | |
| go, | |
| stride=stride, | |
| padding=padding, | |
| dilation=dilation, | |
| groups=groups, | |
| ) | |
| grad_weight = ( | |
| conv2d_weight( | |
| x_cpu, | |
| w_cpu.shape, | |
| go, | |
| stride=stride, | |
| padding=padding, | |
| dilation=dilation, | |
| groups=groups, | |
| ) | |
| if getattr(weight_cpu, "requires_grad", False) | |
| and weight_cpu.dtype.is_floating_point | |
| else None | |
| ) | |
| grad_bias = ( | |
| go.sum(dim=(0, 2, 3)) | |
| if (bias_cpu is not None and getattr(bias_cpu, "requires_grad", False)) | |
| else None | |
| ) | |
| return ( | |
| grad_input.to(grad_out.device), | |
| grad_weight, | |
| grad_bias, | |
| None, | |
| None, | |
| None, | |
| None, | |
| None, | |
| ) | |
| state = _get_device_state(device) | |
| # GPU-side dequant/cast for quantized; float path unchanged | |
| def _materialize_for_bwd(cpu_w): | |
| if _is_quantized_tensor(cpu_w): | |
| w_q_gpu = cpu_w.to(device, non_blocking=True) | |
| try: | |
| w_fp_gpu = w_q_gpu.dequantize() | |
| except Exception: | |
| w_fp_gpu = w_q_gpu.to(dtype=torch.float32, non_blocking=True) | |
| if w_fp_gpu.dtype != target_dtype: | |
| w_fp_gpu = w_fp_gpu.to(target_dtype, non_blocking=True) | |
| return w_fp_gpu | |
| # float path (preserve original behavior: NO dtype cast) | |
| w = cpu_w.to(device, non_blocking=True) | |
| return w | |
| idx, w_bwd = _stage_backward_weight( | |
| state, device, _materialize_for_bwd, weight_cpu | |
| ) | |
| from torch.nn.grad import conv2d_input, conv2d_weight # type: ignore | |
| grad_input = conv2d_input( | |
| x.shape, | |
| w_bwd, | |
| grad_out.to(dtype=target_dtype), | |
| stride=stride, | |
| padding=padding, | |
| dilation=dilation, | |
| groups=groups, | |
| ) | |
| _release_backward_weight_slot(state, idx) | |
| # Compute heavy grads on GPU into staging buffers (frozen bases skip this) | |
| grad_weight = None | |
| grad_bias = None | |
| need_w = ( | |
| getattr(weight_cpu, "requires_grad", False) | |
| and weight_cpu.dtype.is_floating_point | |
| ) | |
| need_b = bias_cpu is not None and getattr(bias_cpu, "requires_grad", False) | |
| if need_w or need_b: | |
| torch.cuda.current_stream().wait_event(state["grad_xfer_done"][idx]) | |
| w_grad_gpu = b_grad_gpu = None | |
| if need_w: | |
| w_grad_gpu = conv2d_weight( | |
| x, | |
| weight_cpu.shape, | |
| grad_out, | |
| stride=stride, | |
| padding=padding, | |
| dilation=dilation, | |
| groups=groups, | |
| ) | |
| state["w_grad_buffers"][idx] = w_grad_gpu | |
| if need_b: | |
| b_grad_gpu = grad_out.sum(dim=(0, 2, 3)) | |
| state["b_grad_buffers"][idx] = b_grad_gpu | |
| grad_weight, grad_bias = _stage_grads_to_cpu( | |
| state, idx, w_grad_gpu, b_grad_gpu, weight_cpu, bias_cpu | |
| ) | |
| return ( | |
| grad_input.to(dtype=grad_out.dtype), | |
| grad_weight, | |
| grad_bias, | |
| None, | |
| None, | |
| None, | |
| None, | |
| None, | |
| ) | |
| class BaseLayerMemoryManager: | |
| def __init__( | |
| self, | |
| module: nn.Module, | |
| manager: "MemoryManager", | |
| ): | |
| self.module: nn.Module = module | |
| self.manager: "MemoryManager" = manager | |
| def attach(cls, module: nn.Module, manager: "MemoryManager"): | |
| if hasattr(module, "_layer_memory_manager"): | |
| return | |
| module._layer_memory_manager = cls(module, manager) | |
| # mark parameters as memory managed | |
| for param in module.parameters(recurse=False): | |
| param._is_memory_managed = True | |
| class EmbeddingLayerMemoryManager(BaseLayerMemoryManager): | |
| """Offloads a (large, frozen) nn.Embedding by keeping the weight on cpu | |
| and doing the row gather THERE: instead of staging a multi-GB vocab table | |
| to the gpu, only the looked-up rows (tokens x dim, ~KBs) cross the bus. | |
| Lookups happen once per prompt, so the cpu gather is free.""" | |
| def __init__(self, module: nn.Module, manager: "MemoryManager"): | |
| super().__init__(module, manager) | |
| # cpu-resident weight; no pinning — the weight never crosses the bus | |
| module.weight.data = module.weight.data.to("cpu") | |
| # subclass buffers (gemma's embed_scale) must join the cpu-side math | |
| for buf_name, buf in module._buffers.items(): | |
| if buf is not None: | |
| module._buffers[buf_name] = buf.to("cpu") | |
| self._original_forward = module.forward | |
| def _mm_forward(input_ids, *args, **kwargs): | |
| if args or kwargs: | |
| return self._original_forward(input_ids, *args, **kwargs) | |
| out_device = ( | |
| input_ids.device | |
| if input_ids.device.type == "cuda" | |
| else self.manager.process_device | |
| ) | |
| # the original forward preserves subclass behavior (scaled word | |
| # embeddings multiply by embed_scale; raw F.embedding would not) | |
| out = self._original_forward(input_ids.to("cpu")) | |
| return out.to(out_device, non_blocking=True) | |
| module.forward = _mm_forward | |
| class LinearLayerMemoryManager(BaseLayerMemoryManager): | |
| def __init__( | |
| self, | |
| module: nn.Module, | |
| manager: "MemoryManager", | |
| ): | |
| super().__init__(module, manager) | |
| # 1) Move params to CPU + pin memory for fast H2D | |
| _move_params_to_cpu_and_pin(self.module) | |
| # 2) Hijack forward | |
| if hasattr(self.module, "ara_lora_ref"): | |
| # ARA, we need to replace the lora forward | |
| self._original_forward = getattr(self.module.ara_lora_ref(), "org_forward") | |
| else: | |
| self._original_forward = getattr(self.module, "forward") | |
| def _mm_forward(x, *args, **kwargs): | |
| # ensure we only use expected signature (Linear: x) | |
| if args or kwargs: | |
| # fall back to original if a custom signature is used | |
| return self._original_forward(x, *args, **kwargs) | |
| weight_cpu = self.module.weight | |
| bias_cpu = getattr(self.module, "bias", None) | |
| device = self.manager.process_device | |
| # NOTE: do NOT move params to device here; autograd fn streams & bounces them | |
| return _BouncingLinearFn.apply(x, weight_cpu, bias_cpu, device) | |
| if hasattr(self.module, "ara_lora_ref"): | |
| self.module.ara_lora_ref().org_forward = _mm_forward | |
| else: | |
| self.module.forward = _mm_forward | |
| self.module._memory_management_device = self.manager.process_device | |
| class OstrisLinearLayerMemoryManager(BaseLayerMemoryManager): | |
| """Offload manager for OstrisLinear (custom-quantized) layers. | |
| The generic linear bounce is wrong for these: module.weight is a property that | |
| fully dequantizes on access, so bouncing it ships a full-precision weight over | |
| PCIe every forward and bypasses the quantizer's hardware kernels. Instead this | |
| keeps the (much smaller) quantized buffers pinned on CPU, stages them H2D into | |
| the same forward ring the float path uses, swaps them onto the module, and runs | |
| the quantizer's own forward on device — so fp4/int8 GEMM paths and the STE | |
| training path work unchanged under offloading. Buffers are read live off the | |
| module each forward (not cached) so requantize_ during merge/reset stays valid. | |
| """ | |
| def __init__( | |
| self, | |
| module: nn.Module, | |
| manager: "MemoryManager", | |
| ): | |
| super().__init__(module, manager) | |
| # 1) Move quantized buffers + bias to CPU and pin for fast async H2D | |
| with torch.no_grad(): | |
| for name, buf in list(module._buffers.items()): | |
| if buf is None: | |
| 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 | |
| module._buffers[name] = buf | |
| bias = module._parameters.get("bias", None) | |
| if bias is not None: | |
| bias.data = _ensure_cpu_pinned(bias.data).detach() | |
| # 2) Hijack forward | |
| if hasattr(self.module, "ara_lora_ref"): | |
| # ARA, we need to replace the lora forward | |
| self._original_forward = getattr(self.module.ara_lora_ref(), "org_forward") | |
| else: | |
| self._original_forward = getattr(self.module, "forward") | |
| def _mm_forward(x, *args, **kwargs): | |
| # ensure we only use expected signature (Linear: x) | |
| if args or kwargs: | |
| return self._original_forward(x, *args, **kwargs) | |
| module = self.module | |
| device = self.manager.process_device | |
| if device.type != "cuda": | |
| return self._original_forward(x) | |
| cpu_bufs = { | |
| n: b | |
| for n, b in module._buffers.items() | |
| if b is not None and b.device.type == "cpu" | |
| } | |
| bias = module._parameters.get("bias", None) | |
| bias_cpu = ( | |
| bias.data | |
| if bias is not None and bias.data.device.type == "cpu" | |
| else None | |
| ) | |
| if not cpu_bufs and bias_cpu is None: | |
| # already resident on device | |
| return self._original_forward(x) | |
| if x.device != device: | |
| # cpu activations from a fully-offloaded pipeline: see | |
| # _BouncingLinearFn.forward | |
| x = x.to(device, non_blocking=True) | |
| state = _get_device_state(device) | |
| d = state["depth"] | |
| idx = state["forward_clk"] | |
| state["forward_clk"] = (idx + 1) % d | |
| ts = state["transfer_stream"] | |
| # the guard makes current_stream() resolve to the process device and | |
| # keeps that device's context active for the quantizer's triton | |
| # kernels (nothing sets the global current device, so it is 0 even | |
| # when training on another gpu) | |
| with torch.cuda.device(device): | |
| with torch.cuda.stream(ts): | |
| ts.wait_event(state["fwd_slot_free"][idx]) | |
| gpu_bufs = { | |
| n: b.to(device, non_blocking=True) for n, b in cpu_bufs.items() | |
| } | |
| gpu_bias = ( | |
| bias_cpu.to(device, non_blocking=True) | |
| if bias_cpu is not None | |
| else None | |
| ) | |
| state["w_buffers"][idx] = gpu_bufs | |
| state["b_buffers"][idx] = gpu_bias | |
| state["fwd_slot_ready"][idx].record() | |
| compute_stream = torch.cuda.current_stream() | |
| compute_stream.wait_event(state["fwd_slot_ready"][idx]) | |
| # These buffers are allocated on the transfer stream but consumed | |
| # on the compute stream. ConvRot's training operators also save the | |
| # quantized buffers for their custom autograd backward, which can | |
| # outlive the forward-only fwd_slot_free event below. Register the | |
| # consuming stream so the caching allocator cannot recycle their | |
| # storage while either forward or backward kernels still use it. | |
| for tensor in gpu_bufs.values(): | |
| tensor.record_stream(compute_stream) | |
| if gpu_bias is not None: | |
| gpu_bias.record_stream(compute_stream) | |
| # swap the quantized state onto the device, run the quantizer's own | |
| # forward, then swap the pinned CPU state back | |
| for n, t in gpu_bufs.items(): | |
| module._buffers[n] = t | |
| if gpu_bias is not None: | |
| bias.data = gpu_bias | |
| try: | |
| out = self._original_forward(x) | |
| finally: | |
| for n, t in cpu_bufs.items(): | |
| module._buffers[n] = t | |
| if bias_cpu is not None: | |
| bias.data = bias_cpu | |
| _release_forward_slot(state, idx) | |
| return out | |
| if hasattr(self.module, "ara_lora_ref"): | |
| self.module.ara_lora_ref().org_forward = _mm_forward | |
| else: | |
| self.module.forward = _mm_forward | |
| self.module._memory_management_device = self.manager.process_device | |
| class ConvLayerMemoryManager(BaseLayerMemoryManager): | |
| def __init__( | |
| self, | |
| module: nn.Module, | |
| manager: "MemoryManager", | |
| ): | |
| super().__init__(module, manager) | |
| # 1) Move params to CPU + pin memory for fast H2D | |
| _move_params_to_cpu_and_pin(self.module) | |
| # Cache static conv attributes from the module | |
| stride = ( | |
| self.module.stride | |
| if isinstance(self.module.stride, tuple) | |
| else (self.module.stride, self.module.stride) | |
| ) | |
| padding = ( | |
| self.module.padding | |
| if isinstance(self.module.padding, tuple) | |
| else (self.module.padding, self.module.padding) | |
| ) | |
| dilation = ( | |
| self.module.dilation | |
| if isinstance(self.module.dilation, tuple) | |
| else (self.module.dilation, self.module.dilation) | |
| ) | |
| groups = self.module.groups | |
| # 2) Hijack forward | |
| if hasattr(self.module, "ara_lora_ref"): | |
| # ARA, we need to replace the lora forward | |
| self._original_forward = getattr(self.module.ara_lora_ref(), "org_forward") | |
| else: | |
| self._original_forward = getattr(self.module, "forward") | |
| def _mm_forward(x, *args, **kwargs): | |
| # Support the typical Conv2d(x) call; if user passes uncommon extras, fallback. | |
| if args or kwargs: | |
| return self._original_forward(x, *args, **kwargs) | |
| weight_cpu = self.module.weight | |
| bias_cpu = getattr(self.module, "bias", None) | |
| device = self.manager.process_device | |
| return _BouncingConv2dFn.apply( | |
| x, weight_cpu, bias_cpu, device, stride, padding, dilation, groups | |
| ) | |
| if hasattr(self.module, "ara_lora_ref"): | |
| self.module.ara_lora_ref().org_forward = _mm_forward | |
| else: | |
| self.module.forward = _mm_forward | |
| self.module._memory_management_device = self.manager.process_device | |