orbit-torch 0.0.4a1__py3-none-any.whl

orbit/plugin/memory_estimator.py

@@ -0,0 +1,234 @@
+import gc
+import torch
+import torch.nn as nn
+from rich.panel import Panel
+from rich.table import Table
+from rich import box
+from typing import TYPE_CHECKING, Optional, Union
+
+from orbit.callback import Callback, Event
+
+if TYPE_CHECKING:
+    from orbit.engine import Engine
+
+class MemoryEstimator(Callback):
+    """
+    显存预估插件。
+    在训练开始前，通过运行一个虚拟 Batch 来预估显存使用峰值。
+    同时支持在训练过程中监控显存使用情况。
+    """
+    def __init__(self, verbose: bool = True, alert_threshold: Union[float, str, int] = 0.8, stop_threshold: Union[float, str, int] = 0.95, clean_interval: Optional[int] = None):
+        '''
+        Args:
+            verbose (bool): 是否打印预估报告。
+            alert_threshold (Union[float, str, int]): 警告阈值。
+                如果是 float <= 1.0，视为总显存的百分比 (例如 0.8 = 80%)。
+                如果是 str (例如 "4GB", "500MB") 或 int (字节数)，视为绝对值。
+            stop_threshold (Union[float, str, int]): 停止阈值。类型同上。
+            clean_interval (Optional[int]): 如果提供，则每隔多少个 Batch 执行一次显存清理 (gc.collect + empty_cache)。
+                这有助于解决显存缓慢泄漏的问题，但可能会轻微影响训练速度。
+        '''
+        super().__init__()
+        self.verbose = verbose
+        self.alert_threshold_arg = alert_threshold
+        self.stop_threshold_arg = stop_threshold
+        self.clean_interval = clean_interval
+        
+        self.alert_bytes = None
+        self.stop_bytes = None
+        
+        self.has_run = False
+        self.has_alerted = False
+
+    def _parse_threshold(self, threshold: Union[float, str, int], total_capacity: int) -> int:
+        if isinstance(threshold, float):
+            return int(threshold * total_capacity)
+        if isinstance(threshold, int):
+            return threshold
+        if isinstance(threshold, str):
+            s = threshold.upper().strip()
+            if s.endswith('GB'):
+                return int(float(s[:-2]) * (1024**3))
+            elif s.endswith('MB'):
+                return int(float(s[:-2]) * (1024**2))
+            elif s.endswith('KB'):
+                return int(float(s[:-2]) * 1024)
+            elif s.endswith('B'):
+                return int(float(s[:-1]))
+            else:
+                try:
+                    return int(float(s))
+                except ValueError:
+                    pass
+        raise ValueError(f"Invalid memory threshold format: {threshold}")
+
+    def on_batch_start(self, event: Event):
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+
+    def on_batch_end(self, event: Event):
+        if not torch.cuda.is_available():
+            return
+        
+        engine = event.engine
+        peak_memory = torch.cuda.max_memory_allocated()
+        total_capacity = torch.cuda.get_device_properties(engine.device).total_memory
+        
+        # 初始化阈值字节数
+        if self.stop_bytes is None:
+            self.stop_bytes = self._parse_threshold(self.stop_threshold_arg, total_capacity)
+        if self.alert_bytes is None:
+            self.alert_bytes = self._parse_threshold(self.alert_threshold_arg, total_capacity)
+        
+        # 格式化辅助函数
+        to_mb = lambda x: x / (1024 ** 2)
+        
+        if peak_memory > self.stop_bytes:
+            engine.print(f"[bold red]Memory usage ({to_mb(peak_memory):.2f} MB) exceeded critical threshold ({to_mb(self.stop_bytes):.2f} MB)! Stopping training.[/]", plugin='MemEst')
+            engine.stop(source="MemoryEstimator", reason=f"Memory usage exceeded critical threshold ({to_mb(self.stop_bytes):.2f} MB)")
+        elif peak_memory > self.alert_bytes and not self.has_alerted:
+            engine.print(f"[yellow]Memory usage ({to_mb(peak_memory):.2f} MB) exceeded warning threshold ({to_mb(self.alert_bytes):.2f} MB).[/]", plugin='MemEst')
+            self.has_alerted = True
+
+        # 定期清理显存
+        if self.clean_interval and (engine.batch_idx + 1) % self.clean_interval == 0:
+            gc.collect()
+            torch.cuda.empty_cache()
+
+    def on_train_start(self, event: Event):
+        if self.has_run:
+            return
+        
+        engine = event.engine
+        if not torch.cuda.is_available():
+            if self.verbose:
+                engine.print("[yellow]CUDA not available. Skipping memory estimation.[/]", plugin='MemEst')
+            return
+
+        # 确保模型在正确的设备上
+        device = engine.device
+        if device.type != 'cuda':
+            return
+
+        try:
+            self._estimate(engine)
+        except Exception as e:
+            engine.print(f"[red]Error during memory estimation: {e}[/]", plugin='MemEst')
+        finally:
+            # 清理
+            if engine.optimizer:
+                engine.optimizer.zero_grad()
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
+            self.has_run = True
+
+    def _estimate(self, engine: 'Engine'):
+        if self.verbose:
+            engine.print("Running dry run for memory estimation...[/]", plugin='MemEst')
+        
+        # 1. 获取一个 Batch 的数据
+        try:
+            batch_data = next(iter(engine.train_loader))
+        except StopIteration:
+            engine.print("[yellow]Train loader is empty. Skipping.[/]", plugin='MemEst')
+            return
+
+        # 2. 准备环境
+        torch.cuda.empty_cache()
+        torch.cuda.reset_peak_memory_stats()
+        initial_memory = torch.cuda.memory_allocated()
+        
+        # 计算模型静态大小 (Weights + Buffers)
+        model_stats = self._get_model_size(engine.model)
+        
+        # 3. 模拟 Forward & Backward
+        try:
+            # 移动数据
+            engine._process_batch_data(batch_data)
+            
+            # Forward
+            with torch.amp.autocast(device_type=engine.device.type, enabled=engine.use_amp):
+                if isinstance(engine.data, (list, tuple)):
+                    output = engine.model(*engine.data)
+                else:
+                    output = engine.model(engine.data)
+                
+                # 构造虚拟 Loss
+                if engine.criterion and engine.target is not None:
+                    loss = engine.criterion(output, engine.target)
+                else:
+                    # 如果没有 target 或 criterion，构造一个标量 loss 用于 backward
+                    if isinstance(output, torch.Tensor):
+                        loss = output.mean()
+                    elif isinstance(output, (list, tuple)) and isinstance(output[0], torch.Tensor):
+                        loss = output[0].mean()
+                    elif isinstance(output, dict):
+                        loss = list(output.values())[0].mean()
+                    else:
+                        loss = torch.tensor(0.0, device=engine.device, requires_grad=True)
+
+            # Backward
+            if engine.use_amp and engine.scaler:
+                engine.scaler.scale(loss).backward()
+            else:
+                loss.backward()
+
+            # 获取峰值显存
+            peak_memory = torch.cuda.max_memory_allocated()
+            total_capacity = torch.cuda.get_device_properties(engine.device).total_memory
+            
+            self._print_report(engine, model_stats, initial_memory, peak_memory, total_capacity)
+
+        except RuntimeError as e:
+            if "out of memory" in str(e):
+                engine.print("[bold red]OOM detected during memory estimation![/]", plugin='MemEst')
+                engine.print(f"[red]Your batch size is likely too large for this device.[/]", plugin='MemEst')
+            else:
+                raise e
+
+    def _get_model_size(self, model: nn.Module) -> float:
+        """计算模型参数和缓冲区的总字节数"""
+        mem_params = sum([param.nelement() * param.element_size() for param in model.parameters()])
+        mem_bufs = sum([buf.nelement() * buf.element_size() for buf in model.buffers()])
+        return mem_params + mem_bufs
+
+    def _print_report(self, engine: 'Engine', model_size: int, initial: int, peak: int, capacity: int):
+        if not self.verbose: return
+
+        # 转换单位为 MB
+        to_mb = lambda x: x / (1024 ** 2)
+        
+        model_mb = to_mb(model_size)
+        peak_mb = to_mb(peak)
+        capacity_mb = to_mb(capacity)
+        usage_percent = (peak / capacity) * 100
+        
+        # 颜色编码
+        if usage_percent < 70:
+            color = "green"
+            status = "Safe"
+        elif usage_percent < 90:
+            color = "yellow"
+            status = "Warning"
+        else:
+            color = "red"
+            status = "Critical"
+
+        table = Table(box=box.SIMPLE, show_header=False)
+        table.add_column("Item", style="cyan")
+        table.add_column("Value", justify="right")
+
+        table.add_row("Model Weights", f"{model_mb:.2f} MB")
+        table.add_row("Est. Peak Memory", f"[{color}]{peak_mb:.2f} MB[/]")
+        table.add_row("Device Capacity", f"{capacity_mb:.2f} MB")
+        table.add_row("Usage", f"[{color}]{usage_percent:.1f}% ({status})[/]")
+        
+        panel = Panel(
+            table,
+            title="[bold]Memory Estimation Report[/]",
+            border_style="blue",
+            expand=False
+        )
+        
+        with engine.out_logs:
+            engine.console.print(panel)

orbit/plugin/mentor.py

@@ -0,0 +1,313 @@
+import numpy as np
+import json
+import os
+import locale
+from typing import Optional, List, TYPE_CHECKING, Tuple, Dict
+from rich.panel import Panel
+from orbit.callback import Callback, Event
+
+if TYPE_CHECKING:
+    from orbit.engine import Engine
+    from orbit.plugin.warmup import Warmup
+
+class Mentor(Callback):
+    """
+    Mentor 插件：监控训练过程，提供改进建议。
+    主要关注 Loss 的异常行为（停滞、发散、过拟合、震荡）并结合状态给出建议。
+    """
+    def __init__(
+        self,
+        patience: int = 3,
+        threshold: float = 1e-4,
+        divergence_threshold: float = 2.0,
+        language: Optional[str] = None,
+        verbose: bool = True
+    ):
+        """
+        Args:
+            patience (int): 连续多少个 Epoch Loss 没有显著改善视为停滞。
+            threshold (float): 视为改善的最小 Loss 变化量。
+            divergence_threshold (float): 当前 Loss 是最小 Loss 的多少倍视为发散。
+            language (str): 语言代码 ('en' 或 'zh')。如果为 None，则自动检测系统语言。
+            verbose (bool): 是否打印建议。
+        """
+        super().__init__()
+        self.patience = patience
+        self.threshold = threshold
+        self.divergence_threshold = divergence_threshold
+        self.language = language if language else self._detect_language()
+        self.verbose = verbose
+        
+        self.loss_history: List[float] = []
+        self.val_loss_history: List[float] = []
+        self.min_loss = np.inf
+        self.stagnation_counter = 0
+        self.increase_counter = 0
+        
+        # 记录 Warmup 相关信息
+        self.has_warmup = False
+        self.warmup_plugin: Optional['Warmup'] = None
+        
+        # 记录其他配置
+        self.has_scheduler = False
+        self.grad_clip_norm = None
+        self.batch_size = 0
+        self.accum_steps = 1
+        
+        # 加载翻译
+        self.i18n = self._load_i18n()
+
+    def _detect_language(self) -> str:
+        try:
+            lang_code, _ = locale.getdefaultlocale()
+            if lang_code and lang_code.startswith('zh'):
+                return 'zh'
+        except:
+            pass
+        return 'en'
+
+    def _load_i18n(self) -> Dict[str, str]:
+        try:
+            # 获取当前文件所在目录
+            current_dir = os.path.dirname(os.path.abspath(__file__))
+            json_path = os.path.join(current_dir, 'data', 'mentor_i18n.json')
+            
+            with open(json_path, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+                
+            return data.get(self.language, data.get('en', {}))
+        except Exception as e:
+            print(f"[Mentor] Warning: Failed to load i18n file: {e}. Fallback to keys.")
+            return {}
+
+    def _t(self, key: str, **kwargs) -> str:
+        """获取翻译并格式化"""
+        template = self.i18n.get(key, key)
+        try:
+            return template.format(**kwargs)
+        except KeyError:
+            return template
+
+    def on_train_start(self, event: Event):
+        engine = event.engine
+        # 1. 检查 Warmup
+        for p in engine.plugins:
+            if p.__class__.__name__ == 'Warmup':
+                self.has_warmup = True
+                self.warmup_plugin = p
+                break
+        
+        # 2. 检查 Scheduler & Gradient Clipping
+        self.has_scheduler = engine.scheduler is not None
+        self.grad_clip_norm = getattr(engine, 'grad_clip_norm', None)
+        
+        # 3. 检查 Batch Size & Accumulation
+        self.accum_steps = engine.accumulation_steps
+        if hasattr(engine, 'train_loader') and hasattr(engine.train_loader, 'batch_size'):
+            self.batch_size = engine.train_loader.batch_size or 1 # Handle None
+        
+        if self.verbose:
+            eff_bs = self.batch_size * self.accum_steps
+            msg = self._t("mentor_watching", eff_bs=eff_bs, bs=self.batch_size, accum=self.accum_steps)
+            engine.print(msg, plugin='Mentor')
+
+    def on_epoch_end(self, event: Event):
+        engine = event.engine
+        current_loss = engine.metrics.get('train_loss')
+        val_loss = engine.metrics.get('val_loss')
+        
+        if current_loss is None:
+            return
+
+        self.loss_history.append(current_loss)
+        if val_loss is not None:
+            self.val_loss_history.append(val_loss)
+        
+        # 收集本 Epoch 的所有建议 (Title, Content, Style)
+        advice_list: List[Tuple[str, str, str]] = []
+
+        # 1. 检查 NaN / Inf
+        if not np.isfinite(current_loss):
+            advice_list.append((
+                self._t("nan_loss_title"),
+                self._t("nan_loss_msg"),
+                "bold red"
+            ))
+        else:
+            # 更新最小 Loss 和 连续上升计数
+            if current_loss < self.min_loss:
+                self.min_loss = current_loss
+                self.stagnation_counter = 0
+                self.increase_counter = 0
+            else:
+                self.stagnation_counter += 1
+                # 检查是否比上一个 epoch 增加
+                if len(self.loss_history) > 1 and current_loss > self.loss_history[-2]:
+                    self.increase_counter += 1
+                else:
+                    self.increase_counter = 0
+
+            # 2. 检查发散 (Divergence)
+            is_diverging = (current_loss > self.min_loss * self.divergence_threshold)
+            is_unstable = (self.increase_counter >= 2)
+            
+            if (is_diverging or is_unstable) and len(self.loss_history) > 1:
+                res = self._analyze_divergence(engine, current_loss, is_unstable)
+                if res: advice_list.append(res)
+                
+                # 重置以避免每个 epoch 都刷屏
+                if is_diverging:
+                    self.min_loss = current_loss 
+                if is_unstable:
+                    self.increase_counter = 0
+
+            # 3. 检查停滞 (Stagnation)
+            if self.stagnation_counter >= self.patience:
+                res = self._analyze_stagnation(engine, current_loss)
+                if res: advice_list.append(res)
+                self.stagnation_counter = 0 # 重置计数器
+
+            # 4. 检查过拟合 (Overfitting)
+            if len(self.val_loss_history) >= 3:
+                res = self._analyze_overfitting(engine)
+                if res: advice_list.append(res)
+
+            # 5. 检查震荡 (Oscillation)
+            if len(self.loss_history) >= 5:
+                res = self._analyze_oscillation(engine)
+                if res: advice_list.append(res)
+
+        # 统一打印
+        if advice_list and self.verbose:
+            with engine.out_logs:
+                for title, content, style in advice_list:
+                    engine.console.print(Panel(content, title=title, border_style=style, expand=False))
+
+    def _analyze_divergence(self, engine: 'Engine', current_loss: float, is_unstable: bool = False) -> Tuple[str, str, str]:
+        in_warmup = engine.is_in_warmup()
+        eff_bs = self.batch_size * self.accum_steps
+        
+        title = self._t("divergence_title")
+        if is_unstable:
+            msg = self._t("divergence_msg_unstable", count=self.increase_counter, loss=current_loss)
+        else:
+            msg = self._t("divergence_msg_spike", loss=current_loss, min_loss=self.min_loss)
+        
+        advice = []
+        
+        # 1. Warmup 建议
+        if not self.has_warmup:
+            advice.append(self._t("advice_add_warmup"))
+            advice.append(self._t("advice_lower_lr"))
+        elif in_warmup:
+            advice.append(self._t("advice_warmup_start_lr"))
+        else:
+            advice.append(self._t("advice_post_warmup_lr"))
+        
+        # 2. Batch Size & Accumulation 建议
+        if eff_bs < 32:
+            advice.append(self._t("advice_small_bs", eff_bs=eff_bs))
+            advice.append(self._t("advice_increase_accum", accum_steps=self.accum_steps))
+        
+        # 3. 通用建议
+        advice.append(self._t("advice_grad_clip"))
+            
+        return title, msg + "\n\n" + "\n".join(advice), "red"
+
+    def _analyze_stagnation(self, engine: 'Engine', current_loss: float) -> Tuple[str, str, str]:
+        in_warmup = engine.is_in_warmup()
+        eff_bs = self.batch_size * self.accum_steps
+        
+        title = self._t("stagnation_title")
+        msg = self._t("stagnation_msg", patience=self.patience)
+        
+        advice = []
+        
+        if in_warmup:
+            # 在 Warmup 期间停滞
+            warmup_epochs = getattr(self.warmup_plugin, 'warmup_epochs', 0)
+            total_epochs = engine.num_epochs
+            
+            advice.append(self._t("advice_warmup_duration", epoch=engine.epoch+1))
+            if warmup_epochs > total_epochs * 0.2:
+                advice.append(self._t("advice_warmup_too_long", warmup_epochs=warmup_epochs))
+            advice.append(self._t("advice_check_start_lr"))
+            
+        else:
+            # 非 Warmup 期间停滞
+            advice.append(self._t("advice_lr_general"))
+            
+            # Scheduler 建议
+            if not self.has_scheduler:
+                advice.append(self._t("advice_add_scheduler"))
+            else:
+                advice.append(self._t("advice_check_scheduler"))
+            
+            # Batch Size 建议
+            if eff_bs > 4096:
+                advice.append(self._t("advice_large_bs", eff_bs=eff_bs))
+                advice.append(self._t("advice_reduce_bs"))
+
+            # 4. 初始化建议 (仅在早期)
+            if engine.epoch < 10:
+                advice.append(self._t("advice_check_init"))
+
+            # 5. 调试与数据建议
+            advice.append(self._t("advice_overfit_single_batch"))
+            advice.append(self._t("advice_data_hard"))
+
+        return title, msg + "\n\n" + "\n".join(advice), "yellow"
+
+    def _analyze_overfitting(self, engine: 'Engine') -> Optional[Tuple[str, str, str]]:
+        """检测过拟合：Train Loss 下降，Val Loss 上升"""
+        if len(self.val_loss_history) < 3 or len(self.loss_history) < 3:
+            return None
+            
+        # 检查最近 3 个 epoch
+        recent_val = self.val_loss_history[-3:]
+        recent_train = self.loss_history[-3:]
+        
+        # Val Loss 持续上升
+        val_rising = (recent_val[-1] > recent_val[-2] > recent_val[-3])
+        # Train Loss 持续下降 (或保持低位)
+        train_dropping = (recent_train[-1] <= recent_train[-2])
+        
+        if val_rising and train_dropping:
+            title = self._t("overfitting_title")
+            msg = self._t("overfitting_msg")
+            advice = [
+                self._t("advice_regularization"),
+                self._t("advice_data_aug"),
+                self._t("advice_early_stopping")
+            ]
+            return title, msg + "\n\n" + "\n".join(advice), "magenta"
+        return None
+
+    def _analyze_oscillation(self, engine: 'Engine') -> Optional[Tuple[str, str, str]]:
+        """检测震荡：Loss 标准差过大"""
+        if len(self.loss_history) < 5:
+            return None
+            
+        recent_loss = self.loss_history[-5:]
+        std_dev = np.std(recent_loss)
+        mean_loss = np.mean(recent_loss)
+        
+        # 如果标准差超过均值的 10% (经验值)，且没有持续下降趋势
+        # 简单的趋势检查：首尾差异不大
+        is_flat_trend = abs(recent_loss[-1] - recent_loss[0]) < std_dev
+        
+        if std_dev > 0.1 * mean_loss and is_flat_trend:
+            title = self._t("oscillation_title")
+            msg = self._t("oscillation_msg", std=std_dev)
+            advice = [
+                self._t("advice_lower_lr_oscillation")
+            ]
+            
+            if not self.has_scheduler:
+                advice.append(self._t("advice_oscillation_scheduler"))
+            
+            if not self.grad_clip_norm:
+                advice.append(self._t("advice_oscillation_grad_clip"))
+                
+            return title, msg + "\n\n" + "\n".join(advice), "cyan"
+        return None

orbit/plugin/overfit.py

@@ -0,0 +1,30 @@
+from orbit.callback import Callback, Event
+from typing import Any
+
+class Overfit(Callback):
+    '''故意重复第一个 Batch 的数据以进行过拟合测试的插件。
+    
+    这对于验证模型架构是否有能力拟合数据非常有用（Sanity Check）。
+    如果模型无法在单个 Batch 上过拟合，则可能存在代码错误或架构问题。
+    '''
+    
+    def __init__(self):
+        self.fixed_data: Any = None
+        self.fixed_target: Any = None
+        self.has_captured = False
+
+    def on_batch_start(self, event: Event):
+        # 仅在训练阶段生效
+        if event.engine.state != "TRAIN":
+            return
+
+        if not self.has_captured:
+            # 捕获第一个 Batch
+            self.fixed_data = event.engine.data
+            self.fixed_target = event.engine.target
+            self.has_captured = True
+            event.engine.print("[yellow]Overfit Plugin: Captured first batch. All subsequent batches will be replaced by this one.[/]", plugin='Overfit')
+        else:
+            # 替换后续 Batch
+            event.engine.data = self.fixed_data
+            event.engine.target = self.fixed_target

orbit/plugin/warmup.py

@@ -0,0 +1,119 @@
+from typing import Optional, List, TYPE_CHECKING
+from orbit.callback import Callback, Event
+
+if TYPE_CHECKING:
+    from orbit.engine import Engine
+
+class Warmup(Callback):
+    """
+    学习率预热 (Warmup) 插件。
+    支持 Linear, Constant, Noam (Transformer) 三种模式。
+    可以在 Batch 粒度上动态调整学习率。
+    """
+    def __init__(
+        self,
+        warmup_steps: int = 0,
+        warmup_epochs: int = 0,
+        mode: str = 'linear',
+        min_lr: float = 0.0,
+        model_dim: Optional[int] = None,
+        scale: float = 1.0,
+    ):
+        """
+        Args:
+            warmup_steps (int): 预热的总步数 (Batch 数)。
+            warmup_epochs (int): 预热的总 Epoch 数。如果设置，优先级高于 warmup_steps。
+            mode (str): 'linear' | 'constant' | 'noam'。
+            min_lr (float): 起始学习率 (linear/constant 模式)。
+            model_dim (int): 模型维度 (仅 noam 模式需要)。
+            scale (float): 缩放因子 (仅 noam 模式需要)。
+        """
+        super().__init__()
+        self.warmup_steps = warmup_steps
+        self.warmup_epochs = warmup_epochs
+        self.mode = mode.lower()
+        self.min_lr = min_lr
+        self.model_dim = model_dim
+        self.scale = scale
+
+        self.total_warmup_steps = 0
+        self.base_lrs: List[float] = []
+        
+        # 验证参数
+        if self.mode == 'noam' and self.model_dim is None:
+            raise ValueError("Noam mode requires 'model_dim' to be specified.")
+
+    def on_train_start(self, event: Event):
+        """
+        训练开始时计算总预热步数并记录初始学习率
+        """
+        engine = event.engine
+        if not engine.optimizer:
+            raise ValueError("Warmup plugin requires an optimizer in the Engine.")
+
+        # 1. 记录优化器的初始学习率 (Base LR)
+        # 如果是从 Checkpoint 恢复，base_lrs 可能会变，但在 Warmup 逻辑里我们通常认为
+        # 预热的目标就是 param_group['initial_lr'] (如果存在) 或者当前的 ['lr']
+        self.base_lrs = []
+        for group in engine.optimizer.param_groups:
+            # 优先使用 initial_lr (由某些 Scheduler 设置)，否则使用当前 lr
+            self.base_lrs.append(group.get('initial_lr', group['lr']))
+
+        # 2. 计算 total_warmup_steps
+        if self.warmup_epochs > 0:
+            if not engine.train_loader:
+                raise ValueError("warmup_epochs requires train_loader to be available.")
+            try:
+                steps_per_epoch = len(engine.train_loader)
+                self.total_warmup_steps = self.warmup_epochs * steps_per_epoch
+            except TypeError:
+                 # 如果 train_loader 无法求 len (例如 iterable dataset)，则必须提供 steps
+                 raise ValueError("Could not determine length of train_loader. Please use 'warmup_steps' instead.")
+        else:
+            self.total_warmup_steps = self.warmup_steps
+
+        # 打印信息
+        if self.total_warmup_steps > 0 or self.mode == 'noam':
+            engine.print(f"[magenta]Strategy activated: {self.mode}[/]", plugin='Warmup')
+            if self.mode != 'noam':
+                engine.print(f"[magenta]Steps: {self.total_warmup_steps} (Epochs: {self.warmup_epochs})[/]", plugin='Warmup')
+
+    def on_batch_start(self, event: Event):
+        """
+        每个 Batch 开始前调整学习率
+        """
+        engine = event.engine
+        # 当前步数 (从 1 开始计算，方便公式)
+        current_step = engine.global_step + 1
+        
+        # 如果超出预热范围且不是 Noam 模式，则不进行干预
+        # Noam 模式是一种全程调度策略，所以它会一直运行
+        if self.mode != 'noam' and current_step > self.total_warmup_steps:
+            return
+
+        for i, group in enumerate(engine.optimizer.param_groups):
+            base_lr = self.base_lrs[i]
+            new_lr = base_lr
+
+            if self.mode == 'noam':
+                # Noam Scheduler: scale * d_model^-0.5 * min(step^-0.5, step * warmup^-1.5)
+                warmup = self.total_warmup_steps
+                # 避免 step=0
+                step = max(1, current_step)
+                
+                term1 = step ** -0.5
+                term2 = step * (warmup ** -1.5)
+                
+                new_lr = self.scale * (self.model_dim ** -0.5) * min(term1, term2)
+                
+            elif self.mode == 'linear':
+                # Linear: min_lr -> base_lr
+                alpha = current_step / self.total_warmup_steps
+                new_lr = self.min_lr + (base_lr - self.min_lr) * alpha
+                
+            elif self.mode == 'constant':
+                # Constant: min_lr
+                new_lr = self.min_lr
+            
+            # 更新学习率
+            group['lr'] = new_lr