sarasa 0.0.2__py3-none-any.whl

sarasa/metrics.py

@@ -0,0 +1,294 @@
+import dataclasses
+import subprocess
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+from loguru import logger
+from torch._utils import _get_device_module
+
+from sarasa.config import Config
+from sarasa.utils import rank
+
+
+# ported from torchtitan
+# hardcoded BF16 type peak flops for NVIDIA A100, H100, H200, B200 GPU and AMD MI250, MI300X, MI325X, MI355X and Intel PVC
+def get_peak_flops(device_name: str) -> float:
+    try:
+        # Run the lspci command and capture the output
+        result = subprocess.run(["lspci"], stdout=subprocess.PIPE, text=True)
+        # Filter the output for lines containing both "NVIDIA" and "H100"
+        filtered_lines = [line for line in result.stdout.splitlines() if "NVIDIA" in line and "H100" in line]
+        # Join all filtered lines into a single string
+        device_name = " ".join(filtered_lines) or device_name
+    except FileNotFoundError as e:
+        logger.warning(f"Error running lspci: {e}, fallback to use device_name")
+    if "A100" in device_name:
+        # data from https://www.nvidia.com/en-us/data-center/a100/
+        return 312e12
+    elif "H100" in device_name:
+        # data from https://www.nvidia.com/en-us/data-center/h100/
+        # NOTE: Specifications are one-half lower without sparsity.
+        if "NVL" in device_name:
+            return 835e12
+        elif "PCIe" in device_name:
+            return 756e12
+        else:  # for H100 SXM and other variants
+            return 989e12
+    elif "H200" in device_name:
+        # data from https://www.nvidia.com/en-us/data-center/h200/
+        return 989e12
+    elif "B200" in device_name:
+        # data from https://nvdam.widen.net/s/wwnsxrhm2w/blackwell-datasheet-3384703
+        return 2.25e15
+    elif "MI355X" in device_name:
+        # MI355X data from https://www.amd.com/en/products/accelerators/instinct/mi350/mi355x.html
+        return 2500e12
+    elif "MI300X" in device_name or "MI325X" in device_name:
+        # MI300X data from https://www.amd.com/en/products/accelerators/instinct/mi300/mi300x.html
+        # MI325X data from https://www.amd.com/en/products/accelerators/instinct/mi300/mi325x.html
+        return 1300e12
+    elif "MI250X" in device_name:
+        # data from https://www.amd.com/en/products/accelerators/instinct/mi200/mi250x.html (per GCD)
+        return 191.5e12
+    elif "Data Center GPU Max 1550" in device_name:
+        # Also known as Ponte Vecchio (PVC).
+        # data from https://www.intel.com/content/www/us/en/docs/oneapi/optimization-guide-gpu/2025-0/intel-xe-gpu-architecture.html
+        # Dot Product Accumulate Systolic (DPAS):
+        # - Freq: 1300MHz
+        # - #ops: 512
+        # Full EU mode (i.e. 512 max compute units): 340.8 TFLOPS (BF16)
+        # Standard EU mode (i.e. 448 max compute units): 298.2 TFLOPS (BF16)
+        max_comp_units = torch.xpu.get_device_properties("xpu").max_compute_units
+        return 512 * max_comp_units * 1300 * 10**6
+    elif "l40s" in device_name:
+        # data from: "https://resources.nvidia.com/en-us-l40s/l40s-datasheet-28413"
+        return 362e12
+
+    else:  # for other GPU types, assume A100
+        logger.warning(f"Peak flops undefined for: {device_name}, fallback to A100")
+        return 312e12
+
+
+@dataclasses.dataclass(slots=True)
+class DevMemStats:
+    max_active_gib: float
+    max_active_perc: float
+    max_reserved_gib: float
+    max_reserved_perc: float
+    num_alloc_retries: int
+    num_ooms: int
+
+
+class DeviceMemoryMonitor:
+    def __init__(
+        self,
+        device: torch.device,
+    ) -> None:
+        self.device = device
+        try:
+            _, self.total_mem = torch.accelerator.get_memory_info(self.device)
+            self.device_name = _get_device_module(self.device.type).get_device_name(self.device)
+        except RuntimeError:
+            if self.device.type == "mps":
+                self.total_mem = torch.mps.recommended_max_memory()
+                self.device_name = "Apple Silicon GPU"
+            else:
+                raise NotImplementedError(f"Device memory monitor not implemented for device type: {self.device.type}")
+
+        self.reset_peak_stats()
+        try:
+            torch.accelerator.empty_cache()
+        except RuntimeError:
+            logger.error(f"Failed to empty cache for device type: {self.device.type}")
+
+    @staticmethod
+    def to_gib(bytes: int) -> float:
+        return bytes / (1024**3)
+
+    def reset_peak_stats(self) -> None:
+        try:
+            torch.accelerator.reset_peak_memory_stats(self.device)
+        except RuntimeError:
+            logger.error(f"Failed to reset peak memory stats for device type: {self.device.type}")
+
+    def get_peak_stats(self) -> DevMemStats:
+        try:
+            info = torch.accelerator.memory_stats(self.device)
+        except RuntimeError:
+            logger.error(f"Failed to get peak memory stats for device type: {self.device.type}")
+            info = {}
+
+        max_active = info.get("active_bytes.all.peak", -1)
+        max_reserved = info.get("reserved_bytes.all.peak", -1)
+        num_retries = info.get("num_alloc_retries", -1)
+        num_ooms = info.get("num_ooms", -1)
+
+        if num_retries > 0:
+            logger.warning(f"{num_retries} {self.device.type.upper()} memory allocation retries.")
+        if num_ooms > 0:
+            logger.warning(f"{num_ooms} {self.device.type.upper()} OOM errors thrown.")
+
+        return DevMemStats(
+            max_active_gib=self.to_gib(max_active),
+            max_active_perc=max_active / self.total_mem * 100,
+            max_reserved_gib=self.to_gib(max_reserved),
+            max_reserved_perc=max_reserved / self.total_mem * 100,
+            num_alloc_retries=num_retries,
+            num_ooms=num_ooms,
+        )
+
+
+class BaseReporter:
+    def config(
+        self,
+        config: dict[str, Any],
+    ) -> None:
+        raise NotImplementedError()
+
+    def log(
+        self,
+        metrics: dict[str, Any],
+        step: int,
+    ) -> None:
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        raise NotImplementedError()
+
+
+class TensorboardReporter(BaseReporter):
+    def __init__(
+        self,
+        log_dir: Path,
+    ) -> None:
+        from torch.utils.tensorboard import SummaryWriter
+
+        self.writer = SummaryWriter(log_dir=log_dir, max_queue=1000)
+
+        logger.info(f"TensorBoard log is available at {log_dir}")
+
+    def config(
+        self,
+        config: dict[str, Any],
+    ) -> None:
+        for k, v in config.items():
+            self.writer.add_text(f"config/{k}", str(v))
+
+    def log(
+        self,
+        metrics: dict[str, float],
+        step: int,
+    ) -> None:
+        for k, v in metrics.items():
+            self.writer.add_scalar(k, v, step)
+
+    def close(self) -> None:
+        self.writer.close()
+
+
+class MetricsProcessor:
+    def __init__(
+        self,
+        config: Config,
+        device: torch.device,
+        flops_per_token: int,
+    ) -> None:
+        self.reporters = []
+        if config.metrics.all_node or rank() == 0:
+            if config.metrics.use_tensorboard:
+                log_dir = config.output_dir / "tensorboard" if config.output_dir else Path("./tensorboard")
+                self.reporters.append(TensorboardReporter(log_dir=log_dir))
+
+        for reporter in self.reporters:
+            reporter.config(config=dataclasses.asdict(config))
+
+        self.device_mem_monitor = DeviceMemoryMonitor(device)
+        self.log_freq = config.metrics.log_freq
+        self.time_last_log = time.perf_counter()
+        gpu_peak_flops = get_peak_flops(self.device_mem_monitor.device_name)
+        logger.info(f"Detected device: {self.device_mem_monitor.device_name}, Peak FLOPS: {gpu_peak_flops}")
+        self.gpu_peak_flops = gpu_peak_flops
+        self.ntokens_since_last_log = 0
+        self.flops_per_token = flops_per_token
+        self.data_load_times: list[float] = []
+        self.reset()
+
+    def should_log(
+        self,
+        step: int,
+    ) -> bool:
+        return step == 1 or step % self.log_freq == 0
+
+    def log(
+        self,
+        step: int,
+        global_avg_loss: float,
+        global_max_loss: float,
+        extra_metrics: dict[str, float] | None = None,
+    ) -> None:
+        time_delta = time.perf_counter() - self.time_last_log
+        device_mem_stats = self.device_mem_monitor.get_peak_stats()
+        time_ete = time_delta / self.log_freq
+        time_data_load = sum(self.data_load_times) / len(self.data_load_times) if self.data_load_times else 0.0
+        time_data_load_perc = 100 * time_data_load / time_ete if time_ete > 0 else 0.0
+
+        metrics = {
+            "loss/avg": global_avg_loss,
+            "loss/max": global_max_loss,
+            "memory/max_active(GiB)": device_mem_stats.max_active_gib,
+            "memory/max_active(%)": device_mem_stats.max_active_perc,
+            "memory/max_reserved(GiB)": device_mem_stats.max_reserved_gib,
+            "memory/max_reserved(%)": device_mem_stats.max_reserved_perc,
+            "memory/num_alloc_retries": device_mem_stats.num_alloc_retries,
+            "memory/num_ooms": device_mem_stats.num_ooms,
+            "time/end-to-end(s)": time_ete,
+            "time/data_load(s)": time_data_load,
+            "time/data_load(%)": time_data_load_perc,
+        }
+
+        log = (
+            f"[Step {step:>10}] loss: {global_avg_loss:.4f}, memory: {device_mem_stats.max_reserved_gib:.2f} GiB, "
+            f"time(s): {time_ete:.2f}sec (data load ratio: {time_data_load_perc:.1f}%)"
+        )
+
+        if extra_metrics is not None:
+            metrics.update(extra_metrics)
+
+        if self.flops_per_token > 0:
+            tps = self.ntokens_since_last_log / time_delta
+            mfu = 100 * self.flops_per_token * tps / self.gpu_peak_flops
+            tflops = self.flops_per_token * tps / 1e12
+
+            metrics.update({
+                "throughput(tps)": tps,
+                "tflops": tflops,
+                "mfu(%)": mfu,
+            })
+            log += f", tflops: {tflops:.2f}, mfu: {mfu:.2f}%"
+
+        for reporter in self.reporters:
+            reporter.log(metrics, step)
+
+        logger.info(log)
+
+        self.reset()
+
+    def reset(self) -> None:
+        self.ntokens_since_last_log = 0
+        self.data_load_times.clear()
+        self.time_last_log = time.perf_counter()
+        self.device_mem_monitor.reset_peak_stats()
+
+    def val_log(
+        self,
+        step: int,
+        val_loss: float,
+        extra_metrics: dict[str, float] | None = None,
+    ) -> None:
+        raise NotImplementedError()
+
+    def close(self) -> None:
+        for reporter in self.reporters:
+            reporter.close()

sarasa/models/__init__.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import abc
+import dataclasses
+from typing import Literal
+
+import torch
+from loguru import logger
+from torch import nn
+
+
+@dataclasses.dataclass
+class ModelConfig:
+    name: Literal["nanochat_gpt", "llama3"] = "nanochat_gpt"
+    num_layers: int = 12
+    head_dim: int = 128
+
+    num_heads: int | None = None  # inferred later if None
+    num_kv_heads: int | None = None  # inferred later if None
+    hidden_dim: int | None = None  # inferred later if None
+    vocab_size: int | None = None  # set later based on tokenizer
+    seq_len: int | None = None  # set later based on data config
+    qk_norm: bool = False  # whether to use RMSNorm on q/k
+
+    def __post_init__(self):
+        # infer hidden_dim, num_heads, num_kv_heads if not provided using the rules presented in nanochat
+        self.hidden_dim = self.hidden_dim or (self.num_layers * 64 + self.head_dim - 1) // self.head_dim * self.head_dim
+        self.num_heads = self.num_heads or self.hidden_dim // self.head_dim
+        self.num_kv_heads = self.num_kv_heads or self.num_heads
+
+        # sanity checks
+        assert self.hidden_dim % self.head_dim == 0
+        assert self.head_dim * self.num_heads == self.hidden_dim
+        assert self.num_kv_heads <= self.num_heads and self.num_heads % self.num_kv_heads == 0
+
+    def create(self) -> BaseModel:
+        if self.vocab_size is None or self.seq_len is None:
+            raise ValueError("vocab_size and seq_len must be set before creating the model")
+
+        match self.name:
+            case "nanochat_gpt":
+                from .nanochat_gpt import GPT
+
+                if not self.qk_norm:
+                    logger.warning("nanochat_gpt model without qk_norm is not recommended")
+
+                return GPT(self)
+
+            case "llama3":
+                from .llama3 import Llama3
+
+                if self.qk_norm:
+                    logger.warning("llama3 model with qk_norm is not standard")
+
+                return Llama3(self)
+
+            case _:
+                raise ValueError(f"Unknown model name: {self.name}")
+
+
+class BaseModel(nn.Module, abc.ABC):
+    # Common interface for all models in Sarasa
+
+    blocks: list[nn.Module]  # TF blocks
+    config: ModelConfig
+
+    @abc.abstractmethod
+    @torch.no_grad()
+    def init_weights(self) -> None:
+        # Actual initialization of model weights
+        pass
+
+    @abc.abstractmethod
+    def param_groups(self) -> dict[str, list[nn.Parameter]]:
+        # Return parameter groups for optimizer
+        pass
+
+    @property
+    def num_params_flops(
+        self,
+    ) -> tuple[int, int]:
+        # Return number of parameters and FLOPs per token (for dense model)
+        config = self.config
+
+        # for tied embeddings, num_params -= num_params_emb
+        num_params = sum(p.numel() for p in self.parameters())
+        num_params_emb = self.token_emb.weight.numel()
+
+        # If forward pass has 1 matmul, then backward pass has 2 matmuls
+        # Each self-attention has 2 matmuls
+        num_flops_per_token = 6 * (
+            (num_params - num_params_emb) + (config.num_layers * config.num_heads * config.head_dim * config.seq_len)
+        )
+
+        return num_params, num_flops_per_token

sarasa/models/attention.py

@@ -0,0 +1,84 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from sarasa.models import ModelConfig
+from sarasa.models.utils import RMSNorm, RoPE
+
+
+class SDPAttention(nn.Module):
+    def __init__(
+        self,
+        is_causal: bool,
+        enable_gqa: bool,
+    ):
+        super().__init__()
+        self.is_causal = is_causal
+        self.enable_gqa = enable_gqa
+
+        if nn.attention.current_flash_attention_impl() == "FA4":
+            self.sdpa_backends = nn.attention.SDPBackend.FLASH_ATTENTION
+        else:
+            self.sdpa_backends = [
+                nn.attention.SDPBackend.CUDNN_ATTENTION,
+                nn.attention.SDPBackend.FLASH_ATTENTION,
+                nn.attention.SDPBackend.MATH,
+            ]
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+    ) -> torch.Tensor:
+        with nn.attention.sdpa_kernel(self.sdpa_backends):
+            return F.scaled_dot_product_attention(
+                query,
+                key,
+                value,
+                is_causal=self.is_causal,
+                enable_gqa=self.enable_gqa,
+            )
+
+
+class CausalSelfAttention(nn.Module):
+    def __init__(
+        self,
+        config: ModelConfig,
+        layer_idx: int | None = None,
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = config.num_heads
+        self.num_kv_heads = config.num_kv_heads
+        self.hidden_dim = config.hidden_dim
+        self.head_dim = self.hidden_dim // self.num_heads
+        self.c_q = nn.Linear(self.hidden_dim, self.num_heads * self.head_dim, bias=False)
+        self.c_k = nn.Linear(self.hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.c_v = nn.Linear(self.hidden_dim, self.num_kv_heads * self.head_dim, bias=False)
+        self.c_proj = nn.Linear(self.hidden_dim, self.hidden_dim, bias=False)
+        self.qk_norm = RMSNorm(self.head_dim) if config.qk_norm else nn.Identity()
+
+        # todo: support varlen etc and kv caching
+        self.attn = SDPAttention(is_causal=True, enable_gqa=self.num_heads != self.num_kv_heads)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos_sin: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        B, T, C = x.size()
+
+        # Project the input to get queries, keys, and values
+        q = self.c_q(x).view(B, T, self.num_heads, self.head_dim)
+        k = self.c_k(x).view(B, T, self.num_kv_heads, self.head_dim)
+        v = self.c_v(x).view(B, T, self.num_kv_heads, self.head_dim)
+
+        # Apply Rotary Embeddings to queries and keys to get relative positional encoding
+        cos, sin = cos_sin
+        q, k = RoPE.apply(q, cos, sin), RoPE.apply(k, cos, sin)
+        q, k = self.qk_norm(q), self.qk_norm(k)
+        y = self.attn(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2))  # (B, n_head, T, head_dim)
+        y = y.transpose(1, 2).contiguous().view(B, T, -1)
+        y = self.c_proj(y)
+        return y

sarasa/models/llama3.py

@@ -0,0 +1,129 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from sarasa.models import BaseModel, ModelConfig
+from sarasa.models.attention import CausalSelfAttention
+from sarasa.models.utils import RMSNorm, RoPE
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        config: ModelConfig,
+        multiple_of: int,
+        ffn_dim_multiplier: float | None,
+    ):
+        super().__init__()
+        hidden_dim = int(8 * config.hidden_dim / 3)
+        # custom dim factor multiplier
+        if ffn_dim_multiplier is not None:
+            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(config.hidden_dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, config.hidden_dim, bias=False)
+        self.w3 = nn.Linear(config.hidden_dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        config: ModelConfig,
+        layer_idx: int,
+        multiple_of: int,
+        ffn_dim_multiplier: float | None,
+    ):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.attention = CausalSelfAttention(config)
+        self.mlp = MLP(config, multiple_of, ffn_dim_multiplier)
+        self.norm = RMSNorm(config.hidden_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos_sin: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        x = x + self.attention(self.norm(x), cos_sin)
+        x = x + self.mlp(self.norm(x))
+        return x
+
+
+class Llama3(BaseModel):
+    def __init__(
+        self,
+        config: ModelConfig,
+        multiple_of: int = 1024,
+        ffn_dim_multiplier: float | None = 1.4,
+    ):
+        super().__init__()
+        self.config = config
+        self.token_emb = nn.Embedding(config.vocab_size, config.hidden_dim)
+        self.max_seq_len = config.seq_len * 16
+        self.head_dim = config.head_dim
+        cos, sin = RoPE.precompute(self.max_seq_len, config.head_dim)
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+        self.blocks = nn.ModuleList([
+            Block(config, layer_idx, multiple_of, ffn_dim_multiplier) for layer_idx in range(config.num_layers)
+        ])
+        self.norm = RMSNorm(config.hidden_dim)
+        self.output = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
+
+    @torch.no_grad()
+    def init_weights(self) -> None:
+        self.cos, self.sin = RoPE.precompute(self.max_seq_len, self.head_dim, device=self.cos.device)
+        torch.nn.init.normal_(self.token_emb.weight)
+        for block in self.blocks:
+            block: Block
+            init_std = 0.02 / (2 * (block.layer_idx + 1)) ** 0.5
+
+            nn.init.trunc_normal_(block.attention.c_q.weight, std=0.02)
+            nn.init.trunc_normal_(block.attention.c_k.weight, std=0.02)
+            nn.init.trunc_normal_(block.attention.c_v.weight, std=0.02)
+            nn.init.trunc_normal_(block.attention.c_proj.weight, std=init_std)
+
+            nn.init.trunc_normal_(block.mlp.w1.weight, std=0.02)
+            nn.init.trunc_normal_(block.mlp.w2.weight, std=init_std)
+            nn.init.trunc_normal_(block.mlp.w3.weight, std=init_std)
+
+        final_out_std = self.output.weight.shape[-1] ** -0.5
+        cutoff_factor = 3
+        nn.init.trunc_normal_(
+            self.output.weight,
+            mean=0.0,
+            std=final_out_std,
+            a=-cutoff_factor * final_out_std,
+            b=cutoff_factor * final_out_std,
+        )
+
+    def param_groups(self) -> dict[str, list[nn.Parameter]]:
+        matrix_params = list(self.blocks.parameters())
+        embedding_params = list(self.token_emb.parameters())
+        lm_head_params = list(self.output.parameters())
+        assert len(list(self.parameters())) == (len(matrix_params) + len(embedding_params) + len(lm_head_params))
+
+        return {
+            "matrix": matrix_params,
+            "embedding": embedding_params,
+            "lm_head": lm_head_params,
+        }
+
+    def forward(
+        self,
+        input: torch.Tensor,
+    ) -> torch.Tensor:
+        B, T = input.size()
+        x = self.token_emb(input)  # (B, T, C)
+        cos_sin = self.cos[:, :T], self.sin[:, :T]
+
+        for block in self.blocks:
+            x = block(x, cos_sin)
+
+        x = self.norm(x)
+        logits = self.output(x)
+        return logits