main.py

"""
Mamba-2 Implementation in PyTorch

This module implements the Mamba-2 architecture, a state-space model based
transformer alternative. It includes robust device handling, type checking,
and comprehensive logging.

Features:
- Multi-device support (CPU/GPU/Multi-GPU)
- Robust error handling and validation
- Comprehensive shape tracking and logging
- Type hints and documentation
"""

import math
import warnings
from dataclasses import dataclass
from enum import Enum, auto
from typing import Optional

import torch
import torch.nn as nn
from loguru import logger
from torch import Tensor


class DeviceType(Enum):
    """Supported device types for model execution."""

    CPU = auto()
    GPU = auto()
    MULTI_GPU = auto()


@dataclass
class Mamba2Config:
    """Configuration for Mamba-2 model.

    Args:
        d_model: Model dimension
        depth: Number of Mamba blocks
        d_state: State dimension for SSM
        d_conv: Convolution kernel size
        expand_factor: Expansion factor for inner dimension
        device_type: Type of device to run on
        dtype: Data type for model parameters
        distributed: Whether to use distributed training
    """

    d_model: int
    depth: int
    d_state: int = 16
    d_conv: int = 4
    expand_factor: int = 2
    device_type: DeviceType = DeviceType.CPU
    dtype: torch.dtype = torch.float32
    distributed: bool = False

    def __post_init__(self):
        """Validate configuration parameters."""
        if self.d_model <= 0:
            raise ValueError(
                f"d_model must be positive, got {self.d_model}"
            )
        if self.depth <= 0:
            raise ValueError(
                f"depth must be positive, got {self.depth}"
            )
        if self.d_state <= 0:
            raise ValueError(
                f"d_state must be positive, got {self.d_state}"
            )
        if self.d_conv <= 0:
            raise ValueError(
                f"d_conv must be positive, got {self.d_conv}"
            )
        if self.expand_factor <= 0:
            raise ValueError(
                f"expand_factor must be positive, got {self.expand_factor}"
            )


class DeviceManager:
    """Manages device placement and data movement for the model."""

    def __init__(self, config: Mamba2Config):
        self.config = config
        self.device = self._setup_device()

    def _setup_device(self) -> torch.device:
        """Set up the appropriate device based on configuration."""
        if self.config.device_type == DeviceType.CPU:
            return torch.device("cpu")

        if not torch.cuda.is_available():
            warnings.warn(
                "GPU requested but CUDA is not available. Falling back to CPU."
            )
            return torch.device("cpu")

        if self.config.device_type == DeviceType.MULTI_GPU:
            if torch.cuda.device_count() < 2:
                warnings.warn(
                    "Multi-GPU requested but less than 2 GPUs available. Using single GPU."
                )
                return torch.device("cuda:0")
            return torch.device("cuda")

        return torch.device("cuda:0")

    def to_device(self, tensor: Tensor) -> Tensor:
        """Move tensor to appropriate device with error handling."""
        try:
            return tensor.to(self.device, dtype=self.config.dtype)
        except RuntimeError as e:
            logger.error(
                f"Failed to move tensor to device {self.device}: {e}"
            )
            raise


class SSM(nn.Module):
    """Structured State Space Model component of Mamba-2.

    Implements the core state space transformation with selective scan.

    Args:
        d_model: Model dimension
        d_state: State dimension
        dt_rank: Rank of Δ projection
        device_manager: Device management instance
    """

    def __init__(
        self,
        d_model: int,
        d_state: int,
        dt_rank: int,
        device_manager: DeviceManager,
    ):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.dt_rank = dt_rank
        self.device_manager = device_manager

        # Initialize parameters
        self.A = nn.Parameter(
            torch.randn(d_state, d_state) / math.sqrt(d_state)
        )
        self.D = nn.Parameter(
            torch.randn(d_model) / math.sqrt(d_model)
        )
        self.dt_projs = nn.Parameter(
            torch.randn(dt_rank, d_model) / math.sqrt(dt_rank)
        )

    def forward(self, x: Tensor, B: Tensor, C: Tensor) -> Tensor:
        """
        Forward pass of SSM.

        Args:
            x: Input tensor (batch, seq_len, d_model)
            B: B matrix (batch, seq_len, d_inner, d_state)
            C: C matrix (batch, seq_len, d_inner, d_state)

        Returns:
            Tensor: Output tensor (batch, seq_len, d_model)

        Raises:
            RuntimeError: If tensor dimensions don't match expected shapes
        """
        self._validate_input_shapes(x, B, C)

        batch, seq_len, d_model = x.shape

        # Compute Δ
        torch.einsum("rd,bsd->bsr", self.dt_projs, x)

        # Discretize A
        A_expanded = self.A.unsqueeze(0).unsqueeze(0)
        A_expanded = A_expanded.expand(batch, seq_len, -1, -1)
        dA = torch.exp(A_expanded)

        # Initialize state
        h = self.device_manager.to_device(
            torch.zeros(batch, self.d_state)
        )
        y = []

        # Selective scan with error checking
        try:
            for t in range(seq_len):
                h = torch.bmm(dA[:, t], h.unsqueeze(-1)).squeeze(-1)
                h = h + torch.einsum("bmd,bm->bd", B[:, t], x[:, t])
                y_t = torch.einsum("bd,bmd->bm", h, C[:, t])
                y.append(y_t)
        except RuntimeError as e:
            logger.error(
                f"Error in selective scan at position {t}: {e}"
            )
            raise

        y = torch.stack(y, dim=1)
        return y + self.D.unsqueeze(0).unsqueeze(0) * x

    def _validate_input_shapes(self, x: Tensor, B: Tensor, C: Tensor):
        """Validate input tensor shapes."""
        if x.dim() != 3:
            raise ValueError(
                f"Expected x to have 3 dimensions, got {x.dim()}"
            )
        if B.dim() != 4:
            raise ValueError(
                f"Expected B to have 4 dimensions, got {B.dim()}"
            )
        if C.dim() != 4:
            raise ValueError(
                f"Expected C to have 4 dimensions, got {C.dim()}"
            )
        if x.size(-1) != self.d_model:
            raise ValueError(
                f"Expected x last dim to be {self.d_model}, got {x.size(-1)}"
            )


class Mamba2Block(nn.Module):
    """Single block of Mamba-2 architecture.

    Args:
        config: Model configuration
        device_manager: Device management instance
    """

    def __init__(
        self, config: Mamba2Config, device_manager: DeviceManager
    ):
        super().__init__()
        self.config = config
        self.device_manager = device_manager
        self.d_inner = config.d_model * config.expand_factor

        # Projections
        self.in_proj_x = nn.Linear(config.d_model, self.d_inner)
        self.in_proj_b = nn.Linear(
            config.d_model, self.d_inner * config.d_state
        )
        self.in_proj_c = nn.Linear(
            config.d_model, self.d_inner * config.d_state
        )

        # Conv1d
        self.conv1d = nn.Conv1d(
            in_channels=self.d_inner,
            out_channels=self.d_inner,
            kernel_size=config.d_conv,
            padding="same",
            groups=self.d_inner,
        )

        # SSM
        self.ssm = SSM(
            d_model=self.d_inner,
            d_state=config.d_state,
            dt_rank=8,
            device_manager=device_manager,
        )

        self.norm = nn.GroupNorm(
            num_groups=1, num_channels=self.d_inner
        )
        self.out_proj = nn.Linear(self.d_inner, config.d_model)

    def forward(self, x: Tensor) -> Tensor:
        """Forward pass of Mamba-2 block."""
        batch, seq_len, _ = x.shape

        # Projections
        x_projected = self.in_proj_x(x)
        b_projected = self.in_proj_b(x)
        c_projected = self.in_proj_c(x)

        # Reshape with dimension checking
        try:
            B = b_projected.view(
                batch, seq_len, self.d_inner, self.config.d_state
            )
            C = c_projected.view(
                batch, seq_len, self.d_inner, self.config.d_state
            )
        except RuntimeError as e:
            logger.error(f"Failed to reshape projections: {e}")
            raise

        # Process
        x_conv = self.conv1d(x_projected.transpose(-1, -2)).transpose(
            -1, -2
        )
        x_ssm = self.ssm(x_conv, B, C)
        x_norm = self.norm(x_ssm.transpose(-1, -2)).transpose(-1, -2)

        return self.out_proj(x_norm)


class Mamba2(nn.Module):
    """Complete Mamba-2 architecture."""

    def __init__(self, config: Mamba2Config):
        super().__init__()
        self.config = config
        self.device_manager = DeviceManager(config)

        # Create blocks
        self.blocks = nn.ModuleList(
            [
                Mamba2Block(config, self.device_manager)
                for _ in range(config.depth)
            ]
        )

        # Move model to appropriate device
        self.to(self.device_manager.device)

        if config.distributed:
            self.blocks = nn.DataParallel(self.blocks)

    def forward(self, x: Tensor) -> Tensor:
        """Forward pass of complete Mamba-2 model."""
        x = self.device_manager.to_device(x)

        for block in self.blocks:
            try:
                x = x + block(x)
            except RuntimeError as e:
                logger.error(f"Error in block forward pass: {e}")
                raise

        return x


def create_mamba2_model(
    config: Mamba2Config, seed: Optional[int] = None
) -> Mamba2:
    """
    Create a Mamba-2 model with specified configuration.

    Args:
        config: Model configuration
        seed: Random seed for reproducibility

    Returns:
        Configured Mamba-2 model

    Raises:
        RuntimeError: If model creation fails
    """
    if seed is not None:
        torch.manual_seed(seed)

    try:
        model = Mamba2(config)
        logger.info(
            f"Created Mamba-2 model: d_model={config.d_model}, "
            f"depth={config.depth}, device={model.device_manager.device}"
        )
        return model
    except Exception as e:
        logger.error(f"Failed to create model: {e}")
        raise


def example_usage():
    """Example usage of Mamba-2 model."""
    # Configure logging
    logger.add("mamba2.log", rotation="500 MB")

    # Create configuration
    config = Mamba2Config(
        d_model=256,
        depth=4,
        device_type=(
            DeviceType.GPU
            if torch.cuda.is_available()
            else DeviceType.CPU
        ),
        distributed=torch.cuda.device_count() > 1,
    )

    # Create model
    model = create_mamba2_model(config, seed=42)

    # Example forward pass
    batch_size, seq_len = 32, 128
    x = torch.randn(batch_size, seq_len, config.d_model)

    logger.info("Starting forward pass")
    with torch.no_grad():
        output = model(x)
    logger.info(
        f"Forward pass complete. Output shape: {output.shape}"
    )


if __name__ == "__main__":
    example_usage()