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/extension.cpp b/areno/accel/csrc/extension.cpp index 20d7f689..d85cf894 100644 --- a/areno/accel/csrc/extension.cpp +++ b/areno/accel/csrc/extension.cpp @@ -192,8 +192,10 @@ 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 scale_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, double beta2, @@ -202,6 +204,33 @@ void areno_adamw_4bit_step_cuda( double eps, double step_size, double bias_correction2_sqrt); +void areno_adamw_4bit_factored_stats_cuda( + torch::Tensor grad, + torch::Tensor factor_sums, + torch::Tensor invalid, + int64_t parameter_shard_start, + 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 factors, + torch::Tensor row_mean, + torch::Tensor invalid, + 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 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, @@ -235,6 +264,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_factored_stats", + &areno_adamw_4bit_factored_stats_cuda, + "ARENO factored AdamW4bit statistics pass"); + m.def( + "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 85e38f31..0a6498a6 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -111,8 +111,10 @@ __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 scale_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, float effective_lr, @@ -120,10 +122,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; @@ -134,20 +141,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 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 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); - 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; } @@ -158,36 +165,67 @@ __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[scale_offset + blockIdx.x] = new_moment_scale; - exp_avg_sq_scale[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 + 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); + 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); - 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; @@ -195,9 +233,148 @@ __global__ void adamw_4bit_kernel( } __syncthreads(); if ((tid & 1) == 0 && local_index < numel) { - const int64_t byte_index = packed_offset + (local_index >> 1); + 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_factored_stats_kernel( + const grad_t* grad, + float* factor_sums, + int* invalid, + int64_t numel, + int64_t parameter_shard_start, + 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; + const float gradient = load_grad(grad, local_index); + const float squared = gradient * gradient; + if (!isfinite(squared)) { + atomicExch(invalid, 1); + continue; + } + atomicAdd(factor_sums + parameter_index / columns, squared); + atomicAdd(factor_sums + rows + parameter_index % columns, squared); + } +} + +template +__global__ void adamw_4bit_factored_step_kernel( + model_t* model, + const grad_t* grad, + uint8_t* exp_avg_q, + float* exp_avg_scale, + const float* factors, + const float* row_mean, + const int* invalid, + int64_t numel, + int64_t moment_packed_offset, + int64_t moment_scale_offset, + int64_t parameter_shard_start, + int64_t rows, + int64_t columns, + float beta1, + float effective_lr, + float weight_decay, + float eps, + float step_size, + float bias_correction2_sqrt) { + 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; + __shared__ float old_moment_scale; + __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; + old_moment_scale = exp_avg_scale[moment_scale_offset + blockIdx.x]; + } + __syncthreads(); + 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 + moment_packed_offset, local_index); + float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; + 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; + 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); + } + local_moment_max = fabsf(moment); + } + 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 (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 + moment_packed_offset, local_index); + float moment = kSigned4bitDynamicMap[moment_code] * old_moment_scale; + 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; + 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)); + store_model(model, local_index, updated_weight); + } else { + moment_codes[tid] = 7; + } + __syncthreads(); + if ((tid & 1) == 0 && local_index < numel) { + 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); } } @@ -376,8 +553,10 @@ 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 scale_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, float beta2, @@ -388,7 +567,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(), @@ -396,8 +576,10 @@ void launch_adamw_4bit( exp_avg_sq_q.data_ptr(), exp_avg_sq_scale.data_ptr(), model.numel(), - packed_offset, - scale_offset, + moment_packed_offset, + moment_scale_offset, + variance_packed_offset, + variance_scale_offset, beta1, beta2, effective_lr, @@ -408,6 +590,78 @@ void launch_adamw_4bit( C10_CUDA_KERNEL_LAUNCH_CHECK(); } +template +void launch_adamw_4bit_factored_stats( + torch::Tensor grad, + torch::Tensor factor_sums, + torch::Tensor invalid, + int64_t parameter_shard_start, + 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_factored_stats_kernel<<>>( + grad.data_ptr(), + factor_sums.data_ptr(), + invalid.data_ptr(), + grad.numel(), + parameter_shard_start, + rows, + columns); + C10_CUDA_KERNEL_LAUNCH_CHECK(); +} + +template +void launch_adamw_4bit_factored_step( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg_q, + torch::Tensor exp_avg_scale, + torch::Tensor factors, + torch::Tensor row_mean, + torch::Tensor invalid, + 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 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(); + 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(), + factors.data_ptr(), + row_mean.data_ptr(), + invalid.data_ptr(), + model.numel(), + moment_packed_offset, + moment_scale_offset, + parameter_shard_start, + rows, + columns, + beta1, + 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, @@ -658,8 +912,10 @@ 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 scale_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, double beta2, @@ -682,7 +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, 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) { @@ -699,6 +956,101 @@ void areno_adamw_4bit_step_cuda( #undef LAUNCH_ADAMW4 } +void areno_adamw_4bit_factored_stats_cuda( + torch::Tensor grad, + torch::Tensor factor_sums, + torch::Tensor invalid, + int64_t parameter_shard_start, + int64_t rows, + int64_t columns) { + c10::cuda::CUDAGuard guard(grad.device()); + TORCH_CHECK( + grad.is_cuda() && factor_sums.is_cuda() && invalid.is_cuda(), + "AdamW4bit factored statistics tensors must be CUDA tensors"); + TORCH_CHECK( + 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_factored_stats( + grad, factor_sums, invalid, parameter_shard_start, rows, columns); + } else if (grad.scalar_type() == at::kFloat) { + launch_adamw_4bit_factored_stats(grad, factor_sums, invalid, parameter_shard_start, rows, columns); + } else { + TORCH_CHECK(false, "AdamW4bit factored gradient must be bfloat16 or float32"); + } +} + +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 factors, + torch::Tensor row_mean, + torch::Tensor invalid, + 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 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() && 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_FACTORED(at::BFloat16, at::BFloat16); + } else if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW4_FACTORED(at::BFloat16, float); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kBFloat16) { + LAUNCH_ADAMW4_FACTORED(float, at::BFloat16); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW4_FACTORED(float, float); + } else { + TORCH_CHECK(false, "AdamW4bit factored model and gradient must be bfloat16 or float32"); + } +#undef LAUNCH_ADAMW4_FACTORED +} + 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..c147a8f9 100644 --- a/areno/accel/optimizer.py +++ b/areno/accel/optimizer.py @@ -205,8 +205,10 @@ def areno_adamw_4bit_step( exp_avg_sq_q: torch.Tensor, exp_avg_sq_scale: torch.Tensor, *, - packed_offset: int, - scale_offset: int, + moment_packed_offset: int, + moment_scale_offset: int, + variance_packed_offset: int, + variance_scale_offset: int, quant_block_size: int, beta1: float, beta2: float, @@ -237,14 +239,14 @@ 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 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_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(): + raise ValueError("AdamW4bit second-moment scale slice is out of bounds") extension().areno_adamw_4bit_step( model, grad, @@ -252,8 +254,10 @@ def areno_adamw_4bit_step( exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, - packed_offset, - scale_offset, + moment_packed_offset, + moment_scale_offset, + variance_packed_offset, + variance_scale_offset, quant_block_size, beta1, beta2, @@ -265,8 +269,134 @@ def areno_adamw_4bit_step( ) +@torch._dynamo.disable +@torch.no_grad() +def areno_adamw_4bit_factored_stats( + grad: torch.Tensor, + factor_sums: torch.Tensor, + invalid: torch.Tensor, + *, + parameter_shard_start: int, + rows: int, + columns: int, +) -> None: + """Accumulate matrix row/column gradient-square sums.""" + + tensors = (grad, factor_sums, invalid) + if any(not tensor.is_cuda for tensor in 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 factored statistics require tensors on one device") + if grad.dtype not in {torch.bfloat16, torch.float32}: + 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 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, + factor_sums, + invalid, + parameter_shard_start, + rows, + columns, + ) + + +@torch._dynamo.disable +@torch.no_grad() +def areno_adamw_4bit_factored_step( + model: torch.Tensor, + grad: torch.Tensor, + exp_avg_q: torch.Tensor, + exp_avg_scale: torch.Tensor, + factors: torch.Tensor, + row_mean: torch.Tensor, + invalid: torch.Tensor, + *, + moment_packed_offset: int, + moment_scale_offset: int, + parameter_shard_start: int, + quant_block_size: int, + rows: int, + columns: int, + beta1: float, + effective_lr: float, + weight_decay: float, + eps: float, + step_size: float, + bias_correction2_sqrt: float, +) -> None: + """Update packed momentum using an Adafactor-style variance estimate.""" + + tensors = ( + model, + grad, + exp_avg_q, + exp_avg_scale, + factors, + row_mean, + invalid, + ) + if any(not tensor.is_cuda for tensor in 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 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 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 factored update requires one int32 validity flag") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("fused AdamW4bit factored update requires contiguous tensors") + if model.numel() != grad.numel(): + 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 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 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, + 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, + ) + + __all__ = [ "areno_adamw_4bit_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 943bca45..3313c04e 100644 --- a/areno/engine/optim/adamw_4bit.py +++ b/areno/engine/optim/adamw_4bit.py @@ -1,9 +1,11 @@ -"""Packed block-wise 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``. Two codes are packed in each -byte and scales are stored per parameter-local block. +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 @@ -14,10 +16,10 @@ 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 = 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 = ( @@ -41,16 +43,21 @@ 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 packed 4-bit momentum and factored matrix variance. + + First-moment quantization blocks restart at every parameter shard. For a + 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/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 @@ -83,6 +90,9 @@ def __init__( dp_group=dp_group, quant_block_size=quant_block_size, ) + self._factored_second_moments: 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 +110,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["factored_second_moments"] = [ + None + 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() return payload @torch.no_grad() @@ -114,6 +132,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_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 @@ -124,16 +149,64 @@ 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) + 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_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) + 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, variance_packed_numel) + state.exp_avg_sq_scale = _load_tensor( + saved, "exp_avg_sq_scale", device, torch.float32, variance_scale_numel + ) + for parameter, saved_factors_for_parameter in zip(self.model_params, saved_factors, strict=True): + if parameter.ndim < 2: + if saved_factors_for_parameter is not None: + raise ValueError("AdamW4bit checkpoint has factored state for a one-dimensional parameter") + continue + if saved_factors_for_parameter is None: + self._factored_second_moments[id(parameter)] = None + continue + 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 factored state.""" + + super().clear_state() + 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 factored state on CPU.""" + + super().offload_state(mode=mode, directory=directory, batch_size=batch_size) + 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 factored state to ``device``.""" + + super().onload_state(device) + 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: @@ -146,67 +219,165 @@ 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) + 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_scale = torch.ones(scale_numel, device=device, dtype=torch.float32) + 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(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]) + 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_scale": (torch.float32, (scale_numel,)), - "exp_avg_sq_q": (torch.uint8, (packed_numel,)), - "exp_avg_sq_scale": (torch.float32, (scale_numel,)), + "exp_avg_q": (torch.uint8, (moment_packed_numel,)), + "exp_avg_scale": (torch.float32, (moment_scale_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_bucket_8bit(self, bucket: _MasterBucket, state) -> None: - """Update a bucket while materializing at most one FP32 block per moment.""" + def step(self, closure=None): + """Update one complete parameter at a time and release ready buckets.""" + + if closure is not None: + with torch.enable_grad(): + closure() + 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 - 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: + started_buckets: set[int] = set() + completed_buckets: set[int] = set() + + for parameter in self.model_params: + parameter_layouts = layouts.get(id(parameter)) + if not parameter_layouts: 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() + 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), + factor_sums, + invalid, + ) + 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] + if index not in started_buckets: + self._ensure_bucket_state(bucket, state) + 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 + + @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( @@ -215,8 +386,10 @@ def _step_param_ref_4bit( ref: _ParamRef, grad: torch.Tensor, state, - packed_offset: int, - scale_offset: int, + moment_packed_offset: int, + moment_scale_offset: int, + variance_packed_offset: int, + variance_scale_offset: int, beta1: float, beta2: float, effective_lr: float, @@ -240,8 +413,10 @@ def _step_param_ref_4bit( state.exp_avg_scale, state.exp_avg_sq_q, state.exp_avg_sq_scale, - packed_offset=packed_offset, - scale_offset=scale_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, beta2=beta2, @@ -254,18 +429,20 @@ 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 - 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), + state.exp_avg_q.narrow(0, moment_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), + state.exp_avg_sq_q.narrow(0, variance_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) @@ -287,23 +464,205 @@ 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, 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, variance_byte_start, byte_count).copy_(variance_q) + state.exp_avg_sq_scale[variance_scale_index].copy_(variance_scale) + + @torch.no_grad() + def _factored_variance_statistics( + self, + bucket: _MasterBucket, + ref: _ParamRef, + grad: torch.Tensor | None, + factor_sums: torch.Tensor, + invalid: torch.Tensor, + ) -> None: + """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_factored_stats + + areno_adamw_4bit_factored_stats( + grad_shard.contiguous(), + factor_sums, + invalid, + parameter_shard_start=parameter_shard_start, + rows=ref.model_param.shape[0], + columns=ref.model_param.numel() // ref.model_param.shape[0], + ) + return + + 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) + flat_start = parameter_shard_start + start + gradient = grad_shard.narrow(0, start, count).to(dtype=torch.float32) + squared = gradient.square() + if not bool(torch.isfinite(squared).all()): + invalid.fill_(1) + return + 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_factored( + self, + bucket: _MasterBucket, + ref: _ParamRef, + grad: torch.Tensor, + state, + moment_packed_offset: int, + moment_scale_offset: int, + updated_factors: torch.Tensor, + row_mean: torch.Tensor, + invalid: torch.Tensor, + beta1: float, + effective_lr: float, + step_size: float, + bias_correction2_sqrt: float, + ) -> None: + """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 + 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_factored_step + + areno_adamw_4bit_factored_step( + model_shard, + grad_shard.contiguous(), + state.exp_avg_q, + state.exp_avg_scale, + updated_factors, + row_mean, + invalid, + 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, + effective_lr=effective_lr, + weight_decay=self.weight_decay, + eps=self.eps, + step_size=step_size, + bias_correction2_sqrt=bias_correction2_sqrt, + ) + return + + 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 = moment_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 + 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) + 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 + moment_q, moment_scale = _quantize_signed_4bit(moment) + 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[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 + state.exp_avg_scale[moment_scale_index].copy_(moment_scale) + + def _ensure_factored_second_moment(self, parameter: torch.nn.Parameter) -> torch.Tensor: + """Materialize row/column second-moment statistics for one parameter.""" + + key = id(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._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 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, 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, scale_offset - packed_offset += (ref.shard_numel + 1) // 2 - scale_offset += (ref.shard_numel + self.quant_block_size - 1) // self.quant_block_size + 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 ref.model_param.ndim >= 2: + 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 +672,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._factored_second_moments.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._factored_second_moments.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 +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 _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 3aac40af..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,43 +156,59 @@ 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 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..326bb6ef 100644 --- a/docs/cli/training.rst +++ b/docs/cli/training.rst @@ -398,8 +398,10 @@ 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 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 bb81e33c..dcbf6094 100644 --- a/docs/reference/adamw-4bit.rst +++ b/docs/reference/adamw-4bit.rst @@ -2,8 +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. 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 ------------ @@ -129,6 +151,12 @@ 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 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``. 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..dcff04d8 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,6 +16,7 @@ from areno.engine.modeling import build_optimizer from areno.engine.optim import AdamW4bit, AdamW8bit, AdamWFP32Master from areno.engine.optim.adamw_4bit import ( + _factored_state_numel_for_parameter, _quantize_positive_4bit, _quantize_signed_4bit, _unpack_positive_4bit, @@ -20,47 +24,158 @@ ) -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, ) -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 _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 _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}", + 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() + gradient_dtype = optimizer.buckets[0].grad_shard.dtype + optimizer.step() + output_queue.put( + ( + rank, + parameter.detach().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, + ) + ) + finally: + dist.destroy_process_group() + + +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, + ) + 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_initializes_and_releases_state_per_parameter() -> 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, ) - assert optimizer.persistent_moment_bytes() <= eight_bit_bytes * 0.6 + 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() + + 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: @@ -84,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"] == 1 + 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_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()) -def test_adamw4bit_disk_offload_preserves_update(tmp_path: Path) -> None: + 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, @@ -144,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) @@ -165,6 +342,126 @@ def test_adamw4bit_nonfinite_gradient_skips_only_affected_block() -> None: assert torch.all(parameter[128:] < 0) +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 + + optimizer.step() + + torch.testing.assert_close(parameter, torch.zeros_like(parameter)) + torch.testing.assert_close( + optimizer._factored_second_moments[id(parameter)], + torch.zeros(parameter.shape[0] + parameter.shape[1]), + ) + + +@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() + + torch.testing.assert_close( + optimizer._factored_second_moments[id(parameter)], + _expected_first_step_factors(gradient), + ) + + +def test_factored_statistics_combine_partial_parameter_chunks() -> None: + parameter = torch.nn.Parameter(torch.zeros(2048, 2049)) + optimizer = _optimizer(parameter, bucket_numel=1) + + 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_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( + [matrix, vector], + lr=3.0e-4, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=1024, + quant_block_size=128, + ) + matrix.grad = torch.ones_like(matrix) + vector.grad = torch.ones_like(vector) + + 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_factored_memory_metrics_match_resident_tensors() -> None: + parameter = torch.nn.Parameter(torch.zeros(1024, 1024)) + optimizer = _optimizer(parameter) + parameter.grad = torch.ones_like(parameter) + optimizer.step() + + metrics = optimizer.state_memory_metrics() + assert metrics["total_bytes"] == optimizer.persistent_moment_bytes() + assert metrics["quantized_state_bytes"] == parameter.numel() // 2 + assert metrics["scale_metadata_bytes"] == (parameter.numel() // 128 + 2048) * 4 + + +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._factored_second_moments[id(parameter)] is not None + optimizer.clear_state() + + 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_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_factored_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_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 + expected_factors = _expected_first_step_factors(averaged_gradient) + assert rank0_model == rank1_model + 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) + + def test_optimizer_config_rejects_multiple_low_bit_modes() -> None: with pytest.raises(ValueError, match="mutually exclusive"): OptimizerConfig(adam_4bit=True, adam_8bit=True) @@ -176,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) 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)