d9d 0.1.0__py3-none-any.whl

d9d/kernel/stochastic/__init__.py

@@ -0,0 +1,11 @@
+"""
+Utilities for stochastic type casting (e.g., FP32 to BF16).
+"""
+
+from .adamw_step import adamw_stochastic_bf16_
+from .copy import copy_fp32_to_bf16_stochastic_
+
+__all__ = [
+    "adamw_stochastic_bf16_",
+    "copy_fp32_to_bf16_stochastic_"
+]

d9d/kernel/stochastic/adamw_step.py

@@ -0,0 +1,204 @@
+import torch
+import triton
+import triton.language as tl
+
+from .ops import fp32_to_bf16_kernel
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
+        triton.Config({"BLOCK_SIZE": 2048}, num_warps=4),
+        triton.Config({"BLOCK_SIZE": 2048}, num_warps=8),
+        triton.Config({"BLOCK_SIZE": 4096}, num_warps=8),
+        triton.Config({"BLOCK_SIZE": 8192}, num_warps=8),
+    ],
+    key=["n_elements"],
+    restore_value=["p_ptr", "m_ptr", "v_ptr"]
+)
+@triton.jit
+def _adamw_stochastic_bf16_kernel(
+        p_ptr: tl.tensor,  # Pointer to parameters (Always BF16 -> read/write)
+        g_ptr: tl.tensor,  # Pointer to gradients  (BF16 or FP32 -> read only)
+        m_ptr: tl.tensor,  # Pointer to exp_avg    (BF16 or FP32 -> read/write)
+        v_ptr: tl.tensor,  # Pointer to exp_avg_sq (BF16 or FP32 -> read/write)
+        n_elements: int,  # Total number of elements
+        lr: float,  # Learning rate
+        beta1: float,
+        beta2: float,
+        eps: float,
+        weight_decay: float,
+        step: int,  # Current step (for bias correction)
+        seed: int,  # Random seed for stochastic rounding
+        BLOCK_SIZE: tl.constexpr,
+        GRAD_IS_BF16: tl.constexpr,  # noqa: N803
+        STATE_IS_BF16: tl.constexpr  # noqa: N803
+):
+    pid = tl.program_id(axis=0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # load parameters
+    p_bf16 = tl.load(p_ptr + offsets, mask=mask)
+    p_fp32 = p_bf16.to(tl.float32)
+
+    # load grad
+    if GRAD_IS_BF16:
+        g_fp32 = tl.load(g_ptr + offsets, mask=mask).to(tl.float32)
+    else:
+        g_fp32 = tl.load(g_ptr + offsets, mask=mask)
+
+    # load states
+    if STATE_IS_BF16:
+        m_curr = tl.load(m_ptr + offsets, mask=mask).to(tl.float32)
+        v_curr = tl.load(v_ptr + offsets, mask=mask).to(tl.float32)
+    else:
+        m_curr = tl.load(m_ptr + offsets, mask=mask)
+        v_curr = tl.load(v_ptr + offsets, mask=mask)
+
+    # now the math goes in fp32
+
+    # do weight decay
+    p_fp32 = p_fp32 * (1.0 - lr * weight_decay)
+
+    # update moments
+    m_next = beta1 * m_curr + (1.0 - beta1) * g_fp32
+    v_next = beta2 * v_curr + (1.0 - beta2) * (g_fp32 * g_fp32)
+
+    # bias correction
+    bias_correction1 = 1.0 - tl.exp(step * tl.log(beta1))
+    bias_correction2 = 1.0 - tl.exp(step * tl.log(beta2))
+
+    m_hat = m_next / bias_correction1
+    v_hat = v_next / bias_correction2
+
+    # compute update
+    update = (lr * m_hat) / (tl.sqrt(v_hat) + eps)
+
+    p_new_fp32 = p_fp32 - update
+
+    # and now we store...
+    # p -> always stochastic fp32 -> bf16
+    # states -> depending on constexprs
+    p_new_bf16 = fp32_to_bf16_kernel(p_new_fp32, offsets, seed)
+    tl.store(p_ptr + offsets, p_new_bf16, mask=mask)
+
+    if STATE_IS_BF16:
+        m_next_bf16 = fp32_to_bf16_kernel(m_next, offsets, seed + 42)
+        v_next_bf16 = fp32_to_bf16_kernel(v_next, offsets, seed + 67)
+
+        tl.store(m_ptr + offsets, m_next_bf16, mask=mask)
+        tl.store(v_ptr + offsets, v_next_bf16, mask=mask)
+    else:
+        tl.store(m_ptr + offsets, m_next, mask=mask)
+        tl.store(v_ptr + offsets, v_next, mask=mask)
+
+
+def adamw_stochastic_bf16_(  # noqa: C901
+        params: torch.Tensor,
+        grads: torch.Tensor,
+        exp_avg: torch.Tensor,
+        exp_avg_sq: torch.Tensor,
+        lr: float,
+        beta1: float,
+        beta2: float,
+        eps: float,
+        weight_decay: float,
+        step: int,
+        generator: torch.Generator | None = None
+) -> None:
+    """
+    Performs a single in-place AdamW optimization step.
+
+    It is specifically designed for scenarios where parameters are stored in BFloat16.
+
+    To mitigate precision loss during the parameter update, it utilizes stochastic rounding when casting
+    FP32 calculation results back to BFloat16.
+
+    This function supports mixed precision for gradients and optimizer states (they can be
+    either FP32 or BFloat16).
+
+    Args:
+        params: The tensor of model parameters to update. Must be BFloat16 and contiguous.
+        grads: The gradient tensor.
+        exp_avg: The exponential moving average of gradient values (first moment).
+        exp_avg_sq: The exponential moving average of squared gradient values (second moment).
+        lr: The learning rate.
+        beta1: Decay rate for the first moment estimate.
+        beta2: Decay rate for the second moment estimate.
+        eps: Term added to the denominator to improve numerical stability.
+        weight_decay: Weight decay coefficient.
+        step: The current optimization step count, used for bias correction.
+        generator: PyTorch random number generator used to create the seed for stochastic rounding.
+
+    Raises:
+        ValueError: If main parameters are not BFloat16, if input tensor shapes do not match,
+            if input tensors are not contiguous (for those that require in-place modification),
+            if the optimizer states (exp_avg, exp_avg_sq) have different dtypes.
+    """
+
+    # check shape equality
+    if grads.shape != params.shape:
+        raise ValueError("Shape mismatch between grads and params.")
+
+    if exp_avg.shape != params.shape:
+        raise ValueError("Shape mismatch between exp_avg state and params.")
+
+    if exp_avg_sq.shape != params.shape:
+        raise ValueError("Shape mismatch between exp_avg_sq state and params.")
+
+    # check params
+    if params.dtype != torch.bfloat16:
+        raise ValueError("Params must be BFloat16 for this kernel.")
+
+    if not params.is_contiguous():
+        raise ValueError("Params must be contiguous since it is an in-place kernel.")
+
+    # check grads
+    if not grads.is_contiguous():
+        grads = grads.contiguous()
+
+    # check states
+    if not exp_avg.is_contiguous():
+        raise ValueError("Exp_avg state must be contiguous since it is an in-place kernel.")
+
+    if not exp_avg_sq.is_contiguous():
+        raise ValueError("Exp_avg_sq state must be contiguous since it is an in-place kernel.")
+
+    if exp_avg.dtype != exp_avg_sq.dtype:
+        raise ValueError("States have different dtypes.")
+
+    n_elements = params.numel()
+
+    grad_is_bf16 = (grads.dtype == torch.bfloat16)
+    state_is_bf16 = (exp_avg.dtype == torch.bfloat16)
+
+    # Generate random seed
+    seed = torch.randint(
+        0, 2 ** 31 - 1, (1,),
+        device="cpu",
+        generator=generator
+    ).item()
+
+    def _grid(meta: dict[str, int]) -> tuple[int, ...]:
+        return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    _adamw_stochastic_bf16_kernel[_grid](
+        params,
+        grads,
+        exp_avg,
+        exp_avg_sq,
+
+        n_elements,
+
+        lr,
+        beta1,
+        beta2,
+        eps,
+        weight_decay,
+        step,
+        seed,
+
+        GRAD_IS_BF16=grad_is_bf16,
+        STATE_IS_BF16=state_is_bf16
+    )

d9d/kernel/stochastic/copy.py

@@ -0,0 +1,104 @@
+import torch
+import triton
+import triton.language as tl
+
+from .ops import fp32_to_bf16_kernel
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
+        triton.Config({"BLOCK_SIZE": 2048}, num_warps=4),
+        triton.Config({"BLOCK_SIZE": 2048}, num_warps=8),
+        triton.Config({"BLOCK_SIZE": 4096}, num_warps=8),
+        triton.Config({"BLOCK_SIZE": 8192}, num_warps=8),
+    ],
+    key=["n_elements"]
+)
+@triton.jit
+def _copy_fp32_to_bf16_kernel(
+        source_ptr: torch.Tensor,
+        target_ptr: torch.Tensor,
+        n_elements: int,
+        seed: int,
+        BLOCK_SIZE: tl.constexpr
+):
+    pid = tl.program_id(axis=0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # load source value (fp32)
+    val_fp32 = tl.load(source_ptr + offsets, mask=mask)
+
+    val_bf16 = fp32_to_bf16_kernel(
+        val_fp32=val_fp32,
+        offsets=offsets,
+        seed=seed
+    )
+
+    tl.store(target_ptr + offsets, val_bf16, mask=mask)
+
+
+def copy_fp32_to_bf16_stochastic_(
+        target: torch.Tensor,
+        source: torch.Tensor,
+        generator: torch.Generator | None = None
+) -> torch.Tensor:
+    """
+    Copies elements from a Float32 tensor to a BFloat16 tensor using stochastic rounding.
+
+    Unlike standard round-to-nearest casting, stochastic rounding probabilistically rounds
+    numbers up or down based on the value of the bits being truncated. This preserves the
+    expected value of the tensor (E[round(x)] = x), which is crucial for accumulating
+    gradients or parameters in low precision without stagnation.
+
+    This operation is performed in-place on the target tensor.
+
+    Args:
+        target: The output tensor where results are written. Must be of type BFloat16
+            and contiguous.
+        source: The input tensor containing values to copy. Must be of type Float32.
+        generator: An optional PyTorch RNG generator to strictly control the random
+            noise used for rounding.
+
+    Returns:
+        The target tensor, modified in-place.
+
+    Raises:
+        ValueError: If target is not contiguous, if source/target shapes do not match,
+            or if dtypes are not FP32 and BF16 respectively.
+    """
+
+    if not source.is_contiguous():
+        source = source.contiguous()
+
+    if not target.is_contiguous():
+        raise ValueError("Since this is an in-place operation, target should be a contiguous tensor!")
+
+    if source.shape != target.shape:
+        raise ValueError("Source and Target Tensors are of different shapes")
+
+    if source.dtype != torch.float32:
+        raise ValueError("Source must be Float32")
+    if target.dtype != torch.bfloat16:
+        raise ValueError("Target must be BFloat16")
+
+    n_elements = source.numel()
+
+    # Generate a random seed for this specific kernel launch
+    seed = torch.randint(
+        0, 2 ** 31 - 1, (1,),
+        device="cpu",
+        generator=generator
+    ).item()
+
+    def _grid(meta: dict[str, int]) -> tuple[int, ...]:
+        return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    _copy_fp32_to_bf16_kernel[_grid](
+        source,
+        target,
+        n_elements,
+        seed
+    )
+    return target

d9d/kernel/stochastic/ops/__init__.py

@@ -0,0 +1,5 @@
+from .round import fp32_to_bf16_kernel
+
+__all__ = [
+    "fp32_to_bf16_kernel"
+]

d9d/kernel/stochastic/ops/round.py

@@ -0,0 +1,22 @@
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def fp32_to_bf16_kernel(
+        val_fp32: tl.tensor,
+        offsets: tl.tensor,
+        seed: int,
+) -> tl.tensor:
+    val_ui32 = val_fp32.to(tl.uint32, bitcast=True)
+
+    # create random noise for last bits
+    rand_val = tl.randint(seed, offsets)
+    noise = rand_val.to(tl.uint32) & 0xFFFF
+
+    # add this noise (FP32)
+    val_ui32_noisy = val_ui32 + noise
+
+    # save in 16 bits
+    bf16_bits = (val_ui32_noisy >> 16).to(tl.int16)
+    return bf16_bits.to(tl.bfloat16, bitcast=True)

d9d/kernel/swiglu/__init__.py

@@ -0,0 +1,5 @@
+from .function import silu_mul
+
+__all__ = [
+    "silu_mul"
+]

d9d/kernel/swiglu/function.py

@@ -0,0 +1,36 @@
+from typing import Any
+
+import torch
+from torch.autograd import Function
+
+from .op import silu_mul_backward, silu_mul_forward
+
+
+class SiLUMulFunction(Function):
+    """
+    Autograd function for the fused silu(x)*y operation.
+    """
+
+    @staticmethod
+    def forward(ctx: Any, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        ctx.save_for_backward(x, y)
+        return silu_mul_forward(x, y)
+
+    @staticmethod
+    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x, y = ctx.saved_tensors
+        return silu_mul_backward(grad_output, x, y)
+
+
+def silu_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    """
+    Applies the SiLU multiplication operation: SiLU(x) * y.
+
+    Args:
+        x: Input tensor x.
+        y: Input tensor y.
+
+    Returns:
+        The resulting tensor of the same shape as inputs.
+    """
+    return SiLUMulFunction.apply(x, y)

d9d/kernel/swiglu/op.py

@@ -0,0 +1,167 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
+        triton.Config({"BLOCK_SIZE": 2048}, num_warps=4),
+        triton.Config({"BLOCK_SIZE": 2048}, num_warps=8),
+        triton.Config({"BLOCK_SIZE": 4096}, num_warps=8),
+        triton.Config({"BLOCK_SIZE": 8192}, num_warps=8),
+    ],
+    key=["n_elements"]
+)
+@triton.jit
+def _silu_mul_kernel(
+        x_ptr: torch.Tensor,
+        y_ptr: torch.Tensor,
+        out_ptr: torch.Tensor,
+        n_elements: int,
+        BLOCK_SIZE: tl.constexpr,
+):
+    # prepare
+    pid = tl.program_id(axis=0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # read
+    x = tl.load(x_ptr + offsets, mask=mask)
+    x_fp32 = x.to(tl.float32)  # sigmoid wants fp32
+    y = tl.load(y_ptr + offsets, mask=mask)
+
+    # compute
+    # cast back to match with torch
+    silu_x = (x_fp32 * tl.sigmoid(x_fp32)).cast(y.dtype)
+    out = silu_x * y
+
+    # write
+    tl.store(out_ptr + offsets, out, mask=mask)
+
+
+def silu_mul_forward(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    """
+    Computes the forward pass of silu(x)*y using Triton.
+
+    Args:
+        x: Input tensor x.
+        y: Input tensor y.
+
+    Returns:
+        The output tensor.
+
+    Raises:
+        ValueError: If inputs x and y do not match in shape or device.
+    """
+
+    if x.shape != y.shape or x.device != y.device:
+        raise ValueError("Inputs x and y must have the same shape, be on same device.")
+
+    if not x.is_contiguous():
+        x = x.contiguous()
+    if not y.is_contiguous():
+        y = y.contiguous()
+
+    n_elements = x.numel()
+    out = torch.empty_like(x)
+
+    def _grid(meta: dict[str, int]) -> tuple[int, ...]:
+        return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    _silu_mul_kernel[_grid](
+        x, y, out,
+        n_elements
+    )
+
+    return out
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_SIZE": 1024}, num_warps=4),
+        triton.Config({"BLOCK_SIZE": 2048}, num_warps=4),
+        triton.Config({"BLOCK_SIZE": 2048}, num_warps=8),
+        triton.Config({"BLOCK_SIZE": 4096}, num_warps=8),
+        triton.Config({"BLOCK_SIZE": 8192}, num_warps=8),
+    ],
+    key=["n_elements"]
+)
+@triton.jit
+def _silu_mul_backward_kernel(
+        grad_out_ptr: torch.Tensor,
+        x_ptr: torch.Tensor,
+        y_ptr: torch.Tensor,
+        grad_x_ptr: torch.Tensor,
+        grad_y_ptr: torch.Tensor,
+        n_elements: int,
+        BLOCK_SIZE: tl.constexpr
+):
+    # prepare
+    pid = tl.program_id(0)
+    block_start = pid * BLOCK_SIZE
+    offsets = block_start + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < n_elements
+
+    # read
+    dout = tl.load(grad_out_ptr + offsets, mask=mask)
+    x = tl.load(x_ptr + offsets, mask=mask).to(tl.float32)  # sigmoid wants fp32
+    y = tl.load(y_ptr + offsets, mask=mask)
+
+    # Recompute Silu components
+    sig_x = tl.sigmoid(x)
+    silu_x = x * sig_x
+
+    # Compute grad_y
+    # dy = dout * silu(x)
+    dx_silu_x = dout * silu_x  # Reuse this variable name logic
+    tl.store(grad_y_ptr + offsets, dx_silu_x, mask=mask)
+
+    # Compute grad_x
+    # silu'(x) = sigmoid(x) + x * sigmoid(x) * (1 - sigmoid(x))
+    #          = sigmoid(x) + silu(x) * (1 - sigmoid(x))
+    d_silu = sig_x + silu_x * (1.0 - sig_x)
+
+    # dx = dout * y * silu'(x)
+    dx = dout * y * d_silu
+    tl.store(grad_x_ptr + offsets, dx, mask=mask)
+
+
+def silu_mul_backward(
+        grad_output: torch.Tensor, x: torch.Tensor, y: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Computes the backward pass of silu(x)*y using Triton.
+
+    Args:
+        grad_output: Gradient of the loss with respect to the output.
+        x: Original input tensor x.
+        y: Original input tensor y.
+
+    Returns:
+        A tuple of (grad_x, grad_y).
+    """
+
+    if not grad_output.is_contiguous():
+        grad_output = grad_output.contiguous()
+    if not x.is_contiguous():
+        x = x.contiguous()
+    if not y.is_contiguous():
+        y = y.contiguous()
+
+    n_elements = x.numel()
+
+    grad_x = torch.empty_like(x)
+    grad_y = torch.empty_like(y)
+
+    def _grid(meta: dict[str, int]) -> tuple[int, ...]:
+        return (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),)
+
+    _silu_mul_backward_kernel[_grid](
+        grad_output, x, y,
+        grad_x, grad_y,
+        n_elements
+    )
+
+    return grad_x, grad_y

d9d/loop/auto/__init__.py

@@ -0,0 +1,9 @@
+from .auto_lr_scheduler import AutoLRSchedulerConfig, AutoLRSchedulerProvider
+from .auto_optimizer import AutoOptimizerConfig, AutoOptimizerProvider
+
+__all__ = [
+    "AutoLRSchedulerConfig",
+    "AutoLRSchedulerProvider",
+    "AutoOptimizerConfig",
+    "AutoOptimizerProvider"
+]

d9d/loop/auto/auto_lr_scheduler.py

@@ -0,0 +1,46 @@
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from d9d.core.protocol import LRSchedulerProtocol
+from d9d.loop.control import InitializeLRSchedulerContext, LRSchedulerProvider
+from d9d.lr_scheduler.piecewise import PiecewiseSchedulerConfig, piecewise_scheduler_from_config
+
+
+class PiecewiseConfig(BaseModel):
+    """
+    Configuration for the piecewise learning rate scheduler.
+
+    Attributes:
+        name: Discriminator tag, must be "piecewise".
+        scheduler: Detailed configuration for the piecewise schedule.
+    """
+    name: Literal["piecewise"] = "piecewise"
+
+    scheduler: PiecewiseSchedulerConfig
+
+
+AutoLRSchedulerConfig = Annotated[
+    PiecewiseConfig,
+    Field(discriminator="name")
+]
+
+
+class AutoLRSchedulerProvider(LRSchedulerProvider):
+    """
+    LRSchedulerProvider that builds a learning rate scheduler based on a configuration object.
+    """
+
+    def __init__(self, config: AutoLRSchedulerConfig):
+        """Constructs the AutoLRSchedulerProvider object."""
+
+        self._config = config
+
+    def __call__(self, context: InitializeLRSchedulerContext) -> LRSchedulerProtocol:
+        match self._config:
+            case PiecewiseConfig():
+                return piecewise_scheduler_from_config(
+                    self._config.scheduler,
+                    optimizer=context.optimizer,
+                    total_steps=context.total_steps
+                )