From afcc9aeb64da1dcb4c952345ec43959fb7d38a65 Mon Sep 17 00:00:00 2001 From: xsuler Date: Tue, 1 Sep 2026 11:31:24 +0800 Subject: [PATCH 1/2] feat(optimizer): add packed 4-bit AdamW states --- areno/accel/__init__.py | 4 +- areno/accel/csrc/extension.cpp | 34 ++ areno/accel/csrc/optimizer.cu | 417 +++++++++++++++++++++++ areno/accel/optimizer.py | 137 +++++++- areno/api/backend/cuda/roles.py | 4 +- areno/api/trainer_config.py | 6 + areno/cli/train.py | 12 +- areno/dashboard/server.py | 2 + areno/engine/config.py | 5 + areno/engine/modeling.py | 9 +- areno/engine/optim/__init__.py | 3 +- areno/engine/optim/adamw_4bit.py | 372 ++++++++++++++++++++ areno/engine/optim/adamw_8bit.py | 155 +++++++-- tests/test_adamw_4bit_cpu.py | 202 +++++++++++ tests/test_adamw_8bit_blockwise_cpu.py | 55 +++ tests/test_fp32_master_optimizer_cuda.py | 27 +- 16 files changed, 1399 insertions(+), 45 deletions(-) create mode 100644 areno/engine/optim/adamw_4bit.py create mode 100644 tests/test_adamw_4bit_cpu.py create mode 100644 tests/test_adamw_8bit_blockwise_cpu.py diff --git a/areno/accel/__init__.py b/areno/accel/__init__.py index 8f61fe21..bcd806c0 100644 --- a/areno/accel/__init__.py +++ b/areno/accel/__init__.py @@ -30,7 +30,7 @@ from areno.accel.linear import areno_grouped_linear, areno_linear from areno.accel.moe import areno_moe_permute, areno_moe_topk_permute, areno_moe_unpermute from areno.accel.normalization import areno_optional_scale_rmsnorm, areno_rmsnorm, areno_rmsnorm_silu_gate -from areno.accel.optimizer import areno_adamw_fp32_master_step +from areno.accel.optimizer import areno_adamw_4bit_step, areno_adamw_8bit_step, areno_adamw_fp32_master_step from areno.accel.router import areno_grouped_topk_router from areno.accel.routing import areno_moe_align from areno.accel.topk import areno_topk_softmax @@ -50,6 +50,8 @@ "areno_moe_permute", "areno_moe_topk_permute", "areno_moe_unpermute", + "areno_adamw_4bit_step", + "areno_adamw_8bit_step", "areno_adamw_fp32_master_step", "areno_optional_scale_rmsnorm", "areno_rmsnorm", diff --git a/areno/accel/csrc/extension.cpp b/areno/accel/csrc/extension.cpp index 6cee02e9..494233b8 100644 --- a/areno/accel/csrc/extension.cpp +++ b/areno/accel/csrc/extension.cpp @@ -185,9 +185,43 @@ void areno_adamw_fp32_master_step_cuda( double eps, double step_size, double bias_correction2_sqrt); +void areno_adamw_4bit_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 exp_avg_sq_scale, + int64_t packed_offset, + int64_t 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_8bit_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 exp_avg_sq_scale, + int64_t quant_block_size, + double beta1, + double beta2, + double effective_lr, + double weight_decay, + double eps, + double step_size, + double bias_correction2_sqrt); 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_8bit_step", &areno_adamw_8bit_step_cuda, "ARENO block-wise 8-bit AdamW step"); m.def("areno_silu_and_mul", &areno_silu_and_mul_cuda, "ARENO SiLU and multiply"); m.def("areno_gelu_tanh_and_mul", &areno_gelu_tanh_and_mul_cuda, "ARENO tanh GELU and multiply"); m.def("areno_silu", &areno_silu_cuda, "ARENO SiLU"); diff --git a/areno/accel/csrc/optimizer.cu b/areno/accel/csrc/optimizer.cu index 989da837..4c885b56 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -8,6 +8,29 @@ namespace { +__device__ __constant__ float kSigned4bitDynamicMap[16] = { + -0.8875f, -0.6625f, -0.4375f, -0.2125f, -0.0775f, -0.0325f, -0.0055f, 0.0f, + 0.0055f, 0.0325f, 0.0775f, 0.2125f, 0.4375f, 0.6625f, 0.8875f, 1.0f}; + +__device__ __forceinline__ uint8_t load_nibble(const uint8_t* packed, int64_t index) { + const uint8_t byte = packed[index >> 1]; + return (index & 1) == 0 ? byte & 0x0Fu : byte >> 4; +} + +__device__ __forceinline__ uint8_t nearest_signed_dynamic_code(float normalized) { + uint8_t best = 0; + float best_distance = fabsf(normalized - kSigned4bitDynamicMap[0]); +#pragma unroll + for (uint8_t code = 1; code < 16; ++code) { + const float distance = fabsf(normalized - kSigned4bitDynamicMap[code]); + if (distance < best_distance) { + best = code; + best_distance = distance; + } + } + return best; +} + __device__ __forceinline__ float adamw_update( float master, float grad, @@ -40,6 +63,125 @@ __device__ __forceinline__ float load_grad(const at::BFloat16* gra return __bfloat162float(raw[index]); } +template +__device__ __forceinline__ float load_model(const model_t* model, int64_t index) { + return static_cast(model[index]); +} + +template <> +__device__ __forceinline__ float load_model(const at::BFloat16* model, int64_t index) { + return __bfloat162float(reinterpret_cast(model)[index]); +} + +template +__device__ __forceinline__ void store_model(model_t* model, int64_t index, float value) { + model[index] = static_cast(value); +} + +template <> +__device__ __forceinline__ void store_model(at::BFloat16* model, int64_t index, float value) { + reinterpret_cast<__nv_bfloat16*>(model)[index] = __float2bfloat16_rn(value); +} + +template +__global__ void adamw_4bit_kernel( + model_t* model, + const grad_t* grad, + uint8_t* exp_avg_q, + float* exp_avg_scale, + uint8_t* exp_avg_sq_q, + float* exp_avg_sq_scale, + int64_t numel, + int64_t packed_offset, + int64_t scale_offset, + 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__ float variance_max[1024]; + __shared__ uint8_t moment_codes[1024]; + __shared__ uint8_t variance_codes[1024]; + __shared__ float new_moment_scale; + __shared__ float new_variance_scale; + __shared__ int invalid_block; + + const int tid = threadIdx.x; + const int64_t block_start = static_cast(blockIdx.x) * blockDim.x; + const int64_t local_index = block_start + 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 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 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; + 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(); + } + if (invalid_block != 0) { + return; + } + 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; + } + __syncthreads(); + + if (active) { + 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); + 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, @@ -207,6 +349,182 @@ void launch_adamw( C10_CUDA_KERNEL_LAUNCH_CHECK(); } +template +void launch_adamw_4bit( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg_q, + torch::Tensor exp_avg_scale, + torch::Tensor exp_avg_sq_q, + torch::Tensor exp_avg_sq_scale, + int64_t packed_offset, + int64_t scale_offset, + 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_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(), + exp_avg_sq_scale.data_ptr(), + model.numel(), + packed_offset, + scale_offset, + beta1, + beta2, + effective_lr, + weight_decay, + eps, + step_size, + bias_correction2_sqrt); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +__global__ void adamw_8bit_blockwise_kernel( + model_t* model, + const grad_t* grad, + uint8_t* exp_avg_q, + float* exp_avg_scale, + uint8_t* exp_avg_sq_q, + float* exp_avg_sq_scale, + int64_t numel, + int64_t quant_block_size, + float beta1, + float beta2, + float effective_lr, + float weight_decay, + float eps, + float step_size, + float bias_correction2_sqrt) { + constexpr int warp_size = 32; + constexpr int max_warps = 8; + __shared__ float warp_moment_maxima[max_warps]; + __shared__ float warp_variance_maxima[max_warps]; + + const int64_t block_start = static_cast(blockIdx.x) * quant_block_size; + const int64_t remaining = numel - block_start; + const int64_t block_numel = quant_block_size < remaining ? quant_block_size : remaining; + const float old_moment_scale = exp_avg_scale[blockIdx.x]; + const float old_variance_scale = exp_avg_sq_scale[blockIdx.x]; + float local_moment_max = 0.0f; + float local_variance_max = 0.0f; + + for (int64_t offset = threadIdx.x; offset < block_numel; offset += blockDim.x) { + const int64_t index = block_start + offset; + const float gradient = load_grad(grad, index); + float moment = (static_cast(exp_avg_q[index]) - 128) * old_moment_scale; + float variance = static_cast(exp_avg_sq_q[index]) * old_variance_scale; + float weight = load_model(model, index); + weight = adamw_update( + weight, + gradient, + moment, + variance, + beta1, + beta2, + effective_lr, + weight_decay, + eps, + step_size, + bias_correction2_sqrt); + store_model(model, index, weight); + local_moment_max = fmaxf(local_moment_max, fabsf(moment)); + local_variance_max = fmaxf(local_variance_max, variance); + } + + 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 = threadIdx.x & (warp_size - 1); + const int warp = threadIdx.x / warp_size; + if (lane == 0) { + warp_moment_maxima[warp] = local_moment_max; + warp_variance_maxima[warp] = local_variance_max; + } + __syncthreads(); + + 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) { + exp_avg_scale[blockIdx.x] = fmaxf(block_moment_max / 127.0f, 1.0e-30f); + exp_avg_sq_scale[blockIdx.x] = fmaxf(block_variance_max / 255.0f, 1.0e-30f); + } + } + __syncthreads(); + const float new_moment_scale = exp_avg_scale[blockIdx.x]; + const float new_variance_scale = exp_avg_sq_scale[blockIdx.x]; + for (int64_t offset = threadIdx.x; offset < block_numel; offset += blockDim.x) { + const int64_t index = block_start + offset; + const float gradient = load_grad(grad, index); + const float moment = + beta1 * (static_cast(exp_avg_q[index]) - 128) * old_moment_scale + (1.0f - beta1) * gradient; + const float variance = beta2 * static_cast(exp_avg_sq_q[index]) * old_variance_scale + + (1.0f - beta2) * gradient * gradient; + const float moment_q = nearbyintf(moment / new_moment_scale) + 128.0f; + const float variance_q = nearbyintf(variance / new_variance_scale); + exp_avg_q[index] = static_cast(fminf(fmaxf(moment_q, 0.0f), 255.0f)); + exp_avg_sq_q[index] = static_cast(fminf(fmaxf(variance_q, 0.0f), 255.0f)); + } +} + +template +void launch_adamw_8bit( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg_q, + torch::Tensor exp_avg_scale, + torch::Tensor exp_avg_sq_q, + torch::Tensor exp_avg_sq_scale, + int64_t quant_block_size, + float beta1, + float beta2, + float effective_lr, + float weight_decay, + float eps, + float step_size, + float bias_correction2_sqrt) { + constexpr int threads = 256; + const int blocks = static_cast((model.numel() + quant_block_size - 1) / quant_block_size); + const auto stream = at::cuda::getCurrentCUDAStream(); + adamw_8bit_blockwise_kernel<<>>( + model.data_ptr(), + grad.data_ptr(), + exp_avg_q.data_ptr(), + exp_avg_scale.data_ptr(), + exp_avg_sq_q.data_ptr(), + exp_avg_sq_scale.data_ptr(), + model.numel(), + quant_block_size, + beta1, + beta2, + effective_lr, + weight_decay, + eps, + step_size, + bias_correction2_sqrt); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + } // namespace void areno_adamw_fp32_master_step_cuda( @@ -240,3 +558,102 @@ void areno_adamw_fp32_master_step_cuda( TORCH_CHECK(false, "gradient must be bfloat16 or float32"); } } + +void areno_adamw_4bit_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 exp_avg_sq_scale, + int64_t packed_offset, + int64_t 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) { + c10::cuda::CUDAGuard guard(model.device()); + TORCH_CHECK( + model.is_cuda() && grad.is_cuda() && exp_avg_q.is_cuda() && exp_avg_scale.is_cuda() && + exp_avg_sq_q.is_cuda() && exp_avg_sq_scale.is_cuda(), + "AdamW4bit tensors must be CUDA tensors"); + TORCH_CHECK(model.is_contiguous() && grad.is_contiguous(), "model and gradient must be contiguous"); + TORCH_CHECK(model.numel() == grad.numel(), "model and gradient sizes must match"); + 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"); + +#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, \ + 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(at::BFloat16, at::BFloat16); + } else if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW4(at::BFloat16, float); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kBFloat16) { + LAUNCH_ADAMW4(float, at::BFloat16); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW4(float, float); + } else { + TORCH_CHECK(false, "model and gradient must be bfloat16 or float32"); + } +#undef LAUNCH_ADAMW4 +} + +void areno_adamw_8bit_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 exp_avg_sq_scale, + 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()); + TORCH_CHECK( + model.is_cuda() && grad.is_cuda() && exp_avg_q.is_cuda() && exp_avg_scale.is_cuda() && + exp_avg_sq_q.is_cuda() && exp_avg_sq_scale.is_cuda(), + "all 8-bit AdamW inputs must be CUDA tensors"); + TORCH_CHECK( + model.is_contiguous() && grad.is_contiguous() && exp_avg_q.is_contiguous() && exp_avg_scale.is_contiguous() && + exp_avg_sq_q.is_contiguous() && exp_avg_sq_scale.is_contiguous(), + "all 8-bit AdamW inputs must be contiguous"); + TORCH_CHECK(model.numel() == grad.numel(), "model and gradient sizes must match"); + TORCH_CHECK(model.numel() == exp_avg_q.numel(), "model and first-moment sizes must match"); + TORCH_CHECK(model.numel() == exp_avg_sq_q.numel(), "model and second-moment sizes must match"); + TORCH_CHECK(quant_block_size >= 1 && quant_block_size <= 4096, "quantization block size must be in [1, 4096]"); + const int64_t block_count = (model.numel() + quant_block_size - 1) / quant_block_size; + TORCH_CHECK(exp_avg_scale.numel() == block_count, "first-moment scale count must match quantization blocks"); + TORCH_CHECK(exp_avg_sq_scale.numel() == block_count, "second-moment scale count must match quantization blocks"); + +#define LAUNCH_ADAMW8(MODEL_T, GRAD_T) \ + launch_adamw_8bit( \ + model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, 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_ADAMW8(at::BFloat16, at::BFloat16); + } else if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW8(at::BFloat16, float); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kBFloat16) { + LAUNCH_ADAMW8(float, at::BFloat16); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW8(float, float); + } else { + TORCH_CHECK(false, "model and gradient must be bfloat16 or float32"); + } +#undef LAUNCH_ADAMW8 +} diff --git a/areno/accel/optimizer.py b/areno/accel/optimizer.py index 1e815787..0bba12ba 100644 --- a/areno/accel/optimizer.py +++ b/areno/accel/optimizer.py @@ -1,4 +1,4 @@ -"""Fused CUDA update for AReno's compact FP32-master AdamW.""" +"""Fused CUDA updates for AReno optimizers.""" from __future__ import annotations @@ -67,4 +67,137 @@ def areno_adamw_fp32_master_step( ) -__all__ = ["areno_adamw_fp32_master_step"] +@torch._dynamo.disable +@torch.no_grad() +def areno_adamw_8bit_step( + model: torch.Tensor, + grad: torch.Tensor, + exp_avg_q: torch.Tensor, + exp_avg_scale: torch.Tensor, + exp_avg_sq_q: torch.Tensor, + exp_avg_sq_scale: torch.Tensor, + *, + block_size: int, + beta1: float, + beta2: float, + effective_lr: float, + weight_decay: float, + eps: float, + step_size: float, + bias_correction2_sqrt: float, +) -> None: + """Update block-quantized AdamW state without full FP32 moments.""" + + tensors = (model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale) + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("fused 8-bit AdamW requires CUDA tensors") + if any(tensor.device != model.device for tensor in tensors[1:]): + raise ValueError("fused 8-bit AdamW requires every tensor on the model device") + if model.dtype not in {torch.bfloat16, torch.float32}: + raise TypeError(f"fused 8-bit AdamW requires bfloat16 or float32 model weights, got {model.dtype}") + if grad.dtype not in {torch.bfloat16, torch.float32}: + raise TypeError(f"fused 8-bit AdamW requires bfloat16 or float32 gradients, got {grad.dtype}") + if exp_avg_q.dtype != torch.uint8 or exp_avg_sq_q.dtype != torch.uint8: + raise TypeError("quantized Adam moments must use uint8") + if exp_avg_scale.dtype != torch.float32 or exp_avg_sq_scale.dtype != torch.float32: + raise TypeError("quantized Adam scales must use float32") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("fused 8-bit AdamW requires contiguous tensors") + if model.numel() != grad.numel() or model.numel() != exp_avg_q.numel(): + raise ValueError("model, gradient, and quantized moment tensors must have the same number of elements") + if exp_avg_q.numel() != exp_avg_sq_q.numel(): + raise ValueError("first- and second-moment tensors must have the same number of elements") + if block_size < 1 or block_size > 4096: + raise ValueError(f"block_size must be between 1 and 4096, got {block_size}") + block_count = (model.numel() + block_size - 1) // block_size + if exp_avg_scale.numel() != block_count or exp_avg_sq_scale.numel() != block_count: + raise ValueError("scale tensors must contain one value per quantization block") + extension().areno_adamw_8bit_step( + model, + grad, + exp_avg_q, + exp_avg_scale, + exp_avg_sq_q, + exp_avg_sq_scale, + block_size, + beta1, + beta2, + effective_lr, + weight_decay, + eps, + step_size, + bias_correction2_sqrt, + ) + + +@torch._dynamo.disable +@torch.no_grad() +def areno_adamw_4bit_step( + model: torch.Tensor, + grad: torch.Tensor, + exp_avg_q: torch.Tensor, + exp_avg_scale: torch.Tensor, + exp_avg_sq_q: torch.Tensor, + exp_avg_sq_scale: torch.Tensor, + *, + packed_offset: int, + scale_offset: 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 one model shard directly from packed block-wise moments.""" + + tensors = (model, grad, exp_avg_q, exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale) + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("fused AdamW4bit requires CUDA tensors") + if model.dtype not in {torch.bfloat16, torch.float32}: + raise TypeError(f"fused AdamW4bit requires bfloat16 or float32 model weights, got {model.dtype}") + if grad.dtype not in {torch.bfloat16, torch.float32}: + raise TypeError(f"fused AdamW4bit requires bfloat16 or float32 gradients, got {grad.dtype}") + if exp_avg_q.dtype != torch.uint8 or exp_avg_sq_q.dtype != torch.uint8: + raise TypeError("packed AdamW4bit moments must use uint8 storage") + if exp_avg_scale.dtype != torch.float32 or exp_avg_sq_scale.dtype != torch.float32: + raise TypeError("AdamW4bit scales must use float32 storage") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("fused AdamW4bit requires contiguous tensors") + if model.numel() != grad.numel(): + raise ValueError("model and gradient shards must have the same number of elements") + 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") + 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 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") + extension().areno_adamw_4bit_step( + model, + grad, + exp_avg_q, + exp_avg_scale, + exp_avg_sq_q, + exp_avg_sq_scale, + packed_offset, + scale_offset, + quant_block_size, + beta1, + beta2, + effective_lr, + weight_decay, + eps, + step_size, + bias_correction2_sqrt, + ) + + +__all__ = ["areno_adamw_4bit_step", "areno_adamw_8bit_step", "areno_adamw_fp32_master_step"] diff --git a/areno/api/backend/cuda/roles.py b/areno/api/backend/cuda/roles.py index 452da4e1..99bd9521 100644 --- a/areno/api/backend/cuda/roles.py +++ b/areno/api/backend/cuda/roles.py @@ -11,7 +11,7 @@ from areno.engine.config import EngineConfig from areno.engine.data import to_device from areno.engine.modeling import build_model_on_device, build_optimizer, canonical_model_path, param_grad, unwrap_model -from areno.engine.optim import AdamW8bit, AdamWFP32Master +from areno.engine.optim import AdamW4bit, AdamW8bit, AdamWFP32Master from areno.engine.parallel.collectives import gather_from_sequence_parallel_region from areno.engine.parallel.context import get_tp_context from areno.engine.protocol import EnsureRolesPayload, ScorePayload, TrainValuesPayload @@ -172,7 +172,7 @@ def __init__( self, path: str, model: torch.nn.Module, - optimizer: AdamW8bit | AdamWFP32Master | None, + optimizer: AdamW4bit | AdamW8bit | AdamWFP32Master | None, value_head: torch.nn.Module | None, sequence_parallel: bool = False, *, diff --git a/areno/api/trainer_config.py b/areno/api/trainer_config.py index 6ec5dd3d..7db26af3 100644 --- a/areno/api/trainer_config.py +++ b/areno/api/trainer_config.py @@ -57,6 +57,7 @@ class TrainerConfig: weight_decay: float = 1.0e-2 grad_clip_norm: float = 1.0 adam_8bit: bool = False + adam_4bit: bool = False unfreeze_multimodal_tower: bool = False unfreeze_multimodal_projector: bool = False multimodal_tower_lr: float | None = None @@ -90,6 +91,10 @@ def __post_init__(self) -> None: self.backend = self.backend.lower() if self.backend not in {"cuda", "mlx"}: raise ValueError("backend must be one of: cuda, mlx") + if self.adam_4bit and self.adam_8bit: + raise ValueError("adam_4bit and adam_8bit are mutually exclusive") + if self.adam_4bit and self.backend != "cuda": + raise ValueError("adam_4bit is only supported by the CUDA backend") if self.attn_backend not in {"flash", "native"}: raise ValueError("attn_backend must be one of: flash, native") if self.model_hub not in {"hf", "modelscope"}: @@ -155,6 +160,7 @@ def optimizer_config(self) -> dict: "weight_decay": self.weight_decay, "grad_clip_norm": self.grad_clip_norm, "adam_8bit": self.adam_8bit, + "adam_4bit": self.adam_4bit, "unfreeze_multimodal_tower": self.unfreeze_multimodal_tower, "unfreeze_multimodal_projector": self.unfreeze_multimodal_projector, "multimodal_tower_lr": self.multimodal_tower_lr, diff --git a/areno/cli/train.py b/areno/cli/train.py index c8f854d9..1de993a3 100644 --- a/areno/cli/train.py +++ b/areno/cli/train.py @@ -145,6 +145,7 @@ def flash_attention_unsupported_model_reason(model_config): "adam_beta1", "adam_beta2", "adam_8bit", + "adam_4bit", "optimizer_state_offload", "optimizer_state_offload_dir", "optimizer_state_offload_batch_size", @@ -238,6 +239,7 @@ def _trainer_config_from_options(**options) -> TrainerConfig: args.rollout_tp_size = getattr(args, "rollout_tp_size", None) args.rollout_devices = getattr(args, "rollout_devices", None) args.policy_sync_bucket_mb = getattr(args, "policy_sync_bucket_mb", 64) + args.adam_4bit = getattr(args, "adam_4bit", False) args.optimizer_state_offload = getattr(args, "optimizer_state_offload", "none") args.optimizer_state_offload_dir = getattr(args, "optimizer_state_offload_dir", None) args.optimizer_state_offload_batch_size = getattr(args, "optimizer_state_offload_batch_size", 1) @@ -543,7 +545,8 @@ def _format_training_config_summary( f"lr={config.optimizer_lr}, min_lr={config.optimizer_min_lr}, " f"decay={config.lr_decay_style}/{config.lr_decay_steps}, " f"betas=({config.optimizer_beta1}, {config.optimizer_beta2}), " - f"weight_decay={config.weight_decay}, adam_8bit={_format_bool(config.adam_8bit)}" + f"weight_decay={config.weight_decay}, adam_8bit={_format_bool(config.adam_8bit)}, " + f"adam_4bit={_format_bool(config.adam_4bit)}" ), ), ("optimizer_state_offload", str(config.optimizer_state_offload)), @@ -844,6 +847,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: args.rollout_tp_size = getattr(args, "rollout_tp_size", None) args.rollout_devices = getattr(args, "rollout_devices", None) args.policy_sync_bucket_mb = getattr(args, "policy_sync_bucket_mb", 64) + args.adam_4bit = getattr(args, "adam_4bit", False) args.unfreeze_multimodal_tower = getattr(args, "unfreeze_multimodal_tower", False) args.unfreeze_multimodal_projector = getattr(args, "unfreeze_multimodal_projector", False) args.multimodal_tower_lr = getattr(args, "multimodal_tower_lr", None) @@ -891,6 +895,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: weight_decay=args.weight_decay, grad_clip_norm=args.grad_clip_norm, adam_8bit=args.adam_8bit, + adam_4bit=args.adam_4bit, unfreeze_multimodal_tower=args.unfreeze_multimodal_tower, unfreeze_multimodal_projector=args.unfreeze_multimodal_projector, multimodal_tower_lr=args.multimodal_tower_lr, @@ -950,6 +955,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: weight_decay=args.weight_decay, grad_clip_norm=args.grad_clip_norm, adam_8bit=args.adam_8bit, + adam_4bit=args.adam_4bit, unfreeze_multimodal_tower=args.unfreeze_multimodal_tower, unfreeze_multimodal_projector=args.unfreeze_multimodal_projector, multimodal_tower_lr=args.multimodal_tower_lr, @@ -1017,6 +1023,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: weight_decay=args.weight_decay, grad_clip_norm=args.grad_clip_norm, adam_8bit=args.adam_8bit, + adam_4bit=args.adam_4bit, unfreeze_multimodal_tower=args.unfreeze_multimodal_tower, unfreeze_multimodal_projector=args.unfreeze_multimodal_projector, multimodal_tower_lr=args.multimodal_tower_lr, @@ -1085,6 +1092,7 @@ def _trainer_config_from_args(args) -> TrainerConfig: weight_decay=args.weight_decay, grad_clip_norm=args.grad_clip_norm, adam_8bit=args.adam_8bit, + adam_4bit=args.adam_4bit, unfreeze_multimodal_tower=args.unfreeze_multimodal_tower, unfreeze_multimodal_projector=args.unfreeze_multimodal_projector, multimodal_tower_lr=args.multimodal_tower_lr, @@ -1270,6 +1278,7 @@ def section(title: str, names: list[str]) -> dict: "weight_decay", "grad_clip_norm", "adam_8bit", + "adam_4bit", ], ), section( @@ -1676,6 +1685,7 @@ def _dataset_builder_for_suffix(suffix: str) -> str: @click.option("--adam-beta1", type=float, default=0.9, show_default=True, help="Policy optimizer Adam beta1.") @click.option("--adam-beta2", type=float, default=0.999, show_default=True, help="Policy optimizer Adam beta2.") @click.option("--adam-8bit", is_flag=True, help="Use 8-bit Adam moment states instead of FP32 Adam states.") +@click.option("--adam-4bit", is_flag=True, help="Use packed block-wise 4-bit Adam moment states.") @click.option("--unfreeze-mm-tower", "unfreeze_multimodal_tower", is_flag=True, help="Train multimodal encoder towers.") @click.option( "--unfreeze-mm-projector", diff --git a/areno/dashboard/server.py b/areno/dashboard/server.py index 0f06cc57..9fd9c183 100644 --- a/areno/dashboard/server.py +++ b/areno/dashboard/server.py @@ -710,6 +710,7 @@ def build_train_command(config: dict[str, Any]) -> list[str]: "--tune-params": config.get("tune_params"), "--greedy": config.get("greedy"), "--adam-8bit": config.get("adam_8bit"), + "--adam-4bit": config.get("adam_4bit"), "--unfreeze-mm-tower": config.get("unfreeze_multimodal_tower"), "--unfreeze-mm-projector": config.get("unfreeze_multimodal_projector"), "--drop-rollout-state": config.get("drop_rollout_state"), @@ -1388,6 +1389,7 @@ def launcher_preflight_action(payload: dict[str, Any]) -> dict[str, Any]: "mini_bs": mini_bs, "max_running_prompts": 1, "adam_8bit": bool_like(config.get("adam_8bit")), + "adam_4bit": bool_like(config.get("adam_4bit")), "keep_rollout_state": False, } ) diff --git a/areno/engine/config.py b/areno/engine/config.py index 763960e1..fbc6e0f6 100644 --- a/areno/engine/config.py +++ b/areno/engine/config.py @@ -38,6 +38,7 @@ class OptimizerConfig: weight_decay: float = 0.0 grad_clip_norm: float | None = None adam_8bit: bool = False + adam_4bit: bool = False fp32_master_bucket_numel: int = 16 * 1024 * 1024 unfreeze_multimodal_tower: bool = False unfreeze_multimodal_projector: bool = False @@ -50,6 +51,10 @@ class OptimizerConfig: multimodal_projector_lr_decay_steps: int | None = None multimodal_projector_lr_decay_style: Literal["constant", "linear", "cosine"] | None = None + def __post_init__(self) -> None: + if self.adam_4bit and self.adam_8bit: + raise ValueError("optimizer.adam_4bit and optimizer.adam_8bit are mutually exclusive") + @dataclass(slots=True) class RuntimeConfig: diff --git a/areno/engine/modeling.py b/areno/engine/modeling.py index b68413a3..471161ca 100644 --- a/areno/engine/modeling.py +++ b/areno/engine/modeling.py @@ -7,7 +7,7 @@ import torch from areno.engine.config import EngineConfig -from areno.engine.optim import AdamW8bit, AdamWFP32Master +from areno.engine.optim import AdamW4bit, AdamW8bit, AdamWFP32Master from areno.models.registry import build_model @@ -97,7 +97,12 @@ def canonical_model_path(path: str | None) -> str | None: def build_optimizer(params, optimizer_config, ctx, *, lr: float | None = None): """Construct the configured DP-sharded optimizer implementation.""" - optimizer_cls = AdamW8bit if optimizer_config.adam_8bit else AdamWFP32Master + if optimizer_config.adam_4bit: + optimizer_cls = AdamW4bit + elif optimizer_config.adam_8bit: + optimizer_cls = AdamW8bit + else: + optimizer_cls = AdamWFP32Master return optimizer_cls( params, lr=optimizer_config.lr if lr is None else float(lr), diff --git a/areno/engine/optim/__init__.py b/areno/engine/optim/__init__.py index 9db6cf25..cd27ed45 100644 --- a/areno/engine/optim/__init__.py +++ b/areno/engine/optim/__init__.py @@ -5,7 +5,8 @@ re-exported through `__all__`. """ +from areno.engine.optim.adamw_4bit import AdamW4bit from areno.engine.optim.adamw_8bit import AdamW8bit from areno.engine.optim.adamw_fp32_master import AdamWFP32Master -__all__ = ["AdamW8bit", "AdamWFP32Master"] +__all__ = ["AdamW4bit", "AdamW8bit", "AdamWFP32Master"] diff --git a/areno/engine/optim/adamw_4bit.py b/areno/engine/optim/adamw_4bit.py new file mode 100644 index 00000000..939da6ad --- /dev/null +++ b/areno/engine/optim/adamw_4bit.py @@ -0,0 +1,372 @@ +"""Packed block-wise 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. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator + +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 + +_DEFAULT_QUANT_BLOCK_SIZE = 128 +_STATE_FORMAT_VERSION = 1 +# 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 = ( + -0.8875, + -0.6625, + -0.4375, + -0.2125, + -0.0775, + -0.0325, + -0.0055, + 0.0, + 0.0055, + 0.0325, + 0.0775, + 0.2125, + 0.4375, + 0.6625, + 0.8875, + 1.0, +) + + +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. + """ + + def __init__( + self, + params: Iterable[torch.nn.Parameter], + *, + lr: float, + betas: tuple[float, float], + weight_decay: float, + bucket_numel: int = _DEFAULT_BUCKET_NUMEL, + quant_block_size: int = _DEFAULT_QUANT_BLOCK_SIZE, + dp_rank: int = 0, + dp_size: int = 1, + dp_group: dist.ProcessGroup | None = None, + ): + 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") + self.quant_block_size = quant_block_size + super().__init__( + params, + lr=lr, + betas=betas, + weight_decay=weight_decay, + bucket_numel=bucket_numel, + dp_rank=dp_rank, + dp_size=dp_size, + dp_group=dp_group, + quant_block_size=quant_block_size, + ) + + def state_dict(self) -> dict: + """Return the versioned packed state for this DP rank.""" + + payload = super().state_dict() + payload.pop("adam_8bit", None) + payload["adam_4bit"] = True + payload["state_format_version"] = _STATE_FORMAT_VERSION + payload["quant_block_size"] = self.quant_block_size + return payload + + @torch.no_grad() + def load_state_dict(self, state_dict: dict) -> None: + """Restore packed state, rejecting incompatible layouts explicitly.""" + + version = int(state_dict.get("state_format_version", 0)) + if version != _STATE_FORMAT_VERSION: + raise ValueError(f"unsupported AdamW4bit state format version: {version}") + saved_block_size = int(state_dict.get("quant_block_size", 0)) + if saved_block_size != self.quant_block_size: + raise ValueError( + f"AdamW4bit quant_block_size mismatch: checkpoint={saved_block_size}, optimizer={self.quant_block_size}" + ) + self._cleanup_disk_offload() + self._active_offload_mode = "none" + self._disk_offload_root = None + self._active_offload_batch_size = 1 + for state in self._states: + state.offload_file = None + state.offload_index = None + state.offload_group = None + state.offload_ready_events = () + saved_states = state_dict.get("state", []) + 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) + 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_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) + + @torch.no_grad() + def _ensure_bucket_state(self, bucket: _MasterBucket, state) -> None: + """Materialize packed moments and per-block scales for one bucket.""" + + device = bucket.refs[0].model_param.device + if state.offload_file is not None: + self._load_state_offload(state, device) + for name in ("exp_avg_q", "exp_avg_scale", "exp_avg_sq_q", "exp_avg_sq_scale"): + 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) + 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) + 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) + + 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]) + specs[index] = { + "exp_avg_q": (torch.uint8, (packed_numel,)), + "exp_avg_scale": (torch.float32, (scale_numel,)), + "exp_avg_sq_q": (torch.uint8, (packed_numel,)), + "exp_avg_sq_scale": (torch.float32, (scale_numel,)), + } + return specs + + @torch.no_grad() + def _step_bucket_8bit(self, bucket: _MasterBucket, state) -> 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, scale_offset in self._iter_ref_layout(bucket): + grad = self._gradient_for_ref(bucket, ref) + if grad is None: + continue + 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 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() + + @torch.no_grad() + def _step_param_ref_4bit( + self, + bucket: _MasterBucket, + ref: _ParamRef, + grad: torch.Tensor, + state, + packed_offset: int, + scale_offset: int, + beta1: float, + beta2: float, + effective_lr: float, + step_size: float, + bias_correction2_sqrt: float, + ) -> None: + """Apply AdamW to one parameter shard in block-sized work buffers.""" + + 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) + model_chunk = ref.model_param.detach().reshape(-1).narrow(0, ref.param_start, ref.numel) + model_shard = model_chunk.narrow(0, ref.shard_start, ref.shard_numel) + if model_shard.is_cuda: + from areno.accel.optimizer import areno_adamw_4bit_step + + areno_adamw_4bit_step( + model_shard, + grad_shard.contiguous(), + state.exp_avg_q, + state.exp_avg_scale, + state.exp_avg_sq_q, + state.exp_avg_sq_scale, + packed_offset=packed_offset, + scale_offset=scale_offset, + 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 + 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 + scale_index = 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], + ) + variance = _unpack_positive_4bit( + state.exp_avg_sq_q.narrow(0, byte_start, byte_count), + count, + state.exp_avg_sq_scale[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) + if not ( + torch.isfinite(grad_block).all() + and torch.isfinite(weight).all() + and torch.isfinite(moment).all() + and torch.isfinite(variance).all() + ): + # Match the fused CUDA path: a bad block must not poison its + # packed state or any neighboring block. + continue + if self.weight_decay != 0.0: + weight.mul_(1.0 - effective_lr * self.weight_decay) + moment.mul_(beta1).add_(grad_block, alpha=1.0 - beta1) + variance.mul_(beta2).addcmul_(grad_block, grad_block, value=1.0 - beta2) + denom = variance.sqrt().div_(bias_correction2_sqrt).add_(self.eps) + weight.addcdiv_(moment, denom, value=-step_size) + 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_scale[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) + + def _bucket_state_sizes(self, bucket: _MasterBucket) -> tuple[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 + + def _iter_ref_layout(self, bucket: _MasterBucket) -> Iterator[tuple[_ParamRef, int, int]]: + packed_offset = 0 + scale_offset = 0 + for ref in bucket.refs: + yield ref, packed_offset, scale_offset + packed_offset += (ref.shard_numel + 1) // 2 + scale_offset += (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.""" + + total = 0 + for state in self._states: + 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() + return total + + +def _load_tensor( + saved: dict, + name: str, + device: torch.device, + dtype: torch.dtype, + expected_numel: int, +) -> torch.Tensor | None: + value = saved.get(name) + if value is None: + return None + result = value.detach().to(device=device, dtype=dtype).view(-1).clone() + if result.numel() != expected_numel: + raise ValueError(f"AdamW4bit {name} has {result.numel()} values, expected {expected_numel}") + return result + + +def _pack_nibbles(codes: torch.Tensor) -> torch.Tensor: + """Pack uint8 values in [0, 15], low nibble first.""" + + if codes.numel() == 0: + return codes.to(dtype=torch.uint8) + codes = codes.to(dtype=torch.uint8).view(-1) + if codes.numel() % 2: + codes = torch.cat((codes, torch.zeros(1, device=codes.device, dtype=torch.uint8))) + return codes[0::2] | (codes[1::2] << 4) + + +def _unpack_nibbles(packed: torch.Tensor, numel: int) -> torch.Tensor: + """Unpack low/high nibbles into a uint8 vector.""" + + result = torch.empty(packed.numel() * 2, device=packed.device, dtype=torch.uint8) + result[0::2] = packed & 0x0F + result[1::2] = packed >> 4 + return result[:numel] + + +def _quantize_signed_4bit(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize a signed block with the paper's dynamic-exponent map.""" + + if tensor.numel() == 0: + return tensor.to(dtype=torch.uint8), torch.ones((), device=tensor.device, dtype=torch.float32) + scale = tensor.abs().amax().to(dtype=torch.float32) + mapping = tensor.new_tensor(_SIGNED_DE_MAP, dtype=torch.float32) + normalized = tensor / scale.clamp_min(1.0e-30) + codes = torch.argmin((normalized.unsqueeze(-1) - mapping).abs(), dim=-1).to(dtype=torch.uint8) + return _pack_nibbles(codes), scale.to(dtype=torch.float32) + + +def _unpack_signed_4bit(packed: torch.Tensor, numel: int, scale: torch.Tensor) -> torch.Tensor: + mapping = packed.new_tensor(_SIGNED_DE_MAP, dtype=torch.float32) + return mapping[_unpack_nibbles(packed, numel).to(dtype=torch.long)].mul_(scale) + + +def _quantize_positive_4bit(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize with T(i)=(i+1)/16, deliberately excluding zero.""" + + if tensor.numel() == 0: + return tensor.to(dtype=torch.uint8), torch.zeros((), device=tensor.device, dtype=torch.float32) + scale = tensor.amax().to(dtype=torch.float32) + safe_scale = scale.clamp_min(1.0e-30) + codes = torch.clamp(torch.round(tensor / safe_scale * 16.0 - 1.0), 0.0, 15.0).to(dtype=torch.uint8) + return _pack_nibbles(codes), scale + + +def _unpack_positive_4bit(packed: torch.Tensor, numel: int, scale: torch.Tensor) -> torch.Tensor: + return (_unpack_nibbles(packed, numel).to(dtype=torch.float32) + 1.0).mul_(scale / 16.0) + + +__all__ = ["AdamW4bit"] diff --git a/areno/engine/optim/adamw_8bit.py b/areno/engine/optim/adamw_8bit.py index a53014b0..632b2136 100644 --- a/areno/engine/optim/adamw_8bit.py +++ b/areno/engine/optim/adamw_8bit.py @@ -18,6 +18,9 @@ _ParamRef, ) +_DEFAULT_QUANT_BLOCK_SIZE = 128 +_MAX_FUSED_QUANT_BLOCK_SIZE = 4096 + @dataclass(slots=True) class _Adam8bitBucketState: @@ -53,7 +56,12 @@ def __init__( dp_rank: int = 0, dp_size: int = 1, dp_group: dist.ProcessGroup | None = None, + quant_block_size: int = _DEFAULT_QUANT_BLOCK_SIZE, ): + if quant_block_size < 1 or quant_block_size > _MAX_FUSED_QUANT_BLOCK_SIZE: + raise ValueError( + f"quant_block_size must be between 1 and {_MAX_FUSED_QUANT_BLOCK_SIZE}, got {quant_block_size}" + ) super().__init__( params, lr=lr, @@ -64,6 +72,7 @@ def __init__( dp_size=dp_size, dp_group=dp_group, ) + self.quant_block_size = quant_block_size self._states = [_Adam8bitBucketState() for _ in self.buckets] @torch.no_grad() @@ -181,6 +190,7 @@ def state_dict(self) -> dict: "dp_rank": self.dp_rank, "dp_size": self.dp_size, "adam_8bit": True, + "quant_block_size": self.quant_block_size, "state": [ { "step": state.step, @@ -207,6 +217,11 @@ def load_state_dict(self, state_dict: dict) -> None: state.offload_group = None state.offload_ready_events = () saved_states = state_dict.get("state", []) + if "quant_block_size" in state_dict: + saved_block_size = int(state_dict["quant_block_size"]) + if saved_block_size < 1 or saved_block_size > _MAX_FUSED_QUANT_BLOCK_SIZE: + raise ValueError(f"invalid saved AdamW8bit quant_block_size: {saved_block_size}") + self.quant_block_size = saved_block_size for saved, bucket, state in zip(saved_states[: len(self.buckets)], self.buckets, self._states, strict=False): if saved is None: continue @@ -219,22 +234,32 @@ def load_state_dict(self, state_dict: dict) -> None: state.exp_avg_q = ( None if exp_avg_q is None else exp_avg_q.detach().to(device=device, dtype=torch.uint8).view(-1).clone() ) - state.exp_avg_scale = ( - None - if exp_avg_scale is None - else exp_avg_scale.detach().to(device=device, dtype=torch.float32).view(()).clone() - ) + state.exp_avg_scale = None if exp_avg_scale is None else self._restore_scales(exp_avg_scale, bucket, device) state.exp_avg_sq_q = ( None if exp_avg_sq_q is None else exp_avg_sq_q.detach().to(device=device, dtype=torch.uint8).view(-1).clone() ) state.exp_avg_sq_scale = ( - None - if exp_avg_sq_scale is None - else exp_avg_sq_scale.detach().to(device=device, dtype=torch.float32).view(()).clone() + None if exp_avg_sq_scale is None else self._restore_scales(exp_avg_sq_scale, bucket, device) ) + def _restore_scales( + self, + saved: torch.Tensor, + bucket: _MasterBucket, + device: torch.device, + ) -> torch.Tensor: + """Restore block scales, expanding legacy bucket-level scalar scales.""" + + expected = self._bucket_scale_count(bucket) + scales = saved.detach().to(device=device, dtype=torch.float32).view(-1) + if scales.numel() == 1 and expected != 1: + return scales.expand(expected).clone() + if scales.numel() != expected: + raise ValueError(f"AdamW8bit checkpoint has {scales.numel()} scales for a bucket requiring {expected}") + return scales.clone() + @torch.no_grad() def _ensure_bucket_state(self, bucket: _MasterBucket, state: _Adam8bitBucketState) -> None: """Materialize or onload quantized moments for one bucket.""" @@ -252,10 +277,26 @@ def _ensure_bucket_state(self, bucket: _MasterBucket, state: _Adam8bitBucketStat state.exp_avg_sq_scale = state.exp_avg_sq_scale.to(device=device) if state.exp_avg_q is None: state.exp_avg_q = torch.full((bucket.shard_numel,), 128, device=device, dtype=torch.uint8) - state.exp_avg_scale = torch.ones((), device=device, dtype=torch.float32) + state.exp_avg_scale = torch.ones(self._bucket_scale_count(bucket), device=device, dtype=torch.float32) if state.exp_avg_sq_q is None: state.exp_avg_sq_q = torch.zeros(bucket.shard_numel, device=device, dtype=torch.uint8) - state.exp_avg_sq_scale = torch.ones((), device=device, dtype=torch.float32) + state.exp_avg_sq_scale = torch.ones(self._bucket_scale_count(bucket), device=device, dtype=torch.float32) + + def _bucket_scale_count(self, bucket: _MasterBucket) -> int: + """Return the number of independently scaled blocks in one DP shard.""" + + return sum(_ceil_div(ref.shard_numel, self.quant_block_size) for ref in bucket.refs) + + def _ref_scale_layout(self, bucket: _MasterBucket) -> list[tuple[_ParamRef, int, int]]: + """Map each parameter ref to its contiguous range in the scale tensors.""" + + layout: list[tuple[_ParamRef, int, int]] = [] + scale_offset = 0 + for ref in bucket.refs: + block_count = _ceil_div(ref.shard_numel, self.quant_block_size) + layout.append((ref, scale_offset, block_count)) + scale_offset += block_count + return layout def _load_state_offload(self, state: _Adam8bitBucketState, device: torch.device) -> None: """Copy one quantized bucket from its persistent raw mmap.""" @@ -287,9 +328,9 @@ def _state_mmap_specs(self, indices: list[int]) -> dict[int, dict[str, tuple[tor return { index: { "exp_avg_q": (torch.uint8, (self.buckets[index].shard_numel,)), - "exp_avg_scale": (torch.float32, ()), + "exp_avg_scale": (torch.float32, (self._bucket_scale_count(self.buckets[index]),)), "exp_avg_sq_q": (torch.uint8, (self.buckets[index].shard_numel,)), - "exp_avg_sq_scale": (torch.float32, ()), + "exp_avg_sq_scale": (torch.float32, (self._bucket_scale_count(self.buckets[index]),)), } for index in indices } @@ -371,7 +412,7 @@ def _state_cpu_payload( @torch.no_grad() def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState) -> None: - """Update all parameter chunks in one bucket using dequantized moments.""" + """Update one bucket without materializing full FP32 moment tensors.""" assert state.exp_avg_q is not None assert state.exp_avg_scale is not None @@ -383,9 +424,7 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState) bias_correction2 = 1.0 - beta2**state.step bias_correction2_sqrt = bias_correction2**0.5 - exp_avg = _dequantize_symmetric(state.exp_avg_q, state.exp_avg_scale) - exp_avg_sq = _dequantize_positive(state.exp_avg_sq_q, state.exp_avg_sq_scale) - for ref in bucket.refs: + for ref, scale_offset, block_count in self._ref_scale_layout(bucket): grad = self._gradient_for_ref(bucket, ref) if grad is None: continue @@ -395,8 +434,9 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState) bucket, ref, grad, - exp_avg, - exp_avg_sq, + state, + scale_offset, + block_count, beta1, beta2, effective_lr, @@ -407,8 +447,6 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState) ref.model_param.grad = None if isinstance(getattr(ref.model_param, "main_grad", None), torch.Tensor): ref.model_param.main_grad = None - state.exp_avg_q, state.exp_avg_scale = _quantize_symmetric(exp_avg) - state.exp_avg_sq_q, state.exp_avg_sq_scale = _quantize_positive(exp_avg_sq) # Collective order is bucket-global, not rank-local. A rank can own # no values from a small DP bucket and must still join the gather that # refreshes every replicated model parameter. @@ -422,8 +460,9 @@ def _step_param_ref_8bit( bucket: _MasterBucket, ref: _ParamRef, grad: torch.Tensor, - exp_avg: torch.Tensor, - exp_avg_sq: torch.Tensor, + state: _Adam8bitBucketState, + scale_offset: int, + block_count: int, beta1: float, beta2: float, effective_lr: float, @@ -434,22 +473,64 @@ def _step_param_ref_8bit( if ref.shard_numel == 0: return + 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 if bucket.grad_shard is not None: - grad_shard = grad.to(dtype=torch.float32) + grad_shard = grad else: - grad_shard = grad.narrow(0, ref.shard_start, ref.shard_numel).to(dtype=torch.float32) + grad_shard = grad.narrow(0, ref.shard_start, ref.shard_numel) model_chunk = ref.model_param.detach().reshape(-1).narrow(0, ref.param_start, ref.numel) model_shard = model_chunk.narrow(0, ref.shard_start, ref.shard_numel) - weight = model_shard.to(dtype=torch.float32) - if self.weight_decay != 0.0: - weight.mul_(1.0 - effective_lr * self.weight_decay) - moment = exp_avg.narrow(0, ref.shard_bucket_start, ref.shard_numel) - variance = exp_avg_sq.narrow(0, ref.shard_bucket_start, ref.shard_numel) - moment.mul_(beta1).add_(grad_shard, alpha=1.0 - beta1) - variance.mul_(beta2).addcmul_(grad_shard, grad_shard, value=1.0 - beta2) - denom = variance.sqrt().div_(bias_correction2_sqrt).add_(self.eps) - weight.addcdiv_(moment, denom, value=-step_size) - model_shard.copy_(weight) + moment_q = state.exp_avg_q.narrow(0, ref.shard_bucket_start, ref.shard_numel) + variance_q = state.exp_avg_sq_q.narrow(0, ref.shard_bucket_start, ref.shard_numel) + moment_scales = state.exp_avg_scale.narrow(0, scale_offset, block_count) + variance_scales = state.exp_avg_sq_scale.narrow(0, scale_offset, block_count) + + if model_shard.is_cuda: + from areno.accel.optimizer import areno_adamw_8bit_step + + areno_adamw_8bit_step( + model_shard, + grad_shard.contiguous(), + moment_q, + moment_scales, + variance_q, + variance_scales, + 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 + + for block_index in range(block_count): + start = block_index * self.quant_block_size + numel = min(self.quant_block_size, ref.shard_numel - start) + weight = model_shard.narrow(0, start, numel).to(dtype=torch.float32) + block_grad = grad_shard.narrow(0, start, numel).to(dtype=torch.float32) + block_moment_q = moment_q.narrow(0, start, numel) + block_variance_q = variance_q.narrow(0, start, numel) + moment = _dequantize_symmetric(block_moment_q, moment_scales[block_index]) + variance = _dequantize_positive(block_variance_q, variance_scales[block_index]) + if self.weight_decay != 0.0: + weight.mul_(1.0 - effective_lr * self.weight_decay) + moment.mul_(beta1).add_(block_grad, alpha=1.0 - beta1) + variance.mul_(beta2).addcmul_(block_grad, block_grad, value=1.0 - beta2) + denom = variance.sqrt().div_(bias_correction2_sqrt).add_(self.eps) + weight.addcdiv_(moment, denom, value=-step_size) + model_shard.narrow(0, start, numel).copy_(weight) + quantized_moment, moment_scale = _quantize_symmetric(moment) + quantized_variance, variance_scale = _quantize_positive(variance) + block_moment_q.copy_(quantized_moment) + block_variance_q.copy_(quantized_variance) + moment_scales[block_index].copy_(moment_scale) + variance_scales[block_index].copy_(variance_scale) def _quantize_symmetric(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -469,7 +550,7 @@ def _dequantize_symmetric(quantized: torch.Tensor, scale: torch.Tensor) -> torch def _quantize_positive(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """Quantize a non-negative FP32 tensor to uint8 with one bucket-level scale.""" + """Quantize one non-negative FP32 block to uint8.""" if tensor.numel() == 0: return tensor.to(dtype=torch.uint8), torch.ones((), device=tensor.device, dtype=torch.float32) @@ -488,3 +569,7 @@ def _cpu_clone(value: torch.Tensor | None) -> torch.Tensor | None: """Return an independent CPU copy of an optional quantized-state tensor.""" return None if value is None else value.detach().to(device="cpu").clone() + + +def _ceil_div(numerator: int, denominator: int) -> int: + return (numerator + denominator - 1) // denominator diff --git a/tests/test_adamw_4bit_cpu.py b/tests/test_adamw_4bit_cpu.py new file mode 100644 index 00000000..d340ab1b --- /dev/null +++ b/tests/test_adamw_4bit_cpu.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +import copy +from pathlib import Path + +import pytest +import torch +from click.testing import CliRunner + +from areno.api.trainer_config import TrainerConfig +from areno.cli.train import train_command +from areno.engine.config import OptimizerConfig +from areno.engine.modeling import build_optimizer +from areno.engine.optim import AdamW4bit, AdamW8bit, AdamWFP32Master +from areno.engine.optim.adamw_4bit import ( + _quantize_positive_4bit, + _quantize_signed_4bit, + _unpack_positive_4bit, + _unpack_signed_4bit, +) + + +def _optimizer(param: torch.nn.Parameter, *, block_size: int = 128) -> AdamW4bit: + return AdamW4bit( + [param], + lr=3.0e-4, + betas=(0.9, 0.99), + weight_decay=0.01, + bucket_numel=max(param.numel(), 1), + quant_block_size=block_size, + ) + + +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()) + + 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 + + eight_bit_param = torch.nn.Parameter(torch.zeros_like(param)) + eight_bit = AdamW8bit( + [eight_bit_param], + lr=3.0e-4, + betas=(0.9, 0.99), + weight_decay=0.01, + bucket_numel=param.numel(), + quant_block_size=128, + ) + eight_bit_param.grad = torch.ones_like(eight_bit_param) + 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_bytes * 0.6 + + +def test_adamw4bit_second_moment_mapping_excludes_zero() -> None: + values = torch.tensor([0.0, 1.0 / 16.0, 0.5, 1.0]) + + packed, scale = _quantize_positive_4bit(values) + restored = _unpack_positive_4bit(packed, values.numel(), scale) + + assert scale.item() == 1.0 + assert restored[0].item() == pytest.approx(1.0 / 16.0) + assert torch.all(restored > 0) + torch.testing.assert_close(restored[1:], values[1:]) + + +def test_adamw4bit_signed_quantizer_preserves_dynamic_map_points() -> None: + values = torch.tensor([-0.8875, -0.2125, -0.0055, 0.0, 0.0325, 0.4375, 1.0]) + + packed, scale = _quantize_signed_4bit(values) + restored = _unpack_signed_4bit(packed, values.numel(), scale) + + torch.testing.assert_close(restored, values, rtol=1.0e-6, atol=1.0e-6) + + +def test_adamw4bit_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.step() + checkpoint = copy.deepcopy(first.state_dict()) + + restored_param = torch.nn.Parameter(first_param.detach().clone()) + restored = _optimizer(restored_param) + 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() + 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"] == 1 + + +def test_adamw4bit_disk_offload_preserves_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.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.step() + reference.step() + + torch.testing.assert_close(candidate_param, reference_param, 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: + 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) + reference = AdamWFP32Master( + [reference_param], + lr=3.0e-4, + betas=(0.9, 0.99), + weight_decay=0.01, + bucket_numel=initial.numel(), + ) + + 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.step() + reference.step() + + torch.testing.assert_close(quantized_param, reference_param, rtol=3.0e-3, atol=3.0e-3) + + +def test_adamw4bit_nonfinite_gradient_skips_only_affected_block() -> None: + parameter = torch.nn.Parameter(torch.zeros(256)) + optimizer = _optimizer(parameter, block_size=128) + gradient = torch.ones_like(parameter) + gradient[4] = torch.inf + parameter.grad = gradient + + optimizer.step() + + torch.testing.assert_close(parameter[:128], torch.zeros(128)) + assert torch.all(parameter[128:] < 0) + + +def test_optimizer_config_rejects_multiple_low_bit_modes() -> None: + with pytest.raises(ValueError, match="mutually exclusive"): + OptimizerConfig(adam_4bit=True, adam_8bit=True) + + +def test_build_optimizer_selects_adamw4bit() -> None: + class Context: + dp_rank = 0 + dp_size = 1 + dp_group = None + + param = torch.nn.Parameter(torch.ones(4)) + optimizer = build_optimizer([param], OptimizerConfig(adam_4bit=True), Context()) + + assert isinstance(optimizer, AdamW4bit) + + +def test_trainer_config_propagates_adamw4bit() -> None: + config = TrainerConfig( + algo="sft", + ckpt="unused", + dataset_path="unused", + backend="cuda", + adam_4bit=True, + ) + + assert config.optimizer_config()["adam_4bit"] is True + assert config.cuda_config().optimizer["adam_4bit"] is True + + +def test_train_cli_exposes_adamw4bit_flag() -> None: + result = CliRunner().invoke(train_command, ["--help"]) + + assert result.exit_code == 0 + assert "--adam-4bit" in result.output diff --git a/tests/test_adamw_8bit_blockwise_cpu.py b/tests/test_adamw_8bit_blockwise_cpu.py new file mode 100644 index 00000000..c1a4fd9d --- /dev/null +++ b/tests/test_adamw_8bit_blockwise_cpu.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import torch + +from areno.engine.optim import AdamW8bit +from areno.engine.optim.adamw_8bit import _dequantize_positive, _quantize_positive + + +def test_adamw8bit_uses_parameter_local_block_scales() -> None: + first = torch.nn.Parameter(torch.zeros(9)) + second = torch.nn.Parameter(torch.zeros(3)) + optimizer = AdamW8bit( + [first, second], + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=32, + quant_block_size=4, + ) + first.grad = torch.tensor([1000.0, 1.0, -1.0, 0.5, 0.25, -0.5, 0.75, -0.25, 0.125]) + second.grad = torch.tensor([0.01, -0.02, 0.03]) + + optimizer.step() + + state = optimizer.state_dict()["state"][0] + assert state["exp_avg_scale"].shape == (4,) + assert state["exp_avg_sq_scale"].shape == (4,) + assert state["exp_avg_scale"][0] > 1000 * state["exp_avg_scale"][-1] + assert state["exp_avg_sq_scale"][0] > 1000 * state["exp_avg_sq_scale"][-1] + torch.testing.assert_close(second, torch.tensor([-1.0e-3, 1.0e-3, -1.0e-3]), atol=1.0e-6, rtol=0.0) + + +def test_adamw8bit_second_moment_uses_full_linear_range_per_block() -> None: + values = torch.tensor([0.0, 1.0 / 255.0, 128.0 / 255.0, 1.0]) + + quantized, scale = _quantize_positive(values) + restored = _dequantize_positive(quantized, scale) + + torch.testing.assert_close(restored, values) + + +def test_adamw8bit_same_lr_does_not_amplify_constant_gradient_step() -> None: + parameter = torch.nn.Parameter(torch.ones(4096)) + optimizer = AdamW8bit( + [parameter], + lr=1.0e-3, + betas=(0.0, 0.0), + weight_decay=0.0, + quant_block_size=2048, + ) + parameter.grad = torch.full_like(parameter, 0.25) + + optimizer.step() + + torch.testing.assert_close(parameter, torch.full_like(parameter, 0.999), rtol=0.0, atol=1.0e-6) diff --git a/tests/test_fp32_master_optimizer_cuda.py b/tests/test_fp32_master_optimizer_cuda.py index 41458857..5e7853fe 100644 --- a/tests/test_fp32_master_optimizer_cuda.py +++ b/tests/test_fp32_master_optimizer_cuda.py @@ -5,11 +5,36 @@ import pytest import torch -from areno.engine.optim import AdamW8bit, AdamWFP32Master +from areno.engine.optim import AdamW4bit, AdamW8bit, AdamWFP32Master pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("optimizer_cls", [AdamW4bit, AdamW8bit]) +def test_fused_low_bit_adamw_matches_cpu_reference(optimizer_cls) -> None: + initial = torch.linspace(-1.0, 1.0, 4099).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": 8192, + "quant_block_size": 128 if optimizer_cls is AdamW4bit else 2048, + } + cuda_optimizer = optimizer_cls([cuda_param], **kwargs) + cpu_optimizer = optimizer_cls([cpu_param], **kwargs) + + for step in range(4): + gradient = torch.sin(torch.linspace(-2.0, 2.0, initial.numel()) + step * 0.2).to(torch.bfloat16) + cuda_param.grad = gradient.cuda() + cpu_param.grad = gradient.clone() + cuda_optimizer.step() + cpu_optimizer.step() + + torch.testing.assert_close(cuda_param.cpu(), cpu_param, rtol=0.0, atol=2.0e-3) + + 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 a60d8dc5475ef55a1f4c23537aa2516cfedd0ed6 Mon Sep 17 00:00:00 2001 From: xsuler Date: Tue, 1 Sep 2026 12:43:13 +0800 Subject: [PATCH 2/2] docs(optimizer): document AdamW4bit usage --- docs/cli/training.rst | 7 +- docs/index.rst | 1 + docs/reference/adamw-4bit.rst | 137 ++++++++++++++++++++++++++++++++++ docs/reference/index.rst | 1 + 4 files changed, 145 insertions(+), 1 deletion(-) create mode 100644 docs/reference/adamw-4bit.rst diff --git a/docs/cli/training.rst b/docs/cli/training.rst index 5727d42e..5034c984 100644 --- a/docs/cli/training.rst +++ b/docs/cli/training.rst @@ -346,7 +346,7 @@ in its description; flags for other algorithms are ignored. persistent writable raw-mmap files, copying each bucket back only when the next optimizer step needs it. Default: ``none``. This option is supported only by the CUDA backend and applies to both FP32-master AdamW and - ``--adam-8bit``. + ``--adam-8bit`` or ``--adam-4bit``. ``--optimizer-state-offload-dir DIRECTORY`` Required when ``--optimizer-state-offload disk`` is selected. Use a fast @@ -383,6 +383,11 @@ in its description; flags for other algorithms are ignored. Use 8-bit Adam moment states instead of FP32 Adam states. Supported by both native 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 + :doc:`../reference/adamw-4bit` for complete usage examples. + ``--unfreeze-mm-tower`` Train recognized vision/audio encoder tower parameters. Towers are frozen by default. diff --git a/docs/index.rst b/docs/index.rst index 7cc230ed..ae008795 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -46,6 +46,7 @@ AReno documentation :caption: Reference CLI Reference + 4-bit AdamW SDK Reference Supported Models diff --git a/docs/reference/adamw-4bit.rst b/docs/reference/adamw-4bit.rst new file mode 100644 index 00000000..bb81e33c --- /dev/null +++ b/docs/reference/adamw-4bit.rst @@ -0,0 +1,137 @@ +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. + +Command line +------------ + +Add ``--adam-4bit`` to any CUDA ``areno train`` command: + +.. code-block:: bash + + areno train \ + --ckpt Qwen/Qwen3-0.6B \ + --dataset-path gsm8k:main \ + --dataset-loader-fn examples/math/dataset_loader.py \ + --reward-fn-path examples/math/math_verify_reward.py \ + --algo gspo \ + --world-size 1 \ + --tp-size 1 \ + --batch-size 2 \ + --n-samples 2 \ + --mini-bs 1 \ + --adam-4bit + +Keep the existing learning-rate and Adam settings unless an experiment calls +for different values: + +.. code-block:: bash + + areno train \ + ... \ + --adam-4bit \ + --lr 1e-6 \ + --adam-beta1 0.9 \ + --adam-beta2 0.999 + +``--adam-4bit`` and ``--adam-8bit`` are mutually exclusive. Passing both +causes configuration validation to fail before training starts. + +Optimizer-state offload +----------------------- + +The 4-bit optimizer supports the same CUDA optimizer-state residency options +as the other CUDA AdamW implementations. + +Keep state on the training device: + +.. code-block:: bash + + areno train ... --adam-4bit + +Offload state to CPU memory between train calls: + +.. code-block:: bash + + areno train ... \ + --adam-4bit \ + --optimizer-state-offload cpu + +Stream state through a local disk directory: + +.. code-block:: bash + + areno train ... \ + --adam-4bit \ + --optimizer-state-offload disk \ + --optimizer-state-offload-dir /local/nvme/areno-optimizer \ + --optimizer-state-offload-batch-size 1 + +Use a fast local NVMe path for disk offload. Runtime mmap files are temporary +scratch files and are not restartable checkpoints. + +Trainer configuration +--------------------- + +Set ``adam_4bit=True`` on a CLI trainer configuration: + +.. code-block:: python + + from areno.api.trainer_config import PolicyTrainerConfig + + config = PolicyTrainerConfig( + algo="gspo", + ckpt="Qwen/Qwen3-0.6B", + dataset_path="gsm8k:main", + dataset_loader_fn="examples/math/dataset_loader.py", + reward_fn_path="examples/math/math_verify_reward.py", + backend="cuda", + world_size=1, + tp_size=1, + adam_4bit=True, + ) + +For the lower-level Trainer SDK, pass the optimizer option through +``CudaConfig``: + +.. code-block:: python + + from areno import Trainer + from areno.api import CUDA, CudaConfig + + trainer = Trainer( + world_size=1, + model_path="Qwen/Qwen3-0.6B", + backend_type=CUDA, + custom_config=CudaConfig( + tp_size=1, + optimizer={ + "adam_4bit": True, + "lr": 1e-6, + "betas": (0.9, 0.999), + "weight_decay": 0.01, + }, + ), + ) + +The equivalent engine-level setting is +``OptimizerConfig(adam_4bit=True)``. + +Requirements and errors +----------------------- + +* Use the CUDA backend. MLX configurations reject ``adam_4bit=True``. +* Do not enable ``adam_8bit`` at the same time. +* Rebuild or reinstall AReno after switching to a revision that adds the + 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``. + +Confirm that the option is available with: + +.. code-block:: bash + + areno train --help | grep adam-4bit diff --git a/docs/reference/index.rst b/docs/reference/index.rst index 32f1540e..951a5c4a 100644 --- a/docs/reference/index.rst +++ b/docs/reference/index.rst @@ -13,5 +13,6 @@ tasks, use :doc:`/cookbook/math-rlvr`. Reference pages: * :doc:`cli` +* :doc:`adamw-4bit` * :doc:`/sdk/trainer` * :doc:`/models/supported`