sarasa 0.0.3__py3-none-any.whl → 0.0.4__py3-none-any.whl

sarasa/__init__.py

@@ -1,2 +1,11 @@
+from .config import DDP as DDP
+from .config import FSDP as FSDP
+from .config import AdamW as AdamW
+from .config import Checkpoint as Checkpoint
 from .config import Config as Config
-from .trainer import Trainer as Trainer
+from .config import Data as Data
+from .config import LRScheduler as LRScheduler
+from .config import Metrics as Metrics
+from .config import Model as Model
+from .config import Train as Train
+from .train import Trainer as Trainer

sarasa/checkpoint.py

@@ -110,3 +110,13 @@ class Checkpointer:
     def close(self) -> None:
         if self.stager is not None:
             self.stager.close()
+
+        if self.save_future is not None:
+            self.save_future.result()
+
+        if self.pg is not None:
+            dist.destroy_process_group(self.pg)
+            self.pg = None
+
+    def __del__(self):
+        self.close()

sarasa/models/__init__.py

@@ -21,6 +21,8 @@ class ModelConfig:
     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
+    rms_eps: float | None = None  # epsilon for RMSNorm, default to library default if None
+    rms_learnable: bool = False  # whether RMSNorm has learnable scale parameter
 
     def __post_init__(self):
         # infer hidden_dim, num_heads, num_kv_heads if not provided using the rules presented in nanochat

sarasa/models/attention.py

@@ -3,7 +3,7 @@ from torch import nn
 from torch.nn import functional as F
 
 from sarasa.models import ModelConfig
-from sarasa.models.utils import RMSNorm, RoPE
+from sarasa.models.utils import RoPE
 
 
 class SDPAttention(nn.Module):
@@ -57,7 +57,11 @@ class CausalSelfAttention(nn.Module):
         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()
+        self.qk_norm = (
+            nn.RMSNorm(self.head_dim, eps=config.rms_eps, elementwise_affine=config.rms_learnable)
+            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)

sarasa/models/llama3.py

@@ -4,7 +4,7 @@ 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
+from sarasa.models.utils import RoPE
 
 
 class MLP(nn.Module):
@@ -41,15 +41,16 @@ class Block(nn.Module):
         self.layer_idx = layer_idx
         self.attention = CausalSelfAttention(config)
         self.mlp = MLP(config, multiple_of, ffn_dim_multiplier)
-        self.norm = RMSNorm(config.hidden_dim)
+        self.attn_norm = nn.RMSNorm(config.hidden_dim, eps=config.rms_eps)
+        self.mlp_norm = nn.RMSNorm(config.hidden_dim, eps=config.rms_eps)
 
     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))
+        x = x + self.attention(self.attn_norm(x), cos_sin)
+        x = x + self.mlp(self.mlp_norm(x))
         return x
 
 
@@ -71,7 +72,7 @@ class Llama3(BaseModel):
         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.norm = nn.RMSNorm(config.hidden_dim, eps=config.rms_eps)
         self.output = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
 
     @torch.no_grad()
@@ -101,16 +102,26 @@ class Llama3(BaseModel):
             b=cutoff_factor * final_out_std,
         )
 
+        for mod in self.modules():
+            if isinstance(mod, nn.RMSNorm):
+                mod.reset_parameters()
+
     def param_groups(self) -> dict[str, list[nn.Parameter]]:
-        matrix_params = list(self.blocks.parameters())
+        matrix_params = [param for param in self.blocks.parameters() if param.ndim == 2]
         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))
+        rms_norm_params = [
+            mod.weight for mod in self.modules() if isinstance(mod, nn.RMSNorm) and mod.elementwise_affine
+        ]
+        assert len(list(self.parameters())) == (
+            len(matrix_params) + len(embedding_params) + len(lm_head_params) + len(rms_norm_params)
+        )
 
         return {
             "matrix": matrix_params,
             "embedding": embedding_params,
             "lm_head": lm_head_params,
+            "rms_norm": rms_norm_params,
         }
 
     def forward(

sarasa/models/nanochat_gpt.py

@@ -8,7 +8,16 @@ 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
+from sarasa.models.utils import RoPE
+
+
+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 MLP(nn.Module):

sarasa/models/utils.py

@@ -1,15 +1,6 @@
 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(

sarasa/{trainer.py → train.py}

@@ -7,6 +7,7 @@ import torch
 import torch.distributed as dist
 from loguru import logger
 from torch.distributed.elastic.multiprocessing.errors import record
+from torch.nn import functional as F
 
 from sarasa.activation_checkpoint import apply_op_sac
 from sarasa.checkpoint import Checkpointer
@@ -48,9 +49,6 @@ class Trainer:
         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()
@@ -68,9 +66,9 @@ class Trainer:
         if config.train.compile:
             logger.info("Compiling the model")
             for block in self.model.blocks:
-                block.compile(fullgraph=True)
+                block.compile(fullgraph=True, dynamic=False)
             self.model.compile(dynamic=False)
-            self.loss_fn.compile()
+            self.loss_fn = torch.compile(self.loss_fn, fullgraph=True, dynamic=False)
 
         if world_size() > 1:
             apply_distributed(
@@ -119,14 +117,6 @@ class Trainer:
                     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):
         try:
@@ -197,7 +187,7 @@ class Trainer:
 
             with self.amp_context:
                 pred = self.model(**input_dict)
-                loss = self.loss_fn(pred.flatten(0, 1), target.flatten(0, 1)) / valid_tokens
+                loss = self.loss_fn(pred, target) / valid_tokens
 
             del pred
             loss.backward()
@@ -241,6 +231,18 @@ class Trainer:
             },
         )
 
+    def loss_fn(
+        self,
+        pred: torch.Tensor,
+        target: torch.Tensor,
+    ) -> torch.Tensor:
+        return F.cross_entropy(
+            pred.flatten(0, 1).float(),
+            target.flatten(0, 1),
+            ignore_index=IGNORE_INDEX,
+            reduction="sum",
+        )
+
     def evaluate(self):
         raise NotImplementedError
 
@@ -256,3 +258,13 @@ class Trainer:
 
         if self.metrics_processor is not None:
             self.metrics_processor.close()
+
+        # cleanup distributed
+        if world_size() > 1:
+            try:
+                dist.destroy_process_group()
+            except Exception as e:
+                logger.warning(f"Failed to destroy process group: {e}")
+
+    def __del__(self):
+        self.close()

sarasa/utils.py

@@ -3,6 +3,7 @@ import gc
 import os
 import sys
 import time
+import typing
 from datetime import timedelta
 from functools import cache
 
@@ -11,8 +12,11 @@ from loguru import logger
 from torch import distributed as dist
 from torch import nn
 
+if typing.TYPE_CHECKING:
+    from sarasa.config import Config, Distributed
 
-def setup_logger(config) -> None:
+
+def setup_logger(config: Config) -> None:
     logger.remove()
     if config.debug:
         logger_format = f"<blue>RANK={rank()}</blue> | " + (
@@ -128,7 +132,7 @@ def update_timeout(
 
 
 def apply_distributed(
-    config,
+    config: Distributed,
     model: nn.Module,
     device: torch.device,
     compile: bool,

{sarasa-0.0.3.dist-info → sarasa-0.0.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sarasa
-Version: 0.0.3
+Version: 0.0.4
 Summary: Add your description here
 License-File: LICENSE
 Requires-Python: >=3.13
@@ -46,6 +46,8 @@ uv add sarasa[cpu|cu128|cu130]
 - Async distributed checkpoint saving
 
 - [ ] Checkpoint loading
+- [ ] FP8 training
+- [ ] Profiling
 
 ## Usage
 
@@ -100,18 +102,22 @@ if __name__ == "__main__":
     trainer.train()
 ```
 
+Thanks to [tyro](https://github.com/brentyi/tyro)'s type support, Sarasa can automatically recognize multiple custom optimizer types.
 From the command line, you can specify which custom optimizer to use:
 
 ```bash
 python script.py optim:custom_optim --optim.lr 0.001 ...
 ```
 
+(As tyro automatically converts config class names from CamelCase to snake_case, config class names are recommended not to include `Config` suffixes.)
+
 ### Config File Example
 
-It's very simple. IDE autocompletion will help you.
+It's very simple.
+IDE autocompletion will help you.
 
 ```python
-from sarasa.config import Config, Data, LRScheduler, Model, Train, LRScheduler
+from sarasa import Config, Data, LRScheduler, Model, Train, LRScheduler
 from custom_optim import CustomOptim
 
 # only one Config instance should be defined in each config file
@@ -135,4 +141,4 @@ config = Config.create(
 
 ## Acknowledgements
 
-This project is heavily inspired by and borrows code from `torchtitan`.
+This project is heavily inspired by and borrows code from [torchtitan](https://github.com/pytorch/torchtitan).

sarasa-0.0.4.dist-info/RECORD

@@ -0,0 +1,21 @@
+sarasa/__init__.py,sha256=cCbISDgbexMLgX9ff5zGZXcyXym3qN6tbzbH0rh2a1k,408
+sarasa/activation_checkpoint.py,sha256=iGib3e2GFxBLOtgcPLQZnzw0Ru6Gd_yFqWZmUw0Cfa4,3056
+sarasa/checkpoint.py,sha256=EjntYpZDGe5fWJhlmyeHcgLA0vjJUR7gTpziT0xWcKw,3849
+sarasa/config.py,sha256=7MNVs5GVCZ2ezpwGIGiyDCIqyRTi7vqnp2D6fWqHv5s,9299
+sarasa/metrics.py,sha256=OzTuK3Oed-I_2FC7rrE9FYi3NgTdKsDsVkWlGgJGh0M,10636
+sarasa/train.py,sha256=cRWUfNGGTNVUMZq77ltRoQaw5uZdn7N8zeWLXiP2F6M,9892
+sarasa/utils.py,sha256=3bSTXjAD5hRXIQxPmEbu7Bt709J-ha0CuHU_6Ieqvjc,4542
+sarasa/data/__init__.py,sha256=I0JOb9QrHEj9zXUX8kLir6ONAyiozeagzApig0WcSt8,1150
+sarasa/data/hf_datasets.py,sha256=DUlCpBOcDtZNEGrx4AtZTPW5IMtxIXMX_pKfnQEqzEg,3966
+sarasa/data/tokenizer.py,sha256=JhUOl9USJRM-DVPY02ouiaNUhAu1w2LLGquMnAyyA68,1752
+sarasa/models/__init__.py,sha256=1jxzxW-lg1Am38CHQhC9K1VD3U2J0Ca70kGCifhxlRA,3396
+sarasa/models/attention.py,sha256=9mI7j_k98OJlaUgRzWNvFHeVqBs07YJpMA9oefYJvaE,3109
+sarasa/models/llama3.py,sha256=2WACwUiS5y8rEmO7EcbXvPR3JtxhL6Jz9qRB28qzPvM,5018
+sarasa/models/nanochat_gpt.py,sha256=zKLgwYOHjohZKQaGwrRp-7uEobZhHKoWT1lvJiYJzZA,8936
+sarasa/models/utils.py,sha256=4aJD8HrJ1JPsR9615UmIFiDy6bwnjRffKIlGoYwolUg,887
+sarasa/optimizers/__init__.py,sha256=TH7CV-dexzVIm_NJKpo3VxnzwLvUjPyckD_oXMT48xo,1844
+sarasa/optimizers/utils.py,sha256=yI1_yHllJFyGbFW8jdMbvLfa5k7zUXjdgkvbr64mFOI,705
+sarasa-0.0.4.dist-info/METADATA,sha256=1D5-WB0KHgWhxDu8UcYp_wHHOCTvd7jXghzNq5CAvvo,3805
+sarasa-0.0.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+sarasa-0.0.4.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sarasa-0.0.4.dist-info/RECORD,,

sarasa-0.0.3.dist-info/RECORD

@@ -1,21 +0,0 @@
-sarasa/__init__.py,sha256=PrMUKyqcTKgKR3R3VMTzfuQ9EwMNVyw7qRl84Dz3d28,77
-sarasa/activation_checkpoint.py,sha256=iGib3e2GFxBLOtgcPLQZnzw0Ru6Gd_yFqWZmUw0Cfa4,3056
-sarasa/checkpoint.py,sha256=nZNo-qv3hvtzZuN0xw4SsCP4QmIM7E5nqGBxGfGYZo0,3616
-sarasa/config.py,sha256=7MNVs5GVCZ2ezpwGIGiyDCIqyRTi7vqnp2D6fWqHv5s,9299
-sarasa/metrics.py,sha256=OzTuK3Oed-I_2FC7rrE9FYi3NgTdKsDsVkWlGgJGh0M,10636
-sarasa/trainer.py,sha256=aRqa4ycbMeyLROP_hcehUpYif8jpGqCdNBHkt-TOTaw,9645
-sarasa/utils.py,sha256=Lee6wryHbob3PGpvR71gvpGjYNn-YsUaqoXvx6lx_tU,4431
-sarasa/data/__init__.py,sha256=I0JOb9QrHEj9zXUX8kLir6ONAyiozeagzApig0WcSt8,1150
-sarasa/data/hf_datasets.py,sha256=DUlCpBOcDtZNEGrx4AtZTPW5IMtxIXMX_pKfnQEqzEg,3966
-sarasa/data/tokenizer.py,sha256=JhUOl9USJRM-DVPY02ouiaNUhAu1w2LLGquMnAyyA68,1752
-sarasa/models/__init__.py,sha256=w9p4lZ0oEH2kRMxPh88Ogphwqx_o_Ik8Upfv9SrW7hA,3223
-sarasa/models/attention.py,sha256=rWm6NurkS5wdnzP_LPonCMJt_gQulySQDPFpVZtpGWU,3006
-sarasa/models/llama3.py,sha256=jGrrC2AQJvdyo_YvbGC4vmy23k-19Itsj_fSHUY2QTc,4509
-sarasa/models/nanochat_gpt.py,sha256=cpyoXwlWhqeUtgYSQ673AW6lh9BQ-64_9G9XURfZ1MY,8721
-sarasa/models/utils.py,sha256=_0F8yFVB2ZVClr8YppBhPOWtpNF0dSkmaElCYCkO_co,1111
-sarasa/optimizers/__init__.py,sha256=TH7CV-dexzVIm_NJKpo3VxnzwLvUjPyckD_oXMT48xo,1844
-sarasa/optimizers/utils.py,sha256=yI1_yHllJFyGbFW8jdMbvLfa5k7zUXjdgkvbr64mFOI,705
-sarasa-0.0.3.dist-info/METADATA,sha256=WXzEK22tSsKV1ZyWhnfsFTg5Ze9zYWJpe8MYjUH_V6M,3452
-sarasa-0.0.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-sarasa-0.0.3.dist-info/licenses/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sarasa-0.0.3.dist-info/RECORD,,