From 069498748f8abef3bd7b8b69898b7b9d83783e61 Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 17:19:42 +0800 Subject: [PATCH 1/7] feat(optimizer): add rank-1 AdamW4bit normalization --- areno/accel/csrc/extension.cpp | 45 +- areno/accel/csrc/optimizer.cu | 358 ++++++++++++- areno/accel/optimizer.py | 173 ++++++- areno/engine/optim/adamw_4bit.py | 617 +++++++++++++++++++++-- areno/engine/training.py | 8 +- docs/cli/training.rst | 5 +- docs/reference/adamw-4bit.rst | 22 + tests/test_adamw_4bit_cpu.py | 258 +++++++++- tests/test_fp32_master_optimizer_cuda.py | 32 ++ 9 files changed, 1442 insertions(+), 76 deletions(-) diff --git a/areno/accel/csrc/extension.cpp b/areno/accel/csrc/extension.cpp index 20d7f689..e5ccdc1f 100644 --- a/areno/accel/csrc/extension.cpp +++ b/areno/accel/csrc/extension.cpp @@ -193,7 +193,42 @@ void areno_adamw_4bit_step_cuda( torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, int64_t packed_offset, - int64_t scale_offset, + int64_t moment_scale_offset, + int64_t variance_scale_offset, + int64_t quant_block_size, + double beta1, + double beta2, + double effective_lr, + double weight_decay, + double eps, + double step_size, + double bias_correction2_sqrt); +void areno_adamw_4bit_rank1_stats_cuda( + torch::Tensor grad, + torch::Tensor exp_avg_sq_q, + torch::Tensor previous_scales, + torch::Tensor updated_scales, + torch::Tensor invalid, + torch::Tensor shape, + torch::Tensor strides, + int64_t packed_offset, + int64_t parameter_shard_start, + int64_t quant_block_size, + double beta2); +void areno_adamw_4bit_rank1_step_cuda( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg_q, + torch::Tensor exp_avg_scale, + torch::Tensor exp_avg_sq_q, + torch::Tensor previous_scales, + torch::Tensor updated_scales, + torch::Tensor invalid, + torch::Tensor shape, + torch::Tensor strides, + int64_t packed_offset, + int64_t moment_scale_offset, + int64_t parameter_shard_start, int64_t quant_block_size, double beta1, double beta2, @@ -235,6 +270,14 @@ void areno_adamw_fp32_state_step_cuda( PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("areno_adamw_fp32_master_step", &areno_adamw_fp32_master_step_cuda, "ARENO compact FP32-master AdamW step"); m.def("areno_adamw_4bit_step", &areno_adamw_4bit_step_cuda, "ARENO packed block-wise AdamW4bit step"); + m.def( + "areno_adamw_4bit_rank1_stats", + &areno_adamw_4bit_rank1_stats_cuda, + "ARENO rank-1 AdamW4bit statistics pass"); + m.def( + "areno_adamw_4bit_rank1_step", + &areno_adamw_4bit_rank1_step_cuda, + "ARENO rank-1 AdamW4bit update pass"); m.def("areno_adamw_8bit_step", &areno_adamw_8bit_step_cuda, "ARENO block-wise 8-bit AdamW step"); m.def("areno_adamw_fp32_state_step", &areno_adamw_fp32_state_step_cuda, "ARENO FP32-state AdamW step"); m.def("areno_silu_and_mul", &areno_silu_and_mul_cuda, "ARENO SiLU and multiply"); diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu index 85e38f31..4004c23d 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -31,6 +31,28 @@ __device__ __forceinline__ uint8_t nearest_signed_dynamic_code(float normalized) return best; } +__device__ __forceinline__ float atomic_max_nonnegative(float* address, float value) { + int* address_as_int = reinterpret_cast(address); + const int old = atomicMax(address_as_int, __float_as_int(value)); + return __int_as_float(old); +} + +__device__ __forceinline__ float rank1_scale( + const float* axis_scales, + const int64_t* shape, + const int64_t* strides, + int64_t ndim, + int64_t flat_index) { + float scale = CUDART_INF_F; + int64_t axis_offset = 0; + for (int64_t axis = 0; axis < ndim; ++axis) { + const int64_t coordinate = (flat_index / strides[axis]) % shape[axis]; + scale = fminf(scale, axis_scales[axis_offset + coordinate]); + axis_offset += shape[axis]; + } + return scale; +} + __device__ __forceinline__ uint8_t nearest_dynamic_code(float value, const float* codebook) { int lower = 0; int upper = 255; @@ -112,7 +134,8 @@ __global__ void adamw_4bit_kernel( float* exp_avg_sq_scale, int64_t numel, int64_t packed_offset, - int64_t scale_offset, + int64_t moment_scale_offset, + int64_t variance_scale_offset, float beta1, float beta2, float effective_lr, @@ -143,9 +166,10 @@ __global__ void adamw_4bit_kernel( if (active) { const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); - const int64_t block_scale_index = scale_offset + blockIdx.x; - moment = kSigned4bitDynamicMap[moment_code] * exp_avg_scale[block_scale_index]; - variance = (static_cast(variance_code) + 1.0f) * exp_avg_sq_scale[block_scale_index] / 16.0f; + const int64_t moment_scale_index = moment_scale_offset + blockIdx.x; + const int64_t variance_scale_index = variance_scale_offset + blockIdx.x; + moment = kSigned4bitDynamicMap[moment_code] * exp_avg_scale[moment_scale_index]; + variance = (static_cast(variance_code) + 1.0f) * exp_avg_sq_scale[variance_scale_index] / 16.0f; const float gradient = load_grad(grad, local_index); updated_weight = load_model(model, local_index); if (weight_decay != 0.0f) { @@ -176,8 +200,8 @@ __global__ void adamw_4bit_kernel( if (tid == 0) { new_moment_scale = moment_max[0]; new_variance_scale = variance_max[0]; - exp_avg_scale[scale_offset + blockIdx.x] = new_moment_scale; - exp_avg_sq_scale[scale_offset + blockIdx.x] = new_variance_scale; + exp_avg_scale[moment_scale_offset + blockIdx.x] = new_moment_scale; + exp_avg_sq_scale[variance_scale_offset + blockIdx.x] = new_variance_scale; } __syncthreads(); @@ -201,6 +225,158 @@ __global__ void adamw_4bit_kernel( } } +template +__global__ void adamw_4bit_rank1_stats_kernel( + const grad_t* grad, + const uint8_t* exp_avg_sq_q, + const float* previous_scales, + float* updated_scales, + int* invalid, + const int64_t* shape, + const int64_t* strides, + int64_t ndim, + int64_t numel, + int64_t packed_offset, + int64_t parameter_shard_start, + float beta2) { + __shared__ float variance_values[1024]; + __shared__ int invalid_block; + const int tid = threadIdx.x; + const int64_t local_index = static_cast(blockIdx.x) * blockDim.x + tid; + const bool active = local_index < numel; + if (tid == 0) { + invalid_block = 0; + } + __syncthreads(); + float variance = 0.0f; + if (active) { + const int64_t parameter_index = parameter_shard_start + local_index; + const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); + const uint8_t code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + variance = (static_cast(code) + 1.0f) * old_scale / 16.0f; + const float gradient = load_grad(grad, local_index); + variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; + if (!isfinite(gradient) || !isfinite(variance)) { + atomicExch(&invalid_block, 1); + } + } + variance_values[tid] = variance; + __syncthreads(); + if (invalid_block != 0) { + if (tid == 0) { + atomicExch(invalid, 1); + } + return; + } + if (active) { + const int64_t parameter_index = parameter_shard_start + local_index; + int64_t axis_offset = 0; + for (int64_t axis = 0; axis < ndim; ++axis) { + const int64_t coordinate = (parameter_index / strides[axis]) % shape[axis]; + atomic_max_nonnegative(updated_scales + axis_offset + coordinate, variance_values[tid]); + axis_offset += shape[axis]; + } + } +} + +template +__global__ void adamw_4bit_rank1_step_kernel( + model_t* model, + const grad_t* grad, + uint8_t* exp_avg_q, + float* exp_avg_scale, + uint8_t* exp_avg_sq_q, + const float* previous_scales, + const float* updated_scales, + const int* invalid, + const int64_t* shape, + const int64_t* strides, + int64_t ndim, + int64_t numel, + int64_t packed_offset, + int64_t moment_scale_offset, + int64_t parameter_shard_start, + float beta1, + float beta2, + float effective_lr, + float weight_decay, + float eps, + float step_size, + float bias_correction2_sqrt) { + __shared__ float moment_max[1024]; + __shared__ uint8_t moment_codes[1024]; + __shared__ uint8_t variance_codes[1024]; + __shared__ float new_moment_scale; + __shared__ int invalid_block; + const int tid = threadIdx.x; + if (*invalid != 0) { + return; + } + const int64_t local_index = static_cast(blockIdx.x) * blockDim.x + tid; + const bool active = local_index < numel; + if (tid == 0) { + invalid_block = 0; + } + __syncthreads(); + float moment = 0.0f; + float variance = 0.0f; + float updated_weight = 0.0f; + if (active) { + const int64_t parameter_index = parameter_shard_start + local_index; + const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); + const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + moment = kSigned4bitDynamicMap[moment_code] * exp_avg_scale[moment_scale_offset + blockIdx.x]; + const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); + variance = (static_cast(variance_code) + 1.0f) * old_scale / 16.0f; + const float gradient = load_grad(grad, local_index); + updated_weight = load_model(model, local_index); + if (weight_decay != 0.0f) { + updated_weight *= 1.0f - effective_lr * weight_decay; + } + moment = beta1 * moment + (1.0f - beta1) * gradient; + variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; + const float denom = sqrtf(variance) / bias_correction2_sqrt + eps; + updated_weight -= step_size * moment / denom; + if (!isfinite(gradient) || !isfinite(moment) || !isfinite(variance) || !isfinite(updated_weight)) { + atomicExch(&invalid_block, 1); + } + } + moment_max[tid] = active ? fabsf(moment) : 0.0f; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + moment_max[tid] = fmaxf(moment_max[tid], moment_max[tid + stride]); + } + __syncthreads(); + } + if (invalid_block != 0) { + return; + } + if (tid == 0) { + new_moment_scale = moment_max[0]; + exp_avg_scale[moment_scale_offset + blockIdx.x] = new_moment_scale; + } + __syncthreads(); + if (active) { + const int64_t parameter_index = parameter_shard_start + local_index; + moment_codes[tid] = nearest_signed_dynamic_code(moment / fmaxf(new_moment_scale, 1.0e-30f)); + const float new_scale = rank1_scale(updated_scales, shape, strides, ndim, parameter_index); + int variance_code = __float2int_rn(variance / fmaxf(new_scale, 1.0e-30f) * 16.0f - 1.0f); + variance_code = variance_code < 0 ? 0 : (variance_code > 15 ? 15 : variance_code); + variance_codes[tid] = static_cast(variance_code); + store_model(model, local_index, updated_weight); + } else { + moment_codes[tid] = 7; + variance_codes[tid] = 0; + } + __syncthreads(); + if ((tid & 1) == 0 && local_index < numel) { + const int64_t byte_index = packed_offset + (local_index >> 1); + exp_avg_q[byte_index] = moment_codes[tid] | static_cast(moment_codes[tid + 1] << 4); + exp_avg_sq_q[byte_index] = variance_codes[tid] | static_cast(variance_codes[tid + 1] << 4); + } +} + template __global__ void adamw_bf16_master_kernel( at::BFloat16* model, @@ -377,7 +553,8 @@ void launch_adamw_4bit( torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, int64_t packed_offset, - int64_t scale_offset, + int64_t moment_scale_offset, + int64_t variance_scale_offset, int64_t quant_block_size, float beta1, float beta2, @@ -397,7 +574,90 @@ void launch_adamw_4bit( exp_avg_sq_scale.data_ptr(), model.numel(), packed_offset, - scale_offset, + moment_scale_offset, + variance_scale_offset, + beta1, + beta2, + effective_lr, + weight_decay, + eps, + step_size, + bias_correction2_sqrt); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void launch_adamw_4bit_rank1_stats( + torch::Tensor grad, + torch::Tensor exp_avg_sq_q, + torch::Tensor previous_scales, + torch::Tensor updated_scales, + torch::Tensor invalid, + torch::Tensor shape, + torch::Tensor strides, + int64_t packed_offset, + int64_t parameter_shard_start, + int64_t quant_block_size, + float beta2) { + const int blocks = static_cast((grad.numel() + quant_block_size - 1) / quant_block_size); + const auto stream = at::cuda::getCurrentCUDAStream(); + adamw_4bit_rank1_stats_kernel<<(quant_block_size), 0, stream>>>( + grad.data_ptr(), + exp_avg_sq_q.data_ptr(), + previous_scales.data_ptr(), + updated_scales.data_ptr(), + invalid.data_ptr(), + shape.data_ptr(), + strides.data_ptr(), + shape.numel(), + grad.numel(), + packed_offset, + parameter_shard_start, + beta2); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void launch_adamw_4bit_rank1_step( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg_q, + torch::Tensor exp_avg_scale, + torch::Tensor exp_avg_sq_q, + torch::Tensor previous_scales, + torch::Tensor updated_scales, + torch::Tensor invalid, + torch::Tensor shape, + torch::Tensor strides, + int64_t packed_offset, + int64_t moment_scale_offset, + int64_t parameter_shard_start, + int64_t quant_block_size, + float beta1, + float beta2, + float effective_lr, + float weight_decay, + float eps, + float step_size, + float bias_correction2_sqrt) { + const int blocks = static_cast((model.numel() + quant_block_size - 1) / quant_block_size); + const auto stream = at::cuda::getCurrentCUDAStream(); + adamw_4bit_rank1_step_kernel<<(quant_block_size), 0, stream>>>( + model.data_ptr(), + grad.data_ptr(), + exp_avg_q.data_ptr(), + exp_avg_scale.data_ptr(), + exp_avg_sq_q.data_ptr(), + previous_scales.data_ptr(), + updated_scales.data_ptr(), + invalid.data_ptr(), + shape.data_ptr(), + strides.data_ptr(), + shape.numel(), + model.numel(), + packed_offset, + moment_scale_offset, + parameter_shard_start, beta1, beta2, effective_lr, @@ -659,7 +919,8 @@ void areno_adamw_4bit_step_cuda( torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, int64_t packed_offset, - int64_t scale_offset, + int64_t moment_scale_offset, + int64_t variance_scale_offset, int64_t quant_block_size, double beta1, double beta2, @@ -682,7 +943,8 @@ void areno_adamw_4bit_step_cuda( #define LAUNCH_ADAMW4(MODEL_T, GRAD_T) \ launch_adamw_4bit( \ - model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, packed_offset, scale_offset, \ + model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, packed_offset, \ + moment_scale_offset, variance_scale_offset, \ quant_block_size, beta1, beta2, effective_lr, weight_decay, eps, step_size, bias_correction2_sqrt) if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kBFloat16) { @@ -699,6 +961,82 @@ void areno_adamw_4bit_step_cuda( #undef LAUNCH_ADAMW4 } +void areno_adamw_4bit_rank1_stats_cuda( + torch::Tensor grad, + torch::Tensor exp_avg_sq_q, + torch::Tensor previous_scales, + torch::Tensor updated_scales, + torch::Tensor invalid, + torch::Tensor shape, + torch::Tensor strides, + int64_t packed_offset, + int64_t parameter_shard_start, + int64_t quant_block_size, + double beta2) { + c10::cuda::CUDAGuard guard(grad.device()); + TORCH_CHECK( + grad.is_cuda() && exp_avg_sq_q.is_cuda() && previous_scales.is_cuda() && updated_scales.is_cuda() && + invalid.is_cuda() && shape.is_cuda() && strides.is_cuda(), + "AdamW4bit rank-1 statistics tensors must be CUDA tensors"); + TORCH_CHECK( + quant_block_size >= 32 && quant_block_size <= 1024 && + (quant_block_size & (quant_block_size - 1)) == 0, + "AdamW4bit block size must be a power of two between 32 and 1024"); + if (grad.scalar_type() == at::kBFloat16) { + launch_adamw_4bit_rank1_stats( + grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides, packed_offset, + parameter_shard_start, quant_block_size, beta2); + } else if (grad.scalar_type() == at::kFloat) { + launch_adamw_4bit_rank1_stats( + grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides, packed_offset, + parameter_shard_start, quant_block_size, beta2); + } else { + TORCH_CHECK(false, "AdamW4bit rank-1 gradient must be bfloat16 or float32"); + } +} + +void areno_adamw_4bit_rank1_step_cuda( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg_q, + torch::Tensor exp_avg_scale, + torch::Tensor exp_avg_sq_q, + torch::Tensor previous_scales, + torch::Tensor updated_scales, + torch::Tensor invalid, + torch::Tensor shape, + torch::Tensor strides, + int64_t packed_offset, + int64_t moment_scale_offset, + int64_t parameter_shard_start, + int64_t quant_block_size, + double beta1, + double beta2, + double effective_lr, + double weight_decay, + double eps, + double step_size, + double bias_correction2_sqrt) { + c10::cuda::CUDAGuard guard(model.device()); +#define LAUNCH_ADAMW4_RANK1(MODEL_T, GRAD_T) \ + launch_adamw_4bit_rank1_step( \ + model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, previous_scales, updated_scales, invalid, \ + shape, strides, packed_offset, moment_scale_offset, parameter_shard_start, quant_block_size, beta1, \ + beta2, effective_lr, weight_decay, eps, step_size, bias_correction2_sqrt) + if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kBFloat16) { + LAUNCH_ADAMW4_RANK1(at::BFloat16, at::BFloat16); + } else if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW4_RANK1(at::BFloat16, float); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kBFloat16) { + LAUNCH_ADAMW4_RANK1(float, at::BFloat16); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW4_RANK1(float, float); + } else { + TORCH_CHECK(false, "AdamW4bit rank-1 model and gradient must be bfloat16 or float32"); + } +#undef LAUNCH_ADAMW4_RANK1 +} + void areno_adamw_8bit_step_cuda( torch::Tensor model, torch::Tensor grad, diff --git a/areno/accel/optimizer.py b/areno/accel/optimizer.py index a339131b..aeed65ed 100644 --- a/areno/accel/optimizer.py +++ b/areno/accel/optimizer.py @@ -206,7 +206,8 @@ def areno_adamw_4bit_step( exp_avg_sq_scale: torch.Tensor, *, packed_offset: int, - scale_offset: int, + moment_scale_offset: int, + variance_scale_offset: int, quant_block_size: int, beta1: float, beta2: float, @@ -241,10 +242,10 @@ def areno_adamw_4bit_step( raise ValueError("packed AdamW4bit state slice is out of bounds") if exp_avg_q.numel() != exp_avg_sq_q.numel(): raise ValueError("packed AdamW4bit moments must have the same length") - if scale_offset < 0 or scale_offset + scale_numel > exp_avg_scale.numel(): - raise ValueError("AdamW4bit scale slice is out of bounds") - if exp_avg_scale.numel() != exp_avg_sq_scale.numel(): - raise ValueError("AdamW4bit scale tensors must have the same length") + if moment_scale_offset < 0 or moment_scale_offset + scale_numel > exp_avg_scale.numel(): + raise ValueError("AdamW4bit first-moment scale slice is out of bounds") + if variance_scale_offset < 0 or variance_scale_offset + scale_numel > exp_avg_sq_scale.numel(): + raise ValueError("AdamW4bit second-moment scale slice is out of bounds") extension().areno_adamw_4bit_step( model, grad, @@ -253,7 +254,165 @@ def areno_adamw_4bit_step( exp_avg_sq_q, exp_avg_sq_scale, packed_offset, - scale_offset, + moment_scale_offset, + variance_scale_offset, + quant_block_size, + beta1, + beta2, + effective_lr, + weight_decay, + eps, + step_size, + bias_correction2_sqrt, + ) + + +def _validate_rank1_metadata(shape: torch.Tensor, strides: torch.Tensor, scales: torch.Tensor) -> None: + if shape.dtype != torch.int64 or strides.dtype != torch.int64: + raise TypeError("AdamW4bit rank-1 shape and strides must use int64") + if shape.ndim != 1 or strides.ndim != 1 or shape.numel() != strides.numel() or shape.numel() < 2: + raise ValueError("AdamW4bit rank-1 metadata must describe a tensor with rank >= 2") + if scales.ndim != 1: + raise ValueError("AdamW4bit rank-1 scales must be flat") + + +@torch._dynamo.disable +@torch.no_grad() +def areno_adamw_4bit_rank1_stats( + grad: torch.Tensor, + exp_avg_sq_q: torch.Tensor, + previous_scales: torch.Tensor, + updated_scales: torch.Tensor, + invalid: torch.Tensor, + shape: torch.Tensor, + strides: torch.Tensor, + *, + packed_offset: int, + parameter_shard_start: int, + quant_block_size: int, + beta2: float, +) -> None: + """Accumulate updated rank-1 second-moment maxima in bounded CUDA blocks.""" + + tensors = (grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides) + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("fused AdamW4bit rank-1 statistics require CUDA tensors") + if any(tensor.device != grad.device for tensor in tensors[1:]): + raise ValueError("fused AdamW4bit rank-1 statistics require tensors on one device") + if grad.dtype not in {torch.bfloat16, torch.float32}: + raise TypeError("AdamW4bit rank-1 gradients must be bfloat16 or float32") + if exp_avg_sq_q.dtype != torch.uint8 or previous_scales.dtype != torch.float32: + raise TypeError("AdamW4bit rank-1 state must use packed uint8 codes and float32 scales") + if updated_scales.dtype != torch.float32 or invalid.dtype != torch.int32 or invalid.numel() != 1: + raise TypeError("AdamW4bit rank-1 outputs must use float32 scales and one int32 validity flag") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("fused AdamW4bit rank-1 statistics require contiguous tensors") + _validate_rank1_metadata(shape, strides, previous_scales) + if previous_scales.numel() != updated_scales.numel(): + raise ValueError("AdamW4bit previous and updated rank-1 scale layouts must match") + packed_numel = (grad.numel() + 1) // 2 + if packed_offset < 0 or packed_offset + packed_numel > exp_avg_sq_q.numel(): + raise ValueError("AdamW4bit rank-1 packed variance slice is out of bounds") + if parameter_shard_start < 0: + raise ValueError("AdamW4bit rank-1 parameter shard start must be non-negative") + extension().areno_adamw_4bit_rank1_stats( + grad, + exp_avg_sq_q, + previous_scales, + updated_scales, + invalid, + shape, + strides, + packed_offset, + parameter_shard_start, + quant_block_size, + beta2, + ) + + +@torch._dynamo.disable +@torch.no_grad() +def areno_adamw_4bit_rank1_step( + model: torch.Tensor, + grad: torch.Tensor, + exp_avg_q: torch.Tensor, + exp_avg_scale: torch.Tensor, + exp_avg_sq_q: torch.Tensor, + previous_scales: torch.Tensor, + updated_scales: torch.Tensor, + invalid: torch.Tensor, + shape: torch.Tensor, + strides: torch.Tensor, + *, + packed_offset: int, + moment_scale_offset: int, + parameter_shard_start: int, + quant_block_size: int, + beta1: float, + beta2: float, + effective_lr: float, + weight_decay: float, + eps: float, + step_size: float, + bias_correction2_sqrt: float, +) -> None: + """Update packed AdamW4bit state using precomputed rank-1 scales.""" + + tensors = ( + model, + grad, + exp_avg_q, + exp_avg_scale, + exp_avg_sq_q, + previous_scales, + updated_scales, + invalid, + shape, + strides, + ) + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("fused AdamW4bit rank-1 update requires CUDA tensors") + if any(tensor.device != model.device for tensor in tensors[1:]): + raise ValueError("fused AdamW4bit rank-1 update requires tensors on one device") + if model.dtype not in {torch.bfloat16, torch.float32} or grad.dtype not in {torch.bfloat16, torch.float32}: + raise TypeError("AdamW4bit rank-1 model and gradient must be bfloat16 or float32") + if exp_avg_q.dtype != torch.uint8 or exp_avg_sq_q.dtype != torch.uint8: + raise TypeError("AdamW4bit rank-1 moments must use packed uint8 storage") + if any(scale.dtype != torch.float32 for scale in (exp_avg_scale, previous_scales, updated_scales)): + raise TypeError("AdamW4bit rank-1 scales must use float32") + if invalid.dtype != torch.int32 or invalid.numel() != 1: + raise TypeError("AdamW4bit rank-1 update requires one int32 validity flag") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("fused AdamW4bit rank-1 update requires contiguous tensors") + if model.numel() != grad.numel(): + raise ValueError("AdamW4bit rank-1 model and gradient sizes must match") + _validate_rank1_metadata(shape, strides, previous_scales) + if previous_scales.numel() != updated_scales.numel(): + raise ValueError("AdamW4bit previous and updated rank-1 scale layouts must match") + packed_numel = (model.numel() + 1) // 2 + if packed_offset < 0 or packed_offset + packed_numel > exp_avg_q.numel(): + raise ValueError("AdamW4bit rank-1 packed moment slice is out of bounds") + if exp_avg_q.numel() != exp_avg_sq_q.numel(): + raise ValueError("AdamW4bit rank-1 packed moments must have the same length") + scale_numel = (model.numel() + quant_block_size - 1) // quant_block_size + if moment_scale_offset < 0 or moment_scale_offset + scale_numel > exp_avg_scale.numel(): + raise ValueError("AdamW4bit rank-1 first-moment scale slice is out of bounds") + if parameter_shard_start < 0: + raise ValueError("AdamW4bit rank-1 parameter shard start must be non-negative") + extension().areno_adamw_4bit_rank1_step( + model, + grad, + exp_avg_q, + exp_avg_scale, + exp_avg_sq_q, + previous_scales, + updated_scales, + invalid, + shape, + strides, + packed_offset, + moment_scale_offset, + parameter_shard_start, quant_block_size, beta1, beta2, @@ -267,6 +426,8 @@ def areno_adamw_4bit_step( __all__ = [ "areno_adamw_4bit_step", + "areno_adamw_4bit_rank1_stats", + "areno_adamw_4bit_rank1_step", "areno_adamw_8bit_step", "areno_adamw_fp32_master_step", "areno_adamw_fp32_state_step", diff --git a/areno/engine/optim/adamw_4bit.py b/areno/engine/optim/adamw_4bit.py index 943bca45..e5b734dd 100644 --- a/areno/engine/optim/adamw_4bit.py +++ b/areno/engine/optim/adamw_4bit.py @@ -1,23 +1,25 @@ -"""Packed block-wise 4-bit-state AdamW. +"""Packed rank-1-normalized 4-bit-state AdamW. The first moment uses signed dynamic-exponent quantization. The non-negative second moment uses the zero-excluding linear map from Li et al. (NeurIPS 2023): -code ``i`` represents ``scale * (i + 1) / 16``. Two codes are packed in each -byte and scales are stored per parameter-local block. +code ``i`` represents ``scale * (i + 1) / 16``. Matrix and higher-rank second +moments use the paper's rank-1 normalization; vectors retain B=128 block +normalization. Two codes are packed in each byte. """ from __future__ import annotations from collections.abc import Iterable, Iterator +from math import prod import torch import torch.distributed as dist from areno.engine.optim.adamw_8bit import AdamW8bit -from areno.engine.optim.adamw_fp32_master import _DEFAULT_BUCKET_NUMEL, _MasterBucket, _ParamRef +from areno.engine.optim.adamw_fp32_master import _DEFAULT_BUCKET_NUMEL, _MasterBucket, _param_grad, _ParamRef _DEFAULT_QUANT_BLOCK_SIZE = 128 -_STATE_FORMAT_VERSION = 1 +_STATE_FORMAT_VERSION = 2 # The signed 4-bit dynamic-exponent map used by the reference implementation # of Li et al. Values are normalized by each block's absolute maximum. _SIGNED_DE_MAP = ( @@ -41,16 +43,18 @@ class AdamW4bit(AdamW8bit): - """AdamW with two packed 4-bit moments and FP32 block scales. - - Quantization blocks restart at every parameter shard. This prevents one - tensor's outlier from setting another tensor's scale and keeps temporary - FP32 state bounded by ``quant_block_size`` on CPU. CUDA updates packed - state directly with a fused block-wise kernel. + """AdamW with two packed 4-bit moments and shape-aware FP32 scales. + + First-moment quantization blocks restart at every parameter shard. For a + tensor with rank >= 2, second-moment scales are the minimum of per-axis + maxima over the original local model-tensor shape. DP ranks combine their + partial axis maxima with MAX. One-dimensional tensors retain parameter- + local block scales. CPU and CUDA updates keep FP32 moment work bounded by + ``quant_block_size``. """ _embedding_fp32_state = False - state_quantizer = "signed-de4/zero-excluding-linear4" + state_quantizer = "signed-de4/rank1-zero-excluding-linear4" def _precision_for_parameter(self, parameter: torch.nn.Parameter) -> str: del parameter @@ -83,6 +87,10 @@ def __init__( dp_group=dp_group, quant_block_size=quant_block_size, ) + self._rank1_metadata_cache: dict[tuple[int, torch.device], tuple[torch.Tensor, torch.Tensor]] = {} + self._rank1_scales: dict[int, torch.Tensor | None] = { + id(parameter): None for parameter in self.model_params if parameter.ndim >= 2 + } def state_dict(self) -> dict: """Return the versioned packed state for this DP rank.""" @@ -100,6 +108,14 @@ def state_dict(self) -> dict: payload["adam_4bit"] = True payload["state_format_version"] = _STATE_FORMAT_VERSION payload["quant_block_size"] = self.quant_block_size + payload["parameter_shapes"] = [tuple(parameter.shape) for parameter in self.model_params] + payload["rank1_scales"] = [ + None + if self._rank1_scales.get(id(parameter)) is None + else self._rank1_scales[id(parameter)].detach().to(device="cpu").clone() + for parameter in self.model_params + ] + payload["state_memory"] = self.state_memory_metrics() return payload @torch.no_grad() @@ -114,6 +130,13 @@ def load_state_dict(self, state_dict: dict) -> None: raise ValueError( f"AdamW4bit quant_block_size mismatch: checkpoint={saved_block_size}, optimizer={self.quant_block_size}" ) + saved_shapes = [tuple(shape) for shape in state_dict.get("parameter_shapes", ())] + current_shapes = [tuple(parameter.shape) for parameter in self.model_params] + if saved_shapes != current_shapes: + raise ValueError("AdamW4bit checkpoint parameter shapes do not match the optimizer") + saved_rank1_scales = state_dict.get("rank1_scales") + if not isinstance(saved_rank1_scales, list) or len(saved_rank1_scales) != len(self.model_params): + raise ValueError("AdamW4bit checkpoint rank-1 scales do not match the optimizer parameters") self._cleanup_disk_offload() self._active_offload_mode = "none" self._disk_offload_root = None @@ -124,16 +147,63 @@ def load_state_dict(self, state_dict: dict) -> None: state.offload_group = None state.offload_ready_events = () saved_states = state_dict.get("state", []) + if len(saved_states) != len(self.buckets): + raise ValueError( + "AdamW4bit checkpoint bucket count does not match the optimizer layout: " + f"checkpoint={len(saved_states)}, optimizer={len(self.buckets)}" + ) for saved, bucket, state in zip(saved_states[: len(self.buckets)], self.buckets, self._states, strict=False): if saved is None: continue device = bucket.refs[0].model_param.device - packed_numel, scale_numel = self._bucket_state_sizes(bucket) + packed_numel, moment_scale_numel, variance_scale_numel = self._bucket_state_sizes(bucket) state.step = int(saved.get("step", 0)) state.exp_avg_q = _load_tensor(saved, "exp_avg_q", device, torch.uint8, packed_numel) - state.exp_avg_scale = _load_tensor(saved, "exp_avg_scale", device, torch.float32, scale_numel) + state.exp_avg_scale = _load_tensor(saved, "exp_avg_scale", device, torch.float32, moment_scale_numel) state.exp_avg_sq_q = _load_tensor(saved, "exp_avg_sq_q", device, torch.uint8, packed_numel) - state.exp_avg_sq_scale = _load_tensor(saved, "exp_avg_sq_scale", device, torch.float32, scale_numel) + state.exp_avg_sq_scale = _load_tensor( + saved, "exp_avg_sq_scale", device, torch.float32, variance_scale_numel + ) + for parameter, saved_scales in zip(self.model_params, saved_rank1_scales, strict=True): + if parameter.ndim < 2: + if saved_scales is not None: + raise ValueError("AdamW4bit checkpoint has rank-1 scales for a one-dimensional parameter") + continue + if saved_scales is None: + self._rank1_scales[id(parameter)] = None + continue + restored_scales = saved_scales.detach().to(device=parameter.device, dtype=torch.float32).view(-1).clone() + if restored_scales.numel() != sum(parameter.shape): + raise ValueError("AdamW4bit checkpoint rank-1 scale length does not match the parameter shape") + self._rank1_scales[id(parameter)] = restored_scales + + def clear_state(self) -> None: + """Drop packed moments and parameter-level rank-1 metadata.""" + + super().clear_state() + for parameter_id in self._rank1_scales: + self._rank1_scales[parameter_id] = None + self._rank1_metadata_cache.clear() + + @torch.no_grad() + def offload_state(self, mode: str = "cpu", directory: str | None = None, batch_size: int = 1) -> None: + """Offload packed buckets and keep small rank-1 metadata on CPU.""" + + super().offload_state(mode=mode, directory=directory, batch_size=batch_size) + for parameter_id, scales in self._rank1_scales.items(): + if scales is not None and scales.device.type != "cpu": + self._rank1_scales[parameter_id] = scales.to(device="cpu") + self._rank1_metadata_cache.clear() + + @torch.no_grad() + def onload_state(self, device: torch.device) -> None: + """Restore packed buckets and shared rank-1 metadata to ``device``.""" + + super().onload_state(device) + for parameter_id, scales in self._rank1_scales.items(): + if scales is not None and scales.device != device: + self._rank1_scales[parameter_id] = scales.to(device=device) + self._rank1_metadata_cache.clear() @torch.no_grad() def _ensure_bucket_state(self, bucket: _MasterBucket, state) -> None: @@ -146,32 +216,142 @@ def _ensure_bucket_state(self, bucket: _MasterBucket, state) -> None: value = getattr(state, name) if value is not None and value.device != device: setattr(state, name, value.to(device=device)) - packed_numel, scale_numel = self._bucket_state_sizes(bucket) + packed_numel, moment_scale_numel, variance_scale_numel = self._bucket_state_sizes(bucket) if state.exp_avg_q is None: # Signed dynamic-exponent zero has code 7, hence byte 0x77. state.exp_avg_q = torch.full((packed_numel,), 0x77, device=device, dtype=torch.uint8) - state.exp_avg_scale = torch.ones(scale_numel, device=device, dtype=torch.float32) + state.exp_avg_scale = torch.ones(moment_scale_numel, device=device, dtype=torch.float32) if state.exp_avg_sq_q is None: state.exp_avg_sq_q = torch.zeros(packed_numel, device=device, dtype=torch.uint8) # A zero scale makes the zero-excluding code initially decode to 0. - state.exp_avg_sq_scale = torch.zeros(scale_numel, device=device, dtype=torch.float32) + state.exp_avg_sq_scale = torch.zeros(variance_scale_numel, device=device, dtype=torch.float32) def _state_mmap_specs(self, indices: list[int]) -> dict[int, dict[str, tuple[torch.dtype, tuple[int, ...]]]]: """Return fixed raw-mmap layouts for packed state and block scales.""" specs: dict[int, dict[str, tuple[torch.dtype, tuple[int, ...]]]] = {} for index in indices: - packed_numel, scale_numel = self._bucket_state_sizes(self.buckets[index]) + packed_numel, moment_scale_numel, variance_scale_numel = self._bucket_state_sizes(self.buckets[index]) specs[index] = { "exp_avg_q": (torch.uint8, (packed_numel,)), - "exp_avg_scale": (torch.float32, (scale_numel,)), + "exp_avg_scale": (torch.float32, (moment_scale_numel,)), "exp_avg_sq_q": (torch.uint8, (packed_numel,)), - "exp_avg_sq_scale": (torch.float32, (scale_numel,)), + "exp_avg_sq_scale": (torch.float32, (variance_scale_numel,)), } return specs @torch.no_grad() - def _step_bucket_8bit(self, bucket: _MasterBucket, state) -> None: + def step(self, closure=None): + """Apply a streaming statistics pass before rank-1 parameter updates.""" + + if closure is not None: + with torch.enable_grad(): + closure() + active_rank1_parameters = { + id(ref.model_param) + for bucket in self.buckets + for ref in bucket.refs + if _uses_rank1_normalization(ref) and self._ref_has_gradient(bucket, ref) + } + if not active_rank1_parameters: + return super().step() + + rank1_work: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} + for parameter in self.model_params: + if id(parameter) in active_rank1_parameters: + rank1_work[id(parameter)] = ( + torch.zeros(_rank1_scale_numel_for_parameter(parameter), device=parameter.device), + torch.zeros((), device=parameter.device, dtype=torch.int32), + ) + self._ensure_rank1_scales(parameter) + + # First streaming pass: each chunk contributes to one parameter-level + # set of axis maxima. Disk-offloaded packed state is returned to disk + # after each group, so this pass does not pin all buckets on the GPU. + for indices in self._bucket_groups(): + group_changed = False + for index in indices: + bucket = self.buckets[index] + state = self._states[index] + if not any( + _uses_rank1_normalization(ref) and self._ref_has_gradient(bucket, ref) for ref in bucket.refs + ): + continue + self._ensure_bucket_state(bucket, state) + for ref, packed_offset, _moment_scale_offset, _variance_scale_offset in self._iter_ref_layout(bucket): + if not _uses_rank1_normalization(ref) or not self._ref_has_gradient(bucket, ref): + continue + updated_scales, invalid = rank1_work[id(ref.model_param)] + self._rank1_variance_statistics( + bucket, + ref, + self._gradient_for_ref(bucket, ref), + state, + packed_offset, + self._ensure_rank1_scales(ref.model_param), + updated_scales, + invalid, + self.betas[1], + ) + group_changed = True + if self._active_offload_mode == "disk": + self._stage_8bit_state_on_cpu(state) + if self._active_offload_mode == "disk" and group_changed: + self._offload_8bit_group_to_disk(indices) + + for parameter in self.model_params: + work = rank1_work.get(id(parameter)) + if work is None: + continue + updated_scales, invalid = work + if self.dp_size > 1: + if self.dp_group is None: + raise RuntimeError("AdamW4bit DP rank-1 normalization requires a DP process group") + dist.all_reduce(updated_scales, op=dist.ReduceOp.MAX, group=self.dp_group) + dist.all_reduce(invalid, op=dist.ReduceOp.MAX, group=self.dp_group) + + # Second streaming pass: recompute each bounded block, update weights, + # and write packed moments using the now-global axis scales. + for indices in self._bucket_groups(): + group_changed = False + for index in indices: + bucket = self.buckets[index] + state = self._states[index] + has_grad = bucket.grad_shard is not None or any( + _param_grad(ref.model_param) is not None for ref in bucket.refs + ) + if has_grad: + self._ensure_bucket_state(bucket, state) + self._step_bucket_8bit(bucket, state, rank1_work) + group_changed = True + if self._active_offload_mode == "disk": + self._stage_8bit_state_on_cpu(state) + elif self._active_offload_mode == "disk": + self._discard_disk_prefetch(index) + if self._active_offload_mode == "disk" and group_changed: + self._offload_8bit_group_to_disk(indices) + + for parameter in self.model_params: + work = rank1_work.get(id(parameter)) + if work is None: + continue + updated_scales, invalid = work + scale_storage = self._ensure_rank1_scales(parameter) + if invalid.is_cuda: + scale_storage.copy_(torch.where(invalid == 0, updated_scales, scale_storage)) + elif int(invalid.item()) == 0: + scale_storage.copy_(updated_scales) + if self._active_offload_mode == "disk" and scale_storage.device.type != "cpu": + self._rank1_scales[id(parameter)] = scale_storage.to(device="cpu") + return None + + @torch.no_grad() + def _step_bucket_8bit( + self, + bucket: _MasterBucket, + state, + rank1_work: dict[int, tuple[torch.Tensor, torch.Tensor]] | None = None, + ) -> None: """Update a bucket while materializing at most one FP32 block per moment.""" assert state.exp_avg_q is not None @@ -182,24 +362,53 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state) -> None: state.step += 1 bias_correction1 = 1.0 - beta1**state.step bias_correction2_sqrt = (1.0 - beta2**state.step) ** 0.5 - for ref, packed_offset, scale_offset in self._iter_ref_layout(bucket): - grad = self._gradient_for_ref(bucket, ref) - if grad is None: + for ref, packed_offset, moment_scale_offset, variance_scale_offset in self._iter_ref_layout(bucket): + has_parameter_grad = ( + id(ref.model_param) in bucket.grad_param_ids + if bucket.grad_shard is not None + else ref.model_param.grad is not None + or isinstance(getattr(ref.model_param, "main_grad", None), torch.Tensor) + ) + if not has_parameter_grad: continue + grad = self._gradient_for_ref(bucket, ref) effective_lr = float(getattr(ref.model_param, "_areno_lr", self.lr)) - self._step_param_ref_4bit( - bucket, - ref, - grad, - state, - packed_offset, - scale_offset, - beta1, - beta2, - effective_lr, - effective_lr / bias_correction1, - bias_correction2_sqrt, - ) + if _uses_rank1_normalization(ref): + if rank1_work is None: + raise RuntimeError("AdamW4bit rank-1 update is missing its parameter statistics pass") + updated_scales, invalid = rank1_work[id(ref.model_param)] + if grad is not None and (invalid.is_cuda or int(invalid.item()) == 0): + self._step_param_ref_rank1( + bucket, + ref, + grad, + state, + packed_offset, + moment_scale_offset, + self._ensure_rank1_scales(ref.model_param), + updated_scales, + invalid, + beta1, + beta2, + effective_lr, + effective_lr / bias_correction1, + bias_correction2_sqrt, + ) + elif grad is not None: + self._step_param_ref_4bit( + bucket, + ref, + grad, + state, + packed_offset, + moment_scale_offset, + variance_scale_offset, + beta1, + beta2, + effective_lr, + effective_lr / bias_correction1, + bias_correction2_sqrt, + ) if ref.param_start + ref.numel == ref.model_param.numel(): ref.model_param.grad = None if isinstance(getattr(ref.model_param, "main_grad", None), torch.Tensor): @@ -208,6 +417,12 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state) -> None: bucket.grad_shard = None bucket.grad_param_ids = frozenset() + @staticmethod + def _ref_has_gradient(bucket: _MasterBucket, ref: _ParamRef) -> bool: + if bucket.grad_shard is not None: + return id(ref.model_param) in bucket.grad_param_ids + return _param_grad(ref.model_param) is not None + @torch.no_grad() def _step_param_ref_4bit( self, @@ -216,7 +431,8 @@ def _step_param_ref_4bit( grad: torch.Tensor, state, packed_offset: int, - scale_offset: int, + moment_scale_offset: int, + variance_scale_offset: int, beta1: float, beta2: float, effective_lr: float, @@ -241,7 +457,8 @@ def _step_param_ref_4bit( state.exp_avg_sq_q, state.exp_avg_sq_scale, packed_offset=packed_offset, - scale_offset=scale_offset, + moment_scale_offset=moment_scale_offset, + variance_scale_offset=variance_scale_offset, quant_block_size=self.quant_block_size, beta1=beta1, beta2=beta2, @@ -256,16 +473,17 @@ def _step_param_ref_4bit( count = min(self.quant_block_size, ref.shard_numel - start) byte_start = packed_offset + start // 2 byte_count = (count + 1) // 2 - scale_index = scale_offset + block_index + moment_scale_index = moment_scale_offset + block_index + variance_scale_index = variance_scale_offset + block_index moment = _unpack_signed_4bit( state.exp_avg_q.narrow(0, byte_start, byte_count), count, - state.exp_avg_scale[scale_index], + state.exp_avg_scale[moment_scale_index], ) variance = _unpack_positive_4bit( state.exp_avg_sq_q.narrow(0, byte_start, byte_count), count, - state.exp_avg_sq_scale[scale_index], + state.exp_avg_sq_scale[variance_scale_index], ) grad_block = grad_shard.narrow(0, start, count).to(dtype=torch.float32) weight = model_shard.narrow(0, start, count).to(dtype=torch.float32) @@ -288,22 +506,215 @@ def _step_param_ref_4bit( moment_q, moment_scale = _quantize_signed_4bit(moment) variance_q, variance_scale = _quantize_positive_4bit(variance) state.exp_avg_q.narrow(0, byte_start, byte_count).copy_(moment_q) - state.exp_avg_scale[scale_index].copy_(moment_scale) + state.exp_avg_scale[moment_scale_index].copy_(moment_scale) state.exp_avg_sq_q.narrow(0, byte_start, byte_count).copy_(variance_q) - state.exp_avg_sq_scale[scale_index].copy_(variance_scale) + state.exp_avg_sq_scale[variance_scale_index].copy_(variance_scale) + + @torch.no_grad() + def _rank1_variance_statistics( + self, + bucket: _MasterBucket, + ref: _ParamRef, + grad: torch.Tensor | None, + state, + packed_offset: int, + previous_scales: torch.Tensor, + updated_scales: torch.Tensor, + invalid: torch.Tensor, + beta2: float, + ) -> None: + """Compute updated per-axis maxima without a full FP32 moment.""" + + assert state.exp_avg_sq_q is not None + if grad is None or ref.shard_numel == 0: + return + grad_shard = grad if bucket.grad_shard is not None else grad.narrow(0, ref.shard_start, ref.shard_numel) + parameter_shard_start = ref.param_start + ref.shard_start + if grad_shard.is_cuda: + from areno.accel.optimizer import areno_adamw_4bit_rank1_stats + + shape, strides = self._rank1_metadata(ref) + areno_adamw_4bit_rank1_stats( + grad_shard.contiguous(), + state.exp_avg_sq_q, + previous_scales, + updated_scales, + invalid, + shape, + strides, + packed_offset=packed_offset, + parameter_shard_start=parameter_shard_start, + quant_block_size=self.quant_block_size, + beta2=beta2, + ) + return + + shape_tuple = tuple(ref.model_param.shape) + for start in range(0, ref.shard_numel, self.quant_block_size): + count = min(self.quant_block_size, ref.shard_numel - start) + byte_start = packed_offset + start // 2 + byte_count = (count + 1) // 2 + flat_start = parameter_shard_start + start + element_scales = _rank1_element_scales(previous_scales, shape_tuple, flat_start, count) + variance = _unpack_positive_4bit_elementwise( + state.exp_avg_sq_q.narrow(0, byte_start, byte_count), count, element_scales + ) + gradient = grad_shard.narrow(0, start, count).to(dtype=torch.float32) + variance.mul_(beta2).addcmul_(gradient, gradient, value=1.0 - beta2) + if not bool(torch.isfinite(gradient).all() & torch.isfinite(variance).all()): + invalid.fill_(1) + return + _accumulate_rank1_maxima(updated_scales, variance, shape_tuple, flat_start) - def _bucket_state_sizes(self, bucket: _MasterBucket) -> tuple[int, int]: + @torch.no_grad() + def _step_param_ref_rank1( + self, + bucket: _MasterBucket, + ref: _ParamRef, + grad: torch.Tensor, + state, + packed_offset: int, + moment_scale_offset: int, + previous_scales: torch.Tensor, + updated_scales: torch.Tensor, + invalid: torch.Tensor, + beta1: float, + beta2: float, + effective_lr: float, + step_size: float, + bias_correction2_sqrt: float, + ) -> None: + """Recompute bounded Adam blocks and requantize with rank-1 scales.""" + + assert state.exp_avg_q is not None + assert state.exp_avg_scale is not None + assert state.exp_avg_sq_q is not None + if ref.shard_numel == 0: + return + grad_shard = grad if bucket.grad_shard is not None else grad.narrow(0, ref.shard_start, ref.shard_numel) + parameter_shard_start = ref.param_start + ref.shard_start + model_shard = ref.model_param.detach().reshape(-1).narrow(0, parameter_shard_start, ref.shard_numel) + if model_shard.is_cuda: + from areno.accel.optimizer import areno_adamw_4bit_rank1_step + + shape, strides = self._rank1_metadata(ref) + areno_adamw_4bit_rank1_step( + model_shard, + grad_shard.contiguous(), + state.exp_avg_q, + state.exp_avg_scale, + state.exp_avg_sq_q, + previous_scales, + updated_scales, + invalid, + shape, + strides, + packed_offset=packed_offset, + moment_scale_offset=moment_scale_offset, + parameter_shard_start=parameter_shard_start, + quant_block_size=self.quant_block_size, + beta1=beta1, + beta2=beta2, + effective_lr=effective_lr, + weight_decay=self.weight_decay, + eps=self.eps, + step_size=step_size, + bias_correction2_sqrt=bias_correction2_sqrt, + ) + return + + shape_tuple = tuple(ref.model_param.shape) + for block_index, start in enumerate(range(0, ref.shard_numel, self.quant_block_size)): + count = min(self.quant_block_size, ref.shard_numel - start) + byte_start = packed_offset + start // 2 + byte_count = (count + 1) // 2 + moment_scale_index = moment_scale_offset + block_index + moment = _unpack_signed_4bit( + state.exp_avg_q.narrow(0, byte_start, byte_count), + count, + state.exp_avg_scale[moment_scale_index], + ) + flat_start = parameter_shard_start + start + old_element_scales = _rank1_element_scales(previous_scales, shape_tuple, flat_start, count) + variance = _unpack_positive_4bit_elementwise( + state.exp_avg_sq_q.narrow(0, byte_start, byte_count), count, old_element_scales + ) + gradient = grad_shard.narrow(0, start, count).to(dtype=torch.float32) + weight = model_shard.narrow(0, start, count).to(dtype=torch.float32).clone() + if self.weight_decay != 0.0: + weight.mul_(1.0 - effective_lr * self.weight_decay) + moment.mul_(beta1).add_(gradient, alpha=1.0 - beta1) + variance.mul_(beta2).addcmul_(gradient, gradient, value=1.0 - beta2) + denom = variance.sqrt().div_(bias_correction2_sqrt).add_(self.eps) + weight.addcdiv_(moment, denom, value=-step_size) + if not bool( + torch.isfinite(gradient).all() + & torch.isfinite(moment).all() + & torch.isfinite(variance).all() + & torch.isfinite(weight).all() + ): + return + new_element_scales = _rank1_element_scales(updated_scales, shape_tuple, flat_start, count) + moment_q, moment_scale = _quantize_signed_4bit(moment) + variance_q = _quantize_positive_4bit_elementwise(variance, new_element_scales) + model_shard.narrow(0, start, count).copy_(weight) + state.exp_avg_q.narrow(0, byte_start, byte_count).copy_(moment_q) + state.exp_avg_scale[moment_scale_index].copy_(moment_scale) + state.exp_avg_sq_q.narrow(0, byte_start, byte_count).copy_(variance_q) + + def _rank1_metadata(self, ref: _ParamRef) -> tuple[torch.Tensor, torch.Tensor]: + key = (id(ref.model_param), ref.model_param.device) + cached = self._rank1_metadata_cache.get(key) + if cached is not None: + return cached + shape_tuple = tuple(int(dimension) for dimension in ref.model_param.shape) + strides_tuple = _contiguous_strides(shape_tuple) + cached = ( + torch.tensor(shape_tuple, device=ref.model_param.device, dtype=torch.int64), + torch.tensor(strides_tuple, device=ref.model_param.device, dtype=torch.int64), + ) + self._rank1_metadata_cache[key] = cached + return cached + + def _ensure_rank1_scales(self, parameter: torch.nn.Parameter) -> torch.Tensor: + """Materialize one shared axis-scale tensor for an original parameter.""" + + key = id(parameter) + scales = self._rank1_scales.get(key) + if scales is None: + scales = torch.zeros( + _rank1_scale_numel_for_parameter(parameter), + device=parameter.device, + dtype=torch.float32, + ) + self._rank1_scales[key] = scales + elif scales.device != parameter.device: + scales = scales.to(device=parameter.device) + self._rank1_scales[key] = scales + return scales + + def _bucket_state_sizes(self, bucket: _MasterBucket) -> tuple[int, int, int]: packed_numel = sum((ref.shard_numel + 1) // 2 for ref in bucket.refs) - scale_numel = sum((ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size for ref in bucket.refs) - return packed_numel, scale_numel + moment_scale_numel = sum( + (ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size for ref in bucket.refs + ) + variance_scale_numel = sum(self._variance_scale_count(ref) for ref in bucket.refs) + return packed_numel, moment_scale_numel, variance_scale_numel - def _iter_ref_layout(self, bucket: _MasterBucket) -> Iterator[tuple[_ParamRef, int, int]]: + def _iter_ref_layout(self, bucket: _MasterBucket) -> Iterator[tuple[_ParamRef, int, int, int]]: packed_offset = 0 - scale_offset = 0 + moment_scale_offset = 0 + variance_scale_offset = 0 for ref in bucket.refs: - yield ref, packed_offset, scale_offset + yield ref, packed_offset, moment_scale_offset, variance_scale_offset packed_offset += (ref.shard_numel + 1) // 2 - scale_offset += (ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size + moment_scale_offset += (ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size + variance_scale_offset += self._variance_scale_count(ref) + + def _variance_scale_count(self, ref: _ParamRef) -> int: + if _uses_rank1_normalization(ref): + return 0 + return (ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size def persistent_moment_bytes(self) -> int: """Return resident packed-moment and scale storage in bytes.""" @@ -313,8 +724,34 @@ def persistent_moment_bytes(self) -> int: for value in (state.exp_avg_q, state.exp_avg_scale, state.exp_avg_sq_q, state.exp_avg_sq_scale): if value is not None: total += value.numel() * value.element_size() + for value in self._rank1_scales.values(): + if value is not None: + total += value.numel() * value.element_size() return total + def state_memory_metrics(self) -> dict[str, int]: + """Report actual packed moments and shape/block metadata bytes.""" + + quantized_state_bytes = 0 + scale_metadata_bytes = 0 + for state in self._states: + if state.step == 0: + continue + for value in (state.exp_avg_q, state.exp_avg_sq_q): + if value is not None: + quantized_state_bytes += value.numel() * value.element_size() + for value in (state.exp_avg_scale, state.exp_avg_sq_scale): + if value is not None: + scale_metadata_bytes += value.numel() * value.element_size() + for value in self._rank1_scales.values(): + if value is not None: + scale_metadata_bytes += value.numel() * value.element_size() + return { + "quantized_state_bytes": quantized_state_bytes, + "scale_metadata_bytes": scale_metadata_bytes, + "total_bytes": quantized_state_bytes + scale_metadata_bytes, + } + def _load_tensor( saved: dict, @@ -384,4 +821,80 @@ def _unpack_positive_4bit(packed: torch.Tensor, numel: int, scale: torch.Tensor) return (_unpack_nibbles(packed, numel).to(dtype=torch.float32) + 1.0).mul_(scale / 16.0) +def _uses_rank1_normalization(ref: _ParamRef) -> bool: + return ref.model_param.ndim >= 2 + + +def _rank1_scale_numel(ref: _ParamRef) -> int: + return _rank1_scale_numel_for_parameter(ref.model_param) + + +def _rank1_scale_numel_for_parameter(parameter: torch.nn.Parameter) -> int: + return sum(int(dimension) for dimension in parameter.shape) + + +def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: + strides: list[int] = [] + for axis in range(len(shape)): + strides.append(prod(shape[axis + 1 :])) + return tuple(strides) + + +def _rank1_element_scales( + axis_scales: torch.Tensor, + shape: tuple[int, ...], + flat_start: int, + count: int, +) -> torch.Tensor: + """Expand paper rank-1 statistics for only one bounded flat slice.""" + + if count == 0: + return axis_scales.new_empty((0,)) + flat_indices = torch.arange(flat_start, flat_start + count, device=axis_scales.device) + result = torch.full((count,), torch.inf, device=axis_scales.device, dtype=torch.float32) + axis_offset = 0 + for dimension, stride in zip(shape, _contiguous_strides(shape), strict=True): + coordinates = torch.div(flat_indices, stride, rounding_mode="floor").remainder_(dimension) + torch.minimum( + result, + axis_scales.narrow(0, axis_offset, dimension)[coordinates], + out=result, + ) + axis_offset += dimension + return result + + +def _accumulate_rank1_maxima( + axis_maxima: torch.Tensor, + values: torch.Tensor, + shape: tuple[int, ...], + flat_start: int, +) -> None: + """Accumulate per-axis maxima for a bounded flat slice.""" + + flat_indices = torch.arange(flat_start, flat_start + values.numel(), device=values.device) + axis_offset = 0 + for dimension, stride in zip(shape, _contiguous_strides(shape), strict=True): + coordinates = torch.div(flat_indices, stride, rounding_mode="floor").remainder_(dimension) + axis_maxima.narrow(0, axis_offset, dimension).scatter_reduce_( + 0, coordinates, values, reduce="amax", include_self=True + ) + axis_offset += dimension + + +def _unpack_positive_4bit_elementwise( + packed: torch.Tensor, + numel: int, + scales: torch.Tensor, +) -> torch.Tensor: + codes = _unpack_nibbles(packed, numel).to(dtype=torch.float32) + return (codes + 1.0).mul_(scales / 16.0) + + +def _quantize_positive_4bit_elementwise(tensor: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + safe_scales = scales.clamp_min(1.0e-30) + codes = torch.clamp(torch.round(tensor / safe_scales * 16.0 - 1.0), 0.0, 15.0).to(dtype=torch.uint8) + return _pack_nibbles(codes) + + __all__ = ["AdamW4bit"] diff --git a/areno/engine/training.py b/areno/engine/training.py index 3aac40af..569fb7e9 100644 --- a/areno/engine/training.py +++ b/areno/engine/training.py @@ -187,11 +187,11 @@ def _train_step( multimodal_lrs = self._set_multimodal_lrs(worker._global_step + 1) worker.optimizer.step() state_memory_metrics = getattr(worker.optimizer, "state_memory_metrics", None) - if ( - callable(state_memory_metrics) - and getattr(worker.optimizer, "state_quantizer", None) == "dynamic-tree-v1" - ): + state_quantizer = getattr(worker.optimizer, "state_quantizer", None) + if callable(state_memory_metrics) and state_quantizer == "dynamic-tree-v1": optimizer_state_metrics = {f"adam8_{name}": value for name, value in state_memory_metrics().items()} + elif callable(state_memory_metrics) and str(state_quantizer).startswith("signed-de4/"): + optimizer_state_metrics = {f"adam4_{name}": value for name, value in state_memory_metrics().items()} worker.optimizer.zero_grad(set_to_none=True) worker._global_step += 1 if worker.adapter_registry is not None: diff --git a/docs/cli/training.rst b/docs/cli/training.rst index 81d2e46f..f8c284d8 100644 --- a/docs/cli/training.rst +++ b/docs/cli/training.rst @@ -398,8 +398,9 @@ in its description; flags for other algorithms are ignored. backends; validate convergence when changing optimizer precision. ``--adam-4bit`` - Use packed block-wise 4-bit Adam moment states. This option is CUDA-only - and cannot be combined with ``--adam-8bit``. See + Use packed 4-bit Adam moment states with rank-1 second-moment normalization + for tensors of rank two or greater and B128 fallback for vectors. This + option is CUDA-only and cannot be combined with ``--adam-8bit``. See :doc:`../reference/adamw-4bit` for complete usage examples. ``--unfreeze-mm-tower`` diff --git a/docs/reference/adamw-4bit.rst b/docs/reference/adamw-4bit.rst index bb81e33c..fe997208 100644 --- a/docs/reference/adamw-4bit.rst +++ b/docs/reference/adamw-4bit.rst @@ -5,6 +5,23 @@ AReno provides an opt-in packed 4-bit AdamW optimizer for CUDA training. It changes optimizer-state storage only; the model checkpoint and training data format are unchanged. +State representation +-------------------- + +The first moment uses parameter-local blocks and the signed dynamic-exponent +4-bit map. For matrix and higher-rank parameters, the second moment uses +rank-1 normalization from *Memory Efficient Optimizers with 4-bit States*: +AReno records the maximum along every coordinate of every axis and scales each +element by the minimum applicable axis statistic. One-dimensional parameters +use 128-element block normalization. The second-moment 4-bit map excludes zero. + +For data-parallel training, axis statistics are defined over the original +parameter shape and combined with ``MAX`` across the DP group. Tensor-parallel +parameters use each rank's local model-tensor shape; this does not add a TP +collective. The fused update uses bounded block-local FP32 work and does not +materialize a parameter-sized FP32 moment tensor. The internal block size is +configurable for programmatic users and defaults to 128. + Command line ------------ @@ -129,6 +146,11 @@ Requirements and errors fused 4-bit optimizer kernel. * A saved 4-bit optimizer state must be resumed with the 4-bit optimizer. The separately saved model weights remain usable without ``--adam-4bit``. +* Optimizer-state checkpoints from the earlier block-only representation are + intentionally incompatible. Model-weight checkpoints remain portable. + +Initialized state is reported as ``adam4_quantized_state_bytes``, +``adam4_scale_metadata_bytes``, and ``adam4_total_bytes``. Confirm that the option is available with: diff --git a/tests/test_adamw_4bit_cpu.py b/tests/test_adamw_4bit_cpu.py index d340ab1b..aba9c280 100644 --- a/tests/test_adamw_4bit_cpu.py +++ b/tests/test_adamw_4bit_cpu.py @@ -1,10 +1,13 @@ from __future__ import annotations import copy +import multiprocessing as mp +import socket from pathlib import Path import pytest import torch +import torch.distributed as dist from click.testing import CliRunner from areno.api.trainer_config import TrainerConfig @@ -13,9 +16,13 @@ from areno.engine.modeling import build_optimizer from areno.engine.optim import AdamW4bit, AdamW8bit, AdamWFP32Master from areno.engine.optim.adamw_4bit import ( + _accumulate_rank1_maxima, _quantize_positive_4bit, + _quantize_positive_4bit_elementwise, _quantize_signed_4bit, + _rank1_element_scales, _unpack_positive_4bit, + _unpack_positive_4bit_elementwise, _unpack_signed_4bit, ) @@ -31,6 +38,48 @@ def _optimizer(param: torch.nn.Parameter, *, block_size: int = 128) -> AdamW4bit ) +def _free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return int(sock.getsockname()[1]) + + +def _gloo_rank1_worker(rank: int, port: int, output_queue) -> None: + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + rank=rank, + world_size=2, + ) + try: + parameter = torch.nn.Parameter(torch.zeros(5, 7)) + optimizer = AdamW4bit( + [parameter], + lr=3.0e-4, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=16, + quant_block_size=128, + dp_rank=rank, + dp_size=2, + dp_group=dist.group.WORLD, + ) + parameter.grad = torch.arange(1, 36, dtype=torch.float32).reshape_as(parameter) + rank * 3.0 + optimizer.reduce_scatter_gradients() + optimizer.step() + output_queue.put( + ( + rank, + parameter.detach().tolist(), + optimizer._rank1_scales[id(parameter)].tolist(), + optimizer.buckets[0].refs[0].shard_start, + optimizer.buckets[0].refs[0].shard_numel, + ) + ) + finally: + dist.destroy_process_group() + + def test_adamw4bit_packs_two_moments_within_storage_budget() -> None: param = torch.nn.Parameter(torch.zeros(8192)) optimizer = _optimizer(param) @@ -102,7 +151,7 @@ def test_adamw4bit_checkpoint_round_trip_preserves_next_update() -> None: restored.step() torch.testing.assert_close(restored_param, first_param, rtol=0.0, atol=0.0) - assert restored.state_dict()["state_format_version"] == 1 + assert restored.state_dict()["state_format_version"] == 2 def test_adamw4bit_disk_offload_preserves_update(tmp_path: Path) -> None: @@ -165,6 +214,213 @@ def test_adamw4bit_nonfinite_gradient_skips_only_affected_block() -> None: assert torch.all(parameter[128:] < 0) +@pytest.mark.parametrize( + "values", + [ + torch.tensor([[1.0, 1.0, 1.0], [9.0, 9.0, 9.0]]), + torch.tensor([[1.0, 2.0, 8.0], [1.0, 2.0, 8.0]]), + torch.tensor([[1.0, 8.0, 1.0], [8.0, 1.0, 8.0], [1.0, 8.0, 1.0]]), + torch.tensor([[1.0, 1.0, 1.0], [1.0, 1000.0, 1.0]]), + torch.zeros(3, 5), + torch.full((3, 5), 1.0e-20), + ], + ids=["row", "column", "checkerboard", "isolated-outlier", "all-zero", "tiny-positive"], +) +def test_rank1_statistics_match_paper_sm3_algorithm(values: torch.Tensor) -> None: + flattened = values.flatten() + statistics = torch.zeros(sum(values.shape)) + + for start in range(0, flattened.numel(), 4): + _accumulate_rank1_maxima(statistics, flattened[start : start + 4], tuple(values.shape), start) + + expected = torch.cat((values.amax(dim=1), values.amax(dim=0))) + torch.testing.assert_close(statistics, expected) + expanded = _rank1_element_scales(statistics, tuple(values.shape), 0, values.numel()).reshape(values.shape) + torch.testing.assert_close( + expanded, + torch.minimum(expected[: values.shape[0], None], expected[values.shape[0] :]), + ) + + +def test_rank1_statistics_generalize_to_higher_rank() -> None: + values = torch.arange(1, 31, dtype=torch.float32).reshape(2, 3, 5) + statistics = torch.zeros(sum(values.shape)) + _accumulate_rank1_maxima(statistics, values.flatten(), tuple(values.shape), 0) + + expected = torch.cat( + ( + values.amax(dim=(1, 2)), + values.amax(dim=(0, 2)), + values.amax(dim=(0, 1)), + ) + ) + torch.testing.assert_close(statistics, expected) + + +def test_adamw4bit_matrix_uses_rank1_second_moment_scales() -> None: + parameter = torch.nn.Parameter(torch.zeros(3, 5)) + optimizer = _optimizer(parameter) + gradient = torch.arange(1, 16, dtype=torch.float32).reshape_as(parameter) + parameter.grad = gradient.clone() + + optimizer.step() + + state = optimizer._states[0] + variance = (1.0 - optimizer.betas[1]) * gradient.square() + expected = torch.cat((variance.amax(dim=1), variance.amax(dim=0))) + rank1_scales = optimizer._rank1_scales[id(parameter)] + assert rank1_scales is not None + torch.testing.assert_close(rank1_scales, expected) + assert state.exp_avg_scale.numel() == 1 + assert state.exp_avg_sq_scale.numel() == 0 + + +def test_adamw4bit_rank1_second_moment_never_decodes_nonzero_scale_to_zero() -> None: + values = torch.tensor([[0.0, 0.01, 0.5], [0.02, 0.25, 1.0]]) + statistics = torch.cat((values.amax(dim=1), values.amax(dim=0))) + scales = _rank1_element_scales(statistics, tuple(values.shape), 0, values.numel()) + + packed = _quantize_positive_4bit_elementwise(values.flatten(), scales) + restored = _unpack_positive_4bit_elementwise(packed, values.numel(), scales) + + assert torch.all(restored[scales > 0] > 0) + + +def test_adamw4bit_preserves_chunking_and_shares_rank1_scales_across_chunks() -> None: + parameter = torch.nn.Parameter(torch.zeros(2048, 2049)) + optimizer = AdamW4bit( + [parameter], + lr=3.0e-4, + betas=(0.9, 0.99), + weight_decay=0.01, + bucket_numel=1, + quant_block_size=128, + ) + + refs = [ref for bucket in optimizer.buckets for ref in bucket.refs] + assert len(optimizer.buckets) == 2 + assert len(refs) == 2 + assert [ref.param_start for ref in refs] == [0, 4 * 1024 * 1024] + assert sum(ref.numel for ref in refs) == parameter.numel() + assert all(ref.model_param is parameter for ref in refs) + scales = optimizer._ensure_rank1_scales(parameter) + assert scales.numel() == sum(parameter.shape) + assert optimizer._rank1_scales[id(parameter)] is scales + + +def test_adamw4bit_rank1_metadata_stays_within_large_matrix_budget() -> None: + parameter = torch.nn.Parameter(torch.zeros(1024, 1024)) + optimizer = _optimizer(parameter) + optimizer._ensure_bucket_state(optimizer.buckets[0], optimizer._states[0]) + optimizer._ensure_rank1_scales(parameter) + optimizer._states[0].step = 1 + + assert optimizer.persistent_moment_bytes() / parameter.numel() <= 1.25 + metrics = optimizer.state_memory_metrics() + assert metrics["total_bytes"] == optimizer.persistent_moment_bytes() + assert metrics["scale_metadata_bytes"] == (parameter.numel() // 128 + sum(parameter.shape)) * 4 + + +def test_adamw4bit_rank1_checkpoint_round_trip_preserves_next_update() -> None: + initial = torch.linspace(-0.5, 0.5, 35).reshape(5, 7).to(torch.bfloat16) + first_parameter = torch.nn.Parameter(initial.clone()) + first = _optimizer(first_parameter) + first_parameter.grad = torch.linspace(-0.3, 0.7, 35).reshape_as(initial).to(torch.bfloat16) + first.step() + checkpoint = copy.deepcopy(first.state_dict()) + + restored_parameter = torch.nn.Parameter(first_parameter.detach().clone()) + restored = _optimizer(restored_parameter) + restored.load_state_dict(checkpoint) + next_gradient = torch.linspace(0.8, -0.4, 35).reshape_as(initial).to(torch.bfloat16) + first_parameter.grad = next_gradient.clone() + restored_parameter.grad = next_gradient.clone() + first.step() + restored.step() + + torch.testing.assert_close(restored_parameter, first_parameter, rtol=0.0, atol=0.0) + torch.testing.assert_close(restored._rank1_scales[id(restored_parameter)], first._rank1_scales[id(first_parameter)]) + + +def test_adamw4bit_clear_state_drops_rank1_scales() -> None: + parameter = torch.nn.Parameter(torch.zeros(5, 7)) + optimizer = _optimizer(parameter) + parameter.grad = torch.ones_like(parameter) + optimizer.step() + + assert optimizer._rank1_scales[id(parameter)] is not None + optimizer.clear_state() + + assert optimizer._rank1_scales[id(parameter)] is None + assert all(state.step == 0 and state.exp_avg_q is None for state in optimizer._states) + + +def test_adamw4bit_rank1_disk_offload_preserves_axis_scales(tmp_path: Path) -> None: + initial = torch.linspace(-0.5, 0.5, 35).reshape(5, 7).to(torch.bfloat16) + candidate_parameter = torch.nn.Parameter(initial.clone()) + reference_parameter = torch.nn.Parameter(initial.clone()) + candidate = _optimizer(candidate_parameter) + reference = _optimizer(reference_parameter) + candidate.configure_state_offload(mode="disk", directory=str(tmp_path), batch_size=1) + + for gradient in ( + torch.linspace(-0.4, 0.7, 35).reshape_as(initial), + torch.linspace(0.8, -0.2, 35).reshape_as(initial), + ): + candidate_parameter.grad = gradient.to(torch.bfloat16) + reference_parameter.grad = gradient.to(torch.bfloat16) + candidate.step() + reference.step() + + candidate.onload_state(torch.device("cpu")) + torch.testing.assert_close(candidate_parameter, reference_parameter, rtol=0.0, atol=0.0) + torch.testing.assert_close( + candidate._rank1_scales[id(candidate_parameter)], reference._rank1_scales[id(reference_parameter)] + ) + assert not list(tmp_path.rglob("*.mmap")) + + +def test_rank1_partial_shards_combine_to_unsharded_statistics() -> None: + values = torch.arange(1, 36, dtype=torch.float32).reshape(5, 7) + combined = torch.zeros(sum(values.shape)) + for start, count in ((0, 13), (13, 12), (25, 10)): + partial = torch.zeros_like(combined) + _accumulate_rank1_maxima(partial, values.flatten().narrow(0, start, count), tuple(values.shape), start) + torch.maximum(combined, partial, out=combined) + + expected = torch.cat((values.amax(dim=1), values.amax(dim=0))) + torch.testing.assert_close(combined, expected) + + +def test_real_gloo_rank1_statistics_match_unsharded_reference_across_split_row() -> None: + spawn = mp.get_context("spawn") + output_queue = spawn.Queue() + port = _free_port() + processes = [spawn.Process(target=_gloo_rank1_worker, args=(rank, port, output_queue)) for rank in range(2)] + for process in processes: + process.start() + try: + results = dict((item[0], item[1:]) for item in (output_queue.get(timeout=30) for _ in processes)) + finally: + for process in processes: + process.join(timeout=5) + if process.is_alive(): + process.terminate() + process.join(timeout=5) + assert all(process.exitcode == 0 for process in processes) + + rank0_model, rank0_scales, rank0_start, rank0_count = results[0] + rank1_model, rank1_scales, rank1_start, rank1_count = results[1] + averaged_gradient = torch.arange(1, 36, dtype=torch.float32).reshape(5, 7) + 1.5 + variance = 0.01 * averaged_gradient.square() + expected_scales = torch.cat((variance.amax(dim=1), variance.amax(dim=0))) + assert rank0_model == rank1_model + torch.testing.assert_close(torch.tensor(rank0_scales), expected_scales) + torch.testing.assert_close(torch.tensor(rank1_scales), expected_scales) + assert (rank0_start, rank0_count) == (0, 18) + assert (rank1_start, rank1_count) == (18, 17) + + def test_optimizer_config_rejects_multiple_low_bit_modes() -> None: with pytest.raises(ValueError, match="mutually exclusive"): OptimizerConfig(adam_4bit=True, adam_8bit=True) diff --git a/tests/test_fp32_master_optimizer_cuda.py b/tests/test_fp32_master_optimizer_cuda.py index 5e7853fe..640cd307 100644 --- a/tests/test_fp32_master_optimizer_cuda.py +++ b/tests/test_fp32_master_optimizer_cuda.py @@ -35,6 +35,38 @@ def test_fused_low_bit_adamw_matches_cpu_reference(optimizer_cls) -> None: torch.testing.assert_close(cuda_param.cpu(), cpu_param, rtol=0.0, atol=2.0e-3) +def test_fused_adamw4bit_rank1_matches_bounded_cpu_reference() -> None: + shape = (65, 67) + initial = torch.linspace(-1.0, 1.0, shape[0] * shape[1]).reshape(shape).to(torch.bfloat16) + cuda_param = torch.nn.Parameter(initial.cuda()) + cpu_param = torch.nn.Parameter(initial.clone()) + kwargs = { + "lr": 2.0e-4, + "betas": (0.9, 0.99), + "weight_decay": 0.02, + "bucket_numel": 1024, + "quant_block_size": 128, + } + cuda_optimizer = AdamW4bit([cuda_param], **kwargs) + cpu_optimizer = AdamW4bit([cpu_param], **kwargs) + + for step in range(4): + gradient = torch.sin(torch.linspace(-3.0, 2.0, initial.numel()) + step * 0.17).reshape(shape) + cuda_param.grad = gradient.to(torch.bfloat16).cuda() + cpu_param.grad = gradient.to(torch.bfloat16) + cuda_optimizer.step() + cpu_optimizer.step() + + torch.testing.assert_close(cuda_param.cpu(), cpu_param, rtol=0.0, atol=2.0e-3) + torch.testing.assert_close( + cuda_optimizer._states[0].exp_avg_sq_scale.cpu(), + cpu_optimizer._states[0].exp_avg_sq_scale, + rtol=2.0e-5, + atol=1.0e-7, + ) + assert cuda_optimizer._states[0].exp_avg_sq_scale.numel() == sum(shape) + + def test_fused_fp32_master_adamw_matches_torch_reference() -> None: device = torch.device("cuda", 0) initial = torch.linspace(-1.0, 1.0, 4099, device=device, dtype=torch.float32).to(torch.bfloat16) From 21039e2b02e3d26677e935cc7fe23230c462f80e Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 17:25:31 +0800 Subject: [PATCH 2/7] fix(cuda): remove attention build warning --- areno/accel/csrc/attention.cu | 1 - areno/accel/csrc/optimizer.cu | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/areno/accel/csrc/attention.cu b/areno/accel/csrc/attention.cu index f8a4de6c..8e177631 100644 --- a/areno/accel/csrc/attention.cu +++ b/areno/accel/csrc/attention.cu @@ -11,7 +11,6 @@ namespace { constexpr int kAttentionThreads = 256; -constexpr int kAttentionTileN = 16; int64_t attention_tile_n(int64_t head_dim) { if (head_dim <= 256) { diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu index 4004c23d..c38b031b 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -43,7 +43,7 @@ __device__ __forceinline__ float rank1_scale( const int64_t* strides, int64_t ndim, int64_t flat_index) { - float scale = CUDART_INF_F; + float scale = __int_as_float(0x7f800000); int64_t axis_offset = 0; for (int64_t axis = 0; axis < ndim; ++axis) { const int64_t coordinate = (flat_index / strides[axis]) % shape[axis]; From a963666a3c2b990aba4a95a2dc77826cf4a5af70 Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 17:45:52 +0800 Subject: [PATCH 3/7] perf(optimizer): reduce AdamW4bit CUDA workspace --- areno/accel/csrc/optimizer.cu | 164 +++++++++++++++++++++++----------- 1 file changed, 113 insertions(+), 51 deletions(-) diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu index c38b031b..2762c763 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -143,10 +143,15 @@ __global__ void adamw_4bit_kernel( float eps, float step_size, float bias_correction2_sqrt) { - __shared__ float moment_max[1024]; - __shared__ float variance_max[1024]; - __shared__ uint8_t moment_codes[1024]; - __shared__ uint8_t variance_codes[1024]; + constexpr int warp_size = 32; + constexpr int max_warps = 32; + __shared__ float warp_moment_maxima[max_warps]; + __shared__ float warp_variance_maxima[max_warps]; + extern __shared__ uint8_t packed_codes[]; + uint8_t* moment_codes = packed_codes; + uint8_t* variance_codes = packed_codes + blockDim.x; + __shared__ float old_moment_scale; + __shared__ float old_variance_scale; __shared__ float new_moment_scale; __shared__ float new_variance_scale; __shared__ int invalid_block; @@ -157,21 +162,20 @@ __global__ void adamw_4bit_kernel( const bool active = local_index < numel; if (tid == 0) { invalid_block = 0; + old_moment_scale = exp_avg_scale[moment_scale_offset + blockIdx.x]; + old_variance_scale = exp_avg_sq_scale[variance_scale_offset + blockIdx.x]; } __syncthreads(); - float moment = 0.0f; - float variance = 0.0f; - float updated_weight = 0.0f; + float local_moment_max = 0.0f; + float local_variance_max = 0.0f; if (active) { const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); - const int64_t moment_scale_index = moment_scale_offset + blockIdx.x; - const int64_t variance_scale_index = variance_scale_offset + blockIdx.x; - moment = kSigned4bitDynamicMap[moment_code] * exp_avg_scale[moment_scale_index]; - variance = (static_cast(variance_code) + 1.0f) * exp_avg_sq_scale[variance_scale_index] / 16.0f; + float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; + float variance = (static_cast(variance_code) + 1.0f) * old_variance_scale / 16.0f; const float gradient = load_grad(grad, local_index); - updated_weight = load_model(model, local_index); + float updated_weight = load_model(model, local_index); if (weight_decay != 0.0f) { updated_weight *= 1.0f - effective_lr * weight_decay; } @@ -182,30 +186,60 @@ __global__ void adamw_4bit_kernel( if (!isfinite(gradient) || !isfinite(moment) || !isfinite(variance) || !isfinite(updated_weight)) { atomicExch(&invalid_block, 1); } + local_moment_max = fabsf(moment); + local_variance_max = variance; } - moment_max[tid] = active ? fabsf(moment) : 0.0f; - variance_max[tid] = active ? variance : 0.0f; - __syncthreads(); - - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - moment_max[tid] = fmaxf(moment_max[tid], moment_max[tid + stride]); - variance_max[tid] = fmaxf(variance_max[tid], variance_max[tid + stride]); - } - __syncthreads(); + for (int offset = warp_size / 2; offset > 0; offset >>= 1) { + local_moment_max = fmaxf(local_moment_max, __shfl_down_sync(0xFFFFFFFFu, local_moment_max, offset)); + local_variance_max = + fmaxf(local_variance_max, __shfl_down_sync(0xFFFFFFFFu, local_variance_max, offset)); } + const int lane = tid & (warp_size - 1); + const int warp = tid / warp_size; + if (lane == 0) { + warp_moment_maxima[warp] = local_moment_max; + warp_variance_maxima[warp] = local_variance_max; + } + __syncthreads(); if (invalid_block != 0) { return; } - if (tid == 0) { - new_moment_scale = moment_max[0]; - new_variance_scale = variance_max[0]; - exp_avg_scale[moment_scale_offset + blockIdx.x] = new_moment_scale; - exp_avg_sq_scale[variance_scale_offset + blockIdx.x] = new_variance_scale; + if (warp == 0) { + const int warp_count = blockDim.x / warp_size; + float block_moment_max = lane < warp_count ? warp_moment_maxima[lane] : 0.0f; + float block_variance_max = lane < warp_count ? warp_variance_maxima[lane] : 0.0f; + for (int offset = warp_size / 2; offset > 0; offset >>= 1) { + block_moment_max = + fmaxf(block_moment_max, __shfl_down_sync(0xFFFFFFFFu, block_moment_max, offset)); + block_variance_max = + fmaxf(block_variance_max, __shfl_down_sync(0xFFFFFFFFu, block_variance_max, offset)); + } + if (lane == 0) { + new_moment_scale = block_moment_max; + new_variance_scale = block_variance_max; + exp_avg_scale[moment_scale_offset + blockIdx.x] = block_moment_max; + exp_avg_sq_scale[variance_scale_offset + blockIdx.x] = block_variance_max; + } } __syncthreads(); + // Recompute instead of keeping FP32 moment, variance and weight values live + // across the reduction. This mirrors the 8-bit kernel and avoids register + // spills into CUDA local memory for the packed 4-bit update. if (active) { + const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); + const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; + float variance = (static_cast(variance_code) + 1.0f) * old_variance_scale / 16.0f; + const float gradient = load_grad(grad, local_index); + float updated_weight = load_model(model, local_index); + if (weight_decay != 0.0f) { + updated_weight *= 1.0f - effective_lr * weight_decay; + } + moment = beta1 * moment + (1.0f - beta1) * gradient; + variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; + const float denom = sqrtf(variance) / bias_correction2_sqrt + eps; + updated_weight -= step_size * moment / denom; const float normalized_moment = moment / fmaxf(new_moment_scale, 1.0e-30f); moment_codes[tid] = nearest_signed_dynamic_code(normalized_moment); const float normalized_variance = variance / fmaxf(new_variance_scale, 1.0e-30f); @@ -239,7 +273,6 @@ __global__ void adamw_4bit_rank1_stats_kernel( int64_t packed_offset, int64_t parameter_shard_start, float beta2) { - __shared__ float variance_values[1024]; __shared__ int invalid_block; const int tid = threadIdx.x; const int64_t local_index = static_cast(blockIdx.x) * blockDim.x + tid; @@ -260,7 +293,6 @@ __global__ void adamw_4bit_rank1_stats_kernel( atomicExch(&invalid_block, 1); } } - variance_values[tid] = variance; __syncthreads(); if (invalid_block != 0) { if (tid == 0) { @@ -273,7 +305,7 @@ __global__ void adamw_4bit_rank1_stats_kernel( int64_t axis_offset = 0; for (int64_t axis = 0; axis < ndim; ++axis) { const int64_t coordinate = (parameter_index / strides[axis]) % shape[axis]; - atomic_max_nonnegative(updated_scales + axis_offset + coordinate, variance_values[tid]); + atomic_max_nonnegative(updated_scales + axis_offset + coordinate, variance); axis_offset += shape[axis]; } } @@ -303,9 +335,13 @@ __global__ void adamw_4bit_rank1_step_kernel( float eps, float step_size, float bias_correction2_sqrt) { - __shared__ float moment_max[1024]; - __shared__ uint8_t moment_codes[1024]; - __shared__ uint8_t variance_codes[1024]; + constexpr int warp_size = 32; + constexpr int max_warps = 32; + __shared__ float warp_moment_maxima[max_warps]; + extern __shared__ uint8_t packed_codes[]; + uint8_t* moment_codes = packed_codes; + uint8_t* variance_codes = packed_codes + blockDim.x; + __shared__ float old_moment_scale; __shared__ float new_moment_scale; __shared__ int invalid_block; const int tid = threadIdx.x; @@ -316,20 +352,19 @@ __global__ void adamw_4bit_rank1_step_kernel( const bool active = local_index < numel; if (tid == 0) { invalid_block = 0; + old_moment_scale = exp_avg_scale[moment_scale_offset + blockIdx.x]; } __syncthreads(); - float moment = 0.0f; - float variance = 0.0f; - float updated_weight = 0.0f; + float local_moment_max = 0.0f; if (active) { const int64_t parameter_index = parameter_shard_start + local_index; const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); - moment = kSigned4bitDynamicMap[moment_code] * exp_avg_scale[moment_scale_offset + blockIdx.x]; + float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); - variance = (static_cast(variance_code) + 1.0f) * old_scale / 16.0f; + float variance = (static_cast(variance_code) + 1.0f) * old_scale / 16.0f; const float gradient = load_grad(grad, local_index); - updated_weight = load_model(model, local_index); + float updated_weight = load_model(model, local_index); if (weight_decay != 0.0f) { updated_weight *= 1.0f - effective_lr * weight_decay; } @@ -340,25 +375,49 @@ __global__ void adamw_4bit_rank1_step_kernel( if (!isfinite(gradient) || !isfinite(moment) || !isfinite(variance) || !isfinite(updated_weight)) { atomicExch(&invalid_block, 1); } + local_moment_max = fabsf(moment); } - moment_max[tid] = active ? fabsf(moment) : 0.0f; - __syncthreads(); - for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { - if (tid < stride) { - moment_max[tid] = fmaxf(moment_max[tid], moment_max[tid + stride]); - } - __syncthreads(); + for (int offset = warp_size / 2; offset > 0; offset >>= 1) { + local_moment_max = fmaxf(local_moment_max, __shfl_down_sync(0xFFFFFFFFu, local_moment_max, offset)); + } + const int lane = tid & (warp_size - 1); + const int warp = tid / warp_size; + if (lane == 0) { + warp_moment_maxima[warp] = local_moment_max; } + __syncthreads(); if (invalid_block != 0) { return; } - if (tid == 0) { - new_moment_scale = moment_max[0]; - exp_avg_scale[moment_scale_offset + blockIdx.x] = new_moment_scale; + if (warp == 0) { + const int warp_count = blockDim.x / warp_size; + float block_moment_max = lane < warp_count ? warp_moment_maxima[lane] : 0.0f; + for (int offset = warp_size / 2; offset > 0; offset >>= 1) { + block_moment_max = + fmaxf(block_moment_max, __shfl_down_sync(0xFFFFFFFFu, block_moment_max, offset)); + } + if (lane == 0) { + new_moment_scale = block_moment_max; + exp_avg_scale[moment_scale_offset + blockIdx.x] = block_moment_max; + } } __syncthreads(); if (active) { const int64_t parameter_index = parameter_shard_start + local_index; + const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); + const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; + const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); + float variance = (static_cast(variance_code) + 1.0f) * old_scale / 16.0f; + const float gradient = load_grad(grad, local_index); + float updated_weight = load_model(model, local_index); + if (weight_decay != 0.0f) { + updated_weight *= 1.0f - effective_lr * weight_decay; + } + moment = beta1 * moment + (1.0f - beta1) * gradient; + variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; + const float denom = sqrtf(variance) / bias_correction2_sqrt + eps; + updated_weight -= step_size * moment / denom; moment_codes[tid] = nearest_signed_dynamic_code(moment / fmaxf(new_moment_scale, 1.0e-30f)); const float new_scale = rank1_scale(updated_scales, shape, strides, ndim, parameter_index); int variance_code = __float2int_rn(variance / fmaxf(new_scale, 1.0e-30f) * 16.0f - 1.0f); @@ -565,7 +624,8 @@ void launch_adamw_4bit( float bias_correction2_sqrt) { const int blocks = static_cast((model.numel() + quant_block_size - 1) / quant_block_size); const auto stream = at::cuda::getCurrentCUDAStream(); - adamw_4bit_kernel<<(quant_block_size), 0, stream>>>( + const size_t shared_bytes = static_cast(2 * quant_block_size) * sizeof(uint8_t); + adamw_4bit_kernel<<(quant_block_size), shared_bytes, stream>>>( model.data_ptr(), grad.data_ptr(), exp_avg_q.data_ptr(), @@ -642,7 +702,9 @@ void launch_adamw_4bit_rank1_step( float bias_correction2_sqrt) { const int blocks = static_cast((model.numel() + quant_block_size - 1) / quant_block_size); const auto stream = at::cuda::getCurrentCUDAStream(); - adamw_4bit_rank1_step_kernel<<(quant_block_size), 0, stream>>>( + const size_t shared_bytes = static_cast(2 * quant_block_size) * sizeof(uint8_t); + adamw_4bit_rank1_step_kernel + <<(quant_block_size), shared_bytes, stream>>>( model.data_ptr(), grad.data_ptr(), exp_avg_q.data_ptr(), From d72305fee5e9463c028b59776e891cc98cfca2b0 Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 18:01:47 +0800 Subject: [PATCH 4/7] fix(cuda): avoid AdamW4bit code redeclaration --- areno/accel/csrc/optimizer.cu | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu index 2762c763..5f7bc9c2 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -243,9 +243,10 @@ __global__ void adamw_4bit_kernel( const float normalized_moment = moment / fmaxf(new_moment_scale, 1.0e-30f); moment_codes[tid] = nearest_signed_dynamic_code(normalized_moment); const float normalized_variance = variance / fmaxf(new_variance_scale, 1.0e-30f); - int variance_code = __float2int_rn(normalized_variance * 16.0f - 1.0f); - variance_code = variance_code < 0 ? 0 : (variance_code > 15 ? 15 : variance_code); - variance_codes[tid] = static_cast(variance_code); + int updated_variance_code = __float2int_rn(normalized_variance * 16.0f - 1.0f); + updated_variance_code = + updated_variance_code < 0 ? 0 : (updated_variance_code > 15 ? 15 : updated_variance_code); + variance_codes[tid] = static_cast(updated_variance_code); store_model(model, local_index, updated_weight); } else { moment_codes[tid] = 7; @@ -420,9 +421,11 @@ __global__ void adamw_4bit_rank1_step_kernel( updated_weight -= step_size * moment / denom; moment_codes[tid] = nearest_signed_dynamic_code(moment / fmaxf(new_moment_scale, 1.0e-30f)); const float new_scale = rank1_scale(updated_scales, shape, strides, ndim, parameter_index); - int variance_code = __float2int_rn(variance / fmaxf(new_scale, 1.0e-30f) * 16.0f - 1.0f); - variance_code = variance_code < 0 ? 0 : (variance_code > 15 ? 15 : variance_code); - variance_codes[tid] = static_cast(variance_code); + int updated_variance_code = + __float2int_rn(variance / fmaxf(new_scale, 1.0e-30f) * 16.0f - 1.0f); + updated_variance_code = + updated_variance_code < 0 ? 0 : (updated_variance_code > 15 ? 15 : updated_variance_code); + variance_codes[tid] = static_cast(updated_variance_code); store_model(model, local_index, updated_weight); } else { moment_codes[tid] = 7; From e71ac9eaac1314078e6a3976c1b11b8cdd23a74f Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 18:34:09 +0800 Subject: [PATCH 5/7] fix(optimizer): stream fresh AdamW4bit rank1 state --- areno/accel/csrc/extension.cpp | 3 ++- areno/accel/csrc/optimizer.cu | 24 +++++++++++++--------- areno/accel/optimizer.py | 4 +++- areno/engine/optim/adamw_4bit.py | 34 +++++++++++++++++++++++++------- tests/test_adamw_4bit_cpu.py | 29 +++++++++++++++++++++++++++ 5 files changed, 76 insertions(+), 18 deletions(-) diff --git a/areno/accel/csrc/extension.cpp b/areno/accel/csrc/extension.cpp index e5ccdc1f..e28a3f10 100644 --- a/areno/accel/csrc/extension.cpp +++ b/areno/accel/csrc/extension.cpp @@ -214,7 +214,8 @@ void areno_adamw_4bit_rank1_stats_cuda( int64_t packed_offset, int64_t parameter_shard_start, int64_t quant_block_size, - double beta2); + double beta2, + bool has_state); void areno_adamw_4bit_rank1_step_cuda( torch::Tensor model, torch::Tensor grad, diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu index 5f7bc9c2..371d8594 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -273,7 +273,8 @@ __global__ void adamw_4bit_rank1_stats_kernel( int64_t numel, int64_t packed_offset, int64_t parameter_shard_start, - float beta2) { + float beta2, + bool has_state) { __shared__ int invalid_block; const int tid = threadIdx.x; const int64_t local_index = static_cast(blockIdx.x) * blockDim.x + tid; @@ -285,9 +286,11 @@ __global__ void adamw_4bit_rank1_stats_kernel( float variance = 0.0f; if (active) { const int64_t parameter_index = parameter_shard_start + local_index; - const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); - const uint8_t code = load_nibble(exp_avg_sq_q + packed_offset, local_index); - variance = (static_cast(code) + 1.0f) * old_scale / 16.0f; + if (has_state) { + const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); + const uint8_t code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + variance = (static_cast(code) + 1.0f) * old_scale / 16.0f; + } const float gradient = load_grad(grad, local_index); variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; if (!isfinite(gradient) || !isfinite(variance)) { @@ -661,7 +664,8 @@ void launch_adamw_4bit_rank1_stats( int64_t packed_offset, int64_t parameter_shard_start, int64_t quant_block_size, - float beta2) { + float beta2, + bool has_state) { const int blocks = static_cast((grad.numel() + quant_block_size - 1) / quant_block_size); const auto stream = at::cuda::getCurrentCUDAStream(); adamw_4bit_rank1_stats_kernel<<(quant_block_size), 0, stream>>>( @@ -676,7 +680,8 @@ void launch_adamw_4bit_rank1_stats( grad.numel(), packed_offset, parameter_shard_start, - beta2); + beta2, + has_state); C10_CUDA_KERNEL_LAUNCH_CHECK(); } @@ -1037,7 +1042,8 @@ void areno_adamw_4bit_rank1_stats_cuda( int64_t packed_offset, int64_t parameter_shard_start, int64_t quant_block_size, - double beta2) { + double beta2, + bool has_state) { c10::cuda::CUDAGuard guard(grad.device()); TORCH_CHECK( grad.is_cuda() && exp_avg_sq_q.is_cuda() && previous_scales.is_cuda() && updated_scales.is_cuda() && @@ -1050,11 +1056,11 @@ void areno_adamw_4bit_rank1_stats_cuda( if (grad.scalar_type() == at::kBFloat16) { launch_adamw_4bit_rank1_stats( grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides, packed_offset, - parameter_shard_start, quant_block_size, beta2); + parameter_shard_start, quant_block_size, beta2, has_state); } else if (grad.scalar_type() == at::kFloat) { launch_adamw_4bit_rank1_stats( grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides, packed_offset, - parameter_shard_start, quant_block_size, beta2); + parameter_shard_start, quant_block_size, beta2, has_state); } else { TORCH_CHECK(false, "AdamW4bit rank-1 gradient must be bfloat16 or float32"); } diff --git a/areno/accel/optimizer.py b/areno/accel/optimizer.py index aeed65ed..b0a81278 100644 --- a/areno/accel/optimizer.py +++ b/areno/accel/optimizer.py @@ -291,6 +291,7 @@ def areno_adamw_4bit_rank1_stats( parameter_shard_start: int, quant_block_size: int, beta2: float, + has_state: bool = True, ) -> None: """Accumulate updated rank-1 second-moment maxima in bounded CUDA blocks.""" @@ -311,7 +312,7 @@ def areno_adamw_4bit_rank1_stats( if previous_scales.numel() != updated_scales.numel(): raise ValueError("AdamW4bit previous and updated rank-1 scale layouts must match") packed_numel = (grad.numel() + 1) // 2 - if packed_offset < 0 or packed_offset + packed_numel > exp_avg_sq_q.numel(): + if has_state and (packed_offset < 0 or packed_offset + packed_numel > exp_avg_sq_q.numel()): raise ValueError("AdamW4bit rank-1 packed variance slice is out of bounds") if parameter_shard_start < 0: raise ValueError("AdamW4bit rank-1 parameter shard start must be non-negative") @@ -327,6 +328,7 @@ def areno_adamw_4bit_rank1_stats( parameter_shard_start, quant_block_size, beta2, + has_state, ) diff --git a/areno/engine/optim/adamw_4bit.py b/areno/engine/optim/adamw_4bit.py index e5b734dd..be407c25 100644 --- a/areno/engine/optim/adamw_4bit.py +++ b/areno/engine/optim/adamw_4bit.py @@ -88,6 +88,7 @@ def __init__( quant_block_size=quant_block_size, ) self._rank1_metadata_cache: dict[tuple[int, torch.device], tuple[torch.Tensor, torch.Tensor]] = {} + self._rank1_empty_codes: dict[torch.device, torch.Tensor] = {} self._rank1_scales: dict[int, torch.Tensor | None] = { id(parameter): None for parameter in self.model_params if parameter.ndim >= 2 } @@ -184,6 +185,7 @@ def clear_state(self) -> None: for parameter_id in self._rank1_scales: self._rank1_scales[parameter_id] = None self._rank1_metadata_cache.clear() + self._rank1_empty_codes.clear() @torch.no_grad() def offload_state(self, mode: str = "cpu", directory: str | None = None, batch_size: int = 1) -> None: @@ -194,6 +196,7 @@ def offload_state(self, mode: str = "cpu", directory: str | None = None, batch_s if scales is not None and scales.device.type != "cpu": self._rank1_scales[parameter_id] = scales.to(device="cpu") self._rank1_metadata_cache.clear() + self._rank1_empty_codes.clear() @torch.no_grad() def onload_state(self, device: torch.device) -> None: @@ -204,6 +207,7 @@ def onload_state(self, device: torch.device) -> None: if scales is not None and scales.device != device: self._rank1_scales[parameter_id] = scales.to(device=device) self._rank1_metadata_cache.clear() + self._rank1_empty_codes.clear() @torch.no_grad() def _ensure_bucket_state(self, bucket: _MasterBucket, state) -> None: @@ -277,7 +281,13 @@ def step(self, closure=None): _uses_rank1_normalization(ref) and self._ref_has_gradient(bucket, ref) for ref in bucket.refs ): continue - self._ensure_bucket_state(bucket, state) + # A fresh optimizer has an implicit all-zero second moment. + # Do not materialize every packed bucket during the statistics + # pass while every FP32 gradient shard is still resident. The + # update pass below initializes one bucket at a time and drops + # its gradient immediately, preserving streaming peak memory. + if state.step > 0 or state.exp_avg_sq_q is not None or state.offload_file is not None: + self._ensure_bucket_state(bucket, state) for ref, packed_offset, _moment_scale_offset, _variance_scale_offset in self._iter_ref_layout(bucket): if not _uses_rank1_normalization(ref) or not self._ref_has_gradient(bucket, ref): continue @@ -525,7 +535,6 @@ def _rank1_variance_statistics( ) -> None: """Compute updated per-axis maxima without a full FP32 moment.""" - assert state.exp_avg_sq_q is not None if grad is None or ref.shard_numel == 0: return grad_shard = grad if bucket.grad_shard is not None else grad.narrow(0, ref.shard_start, ref.shard_numel) @@ -534,18 +543,26 @@ def _rank1_variance_statistics( from areno.accel.optimizer import areno_adamw_4bit_rank1_stats shape, strides = self._rank1_metadata(ref) + variance_codes = state.exp_avg_sq_q + has_state = variance_codes is not None + if variance_codes is None: + variance_codes = self._rank1_empty_codes.get(grad_shard.device) + if variance_codes is None: + variance_codes = torch.empty(1, device=grad_shard.device, dtype=torch.uint8) + self._rank1_empty_codes[grad_shard.device] = variance_codes areno_adamw_4bit_rank1_stats( grad_shard.contiguous(), - state.exp_avg_sq_q, + variance_codes, previous_scales, updated_scales, invalid, shape, strides, - packed_offset=packed_offset, + packed_offset=packed_offset if has_state else 0, parameter_shard_start=parameter_shard_start, quant_block_size=self.quant_block_size, beta2=beta2, + has_state=has_state, ) return @@ -556,10 +573,13 @@ def _rank1_variance_statistics( byte_count = (count + 1) // 2 flat_start = parameter_shard_start + start element_scales = _rank1_element_scales(previous_scales, shape_tuple, flat_start, count) - variance = _unpack_positive_4bit_elementwise( - state.exp_avg_sq_q.narrow(0, byte_start, byte_count), count, element_scales - ) gradient = grad_shard.narrow(0, start, count).to(dtype=torch.float32) + if state.exp_avg_sq_q is None: + variance = torch.zeros(count, device=gradient.device, dtype=torch.float32) + else: + variance = _unpack_positive_4bit_elementwise( + state.exp_avg_sq_q.narrow(0, byte_start, byte_count), count, element_scales + ) variance.mul_(beta2).addcmul_(gradient, gradient, value=1.0 - beta2) if not bool(torch.isfinite(gradient).all() & torch.isfinite(variance).all()): invalid.fill_(1) diff --git a/tests/test_adamw_4bit_cpu.py b/tests/test_adamw_4bit_cpu.py index aba9c280..4cf1aced 100644 --- a/tests/test_adamw_4bit_cpu.py +++ b/tests/test_adamw_4bit_cpu.py @@ -112,6 +112,35 @@ def test_adamw4bit_packs_two_moments_within_storage_budget() -> None: assert optimizer.persistent_moment_bytes() <= eight_bit_bytes * 0.6 +def test_adamw4bit_fresh_rank1_state_is_initialized_during_streaming_update() -> None: + parameters = [torch.nn.Parameter(torch.zeros(32, 32)) for _ in range(2)] + optimizer = AdamW4bit( + parameters, + lr=3.0e-4, + betas=(0.9, 0.99), + weight_decay=0.01, + bucket_numel=1024, + quant_block_size=128, + ) + for parameter in parameters: + parameter.grad = torch.ones_like(parameter) + + live_gradients_at_initialization: list[int] = [] + ensure_bucket_state = optimizer._ensure_bucket_state + + def tracked_ensure_bucket_state(bucket, state) -> None: + live_gradients_at_initialization.append(sum(parameter.grad is not None for parameter in parameters)) + ensure_bucket_state(bucket, state) + + optimizer._ensure_bucket_state = tracked_ensure_bucket_state + optimizer.step() + + # The statistics pass consumes the implicit zero second moment without + # allocating packed state. State is initialized only in the update pass, + # where each completed parameter releases its gradient before the next. + assert live_gradients_at_initialization == [2, 1] + + def test_adamw4bit_second_moment_mapping_excludes_zero() -> None: values = torch.tensor([0.0, 1.0 / 16.0, 0.5, 1.0]) From d8383741d2c6a074f2f33a268907e9b2e4829283 Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 19:35:54 +0800 Subject: [PATCH 6/7] perf(optimizer): bound AdamW4bit rank1 stats launch --- areno/accel/csrc/optimizer.cu | 40 +++++++++++++---------------------- 1 file changed, 15 insertions(+), 25 deletions(-) diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu index 371d8594..5868ce81 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -275,17 +275,14 @@ __global__ void adamw_4bit_rank1_stats_kernel( int64_t parameter_shard_start, float beta2, bool has_state) { - __shared__ int invalid_block; - const int tid = threadIdx.x; - const int64_t local_index = static_cast(blockIdx.x) * blockDim.x + tid; - const bool active = local_index < numel; - if (tid == 0) { - invalid_block = 0; - } - __syncthreads(); - float variance = 0.0f; - if (active) { + // This pass has no block-local reduction. A bounded grid-stride launch + // avoids provisioning per-thread CUDA local storage for one thread per + // parameter element on very large expert tensors. + for (int64_t local_index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + local_index < numel; + local_index += static_cast(blockDim.x) * gridDim.x) { const int64_t parameter_index = parameter_shard_start + local_index; + float variance = 0.0f; if (has_state) { const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); const uint8_t code = load_nibble(exp_avg_sq_q + packed_offset, local_index); @@ -294,18 +291,9 @@ __global__ void adamw_4bit_rank1_stats_kernel( const float gradient = load_grad(grad, local_index); variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; if (!isfinite(gradient) || !isfinite(variance)) { - atomicExch(&invalid_block, 1); - } - } - __syncthreads(); - if (invalid_block != 0) { - if (tid == 0) { atomicExch(invalid, 1); + continue; } - return; - } - if (active) { - const int64_t parameter_index = parameter_shard_start + local_index; int64_t axis_offset = 0; for (int64_t axis = 0; axis < ndim; ++axis) { const int64_t coordinate = (parameter_index / strides[axis]) % shape[axis]; @@ -663,12 +651,14 @@ void launch_adamw_4bit_rank1_stats( torch::Tensor strides, int64_t packed_offset, int64_t parameter_shard_start, - int64_t quant_block_size, float beta2, bool has_state) { - const int blocks = static_cast((grad.numel() + quant_block_size - 1) / quant_block_size); + constexpr int threads = 256; + constexpr int max_blocks = 4096; + int blocks = static_cast((grad.numel() + threads - 1) / threads); + blocks = blocks < max_blocks ? blocks : max_blocks; const auto stream = at::cuda::getCurrentCUDAStream(); - adamw_4bit_rank1_stats_kernel<<(quant_block_size), 0, stream>>>( + adamw_4bit_rank1_stats_kernel<<>>( grad.data_ptr(), exp_avg_sq_q.data_ptr(), previous_scales.data_ptr(), @@ -1056,11 +1046,11 @@ void areno_adamw_4bit_rank1_stats_cuda( if (grad.scalar_type() == at::kBFloat16) { launch_adamw_4bit_rank1_stats( grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides, packed_offset, - parameter_shard_start, quant_block_size, beta2, has_state); + parameter_shard_start, beta2, has_state); } else if (grad.scalar_type() == at::kFloat) { launch_adamw_4bit_rank1_stats( grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides, packed_offset, - parameter_shard_start, quant_block_size, beta2, has_state); + parameter_shard_start, beta2, has_state); } else { TORCH_CHECK(false, "AdamW4bit rank-1 gradient must be bfloat16 or float32"); } From 48f5b67e7179c5b10f44c7fc37da2b99aae4ad46 Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 20:25:28 +0800 Subject: [PATCH 7/7] perf(optimizer): stream factored AdamW4bit state --- areno/accel/csrc/extension.cpp | 43 +- areno/accel/csrc/optimizer.cu | 293 +++++------ areno/accel/optimizer.py | 161 +++--- areno/engine/optim/adamw_4bit.py | 667 ++++++++++-------------- areno/engine/optim/adamw_fp32_master.py | 15 +- areno/engine/training.py | 33 +- docs/cli/training.rst | 7 +- docs/reference/adamw-4bit.rst | 38 +- tests/test_adamw_4bit_cpu.py | 422 +++++++-------- 9 files changed, 744 insertions(+), 935 deletions(-) diff --git a/areno/accel/csrc/extension.cpp b/areno/accel/csrc/extension.cpp index e28a3f10..d85cf894 100644 --- a/areno/accel/csrc/extension.cpp +++ b/areno/accel/csrc/extension.cpp @@ -192,8 +192,9 @@ void areno_adamw_4bit_step_cuda( torch::Tensor exp_avg_scale, torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, - int64_t packed_offset, + int64_t moment_packed_offset, int64_t moment_scale_offset, + int64_t variance_packed_offset, int64_t variance_scale_offset, int64_t quant_block_size, double beta1, @@ -203,36 +204,28 @@ void areno_adamw_4bit_step_cuda( double eps, double step_size, double bias_correction2_sqrt); -void areno_adamw_4bit_rank1_stats_cuda( +void areno_adamw_4bit_factored_stats_cuda( torch::Tensor grad, - torch::Tensor exp_avg_sq_q, - torch::Tensor previous_scales, - torch::Tensor updated_scales, + torch::Tensor factor_sums, torch::Tensor invalid, - torch::Tensor shape, - torch::Tensor strides, - int64_t packed_offset, int64_t parameter_shard_start, - int64_t quant_block_size, - double beta2, - bool has_state); -void areno_adamw_4bit_rank1_step_cuda( + int64_t rows, + int64_t columns); +void areno_adamw_4bit_factored_step_cuda( torch::Tensor model, torch::Tensor grad, torch::Tensor exp_avg_q, torch::Tensor exp_avg_scale, - torch::Tensor exp_avg_sq_q, - torch::Tensor previous_scales, - torch::Tensor updated_scales, + torch::Tensor factors, + torch::Tensor row_mean, torch::Tensor invalid, - torch::Tensor shape, - torch::Tensor strides, - int64_t packed_offset, + int64_t moment_packed_offset, int64_t moment_scale_offset, int64_t parameter_shard_start, int64_t quant_block_size, + int64_t rows, + int64_t columns, double beta1, - double beta2, double effective_lr, double weight_decay, double eps, @@ -272,13 +265,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("areno_adamw_fp32_master_step", &areno_adamw_fp32_master_step_cuda, "ARENO compact FP32-master AdamW step"); m.def("areno_adamw_4bit_step", &areno_adamw_4bit_step_cuda, "ARENO packed block-wise AdamW4bit step"); m.def( - "areno_adamw_4bit_rank1_stats", - &areno_adamw_4bit_rank1_stats_cuda, - "ARENO rank-1 AdamW4bit statistics pass"); + "areno_adamw_4bit_factored_stats", + &areno_adamw_4bit_factored_stats_cuda, + "ARENO factored AdamW4bit statistics pass"); m.def( - "areno_adamw_4bit_rank1_step", - &areno_adamw_4bit_rank1_step_cuda, - "ARENO rank-1 AdamW4bit update pass"); + "areno_adamw_4bit_factored_step", + &areno_adamw_4bit_factored_step_cuda, + "ARENO factored AdamW4bit update pass"); m.def("areno_adamw_8bit_step", &areno_adamw_8bit_step_cuda, "ARENO block-wise 8-bit AdamW step"); m.def("areno_adamw_fp32_state_step", &areno_adamw_fp32_state_step_cuda, "ARENO FP32-state AdamW step"); m.def("areno_silu_and_mul", &areno_silu_and_mul_cuda, "ARENO SiLU and multiply"); diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu index 5868ce81..0a6498a6 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -31,28 +31,6 @@ __device__ __forceinline__ uint8_t nearest_signed_dynamic_code(float normalized) return best; } -__device__ __forceinline__ float atomic_max_nonnegative(float* address, float value) { - int* address_as_int = reinterpret_cast(address); - const int old = atomicMax(address_as_int, __float_as_int(value)); - return __int_as_float(old); -} - -__device__ __forceinline__ float rank1_scale( - const float* axis_scales, - const int64_t* shape, - const int64_t* strides, - int64_t ndim, - int64_t flat_index) { - float scale = __int_as_float(0x7f800000); - int64_t axis_offset = 0; - for (int64_t axis = 0; axis < ndim; ++axis) { - const int64_t coordinate = (flat_index / strides[axis]) % shape[axis]; - scale = fminf(scale, axis_scales[axis_offset + coordinate]); - axis_offset += shape[axis]; - } - return scale; -} - __device__ __forceinline__ uint8_t nearest_dynamic_code(float value, const float* codebook) { int lower = 0; int upper = 255; @@ -133,8 +111,9 @@ __global__ void adamw_4bit_kernel( uint8_t* exp_avg_sq_q, float* exp_avg_sq_scale, int64_t numel, - int64_t packed_offset, + int64_t moment_packed_offset, int64_t moment_scale_offset, + int64_t variance_packed_offset, int64_t variance_scale_offset, float beta1, float beta2, @@ -170,8 +149,8 @@ __global__ void adamw_4bit_kernel( float local_moment_max = 0.0f; float local_variance_max = 0.0f; if (active) { - const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); - const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + const uint8_t moment_code = load_nibble(exp_avg_q + moment_packed_offset, local_index); + const uint8_t variance_code = load_nibble(exp_avg_sq_q + variance_packed_offset, local_index); float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; float variance = (static_cast(variance_code) + 1.0f) * old_variance_scale / 16.0f; const float gradient = load_grad(grad, local_index); @@ -227,8 +206,8 @@ __global__ void adamw_4bit_kernel( // across the reduction. This mirrors the 8-bit kernel and avoids register // spills into CUDA local memory for the packed 4-bit update. if (active) { - const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); - const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + const uint8_t moment_code = load_nibble(exp_avg_q + moment_packed_offset, local_index); + const uint8_t variance_code = load_nibble(exp_avg_sq_q + variance_packed_offset, local_index); float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; float variance = (static_cast(variance_code) + 1.0f) * old_variance_scale / 16.0f; const float gradient = load_grad(grad, local_index); @@ -254,74 +233,54 @@ __global__ void adamw_4bit_kernel( } __syncthreads(); if ((tid & 1) == 0 && local_index < numel) { - const int64_t byte_index = packed_offset + (local_index >> 1); - exp_avg_q[byte_index] = moment_codes[tid] | static_cast(moment_codes[tid + 1] << 4); - exp_avg_sq_q[byte_index] = variance_codes[tid] | static_cast(variance_codes[tid + 1] << 4); + const int64_t moment_byte_index = moment_packed_offset + (local_index >> 1); + const int64_t variance_byte_index = variance_packed_offset + (local_index >> 1); + exp_avg_q[moment_byte_index] = moment_codes[tid] | static_cast(moment_codes[tid + 1] << 4); + exp_avg_sq_q[variance_byte_index] = + variance_codes[tid] | static_cast(variance_codes[tid + 1] << 4); } } template -__global__ void adamw_4bit_rank1_stats_kernel( +__global__ void adamw_4bit_factored_stats_kernel( const grad_t* grad, - const uint8_t* exp_avg_sq_q, - const float* previous_scales, - float* updated_scales, + float* factor_sums, int* invalid, - const int64_t* shape, - const int64_t* strides, - int64_t ndim, int64_t numel, - int64_t packed_offset, int64_t parameter_shard_start, - float beta2, - bool has_state) { - // This pass has no block-local reduction. A bounded grid-stride launch - // avoids provisioning per-thread CUDA local storage for one thread per - // parameter element on very large expert tensors. + int64_t rows, + int64_t columns) { for (int64_t local_index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; local_index < numel; local_index += static_cast(blockDim.x) * gridDim.x) { const int64_t parameter_index = parameter_shard_start + local_index; - float variance = 0.0f; - if (has_state) { - const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); - const uint8_t code = load_nibble(exp_avg_sq_q + packed_offset, local_index); - variance = (static_cast(code) + 1.0f) * old_scale / 16.0f; - } const float gradient = load_grad(grad, local_index); - variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; - if (!isfinite(gradient) || !isfinite(variance)) { + const float squared = gradient * gradient; + if (!isfinite(squared)) { atomicExch(invalid, 1); continue; } - int64_t axis_offset = 0; - for (int64_t axis = 0; axis < ndim; ++axis) { - const int64_t coordinate = (parameter_index / strides[axis]) % shape[axis]; - atomic_max_nonnegative(updated_scales + axis_offset + coordinate, variance); - axis_offset += shape[axis]; - } + atomicAdd(factor_sums + parameter_index / columns, squared); + atomicAdd(factor_sums + rows + parameter_index % columns, squared); } } template -__global__ void adamw_4bit_rank1_step_kernel( +__global__ void adamw_4bit_factored_step_kernel( model_t* model, const grad_t* grad, uint8_t* exp_avg_q, float* exp_avg_scale, - uint8_t* exp_avg_sq_q, - const float* previous_scales, - const float* updated_scales, + const float* factors, + const float* row_mean, const int* invalid, - const int64_t* shape, - const int64_t* strides, - int64_t ndim, int64_t numel, - int64_t packed_offset, + int64_t moment_packed_offset, int64_t moment_scale_offset, int64_t parameter_shard_start, + int64_t rows, + int64_t columns, float beta1, - float beta2, float effective_lr, float weight_decay, float eps, @@ -332,7 +291,6 @@ __global__ void adamw_4bit_rank1_step_kernel( __shared__ float warp_moment_maxima[max_warps]; extern __shared__ uint8_t packed_codes[]; uint8_t* moment_codes = packed_codes; - uint8_t* variance_codes = packed_codes + blockDim.x; __shared__ float old_moment_scale; __shared__ float new_moment_scale; __shared__ int invalid_block; @@ -350,18 +308,17 @@ __global__ void adamw_4bit_rank1_step_kernel( float local_moment_max = 0.0f; if (active) { const int64_t parameter_index = parameter_shard_start + local_index; - const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); - const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + const uint8_t moment_code = load_nibble(exp_avg_q + moment_packed_offset, local_index); float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; - const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); - float variance = (static_cast(variance_code) + 1.0f) * old_scale / 16.0f; + const int64_t row = parameter_index / columns; + const int64_t column = parameter_index % columns; + const float variance = factors[row] * factors[rows + column] / fmaxf(*row_mean, 1.0e-30f); const float gradient = load_grad(grad, local_index); float updated_weight = load_model(model, local_index); if (weight_decay != 0.0f) { updated_weight *= 1.0f - effective_lr * weight_decay; } moment = beta1 * moment + (1.0f - beta1) * gradient; - variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; const float denom = sqrtf(variance) / bias_correction2_sqrt + eps; updated_weight -= step_size * moment / denom; if (!isfinite(gradient) || !isfinite(moment) || !isfinite(variance) || !isfinite(updated_weight)) { @@ -396,37 +353,28 @@ __global__ void adamw_4bit_rank1_step_kernel( __syncthreads(); if (active) { const int64_t parameter_index = parameter_shard_start + local_index; - const uint8_t moment_code = load_nibble(exp_avg_q + packed_offset, local_index); - const uint8_t variance_code = load_nibble(exp_avg_sq_q + packed_offset, local_index); + const uint8_t moment_code = load_nibble(exp_avg_q + moment_packed_offset, local_index); float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; - const float old_scale = rank1_scale(previous_scales, shape, strides, ndim, parameter_index); - float variance = (static_cast(variance_code) + 1.0f) * old_scale / 16.0f; + const int64_t row = parameter_index / columns; + const int64_t column = parameter_index % columns; + const float variance = factors[row] * factors[rows + column] / fmaxf(*row_mean, 1.0e-30f); const float gradient = load_grad(grad, local_index); float updated_weight = load_model(model, local_index); if (weight_decay != 0.0f) { updated_weight *= 1.0f - effective_lr * weight_decay; } moment = beta1 * moment + (1.0f - beta1) * gradient; - variance = beta2 * variance + (1.0f - beta2) * gradient * gradient; const float denom = sqrtf(variance) / bias_correction2_sqrt + eps; updated_weight -= step_size * moment / denom; moment_codes[tid] = nearest_signed_dynamic_code(moment / fmaxf(new_moment_scale, 1.0e-30f)); - const float new_scale = rank1_scale(updated_scales, shape, strides, ndim, parameter_index); - int updated_variance_code = - __float2int_rn(variance / fmaxf(new_scale, 1.0e-30f) * 16.0f - 1.0f); - updated_variance_code = - updated_variance_code < 0 ? 0 : (updated_variance_code > 15 ? 15 : updated_variance_code); - variance_codes[tid] = static_cast(updated_variance_code); store_model(model, local_index, updated_weight); } else { moment_codes[tid] = 7; - variance_codes[tid] = 0; } __syncthreads(); if ((tid & 1) == 0 && local_index < numel) { - const int64_t byte_index = packed_offset + (local_index >> 1); + const int64_t byte_index = moment_packed_offset + (local_index >> 1); exp_avg_q[byte_index] = moment_codes[tid] | static_cast(moment_codes[tid + 1] << 4); - exp_avg_sq_q[byte_index] = variance_codes[tid] | static_cast(variance_codes[tid + 1] << 4); } } @@ -605,8 +553,9 @@ void launch_adamw_4bit( torch::Tensor exp_avg_scale, torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, - int64_t packed_offset, + int64_t moment_packed_offset, int64_t moment_scale_offset, + int64_t variance_packed_offset, int64_t variance_scale_offset, int64_t quant_block_size, float beta1, @@ -627,8 +576,9 @@ void launch_adamw_4bit( exp_avg_sq_q.data_ptr(), exp_avg_sq_scale.data_ptr(), model.numel(), - packed_offset, + moment_packed_offset, moment_scale_offset, + variance_packed_offset, variance_scale_offset, beta1, beta2, @@ -641,58 +591,45 @@ void launch_adamw_4bit( } template -void launch_adamw_4bit_rank1_stats( +void launch_adamw_4bit_factored_stats( torch::Tensor grad, - torch::Tensor exp_avg_sq_q, - torch::Tensor previous_scales, - torch::Tensor updated_scales, + torch::Tensor factor_sums, torch::Tensor invalid, - torch::Tensor shape, - torch::Tensor strides, - int64_t packed_offset, int64_t parameter_shard_start, - float beta2, - bool has_state) { + int64_t rows, + int64_t columns) { constexpr int threads = 256; constexpr int max_blocks = 4096; int blocks = static_cast((grad.numel() + threads - 1) / threads); blocks = blocks < max_blocks ? blocks : max_blocks; const auto stream = at::cuda::getCurrentCUDAStream(); - adamw_4bit_rank1_stats_kernel<<>>( + adamw_4bit_factored_stats_kernel<<>>( grad.data_ptr(), - exp_avg_sq_q.data_ptr(), - previous_scales.data_ptr(), - updated_scales.data_ptr(), + factor_sums.data_ptr(), invalid.data_ptr(), - shape.data_ptr(), - strides.data_ptr(), - shape.numel(), grad.numel(), - packed_offset, parameter_shard_start, - beta2, - has_state); + rows, + columns); C10_CUDA_KERNEL_LAUNCH_CHECK(); } template -void launch_adamw_4bit_rank1_step( +void launch_adamw_4bit_factored_step( torch::Tensor model, torch::Tensor grad, torch::Tensor exp_avg_q, torch::Tensor exp_avg_scale, - torch::Tensor exp_avg_sq_q, - torch::Tensor previous_scales, - torch::Tensor updated_scales, + torch::Tensor factors, + torch::Tensor row_mean, torch::Tensor invalid, - torch::Tensor shape, - torch::Tensor strides, - int64_t packed_offset, + int64_t moment_packed_offset, int64_t moment_scale_offset, int64_t parameter_shard_start, int64_t quant_block_size, + int64_t rows, + int64_t columns, float beta1, - float beta2, float effective_lr, float weight_decay, float eps, @@ -700,26 +637,23 @@ void launch_adamw_4bit_rank1_step( float bias_correction2_sqrt) { const int blocks = static_cast((model.numel() + quant_block_size - 1) / quant_block_size); const auto stream = at::cuda::getCurrentCUDAStream(); - const size_t shared_bytes = static_cast(2 * quant_block_size) * sizeof(uint8_t); - adamw_4bit_rank1_step_kernel + const size_t shared_bytes = static_cast(quant_block_size) * sizeof(uint8_t); + adamw_4bit_factored_step_kernel <<(quant_block_size), shared_bytes, stream>>>( model.data_ptr(), grad.data_ptr(), exp_avg_q.data_ptr(), exp_avg_scale.data_ptr(), - exp_avg_sq_q.data_ptr(), - previous_scales.data_ptr(), - updated_scales.data_ptr(), + factors.data_ptr(), + row_mean.data_ptr(), invalid.data_ptr(), - shape.data_ptr(), - strides.data_ptr(), - shape.numel(), model.numel(), - packed_offset, + moment_packed_offset, moment_scale_offset, parameter_shard_start, + rows, + columns, beta1, - beta2, effective_lr, weight_decay, eps, @@ -978,8 +912,9 @@ void areno_adamw_4bit_step_cuda( torch::Tensor exp_avg_scale, torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, - int64_t packed_offset, + int64_t moment_packed_offset, int64_t moment_scale_offset, + int64_t variance_packed_offset, int64_t variance_scale_offset, int64_t quant_block_size, double beta1, @@ -1003,8 +938,8 @@ void areno_adamw_4bit_step_cuda( #define LAUNCH_ADAMW4(MODEL_T, GRAD_T) \ launch_adamw_4bit( \ - model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, packed_offset, \ - moment_scale_offset, variance_scale_offset, \ + model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, moment_packed_offset, \ + moment_scale_offset, variance_packed_offset, variance_scale_offset, \ quant_block_size, beta1, beta2, effective_lr, weight_decay, eps, step_size, bias_correction2_sqrt) if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kBFloat16) { @@ -1021,81 +956,99 @@ void areno_adamw_4bit_step_cuda( #undef LAUNCH_ADAMW4 } -void areno_adamw_4bit_rank1_stats_cuda( +void areno_adamw_4bit_factored_stats_cuda( torch::Tensor grad, - torch::Tensor exp_avg_sq_q, - torch::Tensor previous_scales, - torch::Tensor updated_scales, + torch::Tensor factor_sums, torch::Tensor invalid, - torch::Tensor shape, - torch::Tensor strides, - int64_t packed_offset, int64_t parameter_shard_start, - int64_t quant_block_size, - double beta2, - bool has_state) { + int64_t rows, + int64_t columns) { c10::cuda::CUDAGuard guard(grad.device()); TORCH_CHECK( - grad.is_cuda() && exp_avg_sq_q.is_cuda() && previous_scales.is_cuda() && updated_scales.is_cuda() && - invalid.is_cuda() && shape.is_cuda() && strides.is_cuda(), - "AdamW4bit rank-1 statistics tensors must be CUDA tensors"); + grad.is_cuda() && factor_sums.is_cuda() && invalid.is_cuda(), + "AdamW4bit factored statistics tensors must be CUDA tensors"); TORCH_CHECK( - quant_block_size >= 32 && quant_block_size <= 1024 && - (quant_block_size & (quant_block_size - 1)) == 0, - "AdamW4bit block size must be a power of two between 32 and 1024"); + grad.is_contiguous() && factor_sums.is_contiguous() && invalid.is_contiguous(), + "AdamW4bit factored statistics tensors must be contiguous"); + TORCH_CHECK(factor_sums.scalar_type() == at::kFloat, "AdamW4bit factored sums must be float32"); + TORCH_CHECK(invalid.scalar_type() == at::kInt && invalid.numel() == 1, "AdamW4bit invalid flag must be int32"); + TORCH_CHECK(rows > 0 && columns > 0, "AdamW4bit factored dimensions must be positive"); + TORCH_CHECK(factor_sums.numel() == rows + columns, "AdamW4bit factored state size must match matrix shape"); + TORCH_CHECK( + parameter_shard_start >= 0 && parameter_shard_start + grad.numel() <= rows * columns, + "AdamW4bit factored gradient slice is out of bounds"); if (grad.scalar_type() == at::kBFloat16) { - launch_adamw_4bit_rank1_stats( - grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides, packed_offset, - parameter_shard_start, beta2, has_state); + launch_adamw_4bit_factored_stats( + grad, factor_sums, invalid, parameter_shard_start, rows, columns); } else if (grad.scalar_type() == at::kFloat) { - launch_adamw_4bit_rank1_stats( - grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides, packed_offset, - parameter_shard_start, beta2, has_state); + launch_adamw_4bit_factored_stats(grad, factor_sums, invalid, parameter_shard_start, rows, columns); } else { - TORCH_CHECK(false, "AdamW4bit rank-1 gradient must be bfloat16 or float32"); + TORCH_CHECK(false, "AdamW4bit factored gradient must be bfloat16 or float32"); } } -void areno_adamw_4bit_rank1_step_cuda( +void areno_adamw_4bit_factored_step_cuda( torch::Tensor model, torch::Tensor grad, torch::Tensor exp_avg_q, torch::Tensor exp_avg_scale, - torch::Tensor exp_avg_sq_q, - torch::Tensor previous_scales, - torch::Tensor updated_scales, + torch::Tensor factors, + torch::Tensor row_mean, torch::Tensor invalid, - torch::Tensor shape, - torch::Tensor strides, - int64_t packed_offset, + int64_t moment_packed_offset, int64_t moment_scale_offset, int64_t parameter_shard_start, int64_t quant_block_size, + int64_t rows, + int64_t columns, double beta1, - double beta2, double effective_lr, double weight_decay, double eps, double step_size, double bias_correction2_sqrt) { c10::cuda::CUDAGuard guard(model.device()); -#define LAUNCH_ADAMW4_RANK1(MODEL_T, GRAD_T) \ - launch_adamw_4bit_rank1_step( \ - model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, previous_scales, updated_scales, invalid, \ - shape, strides, packed_offset, moment_scale_offset, parameter_shard_start, quant_block_size, beta1, \ - beta2, effective_lr, weight_decay, eps, step_size, bias_correction2_sqrt) + TORCH_CHECK( + model.is_cuda() && grad.is_cuda() && exp_avg_q.is_cuda() && exp_avg_scale.is_cuda() && factors.is_cuda() && + row_mean.is_cuda() && invalid.is_cuda(), + "AdamW4bit factored update tensors must be CUDA tensors"); + TORCH_CHECK( + model.is_contiguous() && grad.is_contiguous() && exp_avg_q.is_contiguous() && exp_avg_scale.is_contiguous() && + factors.is_contiguous() && row_mean.is_contiguous() && invalid.is_contiguous(), + "AdamW4bit factored update tensors must be contiguous"); + TORCH_CHECK(model.numel() == grad.numel(), "AdamW4bit factored model and gradient sizes must match"); + TORCH_CHECK(exp_avg_q.scalar_type() == at::kByte, "AdamW4bit packed momentum must be uint8"); + TORCH_CHECK( + exp_avg_scale.scalar_type() == at::kFloat && factors.scalar_type() == at::kFloat && + row_mean.scalar_type() == at::kFloat, + "AdamW4bit factored scales must be float32"); + TORCH_CHECK(invalid.scalar_type() == at::kInt && invalid.numel() == 1, "AdamW4bit invalid flag must be int32"); + TORCH_CHECK( + quant_block_size >= 32 && quant_block_size <= 1024 && + (quant_block_size & (quant_block_size - 1)) == 0, + "AdamW4bit block size must be a power of two between 32 and 1024"); + TORCH_CHECK(rows > 0 && columns > 0, "AdamW4bit factored dimensions must be positive"); + TORCH_CHECK(factors.numel() == rows + columns && row_mean.numel() == 1, "AdamW4bit factored state shape mismatch"); + TORCH_CHECK( + parameter_shard_start >= 0 && parameter_shard_start + model.numel() <= rows * columns, + "AdamW4bit factored parameter slice is out of bounds"); +#define LAUNCH_ADAMW4_FACTORED(MODEL_T, GRAD_T) \ + launch_adamw_4bit_factored_step( \ + model, grad, exp_avg_q, exp_avg_scale, factors, row_mean, invalid, moment_packed_offset, \ + moment_scale_offset, parameter_shard_start, quant_block_size, rows, columns, beta1, effective_lr, \ + weight_decay, eps, step_size, bias_correction2_sqrt) if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kBFloat16) { - LAUNCH_ADAMW4_RANK1(at::BFloat16, at::BFloat16); + LAUNCH_ADAMW4_FACTORED(at::BFloat16, at::BFloat16); } else if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kFloat) { - LAUNCH_ADAMW4_RANK1(at::BFloat16, float); + LAUNCH_ADAMW4_FACTORED(at::BFloat16, float); } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kBFloat16) { - LAUNCH_ADAMW4_RANK1(float, at::BFloat16); + LAUNCH_ADAMW4_FACTORED(float, at::BFloat16); } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kFloat) { - LAUNCH_ADAMW4_RANK1(float, float); + LAUNCH_ADAMW4_FACTORED(float, float); } else { - TORCH_CHECK(false, "AdamW4bit rank-1 model and gradient must be bfloat16 or float32"); + TORCH_CHECK(false, "AdamW4bit factored model and gradient must be bfloat16 or float32"); } -#undef LAUNCH_ADAMW4_RANK1 +#undef LAUNCH_ADAMW4_FACTORED } void areno_adamw_8bit_step_cuda( diff --git a/areno/accel/optimizer.py b/areno/accel/optimizer.py index b0a81278..c147a8f9 100644 --- a/areno/accel/optimizer.py +++ b/areno/accel/optimizer.py @@ -205,8 +205,9 @@ def areno_adamw_4bit_step( exp_avg_sq_q: torch.Tensor, exp_avg_sq_scale: torch.Tensor, *, - packed_offset: int, + moment_packed_offset: int, moment_scale_offset: int, + variance_packed_offset: int, variance_scale_offset: int, quant_block_size: int, beta1: float, @@ -238,10 +239,10 @@ def areno_adamw_4bit_step( raise ValueError("quant_block_size must be a power of two between 32 and 1024") packed_numel = (model.numel() + 1) // 2 scale_numel = (model.numel() + quant_block_size - 1) // quant_block_size - if packed_offset < 0 or packed_offset + packed_numel > exp_avg_q.numel(): - raise ValueError("packed AdamW4bit state slice is out of bounds") - if exp_avg_q.numel() != exp_avg_sq_q.numel(): - raise ValueError("packed AdamW4bit moments must have the same length") + if moment_packed_offset < 0 or moment_packed_offset + packed_numel > exp_avg_q.numel(): + raise ValueError("packed AdamW4bit first-moment slice is out of bounds") + if variance_packed_offset < 0 or variance_packed_offset + packed_numel > exp_avg_sq_q.numel(): + raise ValueError("packed AdamW4bit second-moment slice is out of bounds") if moment_scale_offset < 0 or moment_scale_offset + scale_numel > exp_avg_scale.numel(): raise ValueError("AdamW4bit first-moment scale slice is out of bounds") if variance_scale_offset < 0 or variance_scale_offset + scale_numel > exp_avg_sq_scale.numel(): @@ -253,8 +254,9 @@ def areno_adamw_4bit_step( exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, - packed_offset, + moment_packed_offset, moment_scale_offset, + variance_packed_offset, variance_scale_offset, quant_block_size, beta1, @@ -267,157 +269,122 @@ def areno_adamw_4bit_step( ) -def _validate_rank1_metadata(shape: torch.Tensor, strides: torch.Tensor, scales: torch.Tensor) -> None: - if shape.dtype != torch.int64 or strides.dtype != torch.int64: - raise TypeError("AdamW4bit rank-1 shape and strides must use int64") - if shape.ndim != 1 or strides.ndim != 1 or shape.numel() != strides.numel() or shape.numel() < 2: - raise ValueError("AdamW4bit rank-1 metadata must describe a tensor with rank >= 2") - if scales.ndim != 1: - raise ValueError("AdamW4bit rank-1 scales must be flat") - - @torch._dynamo.disable @torch.no_grad() -def areno_adamw_4bit_rank1_stats( +def areno_adamw_4bit_factored_stats( grad: torch.Tensor, - exp_avg_sq_q: torch.Tensor, - previous_scales: torch.Tensor, - updated_scales: torch.Tensor, + factor_sums: torch.Tensor, invalid: torch.Tensor, - shape: torch.Tensor, - strides: torch.Tensor, *, - packed_offset: int, parameter_shard_start: int, - quant_block_size: int, - beta2: float, - has_state: bool = True, + rows: int, + columns: int, ) -> None: - """Accumulate updated rank-1 second-moment maxima in bounded CUDA blocks.""" + """Accumulate matrix row/column gradient-square sums.""" - tensors = (grad, exp_avg_sq_q, previous_scales, updated_scales, invalid, shape, strides) + tensors = (grad, factor_sums, invalid) if any(not tensor.is_cuda for tensor in tensors): - raise ValueError("fused AdamW4bit rank-1 statistics require CUDA tensors") + raise ValueError("fused AdamW4bit factored statistics require CUDA tensors") if any(tensor.device != grad.device for tensor in tensors[1:]): - raise ValueError("fused AdamW4bit rank-1 statistics require tensors on one device") + raise ValueError("fused AdamW4bit factored statistics require tensors on one device") if grad.dtype not in {torch.bfloat16, torch.float32}: - raise TypeError("AdamW4bit rank-1 gradients must be bfloat16 or float32") - if exp_avg_sq_q.dtype != torch.uint8 or previous_scales.dtype != torch.float32: - raise TypeError("AdamW4bit rank-1 state must use packed uint8 codes and float32 scales") - if updated_scales.dtype != torch.float32 or invalid.dtype != torch.int32 or invalid.numel() != 1: - raise TypeError("AdamW4bit rank-1 outputs must use float32 scales and one int32 validity flag") + raise TypeError("AdamW4bit factored gradients must be bfloat16 or float32") + if factor_sums.dtype != torch.float32 or invalid.dtype != torch.int32 or invalid.numel() != 1: + raise TypeError("AdamW4bit factored outputs must use float32 sums and one int32 validity flag") if any(not tensor.is_contiguous() for tensor in tensors): - raise ValueError("fused AdamW4bit rank-1 statistics require contiguous tensors") - _validate_rank1_metadata(shape, strides, previous_scales) - if previous_scales.numel() != updated_scales.numel(): - raise ValueError("AdamW4bit previous and updated rank-1 scale layouts must match") - packed_numel = (grad.numel() + 1) // 2 - if has_state and (packed_offset < 0 or packed_offset + packed_numel > exp_avg_sq_q.numel()): - raise ValueError("AdamW4bit rank-1 packed variance slice is out of bounds") - if parameter_shard_start < 0: - raise ValueError("AdamW4bit rank-1 parameter shard start must be non-negative") - extension().areno_adamw_4bit_rank1_stats( + raise ValueError("fused AdamW4bit factored statistics require contiguous tensors") + if rows < 1 or columns < 1 or factor_sums.numel() != rows + columns: + raise ValueError("AdamW4bit factored statistics have an invalid matrix shape") + if parameter_shard_start < 0 or parameter_shard_start + grad.numel() > rows * columns: + raise ValueError("AdamW4bit factored gradient slice is out of bounds") + extension().areno_adamw_4bit_factored_stats( grad, - exp_avg_sq_q, - previous_scales, - updated_scales, + factor_sums, invalid, - shape, - strides, - packed_offset, parameter_shard_start, - quant_block_size, - beta2, - has_state, + rows, + columns, ) @torch._dynamo.disable @torch.no_grad() -def areno_adamw_4bit_rank1_step( +def areno_adamw_4bit_factored_step( model: torch.Tensor, grad: torch.Tensor, exp_avg_q: torch.Tensor, exp_avg_scale: torch.Tensor, - exp_avg_sq_q: torch.Tensor, - previous_scales: torch.Tensor, - updated_scales: torch.Tensor, + factors: torch.Tensor, + row_mean: torch.Tensor, invalid: torch.Tensor, - shape: torch.Tensor, - strides: torch.Tensor, *, - packed_offset: int, + moment_packed_offset: int, moment_scale_offset: int, parameter_shard_start: int, quant_block_size: int, + rows: int, + columns: int, beta1: float, - beta2: float, effective_lr: float, weight_decay: float, eps: float, step_size: float, bias_correction2_sqrt: float, ) -> None: - """Update packed AdamW4bit state using precomputed rank-1 scales.""" + """Update packed momentum using an Adafactor-style variance estimate.""" tensors = ( model, grad, exp_avg_q, exp_avg_scale, - exp_avg_sq_q, - previous_scales, - updated_scales, + factors, + row_mean, invalid, - shape, - strides, ) if any(not tensor.is_cuda for tensor in tensors): - raise ValueError("fused AdamW4bit rank-1 update requires CUDA tensors") + raise ValueError("fused AdamW4bit factored update requires CUDA tensors") if any(tensor.device != model.device for tensor in tensors[1:]): - raise ValueError("fused AdamW4bit rank-1 update requires tensors on one device") + raise ValueError("fused AdamW4bit factored update requires tensors on one device") if model.dtype not in {torch.bfloat16, torch.float32} or grad.dtype not in {torch.bfloat16, torch.float32}: - raise TypeError("AdamW4bit rank-1 model and gradient must be bfloat16 or float32") - if exp_avg_q.dtype != torch.uint8 or exp_avg_sq_q.dtype != torch.uint8: - raise TypeError("AdamW4bit rank-1 moments must use packed uint8 storage") - if any(scale.dtype != torch.float32 for scale in (exp_avg_scale, previous_scales, updated_scales)): - raise TypeError("AdamW4bit rank-1 scales must use float32") + raise TypeError("AdamW4bit factored model and gradient must be bfloat16 or float32") + if exp_avg_q.dtype != torch.uint8: + raise TypeError("AdamW4bit factored momentum must use packed uint8 storage") + if any(value.dtype != torch.float32 for value in (exp_avg_scale, factors, row_mean)): + raise TypeError("AdamW4bit factored statistics must use float32") if invalid.dtype != torch.int32 or invalid.numel() != 1: - raise TypeError("AdamW4bit rank-1 update requires one int32 validity flag") + raise TypeError("AdamW4bit factored update requires one int32 validity flag") if any(not tensor.is_contiguous() for tensor in tensors): - raise ValueError("fused AdamW4bit rank-1 update requires contiguous tensors") + raise ValueError("fused AdamW4bit factored update requires contiguous tensors") if model.numel() != grad.numel(): - raise ValueError("AdamW4bit rank-1 model and gradient sizes must match") - _validate_rank1_metadata(shape, strides, previous_scales) - if previous_scales.numel() != updated_scales.numel(): - raise ValueError("AdamW4bit previous and updated rank-1 scale layouts must match") + raise ValueError("AdamW4bit factored model and gradient sizes must match") + if quant_block_size < 32 or quant_block_size > 1024 or quant_block_size & (quant_block_size - 1): + raise ValueError("quant_block_size must be a power of two between 32 and 1024") + if rows < 1 or columns < 1 or factors.numel() != rows + columns or row_mean.numel() != 1: + raise ValueError("AdamW4bit factored update has an invalid matrix shape") packed_numel = (model.numel() + 1) // 2 - if packed_offset < 0 or packed_offset + packed_numel > exp_avg_q.numel(): - raise ValueError("AdamW4bit rank-1 packed moment slice is out of bounds") - if exp_avg_q.numel() != exp_avg_sq_q.numel(): - raise ValueError("AdamW4bit rank-1 packed moments must have the same length") + if moment_packed_offset < 0 or moment_packed_offset + packed_numel > exp_avg_q.numel(): + raise ValueError("AdamW4bit factored packed moment slice is out of bounds") scale_numel = (model.numel() + quant_block_size - 1) // quant_block_size if moment_scale_offset < 0 or moment_scale_offset + scale_numel > exp_avg_scale.numel(): - raise ValueError("AdamW4bit rank-1 first-moment scale slice is out of bounds") - if parameter_shard_start < 0: - raise ValueError("AdamW4bit rank-1 parameter shard start must be non-negative") - extension().areno_adamw_4bit_rank1_step( + raise ValueError("AdamW4bit factored first-moment scale slice is out of bounds") + if parameter_shard_start < 0 or parameter_shard_start + model.numel() > rows * columns: + raise ValueError("AdamW4bit factored parameter shard is out of bounds") + extension().areno_adamw_4bit_factored_step( model, grad, exp_avg_q, exp_avg_scale, - exp_avg_sq_q, - previous_scales, - updated_scales, + factors, + row_mean, invalid, - shape, - strides, - packed_offset, + moment_packed_offset, moment_scale_offset, parameter_shard_start, quant_block_size, + rows, + columns, beta1, - beta2, effective_lr, weight_decay, eps, @@ -428,8 +395,8 @@ def areno_adamw_4bit_rank1_step( __all__ = [ "areno_adamw_4bit_step", - "areno_adamw_4bit_rank1_stats", - "areno_adamw_4bit_rank1_step", + "areno_adamw_4bit_factored_stats", + "areno_adamw_4bit_factored_step", "areno_adamw_8bit_step", "areno_adamw_fp32_master_step", "areno_adamw_fp32_state_step", diff --git a/areno/engine/optim/adamw_4bit.py b/areno/engine/optim/adamw_4bit.py index be407c25..3313c04e 100644 --- a/areno/engine/optim/adamw_4bit.py +++ b/areno/engine/optim/adamw_4bit.py @@ -1,16 +1,16 @@ -"""Packed rank-1-normalized 4-bit-state AdamW. +"""Packed 4-bit first-moment AdamW with factored second moments. The first moment uses signed dynamic-exponent quantization. The non-negative second moment uses the zero-excluding linear map from Li et al. (NeurIPS 2023): -code ``i`` represents ``scale * (i + 1) / 16``. Matrix and higher-rank second -moments use the paper's rank-1 normalization; vectors retain B=128 block -normalization. Two codes are packed in each byte. +code ``i`` represents ``scale * (i + 1) / 16``. Vectors retain that B=128 +representation. Matrix and higher-rank tensors instead use Adafactor-style +row/column second-moment statistics, eliminating their elementwise variance +state. Two first-moment codes are packed in each byte. """ from __future__ import annotations from collections.abc import Iterable, Iterator -from math import prod import torch import torch.distributed as dist @@ -19,7 +19,7 @@ from areno.engine.optim.adamw_fp32_master import _DEFAULT_BUCKET_NUMEL, _MasterBucket, _param_grad, _ParamRef _DEFAULT_QUANT_BLOCK_SIZE = 128 -_STATE_FORMAT_VERSION = 2 +_STATE_FORMAT_VERSION = 3 # The signed 4-bit dynamic-exponent map used by the reference implementation # of Li et al. Values are normalized by each block's absolute maximum. _SIGNED_DE_MAP = ( @@ -43,18 +43,21 @@ class AdamW4bit(AdamW8bit): - """AdamW with two packed 4-bit moments and shape-aware FP32 scales. + """AdamW with packed 4-bit momentum and factored matrix variance. First-moment quantization blocks restart at every parameter shard. For a - tensor with rank >= 2, second-moment scales are the minimum of per-axis - maxima over the original local model-tensor shape. DP ranks combine their - partial axis maxima with MAX. One-dimensional tensors retain parameter- - local block scales. CPU and CUDA updates keep FP32 moment work bounded by + tensor with rank >= 2, the second moment is represented by row and column + means over the original TP-local tensor, flattened as ``[shape[0], -1]``. + DP ranks combine partial sums before applying the exponential update. + One-dimensional tensors retain parameter-local packed B=128 variance. + CPU and CUDA updates keep elementwise FP32 work bounded by ``quant_block_size``. """ _embedding_fp32_state = False - state_quantizer = "signed-de4/rank1-zero-excluding-linear4" + gradient_shard_dtype = torch.bfloat16 + stream_gradient_shards = True + state_quantizer = "signed-de4/factored-second-moment-v1" def _precision_for_parameter(self, parameter: torch.nn.Parameter) -> str: del parameter @@ -87,9 +90,7 @@ def __init__( dp_group=dp_group, quant_block_size=quant_block_size, ) - self._rank1_metadata_cache: dict[tuple[int, torch.device], tuple[torch.Tensor, torch.Tensor]] = {} - self._rank1_empty_codes: dict[torch.device, torch.Tensor] = {} - self._rank1_scales: dict[int, torch.Tensor | None] = { + self._factored_second_moments: dict[int, torch.Tensor | None] = { id(parameter): None for parameter in self.model_params if parameter.ndim >= 2 } @@ -110,10 +111,10 @@ def state_dict(self) -> dict: payload["state_format_version"] = _STATE_FORMAT_VERSION payload["quant_block_size"] = self.quant_block_size payload["parameter_shapes"] = [tuple(parameter.shape) for parameter in self.model_params] - payload["rank1_scales"] = [ + payload["factored_second_moments"] = [ None - if self._rank1_scales.get(id(parameter)) is None - else self._rank1_scales[id(parameter)].detach().to(device="cpu").clone() + if self._factored_second_moments.get(id(parameter)) is None + else self._factored_second_moments[id(parameter)].detach().to(device="cpu").clone() for parameter in self.model_params ] payload["state_memory"] = self.state_memory_metrics() @@ -135,9 +136,9 @@ def load_state_dict(self, state_dict: dict) -> None: current_shapes = [tuple(parameter.shape) for parameter in self.model_params] if saved_shapes != current_shapes: raise ValueError("AdamW4bit checkpoint parameter shapes do not match the optimizer") - saved_rank1_scales = state_dict.get("rank1_scales") - if not isinstance(saved_rank1_scales, list) or len(saved_rank1_scales) != len(self.model_params): - raise ValueError("AdamW4bit checkpoint rank-1 scales do not match the optimizer parameters") + saved_factors = state_dict.get("factored_second_moments") + if not isinstance(saved_factors, list) or len(saved_factors) != len(self.model_params): + raise ValueError("AdamW4bit checkpoint factored moments do not match the optimizer parameters") self._cleanup_disk_offload() self._active_offload_mode = "none" self._disk_offload_root = None @@ -157,57 +158,55 @@ def load_state_dict(self, state_dict: dict) -> None: if saved is None: continue device = bucket.refs[0].model_param.device - packed_numel, moment_scale_numel, variance_scale_numel = self._bucket_state_sizes(bucket) + moment_packed_numel, moment_scale_numel, variance_packed_numel, variance_scale_numel = ( + self._bucket_state_sizes(bucket) + ) state.step = int(saved.get("step", 0)) - state.exp_avg_q = _load_tensor(saved, "exp_avg_q", device, torch.uint8, packed_numel) + state.exp_avg_q = _load_tensor(saved, "exp_avg_q", device, torch.uint8, moment_packed_numel) state.exp_avg_scale = _load_tensor(saved, "exp_avg_scale", device, torch.float32, moment_scale_numel) - state.exp_avg_sq_q = _load_tensor(saved, "exp_avg_sq_q", device, torch.uint8, packed_numel) + state.exp_avg_sq_q = _load_tensor(saved, "exp_avg_sq_q", device, torch.uint8, variance_packed_numel) state.exp_avg_sq_scale = _load_tensor( saved, "exp_avg_sq_scale", device, torch.float32, variance_scale_numel ) - for parameter, saved_scales in zip(self.model_params, saved_rank1_scales, strict=True): + for parameter, saved_factors_for_parameter in zip(self.model_params, saved_factors, strict=True): if parameter.ndim < 2: - if saved_scales is not None: - raise ValueError("AdamW4bit checkpoint has rank-1 scales for a one-dimensional parameter") + if saved_factors_for_parameter is not None: + raise ValueError("AdamW4bit checkpoint has factored state for a one-dimensional parameter") continue - if saved_scales is None: - self._rank1_scales[id(parameter)] = None + if saved_factors_for_parameter is None: + self._factored_second_moments[id(parameter)] = None continue - restored_scales = saved_scales.detach().to(device=parameter.device, dtype=torch.float32).view(-1).clone() - if restored_scales.numel() != sum(parameter.shape): - raise ValueError("AdamW4bit checkpoint rank-1 scale length does not match the parameter shape") - self._rank1_scales[id(parameter)] = restored_scales + restored_factors = ( + saved_factors_for_parameter.detach().to(device=parameter.device, dtype=torch.float32).view(-1).clone() + ) + if restored_factors.numel() != _factored_state_numel_for_parameter(parameter): + raise ValueError("AdamW4bit checkpoint factored state length does not match the parameter shape") + self._factored_second_moments[id(parameter)] = restored_factors def clear_state(self) -> None: - """Drop packed moments and parameter-level rank-1 metadata.""" + """Drop packed moments and parameter-level factored state.""" super().clear_state() - for parameter_id in self._rank1_scales: - self._rank1_scales[parameter_id] = None - self._rank1_metadata_cache.clear() - self._rank1_empty_codes.clear() + for parameter_id in self._factored_second_moments: + self._factored_second_moments[parameter_id] = None @torch.no_grad() def offload_state(self, mode: str = "cpu", directory: str | None = None, batch_size: int = 1) -> None: - """Offload packed buckets and keep small rank-1 metadata on CPU.""" + """Offload packed buckets and keep small factored state on CPU.""" super().offload_state(mode=mode, directory=directory, batch_size=batch_size) - for parameter_id, scales in self._rank1_scales.items(): - if scales is not None and scales.device.type != "cpu": - self._rank1_scales[parameter_id] = scales.to(device="cpu") - self._rank1_metadata_cache.clear() - self._rank1_empty_codes.clear() + for parameter_id, factors in self._factored_second_moments.items(): + if factors is not None and factors.device.type != "cpu": + self._factored_second_moments[parameter_id] = factors.to(device="cpu") @torch.no_grad() def onload_state(self, device: torch.device) -> None: - """Restore packed buckets and shared rank-1 metadata to ``device``.""" + """Restore packed buckets and shared factored state to ``device``.""" super().onload_state(device) - for parameter_id, scales in self._rank1_scales.items(): - if scales is not None and scales.device != device: - self._rank1_scales[parameter_id] = scales.to(device=device) - self._rank1_metadata_cache.clear() - self._rank1_empty_codes.clear() + for parameter_id, factors in self._factored_second_moments.items(): + if factors is not None and factors.device != device: + self._factored_second_moments[parameter_id] = factors.to(device=device) @torch.no_grad() def _ensure_bucket_state(self, bucket: _MasterBucket, state) -> None: @@ -220,13 +219,15 @@ def _ensure_bucket_state(self, bucket: _MasterBucket, state) -> None: value = getattr(state, name) if value is not None and value.device != device: setattr(state, name, value.to(device=device)) - packed_numel, moment_scale_numel, variance_scale_numel = self._bucket_state_sizes(bucket) + moment_packed_numel, moment_scale_numel, variance_packed_numel, variance_scale_numel = self._bucket_state_sizes( + bucket + ) if state.exp_avg_q is None: # Signed dynamic-exponent zero has code 7, hence byte 0x77. - state.exp_avg_q = torch.full((packed_numel,), 0x77, device=device, dtype=torch.uint8) + state.exp_avg_q = torch.full((moment_packed_numel,), 0x77, device=device, dtype=torch.uint8) state.exp_avg_scale = torch.ones(moment_scale_numel, device=device, dtype=torch.float32) if state.exp_avg_sq_q is None: - state.exp_avg_sq_q = torch.zeros(packed_numel, device=device, dtype=torch.uint8) + state.exp_avg_sq_q = torch.zeros(variance_packed_numel, device=device, dtype=torch.uint8) # A zero scale makes the zero-excluding code initially decode to 0. state.exp_avg_sq_scale = torch.zeros(variance_scale_numel, device=device, dtype=torch.float32) @@ -235,198 +236,143 @@ def _state_mmap_specs(self, indices: list[int]) -> dict[int, dict[str, tuple[tor specs: dict[int, dict[str, tuple[torch.dtype, tuple[int, ...]]]] = {} for index in indices: - packed_numel, moment_scale_numel, variance_scale_numel = self._bucket_state_sizes(self.buckets[index]) + moment_packed_numel, moment_scale_numel, variance_packed_numel, variance_scale_numel = ( + self._bucket_state_sizes(self.buckets[index]) + ) specs[index] = { - "exp_avg_q": (torch.uint8, (packed_numel,)), + "exp_avg_q": (torch.uint8, (moment_packed_numel,)), "exp_avg_scale": (torch.float32, (moment_scale_numel,)), - "exp_avg_sq_q": (torch.uint8, (packed_numel,)), + "exp_avg_sq_q": (torch.uint8, (variance_packed_numel,)), "exp_avg_sq_scale": (torch.float32, (variance_scale_numel,)), } return specs @torch.no_grad() def step(self, closure=None): - """Apply a streaming statistics pass before rank-1 parameter updates.""" + """Update one complete parameter at a time and release ready buckets.""" if closure is not None: with torch.enable_grad(): closure() - active_rank1_parameters = { - id(ref.model_param) - for bucket in self.buckets - for ref in bucket.refs - if _uses_rank1_normalization(ref) and self._ref_has_gradient(bucket, ref) - } - if not active_rank1_parameters: - return super().step() + layouts: dict[int, list[tuple[int, _ParamRef, int, int, int, int]]] = {} + remaining_by_bucket: dict[int, int] = {} + for index, bucket in enumerate(self.buckets): + for ref, moment_packed, moment_scale, variance_packed, variance_scale in self._iter_ref_layout(bucket): + if not self._ref_has_gradient(bucket, ref): + continue + layouts.setdefault(id(ref.model_param), []).append( + (index, ref, moment_packed, moment_scale, variance_packed, variance_scale) + ) + remaining_by_bucket[index] = remaining_by_bucket.get(index, 0) + 1 + if not layouts: + return None + + beta1, beta2 = self.betas + started_buckets: set[int] = set() + completed_buckets: set[int] = set() - rank1_work: dict[int, tuple[torch.Tensor, torch.Tensor]] = {} for parameter in self.model_params: - if id(parameter) in active_rank1_parameters: - rank1_work[id(parameter)] = ( - torch.zeros(_rank1_scale_numel_for_parameter(parameter), device=parameter.device), - torch.zeros((), device=parameter.device, dtype=torch.int32), - ) - self._ensure_rank1_scales(parameter) - - # First streaming pass: each chunk contributes to one parameter-level - # set of axis maxima. Disk-offloaded packed state is returned to disk - # after each group, so this pass does not pin all buckets on the GPU. - for indices in self._bucket_groups(): - group_changed = False - for index in indices: - bucket = self.buckets[index] - state = self._states[index] - if not any( - _uses_rank1_normalization(ref) and self._ref_has_gradient(bucket, ref) for ref in bucket.refs - ): - continue - # A fresh optimizer has an implicit all-zero second moment. - # Do not materialize every packed bucket during the statistics - # pass while every FP32 gradient shard is still resident. The - # update pass below initializes one bucket at a time and drops - # its gradient immediately, preserving streaming peak memory. - if state.step > 0 or state.exp_avg_sq_q is not None or state.offload_file is not None: - self._ensure_bucket_state(bucket, state) - for ref, packed_offset, _moment_scale_offset, _variance_scale_offset in self._iter_ref_layout(bucket): - if not _uses_rank1_normalization(ref) or not self._ref_has_gradient(bucket, ref): - continue - updated_scales, invalid = rank1_work[id(ref.model_param)] - self._rank1_variance_statistics( + parameter_layouts = layouts.get(id(parameter)) + if not parameter_layouts: + continue + factored_work: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None + if parameter.ndim >= 2: + previous_factors = self._ensure_factored_second_moment(parameter) + factor_sums = torch.zeros_like(previous_factors) + invalid = torch.zeros((), device=parameter.device, dtype=torch.int32) + for index, ref, _moment_packed, _moment_scale, _variance_packed, _variance_scale in parameter_layouts: + bucket = self.buckets[index] + self._factored_variance_statistics( bucket, ref, self._gradient_for_ref(bucket, ref), - state, - packed_offset, - self._ensure_rank1_scales(ref.model_param), - updated_scales, + factor_sums, invalid, - self.betas[1], ) - group_changed = True - if self._active_offload_mode == "disk": - self._stage_8bit_state_on_cpu(state) - if self._active_offload_mode == "disk" and group_changed: - self._offload_8bit_group_to_disk(indices) - - for parameter in self.model_params: - work = rank1_work.get(id(parameter)) - if work is None: - continue - updated_scales, invalid = work - if self.dp_size > 1: - if self.dp_group is None: - raise RuntimeError("AdamW4bit DP rank-1 normalization requires a DP process group") - dist.all_reduce(updated_scales, op=dist.ReduceOp.MAX, group=self.dp_group) - dist.all_reduce(invalid, op=dist.ReduceOp.MAX, group=self.dp_group) - - # Second streaming pass: recompute each bounded block, update weights, - # and write packed moments using the now-global axis scales. - for indices in self._bucket_groups(): - group_changed = False - for index in indices: + if self.dp_size > 1: + if self.dp_group is None: + raise RuntimeError("AdamW4bit factored second moments require a DP process group") + dist.all_reduce(factor_sums, op=dist.ReduceOp.SUM, group=self.dp_group) + dist.all_reduce(invalid, op=dist.ReduceOp.MAX, group=self.dp_group) + updated_factors = self._finalize_factored_statistics(parameter, previous_factors, factor_sums, beta2) + row_mean = updated_factors[: parameter.shape[0]].mean() + factored_work = updated_factors, row_mean, invalid + + for index, ref, moment_packed, moment_scale, variance_packed, variance_scale in parameter_layouts: bucket = self.buckets[index] state = self._states[index] - has_grad = bucket.grad_shard is not None or any( - _param_grad(ref.model_param) is not None for ref in bucket.refs - ) - if has_grad: + if index not in started_buckets: self._ensure_bucket_state(bucket, state) - self._step_bucket_8bit(bucket, state, rank1_work) - group_changed = True - if self._active_offload_mode == "disk": - self._stage_8bit_state_on_cpu(state) - elif self._active_offload_mode == "disk": - self._discard_disk_prefetch(index) - if self._active_offload_mode == "disk" and group_changed: - self._offload_8bit_group_to_disk(indices) - - for parameter in self.model_params: - work = rank1_work.get(id(parameter)) - if work is None: - continue - updated_scales, invalid = work - scale_storage = self._ensure_rank1_scales(parameter) - if invalid.is_cuda: - scale_storage.copy_(torch.where(invalid == 0, updated_scales, scale_storage)) - elif int(invalid.item()) == 0: - scale_storage.copy_(updated_scales) - if self._active_offload_mode == "disk" and scale_storage.device.type != "cpu": - self._rank1_scales[id(parameter)] = scale_storage.to(device="cpu") + state.step += 1 + started_buckets.add(index) + grad = self._gradient_for_ref(bucket, ref) + if grad is not None: + effective_lr = float(getattr(parameter, "_areno_lr", self.lr)) + bias_correction1 = 1.0 - beta1**state.step + bias_correction2_sqrt = (1.0 - beta2**state.step) ** 0.5 + if factored_work is None: + self._step_param_ref_4bit( + bucket, + ref, + grad, + state, + moment_packed, + moment_scale, + variance_packed, + variance_scale, + beta1, + beta2, + effective_lr, + effective_lr / bias_correction1, + bias_correction2_sqrt, + ) + else: + updated_factors, row_mean, invalid = factored_work + self._step_param_ref_factored( + bucket, + ref, + grad, + state, + moment_packed, + moment_scale, + updated_factors, + row_mean, + invalid, + beta1, + effective_lr, + effective_lr / bias_correction1, + bias_correction2_sqrt, + ) + remaining_by_bucket[index] -= 1 + if remaining_by_bucket[index] == 0: + self._all_gather_bucket(bucket) + bucket.grad_shard = None + bucket.grad_param_ids = frozenset() + completed_buckets.add(index) + + if factored_work is not None: + updated_factors, _row_mean, invalid = factored_work + factor_storage = self._ensure_factored_second_moment(parameter) + if invalid.is_cuda: + factor_storage.copy_(torch.where(invalid == 0, updated_factors, factor_storage)) + elif int(invalid.item()) == 0: + factor_storage.copy_(updated_factors) + parameter.grad = None + if isinstance(getattr(parameter, "main_grad", None), torch.Tensor): + parameter.main_grad = None + + if self._active_offload_mode == "disk": + for indices in self._bucket_groups(): + changed = [index for index in indices if index in completed_buckets] + for index in changed: + self._stage_8bit_state_on_cpu(self._states[index]) + if changed: + self._offload_8bit_group_to_disk(indices) + for parameter_id, factors in self._factored_second_moments.items(): + if factors is not None and factors.device.type != "cpu": + self._factored_second_moments[parameter_id] = factors.to(device="cpu") return None - @torch.no_grad() - def _step_bucket_8bit( - self, - bucket: _MasterBucket, - state, - rank1_work: dict[int, tuple[torch.Tensor, torch.Tensor]] | None = None, - ) -> None: - """Update a bucket while materializing at most one FP32 block per moment.""" - - assert state.exp_avg_q is not None - assert state.exp_avg_scale is not None - assert state.exp_avg_sq_q is not None - assert state.exp_avg_sq_scale is not None - beta1, beta2 = self.betas - state.step += 1 - bias_correction1 = 1.0 - beta1**state.step - bias_correction2_sqrt = (1.0 - beta2**state.step) ** 0.5 - for ref, packed_offset, moment_scale_offset, variance_scale_offset in self._iter_ref_layout(bucket): - has_parameter_grad = ( - id(ref.model_param) in bucket.grad_param_ids - if bucket.grad_shard is not None - else ref.model_param.grad is not None - or isinstance(getattr(ref.model_param, "main_grad", None), torch.Tensor) - ) - if not has_parameter_grad: - continue - grad = self._gradient_for_ref(bucket, ref) - effective_lr = float(getattr(ref.model_param, "_areno_lr", self.lr)) - if _uses_rank1_normalization(ref): - if rank1_work is None: - raise RuntimeError("AdamW4bit rank-1 update is missing its parameter statistics pass") - updated_scales, invalid = rank1_work[id(ref.model_param)] - if grad is not None and (invalid.is_cuda or int(invalid.item()) == 0): - self._step_param_ref_rank1( - bucket, - ref, - grad, - state, - packed_offset, - moment_scale_offset, - self._ensure_rank1_scales(ref.model_param), - updated_scales, - invalid, - beta1, - beta2, - effective_lr, - effective_lr / bias_correction1, - bias_correction2_sqrt, - ) - elif grad is not None: - self._step_param_ref_4bit( - bucket, - ref, - grad, - state, - packed_offset, - moment_scale_offset, - variance_scale_offset, - beta1, - beta2, - effective_lr, - effective_lr / bias_correction1, - bias_correction2_sqrt, - ) - if ref.param_start + ref.numel == ref.model_param.numel(): - ref.model_param.grad = None - if isinstance(getattr(ref.model_param, "main_grad", None), torch.Tensor): - ref.model_param.main_grad = None - self._all_gather_bucket(bucket) - bucket.grad_shard = None - bucket.grad_param_ids = frozenset() - @staticmethod def _ref_has_gradient(bucket: _MasterBucket, ref: _ParamRef) -> bool: if bucket.grad_shard is not None: @@ -440,8 +386,9 @@ def _step_param_ref_4bit( ref: _ParamRef, grad: torch.Tensor, state, - packed_offset: int, + moment_packed_offset: int, moment_scale_offset: int, + variance_packed_offset: int, variance_scale_offset: int, beta1: float, beta2: float, @@ -466,8 +413,9 @@ def _step_param_ref_4bit( state.exp_avg_scale, state.exp_avg_sq_q, state.exp_avg_sq_scale, - packed_offset=packed_offset, + moment_packed_offset=moment_packed_offset, moment_scale_offset=moment_scale_offset, + variance_packed_offset=variance_packed_offset, variance_scale_offset=variance_scale_offset, quant_block_size=self.quant_block_size, beta1=beta1, @@ -481,17 +429,18 @@ def _step_param_ref_4bit( return for block_index, start in enumerate(range(0, ref.shard_numel, self.quant_block_size)): count = min(self.quant_block_size, ref.shard_numel - start) - byte_start = packed_offset + start // 2 + moment_byte_start = moment_packed_offset + start // 2 + variance_byte_start = variance_packed_offset + start // 2 byte_count = (count + 1) // 2 moment_scale_index = moment_scale_offset + block_index variance_scale_index = variance_scale_offset + block_index moment = _unpack_signed_4bit( - state.exp_avg_q.narrow(0, byte_start, byte_count), + state.exp_avg_q.narrow(0, moment_byte_start, byte_count), count, state.exp_avg_scale[moment_scale_index], ) variance = _unpack_positive_4bit( - state.exp_avg_sq_q.narrow(0, byte_start, byte_count), + state.exp_avg_sq_q.narrow(0, variance_byte_start, byte_count), count, state.exp_avg_sq_scale[variance_scale_index], ) @@ -515,126 +464,117 @@ def _step_param_ref_4bit( model_shard.narrow(0, start, count).copy_(weight) moment_q, moment_scale = _quantize_signed_4bit(moment) variance_q, variance_scale = _quantize_positive_4bit(variance) - state.exp_avg_q.narrow(0, byte_start, byte_count).copy_(moment_q) + state.exp_avg_q.narrow(0, moment_byte_start, byte_count).copy_(moment_q) state.exp_avg_scale[moment_scale_index].copy_(moment_scale) - state.exp_avg_sq_q.narrow(0, byte_start, byte_count).copy_(variance_q) + state.exp_avg_sq_q.narrow(0, variance_byte_start, byte_count).copy_(variance_q) state.exp_avg_sq_scale[variance_scale_index].copy_(variance_scale) @torch.no_grad() - def _rank1_variance_statistics( + def _factored_variance_statistics( self, bucket: _MasterBucket, ref: _ParamRef, grad: torch.Tensor | None, - state, - packed_offset: int, - previous_scales: torch.Tensor, - updated_scales: torch.Tensor, + factor_sums: torch.Tensor, invalid: torch.Tensor, - beta2: float, ) -> None: - """Compute updated per-axis maxima without a full FP32 moment.""" + """Accumulate row/column gradient-square sums for one DP shard.""" if grad is None or ref.shard_numel == 0: return grad_shard = grad if bucket.grad_shard is not None else grad.narrow(0, ref.shard_start, ref.shard_numel) parameter_shard_start = ref.param_start + ref.shard_start if grad_shard.is_cuda: - from areno.accel.optimizer import areno_adamw_4bit_rank1_stats - - shape, strides = self._rank1_metadata(ref) - variance_codes = state.exp_avg_sq_q - has_state = variance_codes is not None - if variance_codes is None: - variance_codes = self._rank1_empty_codes.get(grad_shard.device) - if variance_codes is None: - variance_codes = torch.empty(1, device=grad_shard.device, dtype=torch.uint8) - self._rank1_empty_codes[grad_shard.device] = variance_codes - areno_adamw_4bit_rank1_stats( + from areno.accel.optimizer import areno_adamw_4bit_factored_stats + + areno_adamw_4bit_factored_stats( grad_shard.contiguous(), - variance_codes, - previous_scales, - updated_scales, + factor_sums, invalid, - shape, - strides, - packed_offset=packed_offset if has_state else 0, parameter_shard_start=parameter_shard_start, - quant_block_size=self.quant_block_size, - beta2=beta2, - has_state=has_state, + rows=ref.model_param.shape[0], + columns=ref.model_param.numel() // ref.model_param.shape[0], ) return - shape_tuple = tuple(ref.model_param.shape) + rows = ref.model_param.shape[0] + columns = ref.model_param.numel() // rows + row_sums = factor_sums[:rows] + column_sums = factor_sums[rows:] for start in range(0, ref.shard_numel, self.quant_block_size): count = min(self.quant_block_size, ref.shard_numel - start) - byte_start = packed_offset + start // 2 - byte_count = (count + 1) // 2 flat_start = parameter_shard_start + start - element_scales = _rank1_element_scales(previous_scales, shape_tuple, flat_start, count) gradient = grad_shard.narrow(0, start, count).to(dtype=torch.float32) - if state.exp_avg_sq_q is None: - variance = torch.zeros(count, device=gradient.device, dtype=torch.float32) - else: - variance = _unpack_positive_4bit_elementwise( - state.exp_avg_sq_q.narrow(0, byte_start, byte_count), count, element_scales - ) - variance.mul_(beta2).addcmul_(gradient, gradient, value=1.0 - beta2) - if not bool(torch.isfinite(gradient).all() & torch.isfinite(variance).all()): + squared = gradient.square() + if not bool(torch.isfinite(squared).all()): invalid.fill_(1) return - _accumulate_rank1_maxima(updated_scales, variance, shape_tuple, flat_start) + flat_indices = torch.arange(flat_start, flat_start + count, device=gradient.device) + row_sums.index_add_(0, torch.div(flat_indices, columns, rounding_mode="floor"), squared) + column_sums.index_add_(0, flat_indices.remainder(columns), squared) + + @staticmethod + def _finalize_factored_statistics( + parameter: torch.nn.Parameter, + previous_factors: torch.Tensor, + factor_sums: torch.Tensor, + beta2: float, + ) -> torch.Tensor: + """Convert global sums to EMA row/column means in-place.""" + + rows = parameter.shape[0] + columns = parameter.numel() // rows + updated = factor_sums + updated[:rows].div_(columns).mul_(1.0 - beta2).add_(previous_factors[:rows], alpha=beta2) + updated[rows:].div_(rows).mul_(1.0 - beta2).add_(previous_factors[rows:], alpha=beta2) + return updated @torch.no_grad() - def _step_param_ref_rank1( + def _step_param_ref_factored( self, bucket: _MasterBucket, ref: _ParamRef, grad: torch.Tensor, state, - packed_offset: int, + moment_packed_offset: int, moment_scale_offset: int, - previous_scales: torch.Tensor, - updated_scales: torch.Tensor, + updated_factors: torch.Tensor, + row_mean: torch.Tensor, invalid: torch.Tensor, beta1: float, - beta2: float, effective_lr: float, step_size: float, bias_correction2_sqrt: float, ) -> None: - """Recompute bounded Adam blocks and requantize with rank-1 scales.""" + """Update one shard from packed momentum and factored variance.""" assert state.exp_avg_q is not None assert state.exp_avg_scale is not None - assert state.exp_avg_sq_q is not None if ref.shard_numel == 0: return + if not invalid.is_cuda and int(invalid.item()) != 0: + return grad_shard = grad if bucket.grad_shard is not None else grad.narrow(0, ref.shard_start, ref.shard_numel) parameter_shard_start = ref.param_start + ref.shard_start model_shard = ref.model_param.detach().reshape(-1).narrow(0, parameter_shard_start, ref.shard_numel) if model_shard.is_cuda: - from areno.accel.optimizer import areno_adamw_4bit_rank1_step + from areno.accel.optimizer import areno_adamw_4bit_factored_step - shape, strides = self._rank1_metadata(ref) - areno_adamw_4bit_rank1_step( + areno_adamw_4bit_factored_step( model_shard, grad_shard.contiguous(), state.exp_avg_q, state.exp_avg_scale, - state.exp_avg_sq_q, - previous_scales, - updated_scales, + updated_factors, + row_mean, invalid, - shape, - strides, - packed_offset=packed_offset, + moment_packed_offset=moment_packed_offset, moment_scale_offset=moment_scale_offset, parameter_shard_start=parameter_shard_start, quant_block_size=self.quant_block_size, + rows=ref.model_param.shape[0], + columns=ref.model_param.numel() // ref.model_param.shape[0], beta1=beta1, - beta2=beta2, effective_lr=effective_lr, weight_decay=self.weight_decay, eps=self.eps, @@ -643,10 +583,13 @@ def _step_param_ref_rank1( ) return - shape_tuple = tuple(ref.model_param.shape) + rows = ref.model_param.shape[0] + columns = ref.model_param.numel() // rows + row_factors = updated_factors[:rows] + column_factors = updated_factors[rows:] for block_index, start in enumerate(range(0, ref.shard_numel, self.quant_block_size)): count = min(self.quant_block_size, ref.shard_numel - start) - byte_start = packed_offset + start // 2 + byte_start = moment_packed_offset + start // 2 byte_count = (count + 1) // 2 moment_scale_index = moment_scale_offset + block_index moment = _unpack_signed_4bit( @@ -655,16 +598,14 @@ def _step_param_ref_rank1( state.exp_avg_scale[moment_scale_index], ) flat_start = parameter_shard_start + start - old_element_scales = _rank1_element_scales(previous_scales, shape_tuple, flat_start, count) - variance = _unpack_positive_4bit_elementwise( - state.exp_avg_sq_q.narrow(0, byte_start, byte_count), count, old_element_scales - ) gradient = grad_shard.narrow(0, start, count).to(dtype=torch.float32) + flat_indices = torch.arange(flat_start, flat_start + count, device=gradient.device) + variance = row_factors[torch.div(flat_indices, columns, rounding_mode="floor")] + variance = variance * column_factors[flat_indices.remainder(columns)] / row_mean.clamp_min(1.0e-30) weight = model_shard.narrow(0, start, count).to(dtype=torch.float32).clone() if self.weight_decay != 0.0: weight.mul_(1.0 - effective_lr * self.weight_decay) moment.mul_(beta1).add_(gradient, alpha=1.0 - beta1) - variance.mul_(beta2).addcmul_(gradient, gradient, value=1.0 - beta2) denom = variance.sqrt().div_(bias_correction2_sqrt).add_(self.eps) weight.addcdiv_(moment, denom, value=-step_size) if not bool( @@ -674,65 +615,52 @@ def _step_param_ref_rank1( & torch.isfinite(weight).all() ): return - new_element_scales = _rank1_element_scales(updated_scales, shape_tuple, flat_start, count) moment_q, moment_scale = _quantize_signed_4bit(moment) - variance_q = _quantize_positive_4bit_elementwise(variance, new_element_scales) model_shard.narrow(0, start, count).copy_(weight) state.exp_avg_q.narrow(0, byte_start, byte_count).copy_(moment_q) state.exp_avg_scale[moment_scale_index].copy_(moment_scale) - state.exp_avg_sq_q.narrow(0, byte_start, byte_count).copy_(variance_q) - - def _rank1_metadata(self, ref: _ParamRef) -> tuple[torch.Tensor, torch.Tensor]: - key = (id(ref.model_param), ref.model_param.device) - cached = self._rank1_metadata_cache.get(key) - if cached is not None: - return cached - shape_tuple = tuple(int(dimension) for dimension in ref.model_param.shape) - strides_tuple = _contiguous_strides(shape_tuple) - cached = ( - torch.tensor(shape_tuple, device=ref.model_param.device, dtype=torch.int64), - torch.tensor(strides_tuple, device=ref.model_param.device, dtype=torch.int64), - ) - self._rank1_metadata_cache[key] = cached - return cached - def _ensure_rank1_scales(self, parameter: torch.nn.Parameter) -> torch.Tensor: - """Materialize one shared axis-scale tensor for an original parameter.""" + def _ensure_factored_second_moment(self, parameter: torch.nn.Parameter) -> torch.Tensor: + """Materialize row/column second-moment statistics for one parameter.""" key = id(parameter) - scales = self._rank1_scales.get(key) - if scales is None: - scales = torch.zeros( - _rank1_scale_numel_for_parameter(parameter), + factors = self._factored_second_moments.get(key) + if factors is None: + factors = torch.zeros( + _factored_state_numel_for_parameter(parameter), device=parameter.device, dtype=torch.float32, ) - self._rank1_scales[key] = scales - elif scales.device != parameter.device: - scales = scales.to(device=parameter.device) - self._rank1_scales[key] = scales - return scales - - def _bucket_state_sizes(self, bucket: _MasterBucket) -> tuple[int, int, int]: - packed_numel = sum((ref.shard_numel + 1) // 2 for ref in bucket.refs) + self._factored_second_moments[key] = factors + elif factors.device != parameter.device: + factors = factors.to(device=parameter.device) + self._factored_second_moments[key] = factors + return factors + + def _bucket_state_sizes(self, bucket: _MasterBucket) -> tuple[int, int, int, int]: + moment_packed_numel = sum((ref.shard_numel + 1) // 2 for ref in bucket.refs) moment_scale_numel = sum( (ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size for ref in bucket.refs ) + variance_packed_numel = sum((ref.shard_numel + 1) // 2 for ref in bucket.refs if ref.model_param.ndim < 2) variance_scale_numel = sum(self._variance_scale_count(ref) for ref in bucket.refs) - return packed_numel, moment_scale_numel, variance_scale_numel + return moment_packed_numel, moment_scale_numel, variance_packed_numel, variance_scale_numel - def _iter_ref_layout(self, bucket: _MasterBucket) -> Iterator[tuple[_ParamRef, int, int, int]]: - packed_offset = 0 + def _iter_ref_layout(self, bucket: _MasterBucket) -> Iterator[tuple[_ParamRef, int, int, int, int]]: + moment_packed_offset = 0 moment_scale_offset = 0 + variance_packed_offset = 0 variance_scale_offset = 0 for ref in bucket.refs: - yield ref, packed_offset, moment_scale_offset, variance_scale_offset - packed_offset += (ref.shard_numel + 1) // 2 + yield ref, moment_packed_offset, moment_scale_offset, variance_packed_offset, variance_scale_offset + moment_packed_offset += (ref.shard_numel + 1) // 2 moment_scale_offset += (ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size + if ref.model_param.ndim < 2: + variance_packed_offset += (ref.shard_numel + 1) // 2 variance_scale_offset += self._variance_scale_count(ref) def _variance_scale_count(self, ref: _ParamRef) -> int: - if _uses_rank1_normalization(ref): + if ref.model_param.ndim >= 2: return 0 return (ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size @@ -744,7 +672,7 @@ def persistent_moment_bytes(self) -> int: for value in (state.exp_avg_q, state.exp_avg_scale, state.exp_avg_sq_q, state.exp_avg_sq_scale): if value is not None: total += value.numel() * value.element_size() - for value in self._rank1_scales.values(): + for value in self._factored_second_moments.values(): if value is not None: total += value.numel() * value.element_size() return total @@ -763,7 +691,7 @@ def state_memory_metrics(self) -> dict[str, int]: for value in (state.exp_avg_scale, state.exp_avg_sq_scale): if value is not None: scale_metadata_bytes += value.numel() * value.element_size() - for value in self._rank1_scales.values(): + for value in self._factored_second_moments.values(): if value is not None: scale_metadata_bytes += value.numel() * value.element_size() return { @@ -841,80 +769,9 @@ def _unpack_positive_4bit(packed: torch.Tensor, numel: int, scale: torch.Tensor) return (_unpack_nibbles(packed, numel).to(dtype=torch.float32) + 1.0).mul_(scale / 16.0) -def _uses_rank1_normalization(ref: _ParamRef) -> bool: - return ref.model_param.ndim >= 2 - - -def _rank1_scale_numel(ref: _ParamRef) -> int: - return _rank1_scale_numel_for_parameter(ref.model_param) - - -def _rank1_scale_numel_for_parameter(parameter: torch.nn.Parameter) -> int: - return sum(int(dimension) for dimension in parameter.shape) - - -def _contiguous_strides(shape: tuple[int, ...]) -> tuple[int, ...]: - strides: list[int] = [] - for axis in range(len(shape)): - strides.append(prod(shape[axis + 1 :])) - return tuple(strides) - - -def _rank1_element_scales( - axis_scales: torch.Tensor, - shape: tuple[int, ...], - flat_start: int, - count: int, -) -> torch.Tensor: - """Expand paper rank-1 statistics for only one bounded flat slice.""" - - if count == 0: - return axis_scales.new_empty((0,)) - flat_indices = torch.arange(flat_start, flat_start + count, device=axis_scales.device) - result = torch.full((count,), torch.inf, device=axis_scales.device, dtype=torch.float32) - axis_offset = 0 - for dimension, stride in zip(shape, _contiguous_strides(shape), strict=True): - coordinates = torch.div(flat_indices, stride, rounding_mode="floor").remainder_(dimension) - torch.minimum( - result, - axis_scales.narrow(0, axis_offset, dimension)[coordinates], - out=result, - ) - axis_offset += dimension - return result - - -def _accumulate_rank1_maxima( - axis_maxima: torch.Tensor, - values: torch.Tensor, - shape: tuple[int, ...], - flat_start: int, -) -> None: - """Accumulate per-axis maxima for a bounded flat slice.""" - - flat_indices = torch.arange(flat_start, flat_start + values.numel(), device=values.device) - axis_offset = 0 - for dimension, stride in zip(shape, _contiguous_strides(shape), strict=True): - coordinates = torch.div(flat_indices, stride, rounding_mode="floor").remainder_(dimension) - axis_maxima.narrow(0, axis_offset, dimension).scatter_reduce_( - 0, coordinates, values, reduce="amax", include_self=True - ) - axis_offset += dimension - - -def _unpack_positive_4bit_elementwise( - packed: torch.Tensor, - numel: int, - scales: torch.Tensor, -) -> torch.Tensor: - codes = _unpack_nibbles(packed, numel).to(dtype=torch.float32) - return (codes + 1.0).mul_(scales / 16.0) - - -def _quantize_positive_4bit_elementwise(tensor: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: - safe_scales = scales.clamp_min(1.0e-30) - codes = torch.clamp(torch.round(tensor / safe_scales * 16.0 - 1.0), 0.0, 15.0).to(dtype=torch.uint8) - return _pack_nibbles(codes) +def _factored_state_numel_for_parameter(parameter: torch.nn.Parameter) -> int: + rows = int(parameter.shape[0]) + return rows + parameter.numel() // rows __all__ = ["AdamW4bit"] diff --git a/areno/engine/optim/adamw_fp32_master.py b/areno/engine/optim/adamw_fp32_master.py index b0c99049..891ae517 100644 --- a/areno/engine/optim/adamw_fp32_master.py +++ b/areno/engine/optim/adamw_fp32_master.py @@ -102,6 +102,9 @@ class AdamWFP32Master: BF16 model parameters after each bucket update. """ + gradient_shard_dtype = torch.float32 + stream_gradient_shards = False + def __init__( self, params: Iterable[torch.nn.Parameter], @@ -310,14 +313,14 @@ def state_dict(self) -> dict: @torch.no_grad() def reduce_scatter_gradients(self) -> None: - """Reduce one microbatch into persistent FP32 DP gradient shards.""" + """Reduce one microbatch into persistent optimizer-selected DP shards.""" for bucket in self.buckets: device = bucket.refs[0].model_param.device - # Preserve the historical FP32-main-grad reduction contract. The - # model gradient may be BF16, but casting only after the collective - # would permanently round the cross-rank sum in BF16. - dtype = torch.float32 + # Optimizers may explicitly accept a compact gradient shard. The + # default remains FP32 so existing optimizers preserve their + # accumulation and collective precision. + dtype = self.gradient_shard_dtype shard_size = self._max_shard_numel(bucket.numel) padded_numel = shard_size * self.dp_size send = self._arena(device, dtype, "grad_reduce_input", padded_numel) @@ -347,7 +350,7 @@ def reduce_scatter_gradients(self) -> None: # collective shard; keep only its valid prefix. reduced = output.narrow(0, 0, bucket.shard_numel) if bucket.grad_shard is None: - bucket.grad_shard = reduced.to(dtype=torch.float32) + bucket.grad_shard = reduced.to(dtype=dtype) else: bucket.grad_shard.add_(reduced) bucket.grad_param_ids = bucket.grad_param_ids.union(present) diff --git a/areno/engine/training.py b/areno/engine/training.py index 569fb7e9..18e275d0 100644 --- a/areno/engine/training.py +++ b/areno/engine/training.py @@ -19,6 +19,7 @@ from areno.engine.runtime.train_step import ( _clip_grad_norm, _grad_norms, + _grad_norms_from_shards, _merge_metrics, _pack_train_data, _train_meta, @@ -155,32 +156,48 @@ def _train_step( if not isinstance(loss, torch.Tensor): raise TypeError("train_loss_fn must return a torch.Tensor") (loss / max(grad_scale, 1)).backward() - # Keep the original full-gradient path for every optimizer residency - # mode, including disk. Disk offload still streams optimizer state, - # but it does not alter gradient accumulation or synchronization. - self._accumulate_main_gradients() + stream_gradient_shards = bool(getattr(worker.optimizer, "stream_gradient_shards", False)) + if stream_gradient_shards: + # AdamW4bit consumes each microbatch directly into compact BF16 DP + # shards. This avoids materializing the full-model FP32 main_grad + # copy that otherwise dominates optimizer-step peak memory. + self._sync_tensor_parallel_replicated_gradients() + worker.optimizer.reduce_scatter_gradients() + else: + # Preserve the established FP32 accumulation path for every other + # optimizer, including AdamW8bit. + self._accumulate_main_gradients() stepped = allow_step grad_norm = None multimodal_grad_metrics = None clipped_grad_norm = None optimizer_state_metrics = None if stepped: - self._sync_data_parallel_gradients() - self._sync_tensor_parallel_replicated_gradients() + if not stream_gradient_shards: + self._sync_data_parallel_gradients() + self._sync_tensor_parallel_replicated_gradients() self._finalize_router_expert_bias() multimodal_groups = tuple( group for group in worker.multimodal_lr_schedules if any(getattr(param, "_areno_lr_group", None) == group for param in worker.model.parameters()) ) - grad_norms = _grad_norms(worker.model.parameters(), multimodal_groups) + if stream_gradient_shards: + grad_norms = _grad_norms_from_shards(worker.optimizer.grad_shards(), multimodal_groups) + else: + grad_norms = _grad_norms(worker.model.parameters(), multimodal_groups) grad_norm = grad_norms.pop("global") multimodal_grad_metrics = ( {f"{group}_grad_norm": grad_norms[group] for group in multimodal_groups} if multimodal_groups else None ) clipped_grad_norm = grad_norm if worker.grad_clip_norm is not None: - _clip_grad_norm(worker.model.parameters(), grad_norm, worker.grad_clip_norm) + if stream_gradient_shards: + clip_coefficient = float(worker.grad_clip_norm) / (grad_norm + 1.0e-6) if grad_norm > 0.0 else 1.0 + if clip_coefficient < 1.0: + worker.optimizer.scale_gradients(clip_coefficient) + else: + _clip_grad_norm(worker.model.parameters(), grad_norm, worker.grad_clip_norm) clipped_grad_norm = min(grad_norm, float(worker.grad_clip_norm)) current_lr = self._lr_for_step(worker._global_step + 1) worker.optimizer.lr = current_lr diff --git a/docs/cli/training.rst b/docs/cli/training.rst index f8c284d8..326bb6ef 100644 --- a/docs/cli/training.rst +++ b/docs/cli/training.rst @@ -398,9 +398,10 @@ in its description; flags for other algorithms are ignored. backends; validate convergence when changing optimizer precision. ``--adam-4bit`` - Use packed 4-bit Adam moment states with rank-1 second-moment normalization - for tensors of rank two or greater and B128 fallback for vectors. This - option is CUDA-only and cannot be combined with ``--adam-8bit``. See + Use packed 4-bit first moments, factored row/column second moments for + tensors of rank two or greater, B128 fallback for vectors, and BF16 streamed + DP gradient shards. This option is CUDA-only and cannot be combined with + ``--adam-8bit``. See :doc:`../reference/adamw-4bit` for complete usage examples. ``--unfreeze-mm-tower`` diff --git a/docs/reference/adamw-4bit.rst b/docs/reference/adamw-4bit.rst index fe997208..dcbf6094 100644 --- a/docs/reference/adamw-4bit.rst +++ b/docs/reference/adamw-4bit.rst @@ -2,25 +2,30 @@ Using 4-bit AdamW ================= AReno provides an opt-in packed 4-bit AdamW optimizer for CUDA training. It -changes optimizer-state storage only; the model checkpoint and training data -format are unchanged. +changes optimizer-state storage and gradient accumulation precision; the model +checkpoint and training data format are unchanged. State representation -------------------- The first moment uses parameter-local blocks and the signed dynamic-exponent -4-bit map. For matrix and higher-rank parameters, the second moment uses -rank-1 normalization from *Memory Efficient Optimizers with 4-bit States*: -AReno records the maximum along every coordinate of every axis and scales each -element by the minimum applicable axis statistic. One-dimensional parameters -use 128-element block normalization. The second-moment 4-bit map excludes zero. - -For data-parallel training, axis statistics are defined over the original -parameter shape and combined with ``MAX`` across the DP group. Tensor-parallel -parameters use each rank's local model-tensor shape; this does not add a TP -collective. The fused update uses bounded block-local FP32 work and does not -materialize a parameter-sized FP32 moment tensor. The internal block size is -configurable for programmatic users and defaults to 128. +4-bit map. Matrix and higher-rank parameters use Adafactor-style row and column +means for the second moment, with every local tensor interpreted as +``[shape[0], -1]``. One-dimensional parameters retain packed 4-bit second +moments with 128-element block normalization. The second-moment 4-bit map used +for those vectors excludes zero. + +For data-parallel training, partial row and column sums are combined across the +DP group before their exponential update. Tensor-parallel parameters use each +rank's local model-tensor shape; this does not add a TP collective. The fused +update uses bounded block-local FP32 work and does not materialize a +parameter-sized FP32 moment tensor. + +Enabling ``--adam-4bit`` also streams every microbatch into BF16 DP gradient +shards instead of retaining a full-model FP32 ``main_grad`` copy. Gradient norm +and clipping operate directly on those shards. This behavior belongs to the +4-bit mode only; AdamW8bit and FP32 AdamW keep their existing FP32 gradient +accumulation path. The internal quantization block size defaults to 128. Command line ------------ @@ -146,8 +151,9 @@ Requirements and errors fused 4-bit optimizer kernel. * A saved 4-bit optimizer state must be resumed with the 4-bit optimizer. The separately saved model weights remain usable without ``--adam-4bit``. -* Optimizer-state checkpoints from the earlier block-only representation are - intentionally incompatible. Model-weight checkpoints remain portable. +* Optimizer-state checkpoints from the earlier block-only or rank-normalized + representations are intentionally incompatible. Model-weight checkpoints + remain portable. Initialized state is reported as ``adam4_quantized_state_bytes``, ``adam4_scale_metadata_bytes``, and ``adam4_total_bytes``. diff --git a/tests/test_adamw_4bit_cpu.py b/tests/test_adamw_4bit_cpu.py index 4cf1aced..dcff04d8 100644 --- a/tests/test_adamw_4bit_cpu.py +++ b/tests/test_adamw_4bit_cpu.py @@ -16,24 +16,21 @@ from areno.engine.modeling import build_optimizer from areno.engine.optim import AdamW4bit, AdamW8bit, AdamWFP32Master from areno.engine.optim.adamw_4bit import ( - _accumulate_rank1_maxima, + _factored_state_numel_for_parameter, _quantize_positive_4bit, - _quantize_positive_4bit_elementwise, _quantize_signed_4bit, - _rank1_element_scales, _unpack_positive_4bit, - _unpack_positive_4bit_elementwise, _unpack_signed_4bit, ) -def _optimizer(param: torch.nn.Parameter, *, block_size: int = 128) -> AdamW4bit: +def _optimizer(param: torch.nn.Parameter, *, block_size: int = 128, bucket_numel: int | None = None) -> AdamW4bit: return AdamW4bit( [param], lr=3.0e-4, betas=(0.9, 0.99), weight_decay=0.01, - bucket_numel=max(param.numel(), 1), + bucket_numel=max(param.numel(), 1) if bucket_numel is None else bucket_numel, quant_block_size=block_size, ) @@ -44,7 +41,12 @@ def _free_port() -> int: return int(sock.getsockname()[1]) -def _gloo_rank1_worker(rank: int, port: int, output_queue) -> None: +def _expected_first_step_factors(gradient: torch.Tensor, beta2: float = 0.99) -> torch.Tensor: + matrix = gradient.float().reshape(gradient.shape[0], -1).square().mul(1.0 - beta2) + return torch.cat((matrix.mean(dim=1), matrix.mean(dim=0))) + + +def _gloo_factored_worker(rank: int, port: int, output_queue) -> None: dist.init_process_group( backend="gloo", init_method=f"tcp://127.0.0.1:{port}", @@ -66,12 +68,14 @@ def _gloo_rank1_worker(rank: int, port: int, output_queue) -> None: ) parameter.grad = torch.arange(1, 36, dtype=torch.float32).reshape_as(parameter) + rank * 3.0 optimizer.reduce_scatter_gradients() + gradient_dtype = optimizer.buckets[0].grad_shard.dtype optimizer.step() output_queue.put( ( rank, parameter.detach().tolist(), - optimizer._rank1_scales[id(parameter)].tolist(), + optimizer._factored_second_moments[id(parameter)].tolist(), + gradient_dtype == torch.bfloat16, optimizer.buckets[0].refs[0].shard_start, optimizer.buckets[0].refs[0].shard_numel, ) @@ -80,39 +84,74 @@ def _gloo_rank1_worker(rank: int, port: int, output_queue) -> None: dist.destroy_process_group() -def test_adamw4bit_packs_two_moments_within_storage_budget() -> None: - param = torch.nn.Parameter(torch.zeros(8192)) - optimizer = _optimizer(param) - param.grad = torch.linspace(-1.0, 1.0, param.numel()) +def test_adamw4bit_vector_packs_two_moments_within_storage_budget() -> None: + parameter = torch.nn.Parameter(torch.zeros(8192)) + optimizer = _optimizer(parameter) + parameter.grad = torch.linspace(-1.0, 1.0, parameter.numel()) optimizer.step() - assert optimizer.persistent_moment_bytes() / param.numel() <= 1.25 state = optimizer._states[0] - assert state.exp_avg_q.numel() == param.numel() // 2 - assert state.exp_avg_sq_q.numel() == param.numel() // 2 - assert state.exp_avg_scale.numel() == param.numel() // 128 + assert optimizer.persistent_moment_bytes() / parameter.numel() <= 1.25 + assert state.exp_avg_q.numel() == parameter.numel() // 2 + assert state.exp_avg_sq_q.numel() == parameter.numel() // 2 + assert state.exp_avg_scale.numel() == parameter.numel() // 128 - eight_bit_param = torch.nn.Parameter(torch.zeros_like(param)) + eight_bit_parameter = torch.nn.Parameter(torch.zeros_like(parameter)) eight_bit = AdamW8bit( - [eight_bit_param], + [eight_bit_parameter], lr=3.0e-4, betas=(0.9, 0.99), weight_decay=0.01, - bucket_numel=param.numel(), + bucket_numel=parameter.numel(), quant_block_size=128, ) - eight_bit_param.grad = torch.ones_like(eight_bit_param) + eight_bit_parameter.grad = torch.ones_like(eight_bit_parameter) eight_bit.step() - eight_bit_state = eight_bit.state_dict()["state"][0] - eight_bit_bytes = sum( - eight_bit_state[name].numel() * eight_bit_state[name].element_size() - for name in ("exp_avg_q", "exp_avg_scale", "exp_avg_sq_q", "exp_avg_sq_scale") + assert optimizer.persistent_moment_bytes() <= eight_bit.state_memory_metrics()["total_bytes"] * 0.6 + + +def test_adamw4bit_matrix_stores_only_packed_first_moment_and_factors() -> None: + parameter = torch.nn.Parameter(torch.zeros(1024, 1024)) + optimizer = _optimizer(parameter) + parameter.grad = torch.ones_like(parameter) + + optimizer.step() + + state = optimizer._states[0] + factors = optimizer._factored_second_moments[id(parameter)] + assert factors is not None and factors.numel() == 2048 + assert state.exp_avg_q.numel() == parameter.numel() // 2 + assert state.exp_avg_sq_q.numel() == 0 + assert state.exp_avg_sq_scale.numel() == 0 + assert optimizer.persistent_moment_bytes() / parameter.numel() < 0.55 + + +def test_adamw4bit_uses_bf16_gradient_shards_without_changing_adam8_default() -> None: + four_bit_parameter = torch.nn.Parameter(torch.ones(257, dtype=torch.bfloat16)) + four_bit = _optimizer(four_bit_parameter) + four_bit_parameter.grad = torch.linspace(-1.0, 1.0, 257).to(torch.bfloat16) + four_bit.reduce_scatter_gradients() + + eight_bit_parameter = torch.nn.Parameter(torch.ones(257, dtype=torch.bfloat16)) + eight_bit = AdamW8bit( + [eight_bit_parameter], + lr=3.0e-4, + betas=(0.9, 0.99), + weight_decay=0.01, + bucket_numel=257, ) - assert optimizer.persistent_moment_bytes() <= eight_bit_bytes * 0.6 + eight_bit_parameter.grad = torch.linspace(-1.0, 1.0, 257).to(torch.bfloat16) + eight_bit.reduce_scatter_gradients() + assert four_bit.stream_gradient_shards is True + assert four_bit_parameter.grad is None + assert all(bucket.grad_shard.dtype == torch.bfloat16 for bucket in four_bit.buckets) + assert eight_bit.stream_gradient_shards is False + assert all(bucket.grad_shard.dtype == torch.float32 for bucket in eight_bit.buckets) -def test_adamw4bit_fresh_rank1_state_is_initialized_during_streaming_update() -> None: + +def test_adamw4bit_initializes_and_releases_state_per_parameter() -> None: parameters = [torch.nn.Parameter(torch.zeros(32, 32)) for _ in range(2)] optimizer = AdamW4bit( parameters, @@ -135,10 +174,8 @@ def tracked_ensure_bucket_state(bucket, state) -> None: optimizer._ensure_bucket_state = tracked_ensure_bucket_state optimizer.step() - # The statistics pass consumes the implicit zero second moment without - # allocating packed state. State is initialized only in the update pass, - # where each completed parameter releases its gradient before the next. assert live_gradients_at_initialization == [2, 1] + assert all(parameter.grad is None for parameter in parameters) def test_adamw4bit_second_moment_mapping_excludes_zero() -> None: @@ -162,58 +199,120 @@ def test_adamw4bit_signed_quantizer_preserves_dynamic_map_points() -> None: torch.testing.assert_close(restored, values, rtol=1.0e-6, atol=1.0e-6) -def test_adamw4bit_checkpoint_round_trip_preserves_next_update() -> None: +def test_adamw4bit_vector_checkpoint_round_trip_preserves_next_update() -> None: initial = torch.linspace(-0.5, 0.5, 257).to(torch.bfloat16) - first_param = torch.nn.Parameter(initial.clone()) - first = _optimizer(first_param) - first_param.grad = torch.linspace(-0.3, 0.7, first_param.numel()).to(torch.bfloat16) + first_parameter = torch.nn.Parameter(initial.clone()) + first = _optimizer(first_parameter) + first_parameter.grad = torch.linspace(-0.3, 0.7, first_parameter.numel()).to(torch.bfloat16) first.step() checkpoint = copy.deepcopy(first.state_dict()) - restored_param = torch.nn.Parameter(first_param.detach().clone()) - restored = _optimizer(restored_param) + restored_parameter = torch.nn.Parameter(first_parameter.detach().clone()) + restored = _optimizer(restored_parameter) restored.load_state_dict(checkpoint) - next_gradient = torch.linspace(0.8, -0.4, first_param.numel()).to(torch.bfloat16) - first_param.grad = next_gradient.clone() - restored_param.grad = next_gradient.clone() + next_gradient = torch.linspace(0.8, -0.4, first_parameter.numel()).to(torch.bfloat16) + first_parameter.grad = next_gradient.clone() + restored_parameter.grad = next_gradient.clone() first.step() restored.step() - torch.testing.assert_close(restored_param, first_param, rtol=0.0, atol=0.0) - assert restored.state_dict()["state_format_version"] == 2 + torch.testing.assert_close(restored_parameter, first_parameter, rtol=0.0, atol=0.0) + assert restored.state_dict()["state_format_version"] == 3 -def test_adamw4bit_disk_offload_preserves_update(tmp_path: Path) -> None: +def test_adamw4bit_factored_checkpoint_round_trip_preserves_next_update() -> None: + initial = torch.linspace(-0.5, 0.5, 35).reshape(5, 7).to(torch.bfloat16) + first_parameter = torch.nn.Parameter(initial.clone()) + first = _optimizer(first_parameter) + first_parameter.grad = torch.linspace(-0.3, 0.7, 35).reshape_as(initial).to(torch.bfloat16) + first.step() + checkpoint = copy.deepcopy(first.state_dict()) + + restored_parameter = torch.nn.Parameter(first_parameter.detach().clone()) + restored = _optimizer(restored_parameter) + restored.load_state_dict(checkpoint) + next_gradient = torch.linspace(0.8, -0.4, 35).reshape_as(initial).to(torch.bfloat16) + first_parameter.grad = next_gradient.clone() + restored_parameter.grad = next_gradient.clone() + first.step() + restored.step() + + torch.testing.assert_close(restored_parameter, first_parameter, rtol=0.0, atol=0.0) + torch.testing.assert_close( + restored._factored_second_moments[id(restored_parameter)], + first._factored_second_moments[id(first_parameter)], + ) + + +def test_adamw4bit_rejects_legacy_optimizer_state_format() -> None: + parameter = torch.nn.Parameter(torch.zeros(8)) + optimizer = _optimizer(parameter) + parameter.grad = torch.ones_like(parameter) + optimizer.step() + checkpoint = optimizer.state_dict() + checkpoint["state_format_version"] = 2 + + with pytest.raises(ValueError, match="unsupported AdamW4bit state format"): + _optimizer(torch.nn.Parameter(torch.zeros(8))).load_state_dict(checkpoint) + + +def test_adamw4bit_disk_offload_preserves_vector_update(tmp_path: Path) -> None: initial = torch.linspace(-0.5, 0.5, 257).to(torch.bfloat16) - candidate_param = torch.nn.Parameter(initial.clone()) - reference_param = torch.nn.Parameter(initial.clone()) - candidate = _optimizer(candidate_param) - reference = _optimizer(reference_param) + candidate_parameter = torch.nn.Parameter(initial.clone()) + reference_parameter = torch.nn.Parameter(initial.clone()) + candidate = _optimizer(candidate_parameter) + reference = _optimizer(reference_parameter) candidate.configure_state_offload(mode="disk", directory=str(tmp_path), batch_size=2) for gradient in ( torch.linspace(-0.4, 0.7, initial.numel()), torch.linspace(0.8, -0.2, initial.numel()), ): - candidate_param.grad = gradient.to(torch.bfloat16) - reference_param.grad = gradient.to(torch.bfloat16) + candidate_parameter.grad = gradient.to(torch.bfloat16) + reference_parameter.grad = gradient.to(torch.bfloat16) candidate.step() reference.step() - torch.testing.assert_close(candidate_param, reference_param, rtol=0.0, atol=0.0) + torch.testing.assert_close(candidate_parameter, reference_parameter, rtol=0.0, atol=0.0) assert all(state.offload_file is not None for state in candidate._states) candidate.onload_state(torch.device("cpu")) assert all(state.exp_avg_q is not None for state in candidate._states) assert not list(tmp_path.rglob("*.mmap")) -def test_adamw4bit_tracks_fp32_adamw_on_smooth_gradients() -> None: +def test_adamw4bit_disk_offload_preserves_factored_update(tmp_path: Path) -> None: + initial = torch.linspace(-0.5, 0.5, 35).reshape(5, 7).to(torch.bfloat16) + candidate_parameter = torch.nn.Parameter(initial.clone()) + reference_parameter = torch.nn.Parameter(initial.clone()) + candidate = _optimizer(candidate_parameter) + reference = _optimizer(reference_parameter) + candidate.configure_state_offload(mode="disk", directory=str(tmp_path), batch_size=1) + + for gradient in ( + torch.linspace(-0.4, 0.7, 35).reshape_as(initial), + torch.linspace(0.8, -0.2, 35).reshape_as(initial), + ): + candidate_parameter.grad = gradient.to(torch.bfloat16) + reference_parameter.grad = gradient.to(torch.bfloat16) + candidate.step() + reference.step() + + candidate.onload_state(torch.device("cpu")) + torch.testing.assert_close(candidate_parameter, reference_parameter, rtol=0.0, atol=0.0) + torch.testing.assert_close( + candidate._factored_second_moments[id(candidate_parameter)], + reference._factored_second_moments[id(reference_parameter)], + ) + assert not list(tmp_path.rglob("*.mmap")) + + +def test_adamw4bit_vector_tracks_fp32_adamw_on_smooth_gradients() -> None: initial = torch.linspace(-1.0, 1.0, 1024) - quantized_param = torch.nn.Parameter(initial.clone()) - reference_param = torch.nn.Parameter(initial.clone()) - quantized = _optimizer(quantized_param) + quantized_parameter = torch.nn.Parameter(initial.clone()) + reference_parameter = torch.nn.Parameter(initial.clone()) + quantized = _optimizer(quantized_parameter) reference = AdamWFP32Master( - [reference_param], + [reference_parameter], lr=3.0e-4, betas=(0.9, 0.99), weight_decay=0.01, @@ -222,15 +321,15 @@ def test_adamw4bit_tracks_fp32_adamw_on_smooth_gradients() -> None: for step in range(20): gradient = torch.sin(torch.linspace(-2.0, 2.0, initial.numel()) + step * 0.1) - quantized_param.grad = gradient.clone() - reference_param.grad = gradient.clone() + quantized_parameter.grad = gradient.clone() + reference_parameter.grad = gradient.clone() quantized.step() reference.step() - torch.testing.assert_close(quantized_param, reference_param, rtol=3.0e-3, atol=3.0e-3) + torch.testing.assert_close(quantized_parameter, reference_parameter, rtol=3.0e-3, atol=3.0e-3) -def test_adamw4bit_nonfinite_gradient_skips_only_affected_block() -> None: +def test_adamw4bit_nonfinite_vector_gradient_skips_only_affected_block() -> None: parameter = torch.nn.Parameter(torch.zeros(256)) optimizer = _optimizer(parameter, block_size=128) gradient = torch.ones_like(parameter) @@ -243,189 +342,102 @@ def test_adamw4bit_nonfinite_gradient_skips_only_affected_block() -> None: assert torch.all(parameter[128:] < 0) -@pytest.mark.parametrize( - "values", - [ - torch.tensor([[1.0, 1.0, 1.0], [9.0, 9.0, 9.0]]), - torch.tensor([[1.0, 2.0, 8.0], [1.0, 2.0, 8.0]]), - torch.tensor([[1.0, 8.0, 1.0], [8.0, 1.0, 8.0], [1.0, 8.0, 1.0]]), - torch.tensor([[1.0, 1.0, 1.0], [1.0, 1000.0, 1.0]]), - torch.zeros(3, 5), - torch.full((3, 5), 1.0e-20), - ], - ids=["row", "column", "checkerboard", "isolated-outlier", "all-zero", "tiny-positive"], -) -def test_rank1_statistics_match_paper_sm3_algorithm(values: torch.Tensor) -> None: - flattened = values.flatten() - statistics = torch.zeros(sum(values.shape)) +def test_adamw4bit_nonfinite_matrix_gradient_skips_whole_parameter() -> None: + parameter = torch.nn.Parameter(torch.zeros(16, 16)) + optimizer = _optimizer(parameter) + gradient = torch.ones_like(parameter) + gradient[4, 4] = torch.inf + parameter.grad = gradient - for start in range(0, flattened.numel(), 4): - _accumulate_rank1_maxima(statistics, flattened[start : start + 4], tuple(values.shape), start) + optimizer.step() - expected = torch.cat((values.amax(dim=1), values.amax(dim=0))) - torch.testing.assert_close(statistics, expected) - expanded = _rank1_element_scales(statistics, tuple(values.shape), 0, values.numel()).reshape(values.shape) + torch.testing.assert_close(parameter, torch.zeros_like(parameter)) torch.testing.assert_close( - expanded, - torch.minimum(expected[: values.shape[0], None], expected[values.shape[0] :]), + optimizer._factored_second_moments[id(parameter)], + torch.zeros(parameter.shape[0] + parameter.shape[1]), ) -def test_rank1_statistics_generalize_to_higher_rank() -> None: - values = torch.arange(1, 31, dtype=torch.float32).reshape(2, 3, 5) - statistics = torch.zeros(sum(values.shape)) - _accumulate_rank1_maxima(statistics, values.flatten(), tuple(values.shape), 0) - - expected = torch.cat( - ( - values.amax(dim=(1, 2)), - values.amax(dim=(0, 2)), - values.amax(dim=(0, 1)), - ) - ) - torch.testing.assert_close(statistics, expected) - - -def test_adamw4bit_matrix_uses_rank1_second_moment_scales() -> None: - parameter = torch.nn.Parameter(torch.zeros(3, 5)) - optimizer = _optimizer(parameter) - gradient = torch.arange(1, 16, dtype=torch.float32).reshape_as(parameter) +@pytest.mark.parametrize("shape", [(2, 3), (2, 3, 5), (3, 5, 7)]) +def test_factored_statistics_match_row_column_means(shape: tuple[int, ...]) -> None: + parameter = torch.nn.Parameter(torch.zeros(shape)) + optimizer = _optimizer(parameter, bucket_numel=7) + gradient = torch.arange(1, parameter.numel() + 1, dtype=torch.float32).reshape(shape) parameter.grad = gradient.clone() optimizer.step() - state = optimizer._states[0] - variance = (1.0 - optimizer.betas[1]) * gradient.square() - expected = torch.cat((variance.amax(dim=1), variance.amax(dim=0))) - rank1_scales = optimizer._rank1_scales[id(parameter)] - assert rank1_scales is not None - torch.testing.assert_close(rank1_scales, expected) - assert state.exp_avg_scale.numel() == 1 - assert state.exp_avg_sq_scale.numel() == 0 - + torch.testing.assert_close( + optimizer._factored_second_moments[id(parameter)], + _expected_first_step_factors(gradient), + ) -def test_adamw4bit_rank1_second_moment_never_decodes_nonzero_scale_to_zero() -> None: - values = torch.tensor([[0.0, 0.01, 0.5], [0.02, 0.25, 1.0]]) - statistics = torch.cat((values.amax(dim=1), values.amax(dim=0))) - scales = _rank1_element_scales(statistics, tuple(values.shape), 0, values.numel()) - packed = _quantize_positive_4bit_elementwise(values.flatten(), scales) - restored = _unpack_positive_4bit_elementwise(packed, values.numel(), scales) +def test_factored_statistics_combine_partial_parameter_chunks() -> None: + parameter = torch.nn.Parameter(torch.zeros(2048, 2049)) + optimizer = _optimizer(parameter, bucket_numel=1) - assert torch.all(restored[scales > 0] > 0) + refs = [ref for bucket in optimizer.buckets for ref in bucket.refs] + assert len(optimizer.buckets) == 2 + assert [ref.param_start for ref in refs] == [0, 4 * 1024 * 1024] + assert sum(ref.numel for ref in refs) == parameter.numel() + factors = optimizer._ensure_factored_second_moment(parameter) + assert factors.numel() == 2048 + 2049 + assert _factored_state_numel_for_parameter(parameter) == factors.numel() -def test_adamw4bit_preserves_chunking_and_shares_rank1_scales_across_chunks() -> None: - parameter = torch.nn.Parameter(torch.zeros(2048, 2049)) +def test_adamw4bit_mixed_matrix_vector_bucket_uses_disjoint_state_layouts() -> None: + matrix = torch.nn.Parameter(torch.zeros(8, 8)) + vector = torch.nn.Parameter(torch.zeros(17)) optimizer = AdamW4bit( - [parameter], + [matrix, vector], lr=3.0e-4, betas=(0.9, 0.99), - weight_decay=0.01, - bucket_numel=1, + weight_decay=0.0, + bucket_numel=1024, quant_block_size=128, ) + matrix.grad = torch.ones_like(matrix) + vector.grad = torch.ones_like(vector) - refs = [ref for bucket in optimizer.buckets for ref in bucket.refs] - assert len(optimizer.buckets) == 2 - assert len(refs) == 2 - assert [ref.param_start for ref in refs] == [0, 4 * 1024 * 1024] - assert sum(ref.numel for ref in refs) == parameter.numel() - assert all(ref.model_param is parameter for ref in refs) - scales = optimizer._ensure_rank1_scales(parameter) - assert scales.numel() == sum(parameter.shape) - assert optimizer._rank1_scales[id(parameter)] is scales + optimizer.step() + + state = optimizer._states[0] + assert state.exp_avg_q.numel() == 32 + 9 + assert state.exp_avg_sq_q.numel() == 9 + assert state.exp_avg_sq_scale.numel() == 1 -def test_adamw4bit_rank1_metadata_stays_within_large_matrix_budget() -> None: +def test_adamw4bit_factored_memory_metrics_match_resident_tensors() -> None: parameter = torch.nn.Parameter(torch.zeros(1024, 1024)) optimizer = _optimizer(parameter) - optimizer._ensure_bucket_state(optimizer.buckets[0], optimizer._states[0]) - optimizer._ensure_rank1_scales(parameter) - optimizer._states[0].step = 1 + parameter.grad = torch.ones_like(parameter) + optimizer.step() - assert optimizer.persistent_moment_bytes() / parameter.numel() <= 1.25 metrics = optimizer.state_memory_metrics() assert metrics["total_bytes"] == optimizer.persistent_moment_bytes() - assert metrics["scale_metadata_bytes"] == (parameter.numel() // 128 + sum(parameter.shape)) * 4 + assert metrics["quantized_state_bytes"] == parameter.numel() // 2 + assert metrics["scale_metadata_bytes"] == (parameter.numel() // 128 + 2048) * 4 -def test_adamw4bit_rank1_checkpoint_round_trip_preserves_next_update() -> None: - initial = torch.linspace(-0.5, 0.5, 35).reshape(5, 7).to(torch.bfloat16) - first_parameter = torch.nn.Parameter(initial.clone()) - first = _optimizer(first_parameter) - first_parameter.grad = torch.linspace(-0.3, 0.7, 35).reshape_as(initial).to(torch.bfloat16) - first.step() - checkpoint = copy.deepcopy(first.state_dict()) - - restored_parameter = torch.nn.Parameter(first_parameter.detach().clone()) - restored = _optimizer(restored_parameter) - restored.load_state_dict(checkpoint) - next_gradient = torch.linspace(0.8, -0.4, 35).reshape_as(initial).to(torch.bfloat16) - first_parameter.grad = next_gradient.clone() - restored_parameter.grad = next_gradient.clone() - first.step() - restored.step() - - torch.testing.assert_close(restored_parameter, first_parameter, rtol=0.0, atol=0.0) - torch.testing.assert_close(restored._rank1_scales[id(restored_parameter)], first._rank1_scales[id(first_parameter)]) - - -def test_adamw4bit_clear_state_drops_rank1_scales() -> None: +def test_adamw4bit_clear_state_drops_factored_state() -> None: parameter = torch.nn.Parameter(torch.zeros(5, 7)) optimizer = _optimizer(parameter) parameter.grad = torch.ones_like(parameter) optimizer.step() - assert optimizer._rank1_scales[id(parameter)] is not None + assert optimizer._factored_second_moments[id(parameter)] is not None optimizer.clear_state() - assert optimizer._rank1_scales[id(parameter)] is None + assert optimizer._factored_second_moments[id(parameter)] is None assert all(state.step == 0 and state.exp_avg_q is None for state in optimizer._states) -def test_adamw4bit_rank1_disk_offload_preserves_axis_scales(tmp_path: Path) -> None: - initial = torch.linspace(-0.5, 0.5, 35).reshape(5, 7).to(torch.bfloat16) - candidate_parameter = torch.nn.Parameter(initial.clone()) - reference_parameter = torch.nn.Parameter(initial.clone()) - candidate = _optimizer(candidate_parameter) - reference = _optimizer(reference_parameter) - candidate.configure_state_offload(mode="disk", directory=str(tmp_path), batch_size=1) - - for gradient in ( - torch.linspace(-0.4, 0.7, 35).reshape_as(initial), - torch.linspace(0.8, -0.2, 35).reshape_as(initial), - ): - candidate_parameter.grad = gradient.to(torch.bfloat16) - reference_parameter.grad = gradient.to(torch.bfloat16) - candidate.step() - reference.step() - - candidate.onload_state(torch.device("cpu")) - torch.testing.assert_close(candidate_parameter, reference_parameter, rtol=0.0, atol=0.0) - torch.testing.assert_close( - candidate._rank1_scales[id(candidate_parameter)], reference._rank1_scales[id(reference_parameter)] - ) - assert not list(tmp_path.rglob("*.mmap")) - - -def test_rank1_partial_shards_combine_to_unsharded_statistics() -> None: - values = torch.arange(1, 36, dtype=torch.float32).reshape(5, 7) - combined = torch.zeros(sum(values.shape)) - for start, count in ((0, 13), (13, 12), (25, 10)): - partial = torch.zeros_like(combined) - _accumulate_rank1_maxima(partial, values.flatten().narrow(0, start, count), tuple(values.shape), start) - torch.maximum(combined, partial, out=combined) - - expected = torch.cat((values.amax(dim=1), values.amax(dim=0))) - torch.testing.assert_close(combined, expected) - - -def test_real_gloo_rank1_statistics_match_unsharded_reference_across_split_row() -> None: +def test_real_gloo_factored_statistics_match_unsharded_reference() -> None: spawn = mp.get_context("spawn") output_queue = spawn.Queue() port = _free_port() - processes = [spawn.Process(target=_gloo_rank1_worker, args=(rank, port, output_queue)) for rank in range(2)] + processes = [spawn.Process(target=_gloo_factored_worker, args=(rank, port, output_queue)) for rank in range(2)] for process in processes: process.start() try: @@ -438,14 +450,14 @@ def test_real_gloo_rank1_statistics_match_unsharded_reference_across_split_row() process.join(timeout=5) assert all(process.exitcode == 0 for process in processes) - rank0_model, rank0_scales, rank0_start, rank0_count = results[0] - rank1_model, rank1_scales, rank1_start, rank1_count = results[1] + rank0_model, rank0_factors, rank0_bf16, rank0_start, rank0_count = results[0] + rank1_model, rank1_factors, rank1_bf16, rank1_start, rank1_count = results[1] averaged_gradient = torch.arange(1, 36, dtype=torch.float32).reshape(5, 7) + 1.5 - variance = 0.01 * averaged_gradient.square() - expected_scales = torch.cat((variance.amax(dim=1), variance.amax(dim=0))) + expected_factors = _expected_first_step_factors(averaged_gradient) assert rank0_model == rank1_model - torch.testing.assert_close(torch.tensor(rank0_scales), expected_scales) - torch.testing.assert_close(torch.tensor(rank1_scales), expected_scales) + assert rank0_bf16 and rank1_bf16 + torch.testing.assert_close(torch.tensor(rank0_factors), expected_factors) + torch.testing.assert_close(torch.tensor(rank1_factors), expected_factors) assert (rank0_start, rank0_count) == (0, 18) assert (rank1_start, rank1_count) == (18, 17) @@ -461,8 +473,8 @@ class Context: dp_size = 1 dp_group = None - param = torch.nn.Parameter(torch.ones(4)) - optimizer = build_optimizer([param], OptimizerConfig(adam_4bit=True), Context()) + parameter = torch.nn.Parameter(torch.ones(4)) + optimizer = build_optimizer([parameter], OptimizerConfig(adam_4bit=True), Context()) assert isinstance(optimizer, AdamW4bit)