sarasa 0.0.2__py3-none-any.whl

sarasa/models/nanochat_gpt.py

@@ -0,0 +1,192 @@
+# NanoChat's GPT model, adapted from https://github.com/karpathy/nanochat
+
+
+import torch
+from loguru import logger
+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,
+    ):
+        super().__init__()
+        self.c_fc = nn.Linear(config.hidden_dim, 4 * config.hidden_dim, bias=False)
+        self.c_proj = nn.Linear(4 * config.hidden_dim, config.hidden_dim, bias=False)
+
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = F.relu(x).square()
+        x = self.c_proj(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        config: ModelConfig,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.attn = CausalSelfAttention(config, layer_idx)
+        self.mlp = MLP(config)
+        self.norm = RMSNorm(config.hidden_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos_sin: tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        x = x + self.attn(self.norm(x), cos_sin)
+        x = x + self.mlp(self.norm(x))
+        return x
+
+
+class GPT(BaseModel):
+    def __init__(
+        self,
+        config: ModelConfig,
+        pad_vocab_size_to=64,
+    ):
+        """
+        NOTE a major footgun: this __init__ function runs in meta device context (!!)
+        Therefore, any calculations inside here are shapes and dtypes only, no actual data.
+        => We actually initialize all data (parameters, buffers, etc.) in init_weights() instead.
+        """
+        super().__init__()
+        self.config = config
+        self.num_heads = config.num_heads
+        self.hidden_dim = config.hidden_dim
+        self.seq_len = config.seq_len
+        self.vocab_size = config.vocab_size
+        self.num_layers = config.num_layers
+        # For DDP, we want vocab_size divisible by world_size. Also, there are potential performance benefits, see:
+        # https://huggingface.co/docs/transformers/main_classes/model#transformers.PreTrainedModel.resize_token_embeddings
+        padded_vocab_size = ((self.vocab_size + pad_vocab_size_to - 1) // pad_vocab_size_to) * pad_vocab_size_to
+        if padded_vocab_size != self.vocab_size:
+            logger.warning(
+                f"Padding vocab_size from {self.vocab_size} to {padded_vocab_size} to be divisible by {pad_vocab_size_to}"
+            )
+        self.token_emb = nn.Embedding(padded_vocab_size, self.hidden_dim)
+        self.blocks = nn.ModuleList([Block(config, layer_idx) for layer_idx in range(self.num_layers)])
+        self.lm_head = nn.Linear(self.hidden_dim, padded_vocab_size, bias=False)
+        self.norm = RMSNorm(self.hidden_dim)
+        # Per-layer learnable scalars (inspired by modded-nanogpt)
+        # resid_lambdas: scales the residual stream at each layer (init 1.0 = neutral)
+        # x0_lambdas: blends initial embedding back in at each layer (init 0.0 = disabled)
+        # Separate parameters so they can have different optimizer treatment
+        self.resid_lambdas = nn.Parameter(torch.ones(self.num_layers))  # fake init, real init in init_weights()
+        self.x0_lambdas = nn.Parameter(torch.zeros(self.num_layers))  # fake init, real init in init_weights()
+        # To support meta device initialization, we init the rotary embeddings here, but it's just "fake" meta tensors only.
+        # As for rotary_seq_len, these rotary embeddings are pretty small/cheap in memory,
+        # so let's just over-compute them by 10X, but assert fail if we ever reach that amount.
+        # In the future we can dynamically grow the cache, for now it's fine.
+        self.rotary_seq_len = self.seq_len * 16  # 10X over-compute should be enough, TODO make nicer?
+        cos, sin = RoPE.precompute(self.rotary_seq_len, config.head_dim)
+        self.register_buffer("cos", cos, persistent=False)  # persistent=False means it's not saved to the checkpoint
+        self.register_buffer("sin", sin, persistent=False)
+
+    @torch.no_grad()
+    def init_weights(self):
+        """
+        Initialize the full model in this one function for maximum clarity.
+
+        wte (embedding):     normal, std=1.0
+        lm_head:             normal, std=0.001
+        for each block:
+            attn.c_q:        uniform, std=1/sqrt(n_embd)
+            attn.c_k:        uniform, std=1/sqrt(n_embd)
+            attn.c_v:        uniform, std=1/sqrt(n_embd)
+            attn.c_proj:     zeros
+            mlp.c_fc:        uniform, std=1/sqrt(n_embd)
+            mlp.c_proj:      zeros
+        """
+
+        # Embedding and unembedding
+        torch.nn.init.normal_(self.token_emb.weight, mean=0.0, std=1.0)
+        torch.nn.init.normal_(self.lm_head.weight, mean=0.0, std=0.001)
+
+        # Transformer blocks: uniform init with bound = sqrt(3) * std (same standard deviation as normal)
+        n_embd = self.hidden_dim
+        s = 3**0.5 * n_embd**-0.5  # sqrt(3) multiplier makes sure Uniform achieves the same std as Normal
+        for block in self.blocks:
+            torch.nn.init.uniform_(block.attn.c_q.weight, -s, s)  # weights use Uniform to avoid outliers
+            torch.nn.init.uniform_(block.attn.c_k.weight, -s, s)
+            torch.nn.init.uniform_(block.attn.c_v.weight, -s, s)
+            torch.nn.init.zeros_(block.attn.c_proj.weight)  # projections are zero
+            torch.nn.init.uniform_(block.mlp.c_fc.weight, -s, s)
+            torch.nn.init.zeros_(block.mlp.c_proj.weight)
+
+        # Per-layer scalars
+        self.resid_lambdas.fill_(1.0)  # 1.0 => typical residual connections at init
+        self.x0_lambdas.fill_(0.0)  # 0.0 => skip connection to input is disabled at init
+
+        # Rotary embeddings
+        head_dim = self.hidden_dim // self.num_heads
+        self.cos, self.sin = RoPE.precompute(self.rotary_seq_len, head_dim, device=self.cos.device)
+
+        # Cast token embeddings to bf16: optimizer can tolerate it and it saves memory
+        if self.token_emb.weight.device.type == "cuda":
+            self.token_emb.to(dtype=torch.bfloat16)
+
+    def param_groups(
+        self,
+    ) -> dict[str, list[torch.nn.Parameter]]:
+        # Separate out all parameters into 5 groups (matrix, embedding, lm_head, resid_lambdas, x0_lambdas)
+        matrix_params = list(self.blocks.parameters())
+        embedding_params = list(self.token_emb.parameters())
+        lm_head_params = list(self.lm_head.parameters())
+        resid_params = [self.resid_lambdas]
+        x0_params = [self.x0_lambdas]
+        assert len(list(self.parameters())) == (
+            len(matrix_params) + len(embedding_params) + len(lm_head_params) + len(resid_params) + len(x0_params)
+        )
+
+        return {
+            "matrix": matrix_params,
+            "embedding": embedding_params,
+            "lm_head": lm_head_params,
+            "resid_lambdas": resid_params,
+            "x0_lambdas": x0_params,
+        }
+
+    def forward(
+        self,
+        input: torch.Tensor,
+    ) -> torch.Tensor:
+        B, T = input.size()
+
+        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim/2))
+        assert T <= self.cos.size(1), (
+            f"Sequence length grew beyond the rotary embeddings cache: {T} > {self.cos.size(1)}"
+        )
+        assert input.device == self.cos.device, (
+            f"Rotary embeddings and idx are on different devices: {input.device} != {self.cos.device}"
+        )
+        assert self.cos.dtype == torch.bfloat16, "Rotary embeddings must be in bfloat16"
+        # if kv cache exists, we need to offset the rotary embeddings to the current position in the cache
+        cos_sin = self.cos[:, :T], self.sin[:, :T]  # truncate cache to current sequence length
+
+        # Forward the trunk of the Transformer
+        x = self.token_emb(input)
+        x = self.norm(x)
+        x0 = x  # save initial normalized embedding for x0 residual
+        for block, resid_lambda, x0_lambda in zip(self.blocks, self.resid_lambdas, self.x0_lambdas):
+            x = resid_lambda * x + x0_lambda * x0
+            x = block(x, cos_sin)
+        x = self.norm(x)
+
+        # Forward the lm_head (compute logits)
+        softcap = 15  # smoothly cap the logits to the range [-softcap, softcap]
+        logits = self.lm_head(x)  # (B, T, padded_vocab_size) <- very big tensor, large amount of memory
+        logits = logits[..., : self.vocab_size]  # slice to remove padding
+        logits = logits.float()  # switch to fp32 for logit softcap and loss computation
+        logits = softcap * torch.tanh(logits / softcap)  # squash the logits
+
+        return logits

sarasa/models/utils.py

@@ -0,0 +1,39 @@
+import torch
+
+
+class RMSNorm(torch.nn.RMSNorm):
+    # RMSNorm without affine parameters
+    def __init__(
+        self,
+        normalized_shape: int,
+    ):
+        super().__init__(normalized_shape, eps=None, elementwise_affine=False)
+
+
+class RoPE:
+    @staticmethod
+    def precompute(
+        seq_len: int,
+        head_dim: int,
+        device: torch.device = None,
+        base: float = 10000,
+    ):
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)[None, :, None, :]
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.bfloat16(), sin.bfloat16()
+        return cos, sin
+
+    @staticmethod
+    def apply(
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+    ) -> torch.Tensor:
+        assert x.ndim == 4
+        x1, x2 = x.chunk(2, dim=-1)
+        y1 = x1 * cos + x2 * sin
+        y2 = x1 * (-sin) + x2 * cos
+        return torch.cat([y1, y2], 3)

sarasa/optimizers/__init__.py

@@ -0,0 +1,77 @@
+import dataclasses
+from typing import Literal
+
+import torch
+
+from sarasa.models import BaseModel
+from sarasa.optimizers.utils import GroupedOptimizer
+
+
+@dataclasses.dataclass
+class AdamW:
+    """
+    Default optimizer
+    """
+
+    lr: float = 1e-4
+    weight_decay: float = 0.1
+    betas: tuple[float, float] = (0.9, 0.95)
+
+    def create(
+        self,
+        model: BaseModel,
+    ) -> torch.optim.Optimizer:
+        param_groups = model.param_groups()
+        params = sum(param_groups.values(), [])
+        optimizer = torch.optim.AdamW(
+            params,
+            lr=torch.tensor(self.lr, dtype=torch.float32),
+            weight_decay=self.weight_decay,
+            betas=self.betas,
+            fused=True,
+        )
+        return optimizer
+
+
+@dataclasses.dataclass
+class Muon:
+    """
+    Muon optimizer
+    """
+
+    lr: float = 1e-4
+    weight_decay: float = 0.1
+    momentum: float = 0.9
+
+    adam_lr: float | None = None
+    adam_betas: tuple[float, float] = (0.9, 0.95)
+    adam_weight_decay: float = 0
+
+    adjust_lr_fn: Literal["original", "match_rms_adamw"] = "match_rms_adamw"
+
+    def __post_init__(self):
+        self.adam_lr = self.adam_lr or self.lr
+
+    def create(
+        self,
+        model: BaseModel,
+    ) -> torch.optim.Optimizer:
+        param_groups = model.param_groups()
+
+        muon = torch.optim.Muon(
+            param_groups["matrix"],
+            lr=self.lr,
+            weight_decay=self.weight_decay,
+            momentum=self.momentum,
+            adjust_lr_fn=self.adjust_lr_fn,
+        )
+
+        adam = torch.optim.AdamW(
+            sum([param_groups[k] for k in param_groups if k != "matrix"], []),
+            lr=self.adam_lr,
+            betas=self.adam_betas,
+            weight_decay=self.adam_weight_decay,
+            fused=True,
+        )
+
+        return GroupedOptimizer(muon, adam)

sarasa/optimizers/utils.py

@@ -0,0 +1,27 @@
+import torch
+
+
+class GroupedOptimizer(torch.optim.Optimizer):
+    def __init__(
+        self,
+        *optimizers: torch.optim.Optimizer,
+    ):
+        super().__init__(sum([optim.param_groups for optim in optimizers], []), {})
+        self.optimizers = optimizers
+
+    def step(self) -> None:
+        for optim in self.optimizers:
+            optim.step()
+
+    def zero_grad(
+        self,
+        set_to_none: bool = True,
+    ) -> None:
+        for optim in self.optimizers:
+            optim.zero_grad(set_to_none=set_to_none)
+
+    def state_dict(self) -> dict:
+        return super().state_dict()
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        super().load_state_dict(state_dict)

sarasa/trainer.py

@@ -0,0 +1,244 @@
+import contextlib
+import os
+import time
+from collections.abc import Iterable
+
+import torch
+import torch.distributed as dist
+from loguru import logger
+from torch.distributed.elastic.multiprocessing.errors import record
+
+from sarasa.activation_checkpoint import apply_op_sac
+from sarasa.checkpoint import Checkpointer
+from sarasa.config import Config
+from sarasa.metrics import MetricsProcessor
+from sarasa.utils import GarbageCollector, apply_distributed, init_distributed, set_dtype, update_timeout, world_size
+
+IGNORE_INDEX = -100
+
+
+class Trainer:
+    @record
+    def __init__(
+        self,
+        config: Config,
+    ) -> None:
+        self.config = config
+        logger.info(f"Initializing Trainer with config: {self.config}")
+
+        # set seed
+        torch.manual_seed(config.seed)
+        os.environ["PYTHONHASHSEED"] = str(config.seed % 2**32)
+
+        # setup device
+        torch.accelerator.set_device_index(int(os.environ.get("LOCAL_RANK", 0)))
+        self.device = torch.accelerator.current_accelerator(check_available=True)
+
+        self.gc = GarbageCollector(config.train.gc_freq)
+
+        # setup distributed
+        init_distributed(config.distributed.backend, config.distributed.init_timeout_seconds)
+
+        # setup data and tokenizer -> use vocab size for model setup
+        data = config.data.create(batch_size=config.train.local_batch_size)
+        self.data_loader = data["train_loader"]  # setup data loader
+        self.val_loader = data.get("val_loader", None)  # setup eval data loader
+        self.tokenizer = data["tokenizer"]  # setup tokenizer
+
+        vocab_size = len(self.tokenizer)
+        self.config.model.vocab_size = vocab_size
+
+        # todo: support other loss functions
+        self.loss_fn = torch.nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX, reduction="sum")
+
+        # setup model, optimizer, lr scheduler
+        with torch.device("meta"), set_dtype(getattr(torch, config.train.dtype)):
+            self.model = self.config.model.create()
+            num_params, flops_per_token = self.model.num_params_flops
+            model_size = num_params / 1e9
+            model_size, unit = (num_params / 1e6, "M") if model_size < 1 else (model_size, "B")
+            logger.info(f"Model created with {model_size:.2f}{unit} parameters")
+
+        # following torchtitan, (S)AC -> compilation -> distributed wrapping
+        if config.train.use_sac:
+            logger.info("Applying Selective Activation Checkpointing (SAC)")
+            for i, block in enumerate(self.model.blocks):
+                self.model.blocks[i] = apply_op_sac(block)
+
+        if config.train.compile:
+            logger.info("Compiling the model")
+            for block in self.model.blocks:
+                block.compile(fullgraph=True)
+            self.model.compile(dynamic=False)
+            self.loss_fn.compile()
+
+        if world_size() > 1:
+            apply_distributed(
+                config.distributed,
+                self.model,
+                device=self.device,
+                compile=config.train.compile,
+            )
+
+        self.model.to_empty(device=self.device)
+        self.model.init_weights()
+
+        self.optimizer = self.config.optim.create(self.model)
+        self.lr_scheduler = self.config.lr_scheduler.create(self.optimizer, config.train.steps)
+
+        # setup metrics and checkpointer
+        # todo: configure num_flops_per_token
+        self.metrics_processor = MetricsProcessor(config, self.device, flops_per_token)
+        self.checkpointer = Checkpointer(config, self.model) if config.checkpoint.save_freq > 0 else None
+
+        dev_mem_stats = self.metrics_processor.device_mem_monitor.get_peak_stats()
+        logger.info(
+            f"{self.device.type.upper()} memory: {dev_mem_stats.max_reserved_gib:.2f} GiB for model initialization"
+        )
+
+        self.step = 0
+        self.grad_accum_steps = config.train.global_batch_size // (config.train.local_batch_size * world_size())
+        logger.info(f"Gradient accumulation step is set to: {self.grad_accum_steps}")
+
+        self.amp_context = contextlib.nullcontext()
+        if config.distributed.name != "fsdp":
+            self.amp_context = torch.autocast(device_type=self.device.type, dtype=getattr(torch, config.train.dtype))
+
+        # todo: setup profiler context
+        self.profile_context = contextlib.nullcontext()
+
+        if config.train.use_fa4:
+            logger.info("Using FA4 flash attention")
+            try:
+                torch.nn.attention.activate_flash_attention_impl("FA4")
+            except Exception as e:
+                logger.warning(
+                    f"Failed to activate FA4 flash attention: {e}. Install sarasa with `flash_attn` extra for better performance."
+                )
+
+    def __del__(self) -> None:
+        # cleanup distributed
+        if world_size() > 1:
+            try:
+                dist.destroy_process_group()
+            except Exception as e:
+                logger.warning(f"Failed to destroy process group: {e}")
+
+    @record
+    def train(self):
+        logger.info("Starting training...")
+
+        self.model.train()
+        with self.profile_context:
+            data_iter = self.batch_generator(self.data_loader)
+            for _ in range(self.config.train.steps):
+                self.step += 1
+                self.gc.collect(self.step)
+                try:
+                    self.train_step(data_iter)
+                except StopIteration:
+                    logger.warning("Data loader exhausted during training.")
+                    break
+
+                if self.checkpointer is not None:
+                    self.checkpointer.save(self.step)
+
+                if self.config.train.val_freq > 0 and self.step % self.config.train.val_freq == 0:
+                    self.evaluate()
+
+                if world_size() > 1 and self.step == 1:
+                    update_timeout(self.config.distributed.train_timeout_seconds, self.device)
+
+        logger.info("Training completed.")
+
+    def batch_generator(
+        self,
+        data_iter: Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]],
+    ) -> Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]]:
+        data_iter = iter(data_iter)
+        while True:
+            begin = time.perf_counter()
+            batch = next(data_iter)
+            input_dict, target = batch
+            self.metrics_processor.ntokens_since_last_log += target.numel()
+            self.metrics_processor.data_load_times.append(time.perf_counter() - begin)
+            yield input_dict, target
+
+    def train_step(
+        self,
+        batch_iter: Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]],
+    ) -> None:
+        self.optimizer.zero_grad()
+
+        micro_batches = []
+        valid_tokens = torch.tensor(0, dtype=torch.long)
+        for _ in range(self.grad_accum_steps):
+            input_dict, target = next(batch_iter)
+            valid_tokens += (target != IGNORE_INDEX).sum()
+            micro_batches.append((input_dict, target))
+
+        valid_tokens = valid_tokens.to(self.device)
+        if world_size() > 1:
+            dist.all_reduce(valid_tokens, op=dist.ReduceOp.SUM)
+
+        losses = []
+        for input_dict, target in micro_batches:
+            input_dict = {
+                k: v.to(self.device, non_blocking=(self.device.type == "cuda")) for k, v in input_dict.items()
+            }
+            target = target.to(self.device, non_blocking=(self.device.type == "cuda"))
+
+            with self.amp_context:
+                pred = self.model(**input_dict)
+                loss = self.loss_fn(pred.flatten(0, 1), target.flatten(0, 1)) / valid_tokens
+
+            del pred
+            loss.backward()
+            losses.append(loss.detach())
+
+        if self.config.train.grad_clip is not None:
+            torch.nn.utils.clip_grad_norm_(
+                self.model.parameters(), self.config.train.grad_clip, foreach=self.device.type == "cuda"
+            )
+
+        if self.checkpointer is not None:
+            self.checkpointer.wait_for_staging()
+
+        self.optimizer.step()
+        self.lr_scheduler.step()
+
+        loss = torch.stack(losses).sum()
+
+        if not self.metrics_processor.should_log(self.step):
+            return
+
+        if world_size() > 1:
+            avg_loss = loss.clone()
+            dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
+            max_loss = loss.clone()
+            dist.all_reduce(max_loss, op=dist.ReduceOp.MAX)
+        else:
+            avg_loss = max_loss = loss
+
+        with torch.no_grad():
+            grad_norm = torch.nn.utils.get_total_norm(self.model.parameters(), foreach=self.device.type == "cuda")
+
+        lr = self.lr_scheduler.get_last_lr()[0]
+        self.metrics_processor.log(
+            self.step,
+            global_avg_loss=avg_loss.item(),
+            global_max_loss=max_loss.item(),
+            extra_metrics={
+                "grad_norm": grad_norm.item() if grad_norm >= 0 else float("nan"),
+                "lr": lr,
+            },
+        )
+
+    def evaluate(self):
+        raise NotImplementedError
+
+    def evaluation_step(
+        self,
+        batch_iter: Iterable[tuple[dict[str, torch.Tensor], torch.Tensor]],
+    ) -> None:
+        raise NotImplementedError