diff --git a/benchmarking/optimizer_current_stream.py b/benchmarking/optimizer_current_stream.py new file mode 100644 index 000000000..9748d56d7 --- /dev/null +++ b/benchmarking/optimizer_current_stream.py @@ -0,0 +1,432 @@ +#!/usr/bin/env python3 + +import argparse +import ctypes as ct +import json +import os +from pathlib import Path +import statistics +import time + +import torch + +import bitsandbytes as bnb +from bitsandbytes.backends.cuda import ops as cuda_ops +from bitsandbytes.cextension import lib +from bitsandbytes.utils import sync_gpu + +OPTIMIZER_DTYPES_32 = { + "adam": ("fp32", "fp16", "bf16"), + "momentum": ("32", "16"), + "rmsprop": ("32", "16"), + "lion": ("fp32", "fp16", "bf16"), + "adagrad": ("32", "16"), + "ademamix": ("fp32", "fp16", "bf16"), +} +OPTIMIZER_NAMES_8 = ("adam", "momentum", "rmsprop", "lion", "adagrad", "ademamix") +DTYPES = {"fp16": torch.float16, "bf16": torch.bfloat16} + + +class LegacyFunction: + def __init__(self, function, stream_function): + function.argtypes = stream_function.argtypes[:-1] + function.restype = stream_function.restype + self.function = function + + def __call__(self, *args): + return self.function(*args[:-1]) + + +def parse_csv(value, cast=str): + return [cast(item) for item in value.split(",") if item] + + +def parse_args(): + parser = argparse.ArgumentParser(description="Compare legacy and current-stream non-paged optimizer updates") + parser.add_argument("--baseline-library", type=Path, required=True) + parser.add_argument("--expected-library", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--warmups", type=int, default=10) + parser.add_argument("--rounds", type=int, default=7) + parser.add_argument("--inventories", default="32,256,1024") + parser.add_argument("--dtypes", default="fp16,bf16") + parser.add_argument("--bits", default="8,32") + parser.add_argument("--single-large-numel", type=int, default=16 * 1024 * 1024) + parser.add_argument("--peft-layers", type=int, default=32) + parser.add_argument("--peft-hidden", type=int, default=512) + parser.add_argument("--peft-rank", type=int, default=8) + parser.add_argument("--peft-batch", type=int, default=16) + parser.add_argument("--seed", type=int, default=20260821) + return parser.parse_args() + + +def percentile(values, fraction): + ordered = sorted(values) + position = (len(ordered) - 1) * fraction + lower = int(position) + upper = min(lower + 1, len(ordered) - 1) + weight = position - lower + return ordered[lower] * (1.0 - weight) + ordered[upper] * weight + + +def summarize(values): + return { + "median": statistics.median(values), + "p10": percentile(values, 0.1), + "p90": percentile(values, 0.9), + "samples": values, + } + + +def build_symbol_maps(baseline_library): + stream32 = cuda_ops.str2optimizer32bit.copy() + stream8 = cuda_ops.str2optimizer8bit_blockwise.copy() + legacy32 = {} + for name, dtypes in OPTIMIZER_DTYPES_32.items(): + functions = [] + for dtype in dtypes: + legacy = getattr(baseline_library, f"c{name}32bit_grad_{dtype}") + current = getattr(lib, f"c{name}32bit_grad_{dtype}_with_stream") + functions.append(LegacyFunction(legacy, current)) + legacy32[name] = tuple(functions) + legacy32["lamb"] = legacy32["adam"] + legacy32["lars"] = legacy32["momentum"] + + legacy8 = {} + for name in OPTIMIZER_NAMES_8: + functions = [] + for dtype in ("fp32", "fp16", "bf16"): + legacy = getattr(baseline_library, f"c{name}_8bit_blockwise_grad_{dtype}") + current = getattr(lib, f"c{name}_8bit_blockwise_grad_{dtype}_with_stream") + functions.append(LegacyFunction(legacy, current)) + legacy8[name] = tuple(functions) + return stream32, stream8, legacy32, legacy8 + + +@torch.no_grad() +def legacy_step(optimizer): + if not optimizer.initialized: + optimizer.check_overrides() + optimizer.to_gpu() + optimizer.initialized = True + for group_index, group in enumerate(optimizer.param_groups): + for parameter_index, parameter in enumerate(group["params"]): + if parameter.grad is None: + continue + state = optimizer.state[parameter] + if not state: + optimizer.init_state(group, parameter, group_index, parameter_index) + optimizer.prefetch_state(parameter) + optimizer.update_step(group, parameter, group_index, parameter_index) + sync_gpu(parameter) + + +def activate_maps(maps, variant): + stream32, stream8, legacy32, legacy8 = maps + if variant == "legacy": + cuda_ops.str2optimizer32bit = legacy32 + cuda_ops.str2optimizer8bit_blockwise = legacy8 + else: + cuda_ops.str2optimizer32bit = stream32 + cuda_ops.str2optimizer8bit_blockwise = stream8 + + +def completed_measure(label, function): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.nvtx.range_push(label) + started = time.perf_counter() + start.record() + function() + end.record() + end.synchronize() + wall_ms = (time.perf_counter() - started) * 1e3 + event_ms = start.elapsed_time(end) + torch.cuda.nvtx.range_pop() + return event_ms, wall_ms + + +def make_optimizer(parameters, state_bits): + if state_bits == 8: + return bnb.optim.AdamW8bit(parameters, lr=1e-3, min_8bit_size=4096) + return bnb.optim.AdamW32bit(parameters, lr=1e-3) + + +def make_parameter_pair(sizes, dtype, seed): + generator = torch.Generator(device="cuda").manual_seed(seed) + baseline = [] + candidate = [] + for size in sizes: + value = torch.randn(size, dtype=dtype, device="cuda", generator=generator) * 0.01 + gradient = torch.randn(size, dtype=dtype, device="cuda", generator=generator) * 0.001 + baseline_parameter = torch.nn.Parameter(value.clone()) + candidate_parameter = torch.nn.Parameter(value.clone()) + baseline_parameter.grad = gradient.clone() + candidate_parameter.grad = gradient.clone() + baseline.append(baseline_parameter) + candidate.append(candidate_parameter) + return baseline, candidate + + +def assert_optimizer_equal(baseline_optimizer, candidate_optimizer, baseline_params, candidate_params): + for baseline_parameter, candidate_parameter in zip(baseline_params, candidate_params): + if not torch.equal(baseline_parameter, candidate_parameter): + raise AssertionError("parameter mismatch between legacy and current-stream variants") + if not torch.equal(baseline_parameter.grad, candidate_parameter.grad): + raise AssertionError("gradient changed or mismatched") + baseline_state = baseline_optimizer.state[baseline_parameter] + candidate_state = candidate_optimizer.state[candidate_parameter] + if baseline_state.keys() != candidate_state.keys(): + raise AssertionError("optimizer state keys differ") + for key in baseline_state: + left = baseline_state[key] + right = candidate_state[key] + if isinstance(left, torch.Tensor): + if not torch.equal(left, right): + raise AssertionError(f"optimizer state mismatch for {key}") + elif left != right: + raise AssertionError(f"optimizer metadata mismatch for {key}") + + +def optimizer_state_bytes(optimizer): + seen = set() + total = 0 + for state in optimizer.state.values(): + for value in state.values(): + if not isinstance(value, torch.Tensor) or value.data_ptr() in seen: + continue + seen.add(value.data_ptr()) + total += value.numel() * value.element_size() + return total + + +def run_inventory(maps, sizes, dtype_name, state_bits, warmups, rounds, seed, label): + dtype = DTYPES[dtype_name] + baseline_params, candidate_params = make_parameter_pair(sizes, dtype, seed) + baseline_optimizer = make_optimizer(baseline_params, state_bits) + candidate_optimizer = make_optimizer(candidate_params, state_bits) + + def run(variant): + activate_maps(maps, variant) + if variant == "legacy": + legacy_step(baseline_optimizer) + else: + candidate_optimizer.step() + + for index in range(warmups): + order = ("legacy", "current_stream") if index % 2 == 0 else ("current_stream", "legacy") + for variant in order: + completed_measure(f"warmup_{label}_{variant}", lambda variant=variant: run(variant)) + + event_samples = {"legacy": [], "current_stream": []} + wall_samples = {"legacy": [], "current_stream": []} + for round_index in range(rounds): + order = ("legacy", "current_stream") if round_index % 2 == 0 else ("current_stream", "legacy") + for variant in order: + event_ms, wall_ms = completed_measure(f"timed_{label}_{variant}", lambda variant=variant: run(variant)) + event_samples[variant].append(event_ms) + wall_samples[variant].append(wall_ms) + + torch.cuda.synchronize() + assert_optimizer_equal(baseline_optimizer, candidate_optimizer, baseline_params, candidate_params) + legacy_event = summarize(event_samples["legacy"]) + current_event = summarize(event_samples["current_stream"]) + legacy_wall = summarize(wall_samples["legacy"]) + current_wall = summarize(wall_samples["current_stream"]) + return { + "type": "inventory", + "label": label, + "optimizer": "AdamW", + "state_bits": state_bits, + "dtype": dtype_name, + "parameter_count": len(sizes), + "total_numel": sum(sizes), + "min_numel": min(sizes), + "max_numel": max(sizes), + "warmups": warmups, + "rounds": rounds, + "bitwise_equal": True, + "state_bytes": optimizer_state_bytes(candidate_optimizer), + "legacy_event_ms": legacy_event, + "current_stream_event_ms": current_event, + "legacy_wall_ms": legacy_wall, + "current_stream_wall_ms": current_wall, + "event_ratio": legacy_event["median"] / current_event["median"], + "wall_ratio": legacy_wall["median"] / current_wall["median"], + } + + +class AdapterStack(torch.nn.Module): + def __init__(self, layers, hidden, rank, dtype): + super().__init__() + self.left = torch.nn.ParameterList() + self.right = torch.nn.ParameterList() + for _ in range(layers): + self.left.append(torch.nn.Parameter(torch.randn(hidden, rank, device="cuda", dtype=dtype) * 0.01)) + self.right.append(torch.nn.Parameter(torch.randn(rank, hidden, device="cuda", dtype=dtype) * 0.01)) + + def forward(self, value): + for left, right in zip(self.left, self.right): + value = value + (value @ left) @ right + return value + + +def run_peft(maps, layers, hidden, rank, batch, warmups, rounds, seed): + torch.manual_seed(seed) + baseline_model = AdapterStack(layers, hidden, rank, torch.float16) + candidate_model = AdapterStack(layers, hidden, rank, torch.float16) + candidate_model.load_state_dict(baseline_model.state_dict()) + baseline_optimizer = bnb.optim.AdamW8bit(baseline_model.parameters(), lr=1e-3, min_8bit_size=4096) + candidate_optimizer = bnb.optim.AdamW8bit(candidate_model.parameters(), lr=1e-3, min_8bit_size=4096) + inputs = torch.randn(batch, hidden, device="cuda", dtype=torch.float16) + + def run(variant): + if variant == "legacy": + activate_maps(maps, variant) + model = baseline_model + optimizer = baseline_optimizer + else: + activate_maps(maps, variant) + model = candidate_model + optimizer = candidate_optimizer + optimizer.zero_grad(set_to_none=True) + model(inputs).float().square().mean().backward() + if variant == "legacy": + legacy_step(optimizer) + else: + optimizer.step() + + for index in range(warmups): + order = ("legacy", "current_stream") if index % 2 == 0 else ("current_stream", "legacy") + for variant in order: + completed_measure(f"warmup_peft_{variant}", lambda variant=variant: run(variant)) + + event_samples = {"legacy": [], "current_stream": []} + wall_samples = {"legacy": [], "current_stream": []} + for round_index in range(rounds): + order = ("legacy", "current_stream") if round_index % 2 == 0 else ("current_stream", "legacy") + for variant in order: + event_ms, wall_ms = completed_measure(f"timed_peft_{variant}", lambda variant=variant: run(variant)) + event_samples[variant].append(event_ms) + wall_samples[variant].append(wall_ms) + + torch.cuda.synchronize() + assert_optimizer_equal( + baseline_optimizer, + candidate_optimizer, + list(baseline_model.parameters()), + list(candidate_model.parameters()), + ) + legacy_event = summarize(event_samples["legacy"]) + current_event = summarize(event_samples["current_stream"]) + legacy_wall = summarize(wall_samples["legacy"]) + current_wall = summarize(wall_samples["current_stream"]) + return { + "type": "peft", + "optimizer": "AdamW8bit", + "dtype": "fp16", + "layers": layers, + "hidden": hidden, + "rank": rank, + "batch": batch, + "parameter_count": sum(1 for _ in baseline_model.parameters()), + "warmups": warmups, + "rounds": rounds, + "bitwise_equal": True, + "legacy_event_ms": legacy_event, + "current_stream_event_ms": current_event, + "legacy_wall_ms": legacy_wall, + "current_stream_wall_ms": current_wall, + "event_ratio": legacy_event["median"] / current_event["median"], + "wall_ratio": legacy_wall["median"] / current_wall["median"], + } + + +def main(): + args = parse_args() + if args.warmups < 0 or args.rounds < 1: + raise ValueError("warmups must be nonnegative and rounds must be positive") + expected_library = args.expected_library.resolve() + baseline_library_path = args.baseline_library.resolve() + loaded_library = Path(lib._lib._name).resolve() + if loaded_library != expected_library: + raise AssertionError(f"loaded {loaded_library}, expected {expected_library}") + props = torch.cuda.get_device_properties(0) + if "B300" not in props.name or (props.major, props.minor) != (10, 3): + raise AssertionError(f"expected B300/SM103, got {props.name} CC {props.major}.{props.minor}") + + baseline_library = ct.CDLL(str(baseline_library_path)) + maps = build_symbol_maps(baseline_library) + metadata = { + "type": "metadata", + "job_id": os.environ.get("SLURM_JOB_ID"), + "gpu": props.name, + "compute_capability": f"{props.major}.{props.minor}", + "sms": props.multi_processor_count, + "torch": torch.__version__, + "torch_cuda": torch.version.cuda, + "loaded_library": str(loaded_library), + "baseline_library": str(baseline_library_path), + "warmups": args.warmups, + "rounds": args.rounds, + } + records = [metadata] + pattern = (4096, 8192, 16384, 32768) + inventories = parse_csv(args.inventories, int) + dtypes = parse_csv(args.dtypes) + state_bits = parse_csv(args.bits, int) + for count in inventories: + sizes = [pattern[index % len(pattern)] for index in range(count)] + for dtype_name in dtypes: + for bits in state_bits: + records.append( + run_inventory( + maps, + sizes, + dtype_name, + bits, + args.warmups, + args.rounds, + args.seed + count + bits, + f"inventory_{count}", + ) + ) + + for dtype_name in dtypes: + for bits in state_bits: + records.append( + run_inventory( + maps, + [args.single_large_numel], + dtype_name, + bits, + args.warmups, + args.rounds, + args.seed + bits, + "single_large", + ) + ) + + records.append( + run_peft( + maps, + args.peft_layers, + args.peft_hidden, + args.peft_rank, + args.peft_batch, + args.warmups, + args.rounds, + args.seed, + ) + ) + activate_maps(maps, "current_stream") + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w") as output: + for record in records: + line = json.dumps(record, sort_keys=True) + output.write(line + "\n") + print(line) + + +if __name__ == "__main__": + main() diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index a0d9ffe83..02d62dd19 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -65,6 +65,38 @@ def _setup_ctypes(names, argtypes, restype=None): [ct.c_void_p] * 4 + [ct.c_int32, ct.c_int32], ) +# 32-bit optimizer update: (g, p, state1, state2, unorm, optimizer scalars, step, lr, gnorm, skip, n, stream) +_setup_ctypes( + [ + f"c{name}32bit_grad_{dtype}_with_stream" + for name, dtypes in ( + ("adam", ("fp32", "fp16", "bf16")), + ("momentum", ("32", "16")), + ("rmsprop", ("32", "16")), + ("lion", ("fp32", "fp16", "bf16")), + ("adagrad", ("32", "16")), + ("ademamix", ("fp32", "fp16", "bf16")), + ) + for dtype in dtypes + ], + [ct.c_void_p] * 5 + [ct.c_float] * 8 + [ct.c_int32, ct.c_float, ct.c_float, ct.c_bool, ct.c_int32, ct.c_void_p], +) + +# Blockwise 8-bit optimizer update: (p, g, states, scalars, step, lr, maps, absmax, weight decay, gnorm, skip, n, stream) +_setup_ctypes( + [ + f"c{name}_8bit_blockwise_grad_{dtype}_with_stream" + for name in ("adam", "momentum", "rmsprop", "lion", "adagrad", "ademamix") + for dtype in ("fp32", "fp16", "bf16") + ], + [ct.c_void_p] * 4 + + [ct.c_float] * 5 + + [ct.c_int32, ct.c_float] + + [ct.c_void_p] * 4 + + [ct.c_float] * 2 + + [ct.c_bool, ct.c_int32, ct.c_void_p], +) + _get_raw_stream = torch._C._cuda_getCurrentRawStream @@ -985,73 +1017,73 @@ def _( """C FUNCTIONS FOR OPTIMIZERS""" str2optimizer32bit = { "adam": ( - lib.cadam32bit_grad_fp32, - lib.cadam32bit_grad_fp16, - lib.cadam32bit_grad_bf16, + lib.cadam32bit_grad_fp32_with_stream, + lib.cadam32bit_grad_fp16_with_stream, + lib.cadam32bit_grad_bf16_with_stream, ), "momentum": ( - lib.cmomentum32bit_grad_32, - lib.cmomentum32bit_grad_16, + lib.cmomentum32bit_grad_32_with_stream, + lib.cmomentum32bit_grad_16_with_stream, ), "rmsprop": ( - lib.crmsprop32bit_grad_32, - lib.crmsprop32bit_grad_16, + lib.crmsprop32bit_grad_32_with_stream, + lib.crmsprop32bit_grad_16_with_stream, ), "lion": ( - lib.clion32bit_grad_fp32, - lib.clion32bit_grad_fp16, - lib.clion32bit_grad_bf16, + lib.clion32bit_grad_fp32_with_stream, + lib.clion32bit_grad_fp16_with_stream, + lib.clion32bit_grad_bf16_with_stream, ), "adagrad": ( - lib.cadagrad32bit_grad_32, - lib.cadagrad32bit_grad_16, + lib.cadagrad32bit_grad_32_with_stream, + lib.cadagrad32bit_grad_16_with_stream, ), "lamb": ( - lib.cadam32bit_grad_fp32, - lib.cadam32bit_grad_fp16, - lib.cadam32bit_grad_bf16, + lib.cadam32bit_grad_fp32_with_stream, + lib.cadam32bit_grad_fp16_with_stream, + lib.cadam32bit_grad_bf16_with_stream, ), "ademamix": ( - lib.cademamix32bit_grad_fp32, - lib.cademamix32bit_grad_fp16, - lib.cademamix32bit_grad_bf16, + lib.cademamix32bit_grad_fp32_with_stream, + lib.cademamix32bit_grad_fp16_with_stream, + lib.cademamix32bit_grad_bf16_with_stream, ), "lars": ( - lib.cmomentum32bit_grad_32, - lib.cmomentum32bit_grad_16, + lib.cmomentum32bit_grad_32_with_stream, + lib.cmomentum32bit_grad_16_with_stream, ), } str2optimizer8bit_blockwise = { "adam": ( - lib.cadam_8bit_blockwise_grad_fp32, - lib.cadam_8bit_blockwise_grad_fp16, - lib.cadam_8bit_blockwise_grad_bf16, + lib.cadam_8bit_blockwise_grad_fp32_with_stream, + lib.cadam_8bit_blockwise_grad_fp16_with_stream, + lib.cadam_8bit_blockwise_grad_bf16_with_stream, ), "momentum": ( - lib.cmomentum_8bit_blockwise_grad_fp32, - lib.cmomentum_8bit_blockwise_grad_fp16, - lib.cmomentum_8bit_blockwise_grad_bf16, + lib.cmomentum_8bit_blockwise_grad_fp32_with_stream, + lib.cmomentum_8bit_blockwise_grad_fp16_with_stream, + lib.cmomentum_8bit_blockwise_grad_bf16_with_stream, ), "rmsprop": ( - lib.crmsprop_8bit_blockwise_grad_fp32, - lib.crmsprop_8bit_blockwise_grad_fp16, - lib.crmsprop_8bit_blockwise_grad_bf16, + lib.crmsprop_8bit_blockwise_grad_fp32_with_stream, + lib.crmsprop_8bit_blockwise_grad_fp16_with_stream, + lib.crmsprop_8bit_blockwise_grad_bf16_with_stream, ), "lion": ( - lib.clion_8bit_blockwise_grad_fp32, - lib.clion_8bit_blockwise_grad_fp16, - lib.clion_8bit_blockwise_grad_bf16, + lib.clion_8bit_blockwise_grad_fp32_with_stream, + lib.clion_8bit_blockwise_grad_fp16_with_stream, + lib.clion_8bit_blockwise_grad_bf16_with_stream, ), "adagrad": ( - lib.cadagrad_8bit_blockwise_grad_fp32, - lib.cadagrad_8bit_blockwise_grad_fp16, - lib.cadagrad_8bit_blockwise_grad_bf16, + lib.cadagrad_8bit_blockwise_grad_fp32_with_stream, + lib.cadagrad_8bit_blockwise_grad_fp16_with_stream, + lib.cadagrad_8bit_blockwise_grad_bf16_with_stream, ), "ademamix": ( - lib.cademamix_8bit_blockwise_grad_fp32, - lib.cademamix_8bit_blockwise_grad_fp16, - lib.cademamix_8bit_blockwise_grad_bf16, + lib.cademamix_8bit_blockwise_grad_fp32_with_stream, + lib.cademamix_8bit_blockwise_grad_fp16_with_stream, + lib.cademamix_8bit_blockwise_grad_bf16_with_stream, ), } @@ -1092,7 +1124,11 @@ def _optimizer_update_32bit_impl( f"Gradient+optimizer bit data type combination not supported: grad {g.dtype}, optimizer {state1.dtype}", ) + is_paged = getattr(state1, "is_paged", False) or (state2 is not None and getattr(state2, "is_paged", False)) + with _cuda_device_of(g): + # Managed-state prefetches use stream 0, so keep actual paged updates ordered behind them. + stream = None if is_paged else _get_raw_stream(g.device.index) optim_func( get_ptr(g), get_ptr(p), @@ -1112,6 +1148,7 @@ def _optimizer_update_32bit_impl( ct.c_float(gnorm_scale), ct.c_bool(skip_zeros), ct.c_int32(g.numel()), + ct.c_void_p(stream), ) @@ -1184,7 +1221,11 @@ def _optimizer_update_8bit_blockwise_impl( f"Unsupported gradient dtype: {g.dtype}. Supported dtypes: torch.float32, torch.float16, torch.bfloat16" ) + is_paged = getattr(state1, "is_paged", False) or (state2 is not None and getattr(state2, "is_paged", False)) + with _cuda_device_of(g): + # Managed-state prefetches use stream 0, so keep actual paged updates ordered behind them. + stream = None if is_paged else _get_raw_stream(g.device.index) optimizer_fn( get_ptr(p), get_ptr(g), @@ -1205,6 +1246,7 @@ def _optimizer_update_8bit_blockwise_impl( ct.c_float(gnorm_scale), ct.c_bool(skip_zeros), ct.c_int32(g.numel()), + ct.c_void_p(stream), ) diff --git a/bitsandbytes/optim/optimizer.py b/bitsandbytes/optim/optimizer.py index dfc6e5d65..242953094 100644 --- a/bitsandbytes/optim/optimizer.py +++ b/bitsandbytes/optim/optimizer.py @@ -332,7 +332,8 @@ def step(self, closure=None): self.prefetch_state(p) self.update_step(group, p, gindex, pindex) - sync_gpu(p) + if self.is_paged or p.device.type != "cuda" or torch.version.hip is not None: + sync_gpu(p) if self.is_paged and p is not None: # all paged operations are asynchronous, we need # to sync to make sure all tensors are in the right state diff --git a/csrc/compat.cuh b/csrc/compat.cuh index f8c307c2a..6ce84a414 100644 --- a/csrc/compat.cuh +++ b/csrc/compat.cuh @@ -58,6 +58,7 @@ using bnb_error_t = hipError_t; #define BNB_DEVICE_MALLOC(p, s) hipMalloc(p, s) #define BNB_DEVICE_FREE(p) hipFree(p) #define BNB_DEVICE_MEMSET(p, v, s) hipMemset(p, v, s) +#define BNB_DEVICE_MEMSET_ASYNC(p, v, s, stream) hipMemsetAsync(p, v, s, stream) #else // CUDA @@ -70,6 +71,7 @@ using bnb_error_t = cudaError_t; #define BNB_DEVICE_MALLOC(p, s) cudaMalloc(p, s) #define BNB_DEVICE_FREE(p) cudaFree(p) #define BNB_DEVICE_MEMSET(p, v, s) cudaMemset(p, v, s) +#define BNB_DEVICE_MEMSET_ASYNC(p, v, s, stream) cudaMemsetAsync(p, v, s, stream) #endif diff --git a/csrc/ops.cu b/csrc/ops.cu index 16eed4e81..e3631379c 100644 --- a/csrc/ops.cu +++ b/csrc/ops.cu @@ -97,7 +97,7 @@ template void optimizer32bit( T* g, T* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, const float beta1, const float beta2, const float beta3, const float alpha, const float eps, const float weight_decay, const int step, - const float lr, const float gnorm_scale, bool skip_zeros, const int n + const float lr, const float gnorm_scale, bool skip_zeros, const int n, bnb_stream_t stream ) { int num_blocks = n / 4096; num_blocks = n % 4096 == 0 ? num_blocks : num_blocks + 1; @@ -105,13 +105,13 @@ void optimizer32bit( case ADAM: case ADEMAMIX: if (max_unorm > 0.0f) { - BNB_CHECK_RETURN(BNB_DEVICE_MEMSET(unorm, 0, 1 * sizeof(float))); - kPreconditionOptimizer32bit2State<<>>( + BNB_CHECK_RETURN(BNB_DEVICE_MEMSET_ASYNC(unorm, 0, 1 * sizeof(float), stream)); + kPreconditionOptimizer32bit2State<<>>( g, p, state1, state2, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n ); BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } - kOptimizer32bit2State<<>>( + kOptimizer32bit2State<<>>( g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n ); @@ -121,13 +121,14 @@ void optimizer32bit( case RMSPROP: case ADAGRAD: if (max_unorm > 0.0f) { - BNB_CHECK_RETURN(BNB_DEVICE_MEMSET(unorm, 0, 1 * sizeof(float))); - kPreconditionOptimizer32bit1State - <<>>(g, p, state1, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n); + BNB_CHECK_RETURN(BNB_DEVICE_MEMSET_ASYNC(unorm, 0, 1 * sizeof(float), stream)); + kPreconditionOptimizer32bit1State<<>>( + g, p, state1, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n + ); BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } - kOptimizer32bit1State<<>>( + kOptimizer32bit1State<<>>( g, p, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n ); @@ -135,16 +136,17 @@ void optimizer32bit( break; case LION: // in lion, the momentum update after the parameter update - kOptimizer32bit1State<<>>( + kOptimizer32bit1State<<>>( g, p, state1, unorm, max_unorm, param_norm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, skip_zeros, n ); BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); if (max_unorm > 0.0f) { - BNB_CHECK_RETURN(BNB_DEVICE_MEMSET(unorm, 0, 1 * sizeof(float))); - kPreconditionOptimizer32bit1State - <<>>(g, p, state1, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n); + BNB_CHECK_RETURN(BNB_DEVICE_MEMSET_ASYNC(unorm, 0, 1 * sizeof(float), stream)); + kPreconditionOptimizer32bit1State<<>>( + g, p, state1, unorm, beta1, beta2, eps, weight_decay, step, lr, gnorm_scale, n + ); BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); } break; @@ -160,7 +162,7 @@ template void optimizerStatic8bitBlockwise( T* p, T* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, - float weight_decay, const float gnorm_scale, bool skip_zeros, int n + float weight_decay, const float gnorm_scale, bool skip_zeros, int n, bnb_stream_t stream ) { int num_blocks = 0; @@ -170,7 +172,7 @@ void optimizerStatic8bitBlockwise( num_blocks = n / BLOCKSIZE_2STATE; num_blocks = n % BLOCKSIZE_2STATE == 0 ? num_blocks : num_blocks + 1; kOptimizerStatic8bit2StateBlockwise - <<>>( + <<>>( p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, weight_decay, gnorm_scale, skip_zeros, n ); @@ -183,7 +185,7 @@ void optimizerStatic8bitBlockwise( num_blocks = n / BLOCKSIZE_1STATE; num_blocks = n % BLOCKSIZE_1STATE == 0 ? num_blocks : num_blocks + 1; kOptimizerStatic8bit1StateBlockwise - <<>>( + <<>>( p, g, state1, beta1, beta2, eps, step, lr, quantiles1, absmax1, weight_decay, gnorm_scale, skip_zeros, n ); BNB_CHECK_RETURN(BNB_PEEK_LAST_ERROR()); @@ -568,7 +570,7 @@ template void dequantizeBlockwise( gtype * g, gtype * p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ const float beta1, const float beta2, const float beta3, const float alpha, const float eps, \ const float weight_decay, const int step, const float lr, const float gnorm_scale, const bool skip_zeros, \ - const int n \ + const int n, bnb_stream_t stream \ ); MAKE_optimizer32bit(ADAM, half) MAKE_optimizer32bit(ADAM, float) MAKE_optimizer32bit(ADAM, bnb_bfloat16) MAKE_optimizer32bit(MOMENTUM, half) MAKE_optimizer32bit(MOMENTUM, float) MAKE_optimizer32bit(MOMENTUM, bnb_bfloat16) MAKE_optimizer32bit(RMSPROP, half) MAKE_optimizer32bit(RMSPROP, float) MAKE_optimizer32bit(RMSPROP, bnb_bfloat16) MAKE_optimizer32bit( @@ -579,7 +581,7 @@ MAKE_optimizer32bit(ADAM, half) MAKE_optimizer32bit(ADAM, float) MAKE_optimizer3 template void optimizerStatic8bitBlockwise( \ gtype * p, gtype * g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, \ float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, \ - float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n \ + float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n, bnb_stream_t stream \ ); MAKE_optimizerStatic8bitBlockwise(half, ADAM); diff --git a/csrc/ops.cuh b/csrc/ops.cuh index c7114bcaa..ed1c01adb 100644 --- a/csrc/ops.cuh +++ b/csrc/ops.cuh @@ -105,14 +105,14 @@ template void optimizer32bit( T* g, T* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, float beta1, float beta2, float beta3, float alpha, float eps, float weight_decay, int step, float lr, const float gnorm_scale, - bool skip_zeros, int n + bool skip_zeros, int n, bnb_stream_t stream ); template void optimizerStatic8bitBlockwise( T* p, T* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, float* absmax2, - float weight_decay, const float gnorm_scale, bool skip_zeros, int n + float weight_decay, const float gnorm_scale, bool skip_zeros, int n, bnb_stream_t stream ); void gemmex( diff --git a/csrc/pythonInterface.cpp b/csrc/pythonInterface.cpp index 24431603b..78c39fba0 100644 --- a/csrc/pythonInterface.cpp +++ b/csrc/pythonInterface.cpp @@ -73,7 +73,18 @@ MAKE_ELEMENTWISE_FUNC(_mul, fp32, float, _MUL) ) { \ optimizer32bit( \ g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, \ - lr, gnorm_scale, skip_zeros, n \ + lr, gnorm_scale, skip_zeros, n, nullptr \ + ); \ + } \ + void fname##32bit_grad_##gbits##_with_stream( \ + gtype* g, gtype* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ + const float beta1, const float beta2, const float beta3, const float alpha, const float eps, \ + const float weight_decay, const int step, const float lr, float gnorm_scale, bool skip_zeros, const int n, \ + bnb_stream_t stream \ + ) { \ + optimizer32bit( \ + g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, \ + lr, gnorm_scale, skip_zeros, n, stream \ ); \ } @@ -101,7 +112,17 @@ MAKE_FUNC32(ademamix, ADEMAMIX, bnb_bfloat16, bf16) ) { \ optimizerStatic8bitBlockwise( \ p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, \ - weight_decay, gnorm_scale, skip_zeros, n \ + weight_decay, gnorm_scale, skip_zeros, n, nullptr \ + ); \ + } \ + void fname##_8bit_blockwise_grad_##gbits##_with_stream( \ + gtype* p, gtype* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, \ + float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, \ + float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n, bnb_stream_t stream \ + ) { \ + optimizerStatic8bitBlockwise( \ + p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, \ + weight_decay, gnorm_scale, skip_zeros, n, stream \ ); \ } @@ -454,6 +475,17 @@ void cdequantize_blockwise_bf16_nf4( g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, \ lr, gnorm_scale, skip_zeros, n \ ); \ + } \ + void c##name##32bit_grad_##gbits##_with_stream( \ + gtype* g, gtype* p, float* state1, float* state2, float* unorm, float max_unorm, float param_norm, \ + const float beta1, const float beta2, const float beta3, const float alpha, const float eps, \ + const float weight_decay, const int step, const float lr, const float gnorm_scale, bool skip_zeros, \ + const int n, bnb_stream_t stream \ + ) { \ + name##32bit_grad_##gbits##_with_stream( \ + g, p, state1, state2, unorm, max_unorm, param_norm, beta1, beta2, beta3, alpha, eps, weight_decay, step, \ + lr, gnorm_scale, skip_zeros, n, stream \ + ); \ } MAKE_CFUNC32(adam, float, fp32) @@ -482,6 +514,16 @@ MAKE_CFUNC32(ademamix, bnb_bfloat16, bf16) p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, \ weight_decay, gnorm_scale, skip_zeros, n \ ); \ + } \ + void c##fname##_8bit_blockwise_grad_##gbits##_with_stream( \ + gtype* p, gtype* g, unsigned char* state1, unsigned char* state2, float beta1, float beta2, float beta3, \ + float alpha, float eps, int step, float lr, float* quantiles1, float* quantiles2, float* absmax1, \ + float* absmax2, float weight_decay, const float gnorm_scale, bool skip_zeros, int n, bnb_stream_t stream \ + ) { \ + fname##_8bit_blockwise_grad_##gbits##_with_stream( \ + p, g, state1, state2, beta1, beta2, beta3, alpha, eps, step, lr, quantiles1, quantiles2, absmax1, absmax2, \ + weight_decay, gnorm_scale, skip_zeros, n, stream \ + ); \ } MAKE_CBLOCKWISE8(adam, ADAM, half, fp16) diff --git a/tests/test_optimizer_stream.py b/tests/test_optimizer_stream.py new file mode 100644 index 000000000..8547e02a7 --- /dev/null +++ b/tests/test_optimizer_stream.py @@ -0,0 +1,252 @@ +import importlib + +import pytest +import torch + +import bitsandbytes as bnb + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is not None, + reason="requires NVIDIA CUDA", +) + + +def _optimizer_update_32bit( + g, + p, + state1, + state2, + *, + optimizer_name="adam", + unorm_vec=None, + max_unorm=0.0, + param_norm=0.0, +): + torch.ops.bitsandbytes.optimizer_update_32bit( + optimizer_name, + g, + p, + state1, + state2, + unorm_vec, + max_unorm, + param_norm, + 0.9, + 0.999, + 0.0, + 0.0, + 1e-8, + 0.0, + 1, + 1e-3, + 1.0, + False, + ) + + +def _optimizer_update_8bit(g, p, state1, state2, qmap1, qmap2, absmax1, absmax2): + torch.ops.bitsandbytes.optimizer_update_8bit_blockwise( + "adam", + g, + p, + state1, + state2, + 0.9, + 0.999, + 0.0, + 0.0, + 1e-8, + 1, + 1e-3, + qmap1, + qmap2, + absmax1, + absmax2, + 0.0, + 1.0, + False, + ) + + +def _make_inputs(family): + torch.manual_seed(0) + n = 4097 + p = torch.randn(n, device="cuda", dtype=torch.float32) + g = torch.zeros_like(p) + final_g = torch.randn_like(p) * 0.01 + if family == "32bit": + return [g, p, torch.zeros_like(p), torch.zeros_like(p)], final_g + + blocks = (n + 255) // 256 + qmap1 = bnb.functional.create_dynamic_map(signed=True).to("cuda") + qmap2 = bnb.functional.create_dynamic_map(signed=False).to("cuda") + return [ + g, + p, + torch.zeros(n, device="cuda", dtype=torch.uint8), + torch.zeros(n, device="cuda", dtype=torch.uint8), + qmap1, + qmap2, + torch.zeros(blocks, device="cuda", dtype=torch.float32), + torch.zeros(blocks, device="cuda", dtype=torch.float32), + ], final_g + + +def _run_update(family, values): + if family == "32bit": + _optimizer_update_32bit(*values) + else: + _optimizer_update_8bit(*values) + + +@pytest.mark.parametrize("family", ["32bit", "8bit"]) +def test_optimizer_update_uses_current_stream(family): + for _ in range(2): + inputs, final_g = _make_inputs(family) + reference = [value.clone() for value in inputs] + reference[0].copy_(final_g) + _run_update(family, reference) + torch.cuda.synchronize() + + blocker = torch.cuda.Stream() + caller = torch.cuda.Stream() + gate = torch.cuda.Event() + done = torch.cuda.Event() + with torch.cuda.stream(blocker): + torch.cuda._sleep(20_000_000) + gate.record() + assert not gate.query() + + with torch.cuda.stream(caller): + caller.wait_event(gate) + inputs[0].copy_(final_g) + _run_update(family, inputs) + done.record() + done.synchronize() + + for actual, expected in zip(inputs, reference): + assert torch.equal(actual, expected) + + +def test_optimizer_stream_symbols_are_additive_and_selected(): + from bitsandbytes.backends.cuda import ops as cuda_ops + from bitsandbytes.cextension import lib + + symbol_pairs = ( + ("cadam32bit_grad_fp32", "cadam32bit_grad_fp32_with_stream"), + ("clion32bit_grad_fp16", "clion32bit_grad_fp16_with_stream"), + ("cadam_8bit_blockwise_grad_fp32", "cadam_8bit_blockwise_grad_fp32_with_stream"), + ("clion_8bit_blockwise_grad_bf16", "clion_8bit_blockwise_grad_bf16_with_stream"), + ) + for legacy, stream_aware in symbol_pairs: + assert getattr(lib, legacy) is not None + assert getattr(lib, stream_aware) is not None + + assert cuda_ops.str2optimizer32bit["adam"][0] is lib.cadam32bit_grad_fp32_with_stream + assert cuda_ops.str2optimizer8bit_blockwise["adam"][0] is lib.cadam_8bit_blockwise_grad_fp32_with_stream + + +def test_paged_optimizer_preserves_default_stream_prefetch_order(monkeypatch): + from bitsandbytes.backends.cuda import ops as cuda_ops + + p = torch.nn.Parameter(torch.randn(100_000, device="cuda")) + optimizer = bnb.optim.PagedAdamW32bit([p], lr=1e-3) + p.grad = torch.randn_like(p) + optimizer.step() + + state = optimizer.state[p] + assert getattr(state["state1"], "is_paged", False) + assert getattr(state["state2"], "is_paged", False) + + calls = [] + real_prefetch = bnb.functional.prefetch_tensor + real_optimizers = cuda_ops.str2optimizer32bit["adam"] + + def recording_prefetch(tensor): + calls.append(("prefetch", tensor)) + real_prefetch(tensor) + + def recording_update(*args): + calls.append(("update", args[-1].value)) + return real_optimizers[0](*args) + + monkeypatch.setattr(bnb.functional, "prefetch_tensor", recording_prefetch) + monkeypatch.setitem(cuda_ops.str2optimizer32bit, "adam", (recording_update, *real_optimizers[1:])) + + p.grad = torch.randn_like(p) + torch.cuda.synchronize() + caller = torch.cuda.Stream() + with torch.cuda.stream(caller): + optimizer.step() + + assert [kind for kind, _ in calls] == ["prefetch", "prefetch", "update"] + assert calls[-1][1] is None + + +def test_optimizer_max_unorm_uses_current_stream(): + torch.manual_seed(0) + n = 4097 + p = torch.randn(n, device="cuda", dtype=torch.float32) + g = torch.zeros_like(p) + final_g = torch.randn_like(p) * 0.01 + values = [g, p, torch.zeros_like(p), torch.zeros_like(p), torch.zeros(1, device="cuda")] + reference = [value.clone() for value in values] + param_norm = p.norm().item() + + reference[0].copy_(final_g) + _optimizer_update_32bit( + *reference[:4], + optimizer_name="lamb", + unorm_vec=reference[4], + max_unorm=0.5, + param_norm=param_norm, + ) + torch.cuda.synchronize() + + blocker = torch.cuda.Stream() + caller = torch.cuda.Stream() + gate = torch.cuda.Event() + done = torch.cuda.Event() + with torch.cuda.stream(blocker): + torch.cuda._sleep(20_000_000) + gate.record() + assert not gate.query() + + with torch.cuda.stream(caller): + caller.wait_event(gate) + values[0].copy_(final_g) + _optimizer_update_32bit( + *values[:4], + optimizer_name="lamb", + unorm_vec=values[4], + max_unorm=0.5, + param_norm=param_norm, + ) + done.record() + done.synchronize() + + assert torch.equal(values[0], reference[0]) + for actual, expected in zip(values[1:], reference[1:]): + torch.testing.assert_close(actual, expected, atol=1e-6, rtol=1e-5) + + +@pytest.mark.parametrize("paged,expected_syncs", [(False, 0), (True, 3)]) +def test_optimizer_step_sync_policy(monkeypatch, paged, expected_syncs): + optimizer_module = importlib.import_module("bitsandbytes.optim.optimizer") + real_sync = optimizer_module.sync_gpu + calls = [] + + def recording_sync(tensor): + calls.append(tensor) + real_sync(tensor) + + monkeypatch.setattr(optimizer_module, "sync_gpu", recording_sync) + params = [torch.nn.Parameter(torch.randn(512, device="cuda")) for _ in range(2)] + optimizer_cls = bnb.optim.PagedAdamW32bit if paged else bnb.optim.AdamW32bit + optimizer = optimizer_cls(params, lr=1e-3) + for param in params: + param.grad = torch.randn_like(param) + + optimizer.step() + torch.cuda.synchronize() + assert len(calls) == expected_syncs