From fda56b724162be9396e075375aaa38c518adf8d6 Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 14:21:14 +0800 Subject: [PATCH 1/3] feat(optimizer): complete dynamic AdamW8bit state routing --- areno/accel/__init__.py | 8 +- areno/accel/csrc/extension.cpp | 15 + areno/accel/csrc/optimizer.cu | 179 +++++++- areno/accel/optimizer.py | 74 +++- areno/api/backend/mlx/backend.py | 5 +- areno/api/backend/mlx/optimizer.py | 122 ++++-- areno/api/backend/mlx/provider.py | 37 ++ areno/engine/layers/vocab.py | 3 + areno/engine/optim/__init__.py | 4 +- areno/engine/optim/adamw_4bit.py | 15 + areno/engine/optim/adamw_8bit.py | 555 +++++++++++++++++++++---- areno/engine/optim/dynamic_quant.py | 61 +++ areno/engine/training.py | 8 + docs/cli/training.rst | 18 +- docs/getting-started/mlx.rst | 8 +- tests/test_adamw_8bit_blockwise_cpu.py | 275 +++++++++++- tests/test_mlx_training_cpu.py | 60 +++ 17 files changed, 1302 insertions(+), 145 deletions(-) create mode 100644 areno/engine/optim/dynamic_quant.py diff --git a/areno/accel/__init__.py b/areno/accel/__init__.py index bcd806c0..590dad00 100644 --- a/areno/accel/__init__.py +++ b/areno/accel/__init__.py @@ -30,7 +30,12 @@ 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_4bit_step, areno_adamw_8bit_step, areno_adamw_fp32_master_step +from areno.accel.optimizer import ( + areno_adamw_4bit_step, + areno_adamw_8bit_step, + areno_adamw_fp32_master_step, + areno_adamw_fp32_state_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 @@ -53,6 +58,7 @@ "areno_adamw_4bit_step", "areno_adamw_8bit_step", "areno_adamw_fp32_master_step", + "areno_adamw_fp32_state_step", "areno_optional_scale_rmsnorm", "areno_rmsnorm", "areno_rmsnorm_silu_gate", diff --git a/areno/accel/csrc/extension.cpp b/areno/accel/csrc/extension.cpp index 494233b8..20d7f689 100644 --- a/areno/accel/csrc/extension.cpp +++ b/areno/accel/csrc/extension.cpp @@ -209,6 +209,8 @@ void areno_adamw_8bit_step_cuda( torch::Tensor exp_avg_scale, torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, + torch::Tensor signed_codebook, + torch::Tensor unsigned_codebook, int64_t quant_block_size, double beta1, double beta2, @@ -217,11 +219,24 @@ void areno_adamw_8bit_step_cuda( double eps, double step_size, double bias_correction2_sqrt); +void areno_adamw_fp32_state_step_cuda( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg, + torch::Tensor exp_avg_sq, + 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_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"); 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 4c885b56..85e38f31 100644 --- a/areno/accel/csrc/optimizer.cu +++ b/areno/accel/csrc/optimizer.cu @@ -31,6 +31,25 @@ __device__ __forceinline__ uint8_t nearest_signed_dynamic_code(float normalized) return best; } +__device__ __forceinline__ uint8_t nearest_dynamic_code(float value, const float* codebook) { + int lower = 0; + int upper = 255; + while (lower < upper) { + const int middle = (lower + upper) >> 1; + if (codebook[middle] < value) { + lower = middle + 1; + } else { + upper = middle; + } + } + if (lower == 0) { + return 0; + } + const float left_distance = fabsf(value - codebook[lower - 1]); + const float right_distance = fabsf(codebook[lower] - value); + return static_cast(left_distance <= right_distance ? lower - 1 : lower); +} + __device__ __forceinline__ float adamw_update( float master, float grad, @@ -397,6 +416,8 @@ __global__ void adamw_8bit_blockwise_kernel( float* exp_avg_scale, uint8_t* exp_avg_sq_q, float* exp_avg_sq_scale, + const float* signed_codebook, + const float* unsigned_codebook, int64_t numel, int64_t quant_block_size, float beta1, @@ -410,6 +431,7 @@ __global__ void adamw_8bit_blockwise_kernel( constexpr int max_warps = 8; __shared__ float warp_moment_maxima[max_warps]; __shared__ float warp_variance_maxima[max_warps]; + __shared__ int invalid_block; const int64_t block_start = static_cast(blockIdx.x) * quant_block_size; const int64_t remaining = numel - block_start; @@ -418,12 +440,16 @@ __global__ void adamw_8bit_blockwise_kernel( const float old_variance_scale = exp_avg_sq_scale[blockIdx.x]; float local_moment_max = 0.0f; float local_variance_max = 0.0f; + if (threadIdx.x == 0) { + invalid_block = 0; + } + __syncthreads(); 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 moment = signed_codebook[exp_avg_q[index]] * old_moment_scale; + float variance = unsigned_codebook[exp_avg_sq_q[index]] * old_variance_scale; float weight = load_model(model, index); weight = adamw_update( weight, @@ -437,7 +463,9 @@ __global__ void adamw_8bit_blockwise_kernel( eps, step_size, bias_correction2_sqrt); - store_model(model, index, weight); + if (!isfinite(gradient) || !isfinite(moment) || !isfinite(variance) || !isfinite(weight)) { + atomicExch(&invalid_block, 1); + } local_moment_max = fmaxf(local_moment_max, fabsf(moment)); local_variance_max = fmaxf(local_variance_max, variance); } @@ -454,6 +482,9 @@ __global__ void adamw_8bit_blockwise_kernel( warp_variance_maxima[warp] = local_variance_max; } __syncthreads(); + if (invalid_block != 0) { + return; + } if (warp == 0) { const int warp_count = blockDim.x / warp_size; @@ -466,8 +497,8 @@ __global__ void adamw_8bit_blockwise_kernel( 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); + exp_avg_scale[blockIdx.x] = block_moment_max; + exp_avg_sq_scale[blockIdx.x] = block_variance_max; } } __syncthreads(); @@ -476,14 +507,49 @@ __global__ void adamw_8bit_blockwise_kernel( 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)); + const float previous_moment = signed_codebook[exp_avg_q[index]] * old_moment_scale; + const float previous_variance = unsigned_codebook[exp_avg_sq_q[index]] * old_variance_scale; + const float moment = beta1 * previous_moment + (1.0f - beta1) * gradient; + const float variance = beta2 * previous_variance + (1.0f - beta2) * gradient * gradient; + float weight = load_model(model, index); + if (weight_decay != 0.0f) { + weight *= 1.0f - effective_lr * weight_decay; + } + const float denom = sqrtf(variance) / bias_correction2_sqrt + eps; + weight -= step_size * moment / denom; + const float normalized_moment = moment / fmaxf(new_moment_scale, 1.0e-30f); + const float normalized_variance = variance / fmaxf(new_variance_scale, 1.0e-30f); + exp_avg_q[index] = nearest_dynamic_code(normalized_moment, signed_codebook); + exp_avg_sq_q[index] = nearest_dynamic_code(normalized_variance, unsigned_codebook); + store_model(model, index, weight); + } +} + +template +__global__ void adamw_fp32_state_kernel( + model_t* model, + const grad_t* grad, + float* exp_avg, + float* exp_avg_sq, + int64_t numel, + float beta1, + float beta2, + float effective_lr, + float weight_decay, + float eps, + float step_size, + float bias_correction2_sqrt) { + for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < numel; + index += static_cast(blockDim.x) * gridDim.x) { + float moment = exp_avg[index]; + float variance = exp_avg_sq[index]; + const float weight = adamw_update( + load_model(model, index), load_grad(grad, index), moment, variance, beta1, beta2, + effective_lr, weight_decay, eps, step_size, bias_correction2_sqrt); + exp_avg[index] = moment; + exp_avg_sq[index] = variance; + store_model(model, index, weight); } } @@ -495,6 +561,8 @@ void launch_adamw_8bit( torch::Tensor exp_avg_scale, torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, + torch::Tensor signed_codebook, + torch::Tensor unsigned_codebook, int64_t quant_block_size, float beta1, float beta2, @@ -513,6 +581,8 @@ void launch_adamw_8bit( exp_avg_scale.data_ptr(), exp_avg_sq_q.data_ptr(), exp_avg_sq_scale.data_ptr(), + signed_codebook.data_ptr(), + unsigned_codebook.data_ptr(), model.numel(), quant_block_size, beta1, @@ -525,6 +595,28 @@ void launch_adamw_8bit( C10_CUDA_KERNEL_LAUNCH_CHECK(); } +template +void launch_adamw_fp32_state( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg, + torch::Tensor exp_avg_sq, + 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() + threads - 1) / threads); + const auto stream = at::cuda::getCurrentCUDAStream(); + adamw_fp32_state_kernel<<>>( + model.data_ptr(), grad.data_ptr(), exp_avg.data_ptr(), exp_avg_sq.data_ptr(), + model.numel(), 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( @@ -614,6 +706,8 @@ void areno_adamw_8bit_step_cuda( torch::Tensor exp_avg_scale, torch::Tensor exp_avg_sq_q, torch::Tensor exp_avg_sq_scale, + torch::Tensor signed_codebook, + torch::Tensor unsigned_codebook, int64_t quant_block_size, double beta1, double beta2, @@ -625,11 +719,13 @@ void areno_adamw_8bit_step_cuda( 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(), + exp_avg_sq_q.is_cuda() && exp_avg_sq_scale.is_cuda() && signed_codebook.is_cuda() && + unsigned_codebook.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(), + exp_avg_sq_q.is_contiguous() && exp_avg_sq_scale.is_contiguous() && signed_codebook.is_contiguous() && + unsigned_codebook.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"); @@ -638,11 +734,14 @@ void areno_adamw_8bit_step_cuda( 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"); + TORCH_CHECK(signed_codebook.numel() == 256, "signed dynamic codebook must have 256 entries"); + TORCH_CHECK(unsigned_codebook.numel() == 256, "unsigned dynamic codebook must have 256 entries"); -#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) +#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, signed_codebook, \ + unsigned_codebook, 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); @@ -657,3 +756,45 @@ void areno_adamw_8bit_step_cuda( } #undef LAUNCH_ADAMW8 } + +void areno_adamw_fp32_state_step_cuda( + torch::Tensor model, + torch::Tensor grad, + torch::Tensor exp_avg, + torch::Tensor exp_avg_sq, + 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.is_cuda() && exp_avg_sq.is_cuda(), + "all FP32-state AdamW inputs must be CUDA tensors"); + TORCH_CHECK( + model.is_contiguous() && grad.is_contiguous() && exp_avg.is_contiguous() && exp_avg_sq.is_contiguous(), + "all FP32-state AdamW inputs must be contiguous"); + TORCH_CHECK(model.numel() == grad.numel(), "model and gradient sizes must match"); + TORCH_CHECK(model.numel() == exp_avg.numel(), "model and first-moment sizes must match"); + TORCH_CHECK(model.numel() == exp_avg_sq.numel(), "model and second-moment sizes must match"); + +#define LAUNCH_ADAMW_FP32_STATE(MODEL_T, GRAD_T) \ + launch_adamw_fp32_state( \ + model, grad, exp_avg, exp_avg_sq, beta1, beta2, effective_lr, weight_decay, eps, step_size, \ + bias_correction2_sqrt) + + if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kBFloat16) { + LAUNCH_ADAMW_FP32_STATE(at::BFloat16, at::BFloat16); + } else if (model.scalar_type() == at::kBFloat16 && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW_FP32_STATE(at::BFloat16, float); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kBFloat16) { + LAUNCH_ADAMW_FP32_STATE(float, at::BFloat16); + } else if (model.scalar_type() == at::kFloat && grad.scalar_type() == at::kFloat) { + LAUNCH_ADAMW_FP32_STATE(float, float); + } else { + TORCH_CHECK(false, "model and gradient must be bfloat16 or float32"); + } +#undef LAUNCH_ADAMW_FP32_STATE +} diff --git a/areno/accel/optimizer.py b/areno/accel/optimizer.py index 0bba12ba..a339131b 100644 --- a/areno/accel/optimizer.py +++ b/areno/accel/optimizer.py @@ -76,6 +76,8 @@ def areno_adamw_8bit_step( exp_avg_scale: torch.Tensor, exp_avg_sq_q: torch.Tensor, exp_avg_sq_scale: torch.Tensor, + signed_codebook: torch.Tensor, + unsigned_codebook: torch.Tensor, *, block_size: int, beta1: float, @@ -88,7 +90,16 @@ def areno_adamw_8bit_step( ) -> 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) + tensors = ( + model, + grad, + exp_avg_q, + exp_avg_scale, + exp_avg_sq_q, + exp_avg_sq_scale, + signed_codebook, + unsigned_codebook, + ) 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:]): @@ -101,6 +112,10 @@ def areno_adamw_8bit_step( 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 signed_codebook.dtype != torch.float32 or unsigned_codebook.dtype != torch.float32: + raise TypeError("dynamic quantization codebooks must use float32") + if signed_codebook.numel() != 256 or unsigned_codebook.numel() != 256: + raise ValueError("dynamic quantization codebooks must contain 256 entries") 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(): @@ -119,6 +134,8 @@ def areno_adamw_8bit_step( exp_avg_scale, exp_avg_sq_q, exp_avg_sq_scale, + signed_codebook, + unsigned_codebook, block_size, beta1, beta2, @@ -130,6 +147,54 @@ def areno_adamw_8bit_step( ) +@torch._dynamo.disable +@torch.no_grad() +def areno_adamw_fp32_state_step( + model: torch.Tensor, + grad: torch.Tensor, + exp_avg: torch.Tensor, + exp_avg_sq: torch.Tensor, + *, + beta1: float, + beta2: float, + effective_lr: float, + weight_decay: float, + eps: float, + step_size: float, + bias_correction2_sqrt: float, +) -> None: + """Update BF16/FP32 weights with persistent FP32 Adam moments.""" + + tensors = (model, grad, exp_avg, exp_avg_sq) + if any(not tensor.is_cuda for tensor in tensors): + raise ValueError("fused FP32-state AdamW requires CUDA tensors") + if any(tensor.device != model.device for tensor in tensors[1:]): + raise ValueError("fused FP32-state AdamW requires every tensor on the model device") + if model.dtype not in {torch.bfloat16, torch.float32}: + raise TypeError(f"fused FP32-state AdamW requires bfloat16 or float32 model weights, got {model.dtype}") + if grad.dtype not in {torch.bfloat16, torch.float32}: + raise TypeError(f"fused FP32-state AdamW requires bfloat16 or float32 gradients, got {grad.dtype}") + if exp_avg.dtype != torch.float32 or exp_avg_sq.dtype != torch.float32: + raise TypeError("FP32-state Adam moments must use float32") + if any(not tensor.is_contiguous() for tensor in tensors): + raise ValueError("fused FP32-state AdamW requires contiguous tensors") + if any(tensor.numel() != model.numel() for tensor in tensors[1:]): + raise ValueError("model, gradient, and FP32 moments must have the same number of elements") + extension().areno_adamw_fp32_state_step( + model, + grad, + exp_avg, + exp_avg_sq, + beta1, + beta2, + effective_lr, + weight_decay, + eps, + step_size, + bias_correction2_sqrt, + ) + + @torch._dynamo.disable @torch.no_grad() def areno_adamw_4bit_step( @@ -200,4 +265,9 @@ def areno_adamw_4bit_step( ) -__all__ = ["areno_adamw_4bit_step", "areno_adamw_8bit_step", "areno_adamw_fp32_master_step"] +__all__ = [ + "areno_adamw_4bit_step", + "areno_adamw_8bit_step", + "areno_adamw_fp32_master_step", + "areno_adamw_fp32_state_step", +] diff --git a/areno/api/backend/mlx/backend.py b/areno/api/backend/mlx/backend.py index c76c70bf..5175b083 100644 --- a/areno/api/backend/mlx/backend.py +++ b/areno/api/backend/mlx/backend.py @@ -92,7 +92,10 @@ def initialize(self, ctx: Context): self._validate_tokenizer(ctx.tokenizer) optimizer_config = self.config.optimizer self.provider.configure_trainability(optimizer_config) - self.optimizer, self._optimizer_groups = build_optimizer(optimizer_config) + self.optimizer, self._optimizer_groups = build_optimizer( + optimizer_config, + state_precision_for_parameter=self.provider.optimizer_state_precision, + ) if self.config.gradient_checkpointing: self._enable_gradient_checkpointing() self.model.train() diff --git a/areno/api/backend/mlx/optimizer.py b/areno/api/backend/mlx/optimizer.py index fcdf1d64..320a52da 100644 --- a/areno/api/backend/mlx/optimizer.py +++ b/areno/api/backend/mlx/optimizer.py @@ -2,12 +2,20 @@ from __future__ import annotations +from collections.abc import Callable from typing import Any from areno.api.backend.mlx.provider import parameter_group +from areno.engine.optim.dynamic_quant import SIGNED_DYNAMIC_MAP, UNSIGNED_DYNAMIC_MAP +_MLX_CODEBOOK_CACHE: dict[bool, Any] = {} -def build_optimizer(config: dict[str, Any]): + +def build_optimizer( + config: dict[str, Any], + *, + state_precision_for_parameter: Callable[[str, Any], str] | None = None, +): """Build AdamW groups matching CUDA policy/tower/projector controls.""" import mlx.optimizers as optim @@ -16,14 +24,14 @@ def build_optimizer(config: dict[str, Any]): filters = [] if config.get("unfreeze_multimodal_tower"): tower_config = _group_config(config, "tower") - groups.append(("tower", _adamw(tower_config), tower_config)) + groups.append(("tower", _adamw(tower_config, state_precision_for_parameter), tower_config)) filters.append(lambda path, _: parameter_group(path) == "tower") if config.get("unfreeze_multimodal_projector"): projector_config = _group_config(config, "projector") - groups.append(("projector", _adamw(projector_config), projector_config)) + groups.append(("projector", _adamw(projector_config, state_precision_for_parameter), projector_config)) filters.append(lambda path, _: parameter_group(path) == "projector") model_config = dict(config) - groups.append(("model", _adamw(model_config), model_config)) + groups.append(("model", _adamw(model_config, state_precision_for_parameter), model_config)) if len(groups) == 1: return groups[0][1], groups optimizer_type = _streaming_multi_optimizer_class() if config.get("adam_8bit") else optim.MultiOptimizer @@ -78,7 +86,10 @@ def _group_config(config: dict[str, Any], group: str) -> dict[str, Any]: return result -def _adamw(config: dict[str, Any]): +def _adamw( + config: dict[str, Any], + state_precision_for_parameter: Callable[[str, Any], str] | None = None, +): import mlx.optimizers as optim kwargs = { @@ -88,7 +99,10 @@ def _adamw(config: dict[str, Any]): } if not config.get("adam_8bit"): return optim.AdamW(**kwargs, bias_correction=True) - return _quantized_adamw_class()(**kwargs) + return _quantized_adamw_class()( + **kwargs, + state_precision_for_parameter=state_precision_for_parameter, + ) def _quantized_adamw_class(): @@ -96,7 +110,7 @@ def _quantized_adamw_class(): from mlx.optimizers import Optimizer class AdamW8bit(Optimizer): - """AdamW with blockwise uint8 first-moment and root-second-moment storage.""" + """Paper-compatible blockwise dynamic AdamW with FP32 embedding states.""" def __init__( self, @@ -104,8 +118,9 @@ def __init__( betas=(0.9, 0.999), eps: float = 1e-8, weight_decay: float = 0.01, - block_size: int = 256, + block_size: int = 128, update_blocks: int = 8192, + state_precision_for_parameter: Callable[[str, Any], str] | None = None, ) -> None: super().__init__() self._maybe_schedule("learning_rate", learning_rate) @@ -114,6 +129,7 @@ def __init__( self.weight_decay = float(weight_decay) self.block_size = int(block_size) self.update_blocks = int(update_blocks) + self.state_precision_for_parameter = state_precision_for_parameter def init_single(self, parameter, state: dict) -> None: size = int(parameter.size) @@ -128,6 +144,28 @@ def apply_single(self, gradient, parameter, state: dict): bias_correction2_sqrt = mx.sqrt(1.0 - beta2**step) size = int(state["size"]) initialized = bool(state["initialized"]) + precision = str(state.get("precision", "8bit")) + if precision == "fp32": + grad = gradient.astype(mx.float32) + values = parameter.astype(mx.float32) + if initialized: + m = state["m"] + v = state["v"] + else: + m = mx.zeros_like(values, dtype=mx.float32) + v = mx.zeros_like(values, dtype=mx.float32) + m = beta1 * m + (1.0 - beta1) * grad + v = beta2 * v + (1.0 - beta2) * mx.square(grad) + denom = mx.sqrt(v) / bias_correction2_sqrt + self.eps + updated = (values * (1.0 - lr * self.weight_decay) - (lr / bias_correction1) * m / denom).astype( + parameter.dtype + ) + state["m"] = m + state["v"] = v + state["initialized"] = True + mx.eval(updated, m, v) + mx.clear_cache() + return updated block_count = (size + self.block_size - 1) // self.block_size grad = gradient.reshape(-1) values = parameter.reshape(-1) @@ -154,19 +192,18 @@ def apply_single(self, gradient, parameter, state: dict): state["m_scale"][block_start:block_end], signed=True, ) - v_root = _dequant_blocks( + v = _dequant_blocks( state["v_q"][value_start:padded_end], state["v_scale"][block_start:block_end], signed=False, ) - v = mx.square(v_root) m = beta1 * m + (1.0 - beta1) * grad_chunk v = beta2 * v + (1.0 - beta2) * mx.square(grad_chunk) else: m = (1.0 - beta1) * grad_chunk v = (1.0 - beta2) * mx.square(grad_chunk) next_m_q, next_m_scale = _quantize_signed(m, self.block_size) - next_v_q, next_v_scale = _quantize_unsigned(mx.sqrt(v), self.block_size) + next_v_q, next_v_scale = _quantize_unsigned(v, self.block_size) denom = mx.sqrt(v) / bias_correction2_sqrt + self.eps updated = (value_chunk * (1.0 - lr * self.weight_decay) - (lr / bias_correction1) * m / denom)[ :actual @@ -194,6 +231,18 @@ def update_streaming(self, model, gradients: dict) -> None: self._begin_streaming_step() _apply_streaming_leaves(model, gradients, lambda _: self) + def prepare_parameter_state(self, path: str, parameter: Any, state: dict) -> None: + if "precision" in state: + return + precision = ( + "8bit" + if self.state_precision_for_parameter is None + else str(self.state_precision_for_parameter(path, parameter)) + ) + if precision not in {"8bit", "fp32"}: + raise ValueError(f"unsupported MLX AdamW8bit state precision: {precision!r}") + state["precision"] = precision + def _begin_streaming_step(self) -> None: for name, scheduler in self._schedulers.items(): self.state[name] = scheduler(self.step) @@ -235,6 +284,9 @@ def _apply_streaming_leaves(model: Any, gradients: dict, optimizer_for_path) -> optimizer = optimizer_for_path(path) state = _tree_get(optimizer.state, path) parameter = _model_parameter(model, path) + prepare = getattr(optimizer, "prepare_parameter_state", None) + if prepare is not None: + prepare(path, parameter, state) updated = optimizer.apply_single(gradient, parameter, state) _set_model_parameter(model, path, updated) _tree_set(gradients, path, None) @@ -245,7 +297,7 @@ def _apply_streaming_leaves(model: Any, gradients: dict, optimizer_for_path) -> def _tree_get(tree: Any, path: str) -> Any: current = tree for part in path.split("."): - current = current[int(part)] if isinstance(current, (list, tuple)) else current[part] + current = current[int(part)] if isinstance(current, list | tuple) else current[part] return current @@ -253,7 +305,7 @@ def _tree_set(tree: Any, path: str, value: Any) -> None: parts = path.split(".") current = tree for part in parts[:-1]: - current = current[int(part)] if isinstance(current, (list, tuple)) else current[part] + current = current[int(part)] if isinstance(current, list | tuple) else current[part] final = parts[-1] if isinstance(current, list): current[int(final)] = value @@ -264,7 +316,7 @@ def _tree_set(tree: Any, path: str, value: Any) -> None: def _model_parameter(model: Any, path: str): current = model for part in path.split("."): - current = current[int(part)] if isinstance(current, (list, tuple)) else getattr(current, part) + current = current[int(part)] if isinstance(current, list | tuple) else getattr(current, part) return current @@ -272,7 +324,7 @@ def _set_model_parameter(model: Any, path: str, value: Any) -> None: parts = path.split(".") current = model for part in parts[:-1]: - current = current[int(part)] if isinstance(current, (list, tuple)) else getattr(current, part) + current = current[int(part)] if isinstance(current, list | tuple) else getattr(current, part) final = parts[-1] if isinstance(current, list): current[int(final)] = value @@ -291,28 +343,44 @@ def _blocked(value, block_size: int): def _quantize_signed(value, block_size: int): - import mlx.core as mx - - blocks = _blocked(value, block_size) - scale = mx.maximum(mx.max(mx.abs(blocks), axis=1, keepdims=True) / 127.0, mx.array(1e-12)) - quantized = mx.clip(mx.round(blocks / scale), -127, 127).astype(mx.int16) + 128 - return quantized.astype(mx.uint8).reshape(-1), scale + return _quantize_dynamic(value, block_size, signed=True) def _quantize_unsigned(value, block_size: int): + return _quantize_dynamic(value, block_size, signed=False) + + +def _quantize_dynamic(value, block_size: int, *, signed: bool): import mlx.core as mx blocks = _blocked(value, block_size) - scale = mx.maximum(mx.max(blocks, axis=1, keepdims=True) / 255.0, mx.array(1e-12)) - quantized = mx.clip(mx.round(blocks / scale), 0, 255).astype(mx.uint8) - return quantized.reshape(-1), scale + scale = mx.max(mx.abs(blocks) if signed else mx.maximum(blocks, 0.0), axis=1, keepdims=True) + normalized = blocks / mx.maximum(scale, mx.array(1.0e-30, dtype=mx.float32)) + if not signed: + normalized = mx.maximum(normalized, 0.0) + codebook = _mlx_dynamic_codebook(signed=signed) + boundaries = (codebook[:-1] + codebook[1:]) * 0.5 + quantized = mx.searchsorted(boundaries, normalized).astype(mx.uint8) + return quantized.reshape(-1), scale.astype(mx.float32) def _dequant_blocks(quantized, scale, *, signed: bool): - values = quantized.reshape(scale.shape[0], -1).astype(scale.dtype) - if signed: - values = values - 128.0 + import mlx.core as mx + + codebook = _mlx_dynamic_codebook(signed=signed) + values = codebook[quantized.reshape(scale.shape[0], -1).astype(mx.uint32)] return (values * scale).reshape(-1) +def _mlx_dynamic_codebook(*, signed: bool): + import mlx.core as mx + + codebook = _MLX_CODEBOOK_CACHE.get(signed) + if codebook is None: + values = SIGNED_DYNAMIC_MAP if signed else UNSIGNED_DYNAMIC_MAP + codebook = mx.array(values, dtype=mx.float32) + _MLX_CODEBOOK_CACHE[signed] = codebook + return codebook + + __all__ = ["apply_optimizer_update", "build_optimizer", "materialize_optimizer_update", "set_group_learning_rates"] diff --git a/areno/api/backend/mlx/provider.py b/areno/api/backend/mlx/provider.py index ffa275a0..c9421ea4 100644 --- a/areno/api/backend/mlx/provider.py +++ b/areno/api/backend/mlx/provider.py @@ -65,6 +65,43 @@ def generation_model(self): def configure_trainability(self, optimizer_config: dict[str, Any]) -> None: del optimizer_config + def optimizer_state_precision(self, path: str, parameter: Any) -> str: + """Return role-aware optimizer-state precision for one MLX parameter.""" + + del path + for embedding in self._token_embedding_modules(): + if getattr(embedding, "weight", None) is parameter: + return "fp32" + return "8bit" + + def _token_embedding_modules(self) -> tuple[Any, ...]: + language_model = self.generation_model + body = getattr(language_model, "model", None) + modules = [] + seen: set[int] = set() + + def add(module: Any) -> None: + if module is None: + return + if isinstance(module, dict): + for child in module.values(): + add(child) + return + if isinstance(module, list | tuple): + for child in module: + add(child) + return + if id(module) not in seen: + modules.append(module) + seen.add(id(module)) + + for owner in (body, language_model): + if owner is None: + continue + for name in ("embed_tokens", "word_embeddings", "embed_tokens_per_layer"): + add(getattr(owner, name, None)) + return tuple(modules) + def prepare_generation_prompt(self, tokens: list[int], features: dict[str, Any] | None) -> dict[str, Any]: if features is not None: raise ValueError("text-only MLX checkpoints cannot consume multimodal prompt features") diff --git a/areno/engine/layers/vocab.py b/areno/engine/layers/vocab.py index df129ba1..4c234540 100644 --- a/areno/engine/layers/vocab.py +++ b/areno/engine/layers/vocab.py @@ -42,6 +42,9 @@ def __init__(self, vocab_size: int, hidden_size: int, *, dtype: torch.dtype | No self.vocab_start, self.vocab_end = _shard_range(vocab_size, ctx.rank, ctx.world_size) self.weight = nn.Parameter(torch.empty(self.vocab_end - self.vocab_start, hidden_size, dtype=dtype)) mark_tensor_parallel_parameter(self.weight, True, sequence_parallel=True) + # Optimizer-role metadata is consumed by AdamW8bit without relying on + # module/parameter name substrings. It does not change model outputs. + self.weight._areno_optimizer_role = "token_embedding" nn.init.normal_(self.weight, mean=0.0, std=0.02) def forward(self, input_ids: torch.Tensor) -> torch.Tensor: diff --git a/areno/engine/optim/__init__.py b/areno/engine/optim/__init__.py index cd27ed45..f2220123 100644 --- a/areno/engine/optim/__init__.py +++ b/areno/engine/optim/__init__.py @@ -6,7 +6,7 @@ """ from areno.engine.optim.adamw_4bit import AdamW4bit -from areno.engine.optim.adamw_8bit import AdamW8bit +from areno.engine.optim.adamw_8bit import AdamW8bit, set_optimizer_state_precision from areno.engine.optim.adamw_fp32_master import AdamWFP32Master -__all__ = ["AdamW4bit", "AdamW8bit", "AdamWFP32Master"] +__all__ = ["AdamW4bit", "AdamW8bit", "AdamWFP32Master", "set_optimizer_state_precision"] diff --git a/areno/engine/optim/adamw_4bit.py b/areno/engine/optim/adamw_4bit.py index 939da6ad..943bca45 100644 --- a/areno/engine/optim/adamw_4bit.py +++ b/areno/engine/optim/adamw_4bit.py @@ -49,6 +49,13 @@ class AdamW4bit(AdamW8bit): state directly with a fused block-wise kernel. """ + _embedding_fp32_state = False + state_quantizer = "signed-de4/zero-excluding-linear4" + + def _precision_for_parameter(self, parameter: torch.nn.Parameter) -> str: + del parameter + return "8bit" + def __init__( self, params: Iterable[torch.nn.Parameter], @@ -82,6 +89,14 @@ def state_dict(self) -> dict: payload = super().state_dict() payload.pop("adam_8bit", None) + payload.pop("quantizer", None) + payload.pop("precision_policy", None) + payload.pop("state_memory", None) + for state in payload["state"]: + state.pop("precision", None) + state.pop("quantizer", None) + state.pop("exp_avg", None) + state.pop("exp_avg_sq", None) payload["adam_4bit"] = True payload["state_format_version"] = _STATE_FORMAT_VERSION payload["quant_block_size"] = self.quant_block_size diff --git a/areno/engine/optim/adamw_8bit.py b/areno/engine/optim/adamw_8bit.py index 632b2136..499efc29 100644 --- a/areno/engine/optim/adamw_8bit.py +++ b/areno/engine/optim/adamw_8bit.py @@ -2,8 +2,9 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Mapping from dataclasses import dataclass +from typing import Any import torch import torch.distributed as dist @@ -17,9 +18,18 @@ _param_grad, _ParamRef, ) +from areno.engine.optim.dynamic_quant import ( + SIGNED_DYNAMIC_MAP, + SIGNED_DYNAMIC_ZERO, + UNSIGNED_DYNAMIC_MAP, + UNSIGNED_DYNAMIC_ZERO, +) _DEFAULT_QUANT_BLOCK_SIZE = 128 _MAX_FUSED_QUANT_BLOCK_SIZE = 4096 +_DYNAMIC_QUANTIZER = "dynamic-tree-v1" +_VALID_STATE_PRECISIONS = frozenset({"8bit", "fp32"}) +_CODEBOOK_CACHE: dict[tuple[torch.device, bool], torch.Tensor] = {} @dataclass(slots=True) @@ -31,6 +41,10 @@ class _Adam8bitBucketState: exp_avg_scale: torch.Tensor | None = None exp_avg_sq_q: torch.Tensor | None = None exp_avg_sq_scale: torch.Tensor | None = None + exp_avg: torch.Tensor | None = None + exp_avg_sq: torch.Tensor | None = None + precision: str = "8bit" + quantizer: str = _DYNAMIC_QUANTIZER offload_file: str | None = None offload_index: int | None = None offload_group: _MmapGroup | None = None @@ -42,12 +56,16 @@ class AdamW8bit(AdamWFP32Master): The model parameters remain BF16 on every DP rank. Adam moments are stored for only this rank's DP shard and re-quantized after every bucket update. - This trades optimizer precision for much lower persistent optimizer memory. + Explicit token-embedding parameters keep FP32 moments for stability, while + other parameters use the paper-compatible dynamic 8-bit codebooks. """ + _embedding_fp32_state = True + state_quantizer = _DYNAMIC_QUANTIZER + def __init__( self, - params: Iterable[torch.nn.Parameter], + params: Iterable[torch.nn.Parameter] | Iterable[Mapping[str, Any]], *, lr: float, betas: tuple[float, float], @@ -62,8 +80,12 @@ def __init__( raise ValueError( f"quant_block_size must be between 1 and {_MAX_FUSED_QUANT_BLOCK_SIZE}, got {quant_block_size}" ) + normalized_params, precision_by_id, role_by_id = _normalize_parameter_policies(params) + self._parameter_state_precision = precision_by_id + self._parameter_roles = role_by_id + self._bucket_numel = max(bucket_numel, 1) super().__init__( - params, + normalized_params, lr=lr, betas=betas, weight_decay=weight_decay, @@ -73,7 +95,40 @@ def __init__( dp_group=dp_group, ) self.quant_block_size = quant_block_size - self._states = [_Adam8bitBucketState() for _ in self.buckets] + self._states = [ + _Adam8bitBucketState( + precision=self._precision_for_parameter(bucket.refs[0].model_param), + quantizer=self.state_quantizer, + ) + for bucket in self.buckets + ] + + @torch.no_grad() + def _build_buckets(self, params: list[torch.nn.Parameter], bucket_numel: int) -> list[_MasterBucket]: + """Keep FP32-exempt and quantized parameters in separate buckets.""" + + buckets: list[_MasterBucket] = [] + pending: list[torch.nn.Parameter] = [] + pending_precision: str | None = None + for parameter in params: + precision = self._precision_for_parameter(parameter) + if pending and precision != pending_precision: + buckets.extend(AdamWFP32Master._build_buckets(self, pending, bucket_numel)) + pending = [] + pending.append(parameter) + pending_precision = precision + if pending: + buckets.extend(AdamWFP32Master._build_buckets(self, pending, bucket_numel)) + return buckets + + def _precision_for_parameter(self, parameter: torch.nn.Parameter) -> str: + explicit = self._parameter_state_precision.get(id(parameter)) + if explicit is not None: + return explicit + role = self._parameter_roles.get(id(parameter), getattr(parameter, "_areno_optimizer_role", None)) + if self._embedding_fp32_state and role == "token_embedding": + return "fp32" + return "8bit" @torch.no_grad() def step(self, closure=None): @@ -115,6 +170,8 @@ def clear_state(self) -> None: state.exp_avg_scale = None state.exp_avg_sq_q = None state.exp_avg_sq_scale = None + state.exp_avg = None + state.exp_avg_sq = None state.offload_file = None state.offload_index = None state.offload_group = None @@ -168,6 +225,10 @@ def onload_state(self, device: torch.device) -> None: state.exp_avg_sq_q = state.exp_avg_sq_q.to(device=device) if state.exp_avg_sq_scale is not None and state.exp_avg_sq_scale.device != device: state.exp_avg_sq_scale = state.exp_avg_sq_scale.to(device=device) + if state.exp_avg is not None and state.exp_avg.device != device: + state.exp_avg = state.exp_avg.to(device=device) + if state.exp_avg_sq is not None and state.exp_avg_sq.device != device: + state.exp_avg_sq = state.exp_avg_sq.to(device=device) state.offload_file = None state.offload_index = None state.offload_group = None @@ -190,14 +251,29 @@ def state_dict(self) -> dict: "dp_rank": self.dp_rank, "dp_size": self.dp_size, "adam_8bit": True, + "quantizer": self.state_quantizer, "quant_block_size": self.quant_block_size, + "precision_policy": [ + { + "parameter_index": index, + "role": self._parameter_roles.get(id(parameter)), + "precision": self._precision_for_parameter(parameter), + "numel": parameter.numel(), + } + for index, parameter in enumerate(self.model_params) + ], + "state_memory": self.state_memory_metrics(), "state": [ { "step": state.step, + "precision": state.precision, + "quantizer": state.quantizer, "exp_avg_q": payload["exp_avg_q"], "exp_avg_scale": payload["exp_avg_scale"], "exp_avg_sq_q": payload["exp_avg_sq_q"], "exp_avg_sq_scale": payload["exp_avg_sq_scale"], + "exp_avg": payload["exp_avg"], + "exp_avg_sq": payload["exp_avg_sq"], } for state, payload in zip(self._states, payloads, strict=True) ], @@ -211,12 +287,31 @@ def load_state_dict(self, state_dict: dict) -> None: self._active_offload_mode = "none" self._disk_offload_root = None self._active_offload_batch_size = 1 + saved_states = state_dict.get("state", []) + saved_quantizer = str(state_dict.get("quantizer", "")) + if saved_quantizer != _DYNAMIC_QUANTIZER: + raise ValueError(f"unsupported AdamW8bit quantizer: {saved_quantizer}") + if state_dict.get("precision_policy") is None: + raise ValueError("AdamW8bit checkpoint is missing its precision policy") + self._restore_precision_policy(state_dict["precision_policy"]) + if len(saved_states) != len(self.buckets): + raise ValueError( + "AdamW8bit checkpoint bucket count does not match the current optimizer layout: " + f"checkpoint={len(saved_states)}, optimizer={len(self.buckets)}" + ) for state in self._states: + state.step = 0 + state.exp_avg_q = None + state.exp_avg_scale = None + state.exp_avg_sq_q = None + state.exp_avg_sq_scale = None + state.exp_avg = None + state.exp_avg_sq = None + state.quantizer = self.state_quantizer state.offload_file = None state.offload_index = 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: @@ -227,35 +322,103 @@ def load_state_dict(self, state_dict: dict) -> None: continue device = bucket.refs[0].model_param.device state.step = int(saved.get("step", 0)) + saved_precision = str(saved.get("precision", "8bit")) + bucket_quantizer = str(saved.get("quantizer", saved_quantizer)) + if bucket_quantizer != _DYNAMIC_QUANTIZER: + raise ValueError(f"unsupported AdamW8bit bucket quantizer: {bucket_quantizer}") + if saved_precision != state.precision: + raise ValueError( + "AdamW8bit state precision policy mismatch: " + f"checkpoint={saved_precision}, optimizer={state.precision}" + ) + state.quantizer = bucket_quantizer + if state.precision == "fp32": + state.exp_avg = _load_optional_state_tensor(saved.get("exp_avg"), bucket, device) + state.exp_avg_sq = _load_optional_state_tensor(saved.get("exp_avg_sq"), bucket, device) + state.exp_avg_q = None + state.exp_avg_scale = None + state.exp_avg_sq_q = None + state.exp_avg_sq_scale = None + continue exp_avg_q = saved.get("exp_avg_q") exp_avg_scale = saved.get("exp_avg_scale") exp_avg_sq_q = saved.get("exp_avg_sq_q") exp_avg_sq_scale = saved.get("exp_avg_sq_scale") - 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_q = _load_optional_quantized_tensor(exp_avg_q, bucket, device) 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_q = _load_optional_quantized_tensor(exp_avg_sq_q, bucket, device) state.exp_avg_sq_scale = ( None if exp_avg_sq_scale is None else self._restore_scales(exp_avg_sq_scale, bucket, device) ) + def _restore_precision_policy(self, saved_policy: Any) -> None: + """Rebuild buckets from the checkpoint's identity-ordered policy.""" + + if not isinstance(saved_policy, list) or len(saved_policy) != len(self.model_params): + raise ValueError( + "AdamW8bit checkpoint precision policy does not match the current parameter count: " + f"checkpoint={len(saved_policy) if isinstance(saved_policy, list) else 'invalid'}, " + f"optimizer={len(self.model_params)}" + ) + restored_precision: dict[int, str] = {} + restored_roles: dict[int, str] = {} + for expected_index, (entry, parameter) in enumerate(zip(saved_policy, self.model_params, strict=True)): + if not isinstance(entry, Mapping) or int(entry.get("parameter_index", -1)) != expected_index: + raise ValueError( + f"AdamW8bit checkpoint has an invalid precision policy entry at index {expected_index}" + ) + saved_numel = int(entry.get("numel", -1)) + if saved_numel != parameter.numel(): + raise ValueError( + "AdamW8bit checkpoint precision policy parameter size mismatch: " + f"index={expected_index}, checkpoint={saved_numel}, optimizer={parameter.numel()}" + ) + restored_precision[id(parameter)] = _normalize_state_precision(str(entry.get("precision", "8bit"))) + role = entry.get("role") + if role is not None: + restored_roles[id(parameter)] = str(role) + self._parameter_state_precision = restored_precision + self._parameter_roles = restored_roles + self.buckets = self._build_buckets(self.model_params, self._bucket_numel) + self._states = [ + _Adam8bitBucketState( + precision=self._precision_for_parameter(bucket.refs[0].model_param), + quantizer=self.state_quantizer, + ) + for bucket in self.buckets + ] + + def state_memory_metrics(self) -> dict[str, int]: + """Report logical persistent moment storage for initialized buckets.""" + + quantized_state_bytes = 0 + fp32_exempt_bytes = 0 + block_metadata_bytes = 0 + for bucket, state in zip(self.buckets, self._states, strict=True): + if state.step == 0: + continue + if state.precision == "fp32": + fp32_exempt_bytes += 2 * bucket.shard_numel * 4 + else: + quantized_state_bytes += 2 * bucket.shard_numel + block_metadata_bytes += 2 * self._bucket_scale_count(bucket) * 4 + return { + "quantized_state_bytes": quantized_state_bytes, + "fp32_exempt_bytes": fp32_exempt_bytes, + "block_metadata_bytes": block_metadata_bytes, + "total_bytes": quantized_state_bytes + fp32_exempt_bytes + block_metadata_bytes, + } + def _restore_scales( self, saved: torch.Tensor, bucket: _MasterBucket, device: torch.device, ) -> torch.Tensor: - """Restore block scales, expanding legacy bucket-level scalar scales.""" + """Restore block scales for one quantized bucket.""" 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() @@ -275,12 +438,24 @@ def _ensure_bucket_state(self, bucket: _MasterBucket, state: _Adam8bitBucketStat state.exp_avg_sq_q = state.exp_avg_sq_q.to(device=device) if state.exp_avg_sq_scale is not None and state.exp_avg_sq_scale.device != device: state.exp_avg_sq_scale = state.exp_avg_sq_scale.to(device=device) + if state.exp_avg is not None and state.exp_avg.device != device: + state.exp_avg = state.exp_avg.to(device=device) + if state.exp_avg_sq is not None and state.exp_avg_sq.device != device: + state.exp_avg_sq = state.exp_avg_sq.to(device=device) + if state.precision == "fp32": + if state.exp_avg is None: + state.exp_avg = torch.zeros(bucket.shard_numel, device=device, dtype=torch.float32) + if state.exp_avg_sq is None: + state.exp_avg_sq = torch.zeros(bucket.shard_numel, device=device, dtype=torch.float32) + return 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(self._bucket_scale_count(bucket), device=device, dtype=torch.float32) + state.exp_avg_q = torch.full((bucket.shard_numel,), SIGNED_DYNAMIC_ZERO, device=device, dtype=torch.uint8) + state.exp_avg_scale = torch.zeros(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(self._bucket_scale_count(bucket), device=device, dtype=torch.float32) + state.exp_avg_sq_q = torch.full( + (bucket.shard_numel,), UNSIGNED_DYNAMIC_ZERO, device=device, dtype=torch.uint8 + ) + state.exp_avg_sq_scale = torch.zeros(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.""" @@ -309,10 +484,9 @@ def _load_state_offload(self, state: _Adam8bitBucketState, device: torch.device) state.offload_index, state.offload_group.tensors[state.offload_index], ) - state.exp_avg_q = _host_tensor_to(saved["exp_avg_q"], device, prefetched=prefetched) - state.exp_avg_scale = _host_tensor_to(saved["exp_avg_scale"], device, prefetched=prefetched) - state.exp_avg_sq_q = _host_tensor_to(saved["exp_avg_sq_q"], device, prefetched=prefetched) - state.exp_avg_sq_scale = _host_tensor_to(saved["exp_avg_sq_scale"], device, prefetched=prefetched) + for name in ("exp_avg_q", "exp_avg_scale", "exp_avg_sq_q", "exp_avg_sq_scale", "exp_avg", "exp_avg_sq"): + value = saved.get(name) + setattr(state, name, None if value is None else _host_tensor_to(value, device, prefetched=prefetched)) if prefetched and device.type == "cuda": self._retain_disk_prefetch(state.offload_index, saved, device) @@ -325,22 +499,33 @@ def _disk_mmap_group_for_index(self, index: int) -> _MmapGroup | None: def _state_mmap_specs(self, indices: list[int]) -> dict[int, dict[str, tuple[torch.dtype, tuple[int, ...]]]]: """Return the fixed raw-mmap layout for quantized Adam state.""" - return { - index: { - "exp_avg_q": (torch.uint8, (self.buckets[index].shard_numel,)), - "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, (self._bucket_scale_count(self.buckets[index]),)), - } - for index in indices - } + specs: dict[int, dict[str, tuple[torch.dtype, tuple[int, ...]]]] = {} + for index in indices: + bucket = self.buckets[index] + if self._states[index].precision == "fp32": + specs[index] = { + "exp_avg": (torch.float32, (bucket.shard_numel,)), + "exp_avg_sq": (torch.float32, (bucket.shard_numel,)), + } + else: + specs[index] = { + "exp_avg_q": (torch.uint8, (bucket.shard_numel,)), + "exp_avg_scale": (torch.float32, (self._bucket_scale_count(bucket),)), + "exp_avg_sq_q": (torch.uint8, (bucket.shard_numel,)), + "exp_avg_sq_scale": (torch.float32, (self._bucket_scale_count(bucket),)), + } + return specs def _offload_8bit_group_to_disk(self, indices: list[int]) -> None: """Persist a bounded group of quantized states in one serialization call.""" if self._disk_offload_root is None: raise RuntimeError("disk optimizer offload is active without a usable directory") - present_indices = [index for index in indices if self._states[index].exp_avg_q is not None] + present_indices = [ + index + for index in indices + if self._states[index].exp_avg_q is not None or self._states[index].exp_avg is not None + ] if not present_indices: return group = self._get_or_create_mmap_group(indices, self._state_mmap_specs(indices)) @@ -348,15 +533,17 @@ def _offload_8bit_group_to_disk(self, indices: list[int]) -> None: ready_events: list[torch.cuda.Event] = [] for index in present_indices: state = self._states[index] - 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 payloads[index] = { - "exp_avg_q": state.exp_avg_q, - "exp_avg_scale": state.exp_avg_scale, - "exp_avg_sq_q": state.exp_avg_sq_q, - "exp_avg_sq_scale": state.exp_avg_sq_scale, + name: value + for name in ( + "exp_avg_q", + "exp_avg_scale", + "exp_avg_sq_q", + "exp_avg_sq_scale", + "exp_avg", + "exp_avg_sq", + ) + if (value := getattr(state, name)) is not None } ready_events.extend(state.offload_ready_events) self._submit_disk_group_write(indices, group, payloads, tuple(ready_events)) @@ -369,6 +556,8 @@ def _offload_8bit_group_to_disk(self, indices: list[int]) -> None: state.exp_avg_scale = None state.exp_avg_sq_q = None state.exp_avg_sq_scale = None + state.exp_avg = None + state.exp_avg_sq = None state.offload_ready_events = () def _stage_8bit_state_on_cpu(self, state: _Adam8bitBucketState) -> None: @@ -381,6 +570,8 @@ def _stage_8bit_state_on_cpu(self, state: _Adam8bitBucketState) -> None: "exp_avg_scale": state.exp_avg_scale, "exp_avg_sq_q": state.exp_avg_sq_q, "exp_avg_sq_scale": state.exp_avg_sq_scale, + "exp_avg": state.exp_avg, + "exp_avg_sq": state.exp_avg_sq, }.items() if tensor is not None } @@ -389,6 +580,8 @@ def _stage_8bit_state_on_cpu(self, state: _Adam8bitBucketState) -> None: state.exp_avg_scale = staged.get("exp_avg_scale") state.exp_avg_sq_q = staged.get("exp_avg_sq_q") state.exp_avg_sq_scale = staged.get("exp_avg_sq_scale") + state.exp_avg = staged.get("exp_avg") + state.exp_avg_sq = staged.get("exp_avg_sq") def _state_cpu_payload( self, @@ -402,22 +595,30 @@ def _state_cpu_payload( assert state.offload_group is not None self._wait_disk_group_write(state.offload_group) saved = state.offload_group.tensors[index] - return {name: tensor.clone() for name, tensor in saved.items()} + return { + name: None if saved.get(name) is None else saved[name].clone() + for name in ( + "exp_avg_q", + "exp_avg_scale", + "exp_avg_sq_q", + "exp_avg_sq_scale", + "exp_avg", + "exp_avg_sq", + ) + } return { "exp_avg_q": _cpu_clone(state.exp_avg_q), "exp_avg_scale": _cpu_clone(state.exp_avg_scale), "exp_avg_sq_q": _cpu_clone(state.exp_avg_sq_q), "exp_avg_sq_scale": _cpu_clone(state.exp_avg_sq_scale), + "exp_avg": _cpu_clone(state.exp_avg), + "exp_avg_sq": _cpu_clone(state.exp_avg_sq), } @torch.no_grad() def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState) -> None: """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 - 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 @@ -430,19 +631,32 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState) continue effective_lr = float(getattr(ref.model_param, "_areno_lr", self.lr)) step_size = effective_lr / bias_correction1 - self._step_param_ref_8bit( - bucket, - ref, - grad, - state, - scale_offset, - block_count, - beta1, - beta2, - effective_lr, - step_size, - bias_correction2_sqrt, - ) + if state.precision == "fp32": + self._step_param_ref_fp32( + bucket, + ref, + grad, + state, + beta1, + beta2, + effective_lr, + step_size, + bias_correction2_sqrt, + ) + else: + self._step_param_ref_8bit( + bucket, + ref, + grad, + state, + scale_offset, + block_count, + beta1, + beta2, + effective_lr, + step_size, + 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): @@ -453,6 +667,58 @@ def _step_bucket_8bit(self, bucket: _MasterBucket, state: _Adam8bitBucketState) self._all_gather_bucket(bucket) bucket.grad_shard = None bucket.grad_param_ids = frozenset() + if state.precision == "8bit": + state.quantizer = self.state_quantizer + + @torch.no_grad() + def _step_param_ref_fp32( + self, + bucket: _MasterBucket, + ref: _ParamRef, + grad: torch.Tensor, + state: _Adam8bitBucketState, + beta1: float, + beta2: float, + effective_lr: float, + step_size: float, + bias_correction2_sqrt: float, + ) -> None: + """Update one FP32-exempt parameter shard without a master-weight copy.""" + + if ref.shard_numel == 0: + return + assert state.exp_avg is not None + assert state.exp_avg_sq is not None + grad_shard = grad if bucket.grad_shard is not None else grad.narrow(0, ref.shard_start, ref.shard_numel) + model_shard = ref.model_param.detach().reshape(-1).narrow(0, ref.param_start + ref.shard_start, ref.shard_numel) + exp_avg = state.exp_avg.narrow(0, ref.shard_bucket_start, ref.shard_numel) + exp_avg_sq = state.exp_avg_sq.narrow(0, ref.shard_bucket_start, ref.shard_numel) + if model_shard.is_cuda: + from areno.accel.optimizer import areno_adamw_fp32_state_step + + areno_adamw_fp32_state_step( + model_shard, + grad_shard.contiguous(), + exp_avg, + exp_avg_sq, + 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 + weight = model_shard.float() + gradient = grad_shard.float() + if self.weight_decay != 0.0: + weight.mul_(1.0 - effective_lr * self.weight_decay) + exp_avg.mul_(beta1).add_(gradient, alpha=1.0 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(gradient, gradient, value=1.0 - beta2) + denom = exp_avg_sq.sqrt().div_(bias_correction2_sqrt).add_(self.eps) + weight.addcdiv_(exp_avg, denom, value=-step_size) + model_shard.copy_(weight) @torch.no_grad() def _step_param_ref_8bit( @@ -491,6 +757,9 @@ def _step_param_ref_8bit( if model_shard.is_cuda: from areno.accel.optimizer import areno_adamw_8bit_step + signed_codebook = _dynamic_codebook(model_shard.device, signed=True) + unsigned_codebook = _dynamic_codebook(model_shard.device, signed=False) + areno_adamw_8bit_step( model_shard, grad_shard.contiguous(), @@ -498,6 +767,8 @@ def _step_param_ref_8bit( moment_scales, variance_q, variance_scales, + signed_codebook, + unsigned_codebook, block_size=self.quant_block_size, beta1=beta1, beta2=beta2, @@ -512,57 +783,173 @@ def _step_param_ref_8bit( 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) + # ``Tensor.to(float32)`` aliases an already-FP32 model shard. Keep + # this speculative update private until the whole block passes the + # finite check, matching the fused CUDA kernel's two-pass commit. + weight = model_shard.narrow(0, start, numel).to(dtype=torch.float32).clone() 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]) + moment = _dequantize_dynamic(block_moment_q, moment_scales[block_index], signed=True) + variance = _dequantize_dynamic(block_variance_q, variance_scales[block_index], signed=False) 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) + if not bool( + torch.isfinite(block_grad).all() + & torch.isfinite(moment).all() + & torch.isfinite(variance).all() + & torch.isfinite(weight).all() + ): + continue model_shard.narrow(0, start, numel).copy_(weight) - quantized_moment, moment_scale = _quantize_symmetric(moment) - quantized_variance, variance_scale = _quantize_positive(variance) + quantized_moment, moment_scale = _quantize_dynamic(moment, signed=True) + quantized_variance, variance_scale = _quantize_dynamic(variance, signed=False) 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]: - """Quantize a signed FP32 tensor to uint8 with one bucket-level scale.""" +def set_optimizer_state_precision( + parameter: torch.nn.Parameter, + precision: str, + *, + role: str | None = None, +) -> torch.nn.Parameter: + """Attach an explicit AdamW8bit state policy to a parameter. + + This is intentionally parameter metadata rather than name matching, so it + survives bucket construction and works for non-standard model layouts. + """ + + normalized = _normalize_state_precision(precision) + parameter._areno_optimizer_state_precision = normalized + if role is not None: + parameter._areno_optimizer_role = str(role) + return parameter + + +def _normalize_parameter_policies( + params: Iterable[torch.nn.Parameter] | Iterable[Mapping[str, Any]], +) -> tuple[list[torch.nn.Parameter], dict[int, str], dict[int, str]]: + values = list(params) + flattened: list[torch.nn.Parameter] = [] + precision_by_id: dict[int, str] = {} + role_by_id: dict[int, str] = {} + seen: set[int] = set() + for value in values: + if isinstance(value, Mapping): + group_params = list(value.get("params", ())) + group_precision = value.get("state_precision") + group_role = value.get("role") + else: + group_params = [value] + group_precision = None + group_role = None + for parameter in group_params: + if not isinstance(parameter, torch.nn.Parameter): + raise TypeError("AdamW8bit params must contain torch.nn.Parameter values") + identity = id(parameter) + explicit = getattr(parameter, "_areno_optimizer_state_precision", group_precision) + if explicit is not None: + precision_by_id[identity] = _normalize_state_precision(str(explicit)) + role = getattr(parameter, "_areno_optimizer_role", group_role) + if role is not None: + role_by_id[identity] = str(role) + if identity not in seen: + flattened.append(parameter) + seen.add(identity) + return flattened, precision_by_id, role_by_id + + +def _normalize_state_precision(precision: str) -> str: + aliases = {"uint8": "8bit", "int8": "8bit", "float32": "fp32"} + normalized = aliases.get(precision.lower(), precision.lower()) + if normalized not in _VALID_STATE_PRECISIONS: + raise ValueError(f"state_precision must be one of {sorted(_VALID_STATE_PRECISIONS)}, got {precision!r}") + return normalized + + +def _dynamic_codebook(device: torch.device, *, signed: bool) -> torch.Tensor: + key = (device, signed) + codebook = _CODEBOOK_CACHE.get(key) + if codebook is None: + values = SIGNED_DYNAMIC_MAP if signed else UNSIGNED_DYNAMIC_MAP + codebook = torch.tensor(values, device=device, dtype=torch.float32) + _CODEBOOK_CACHE[key] = codebook + return codebook + + +def _quantize_dynamic(tensor: torch.Tensor, *, signed: bool) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize one FP32 block with the paper's dynamic-tree codebook.""" if tensor.numel() == 0: - return tensor.to(dtype=torch.uint8), torch.ones((), device=tensor.device, dtype=torch.float32) - scale = tensor.abs().amax().div(127.0).clamp_min(1.0e-30) - quantized = torch.clamp(torch.round(tensor / scale) + 128.0, 0.0, 255.0).to(dtype=torch.uint8) - return quantized, scale.to(dtype=torch.float32) + return tensor.to(dtype=torch.uint8), torch.zeros((), device=tensor.device, dtype=torch.float32) + scale = (tensor.abs().amax() if signed else tensor.clamp_min(0).amax()).to(dtype=torch.float32) + normalized = tensor.float().div(scale.clamp_min(torch.finfo(torch.float32).tiny)) + if not signed: + normalized.clamp_min_(0.0) + codebook = _dynamic_codebook(tensor.device, signed=signed) + boundaries = (codebook[:-1] + codebook[1:]) * 0.5 + codes = torch.bucketize(normalized, boundaries).to(dtype=torch.uint8) + return codes, scale -def _dequantize_symmetric(quantized: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: - """Dequantize signed uint8 moments back to FP32.""" +def _dequantize_dynamic(quantized: torch.Tensor, scale: torch.Tensor, *, signed: bool) -> torch.Tensor: + codebook = _dynamic_codebook(quantized.device, signed=signed) + return codebook[quantized.long()].mul_(scale) - return (quantized.to(dtype=torch.float32) - 128.0).mul_(scale) +# Private convenience aliases used by focused quantization tests. +def _quantize_symmetric(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return _quantize_dynamic(tensor, signed=True) -def _quantize_positive(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - """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) - scale = tensor.amax().div(255.0).clamp_min(1.0e-30) - quantized = torch.clamp(torch.round(tensor / scale), 0.0, 255.0).to(dtype=torch.uint8) - return quantized, scale.to(dtype=torch.float32) +def _dequantize_symmetric(quantized: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + return _dequantize_dynamic(quantized, scale, signed=True) + + +def _quantize_positive(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return _quantize_dynamic(tensor, signed=False) def _dequantize_positive(quantized: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: - """Dequantize non-negative uint8 moments back to FP32.""" + return _dequantize_dynamic(quantized, scale, signed=False) + - return quantized.to(dtype=torch.float32).mul_(scale) +def _load_optional_state_tensor( + saved: Any, + bucket: _MasterBucket, + device: torch.device, +) -> torch.Tensor | None: + if saved is None: + return None + value = saved.detach().to(device=device, dtype=torch.float32).view(-1) + if value.numel() != bucket.shard_numel: + raise ValueError( + f"AdamW8bit checkpoint has {value.numel()} FP32 state values for a bucket requiring {bucket.shard_numel}" + ) + return value.clone() + + +def _load_optional_quantized_tensor( + saved: Any, + bucket: _MasterBucket, + device: torch.device, +) -> torch.Tensor | None: + if saved is None: + return None + value = saved.detach().to(device=device, dtype=torch.uint8).view(-1) + if value.numel() != bucket.shard_numel: + raise ValueError( + f"AdamW8bit checkpoint has {value.numel()} quantized state values " + f"for a bucket requiring {bucket.shard_numel}" + ) + return value.clone() def _cpu_clone(value: torch.Tensor | None) -> torch.Tensor | None: diff --git a/areno/engine/optim/dynamic_quant.py b/areno/engine/optim/dynamic_quant.py new file mode 100644 index 00000000..2874072c --- /dev/null +++ b/areno/engine/optim/dynamic_quant.py @@ -0,0 +1,61 @@ +"""Reference codebooks for Dettmers et al. block-wise dynamic quantization.""" + +from __future__ import annotations + + +def create_dynamic_map(*, signed: bool, max_exponent_bits: int = 7, total_bits: int = 8) -> tuple[float, ...]: + """Build the paper's sorted dynamic-tree codebook without a torch dependency. + + The construction matches the reference implementation released with + ``8-Bit Optimizers via Block-wise Quantization``. Signed maps reserve one + sign bit; unsigned maps reclaim it for an additional fractional bit. + """ + + if total_bits != 8: + raise ValueError("AReno dynamic optimizer states currently require total_bits=8") + if max_exponent_bits < 1 or max_exponent_bits >= total_bits: + raise ValueError("max_exponent_bits must be between 1 and total_bits - 1") + + values: list[float] = [] + non_sign_bits = total_bits - 1 + additional_items = 2 ** (non_sign_bits - max_exponent_bits) - 1 + last_index = 0 + for index in range(max_exponent_bits): + last_index = index + fraction_items = 2 ** (index + non_sign_bits - max_exponent_bits + (0 if signed else 1)) + 1 + step = 0.9 / (fraction_items - 1) + means = (0.1 + (item + 0.5) * step for item in range(fraction_items - 1)) + scale = 10 ** (-(max_exponent_bits - 1) + index) + positive = [scale * mean for mean in means] + values.extend(positive) + if signed: + values.extend(-value for value in positive) + + if additional_items > 0: + step = 0.9 / additional_items + means = (0.1 + (item + 0.5) * step for item in range(additional_items)) + scale = 10 ** (-(max_exponent_bits - 1) + last_index) + positive = [scale * mean for mean in means] + values.extend(positive) + if signed: + values.extend(-value for value in positive) + + values.extend((0.0, 1.0)) + if len(values) != 2**total_bits: + raise AssertionError(f"dynamic codebook has {len(values)} entries, expected {2**total_bits}") + return tuple(sorted(values)) + + +SIGNED_DYNAMIC_MAP = create_dynamic_map(signed=True) +UNSIGNED_DYNAMIC_MAP = create_dynamic_map(signed=False) +SIGNED_DYNAMIC_ZERO = SIGNED_DYNAMIC_MAP.index(0.0) +UNSIGNED_DYNAMIC_ZERO = UNSIGNED_DYNAMIC_MAP.index(0.0) + + +__all__ = [ + "SIGNED_DYNAMIC_MAP", + "SIGNED_DYNAMIC_ZERO", + "UNSIGNED_DYNAMIC_MAP", + "UNSIGNED_DYNAMIC_ZERO", + "create_dynamic_map", +] diff --git a/areno/engine/training.py b/areno/engine/training.py index e0dcf585..3aac40af 100644 --- a/areno/engine/training.py +++ b/areno/engine/training.py @@ -163,6 +163,7 @@ def _train_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() @@ -185,6 +186,12 @@ def _train_step( 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" + ): + optimizer_state_metrics = {f"adam8_{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: @@ -216,6 +223,7 @@ def _train_step( {"grad_norm": grad_norm} if grad_norm is not None else None, multimodal_grad_metrics, {"clipped_grad_norm": clipped_grad_norm} if clipped_grad_norm is not None else None, + optimizer_state_metrics, ), } return None diff --git a/docs/cli/training.rst b/docs/cli/training.rst index 5034c984..81d2e46f 100644 --- a/docs/cli/training.rst +++ b/docs/cli/training.rst @@ -380,8 +380,22 @@ in its description; flags for other algorithms are ignored. Policy optimizer Adam beta2. Default: ``0.999``. ``--adam-8bit`` - Use 8-bit Adam moment states instead of FP32 Adam states. Supported by both - native backends; validate convergence when changing optimizer precision. + Use block-wise 8-bit Adam moment states instead of FP32 Adam states. + CUDA and MLX use the signed dynamic-tree codebook for the first moment and + the unsigned dynamic codebook for the second moment from *8-Bit Optimizers + via Block-wise Quantization*. Token-embedding weights and gradients retain + their normal model precision, while their optimizer moments remain FP32 to + avoid quantizing embedding-gradient outliers. Other initialized moment + state uses two bytes per parameter plus FP32 block scales. + CUDA training metrics expose ``adam8_quantized_state_bytes``, + ``adam8_fp32_exempt_bytes``, ``adam8_block_metadata_bytes``, and + ``adam8_total_bytes`` for initialized DP-local optimizer state. + + This option does not insert a normalization layer or reinitialize a loaded + embedding. The paper's Stable Embedding forward architecture is not enabled + for AReno's current RoPE-only language models because they do not expose the + compatible additive-position-embedding boundary. 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 diff --git a/docs/getting-started/mlx.rst b/docs/getting-started/mlx.rst index 2bd699f7..fdbcfe14 100644 --- a/docs/getting-started/mlx.rst +++ b/docs/getting-started/mlx.rst @@ -92,8 +92,12 @@ options have the largest effect on MLX unified-memory use: instead of retaining it for the next rollout. ``--adam-8bit`` - Stores Adam moment state in the MLX backend's 8-bit representation. This - reduces optimizer memory; validate convergence for the target task. + Stores non-embedding Adam moments in the same block-wise dynamic 8-bit + representation used by the CUDA backend. Token-embedding optimizer moments + stay FP32, selected by parameter identity rather than name matching; model + weights, gradients, and forward behavior are unchanged. This reduces + optimizer memory for the remaining parameters; validate convergence for the + target task. ``--activation-checkpointing`` Recomputes supported decoder activations during backward. It is enabled by diff --git a/tests/test_adamw_8bit_blockwise_cpu.py b/tests/test_adamw_8bit_blockwise_cpu.py index c1a4fd9d..902d56a1 100644 --- a/tests/test_adamw_8bit_blockwise_cpu.py +++ b/tests/test_adamw_8bit_blockwise_cpu.py @@ -1,9 +1,24 @@ from __future__ import annotations +import copy +from types import SimpleNamespace +from unittest.mock import patch + import torch -from areno.engine.optim import AdamW8bit -from areno.engine.optim.adamw_8bit import _dequantize_positive, _quantize_positive +from areno.engine.optim import AdamW8bit, set_optimizer_state_precision +from areno.engine.optim.adamw_8bit import ( + _dequantize_positive, + _dequantize_symmetric, + _quantize_positive, + _quantize_symmetric, +) +from areno.engine.optim.dynamic_quant import ( + SIGNED_DYNAMIC_MAP, + SIGNED_DYNAMIC_ZERO, + UNSIGNED_DYNAMIC_MAP, + UNSIGNED_DYNAMIC_ZERO, +) def test_adamw8bit_uses_parameter_local_block_scales() -> None: @@ -30,13 +45,37 @@ def test_adamw8bit_uses_parameter_local_block_scales() -> None: 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]) +def test_adamw8bit_dynamic_codebooks_match_paper_reference_construction() -> None: + assert len(SIGNED_DYNAMIC_MAP) == len(UNSIGNED_DYNAMIC_MAP) == 256 + assert SIGNED_DYNAMIC_ZERO == 127 + assert UNSIGNED_DYNAMIC_ZERO == 0 + assert SIGNED_DYNAMIC_MAP[0] == -0.99296875 + assert SIGNED_DYNAMIC_MAP[-1] == UNSIGNED_DYNAMIC_MAP[-1] == 1.0 + assert all(left <= right for left, right in zip(SIGNED_DYNAMIC_MAP, SIGNED_DYNAMIC_MAP[1:])) + assert all(left <= right for left, right in zip(UNSIGNED_DYNAMIC_MAP, UNSIGNED_DYNAMIC_MAP[1:])) + + +def test_adamw8bit_dynamic_codebook_golden_round_trip() -> None: + signed = torch.tensor(SIGNED_DYNAMIC_MAP) + unsigned = torch.tensor(UNSIGNED_DYNAMIC_MAP) + + signed_q, signed_scale = _quantize_symmetric(signed) + unsigned_q, unsigned_scale = _quantize_positive(unsigned) + + expected_codes = torch.arange(256, dtype=torch.int64).to(torch.uint8) + assert torch.equal(signed_q, expected_codes) + assert torch.equal(unsigned_q, expected_codes) + torch.testing.assert_close(_dequantize_symmetric(signed_q, signed_scale), signed) + torch.testing.assert_close(_dequantize_positive(unsigned_q, unsigned_scale), unsigned) + + +def test_adamw8bit_unsigned_dynamic_map_preserves_small_second_moments() -> None: + values = torch.tensor([0.0, 7.75e-7, 5.5e-5, 5.5e-3, 0.55, 1.0]) quantized, scale = _quantize_positive(values) restored = _dequantize_positive(quantized, scale) - torch.testing.assert_close(restored, values) + torch.testing.assert_close(restored, values, rtol=0.15, atol=1.0e-8) def test_adamw8bit_same_lr_does_not_amplify_constant_gradient_step() -> None: @@ -53,3 +92,229 @@ def test_adamw8bit_same_lr_does_not_amplify_constant_gradient_step() -> None: optimizer.step() torch.testing.assert_close(parameter, torch.full_like(parameter, 0.999), rtol=0.0, atol=1.0e-6) + + +def test_adamw8bit_routes_embedding_role_to_fp32_state_without_name_matching() -> None: + embedding = torch.nn.Parameter(torch.ones(12)) + embedding._areno_optimizer_role = "token_embedding" + ordinary = torch.nn.Parameter(torch.ones(12)) + optimizer = AdamW8bit( + [embedding, ordinary], + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=64, + quant_block_size=4, + ) + + assert [state.precision for state in optimizer._states] == ["fp32", "8bit"] + embedding.grad = torch.linspace(-1.0, 1.0, 12) + ordinary.grad = embedding.grad.clone() + optimizer.step() + state = optimizer.state_dict() + + assert state["quantizer"] == "dynamic-tree-v1" + assert [item["precision"] for item in state["precision_policy"]] == ["fp32", "8bit"] + assert state["precision_policy"][0]["role"] == "token_embedding" + assert state["state"][0]["exp_avg"] is not None + assert state["state"][0]["exp_avg_q"] is None + assert state["state"][1]["exp_avg"] is None + assert state["state"][1]["exp_avg_q"].dtype == torch.uint8 + + +def test_vocab_parallel_embedding_role_is_safe_for_pretrained_weights_and_dp_sharding() -> None: + from areno.engine.layers.vocab import VocabParallelEmbedding + + with patch( + "areno.engine.layers.vocab.get_tp_context", + return_value=SimpleNamespace(rank=1, world_size=2), + ): + embedding = VocabParallelEmbedding(18, 6) + loaded_weight = torch.linspace(-1.0, 1.0, embedding.weight.numel()).reshape_as(embedding.weight) + embedding.weight.data.copy_(loaded_weight) + optimizer = AdamW8bit( + embedding.parameters(), + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + dp_rank=1, + dp_size=2, + ) + + assert embedding.weight._areno_optimizer_role == "token_embedding" + assert optimizer._states[0].precision == "fp32" + assert optimizer.buckets[0].shard_numel == (embedding.weight.numel() + 1) // 2 + torch.testing.assert_close(embedding.weight, loaded_weight, rtol=0.0, atol=0.0) + + +def test_adamw8bit_explicit_precision_override_beats_embedding_default_and_deduplicates_ties() -> None: + tied = torch.nn.Parameter(torch.ones(8)) + tied._areno_optimizer_role = "token_embedding" + set_optimizer_state_precision(tied, "8bit") + fp32 = torch.nn.Parameter(torch.ones(4)) + optimizer = AdamW8bit( + [ + {"params": [tied], "state_precision": "fp32"}, + {"params": [tied, fp32], "state_precision": "fp32"}, + ], + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=64, + ) + + assert optimizer.model_params == [tied, fp32] + assert [state.precision for state in optimizer._states] == ["8bit", "fp32"] + + +def test_adamw8bit_mixed_state_checkpoint_round_trip_preserves_next_update() -> None: + embedding = torch.nn.Parameter(torch.linspace(-0.5, 0.5, 12)) + embedding._areno_optimizer_role = "token_embedding" + ordinary = torch.nn.Parameter(torch.linspace(0.25, -0.25, 9)) + first = AdamW8bit( + [embedding, ordinary], + lr=4.0e-4, + betas=(0.9, 0.99), + weight_decay=0.01, + bucket_numel=64, + quant_block_size=4, + ) + embedding.grad = torch.linspace(-0.7, 0.3, embedding.numel()) + ordinary.grad = torch.linspace(0.4, -0.8, ordinary.numel()) + first.step() + checkpoint = copy.deepcopy(first.state_dict()) + + restored_embedding = torch.nn.Parameter(embedding.detach().clone()) + restored_embedding._areno_optimizer_role = "token_embedding" + restored_ordinary = torch.nn.Parameter(ordinary.detach().clone()) + restored = AdamW8bit( + [restored_embedding, restored_ordinary], + lr=4.0e-4, + betas=(0.9, 0.99), + weight_decay=0.01, + bucket_numel=64, + quant_block_size=4, + ) + restored.load_state_dict(checkpoint) + + next_embedding_grad = torch.linspace(0.6, -0.2, embedding.numel()) + next_ordinary_grad = torch.linspace(-0.1, 0.9, ordinary.numel()) + embedding.grad = next_embedding_grad.clone() + restored_embedding.grad = next_embedding_grad.clone() + ordinary.grad = next_ordinary_grad.clone() + restored_ordinary.grad = next_ordinary_grad.clone() + first.step() + restored.step() + + torch.testing.assert_close(restored_embedding, embedding, rtol=0.0, atol=0.0) + torch.testing.assert_close(restored_ordinary, ordinary, rtol=0.0, atol=0.0) + for actual, expected in zip(restored.state_dict()["state"], first.state_dict()["state"], strict=True): + for key in ("exp_avg_q", "exp_avg_scale", "exp_avg_sq_q", "exp_avg_sq_scale", "exp_avg", "exp_avg_sq"): + if expected[key] is None: + assert actual[key] is None + else: + torch.testing.assert_close(actual[key], expected[key], rtol=0.0, atol=0.0) + + +def test_adamw8bit_checkpoint_restores_saved_precision_policy_by_parameter_identity() -> None: + embedding = torch.nn.Parameter(torch.ones(6)) + embedding._areno_optimizer_role = "token_embedding" + ordinary = torch.nn.Parameter(torch.ones(6)) + source = AdamW8bit( + [embedding, ordinary], + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=16, + quant_block_size=4, + ) + embedding.grad = torch.ones_like(embedding) + ordinary.grad = torch.ones_like(ordinary) + source.step() + + restored_embedding = torch.nn.Parameter(embedding.detach().clone()) + restored_ordinary = torch.nn.Parameter(ordinary.detach().clone()) + restored = AdamW8bit( + [restored_embedding, restored_ordinary], + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=16, + quant_block_size=4, + ) + assert [state.precision for state in restored._states] == ["8bit"] + + restored.load_state_dict(source.state_dict()) + + assert [state.precision for state in restored._states] == ["fp32", "8bit"] + assert restored._parameter_roles[id(restored_embedding)] == "token_embedding" + + +def test_adamw8bit_reports_mixed_state_storage() -> None: + embedding = torch.nn.Parameter(torch.ones(12)) + embedding._areno_optimizer_role = "token_embedding" + ordinary = torch.nn.Parameter(torch.ones(12)) + optimizer = AdamW8bit( + [embedding, ordinary], + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=64, + quant_block_size=4, + ) + embedding.grad = torch.ones_like(embedding) + ordinary.grad = torch.ones_like(ordinary) + optimizer.step() + + assert optimizer.state_memory_metrics() == { + "quantized_state_bytes": 24, + "fp32_exempt_bytes": 96, + "block_metadata_bytes": 24, + "total_bytes": 144, + } + + +def test_adamw8bit_nonfinite_gradient_skips_only_affected_block() -> None: + parameter = torch.nn.Parameter(torch.zeros(8)) + optimizer = AdamW8bit( + [parameter], + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=16, + quant_block_size=4, + ) + gradient = torch.ones_like(parameter) + gradient[1] = torch.inf + parameter.grad = gradient + optimizer.step() + + torch.testing.assert_close(parameter[:4], torch.zeros(4)) + assert torch.all(parameter[4:] < 0) + + +def test_adamw8bit_disk_offload_supports_mixed_state(tmp_path) -> None: + embedding = torch.nn.Parameter(torch.ones(8)) + embedding._areno_optimizer_role = "token_embedding" + ordinary = torch.nn.Parameter(torch.ones(8)) + candidate = AdamW8bit( + [embedding, ordinary], + lr=1.0e-3, + betas=(0.9, 0.99), + weight_decay=0.0, + bucket_numel=16, + quant_block_size=4, + ) + candidate.configure_state_offload(mode="disk", directory=str(tmp_path), batch_size=2) + embedding.grad = torch.linspace(-1.0, 1.0, embedding.numel()) + ordinary.grad = torch.linspace(1.0, -1.0, ordinary.numel()) + candidate.step() + + assert all(state.offload_file is not None for state in candidate._states) + saved = candidate.state_dict()["state"] + assert saved[0]["exp_avg"] is not None and saved[0]["exp_avg_q"] is None + assert saved[1]["exp_avg"] is None and saved[1]["exp_avg_q"] is not None + candidate.onload_state(torch.device("cpu")) + assert candidate._states[0].exp_avg is not None + assert candidate._states[1].exp_avg_q is not None + assert not list(tmp_path.rglob("*.mmap")) diff --git a/tests/test_mlx_training_cpu.py b/tests/test_mlx_training_cpu.py index b9b81ee4..d9a3cccb 100644 --- a/tests/test_mlx_training_cpu.py +++ b/tests/test_mlx_training_cpu.py @@ -285,3 +285,63 @@ def test_adam8bit_matches_bias_corrected_adamw_for_uniform_moments(): mx.eval(reference_model.parameters(), quantized_model.parameters()) assert bool(mx.allclose(reference_model.weight, quantized_model.weight, atol=2e-5, rtol=2e-5).item()) + + +def test_adam8bit_dynamic_codebooks_match_cuda_reference(): + mx = pytest.importorskip("mlx.core") + + from areno.api.backend.mlx.optimizer import _mlx_dynamic_codebook + from areno.engine.optim.dynamic_quant import SIGNED_DYNAMIC_MAP, UNSIGNED_DYNAMIC_MAP + + _require_mlx_device(mx) + signed = _mlx_dynamic_codebook(signed=True) + unsigned = _mlx_dynamic_codebook(signed=False) + mx.eval(signed, unsigned) + + np.testing.assert_array_equal(np.array(signed), np.asarray(SIGNED_DYNAMIC_MAP, dtype=np.float32)) + np.testing.assert_array_equal(np.array(unsigned), np.asarray(UNSIGNED_DYNAMIC_MAP, dtype=np.float32)) + + +def test_adam8bit_mlx_precision_callback_keeps_fp32_moments(): + mx = pytest.importorskip("mlx.core") + nn = pytest.importorskip("mlx.nn") + from mlx.utils import tree_flatten, tree_unflatten + + from areno.api.backend.mlx.optimizer import _quantized_adamw_class, apply_optimizer_update + + _require_mlx_device(mx) + model = nn.Linear(8, 2, bias=False) + optimizer = _quantized_adamw_class()( + learning_rate=1e-3, + weight_decay=0.0, + state_precision_for_parameter=lambda _path, _parameter: "fp32", + ) + path = tree_flatten(model.trainable_parameters())[0][0] + gradient = mx.ones_like(model.weight) + apply_optimizer_update(model, optimizer, tree_unflatten([(path, gradient)])) + state_names = {name for name, _ in tree_flatten(optimizer.state)} + + assert any(name.endswith("m") for name in state_names) + assert any(name.endswith("v") for name in state_names) + assert not any(name.endswith(("m_q", "v_q", "m_scale", "v_scale")) for name in state_names) + + +def test_mlx_provider_routes_embedding_by_identity_without_path_matching(): + from areno.api.backend.mlx.provider import MlxModelProvider + + embedding_weight = object() + per_layer_weight = object() + ordinary_weight = object() + embedding = SimpleNamespace(weight=embedding_weight) + per_layer_embedding = SimpleNamespace(weight=per_layer_weight) + model = SimpleNamespace( + model=SimpleNamespace( + embed_tokens=embedding, + embed_tokens_per_layer=[per_layer_embedding], + ) + ) + provider = MlxModelProvider(model, tokenizer=None, processor=None, config={}) + + assert provider.optimizer_state_precision("unexpected.path", embedding_weight) == "fp32" + assert provider.optimizer_state_precision("another.unexpected.path", per_layer_weight) == "fp32" + assert provider.optimizer_state_precision("embed_tokens.lookalike", ordinary_weight) == "8bit" From 4ef9db2786009dae2ed838463cbb9be47f19c5c8 Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 14:52:43 +0800 Subject: [PATCH 2/3] test(optimizer): account for dynamic 8-bit quantization error --- tests/test_multimodal_optimizer_cpu.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/test_multimodal_optimizer_cpu.py b/tests/test_multimodal_optimizer_cpu.py index 6f61afda..6c7260c4 100644 --- a/tests/test_multimodal_optimizer_cpu.py +++ b/tests/test_multimodal_optimizer_cpu.py @@ -48,6 +48,11 @@ def test_cuda_adam8bit_matches_fp32_master_bias_corrected_updates(): betas=kwargs["betas"], weight_decay=kwargs["weight_decay"], ) + # Dynamic signed quantization deliberately uses the paper's asymmetric + # codebook, whose negative endpoint is -0.99296875 rather than -1.0. + # Accumulated updates therefore track FP32 within a small quantization + # budget instead of being bit-exact. + quantization_atol = 5e-6 for gradient in (0.25, -0.5, 0.125, 1.0): reference_param.grad = torch.tensor([gradient]) @@ -56,5 +61,5 @@ def test_cuda_adam8bit_matches_fp32_master_bias_corrected_updates(): reference.step() quantized.step() torch_reference.step() - torch.testing.assert_close(quantized_param, reference_param, atol=1e-6, rtol=1e-6) - torch.testing.assert_close(quantized_param, torch_param, atol=1e-6, rtol=1e-6) + torch.testing.assert_close(quantized_param, reference_param, atol=quantization_atol, rtol=1e-6) + torch.testing.assert_close(quantized_param, torch_param, atol=quantization_atol, rtol=1e-6) From 4c4dfb855cda302ce0d7623cd5c16517446db077 Mon Sep 17 00:00:00 2001 From: xsuler Date: Wed, 2 Sep 2026 15:07:31 +0800 Subject: [PATCH 3/3] test(skills): account for single-turn demo skill --- tests/test_agent_skills_cpu.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_agent_skills_cpu.py b/tests/test_agent_skills_cpu.py index a081da37..5e796f74 100644 --- a/tests/test_agent_skills_cpu.py +++ b/tests/test_agent_skills_cpu.py @@ -26,7 +26,7 @@ def test_repository_agent_skills_are_valid(): ) result = json.loads(process.stdout) - assert result["skill_count"] == 10 + assert result["skill_count"] == 11 assert result["script_count"] >= 15