luminav 1.1.5__py3-none-any.whl

luminav/__init__.py

@@ -0,0 +1,3 @@
+from .luminav import LuminaV, __version__, __author__, __license__
+
+__all__ = ["LuminaV", "__version__", "__author__", "__license__"]

luminav/luminav.py

@@ -0,0 +1,651 @@
+# Copyright (c) 2026 Lumina Moon and Contributors.
+# Licensed under the Apache License, Version 2.0 (see LICENSE for details).
+
+"""
+LuminaV: We Were Too Broke for AdamW So We Trapped Gradients in a
+        Hyperbolic Straitjacket and Hired a Traffic Cop to Slap Them
+
+Official Repo : https://huggingface.co/cloverx-id/LuminaV-Optimizer-Paper
+Software DOI  : 10.57967/hf/10365
+Paper PDF     : https://huggingface.co/cloverx-id/XoneLM-1.0-Paper/blob/main/LuminaV.pdf
+Paper DOI     : 10.57967/hf/10270
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+from torch import Tensor
+from torch.optim import Optimizer
+
+__all__ = ["LuminaV"]
+__version__ = "1.1.5"
+__author__ = "Silver Moon (cloverxion)"
+__organization__ = "Lumina Moon"
+__license__ = "Apache-2.0"
+
+logger = logging.getLogger("LuminaV")
+
+
+def _supports_native_bf16() -> bool:
+    return torch.cuda.is_available() and torch.cuda.is_bf16_supported()
+
+
+HAS_TRITON = False
+try:
+    import triton
+    import triton.language as tl
+    HAS_TRITON = True
+except ImportError:
+    HAS_TRITON = False
+
+if HAS_TRITON:
+    @triton.jit
+    def _prng(seed, idx):
+        x = idx * 1103515245 + seed + 12345
+        x = x ^ (x >> 16)
+        return x
+
+    @triton.jit
+    def _store_param(p_ptr, offsets, p_val, mask, seed, dtype_mode: tl.constexpr, use_sr: tl.constexpr):
+        if use_sr:
+            if dtype_mode == 1:  # bfloat16
+                p_int = p_val.to(tl.int32, bitcast=True)
+                noise = (_prng(seed, offsets) & 0x7FFFFFFF) & 0xFFFF
+                p_int = (p_int + noise) & ~0xFFFF
+                p_out = p_int.to(tl.float32, bitcast=True).to(tl.bfloat16)
+                tl.store(p_ptr + offsets, p_out, mask=mask)
+                return
+            elif dtype_mode == 2:  # float16
+                p_int = p_val.to(tl.int32, bitcast=True)
+                noise = (_prng(seed, offsets) & 0x7FFFFFFF) & 0x1FFF
+                p_int = (p_int + noise) & ~0x1FFF
+                p_out = p_int.to(tl.float32, bitcast=True).to(tl.float16)
+                tl.store(p_ptr + offsets, p_out, mask=mask)
+                return
+
+        if dtype_mode == 1:
+            tl.store(p_ptr + offsets, p_val.to(tl.bfloat16), mask=mask)
+        elif dtype_mode == 2:
+            tl.store(p_ptr + offsets, p_val.to(tl.float16), mask=mask)
+        else:
+            tl.store(p_ptr + offsets, p_val, mask=mask)
+
+    @triton.jit
+    def _triton_tanh_fast(x):
+        return 2.0 * tl.sigmoid(2.0 * x) - 1.0
+
+    @triton.jit
+    def _lumina_v2_pass1_kernel(
+        grad_ptr, exp_avg_ptr, exp_avg_sq_ptr, mask_sum_ptr,
+        n_elements, beta1, beta2, c1, c2, BLOCK_SIZE: tl.constexpr
+    ):
+        pid = tl.program_id(axis=0)
+        offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+
+        g = tl.load(grad_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        m = tl.load(exp_avg_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        v = tl.load(exp_avg_sq_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+        m_new = beta1 * m + (1.0 - beta1) * g
+        nes_m = beta1 * m_new + (1.0 - beta1) * g
+        diff = g - m_new
+        v_new = beta2 * v + (1.0 - beta2) * (diff * diff)
+        v_new = tl.minimum(v_new, 65000.0)
+
+        sigma = tl.sqrt(tl.maximum(v_new, 0.0)) * c1 + c2
+        u = _triton_tanh_fast(nes_m / (sigma + 1e-6))
+        m_mask = tl.where((u * g) > 0.0, 1.0, 0.0)
+
+        tl.store(exp_avg_ptr + offsets, m_new, mask=mask)
+        tl.store(exp_avg_sq_ptr + offsets, v_new, mask=mask)
+
+        block_sum = tl.sum(tl.where(mask, m_mask, 0.0), axis=0)
+        tl.atomic_add(mask_sum_ptr, block_sum)
+
+    @triton.jit
+    def _lumina_v2_pass2_kernel(
+        p_ptr, grad_ptr, exp_avg_ptr, exp_avg_sq_ptr, mask_sum_ptr,
+        n_elements, beta1, c1, c2, lr, weight_decay, clamp_min,
+        seed, dtype_mode: tl.constexpr, use_sr: tl.constexpr, BLOCK_SIZE: tl.constexpr
+    ):
+        pid = tl.program_id(axis=0)
+        offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+
+        m_sum = tl.load(mask_sum_ptr)
+        m_bar = tl.minimum(tl.maximum(m_sum / n_elements, clamp_min), 1.0)
+
+        p = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        g = tl.load(grad_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        m = tl.load(exp_avg_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        v = tl.load(exp_avg_sq_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+        nes_m = beta1 * m + (1.0 - beta1) * g
+        sigma = tl.sqrt(tl.maximum(v, 0.0)) * c1 + c2
+        u = _triton_tanh_fast(nes_m / (sigma + 1e-6))
+        m_mask = tl.where((u * g) > 0.0, 1.0, 0.0)
+
+        if weight_decay != 0.0:
+            p = p * (1.0 - lr * weight_decay)
+
+        delta_theta = (u * m_mask) / m_bar
+        p_updated = p - lr * delta_theta
+        _store_param(p_ptr, offsets, p_updated, mask, seed, dtype_mode, use_sr)
+
+    @triton.jit
+    def _lumina_v2_single_pass_kernel(
+        p_ptr, grad_ptr, exp_avg_ptr, exp_avg_sq_ptr,
+        n_elements, beta1, beta2, c1, c2, lr, weight_decay,
+        seed, dtype_mode: tl.constexpr, use_sr: tl.constexpr, BLOCK_SIZE: tl.constexpr
+    ):
+        pid = tl.program_id(axis=0)
+        offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+
+        p = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        g = tl.load(grad_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        m = tl.load(exp_avg_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        v = tl.load(exp_avg_sq_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+        m_new = beta1 * m + (1.0 - beta1) * g
+        nes_m = beta1 * m_new + (1.0 - beta1) * g
+        diff = g - m_new
+        v_new = beta2 * v + (1.0 - beta2) * (diff * diff)
+        v_new = tl.minimum(v_new, 65000.0)
+
+        sigma = tl.sqrt(tl.maximum(v_new, 0.0)) * c1 + c2
+        u = _triton_tanh_fast(nes_m / (sigma + 1e-6))
+
+        tl.store(exp_avg_ptr + offsets, m_new, mask=mask)
+        tl.store(exp_avg_sq_ptr + offsets, v_new, mask=mask)
+
+        if weight_decay != 0.0:
+            p = p * (1.0 - lr * weight_decay)
+
+        p_updated = p - lr * u
+        _store_param(p_ptr, offsets, p_updated, mask, seed, dtype_mode, use_sr)
+
+    @triton.jit
+    def _lumina_v1_pass1_kernel(
+        grad_ptr, exp_avg_ptr, rms_sum_ptr,
+        n_elements, beta1, BLOCK_SIZE: tl.constexpr
+    ):
+        pid = tl.program_id(axis=0)
+        offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+
+        g = tl.load(grad_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        m = tl.load(exp_avg_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+        m_new = beta1 * m + (1.0 - beta1) * g
+        nes_m = beta1 * m_new + (1.0 - beta1) * g
+        tl.store(exp_avg_ptr + offsets, m_new, mask=mask)
+
+        sq_val = nes_m * nes_m
+        block_sq_sum = tl.sum(tl.where(mask, sq_val, 0.0), axis=0)
+        tl.atomic_add(rms_sum_ptr, block_sq_sum)
+
+    @triton.jit
+    def _lumina_v1_pass2_kernel(
+        grad_ptr, exp_avg_ptr, rms_sum_ptr, mask_sum_ptr,
+        n_elements, beta1, tau, eps, c2, alpha_ss, BLOCK_SIZE: tl.constexpr
+    ):
+        pid = tl.program_id(axis=0)
+        offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+
+        sq_sum = tl.load(rms_sum_ptr)
+        rms = tl.sqrt(sq_sum / n_elements + eps)
+        sigma = tau * rms + c2
+
+        g = tl.load(grad_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        m = tl.load(exp_avg_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+        nes_m = beta1 * m + (1.0 - beta1) * g
+        z = nes_m / (sigma + 1e-6)
+        z_soft = z / (1.0 + alpha_ss * tl.abs(z))
+        u = _triton_tanh_fast(z_soft)
+        m_mask = tl.where((u * g) > 0.0, 1.0, 0.0)
+
+        block_sum = tl.sum(tl.where(mask, m_mask, 0.0), axis=0)
+        tl.atomic_add(mask_sum_ptr, block_sum)
+
+    @triton.jit
+    def _lumina_v1_pass3_update_kernel(
+        p_ptr, grad_ptr, exp_avg_ptr, rms_sum_ptr, mask_sum_ptr,
+        n_elements, beta1, tau, eps, c2, alpha_ss, lr, weight_decay, clamp_min,
+        seed, dtype_mode: tl.constexpr, use_sr: tl.constexpr, cautious: tl.constexpr, BLOCK_SIZE: tl.constexpr
+    ):
+        pid = tl.program_id(axis=0)
+        offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+        mask = offsets < n_elements
+
+        sq_sum = tl.load(rms_sum_ptr)
+        rms = tl.sqrt(sq_sum / n_elements + eps)
+        sigma = tau * rms + c2
+
+        m_sum = tl.load(mask_sum_ptr)
+        m_bar = tl.minimum(tl.maximum(m_sum / n_elements, clamp_min), 1.0)
+
+        p = tl.load(p_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        g = tl.load(grad_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+        m = tl.load(exp_avg_ptr + offsets, mask=mask, other=0.0).to(tl.float32)
+
+        nes_m = beta1 * m + (1.0 - beta1) * g
+        z = nes_m / (sigma + 1e-6)
+        z_soft = z / (1.0 + alpha_ss * tl.abs(z))
+        u = _triton_tanh_fast(z_soft)
+        m_mask = tl.where((u * g) > 0.0, 1.0, 0.0)
+
+        if weight_decay != 0.0:
+            p = p * (1.0 - lr * weight_decay)
+
+        delta_theta = (u * m_mask) / m_bar if cautious else u
+        p_updated = p - lr * delta_theta
+        _store_param(p_ptr, offsets, p_updated, mask, seed, dtype_mode, use_sr)
+
+
+def _sr_update(p: Tensor, update: Tensor, lr: float, use_sr: bool):
+    if use_sr and p.dtype == torch.bfloat16:
+        p_fp32 = p.float().add_(update, alpha=-lr)
+        int_view = p_fp32.view(torch.int32)
+        noise = torch.randint(0, 1 << 16, p.shape, dtype=torch.int32, device=p.device)
+        p.copy_(((int_view + noise) & ~0xFFFF).view(torch.float32).to(torch.bfloat16))
+    elif use_sr and p.dtype == torch.float16:
+        p_fp32 = p.float().add_(update, alpha=-lr)
+        int_view = p_fp32.view(torch.int32)
+        noise = torch.randint(0, 1 << 13, p.shape, dtype=torch.int32, device=p.device)
+        p.copy_(((int_view + noise) & ~0x1FFF).view(torch.float32).to(torch.float16))
+    else:
+        p.add_(update, alpha=-lr)
+
+
+class LuminaV(Optimizer):
+    def __init__(
+        self,
+        params,
+        lr: float = 8e-4,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-8,
+        weight_decay: float = 8e-2,
+        tau: float = 0.8,
+        alpha_ss: float = 0.5,
+        cautious: bool = True,
+        cautious_clamp_min: float = 0.2,
+        buffer: int = 2,
+        stochastic_rounding: bool = True,
+        execution: str = "auto",
+    ):
+        if lr < 0.0:
+            raise ValueError(f"Invalid learning rate: {lr}")
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError(f"Invalid beta1 parameter: {betas[0]}")
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError(f"Invalid beta2 parameter: {betas[1]}")
+        if eps <= 0.0:
+            raise ValueError(f"Invalid epsilon value: {eps}")
+        if weight_decay < 0.0:
+            raise ValueError(f"Invalid weight_decay value: {weight_decay}")
+        if tau <= 0.0:
+            raise ValueError(f"Invalid tau parameter: {tau}")
+        if not 0.0 < cautious_clamp_min <= 1.0:
+            raise ValueError(f"Invalid cautious_clamp_min: {cautious_clamp_min}")
+        if buffer not in (1, 2):
+            raise ValueError(f"Buffer count must be 1 (Single) or 2 (Dual), got: {buffer}")
+
+        defaults = dict(
+            lr=lr, betas=betas, eps=eps, weight_decay=weight_decay,
+            tau=tau, alpha_ss=alpha_ss, cautious=cautious,
+            cautious_clamp_min=cautious_clamp_min, buffer=buffer,
+            stochastic_rounding=stochastic_rounding, execution=execution
+        )
+        super().__init__(params, defaults)
+        self._scratch_tensors: Dict[torch.device, torch.Tensor] = {}
+
+    def _get_scratch_buffer(self, device: torch.device, count: int) -> torch.Tensor:
+        if device not in self._scratch_tensors or self._scratch_tensors[device].numel() < count:
+            self._scratch_tensors[device] = torch.zeros((count,), device=device, dtype=torch.float32)
+        buf = self._scratch_tensors[device][:count]
+        buf.zero_()
+        return buf
+
+    @torch.no_grad()
+    def step(self, closure: Optional[Callable[[], float]] = None) -> Optional[float]:
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            steps = []
+            buf_count = group["buffer"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                if p.grad.is_sparse:
+                    raise RuntimeError("LuminaV does not support sparse gradients.")
+
+                params_with_grad.append(p)
+                grads.append(p.grad)
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state["step"] = 0
+                    state["exp_avg"] = torch.zeros_like(p, memory_format=torch.contiguous_format)
+                    if buf_count == 2:
+                        state["exp_avg_sq"] = torch.zeros_like(p, memory_format=torch.contiguous_format)
+
+                exp_avgs.append(state["exp_avg"])
+                if buf_count == 2:
+                    if "exp_avg_sq" not in state:
+                        state["exp_avg_sq"] = torch.zeros_like(p, memory_format=torch.contiguous_format)
+                    exp_avg_sqs.append(state["exp_avg_sq"])
+
+                state["step"] += 1
+                steps.append(state["step"])
+
+            if not params_with_grad:
+                continue
+
+            lr = group["lr"]
+            beta1, beta2 = group["betas"]
+            eps = group["eps"]
+            weight_decay = group["weight_decay"]
+            tau = group["tau"]
+            alpha_ss = group["alpha_ss"]
+            cautious = group["cautious"]
+            clamp_min = group["cautious_clamp_min"]
+            use_sr = group["stochastic_rounding"]
+            exec_mode = group["execution"]
+
+            all_cuda = all(p.is_cuda for p in params_with_grad)
+
+            if exec_mode == "auto":
+                if HAS_TRITON and all_cuda:
+                    exec_mode = "triton"
+                elif hasattr(torch, "_foreach_mul_"):
+                    exec_mode = "foreach"
+                else:
+                    exec_mode = "single"
+
+            if exec_mode == "triton" and not (HAS_TRITON and all_cuda):
+                exec_mode = "foreach"
+
+            if exec_mode == "triton" and HAS_TRITON and all_cuda:
+                try:
+                    if buf_count == 2:
+                        self._triton_step_v2(params_with_grad, grads, exp_avgs, exp_avg_sqs, steps, lr, beta1, beta2, eps, weight_decay, tau, cautious, clamp_min, use_sr)
+                    else:
+                        self._triton_step_v1(params_with_grad, grads, exp_avgs, steps, lr, beta1, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, use_sr)
+                except Exception as e:
+                    logger.warning(f"Triton kernel fallback to C++ foreach engine: {e}")
+                    self._grouped_foreach_step(params_with_grad, grads, exp_avgs, exp_avg_sqs, steps, lr, beta1, beta2, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, buf_count, use_sr)
+            elif exec_mode == "foreach":
+                self._grouped_foreach_step(params_with_grad, grads, exp_avgs, exp_avg_sqs, steps, lr, beta1, beta2, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, buf_count, use_sr)
+            else:
+                if buf_count == 2:
+                    self._single_step_v2(params_with_grad, grads, exp_avgs, exp_avg_sqs, steps, lr, beta1, beta2, eps, weight_decay, tau, cautious, clamp_min, use_sr)
+                else:
+                    self._single_step_v1(params_with_grad, grads, exp_avgs, steps, lr, beta1, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, use_sr)
+
+        return loss
+
+    def _grouped_foreach_step(self, params, grads, exp_avgs, exp_avg_sqs, steps, lr, beta1, beta2, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, buf_count, use_sr):
+        groups: Dict[Tuple[torch.device, torch.dtype], List[int]] = {}
+        for idx, p in enumerate(params):
+            key = (p.device, p.dtype)
+            if key not in groups:
+                groups[key] = []
+            groups[key].append(idx)
+
+        for (dev, dt), indices in groups.items():
+            sub_params = [params[i] for i in indices]
+            sub_grads = [grads[i] for i in indices]
+            sub_exp_avgs = [exp_avgs[i] for i in indices]
+            sub_steps = [steps[i] for i in indices]
+
+            if buf_count == 2:
+                sub_exp_avg_sqs = [exp_avg_sqs[i] for i in indices]
+                self._foreach_step_v2(sub_params, sub_grads, sub_exp_avgs, sub_exp_avg_sqs, sub_steps, lr, beta1, beta2, eps, weight_decay, tau, cautious, clamp_min, use_sr)
+            else:
+                self._foreach_step_v1(sub_params, sub_grads, sub_exp_avgs, sub_steps, lr, beta1, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, use_sr)
+
+    def _triton_step_v2(self, params, grads, exp_avgs, exp_avg_sqs, steps, lr, beta1, beta2, eps, weight_decay, tau, cautious, clamp_min, use_sr):
+        BLOCK_SIZE = 1024
+        device = params[0].device
+        scratch = self._get_scratch_buffer(device, len(params))
+
+        for i in range(len(params)):
+            p, grad, exp_avg, exp_avg_sq, step = params[i], grads[i], exp_avgs[i], exp_avg_sqs[i], steps[i]
+            is_orig_contig = p.is_contiguous()
+            p_contig = p if is_orig_contig else p.contiguous()
+            grad_contig = grad if grad.is_contiguous() else grad.contiguous()
+            exp_avg_contig = exp_avg if exp_avg.is_contiguous() else exp_avg.contiguous()
+            exp_avg_sq_contig = exp_avg_sq if exp_avg_sq.is_contiguous() else exp_avg_sq.contiguous()
+
+            effective_eps = max(eps, 1e-4) if p.dtype == torch.float16 else eps
+            bc1 = 1.0 - (beta1**step)
+            bc2 = 1.0 - (beta2**step)
+            c1 = (bc1 * tau) / math.sqrt(bc2)
+            c2 = effective_eps * bc1 * tau
+
+            if p.dtype == torch.bfloat16:
+                dtype_mode = 1
+            elif p.dtype == torch.float16:
+                dtype_mode = 2
+            else:
+                dtype_mode = 0
+
+            seed = torch.randint(0, 2147483647, (1,)).item()
+            n_elements = p.numel()
+            grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+
+            if not cautious:
+                _lumina_v2_single_pass_kernel[grid](
+                    p_contig, grad_contig, exp_avg_contig, exp_avg_sq_contig,
+                    n_elements, beta1, beta2, c1, c2, lr, weight_decay,
+                    seed, dtype_mode, use_sr, BLOCK_SIZE=BLOCK_SIZE
+                )
+            else:
+                mask_sum_ptr = scratch[i : i + 1]
+                mask_sum_ptr.zero_()
+                _lumina_v2_pass1_kernel[grid](
+                    grad_contig, exp_avg_contig, exp_avg_sq_contig, mask_sum_ptr,
+                    n_elements, beta1, beta2, c1, c2, BLOCK_SIZE=BLOCK_SIZE
+                )
+                _lumina_v2_pass2_kernel[grid](
+                    p_contig, grad_contig, exp_avg_contig, exp_avg_sq_contig, mask_sum_ptr,
+                    n_elements, beta1, c1, c2, lr, weight_decay, clamp_min,
+                    seed, dtype_mode, use_sr, BLOCK_SIZE=BLOCK_SIZE
+                )
+
+            if not is_orig_contig:
+                p.copy_(p_contig)
+
+    def _triton_step_v1(self, params, grads, exp_avgs, steps, lr, beta1, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, use_sr):
+        BLOCK_SIZE = 1024
+        device = params[0].device
+        scratch_rms = self._get_scratch_buffer(device, len(params) * 2)
+
+        for i in range(len(params)):
+            p, grad, exp_avg, step = params[i], grads[i], exp_avgs[i], steps[i]
+            is_orig_contig = p.is_contiguous()
+            p_contig = p if is_orig_contig else p.contiguous()
+            grad_contig = grad if grad.is_contiguous() else grad.contiguous()
+            exp_avg_contig = exp_avg if exp_avg.is_contiguous() else exp_avg.contiguous()
+
+            effective_eps = max(eps, 1e-4) if p.dtype == torch.float16 else eps
+            bc1 = 1.0 - (beta1**step)
+            c2 = effective_eps * bc1 * tau
+
+            if p.dtype == torch.bfloat16:
+                dtype_mode = 1
+            elif p.dtype == torch.float16:
+                dtype_mode = 2
+            else:
+                dtype_mode = 0
+
+            seed = torch.randint(0, 2147483647, (1,)).item()
+            n_elements = p.numel()
+            grid = (triton.cdiv(n_elements, BLOCK_SIZE),)
+            rms_sum_ptr = scratch_rms[2 * i : 2 * i + 1]
+            mask_sum_ptr = scratch_rms[2 * i + 1 : 2 * i + 2]
+            rms_sum_ptr.zero_()
+            mask_sum_ptr.zero_()
+
+            _lumina_v1_pass1_kernel[grid](grad_contig, exp_avg_contig, rms_sum_ptr, n_elements, beta1, BLOCK_SIZE=BLOCK_SIZE)
+            _lumina_v1_pass2_kernel[grid](grad_contig, exp_avg_contig, rms_sum_ptr, mask_sum_ptr, n_elements, beta1, tau, eps, c2, alpha_ss, BLOCK_SIZE=BLOCK_SIZE)
+            _lumina_v1_pass3_update_kernel[grid](
+                p_contig, grad_contig, exp_avg_contig, rms_sum_ptr, mask_sum_ptr,
+                n_elements, beta1, tau, eps, c2, alpha_ss, lr, weight_decay, clamp_min,
+                seed, dtype_mode, use_sr, cautious, BLOCK_SIZE=BLOCK_SIZE
+            )
+
+            if not is_orig_contig:
+                p.copy_(p_contig)
+
+    def _single_step_v2(self, params, grads, exp_avgs, exp_avg_sqs, steps, lr, beta1, beta2, eps, weight_decay, tau, cautious, clamp_min, use_sr):
+        for i in range(len(params)):
+            p, grad, exp_avg, exp_avg_sq, step = params[i], grads[i], exp_avgs[i], exp_avg_sqs[i], steps[i]
+            effective_eps = max(eps, 1e-4) if p.dtype == torch.float16 else eps
+            bc1 = 1.0 - (beta1**step)
+            bc2 = 1.0 - (beta2**step)
+            c1 = (bc1 * tau) / math.sqrt(bc2)
+            c2 = effective_eps * bc1 * tau
+
+            if weight_decay != 0.0:
+                p.mul_(1.0 - lr * weight_decay)
+
+            exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+            nes_m = torch.mul(exp_avg, beta1).add_(grad, alpha=1.0 - beta1)
+            grad_diff = grad - exp_avg
+
+            if p.dtype == torch.float16:
+                diff_f32 = grad_diff.float()
+                v_f32 = exp_avg_sq.float().mul_(beta2).addcmul_(diff_f32, diff_f32, value=1.0 - beta2).clamp_(0.0, 65000.0)
+                exp_avg_sq.copy_(v_f32.to(dtype=p.dtype))
+            else:
+                exp_avg_sq.mul_(beta2).addcmul_(grad_diff, grad_diff, value=1.0 - beta2)
+
+            sigma = exp_avg_sq.float().sqrt().mul_(c1).add_(c2)
+            update = torch.tanh(nes_m.float() / (sigma + 1e-6)).to(dtype=p.dtype)
+
+            if cautious:
+                mask = (update * grad > 0).to(dtype=grad.dtype)
+                mask_scale = mask.float().mean().clamp_(min=clamp_min, max=1.0).to(dtype=grad.dtype)
+                update = update.mul_(mask).div_(mask_scale)
+
+            _sr_update(p, update, lr, use_sr)
+
+    def _single_step_v1(self, params, grads, exp_avgs, steps, lr, beta1, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, use_sr):
+        for i in range(len(params)):
+            p, grad, exp_avg, step = params[i], grads[i], exp_avgs[i], steps[i]
+            if weight_decay != 0.0:
+                p.mul_(1.0 - lr * weight_decay)
+
+            exp_avg.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+            nes_m = torch.mul(exp_avg, beta1).add_(grad, alpha=1.0 - beta1)
+            effective_eps = max(eps, 1e-4) if p.dtype == torch.float16 else eps
+            bc1 = 1.0 - (beta1**step)
+            c2 = effective_eps * bc1 * tau
+
+            rms = torch.sqrt(nes_m.float().square().mean() + effective_eps)
+            sigma = rms * tau + c2
+            z = nes_m.float() / (sigma + 1e-6)
+            z_soft = z / (1.0 + alpha_ss * torch.abs(z))
+            update = torch.tanh(z_soft).to(dtype=p.dtype)
+
+            if cautious:
+                mask = (update * grad > 0).to(dtype=grad.dtype)
+                mask_scale = mask.float().mean().clamp_(min=clamp_min, max=1.0).to(dtype=grad.dtype)
+                update = update.mul_(mask).div_(mask_scale)
+
+            _sr_update(p, update, lr, use_sr)
+
+    def _foreach_step_v2(self, params, grads, exp_avgs, exp_avg_sqs, steps, lr, beta1, beta2, eps, weight_decay, tau, cautious, clamp_min, use_sr):
+        if weight_decay != 0.0:
+            torch._foreach_mul_(params, 1.0 - lr * weight_decay)
+
+        torch._foreach_mul_(exp_avgs, beta1)
+        torch._foreach_add_(exp_avgs, grads, alpha=1.0 - beta1)
+
+        nes_m_list = torch._foreach_mul(exp_avgs, beta1)
+        torch._foreach_add_(nes_m_list, grads, alpha=1.0 - beta1)
+
+        grad_diff_list = torch._foreach_sub(grads, exp_avgs)
+        is_fp16 = params[0].dtype == torch.float16
+        if is_fp16:
+            for i in range(len(params)):
+                d_f32 = grad_diff_list[i].float()
+                v_f32 = exp_avg_sqs[i].float().mul_(beta2).addcmul_(d_f32, d_f32, value=1.0 - beta2).clamp_(0.0, 65000.0)
+                exp_avg_sqs[i].copy_(v_f32.to(dtype=torch.float16))
+        else:
+            torch._foreach_mul_(exp_avg_sqs, beta2)
+            torch._foreach_addcmul_(exp_avg_sqs, grad_diff_list, grad_diff_list, value=1.0 - beta2)
+
+        effective_eps = max(eps, 1e-4) if is_fp16 else eps
+        bias_correction1 = [1.0 - (beta1**st) for st in steps]
+        bias_correction2 = [1.0 - (beta2**st) for st in steps]
+        c1_list = [(bc1 * tau) / math.sqrt(bc2) for bc1, bc2 in zip(bias_correction1, bias_correction2)]
+        c2_list = [effective_eps * bc1 * tau for bc1 in bias_correction1]
+
+        updates = []
+        for i in range(len(params)):
+            sigma = exp_avg_sqs[i].float().sqrt().mul_(c1_list[i]).add_(c2_list[i])
+            u = torch.tanh(nes_m_list[i].float() / (sigma + 1e-6)).to(dtype=params[i].dtype)
+            if cautious:
+                mask = (u * grads[i] > 0).to(dtype=grads[i].dtype)
+                scale = mask.float().mean().clamp_(min=clamp_min, max=1.0).to(dtype=grads[i].dtype)
+                u = u.mul(mask).div(scale)
+            updates.append(u)
+
+        if use_sr and params[0].dtype in (torch.bfloat16, torch.float16):
+            for i in range(len(params)):
+                _sr_update(params[i], updates[i], lr, use_sr=True)
+        else:
+            torch._foreach_add_(params, updates, alpha=-lr)
+
+    def _foreach_step_v1(self, params, grads, exp_avgs, steps, lr, beta1, eps, weight_decay, tau, alpha_ss, cautious, clamp_min, use_sr):
+        if weight_decay != 0.0:
+            torch._foreach_mul_(params, 1.0 - lr * weight_decay)
+
+        torch._foreach_mul_(exp_avgs, beta1)
+        torch._foreach_add_(exp_avgs, grads, alpha=1.0 - beta1)
+
+        nes_m_list = torch._foreach_mul(exp_avgs, beta1)
+        torch._foreach_add_(nes_m_list, grads, alpha=1.0 - beta1)
+
+        is_fp16 = params[0].dtype == torch.float16
+        effective_eps = max(eps, 1e-4) if is_fp16 else eps
+        bias_correction1 = [1.0 - (beta1**st) for st in steps]
+        c2_list = [effective_eps * bc1 * tau for bc1 in bias_correction1]
+
+        updates = []
+        for i in range(len(params)):
+            m = nes_m_list[i]
+            rms = torch.sqrt(m.float().square().mean() + effective_eps)
+            sigma = rms * tau + c2_list[i]
+            z = m.float() / (sigma + 1e-6)
+            z_soft = z / (1.0 + alpha_ss * torch.abs(z))
+            u = torch.tanh(z_soft).to(dtype=params[i].dtype)
+            if cautious:
+                mask = (u * grads[i] > 0).to(dtype=grads[i].dtype)
+                scale = mask.float().mean().clamp_(min=clamp_min, max=1.0).to(dtype=grads[i].dtype)
+                u = u.mul(mask).div(scale)
+            updates.append(u)
+
+        if use_sr and params[0].dtype in (torch.bfloat16, torch.float16):
+            for i in range(len(params)):
+                _sr_update(params[i], updates[i], lr, use_sr=True)
+        else:
+            torch._foreach_add_(params, updates, alpha=-lr)

luminav-1.1.5.dist-info/METADATA

@@ -0,0 +1,202 @@
+Metadata-Version: 2.4
+Name: luminav
+Version: 1.1.5
+Summary: A master-free, memory-efficient adaptive optimizer for low-precision deep learning
+Author: Silver Moon (cloverxion)
+Maintainer: Lumina Moon
+License: Apache-2.0
+Project-URL: Homepage, https://huggingface.co/cloverx-id/LuminaV-Optimizer-Paper
+Project-URL: Author Profile, https://huggingface.co/cloverxion
+Project-URL: Organization, https://huggingface.co/Lumina-Moon
+Project-URL: Software DOI, https://doi.org/10.57967/hf/10365
+Project-URL: Paper DOI, https://doi.org/10.57967/hf/10270
+Project-URL: Bug Tracker, https://huggingface.co/cloverx-id/LuminaV-Optimizer-Paper/discussions
+Project-URL: Changelog, https://huggingface.co/cloverx-id/LuminaV-Optimizer-Paper/blob/main/CHANGELOG.md
+Keywords: pytorch,optimizer,low-vram,fp16,bf16,triton,deep-learning,llm,transformer
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.0.0
+Provides-Extra: triton
+Requires-Dist: triton; extra == "triton"
+Dynamic: license-file
+
+# LuminaV Optimizer
+
+> **We Were Too Broke for AdamW So We Trapped Gradients in a Hyperbolic Straitjacket and Hired a Traffic Cop to Slap Them**
+
+[![PyPI - Version](https://img.shields.io/pypi/v/luminav.svg?style=flat-square)](https://pypi.org/project/luminav/)
+[![Software DOI](https://img.shields.io/badge/Software_DOI-10.57967%2Fhf%2F10365-purple?style=flat-square)](https://doi.org/10.57967/hf/10365)
+[![Paper DOI](https://img.shields.io/badge/Paper_DOI-10.57967%2Fhf%2F10270-darkblue?style=flat-square)](https://doi.org/10.57967/hf/10270)
+[![Python Version](https://img.shields.io/badge/python-3.10+-blue?style=flat-square)](https://pypi.org/project/luminav/)
+[![License](https://img.shields.io/badge/License-Apache%202.0-yellow?style=flat-square)](LICENSE)
+
+---
+
+## Official Repository & Research Paper
+
+- **Official Repository & Issue Tracker:** [https://huggingface.co/cloverx-id/LuminaV-Optimizer-Paper](https://huggingface.co/cloverx-id/LuminaV-Optimizer-Paper)
+- **Primary Research Paper Archive:** [https://huggingface.co/cloverx-id/XoneLM-1.0-Paper](https://huggingface.co/cloverx-id/XoneLM-1.0-Paper)
+- **Software DOI:** [10.57967/hf/10365](https://doi.org/10.57967/hf/10365)
+- **Paper DOI:** [10.57967/hf/10270](https://doi.org/10.57967/hf/10270)
+
+---
+
+## What's New in v1.1.5 (Latest Release)
+
+The **v1.1.5** release introduces the **Zero-VRAM FP16 Numerical Safety Shield**, resolving early pretraining instability and eliminating `NaN` collapses in pure half-precision training without allocating 4-byte master weights:
+
+- **Dynamic Epsilon Floor for FP16:** Automatically floors effective epsilon to `max(eps, 1e-4)` for `torch.float16` parameters, preventing the bias-corrected scale factor `c2` from underflowing into IEEE-754 subnormal/zero limits (~5.96e-8) and eliminating step-1 `0.0 / 0.0 = NaN` errors on zero-gradient parameters (e.g., unselected vocabulary tokens, padding, or dropout paths).
+- **On-Chip Second-Moment Clamping:** Added an upper-bound register ceiling (`65,000.0`) to the centered innovation variance in Triton kernels and PyTorch loops before storing into FP16 pointers. This prevents unclipped initial gradient spikes (`|g| > 256`) from exceeding the FP16 maximum dynamic range (65,504) and permanently saturating state buffers to `+inf`.
+- **Zero-Sigma Division Guard:** Embedded an explicit `sigma + 1e-6` protection across all Triton kernel passes and PyTorch fallbacks to ensure division safety during cold-start iterations.
+- **Native Hardware Capability Detection:** Integrated `_supports_native_bf16()` to automatically detect Nvidia Ampere SM80+ architectures for optimal native execution.
+- **Empirically Proven on Tesla T4:** Empirically verified by pretraining `SmolLM-135M` from scratch on Wikimedia Wikipedia in pure FP16 on an Nvidia Tesla T4 GPU for 100 consecutive steps with zero `NaN` occurrences—even with external gradient clipping completely disabled.
+
+---
+
+## Overview
+
+**LuminaV** is a master-free, memory-efficient adaptive optimizer engineered specifically for deep learning workloads running directly in low precision (`FP16` / `BF16`) without maintaining redundant 4-byte FP32 master weights.
+
+By combining **Centered Innovation Variance**, **Hyperbolic Tangent (tanh) Coordinate Bounding**, a **Directional Traffic-Cop Mask**, and **On-Chip Bitwise Stochastic Rounding**, LuminaV eliminates the standard 16-byte-per-parameter memory tax imposed by AdamW while preventing weight stagnation and numerical explosions.
+
+---
+
+## Key Features
+
+1. **Zero Master-Weight Copies:** Directly mutates parameter weights in native `FP16` or `BF16`, eliminating the redundant 4-byte FP32 master weight allocation.
+2. **On-Chip Bitwise Stochastic Rounding (SR):** Implements in-register bitcast hashing in Triton to provide unbiased stochastic rounding, preventing weight stagnation during fine-grained updates or learning rate decay.
+3. **Hyperbolic tanh Bounding Envelope:** Maps normalized momentum through a `(-1.0, 1.0)` transfer function, guaranteeing coordinate updates cannot explode beyond the step learning rate.
+4. **The Traffic-Cop Directional Gate:** Dynamically eliminates coordinate updates whenever historical momentum conflicts with the incoming mini-batch gradient direction (`u_t · g_t <= 0`).
+5. **Centered Innovation Variance:** Tracks centered innovation dispersion `(g_t - m_t)^2` rather than uncentered raw second moments, suppressing variance inflation during confident descent.
+6. **Dual Execution Engine:** Fully accelerated custom OpenAI Triton kernels for CUDA devices, paired with vectorized C++ `torch._foreach` multi-tensor fallbacks.
+
+---
+
+## Installation
+
+Install directly via PyPI:
+
+```bash
+pip install luminav
+```
+
+For GPU acceleration via OpenAI Triton:
+
+```bash
+pip install luminav[triton]
+```
+
+Or install in editable mode from source:
+
+```bash
+git clone https://huggingface.co/cloverx-id/LuminaV-Optimizer-Paper
+cd LuminaV-Optimizer-Paper
+pip install -e .
+```
+
+---
+
+## Quickstart
+
+```python
+import torch
+import torch.nn as nn
+from luminav import LuminaV
+
+# 1. Instantiate model directly in native half precision (e.g. BF16 or FP16)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = nn.Linear(1024, 1024).to(device=device, dtype=torch.float16)
+
+# 2. Initialize LuminaV
+optimizer = LuminaV(
+    model.parameters(),
+    lr=8e-4,
+    betas=(0.9, 0.999),
+    eps=1e-8,
+    weight_decay=0.08,
+    tau=0.8,
+    buffer=2,                   # 2 = Dual-Buffer (Standard), 1 = Single-Buffer (Low VRAM)
+    stochastic_rounding=True,
+    execution="auto"
+)
+
+# 3. Standard training loop
+data = torch.randn(32, 1024, device=device, dtype=torch.float16)
+target = torch.randn(32, 1024, device=device, dtype=torch.float16)
+criterion = nn.MSELoss()
+
+optimizer.zero_grad(set_to_none=True)
+output = model(data)
+loss = criterion(output, target)
+loss.backward()
+
+optimizer.step()
+```
+
+---
+
+## Parameter Reference
+
+| Parameter | Type | Default | Description |
+| :--- | :--- | :--- | :--- |
+| `params` | `iterable` | *Required* | Iterable of parameters to optimize or dicts defining parameter groups. |
+| `lr` | `float` | `8e-4` | Learning rate (eta). |
+| `betas` | `Tuple[float, float]` | `(0.9, 0.999)` | Coefficients (beta1, beta2) for running momentum and centered innovation variance. |
+| `eps` | `float` | `1e-8` | Numerical stability term (epsilon). Automatically floored to `1e-4` in FP16 to prevent subnormal underflow. |
+| `weight_decay` | `float` | `8e-2` | Decoupled weight decay coefficient (lambda). |
+| `tau` | `float` | `0.8` | Analytical bias correction temperature parameter (tau). |
+| `alpha_ss` | `float` | `0.5` | Softsign dampening factor (alpha_ss) used in single-buffer mode (`buffer=1`). |
+| `cautious` | `bool` | `True` | If `True`, enables Traffic-Cop directional verification masking. |
+| `cautious_clamp_min` | `float` | `0.2` | Safety floor density clamp (gamma_min) preventing division by zero in masked normalization. |
+| `buffer` | `int` | `2` | Buffer mode: `2` (Dual-buffer tracking m_t and v_t) or `1` (Single-buffer scalar RMS tracking). |
+| `stochastic_rounding` | `bool` | `True` | Enables bitwise stochastic rounding on native FP16/BF16 weights. |
+| `execution` | `str` | `"auto"` | Execution engine: `"auto"`, `"triton"`, `"foreach"`, or `"single"`. |
+
+---
+
+## Citation
+
+If you utilize LuminaV in your research or applications, please cite both the foundational paper and this software implementation:
+
+```bibtex
+# 1. To cite the official research paper & theoretical mechanics
+@misc{luminamoon2026luminav_paper,
+  author       = {{Silver Moon (cloverxion)}},
+  organization = {Lumina Moon},
+  title        = {{LuminaV: We Were Too Broke for AdamW So We Trapped Gradients in a Hyperbolic Straitjacket and Hired a Traffic Cop to Slap Them}},
+  year         = {2026},
+  publisher    = {Hugging Face},
+  doi          = {10.57967/hf/10270},
+  url          = {https://huggingface.co/cloverx-id/XoneLM-1.0-Paper}
+}
+
+# 2. To cite this software implementation & standalone codebase
+@software{luminamoon2026luminav_code,
+  author       = {{Silver Moon (cloverxion)}},
+  organization = {Lumina Moon},
+  title        = {{LuminaV Optimizer: Official PyTorch Implementation}},
+  year         = {2026},
+  publisher    = {Hugging Face},
+  version      = {1.1.5},
+  doi          = {10.57967/hf/10365},
+  url          = {https://huggingface.co/cloverx-id/LuminaV-Optimizer-Paper}
+}
+```
+
+---
+
+## License
+
+Apache License 2.0. See [LICENSE](LICENSE) for full terms.

luminav-1.1.5.dist-info/RECORD

@@ -0,0 +1,7 @@
+luminav/__init__.py,sha256=9QXx2WxRhVYNY5FGYkTXGKumRiS6dOvVE5qLlRxaYUc,134
+luminav/luminav.py,sha256=HhZG-1igoI4Q_mpiwtNBOzOWSK16l8prKA3_L_pN7Zg,28360
+luminav-1.1.5.dist-info/licenses/LICENSE,sha256=vfIk0HOJipiGkxSJLASwBoUIc4mkDsspxhxb_5DpKVQ,762
+luminav-1.1.5.dist-info/METADATA,sha256=5wL-XC_rpHWSXOFB8iBmCeRO3QhUlloGDG_bcLeSW9U,10086
+luminav-1.1.5.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+luminav-1.1.5.dist-info/top_level.txt,sha256=8yzfkXhUGr1S4izcM8vOu7nw7MfZRRwNAJu3mto91O0,8
+luminav-1.1.5.dist-info/RECORD,,

luminav-1.1.5.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

luminav-1.1.5.dist-info/licenses/LICENSE

@@ -0,0 +1,17 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   Copyright (c) 2026 Lumina Moon and Contributors.
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

luminav-1.1.5.dist-info/top_level.txt

@@ -0,0 +1 @@
+luminav