adafactor8bit 0.1.0__tar.gz

adafactor8bit-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 WANG YAN
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

adafactor8bit-0.1.0/MANIFEST.in

@@ -0,0 +1,3 @@
+include LICENSE
+include README.md
+recursive-include adafactor8bit *.cu

adafactor8bit-0.1.0/PKG-INFO

@@ -0,0 +1,117 @@
+Metadata-Version: 2.4
+Name: adafactor8bit
+Version: 0.1.0
+Summary: 8-bit Adafactor Optimizer with Fused CUDA Kernels
+Home-page: https://github.com/yanfeiwong/adafactor-8bit
+Author: WANG YAN
+Author-email: yanfeiwong1997@outlook.com
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1
+Requires-Dist: ninja
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+**English** | [中文](./README_ZH.md)
+
+# Adafactor 8-bit with Fused CUDA Kernels
+
+An 8-bit Adafactor optimizer designed for memory-efficient large-scale model training.
+
+It uses fused CUDA kernels and block-wise quantization to reduce optimizer state memory while maintaining training stability, making it suitable for training large models such as LLMs and diffusion models.
+
+## Key Features
+
+- **Fused CUDA Kernel**: Integrates dequantization, EMA updates, Warp-Shuffle reduction, and requantization into a single kernel, utilizing `float4` vectorization to maximize memory bandwidth utilization.
+- **Zero CPU-GPU Sync**: Refactored the control flow to eliminate implicit synchronizations, ensuring the GPU computation pipeline runs asynchronously at high speed.
+- **Cross-Platform JIT**: Utilizes JIT (Just-In-Time) automatic compilation for seamless setup across Windows and Linux environments.
+
+## Algorithm Details
+
+Rebuilt upon the official PyTorch Adafactor, the mathematical logic **aligns more closely with the original paper and `HuggingFace transformers`**. Key differences include:
+
+1. **Safe Injection of `eps1`**: The official PyTorch implementation defaults to `eps1=None` and relies on `clamp`, which can lead to NaNs when encountering zero or extremely small gradients. This project adopts the original `grad_squared + eps1` approach, fundamentally guaranteeing the strict positive definiteness of the second moment and preventing training crashes caused by `rsqrt(0)`.
+2. **Coupled Weight Decay**: Unlike the official PyTorch implementation which decouples Weight Decay from RMS, this project retains the Coupled mechanism from the original paper (Weight Decay multiplied by the effective learning rate that includes RMS scaling).
+3. **Standard Parameter Support**: Fully retains core Adafactor switches such as `relative_step` and `scale_parameter`, ensuring compatibility with existing learning rate scheduling strategies.
+
+## Performance
+
+- **Memory Footprint**: The memory usage of optimizer states is **significantly lower than `AdamW8Bit`** (bitsandbytes), making it an ideal choice for training massive models or when memory-constrained.
+- **Training Speed**: The Fused Kernel and Zero-Sync design enable it to achieve step speeds comparable to mainstream 8-bit optimizers.
+- **Quantization Precision & Stability**: The second moment (variance) in Adafactor is always non-negative, so we map it to `UINT8 (0~255)`. Compared to traditional 8-bit optimizers that map to `INT8 (-127~127)`, providing higher effective quantization precision within the non-negative variance domain.
+
+## Installation
+
+This project uses JIT (Just-In-Time) compilation.
+
+Please ensure torch and ninja are installed, and a CUDA compiler (such as MSVC or GCC) is available in your environment.
+
+If CUDA compilation fails, the optimizer will automatically fall back to the pure PyTorch implementation.
+
+
+```bash
+pip install git+https://github.com/yanfeiwong/adafactor-8bit.git
+```
+
+## Usage Example
+
+It is recommended to use `param_groups` to keep sensitive layers (Embedding, Norm, Bias) in FP32, enabling 8-bit quantization only for large 2D weight matrices.
+
+```python
+import torch
+import torch.nn as nn
+from adafactor8bit import Adafactor8Bit
+
+def get_param_groups(model, weight_decay=1e-2):
+    decay, no_decay = [], []
+    for name, param in model.named_parameters():
+        if not param.requires_grad: continue
+        # Protect 1D tensors, biases, norms, and embeddings
+        if param.ndim <= 1 or "bias" in name or "norm" in name or "embed" in name:
+            no_decay.append(param)
+        else:
+            decay.append(param)
+            
+    return [
+        {"params": decay, "weight_decay": weight_decay, "quantize": True},
+        {"params": no_decay, "weight_decay": 0.0, "quantize": False}
+    ]
+
+model = MyModel().cuda()
+optimizer = Adafactor8Bit(
+    get_param_groups(model), 
+    lr=1e-3, 
+    relative_step=False,
+    block_size=2048,
+    min_8bit_size=4096
+)
+
+# Training loop...
+```
+
+For a complete example, please refer to [basic_usage.py](./examples/basic_usage.py).
+
+## Acknowledgements
+
+Thanks to the large language models Qwen and DeepSeek for valuable technical discussions and code reviews on CUDA low-level optimization, memory safety mechanisms, and cross-platform compilation pipeline design.
+
+Thanks to Tim Dettmers for the inspiration from the paper [8-BIT OPTIMIZERS VIA BLOCK-WISE QUANTIZATION](https://arxiv.org/pdf/2110.02861) and the [bitsandbytes](https://github.com/bitsandbytes-foundation/bitsandbytes) library.
+
+Thanks to the PyTorch team for providing the foundational Optimizer implementation and the C++ Extension toolchain.
+
+## License
+
+[The project is released under the MIT License.](./LICENSE)

adafactor8bit-0.1.0/README.md

@@ -0,0 +1,90 @@
+**English** | [中文](./README_ZH.md)
+
+# Adafactor 8-bit with Fused CUDA Kernels
+
+An 8-bit Adafactor optimizer designed for memory-efficient large-scale model training.
+
+It uses fused CUDA kernels and block-wise quantization to reduce optimizer state memory while maintaining training stability, making it suitable for training large models such as LLMs and diffusion models.
+
+## Key Features
+
+- **Fused CUDA Kernel**: Integrates dequantization, EMA updates, Warp-Shuffle reduction, and requantization into a single kernel, utilizing `float4` vectorization to maximize memory bandwidth utilization.
+- **Zero CPU-GPU Sync**: Refactored the control flow to eliminate implicit synchronizations, ensuring the GPU computation pipeline runs asynchronously at high speed.
+- **Cross-Platform JIT**: Utilizes JIT (Just-In-Time) automatic compilation for seamless setup across Windows and Linux environments.
+
+## Algorithm Details
+
+Rebuilt upon the official PyTorch Adafactor, the mathematical logic **aligns more closely with the original paper and `HuggingFace transformers`**. Key differences include:
+
+1. **Safe Injection of `eps1`**: The official PyTorch implementation defaults to `eps1=None` and relies on `clamp`, which can lead to NaNs when encountering zero or extremely small gradients. This project adopts the original `grad_squared + eps1` approach, fundamentally guaranteeing the strict positive definiteness of the second moment and preventing training crashes caused by `rsqrt(0)`.
+2. **Coupled Weight Decay**: Unlike the official PyTorch implementation which decouples Weight Decay from RMS, this project retains the Coupled mechanism from the original paper (Weight Decay multiplied by the effective learning rate that includes RMS scaling).
+3. **Standard Parameter Support**: Fully retains core Adafactor switches such as `relative_step` and `scale_parameter`, ensuring compatibility with existing learning rate scheduling strategies.
+
+## Performance
+
+- **Memory Footprint**: The memory usage of optimizer states is **significantly lower than `AdamW8Bit`** (bitsandbytes), making it an ideal choice for training massive models or when memory-constrained.
+- **Training Speed**: The Fused Kernel and Zero-Sync design enable it to achieve step speeds comparable to mainstream 8-bit optimizers.
+- **Quantization Precision & Stability**: The second moment (variance) in Adafactor is always non-negative, so we map it to `UINT8 (0~255)`. Compared to traditional 8-bit optimizers that map to `INT8 (-127~127)`, providing higher effective quantization precision within the non-negative variance domain.
+
+## Installation
+
+This project uses JIT (Just-In-Time) compilation.
+
+Please ensure torch and ninja are installed, and a CUDA compiler (such as MSVC or GCC) is available in your environment.
+
+If CUDA compilation fails, the optimizer will automatically fall back to the pure PyTorch implementation.
+
+
+```bash
+pip install git+https://github.com/yanfeiwong/adafactor-8bit.git
+```
+
+## Usage Example
+
+It is recommended to use `param_groups` to keep sensitive layers (Embedding, Norm, Bias) in FP32, enabling 8-bit quantization only for large 2D weight matrices.
+
+```python
+import torch
+import torch.nn as nn
+from adafactor8bit import Adafactor8Bit
+
+def get_param_groups(model, weight_decay=1e-2):
+    decay, no_decay = [], []
+    for name, param in model.named_parameters():
+        if not param.requires_grad: continue
+        # Protect 1D tensors, biases, norms, and embeddings
+        if param.ndim <= 1 or "bias" in name or "norm" in name or "embed" in name:
+            no_decay.append(param)
+        else:
+            decay.append(param)
+            
+    return [
+        {"params": decay, "weight_decay": weight_decay, "quantize": True},
+        {"params": no_decay, "weight_decay": 0.0, "quantize": False}
+    ]
+
+model = MyModel().cuda()
+optimizer = Adafactor8Bit(
+    get_param_groups(model), 
+    lr=1e-3, 
+    relative_step=False,
+    block_size=2048,
+    min_8bit_size=4096
+)
+
+# Training loop...
+```
+
+For a complete example, please refer to [basic_usage.py](./examples/basic_usage.py).
+
+## Acknowledgements
+
+Thanks to the large language models Qwen and DeepSeek for valuable technical discussions and code reviews on CUDA low-level optimization, memory safety mechanisms, and cross-platform compilation pipeline design.
+
+Thanks to Tim Dettmers for the inspiration from the paper [8-BIT OPTIMIZERS VIA BLOCK-WISE QUANTIZATION](https://arxiv.org/pdf/2110.02861) and the [bitsandbytes](https://github.com/bitsandbytes-foundation/bitsandbytes) library.
+
+Thanks to the PyTorch team for providing the foundational Optimizer implementation and the C++ Extension toolchain.
+
+## License
+
+[The project is released under the MIT License.](./LICENSE)

adafactor8bit-0.1.0/adafactor8bit/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) 2026 WANG YAN
+# Licensed under the MIT License.
+
+from .optimizer import Adafactor8Bit
+
+__all__ = ["Adafactor8Bit"]

adafactor8bit-0.1.0/adafactor8bit/kernels.cu

@@ -0,0 +1,139 @@
+// Copyright (c) 2026 WANG YAN
+// Licensed under the MIT License.
+
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+__global__ void fused_quantize_lerp_kernel(
+    unsigned char* __restrict__ q,
+    float* __restrict__ scale,
+    const float* __restrict__ new_val,
+    const float beta,
+    int block_size)
+{
+    int block_id = blockIdx.x;
+    int tid = threadIdx.x;
+    int stride = blockDim.x;
+    int start = block_id * block_size;
+
+    extern __shared__ float shared_mem[];
+    float* local_vals = shared_mem;           
+    float* s_max = &shared_mem[block_size];
+
+    float old_scale = scale[block_id];
+    float thread_max = 0.0f;
+    float one_minus_b = 1.0f - beta;
+
+    const float4* new_val_vec = reinterpret_cast<const float4*>(new_val + start);
+    uchar4* q_vec = reinterpret_cast<uchar4*>(q + start);
+
+    int vec_iters = (block_size / 4) / stride; 
+
+    for (int i = 0; i < vec_iters; i++) {
+        int idx = tid + i * stride; 
+        
+        float4 nv = new_val_vec[idx];
+        uchar4 q_val = q_vec[idx];
+
+        float upd0 = ((float)q_val.x / 255.0f * old_scale) * one_minus_b + nv.x * beta;
+        local_vals[idx * 4 + 0] = upd0;
+        thread_max = fmaxf(thread_max, upd0);
+
+        float upd1 = ((float)q_val.y / 255.0f * old_scale) * one_minus_b + nv.y * beta;
+        local_vals[idx * 4 + 1] = upd1;
+        thread_max = fmaxf(thread_max, upd1);
+
+        float upd2 = ((float)q_val.z / 255.0f * old_scale) * one_minus_b + nv.z * beta;
+        local_vals[idx * 4 + 2] = upd2;
+        thread_max = fmaxf(thread_max, upd2);
+
+        float upd3 = ((float)q_val.w / 255.0f * old_scale) * one_minus_b + nv.w * beta;
+        local_vals[idx * 4 + 3] = upd3;
+        thread_max = fmaxf(thread_max, upd3);
+    }
+
+
+    float val = thread_max;
+    for (int offset = 16; offset > 0; offset /= 2) {
+        val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
+    }
+
+
+    if (tid % 32 == 0) {
+        s_max[tid / 32] = val;
+    }
+    __syncthreads();
+
+
+    if (tid < 32) { 
+
+        val = (tid < 8) ? s_max[tid] : 0.0f; 
+        
+        for (int offset = 4; offset > 0; offset /= 2) {
+            val = fmaxf(val, __shfl_down_sync(0xffffffff, val, offset));
+        }
+        if (tid == 0) {
+            s_max[0] = val;
+        }
+    }
+    __syncthreads();
+
+    float new_scale = fmaxf(s_max[0], 1e-12f);
+    float inv_scale = 255.0f / new_scale;
+
+    for (int i = 0; i < vec_iters; i++) {
+        int idx = tid + i * stride;
+        
+        float v0 = local_vals[idx * 4 + 0];
+        float v1 = local_vals[idx * 4 + 1];
+        float v2 = local_vals[idx * 4 + 2];
+        float v3 = local_vals[idx * 4 + 3];
+
+        uchar4 out_q;
+        out_q.x = (unsigned char)fminf(fmaxf(v0 * inv_scale + 0.5f, 0.0f), 255.0f);
+        out_q.y = (unsigned char)fminf(fmaxf(v1 * inv_scale + 0.5f, 0.0f), 255.0f);
+        out_q.z = (unsigned char)fminf(fmaxf(v2 * inv_scale + 0.5f, 0.0f), 255.0f);
+        out_q.w = (unsigned char)fminf(fmaxf(v3 * inv_scale + 0.5f, 0.0f), 255.0f);
+
+        q_vec[idx] = out_q;
+    }
+
+    if (tid == 0) {
+        scale[block_id] = new_scale;
+    }
+}
+
+torch::Tensor fused_quantize_lerp_cuda(
+    torch::Tensor q,
+    torch::Tensor scale,
+    torch::Tensor new_val, 
+    float beta,
+    int block_size)
+{
+    TORCH_CHECK(q.scalar_type() == at::kByte, "q must be uint8");
+    TORCH_CHECK(scale.scalar_type() == at::kFloat, "scale must be float32");
+    TORCH_CHECK(new_val.scalar_type() == at::kFloat, "new_val must be float32");
+    TORCH_CHECK(q.is_cuda() && scale.is_cuda() && new_val.is_cuda(), "tensors must be on CUDA");
+    TORCH_CHECK(q.is_contiguous() && scale.is_contiguous() && new_val.is_contiguous(), "tensors must be contiguous");
+    
+    int threads = 256;
+    TORCH_CHECK(block_size >= threads, "block_size must be >= threads");
+    TORCH_CHECK(block_size % 4 == 0, "block_size must be multiple of 4 for vectorization");
+
+    int num_blocks = scale.size(0);   
+    size_t shared_mem = (block_size + 8) * sizeof(float); 
+
+    fused_quantize_lerp_kernel<<<num_blocks, threads, shared_mem>>>(
+        q.data_ptr<unsigned char>(),
+        scale.data_ptr<float>(),
+        new_val.data_ptr<float>(),
+        beta, 
+        block_size
+    );
+
+    return q;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fused_quantize_lerp", &fused_quantize_lerp_cuda, "Fused quantize lerp (CUDA)");
+}

adafactor8bit-0.1.0/adafactor8bit/optimizer.py

@@ -0,0 +1,383 @@
+# Copyright (c) 2026 WANG YAN
+# Licensed under the MIT License.
+
+import os
+import sys
+import math
+import logging
+from typing import Tuple, Optional, Union, List, Dict, Any, Iterable
+
+import torch
+from torch import Tensor
+from torch.optim import Optimizer
+from torch.utils.cpp_extension import load
+
+__all__ = ["Adafactor8Bit"]
+
+logger = logging.getLogger(__name__)
+
+# ==========================================
+# 1. CUDA Kernel JIT Loading
+# ==========================================
+_CUDA_MODULE = None
+_CUDA_AVAILABLE = False
+
+
+def _load_cuda_module() -> bool:
+    global _CUDA_MODULE, _CUDA_AVAILABLE
+    if _CUDA_MODULE is not None or _CUDA_AVAILABLE:
+        return _CUDA_AVAILABLE
+
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    cu_path = os.path.join(current_dir, "kernels.cu")
+
+    if not os.path.exists(cu_path):
+        logger.warning("kernels.cu not found. Falling back to pure PyTorch implementation.")
+        return False
+
+    try:
+        is_windows = sys.platform == "win32"
+        extra_cflags = ["/Zc:preprocessor"] if is_windows else []
+        extra_cuda_cflags = ["-O3", "--use_fast_math"]
+        if is_windows:
+            extra_cuda_cflags.extend(["-Xcompiler", "/Zc:preprocessor"])
+
+        _CUDA_MODULE = load(
+            name="adafactor8bit_cuda",
+            sources=[cu_path],
+            extra_cflags=extra_cflags,
+            extra_cuda_cflags=extra_cuda_cflags,
+            verbose=False
+        )
+        _CUDA_AVAILABLE = True
+        logger.info("Adafactor8Bit: CUDA Kernel loaded successfully!")
+    except Exception as e:
+        logger.warning(f"Adafactor8Bit: Failed to load CUDA Kernel. Falling back to PyTorch. Error: {e}")
+
+    return _CUDA_AVAILABLE
+
+
+# ==========================================
+# 2. Quantization Utilities
+# ==========================================
+def _quantize_nonneg(tensor: Tensor, block_size: int = 2048) -> Tuple[Tensor, Tensor, torch.Size, int]:
+    """将非负 FP32 张量分块量化为 UINT8 (0-255)"""
+    shape = tensor.shape
+    flat = tensor.flatten()
+    pad = (block_size - flat.numel() % block_size) % block_size
+    if pad:
+        flat = torch.nn.functional.pad(flat, (0, pad))
+
+    blocks = flat.view(-1, block_size)
+    scale = blocks.amax(dim=1, keepdim=True).clamp(min=1e-12)
+    q = torch.round((blocks / scale * 255.0).clamp_(0, 255)).to(torch.uint8)
+    return q, scale.squeeze(-1), shape, pad
+
+
+def _dequantize_nonneg(q: Tensor, scale: Tensor, shape: torch.Size, pad: int) -> Tensor:
+    blocks = q.float() * scale.unsqueeze(-1) / 255.0
+    flat = blocks.flatten()
+    if pad:
+        flat = flat[:-pad]
+    return flat.view(shape)
+
+
+def _pad_to_block_size(tensor: Tensor, block_size: int) -> Tensor:
+    flat = tensor.contiguous().flatten()
+    pad_len = (block_size - flat.numel() % block_size) % block_size
+    if pad_len > 0:
+        return torch.nn.functional.pad(flat, (0, pad_len))
+    return flat
+
+
+# ==========================================
+# 3. Optimizer Core
+# ==========================================
+class Adafactor8Bit(Optimizer):
+    """
+    8-bit Adafactor Optimizer with Fused CUDA Kernels.
+
+    Implements block-wise quantization for optimizer states, drastically reducing
+    memory footprint while maintaining training stability.
+    """
+
+    def __init__(
+        self,
+        params: Iterable[Union[Tensor, Dict[str, Any]]],
+        lr: float = 1e-2,
+        beta2_decay: float = -0.8,
+        eps: Tuple[Optional[float], float] = (1e-30, 1e-3),
+        d: float = 1.0,
+        weight_decay: float = 0.0,
+        maximize: bool = False,
+        relative_step: bool = False,
+        scale_parameter: bool = True,
+        quantize: bool = True,
+        block_size: int = 2048,
+        min_8bit_size: int = 4096,
+    ):
+        if not 0.0 <= lr:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 >= beta2_decay:
+            raise ValueError(f"Invalid beta2_decay: {beta2_decay}")
+
+        eps1, eps2 = eps
+        if eps1 is not None and not 0.0 <= eps1:
+            raise ValueError(f"Invalid epsilon1: {eps1}")
+        if not 0.0 <= eps2:
+            raise ValueError(f"Invalid epsilon2: {eps2}")
+        if not 1.0 <= d:
+            raise ValueError(f"Invalid clipping threshold d: {d}")
+        if not 0.0 <= weight_decay:
+            raise ValueError(f"Invalid weight_decay: {weight_decay}")
+        
+        if quantize and block_size % 1024 != 0:
+            raise ValueError(
+                f"block_size must be a multiple of 1024 for CUDA float4 vectorization, "
+                f"but got {block_size}. (Recommended: 2048 or 4096)"
+            )
+
+        defaults = dict(
+            lr=lr, beta2_decay=beta2_decay, eps=eps, d=d, weight_decay=weight_decay,
+            maximize=maximize, relative_step=relative_step, scale_parameter=scale_parameter,
+            quantize=quantize, block_size=block_size, min_8bit_size=min_8bit_size,
+        )
+        super().__init__(params, defaults)
+
+    def _init_group(
+        self,
+        group: Dict[str, Any],
+        params_with_grad: List[Tensor],
+        grads: List[Tensor],
+        states: List[Dict[str, Any]],
+        state_steps: List[float]
+    ):
+        group_quantize = group.get("quantize", True)
+        block_size = group.get("block_size", 2048)
+        min_8bit_size = group.get("min_8bit_size", 4096)
+
+        for p in group["params"]:
+            if p.grad is None:
+                continue
+
+            force_fp32 = p.numel() < min_8bit_size
+            use_quant = group_quantize and not force_fp32
+
+            params_with_grad.append(p)
+            grads.append(p.grad)
+            state = self.state[p]
+
+            if len(state) == 0:
+                state["step"] = 0.0
+                state["is_quantized"] = use_quant
+                state["block_size"] = block_size
+
+                if p.grad.dim() > 1:
+                    r_shape = list(p.grad.shape)
+                    r_shape[-1] = 1
+                    c_shape = list(p.grad.shape)
+                    c_shape[-2] = 1
+                    if use_quant:
+                        r_tmp = torch.zeros(r_shape, device=p.device)
+                        q, s, sh, pad = _quantize_nonneg(r_tmp, block_size)
+                        state["row_var_q"], state["row_var_scale"], state["row_var_shape"], state["row_var_pad"] = q, s, sh, pad
+
+                        c_tmp = torch.zeros(c_shape, device=p.device)
+                        q, s, sh, pad = _quantize_nonneg(c_tmp, block_size)
+                        state["col_var_q"], state["col_var_scale"], state["col_var_shape"], state["col_var_pad"] = q, s, sh, pad
+                    else:
+                        state["row_var"] = torch.zeros(r_shape, device=p.device)
+                        state["col_var"] = torch.zeros(c_shape, device=p.device)
+                else:
+                    if use_quant:
+                        v_tmp = torch.zeros_like(p.grad, memory_format=torch.preserve_format)
+                        q, s, sh, pad = _quantize_nonneg(v_tmp, block_size)
+                        state["variance_q"], state["variance_scale"], state["variance_shape"], state["variance_pad"] = q, s, sh, pad
+                    else:
+                        state["variance"] = torch.zeros_like(p.grad, memory_format=torch.preserve_format)
+            else:
+                # Checkpoint Compatibility & Device Alignment
+                if torch.is_tensor(state["step"]):
+                    state["step"] = float(state["step"].item())
+                elif not isinstance(state["step"], (int, float)):
+                    state["step"] = float(state["step"])
+
+                for k in list(state.keys()):
+                    if isinstance(state[k], torch.Tensor):
+                        if state[k].device != p.device:
+                            state[k] = state[k].to(p.device)
+
+                        # Defensive dtype casting for legacy checkpoints
+                        if k.endswith("_q") and state[k].dtype != torch.uint8:
+                            state[k] = state[k].to(torch.uint8)
+                            scale_key = k.replace("_q", "_scale")
+                            if scale_key in state and state[scale_key].dtype != torch.float32:
+                                state[scale_key] = state[scale_key].to(torch.float32)
+
+                curr_block_size = state.get("block_size", block_size)
+
+                # Auto Upgrade: FP32 -> 8-bit
+                if use_quant and not state.get("is_quantized", False):
+                    if p.grad.dim() > 1:
+                        if "row_var" in state and "row_var_q" not in state:
+                            q, s, sh, pad = _quantize_nonneg(state["row_var"], curr_block_size)
+                            state["row_var_q"], state["row_var_scale"], state["row_var_shape"], state["row_var_pad"] = q, s, sh, pad
+                            del state["row_var"]
+                        if "col_var" in state and "col_var_q" not in state:
+                            q, s, sh, pad = _quantize_nonneg(state["col_var"], curr_block_size)
+                            state["col_var_q"], state["col_var_scale"], state["col_var_shape"], state["col_var_pad"] = q, s, sh, pad
+                            del state["col_var"]
+                    else:
+                        if "variance" in state and "variance_q" not in state:
+                            q, s, sh, pad = _quantize_nonneg(state["variance"], curr_block_size)
+                            state["variance_q"], state["variance_scale"], state["variance_shape"], state["variance_pad"] = q, s, sh, pad
+                            del state["variance"]
+                    state["is_quantized"] = True
+
+                # Auto Downgrade: 8-bit -> FP32
+                elif not use_quant and state.get("is_quantized", False):
+                    if p.grad.dim() > 1:
+                        if "row_var_q" in state:
+                            state["row_var"] = _dequantize_nonneg(state.pop("row_var_q"), state.pop("row_var_scale"), state.pop("row_var_shape"), state.pop("row_var_pad"))
+                        if "col_var_q" in state:
+                            state["col_var"] = _dequantize_nonneg(state.pop("col_var_q"), state.pop("col_var_scale"), state.pop("col_var_shape"), state.pop("col_var_pad"))
+                    else:
+                        if "variance_q" in state:
+                            state["variance"] = _dequantize_nonneg(state.pop("variance_q"), state.pop("variance_scale"), state.pop("variance_shape"), state.pop("variance_pad"))
+                    state["is_quantized"] = False
+
+                if "is_quantized" not in state:
+                    state["is_quantized"] = ("row_var_q" in state or "col_var_q" in state or "variance_q" in state)
+                if "block_size" not in state:
+                    state["block_size"] = block_size
+
+            states.append(state)
+            state_steps.append(state["step"])
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step."""
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad, grads, states, state_steps = [], [], [], []
+            eps1, eps2 = group["eps"]
+            self._init_group(group, params_with_grad, grads, states, state_steps)
+
+            for i in range(len(params_with_grad)):
+                _update_param_8bit(
+                    params_with_grad[i], grads[i], states[i],
+                    d=group["d"], lr=group["lr"], beta2_decay=group["beta2_decay"],
+                    weight_decay=group["weight_decay"], eps1=eps1, eps2=eps2,
+                    maximize=group["maximize"], relative_step=group["relative_step"],
+                    scale_parameter=group["scale_parameter"], block_size=group.get("block_size", 2048),
+                )
+        return loss
+
+
+# ==========================================
+# 4. Parameter Update Logic
+# ==========================================
+def _update_param_8bit(
+    param: Tensor, grad: Tensor, 
+    state: Dict[str, Any],
+    d: float, lr: Union[float, Tensor], 
+    beta2_decay: float, weight_decay: float,
+    eps1: Optional[float], 
+    eps2: float, 
+    maximize: bool, 
+    relative_step: bool,
+    scale_parameter: bool, 
+    block_size: int
+):
+    if maximize:
+        grad = -grad
+    if eps1 is None:
+        eps1 = torch.finfo(param.dtype).eps
+
+    quantize = state.get("is_quantized", False)
+    curr_block_size = state.get("block_size", block_size)
+
+    # 1. Pure CPU step increment (Zero Sync)
+    step = state["step"] + 1.0
+    state["step"] = step
+    beta_val = math.pow(step, beta2_decay)
+
+    # 2. Learning Rate scheduling
+    if isinstance(lr, float):
+        rho = min(lr, 1.0 / math.sqrt(step)) if relative_step else lr
+        rho_t = torch.tensor(rho, device=param.device, dtype=torch.float32)
+    else:
+        if relative_step:
+            step_t = torch.tensor(step, device=param.device, dtype=torch.float32)
+            rho_t = torch.minimum(step_t.rsqrt(), lr)
+        else:
+            rho_t = lr
+
+    if scale_parameter:
+        param_rms = param.norm(2) / math.sqrt(param.numel())
+        alpha = torch.clamp(param_rms, min=eps2) * rho_t
+    else:
+        alpha = rho_t
+
+    if weight_decay != 0:
+        param.mul_(1 - alpha * weight_decay)
+
+    # 3. Second Moment Estimation
+    if grad.dim() > 1:
+        row_mean = torch.norm(grad, dim=-1, keepdim=True).square_().div_(grad.size(-1)).add_(eps1)
+        col_mean = torch.norm(grad, dim=-2, keepdim=True).square_().div_(grad.size(-2)).add_(eps1)
+
+        if quantize:
+            if _load_cuda_module():
+                row_mean_padded = _pad_to_block_size(row_mean, curr_block_size)
+                col_mean_padded = _pad_to_block_size(col_mean, curr_block_size)
+
+                _CUDA_MODULE.fused_quantize_lerp(state["row_var_q"], state["row_var_scale"], row_mean_padded, beta_val, curr_block_size)
+                _CUDA_MODULE.fused_quantize_lerp(state["col_var_q"], state["col_var_scale"], col_mean_padded, beta_val, curr_block_size)
+
+                row_var = _dequantize_nonneg(state["row_var_q"], state["row_var_scale"], state["row_var_shape"], state["row_var_pad"])
+                col_var = _dequantize_nonneg(state["col_var_q"], state["col_var_scale"], state["col_var_shape"], state["col_var_pad"])
+            else:
+                row_var = _dequantize_nonneg(state["row_var_q"], state["row_var_scale"], state["row_var_shape"], state["row_var_pad"])
+                col_var = _dequantize_nonneg(state["col_var_q"], state["col_var_scale"], state["col_var_shape"], state["col_var_pad"])
+                row_var.lerp_(row_mean, beta_val)
+                col_var.lerp_(col_mean, beta_val)
+                q, s, sh, pad = _quantize_nonneg(row_var, curr_block_size)
+                state["row_var_q"], state["row_var_scale"], state["row_var_shape"], state["row_var_pad"] = q, s, sh, pad
+                q, s, sh, pad = _quantize_nonneg(col_var, curr_block_size)
+                state["col_var_q"], state["col_var_scale"], state["col_var_shape"], state["col_var_pad"] = q, s, sh, pad
+        else:
+            row_var = state["row_var"]
+            col_var = state["col_var"]
+            row_var.lerp_(row_mean, beta_val)
+            col_var.lerp_(col_mean, beta_val)
+
+        var_estimate = row_var @ col_var
+        var_estimate.div_(row_var.mean(dim=-2, keepdim=True).clamp_(min=torch.finfo(param.dtype).eps))
+
+    else:
+        grad_sq = grad.square().add_(eps1)
+        if quantize:
+            if _load_cuda_module():
+                grad_sq_padded = _pad_to_block_size(grad_sq, curr_block_size)
+                _CUDA_MODULE.fused_quantize_lerp(state["variance_q"], state["variance_scale"], grad_sq_padded, beta_val, curr_block_size)
+                variance = _dequantize_nonneg(state["variance_q"], state["variance_scale"], state["variance_shape"], state["variance_pad"])
+            else:
+                variance = _dequantize_nonneg(state["variance_q"], state["variance_scale"], state["variance_shape"], state["variance_pad"])
+                variance.lerp_(grad_sq, beta_val)
+                q, s, sh, pad = _quantize_nonneg(variance, curr_block_size)
+                state["variance_q"], state["variance_scale"], state["variance_shape"], state["variance_pad"] = q, s, sh, pad
+        else:
+            variance = state["variance"]
+            variance.lerp_(grad_sq, beta_val)
+
+        var_estimate = variance
+
+    # 4. Parameter Update & Gradient Clipping
+    update = var_estimate.clamp_(min=torch.finfo(param.dtype).eps).rsqrt_().mul_(grad)
+    denom = torch.clamp(update.norm(2) / (math.sqrt(update.numel()) * d), min=1.0)
+    param.add_(update, alpha=-alpha / denom)

adafactor8bit-0.1.0/adafactor8bit.egg-info/PKG-INFO

@@ -0,0 +1,117 @@
+Metadata-Version: 2.4
+Name: adafactor8bit
+Version: 0.1.0
+Summary: 8-bit Adafactor Optimizer with Fused CUDA Kernels
+Home-page: https://github.com/yanfeiwong/adafactor-8bit
+Author: WANG YAN
+Author-email: yanfeiwong1997@outlook.com
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1
+Requires-Dist: ninja
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+**English** | [中文](./README_ZH.md)
+
+# Adafactor 8-bit with Fused CUDA Kernels
+
+An 8-bit Adafactor optimizer designed for memory-efficient large-scale model training.
+
+It uses fused CUDA kernels and block-wise quantization to reduce optimizer state memory while maintaining training stability, making it suitable for training large models such as LLMs and diffusion models.
+
+## Key Features
+
+- **Fused CUDA Kernel**: Integrates dequantization, EMA updates, Warp-Shuffle reduction, and requantization into a single kernel, utilizing `float4` vectorization to maximize memory bandwidth utilization.
+- **Zero CPU-GPU Sync**: Refactored the control flow to eliminate implicit synchronizations, ensuring the GPU computation pipeline runs asynchronously at high speed.
+- **Cross-Platform JIT**: Utilizes JIT (Just-In-Time) automatic compilation for seamless setup across Windows and Linux environments.
+
+## Algorithm Details
+
+Rebuilt upon the official PyTorch Adafactor, the mathematical logic **aligns more closely with the original paper and `HuggingFace transformers`**. Key differences include:
+
+1. **Safe Injection of `eps1`**: The official PyTorch implementation defaults to `eps1=None` and relies on `clamp`, which can lead to NaNs when encountering zero or extremely small gradients. This project adopts the original `grad_squared + eps1` approach, fundamentally guaranteeing the strict positive definiteness of the second moment and preventing training crashes caused by `rsqrt(0)`.
+2. **Coupled Weight Decay**: Unlike the official PyTorch implementation which decouples Weight Decay from RMS, this project retains the Coupled mechanism from the original paper (Weight Decay multiplied by the effective learning rate that includes RMS scaling).
+3. **Standard Parameter Support**: Fully retains core Adafactor switches such as `relative_step` and `scale_parameter`, ensuring compatibility with existing learning rate scheduling strategies.
+
+## Performance
+
+- **Memory Footprint**: The memory usage of optimizer states is **significantly lower than `AdamW8Bit`** (bitsandbytes), making it an ideal choice for training massive models or when memory-constrained.
+- **Training Speed**: The Fused Kernel and Zero-Sync design enable it to achieve step speeds comparable to mainstream 8-bit optimizers.
+- **Quantization Precision & Stability**: The second moment (variance) in Adafactor is always non-negative, so we map it to `UINT8 (0~255)`. Compared to traditional 8-bit optimizers that map to `INT8 (-127~127)`, providing higher effective quantization precision within the non-negative variance domain.
+
+## Installation
+
+This project uses JIT (Just-In-Time) compilation.
+
+Please ensure torch and ninja are installed, and a CUDA compiler (such as MSVC or GCC) is available in your environment.
+
+If CUDA compilation fails, the optimizer will automatically fall back to the pure PyTorch implementation.
+
+
+```bash
+pip install git+https://github.com/yanfeiwong/adafactor-8bit.git
+```
+
+## Usage Example
+
+It is recommended to use `param_groups` to keep sensitive layers (Embedding, Norm, Bias) in FP32, enabling 8-bit quantization only for large 2D weight matrices.
+
+```python
+import torch
+import torch.nn as nn
+from adafactor8bit import Adafactor8Bit
+
+def get_param_groups(model, weight_decay=1e-2):
+    decay, no_decay = [], []
+    for name, param in model.named_parameters():
+        if not param.requires_grad: continue
+        # Protect 1D tensors, biases, norms, and embeddings
+        if param.ndim <= 1 or "bias" in name or "norm" in name or "embed" in name:
+            no_decay.append(param)
+        else:
+            decay.append(param)
+            
+    return [
+        {"params": decay, "weight_decay": weight_decay, "quantize": True},
+        {"params": no_decay, "weight_decay": 0.0, "quantize": False}
+    ]
+
+model = MyModel().cuda()
+optimizer = Adafactor8Bit(
+    get_param_groups(model), 
+    lr=1e-3, 
+    relative_step=False,
+    block_size=2048,
+    min_8bit_size=4096
+)
+
+# Training loop...
+```
+
+For a complete example, please refer to [basic_usage.py](./examples/basic_usage.py).
+
+## Acknowledgements
+
+Thanks to the large language models Qwen and DeepSeek for valuable technical discussions and code reviews on CUDA low-level optimization, memory safety mechanisms, and cross-platform compilation pipeline design.
+
+Thanks to Tim Dettmers for the inspiration from the paper [8-BIT OPTIMIZERS VIA BLOCK-WISE QUANTIZATION](https://arxiv.org/pdf/2110.02861) and the [bitsandbytes](https://github.com/bitsandbytes-foundation/bitsandbytes) library.
+
+Thanks to the PyTorch team for providing the foundational Optimizer implementation and the C++ Extension toolchain.
+
+## License
+
+[The project is released under the MIT License.](./LICENSE)

adafactor8bit-0.1.0/adafactor8bit.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+LICENSE
+MANIFEST.in
+README.md
+setup.py
+adafactor8bit/__init__.py
+adafactor8bit/kernels.cu
+adafactor8bit/optimizer.py
+adafactor8bit.egg-info/PKG-INFO
+adafactor8bit.egg-info/SOURCES.txt
+adafactor8bit.egg-info/dependency_links.txt
+adafactor8bit.egg-info/requires.txt
+adafactor8bit.egg-info/top_level.txt

adafactor8bit-0.1.0/adafactor8bit.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

adafactor8bit-0.1.0/adafactor8bit.egg-info/requires.txt

@@ -0,0 +1,2 @@
+torch>=2.1
+ninja

adafactor8bit-0.1.0/adafactor8bit.egg-info/top_level.txt

@@ -0,0 +1 @@
+adafactor8bit

adafactor8bit-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

adafactor8bit-0.1.0/setup.py

@@ -0,0 +1,33 @@
+# Copyright (c) 2026 WANG YAN
+# Licensed under the MIT License.
+
+from pathlib import Path
+from setuptools import setup, find_packages
+
+this_directory = Path(__file__).parent
+long_description = (this_directory / "README.md").read_text(encoding="utf-8")
+
+setup(
+    name="adafactor8bit",
+    version="0.1.0",
+    description="8-bit Adafactor Optimizer with Fused CUDA Kernels",
+    author="WANG YAN",
+    author_email="yanfeiwong1997@outlook.com",
+    url="https://github.com/yanfeiwong/adafactor-8bit",
+    packages=find_packages(),
+    include_package_data=True,
+    package_data={"adafactor8bit": ["*.cu"]}, 
+    install_requires=[
+        "torch>=2.1",
+        "ninja",
+    ],
+    python_requires=">=3.10",
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "License :: OSI Approved :: MIT License",
+        "Operating System :: OS Independent",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+)