Skip to content

feat(optimizer): add opt-in packed 4-bit AdamW states #537

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
xsuler merged 2 commits into from
Sep 1, 2026
Merged

feat(optimizer): add opt-in packed 4-bit AdamW states #537

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
afcc9ae
feat(optimizer): add packed 4-bit AdamW states
xsuler Sep 1, 2026
a60d8dc
docs(optimizer): document AdamW4bit usage
xsuler Sep 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion areno/accel/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@
from areno.accel.linear import areno_grouped_linear, areno_linear
from areno.accel.moe import areno_moe_permute, areno_moe_topk_permute, areno_moe_unpermute
from areno.accel.normalization import areno_optional_scale_rmsnorm, areno_rmsnorm, areno_rmsnorm_silu_gate
from areno.accel.optimizer import areno_adamw_fp32_master_step
from areno.accel.optimizer import areno_adamw_4bit_step, areno_adamw_8bit_step, areno_adamw_fp32_master_step
from areno.accel.router import areno_grouped_topk_router
from areno.accel.routing import areno_moe_align
from areno.accel.topk import areno_topk_softmax
Expand All @@ -50,6 +50,8 @@
"areno_moe_permute",
"areno_moe_topk_permute",
"areno_moe_unpermute",
"areno_adamw_4bit_step",
"areno_adamw_8bit_step",
"areno_adamw_fp32_master_step",
"areno_optional_scale_rmsnorm",
"areno_rmsnorm",
Expand Down
34 changes: 34 additions & 0 deletions areno/accel/csrc/extension.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,9 +185,43 @@ void areno_adamw_fp32_master_step_cuda(
double eps,
double step_size,
double bias_correction2_sqrt);
void areno_adamw_4bit_step_cuda(
torch::Tensor model,
torch::Tensor grad,
torch::Tensor exp_avg_q,
torch::Tensor exp_avg_scale,
torch::Tensor exp_avg_sq_q,
torch::Tensor exp_avg_sq_scale,
int64_t packed_offset,
int64_t scale_offset,
int64_t quant_block_size,
double beta1,
double beta2,
double effective_lr,
double weight_decay,
double eps,
double step_size,
double bias_correction2_sqrt);
void areno_adamw_8bit_step_cuda(
torch::Tensor model,
torch::Tensor grad,
torch::Tensor exp_avg_q,
torch::Tensor exp_avg_scale,
torch::Tensor exp_avg_sq_q,
torch::Tensor exp_avg_sq_scale,
int64_t quant_block_size,
double beta1,
double beta2,
double effective_lr,
double weight_decay,
double eps,
double step_size,
double bias_correction2_sqrt);

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("areno_adamw_fp32_master_step", &areno_adamw_fp32_master_step_cuda, "ARENO compact FP32-master AdamW step");
m.def("areno_adamw_4bit_step", &areno_adamw_4bit_step_cuda, "ARENO packed block-wise AdamW4bit step");
m.def("areno_adamw_8bit_step", &areno_adamw_8bit_step_cuda, "ARENO block-wise 8-bit AdamW step");
m.def("areno_silu_and_mul", &areno_silu_and_mul_cuda, "ARENO SiLU and multiply");
m.def("areno_gelu_tanh_and_mul", &areno_gelu_tanh_and_mul_cuda, "ARENO tanh GELU and multiply");
m.def("areno_silu", &areno_silu_cuda, "ARENO SiLU");
Expand Down
Loading
Loading