orbit-torch 0.0.4a1__py3-none-any.whl

orbit/optim/__init__.py

@@ -0,0 +1,2 @@
+from orbit.optim.sam  import SAM
+from orbit.optim.muon import Muon

orbit/optim/muon.py

@@ -0,0 +1,193 @@
+import math
+import torch
+
+# 此代码片段改编自以下 GitHub 仓库的修改版本：
+# https://github.com/KellerJordan/Muon/blob/master/muon.py
+@torch.compile
+def zeropower_via_newtonschulz5(G, steps):
+    '''
+    使用 Newton-Schulz 迭代计算 G 的零次幂/正交化。我们选择使用五次迭代，
+    其系数被选择为最大化零处的斜率。为了最小化步骤，经验表明，即使在迭代不再在区间上的
+    所有位置完全收敛到 1 的点之后，继续增加零处的斜率也是有效的。因此，此迭代不产生 UV^T，
+    而是产生类似 US'V^T 的结果，其中 S' 是对角矩阵，S_{ii}' ~ Uniform(0.5, 1.5)，
+    这在模型性能方面相对于 UV^T（其中 USV^T = G 是 SVD）完全没有负面影响。
+    '''
+    assert len(G.shape) == 2
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.bfloat16()
+    if G.size(0) > G.size(1):
+        X = X.T
+    # 确保谱范数至多为 1
+    X = X / (X.norm() + 1e-7)
+    # 执行 NS 迭代
+    for _ in range(steps):
+        A = X @ X.T
+        B = (
+            b * A + c * A @ A
+        )  # 改编自 @jxbz, @leloykun 和 @YouJiacheng 的建议
+        X = a * X + B @ X
+
+    if G.size(0) > G.size(1):
+        X = X.T
+    return X
+
+
+class Muon(torch.optim.Optimizer):
+    '''
+    Muon - MomentUm Orthogonalized by Newton-schulz (通过 Newton-Schulz 正交化的动量)
+
+    Muon 内部运行标准的 SGD-momentum，然后执行正交化后处理步骤，
+    其中每个 2D 参数的更新都被替换为最近的正交矩阵。为了有效地正交化每个更新，
+    我们使用 Newton-Schulz 迭代，其优点是可以在 GPU 上以 bfloat16 稳定运行。
+
+    一些警告：
+    - 我们认为此优化器不太可能在小批量训练中表现良好。
+    - 我们认为它可能不适合微调预训练模型，但我们尚未对此进行测试。
+
+    参数:
+        muon_params: 要由 Muon 优化的参数。
+        lr: 学习率。更新的谱范数将为 `lr`。（0.02 是一个很好的默认值）
+        momentum: 内部 SGD 使用的动量。（0.95 是一个很好的默认值）
+        nesterov: 是否在内部 SGD 中使用 Nesterov 风格的动量。（推荐）
+        ns_steps: 要运行的 Newton-Schulz 迭代次数。（6 可能总是足够的）
+        adamw_params: 要由 AdamW 优化的参数。`muon_params` 中任何 {0, 1}-D 参数
+            或被检测为嵌入或 lm_head 的参数也将由 AdamW 优化。
+        adamw_lr: 内部 AdamW 的学习率。
+        adamw_betas: 内部 AdamW 的 beta 参数。
+        adamw_eps: 内部 AdamW 的 epsilon 参数。
+        adamw_wd: 内部 AdamW 的权重衰减。
+    '''
+
+    def __init__(
+        self,
+        lr=1e-3,
+        wd=0.1,
+        muon_params=None,
+        momentum=0.95,
+        nesterov=True,
+        ns_steps=5,
+        adamw_params=None,
+        adamw_betas=(0.9, 0.95),
+        adamw_eps=1e-8,
+    ):
+
+        defaults = dict(
+            lr=lr,
+            wd=wd,
+            momentum=momentum,
+            nesterov=nesterov,
+            ns_steps=ns_steps,
+            adamw_betas=adamw_betas,
+            adamw_eps=adamw_eps,
+        )
+
+        params = list(muon_params)
+        adamw_params = list(adamw_params) if adamw_params is not None else []
+        params.extend(adamw_params)
+        super().__init__(params, defaults)
+        # 将参数分类为我们将使用 Muon 的参数和不使用的参数
+        for p in muon_params:
+            # 对 muon_params 中每个 >= 2D 且看起来不像嵌入或头层的参数使用 Muon
+            assert p.ndim == 2, p.ndim
+            self.state[p]['use_muon'] = True
+        for p in adamw_params:
+            # 对 adamw_params 中的参数不使用 Muon
+            self.state[p]['use_muon'] = False
+
+    def adjust_lr_for_muon(self, lr, param_shape):
+        A, B = param_shape[:2]
+        # 我们根据参数矩阵的大小调整学习率和权重衰减，如论文中所述
+        adjusted_ratio = 0.2 * math.sqrt(max(A, B))
+        adjusted_lr = lr * adjusted_ratio
+        return adjusted_lr
+
+    def step(self, closure=None):
+        '''
+        执行单个优化步骤。
+
+        Args:
+            closure (Callable, optional): 重新评估模型并返回损失的闭包。
+        '''
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+
+            ############################
+            #           Muon           #
+            ############################
+
+            params = [p for p in group['params'] if self.state[p]['use_muon']]
+            # import pdb; pdb.set_trace()
+            lr = group['lr']
+            wd = group['wd']
+            momentum = group['momentum']
+
+            # 生成权重更新
+            for p in params:
+                # 完整性检查
+                g = p.grad
+                if g is None:
+                    continue
+                if g.ndim > 2:
+                    g = g.view(g.size(0), -1)
+                assert g is not None
+
+                # 计算更新
+                state = self.state[p]
+                if 'momentum_buffer' not in state:
+                    state['momentum_buffer'] = torch.zeros_like(g)
+                buf = state['momentum_buffer']
+                buf.mul_(momentum).add_(g)
+                if group['nesterov']:
+                    g = g.add(buf, alpha=momentum)
+                else:
+                    g = buf
+                u = zeropower_via_newtonschulz5(g, steps=group['ns_steps'])
+
+                # 缩放更新
+                adjusted_lr = self.adjust_lr_for_muon(lr, p.shape)
+
+                # 应用权重衰减
+                p.data.mul_(1 - lr * wd)
+
+                # 应用更新
+                p.data.add_(u, alpha=-adjusted_lr)
+
+            ############################
+            #       AdamW backup       #
+            ############################
+
+            params = [p for p in group['params'] if not self.state[p]['use_muon']]
+            lr = group['lr']
+            beta1, beta2 = group['adamw_betas']
+            eps = group['adamw_eps']
+            weight_decay = group['wd']
+
+            for p in params:
+                g = p.grad
+                if g is None:
+                    continue
+                state = self.state[p]
+                if 'step' not in state:
+                    state['step'] = 0
+                    state['moment1'] = torch.zeros_like(g)
+                    state['moment2'] = torch.zeros_like(g)
+                state['step'] += 1
+                step = state['step']
+                buf1 = state['moment1']
+                buf2 = state['moment2']
+                buf1.lerp_(g, 1 - beta1)
+                buf2.lerp_(g.square(), 1 - beta2)
+
+                g = buf1 / (eps + buf2.sqrt())
+
+                bias_correction1 = 1 - beta1**step
+                bias_correction2 = 1 - beta2**step
+                scale = bias_correction1 / bias_correction2**0.5
+                p.data.mul_(1 - lr * weight_decay)
+                p.data.add_(g, alpha=-lr / scale)
+
+        return loss

orbit/optim/sam.py

@@ -0,0 +1,92 @@
+import torch
+
+class SAM(torch.optim.Optimizer):
+    '''锐度感知最小化优化器 (Sharpness-Aware Minimization)。
+
+    该优化器通过寻找损失景观中平坦的局部最小值来提高模型的泛化能力。
+    它执行两步梯度更新：首先寻找使损失最大化的扰动，然后在该点进行梯度更新。
+
+    Attributes:
+        base_optimizer (torch.optim.Optimizer): 基础优化器实例（如 SGD 或 Adam）。
+        param_groups (list): 包含优化器参数组的列表。
+        state (dict): 存储参数状态（如扰动向量 e_w）。
+    '''
+
+    def __init__(self, params, base_optimizer, rho=0.05, **kwargs):
+        '''初始化 SAM 优化器。
+
+        Args:
+            params (iterable): 可迭代的模型参数或定义参数组的字典。
+            base_optimizer (class): 基础优化器类（注意是类名，如 torch.optim.SGD）。
+            rho (float, optional): 邻域大小，用于控制扰动范围。默认为 0.05。
+            **kwargs: 传递给基础优化器的其他超参数（如 lr, momentum, weight_decay）。
+
+        Raises:
+            AssertionError: 如果 rho 小于 0。
+        '''
+        assert rho >= 0.0, f'无效的 rho 值: {rho}，必须是非负数。'
+
+        defaults = dict(rho=rho, **kwargs)
+        super(SAM, self).__init__(params, defaults)
+
+        self.base_optimizer = base_optimizer(self.param_groups, **kwargs)
+        self.param_groups = self.base_optimizer.param_groups
+
+    @torch.no_grad()
+    def first_step(self, zero_grad=False):
+        '''计算并应用参数扰动 epsilon。
+
+        该步骤根据当前梯度将模型参数 $w$ 更新为 $w + \epsilon$。
+
+        Args:
+            zero_grad (bool, optional): 是否在计算后清除梯度。默认为 False。
+        '''
+        grad_norm = self._grad_norm()
+        for group in self.param_groups:
+            scale = group['rho'] / (grad_norm + 1e-12)
+
+            for p in group['params']:
+                if p.grad is None: continue
+                e_w = p.grad * scale.to(p)
+                self.state[p]['e_w'] = e_w  # 存储扰动用于第二步恢复
+                p.add_(e_w)  # 参数移动到 w + e_w
+
+        if zero_grad: self.zero_grad()
+
+    @torch.no_grad()
+    def second_step(self, zero_grad=False):
+        '''恢复参数并执行真正的梯度更新。
+
+        该步骤将参数从 $w + \epsilon$ 恢复回 $w$，并利用扰动位置的梯度更新 $w$。
+
+        Args:
+            zero_grad (bool, optional): 是否在更新后清除梯度。默认为 False。
+        '''
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None: continue
+                p.sub_(self.state[p]['e_w'])  # 恢复到原始参数 w
+
+        self.base_optimizer.step()  # 使用 w + e_w 处的梯度进行实际更新
+
+        if zero_grad: self.zero_grad()
+
+    def _grad_norm(self):
+        '''计算所有参数梯度的全局 L2 范数。
+
+        Returns:
+            torch.Tensor: 梯度的 L2 范数标量。
+        '''
+        shared_device = self.param_groups[0]['params'][0].device
+        norm = torch.norm(
+                    torch.stack([
+                        p.grad.norm(p=2).to(shared_device)
+                        for group in self.param_groups for p in group['params']
+                        if p.grad is not None
+                    ]),
+                    p=2
+               )
+        return norm
+
+    def step(self, closure=None):
+        raise NotImplementedError('SAM requires steps to be run manually: first_step and second_step')

orbit/plugin/__init__.py

@@ -0,0 +1,10 @@
+try: from orbit.plugin.classification import ClassificationReport
+except: pass
+
+from orbit.plugin.warmup import Warmup
+from orbit.plugin.early_stopping import EarlyStopping
+from orbit.plugin.gradient_accumulation import GradientAccumulation
+from orbit.plugin.mentor import Mentor
+from orbit.plugin.ema import EMA # Not tested
+from orbit.plugin.memory_estimator import MemoryEstimator
+from orbit.plugin.overfit import Overfit

orbit/plugin/board.py

@@ -0,0 +1,61 @@
+from torch.utils.tensorboard import SummaryWriter
+from typing import Optional, TYPE_CHECKING
+from orbit.callback import Callback, Event
+
+if TYPE_CHECKING:
+    from ..engine import Engine
+
+class Board(Callback):
+    def __init__(self, name: str, log_dir: str):
+        super().__init__()
+        self.log_dir = log_dir + '/' + name
+        self.writer: Optional[SummaryWriter] = None
+
+    def on_init(self, event: Event):
+        '''初始化 SummaryWriter'''
+        # 如果 log_dir 不存在会自动创建
+        self.writer = SummaryWriter(log_dir=self.log_dir)
+        event.engine.writer = self.writer 
+        event.engine.print(f'[cyan]Initialized. Log dir: {self.log_dir}[/]', plugin='Board')
+
+    def on_batch_end(self, event: Event):
+        '''
+        每个 Batch 结束时记录：
+        1. Train Loss (Batch级)
+        2. Learning Rate
+        '''
+        engine = event.engine
+        if engine.state == 'TRAIN':
+            # 记录 Training Loss
+            if engine.loss is not None:
+                self.writer.add_scalar('Train/Batch_Loss', engine.loss.item(), engine.global_step)
+            
+            # 记录 Learning Rate (取第一个参数组)
+            if engine.optimizer:
+                current_lr = engine.optimizer.param_groups[0]['lr']
+                self.writer.add_scalar('Train/LR', current_lr, engine.global_step)
+
+    def on_epoch_end(self, event: Event):
+        '''
+        每 Epoch 结束时记录：
+        1. Epoch 平均 Loss (Train & Val)
+        2. 其他 Metrics (如果在 engine.metrics 字典里有的话)
+        '''
+        engine = event.engine
+        # engine.metrics: {'train_loss': 0.5, 'val_loss': 0.4, 'acc': 0.9}
+        for key, value in engine.metrics.items():
+            if 'loss' in key.lower():
+                tag = f'Loss/{key}'
+            elif 'acc' in key.lower():
+                tag = f'Accuracy/{key}'
+            else:
+                tag = f'Metrics/{key}'
+            
+            self.writer.add_scalar(tag, value, engine.epoch + 1)
+            
+        self.writer.flush()
+
+    def on_train_end(self, event: Event):
+        '''训练结束关闭 Writer'''
+        if self.writer:
+            self.writer.close()

orbit/plugin/checkpoint.py

@@ -0,0 +1,245 @@
+import os
+import torch
+from typing import TYPE_CHECKING, List, Tuple
+from orbit.callback import Callback, Event
+
+if TYPE_CHECKING:
+    from orbit.engine import Engine
+
+class Checkpoint(Callback):
+    def __init__(
+        self, 
+        name: str, 
+        path: str, 
+        save_weights_only: bool = False,
+        monitor: str = 'val_loss', # 默认监控 val_loss
+        mode: str = 'min',         # 默认 loss 越小越好
+        save_top_k: int = 1,
+        save_last: bool = True,
+        every_n_train_steps: int = None,
+        verbose: bool = True
+    ):
+        """
+        Args:
+            name (str): 模型名称前缀。
+            path (str): 保存目录。
+            save_weights_only (bool): 是否只保存模型权重 (不保存 optimizer 等状态)。
+            monitor (str): 监控指标 (例如 'val_loss', 'val_acc')。默认 'val_loss'。
+            mode (str): 'min' (越小越好) 或 'max' (越大越好)。
+            save_top_k (int): 保存最好的 K 个模型。设为 0 则禁用 Top-K 保存。
+            save_last (bool): 是否总是保存 '{name}_last.pt'。
+            every_n_train_steps (int): 每隔多少个训练步保存一次。
+            verbose (bool): 是否打印保存信息。
+        """
+        super().__init__()
+        self.name = name
+        self.path = path
+        self.save_weights_only = save_weights_only
+        self.monitor = monitor
+        self.mode = mode
+        self.save_top_k = save_top_k
+        self.save_last = save_last
+        self.every_n_train_steps = every_n_train_steps
+        self.verbose = verbose
+        
+        # 维护 Top-K 模型列表: [(score, filename), ...]
+        self.best_k_models: List[Tuple[float, str]] = []
+        
+        # 记录上一个 Step Checkpoint 文件名，用于删除
+        self.last_step_checkpoint: str = None
+        
+        # 内部状态 Key
+        self._meta_key = 'checkpoint_callback'
+
+    def on_init(self, event: Event):
+        """
+        1. 创建文件夹
+        2. 尝试恢复 Checkpoint 状态 (best_k_models)
+        3. 尝试加载 last checkpoint
+        """
+        engine = event.engine
+        if not os.path.exists(self.path):
+            os.makedirs(self.path, exist_ok=True)
+        
+        # 尝试恢复 best_k_models 状态
+        if self._meta_key in engine.meta:
+            self.best_k_models = engine.meta[self._meta_key].get('best_k_models', [])
+
+        load_path = os.path.join(self.path, self.name + "_last.pt").replace("\\", "/")
+        
+        if os.path.exists(load_path):
+            self._load(engine, load_path)
+        else:
+            engine.print(f"[yellow]Warning: Resume checkpoint '{load_path}' not found. Starting from scratch.[/]", plugin='Checkpointing')
+
+    def on_batch_end(self, event: Event):
+        """
+        每个 Batch 结束时：
+        检查是否需要按 Step 保存
+        """
+        if self.every_n_train_steps and event.engine.state == "TRAIN":
+            step = event.engine.global_step
+            if step > 0 and step % self.every_n_train_steps == 0:
+                # 保存 step checkpoint
+                filename = f"{self.name}_step_{step}.pt"
+                
+                # 传递 is_step=True
+                self._save(event.engine, filename, verbose=self.verbose, is_step=True)
+                
+                # 删除旧的 step checkpoint
+                if self.last_step_checkpoint and self.last_step_checkpoint != filename:
+                    self._remove(event.engine, self.last_step_checkpoint)
+                self.last_step_checkpoint = filename
+                
+                # 同时更新 last checkpoint
+                if self.save_last:
+                    self._save(event.engine, f"{self.name}_last.pt", verbose=False, is_step=True)
+
+    def on_epoch_end(self, event: Event):
+        """
+        每个 Epoch 结束时：
+        1. 保存 last
+        2. 如果设置了 monitor，保存 top_k
+        """
+        engine = event.engine
+        # 1. Save Last
+        if self.save_last:
+            self._save(engine, f"{self.name}_last.pt", verbose=False) # last 不需要每次都啰嗦
+        
+        # 2. Save Top K
+        if self.monitor and self.save_top_k > 0:
+            current_score = engine.metrics.get(self.monitor)
+            
+            if current_score is None:
+                if self.verbose:
+                    engine.print(f"[yellow]Metric '{self.monitor}' not found in metrics. Skipping Top-K save.[/]", plugin='Checkpointing')
+                return
+
+            self._check_and_save_top_k(engine, current_score)
+
+    def _check_and_save_top_k(self, engine: 'Engine', current_score: float):
+        """检查并保存 Top-K 模型"""
+        filename = f"{self.name}_ep{engine.epoch+1}_{self.monitor}_{current_score:.4f}.pt"
+        
+        # 逻辑简化：总是先加入，然后排序，如果超过 K 个，删除最差的
+        self.best_k_models.append((current_score, filename))
+        
+        # 排序
+        reverse = (self.mode == 'max')
+        self.best_k_models.sort(key=lambda x: x[0], reverse=reverse)
+        
+        # 如果列表过长，处理溢出
+        if len(self.best_k_models) > self.save_top_k:
+            worst_model = self.best_k_models.pop() # 移除最后一个（最差的）
+            worst_score, worst_filename = worst_model
+            
+            # 如果刚才加入的就是最差的，说明没进 Top K，不需要保存
+            if worst_filename == filename:
+                return 
+            
+            # 否则，保存新的，删除旧的最差的
+            self._save(engine, filename, verbose=self.verbose)
+            self._remove(engine, worst_filename)
+        else:
+            # 列表没满，直接保存
+            self._save(engine, filename, verbose=self.verbose)
+            
+        # 更新 Meta 状态
+        engine.meta[self._meta_key] = {'best_k_models': self.best_k_models}
+
+    def _save(self, engine: 'Engine', filename: str, verbose: bool = True, is_step: bool = False):
+        # 确保 meta 数据是最新的
+        if self.monitor:
+            engine.meta[self._meta_key] = {'best_k_models': self.best_k_models}
+
+        # 获取原始模型 (去除 DataParallel 包装) 以保证 Checkpoint 通用性
+        raw_model = engine.unwrap_model()
+
+        state = {
+            'epoch': engine.epoch,
+            'global_step': engine.global_step,
+            'batch_idx': engine.batch_idx,
+            'is_step': is_step,
+            'model_state_dict': raw_model.state_dict(),
+            'optimizer_state_dict': engine.optimizer.state_dict() if engine.optimizer else None,
+            'scheduler_state_dict': engine.scheduler.state_dict() if engine.scheduler else None,
+            'scaler_state_dict': engine.scaler.state_dict() if engine.scaler else None,
+            'meta': engine.meta,
+        }
+        if self.save_weights_only:
+            state = raw_model.state_dict()
+        
+        file_path = os.path.join(self.path, filename)
+        try:
+            torch.save(state, file_path)
+            if verbose:
+                # 使用相对路径显示，更简洁
+                rel_path = os.path.relpath(file_path)
+                engine.print(f"Saved checkpoint: {rel_path}", plugin='Checkpointing')
+        except Exception as e:
+            engine.print(f"[red]Failed to save checkpoint: {e}[/]", plugin='Checkpointing')
+
+    def _remove(self, engine: 'Engine', filename: str):
+        """删除旧的 Checkpoint 文件"""
+        file_path = os.path.join(self.path, filename)
+        if os.path.exists(file_path):
+            try:
+                os.remove(file_path)
+                if self.verbose:
+                    engine.print(f"[dim]Removed old checkpoint: {filename}[/]", plugin='Checkpointing')
+            except OSError as e:
+                engine.print(f"[red]Failed to remove checkpoint {filename}: {e}[/]", plugin='Checkpointing')
+
+    def _load(self, engine: 'Engine', file_path: str):
+        """加载 Checkpoint 的核心逻辑"""
+        engine.print(f"[cyan]Loading checkpoint from: {file_path}[/]", plugin='Checkpointing')
+        try:
+            # 加载到设备
+            checkpoint = torch.load(file_path, map_location=engine.device)
+            
+            # 获取原始模型以进行加载
+            raw_model = engine.unwrap_model()
+
+            # 1. 加载模型权重
+            if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+                raw_model.load_state_dict(checkpoint['model_state_dict'])
+            else:
+                raw_model.load_state_dict(checkpoint)
+                engine.print("[yellow]Loaded model weights only (legacy format).[/]", plugin='Checkpointing')
+                return 
+            
+            # 2. 恢复训练状态
+            if not self.save_weights_only:
+                if engine.optimizer and 'optimizer_state_dict' in checkpoint:
+                    engine.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+
+                if engine.scheduler and 'scheduler_state_dict' in checkpoint:
+                    engine.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+                
+                if engine.scaler and 'scaler_state_dict' in checkpoint:
+                    engine.scaler.load_state_dict(checkpoint['scaler_state_dict'])
+                
+                if 'meta' in checkpoint:
+                    engine.meta.update(checkpoint['meta'])
+
+                loaded_epoch = checkpoint.get('epoch', 0)
+                loaded_batch_idx = checkpoint.get('batch_idx', -1)
+                is_step = checkpoint.get('is_step', False)
+                
+                if is_step:
+                    # 如果是 Step Checkpoint，从当前 Epoch 的下一个 Batch 继续
+                    engine.start_epoch = loaded_epoch
+                    engine.start_batch_idx = loaded_batch_idx
+                    msg = f"Epoch {engine.start_epoch}, Batch {engine.start_batch_idx + 1}"
+                else:
+                    # 如果是 Epoch Checkpoint，从下一个 Epoch 开始
+                    engine.start_epoch = loaded_epoch + 1
+                    engine.start_batch_idx = -1
+                    msg = f"Epoch {engine.start_epoch}"
+
+                engine.global_step = checkpoint.get('global_step', 0)
+                
+                engine.print(f"[green]Successfully resumed training from {msg}, Global Step {engine.global_step}[/]", plugin='Checkpointing')
+                
+        except Exception as e:
+            engine.print(f"[red]Failed to load checkpoint: {e}[/]", plugin='Checkpointing')