[DOCS][FusedDenseGELUDense]

docs/zeta/nn/modules/fused_gelu_dense.md

@@ -0,0 +1,140 @@
+# `FusedDenseGELUDense`
+
+## Overview
+
+The `FusedDenseGELUDense` module is a versatile neural network layer designed for efficient computation of dense layers with GELU (Gaussian Error Linear Unit) activations. This documentation will provide an in-depth understanding of the module's architecture, purpose, parameters, and usage examples.
+
+## Table of Contents
+
+1. [Introduction](#introduction)
+2. [Architecture](#architecture)
+3. [Purpose](#purpose)
+4. [Class Definition](#class-definition)
+    - [Parameters](#parameters)
+    - [Internal Layers](#internal-layers)
+5. [Functionality and Usage](#functionality-and-usage)
+    - [Forward Pass](#forward-pass)
+6. [Examples](#examples)
+    - [Basic Usage](#basic-usage)
+    - [Custom Configuration](#custom-configuration)
+    - [Quantization with bitsandbytes](#quantization-with-bitsandbytes)
+7. [Additional Information](#additional-information)
+8. [References](#references)
+
+---
+
+## 1. Introduction <a name="introduction"></a>
+
+The `FusedDenseGELUDense` module combines dense layers with GELU activations in a single neural network layer. This fusion improves computational efficiency and is particularly useful in various deep learning applications.
+
+## 2. Architecture <a name="architecture"></a>
+
+The `FusedDenseGELUDense` layer consists of two dense sub-layers, each followed by a GELU activation function. It takes an input tensor and passes it through these sub-layers to produce the final output.
+
+## 3. Purpose <a name="purpose"></a>
+
+The primary purpose of the `FusedDenseGELUDense` layer is to efficiently compute dense transformations with GELU activations. It is designed for use in neural networks, providing a convenient way to incorporate these operations into deep learning models.
+
+## 4. Class Definition <a name="class-definition"></a>
+
+### Parameters <a name="parameters"></a>
+
+- `dim` (int): Input dimension.
+- `dim_out` (int): Output dimension.
+- `bias` (bool, optional): Whether to include bias terms. Defaults to True.
+- `has_fp16_weights` (bool, optional): Whether to use fp16 weights. Defaults to False.
+- `threshold` (float, optional): Threshold for quantization. Defaults to 6.0.
+
+### Internal Layers <a name="internal-layers"></a>
+
+The `FusedDenseGELUDense` layer consists of the following internal layers:
+
+1. `dense1`: The first dense layer.
+2. `act`: The GELU activation function.
+3. `dense2`: The second dense layer.
+
+## 5. Functionality and Usage <a name="functionality-and-usage"></a>
+
+### Forward Pass <a name="forward-pass"></a>
+
+The `forward` method of the `FusedDenseGELUDense` layer performs the following operations:
+
+1. Applies the first dense layer (`dense1`) to the input tensor.
+2. Applies the GELU activation function (`act`) to the result.
+3. Applies the second dense layer (`dense2`) to the GELU-activated output.
+
+## 6. Examples <a name="examples"></a>
+
+### Basic Usage <a name="basic-usage"></a>
+
+Here's a basic example of using the `FusedDenseGELUDense` layer:
+
+```python
+import torch
+from zeta.nn import FusedDenseGELUDense
+
+# Create an instance of FusedDenseGELUDense
+model = FusedDenseGELUDense(dim=512, dim_out=1024)
+
+# Generate random input tensor
+x = torch.randn(1, 512)
+
+# Forward pass
+out = model(x)
+
+# Check the output shape
+print(out.shape)  # torch.Size([1, 512])
+```
+
+### Custom Configuration <a name="custom-configuration"></a>
+
+You can customize the layer by specifying different parameters:
+
+```python
+# Create a custom FusedDenseGELUDense layer
+custom_model = FusedDenseGELUDense(
+    dim=256, dim_out=512, bias=False, has_fp16_weights=True, threshold=4.0
+)
+
+# Generate random input tensor
+x = torch.randn(1, 256)
+
+# Forward pass with the custom configuration
+out = custom_model(x)
+```
+
+### Quantization with bitsandbytes <a name="quantization-with-bitsandbytes"></a>
+
+You can enable quantization using the `bitsandbytes` library by providing a quantized implementation of the dense layers:
+
+```python
+# Install bitsandbytes if not already installed
+# pip install bitsandbytes
+
+import torch
+from zeta.nn import FusedDenseGELUDense
+
+# Create an instance of FusedDenseGELUDense with quantization
+quantized_model = FusedDenseGELUDense(
+    dim=512, dim_out=1024, has_fp16_weights=True, threshold=4.0
+)
+
+# Generate random input tensor
+x = torch.randn(1, 512)
+
+# Forward pass with quantization
+out = quantized_model(x)
+```
+
+## 7. Additional Information <a name="additional-information"></a>
+
+- The `FusedDenseGELUDense` layer efficiently combines dense and GELU activation operations.
+- Custom configurations for bias, weight precision, and threshold are supported.
+- Quantization can be enabled using the `bitsandbytes` library for further efficiency.
+
+## 8. References <a name="references"></a>
+
+For more information on GELU activations and dense layers in PyTorch, refer to the official PyTorch documentation:
+
+- [GELU Activation Function](https://pytorch.org/docs/stable/generated/torch.nn.GELU.html)
+- [Dense Layer](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html)

mkdocs.yml

@@ -109,6 +109,7 @@ nav:
               - MultiModalAdapterDenseNetwork: "zeta/nn/modules/mm_adapter.md"
               - CustomMLP: "zeta/nn/modules/custom_mlp.md"
               - PolymorphicNeuronLayer: "zeta/nn/modules/polymorphic_activation.md"
+              - FusedDenseGELUDense: "zeta/nn/modules/fused_gelu_dense.md"
           - zeta.nn.attention:
               - FlashAttention: "zeta/nn/attention/flash_attention.md"
               - MultiQueryAttention: "zeta/nn/attention/multiquery.md"

tests/nn/modules/test_fused_gelu_dense.py

@@ -0,0 +1,70 @@
+import pytest
+import torch
+from zeta.nn.modules.fused_gelu_dense import FusedDenseGELUDense
+
+def test_class_init():
+    model = FusedDenseGELUDense(512, 1024)
+
+    assert model.dim == 512
+    assert model.dim_out == 1024
+    assert model.bias == True
+    assert model.has_fp16_weights == False
+    assert model.threshold == 6.0
+
+def test_class_init_with_args():
+    model = FusedDenseGELUDense(512, 1024, bias=False, has_fp16_weights=True, threshold=5.0)
+
+    assert model.dim == 512
+    assert model.dim_out == 1024
+    assert model.bias == False
+    assert model.has_fp16_weights == True
+    assert model.threshold == 5.0
+
+def test_forward():
+    model = FusedDenseGELUDense(512, 1024)
+    x = torch.randn(1, 512)
+    out = model(x)
+
+    assert out.shape == torch.Size([1, 512])
+
+def test_forward_with_different_input():
+    model = FusedDenseGELUDense(512, 1024)
+    x = torch.randn(2, 512)
+    out = model(x)
+
+    assert out.shape == torch.Size([2, 512])
+
+def test_forward_with_different_dim():
+    model = FusedDenseGELUDense(256, 512)
+    x = torch.randn(1, 256)
+    out = model(x)
+
+    assert out.shape == torch.Size([1, 256])
+
+def test_forward_with_different_dim_out():
+    model = FusedDenseGELUDense(512, 2048)
+    x = torch.randn(1, 512)
+    out = model(x)
+
+    assert out.shape == torch.Size([1, 512])
+
+def test_forward_with_no_bias():
+    model = FusedDenseGELUDense(512, 1024, bias=False)
+    x = torch.randn(1, 512)
+    out = model(x)
+
+    assert out.shape == torch.Size([1, 512])
+
+def test_forward_with_fp16_weights():
+    model = FusedDenseGELUDense(512, 1024, has_fp16_weights=True)
+    x = torch.randn(1, 512)
+    out = model(x)
+
+    assert out.shape == torch.Size([1, 512])
+
+def test_forward_with_different_threshold():
+    model = FusedDenseGELUDense(512, 1024, threshold=5.0)
+    x = torch.randn(1, 512)
+    out = model(x)
+
+    assert out.shape == torch.Size([1, 512])

zeta/cloud/main.py

@@ -13,7 +13,7 @@
 
 def zetacloud(
     task_name: str = None,
-    cluster_name: str = "[ZetaTrainingRun]",
+    cluster_name: str = "ZetaTrainingRun",
     cloud: Any = AWS(),
     gpus: str = None,
     filename: str = "train.py",

zeta/nn/modules/fused_gelu_dense.py

@@ -0,0 +1,98 @@
+import torch 
+from torch import nn
+
+class FusedDenseGELUDense(nn.Module):
+    """FuseFusedDenseGELUDense
+
+    Args
+        dim (int): Input dimension
+        dim_out (int): Output dimension
+        bias (bool, optional): Bias. Defaults to True.
+        has_fp16_weights (bool, optional): Use fp16 weights. Defaults to False.
+        threshold (float, optional): Threshold for quantization. Defaults to 6.0.
+        
+    Examples:
+        >>> x = torch.randn(1, 512)
+        >>> model = FusedDenseGELUDense(512, 1024)
+        >>> out = model(x)
+        >>> out.shape
+        torch.Size([1, 512])
+    """
+    def __init__(
+        self,
+        dim: int,
+        dim_out: int,
+        bias: bool = True,
+        has_fp16_weights: bool = False,
+        threshold: float = 6.0,
+        *args,
+        **kwargs
+    ):
+        super(FusedDenseGELUDense, self).__init__()
+        self.dim = dim 
+        self.dim_out = dim_out
+        self.bias = bias
+        self.has_fp16_weights = has_fp16_weights
+        self.threshold = threshold
+
+
+        try:
+            import bitsandbytes as bnb
+            # Using bitsandbytes for quantization
+            self.dense1 = bnb.nn.Linear8bitLt(
+                dim,
+                dim_out,
+                bias=bias,
+                has_fp16_weights=has_fp16_weights,
+                threshold=threshold,
+                *args,
+                **kwargs
+            )
+
+            # Reverse
+            self.dense2 = bnb.nn.Linear8bitLt(
+                dim_out,
+                dim,
+                bias=bias,
+                has_fp16_weights=has_fp16_weights,
+                threshold=threshold,
+                *args,
+                **kwargs
+            )
+
+        except ModuleNotFoundError:
+            # Using torch.nn.Linear
+            self.dense1 = nn.Linear(
+                dim,
+                dim_out,
+                bias=bias
+                *args,
+                **kwargs
+            )
+
+            # Dense 2
+            self.dense2 = nn.Linear(
+                dim_out,
+                dim,
+                bias=bias
+                *args,
+                **kwargs
+            )
+
+        # Activation
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass
+
+        Args:
+            x (torch.Tensor): x input
+
+        Returns:
+            torch.Tensor: _description_
+        """
+        x = self.dense1(x)
+        x = self.act(x)
+        x = self.dense2(x)
+        return x
+

zeta/optim/__init__.py

@@ -27,5 +27,5 @@
     "StableAdamWUnfused",
     "GradientAscent",
     "GradientEquilibrum",
-    "DecoupledLionW8Bit"
+    "DecoupledLionW8Bit",
 ]