846 lines
32 KiB
Python
846 lines
32 KiB
Python
"""
|
|
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 _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):
|
|
@staticmethod
|
|
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)
|
|
|
|
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
|
|
|
|
@staticmethod
|
|
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):
|
|
@staticmethod
|
|
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)
|
|
|
|
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
|
|
|
|
@staticmethod
|
|
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
|
|
|
|
@classmethod
|
|
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 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)
|
|
|
|
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()
|
|
torch.cuda.current_stream().wait_event(state["fwd_slot_ready"][idx])
|
|
|
|
# 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
|