RamTorch 0.2.2__tar.gz → 0.2.3__tar.gz

{ramtorch-0.2.2 → ramtorch-0.2.3}/PKG-INFO

@@ -1,16 +1,23 @@
 Metadata-Version: 2.4
 Name: RamTorch
-Version: 0.2.2
+Version: 0.2.3
 Summary: RAM is All You Need
+Home-page: https://github.com/lodestone-rock/RamTorch
+Author: Lodestone
 Author-email: Lodestone <lodestone.rock@gmail.com>
-License: Apache-2.0
 Project-URL: Homepage, https://github.com/lodestone-rock/RamTorch
-Classifier: Programming Language :: Python :: 3
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-License-File: LICENSE
-Dynamic: license-file
+Dynamic: author
+Dynamic: home-page
+Dynamic: requires-python
 
 # RamTorch
 

{ramtorch-0.2.2 → ramtorch-0.2.3}/RamTorch.egg-info/PKG-INFO

@@ -1,16 +1,23 @@
 Metadata-Version: 2.4
 Name: RamTorch
-Version: 0.2.2
+Version: 0.2.3
 Summary: RAM is All You Need
+Home-page: https://github.com/lodestone-rock/RamTorch
+Author: Lodestone
 Author-email: Lodestone <lodestone.rock@gmail.com>
-License: Apache-2.0
 Project-URL: Homepage, https://github.com/lodestone-rock/RamTorch
-Classifier: Programming Language :: Python :: 3
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Education
+Classifier: Intended Audience :: Science/Research
 Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
-License-File: LICENSE
-Dynamic: license-file
+Dynamic: author
+Dynamic: home-page
+Dynamic: requires-python
 
 # RamTorch
 

ramtorch-0.2.3/RamTorch.egg-info/SOURCES.txt

@@ -0,0 +1,29 @@
+README.md
+pyproject.toml
+setup.py
+./env/bin/activate_this.py
+./kernels/int8_matmul.py
+./ramtorch/__init__.py
+./ramtorch/helpers.py
+./ramtorch/multi_gpu.py
+./ramtorch/zero1.py
+./ramtorch/zero2.py
+./ramtorch/modules/__init__.py
+./ramtorch/modules/linear.py
+./ramtorch/stochastic_optimizers/__init__.py
+./ramtorch/stochastic_optimizers/adamw.py
+RamTorch.egg-info/PKG-INFO
+RamTorch.egg-info/SOURCES.txt
+RamTorch.egg-info/dependency_links.txt
+RamTorch.egg-info/top_level.txt
+env/bin/activate_this.py
+kernels/int8_matmul.py
+ramtorch/__init__.py
+ramtorch/helpers.py
+ramtorch/multi_gpu.py
+ramtorch/zero1.py
+ramtorch/zero2.py
+ramtorch/modules/__init__.py
+ramtorch/modules/linear.py
+ramtorch/stochastic_optimizers/__init__.py
+ramtorch/stochastic_optimizers/adamw.py

ramtorch-0.2.3/RamTorch.egg-info/top_level.txt

@@ -0,0 +1,6 @@
+data
+debug_snapshots
+dist
+env
+kernels
+ramtorch

ramtorch-0.2.3/env/bin/activate_this.py

@@ -0,0 +1,36 @@
+"""
+Activate virtualenv for current interpreter:
+
+Use exec(open(this_file).read(), {'__file__': this_file}).
+
+This can be used when you must use an existing Python interpreter, not the virtualenv bin/python.
+"""  # noqa: D415
+from __future__ import annotations
+
+import os
+import site
+import sys
+
+try:
+    abs_file = os.path.abspath(__file__)
+except NameError as exc:
+    msg = "You must use exec(open(this_file).read(), {'__file__': this_file})"
+    raise AssertionError(msg) from exc
+
+bin_dir = os.path.dirname(abs_file)
+base = bin_dir[: -len('bin') - 1]  # strip away the bin part from the __file__, plus the path separator
+
+# prepend bin to PATH (this file is inside the bin directory)
+os.environ["PATH"] = os.pathsep.join([bin_dir, *os.environ.get("PATH", "").split(os.pathsep)])
+os.environ["VIRTUAL_ENV"] = base  # virtual env is right above bin directory
+os.environ["VIRTUAL_ENV_PROMPT"] = '' or os.path.basename(base)
+
+# add the virtual environments libraries to the host python import mechanism
+prev_length = len(sys.path)
+for lib in '../lib/python3.12/site-packages'.split(os.pathsep):
+    path = os.path.realpath(os.path.join(bin_dir, lib))
+    site.addsitedir(path.decode("utf-8") if '' else path)
+sys.path[:] = sys.path[prev_length:] + sys.path[0:prev_length]
+
+sys.real_prefix = sys.prefix
+sys.prefix = base

ramtorch-0.2.3/kernels/int8_matmul.py

@@ -0,0 +1,268 @@
+import torch
+import triton
+import triton.language as tl
+from triton import Config
+from typing import Tuple
+
+
+"""
+simplified explanation of the scaled int8 matmul algorithm
+adopted from deepseek scaled FP8 matmul and jetfire paper
+https://arxiv.org/abs/2403.12422
+https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
+
+                                                     N dimension →  
+                                               INT8 weights                 scaler per block
+                                               ┌-----┬-----┬─────┬─────┐    ┌-----┬-----┬─────┬─────┐
+                                               : b00 : b01 : b02 | b03 |    :     :     :     |     |
+                                               ├-----┼-----┼─────┼─────┤    :b_s00:b_s10:b_s20|b_s30|
+                                           K   : b10 : b11 : b12 | b13 |    :     :     :     |     |
+                                          dim  ├-----┼-----┼─────┼─────┤    ├-----┼-----┼─────┼─────┤
+                                           ↓   | b20 | b21 | b22 | b23 |    |     |     |     |     |
+                                               ├─────┼─────┼─────┼─────┤    |b_s01|b_s11|b_s21|b_s31|
+                                               | b30 | b31 | b32 | b33 |    |     |     |     |     |
+                                               └─────┴─────┴─────┴─────┘    └─────┴─────┴─────┴─────┘
+                                               ┌-----┬-----┐
+                                               : b00 : b01 :
+     ├─── blk ───┤                             ├-----┼-----┤
+                                               : b10 : b11 :
+            K dimension →                      └-----┴-----┘                                
+     INT8 activations
+     ┌-----┬-----┬─────┬─────┐   ┌-----┬-----┐ ┌-----┬-----┐   ┌-----------┐   ┌-----┬-----┐   ┌-----┬-----┐
+     : a00 : a01 : a02 | a03 |   : a00 : a01 : :  @  :  @  :   :   a_s00   :   :     :     :   :acc00:acc01:
+     ├-----┼-----┼─────┼─────┤   ├-----┼-----┤ ├-----┼-----┤ * ├-----------┤ * :b_s00:b_s10: = ├-----┼-----┤ 
+ M   : a10 : a11 : a12 | a13 |   : a10 : a11 : :  @  :  @  :   :   a_s10   :   :     :     :   :acc10:acc11:
+dim  ├-----┼-----┼─────┼─────┤   └-----┴-----┘ └-----┴-----┘   └-----------┘   └-----┴-----┘   └-----┴-----┘
+ ↓   | a20 | a21 | a22 | a23 |   INT8 matmul acc in INT32      rescale the FP32 intermediate   accumulate
+     ├─────┼─────┼─────┼─────┤   then cast to FP32             "rank 1" hadamard scaler        intermediate
+     | a30 | a31 | a32 | a33 |  
+     └─────┴─────┴─────┴─────┘  
+     scaler per block
+     ┌-----------┬───────────┐
+     :   a_s00   :   a_s01   |
+     ├-----------┼───────────┤
+     :   a_s10   :   a_s11   |
+     ├-----------┼───────────┤
+     |   a_s20   |   a_s21   |
+     ├───────────┼───────────┤
+     |   a_s30   |   a_s31   |
+     └───────────┴───────────┘
+"""
+
+
+@triton.jit
+def act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
+    """
+    Quantizes the input tensor `x_ptr` and stores the result in `y_ptr` and the scaling factor in `s_ptr`.
+
+    Args:
+        x_ptr (triton.Pointer): Pointer to the input tensor.
+        y_ptr (triton.Pointer): Pointer to the output tensor where quantized values will be stored.
+        s_ptr (triton.Pointer): Pointer to the output tensor where scaling factors will be stored.
+        BLOCK_SIZE (tl.constexpr): The size of the block to be processed by each program instance.
+
+    Returns:
+        None
+    """
+    pid = tl.program_id(axis=0)
+    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    x = tl.load(x_ptr + offs).to(tl.float32)
+    amax = tl.max(tl.abs(x))  # reduction
+    # amax = tl.maximum(amax, 1e-4) # clamp to 1e-4
+    s = amax / 127.0
+    y = x / s
+    y = y.to(y_ptr.dtype.element_ty)
+    tl.store(y_ptr + offs, y)
+    tl.store(s_ptr + pid, s)
+
+
+def act_quant(
+    x: torch.Tensor, block_size: int = 128
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Quantizes the input tensor `x` using block-wise quantization.
+
+    Args:
+        x (torch.Tensor): The input tensor to be quantized. Must be contiguous and its last dimension size must be divisible by `block_size`.
+        block_size (int, optional): The size of the blocks to be used for quantization. Default is 128.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: A tuple containing:
+            - The quantized tensor with dtype `torch.int8`.
+            - A tensor of scaling factors with dtype `torch.float32`.
+    """
+    assert x.is_contiguous(), "Input tensor must be contiguous"
+    assert (
+        x.size(-1) % block_size == 0
+    ), f"Last dimension size must be divisible by block_size (block_size={block_size})"
+    y = torch.empty_like(x, dtype=torch.int8)
+    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
+    grid = lambda meta: (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)
+    act_quant_kernel[grid](x, y, s, BLOCK_SIZE=block_size)
+    return y, s
+
+
+@triton.jit
+def weight_dequant_kernel(x_ptr, s_ptr, y_ptr, M, N, BLOCK_SIZE: tl.constexpr):
+    """
+    Dequantizes weights using the provided scaling factors and stores the result.
+
+    Args:
+        x_ptr (tl.pointer): Pointer to the quantized weights.
+        s_ptr (tl.pointer): Pointer to the scaling factors.
+        y_ptr (tl.pointer): Pointer to the output buffer for dequantized weights.
+        M (int): Number of rows in the weight matrix.
+        N (int): Number of columns in the weight matrix.
+        BLOCK_SIZE (tl.constexpr): Size of the block for tiling.
+
+    Returns:
+        None
+    """
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    n = tl.cdiv(N, BLOCK_SIZE)
+    offs_m = pid_m * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs_n = pid_n * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    offs = offs_m[:, None] * N + offs_n[None, :]
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    x = tl.load(x_ptr + offs, mask=mask).to(tl.float32)
+    s = tl.load(s_ptr + pid_m * n + pid_n)
+    y = x * s
+    tl.store(y_ptr + offs, y, mask=mask)
+
+
+def weight_dequant(
+    x: torch.Tensor, s: torch.Tensor, block_size: int = 128
+) -> torch.Tensor:
+    """
+    Dequantizes the given weight tensor using the provided scale tensor.
+
+    Args:
+        x (torch.Tensor): The quantized weight tensor of shape (M, N).
+        s (torch.Tensor): The scale tensor of shape (M//block_size, N//block_size).
+        block_size (int, optional): The block size to use for dequantization. Defaults to 128.
+
+    Returns:
+        torch.Tensor: The dequantized weight tensor of the same shape as `x`.
+
+    Raises:
+        AssertionError: If `x` or `s` are not contiguous or if their dimensions are not 2.
+    """
+    assert x.is_contiguous() and s.is_contiguous(), "Input tensors must be contiguous"
+    assert x.dim() == 2 and s.dim() == 2, "Input tensors must have 2 dimensions"
+    M, N = x.size()
+    y = torch.empty_like(x, dtype=torch.get_default_dtype())
+    grid = lambda meta: (
+        triton.cdiv(M, meta["BLOCK_SIZE"]),
+        triton.cdiv(N, meta["BLOCK_SIZE"]),
+    )
+    weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE=block_size)
+    return y
+
+
+# matmul intermediate block size is hardcoded to 128
+int8_gemm_configs = [
+    Config(
+        {"BLOCK_SIZE_M": block_m, "BLOCK_SIZE_N": block_n, "BLOCK_SIZE_K": 128},
+        num_stages=num_stages,
+        num_warps=8,
+    )
+    for block_m in [16, 32, 64]
+    for block_n in [32, 64, 128]
+    for num_stages in [3, 4, 5, 6]
+]
+
+
+@triton.autotune(configs=int8_gemm_configs, key=["N", "K"])
+@triton.jit
+def int8_gemm_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    a_s_ptr,
+    b_s_ptr,
+    M,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    """
+    Performs a matrix multiplication operation on INT8 matrices with scaling factors.
+
+    Args:
+        a_ptr (tl.tensor): Pointer to the first input matrix A.
+        b_ptr (tl.tensor): Pointer to the second input matrix B.
+        c_ptr (tl.tensor): Pointer to the output matrix C.
+        a_s_ptr (tl.tensor): Pointer to the scaling factors for matrix A.
+        b_s_ptr (tl.tensor): Pointer to the scaling factors for matrix B.
+        M (int): Number of rows in matrix A and C.
+        N (tl.constexpr): Number of columns in matrix B and C.
+        K (tl.constexpr): Number of columns in matrix A and rows in matrix B.
+        BLOCK_SIZE_M (tl.constexpr): Block size for the M dimension.
+        BLOCK_SIZE_N (tl.constexpr): Block size for the N dimension.
+        BLOCK_SIZE_K (tl.constexpr): Block size for the K dimension.
+
+    Returns:
+        None
+    """
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    k = tl.cdiv(K, BLOCK_SIZE_K)
+    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = a_ptr + offs_m[:, None] * K + offs_k[None, :]
+    b_ptrs = b_ptr + offs_n[None, :] * K + offs_k[:, None]
+    a_s_ptrs = a_s_ptr + offs_m * k
+    b_s_ptrs = b_s_ptr + offs_n * k
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for i in range(k):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - i * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - i * BLOCK_SIZE_K, other=0.0)
+        a_s = tl.load(a_s_ptrs)
+        b_s = tl.load(b_s_ptrs)
+        # Cast to float32 before multiplying with scaling factors.
+        dot_prod = tl.dot(a, b)
+        accumulator += dot_prod.to(tl.float32) * a_s[:, None] * b_s[None, :]
+        a_ptrs += BLOCK_SIZE_K
+        b_ptrs += BLOCK_SIZE_K
+        a_s_ptrs += 1
+        b_s_ptrs += 1
+    c = accumulator.to(c_ptr.dtype.element_ty)
+    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = c_ptr + offs_m[:, None] * N + offs_n[None, :]
+    mask = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    tl.store(c_ptrs, c, mask=mask)
+
+
+def int8_gemm(a: torch.Tensor, a_s: torch.Tensor, b: torch.Tensor, b_s: torch.Tensor):
+    """
+    Perform a matrix multiplication using INT8 precision.
+
+    Args:
+        a (torch.Tensor): The first input matrix, must be contiguous.
+        a_s (torch.Tensor): The scaling factor for the first input matrix, must be contiguous.
+        b (torch.Tensor): The second input matrix, must be contiguous.
+        b_s (torch.Tensor): The scaling factor for the second input matrix, must be contiguous.
+
+    Returns:
+        torch.Tensor: The result of the matrix multiplication.
+    """
+    assert a.is_contiguous() and b.is_contiguous(), "Input tensors must be contiguous"
+    assert (
+        a_s.is_contiguous() and b_s.is_contiguous()
+    ), "Scaling factor tensors must be contiguous"
+    K = a.size(-1)
+    M = a.numel() // K
+    N = b.size(0)
+    c = a.new_empty(*a.size()[:-1], N, dtype=torch.get_default_dtype())
+    grid = lambda META: (
+        triton.cdiv(M, META["BLOCK_SIZE_M"]),
+        triton.cdiv(N, META["BLOCK_SIZE_N"]),
+    )
+    int8_gemm_kernel[grid](a, b, c, a_s, b_s, M, N, K)
+    return c

ramtorch-0.2.3/pyproject.toml

@@ -0,0 +1,25 @@
+[project]
+name = "RamTorch"
+version = "0.2.3"
+authors = [{name = "Lodestone", email = "lodestone.rock@gmail.com"}]
+description = "RAM is All You Need"
+readme = "README.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Programming Language :: Python :: 3",
+]
+
+[project.urls]
+Homepage = "https://github.com/lodestone-rock/RamTorch"
+
+[tool.setuptools.packages.find]
+exclude = ["examples*"]
+
+[tool.setuptools]
+license-files = []