orbit-torch 0.0.4a1__py3-none-any.whl → 0.1.0b1__py3-none-any.whl

orbit/plugin/lora.py

@@ -0,0 +1,346 @@
+import os
+import torch
+import inspect
+import json
+
+from typing import Optional, List, TYPE_CHECKING
+from orbit.callback import Event
+from orbit.plugin.checkpoint import Checkpoint, safe_save_file, safe_load_file, safe_open
+from orbit.utils.lora import inject_lora, freeze_backbone_only
+
+if TYPE_CHECKING: from orbit.engine import Engine
+
+class LoRA(Checkpoint):
+    '''LoRA 插件：集成 LoRA 注入、冻结、轻量化保存与加载功能。
+
+    该插件继承自 Checkpoint，在训练开始时自动将模型转换为 LoRA 模型，
+    重建优化器以适应新参数，并提供仅保存训练参数的轻量化 Checkpoint 功能。
+    '''
+    def __init__(
+        self,
+        name: str = "lora_model",
+        path: str = "checkpoints",
+        # LoRA 参数
+        r: int = 8,
+        lora_alpha: int = 16,
+        lora_dropout: float = 0.05,
+        target_names: Optional[List[str]] = None,
+        exclude_names: Optional[List[str]] = None,
+        unlock_head_keywords: Optional[List[str]] = None,
+        gate: bool = False,
+        dora: bool = False,
+        # Checkpoint 参数
+        monitor: str = 'val_loss',
+        mode: str = 'min',
+        save_top_k: int = 1,
+        save_last: bool = True,
+        every_n_train_steps: Optional[int] = None,
+        use_safetensors: bool = False,
+        verbose: bool = True
+    ):
+        '''初始化 LoRA 插件。
+
+        Args:
+            name (str): 模型名称前缀。
+            path (str): Checkpoint 保存目录。
+            r (int): LoRA 秩。
+            lora_alpha (int): LoRA 缩放系数。
+            lora_dropout (float): LoRA Dropout。
+            target_names (list, optional): 仅注入包含这些名称的层。
+            exclude_names (list, optional): 排除包含这些名称的层。
+            unlock_head_keywords (list, optional): 除了 LoRA 层外，还需要解冻的层关键字（如分类头）。
+            gate (bool): 是否使用 Gated LoRA。
+            dora (bool): 是否使用 DoRA。
+            monitor (str): 监控指标。
+            mode (str): 监控指标模式 ('min'/'max')。
+            save_top_k (int): 保存最佳模型数量。
+            save_last (bool): 是否保存最后的模型。
+            every_n_train_steps (int, optional): 每 N 步保存一次。
+            use_safetensors (bool): 是否使用 safetensors 格式保存。
+            verbose (bool): 是否打印日志。
+        '''
+        # 初始化 Checkpoint，强制 save_weights_only=False 以保留训练状态，
+        # 但我们会在 _save 中自定义过滤逻辑。
+        super().__init__(
+            name=name,
+            path=path,
+            save_weights_only=False,
+            monitor=monitor,
+            mode=mode,
+            save_top_k=save_top_k,
+            save_last=save_last,
+            every_n_train_steps=every_n_train_steps,
+            use_safetensors=use_safetensors,
+            verbose=verbose
+        )
+        
+        self.r = r
+        self.lora_alpha = lora_alpha
+        self.lora_dropout = lora_dropout
+        self.target_names = target_names
+        self.exclude_names = exclude_names
+        self.unlock_head_keywords = unlock_head_keywords
+        self.gate = gate
+        self.dora = dora
+        
+        self.injected = False
+
+    def on_init(self, event: Event):
+        engine = event.engine
+        
+        # 0. 冲突检测：检查是否存在其他 Checkpoint 插件
+        other_checkpoints = [p for p in engine.plugins if isinstance(p, Checkpoint) and p is not self]
+        if other_checkpoints:
+            engine.print("[yellow]Warning: Multiple Checkpoint plugins detected. Since 'LoRA' inherits from 'Checkpoint', using both may cause conflicts (e.g. double saving). Suggest removing the standard 'Checkpoint' plugin.[/]", plugin='LoRA')
+
+        model = engine.unwrap_model()
+        
+        # 1. 注入 LoRA 并冻结骨干
+        if not self.injected:
+            engine.print(f"[cyan]Injecting LoRA (r={self.r}, alpha={self.lora_alpha})...[/]", plugin='LoRA')
+            inject_lora(
+                model, 
+                r=self.r, 
+                lora_alpha=self.lora_alpha, 
+                lora_dropout=self.lora_dropout,
+                gate=self.gate,
+                dora=self.dora,
+                target_names=self.target_names,
+                exclude_names=self.exclude_names
+            )
+            
+            freeze_backbone_only(
+                model, 
+                unlock_head_keywords=self.unlock_head_keywords,
+                verbose=self.verbose
+            )
+            
+            model.to(engine.device)
+            
+            # 2. 重建 Optimizer
+            # 由于参数集发生变化（新增了 LoRA 参数，冻结了大部分参数），旧的优化器无法使用。
+            # 我们尝试使用旧优化器的类和 defaults 重新初始化。
+            if engine.optimizer:
+                old_opt = engine.optimizer
+                trainable_params = [p for p in model.parameters() if p.requires_grad]
+                
+                if not trainable_params:
+                    engine.print("[red]Warning: No trainable parameters found after LoRA injection![/]", plugin='LoRA')
+                else:
+                    opt_cls = old_opt.__class__
+                    defaults = old_opt.defaults
+                    
+                    # 过滤掉不在构造函数中的参数，防止某些库（如 transformers）在 defaults 中添加了额外元数据导致重建失败
+                    sig = inspect.signature(opt_cls.__init__)
+                    has_kwargs = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values())
+                    if has_kwargs:
+                        filtered_defaults = defaults
+                    else:
+                        filtered_defaults = {k: v for k, v in defaults.items() if k in sig.parameters}
+
+                    engine.print(f"[cyan]Re-initializing Optimizer {opt_cls.__name__} for {len(trainable_params)} trainable groups...[/]", plugin='LoRA')
+                    
+                    # 创建新优化器
+                    if opt_cls.__name__ == 'SAM':
+                        base_opt_cls = old_opt.base_optimizer.__class__
+                        new_opt = opt_cls(trainable_params, base_optimizer=base_opt_cls, **filtered_defaults)
+                    elif opt_cls.__name__ == 'Muon':
+                        muon_params = [p for p in trainable_params if p.ndim == 2]
+                        adamw_params = [p for p in trainable_params if p.ndim != 2]
+                        new_opt = opt_cls(muon_params=muon_params, adamw_params=adamw_params, **filtered_defaults)
+                    else:
+                        new_opt = opt_cls(trainable_params, **filtered_defaults)
+                    
+                    # 替换 Engine 中的优化器
+                    engine.optimizer = new_opt
+                    
+                    # 处理 Scheduler 的关联
+                    if engine.scheduler:
+                        if hasattr(engine.scheduler, 'optimizer'):
+                            engine.scheduler.optimizer = new_opt
+                            engine.print("[yellow]Note: Scheduler optimizer reference updated.[/]", plugin='LoRA')
+                        else:
+                            engine.print("[yellow]Warning: Scheduler exists but cannot auto-update optimizer reference. Please verify scheduler behavior.[/]", plugin='LoRA')
+
+            self.injected = True
+            
+        # 3. 执行 Checkpoint 的初始化 (创建目录, 尝试 Resume)
+        super().on_init(event)
+
+    def _save(self, engine: 'Engine', filename: str, verbose: bool = True, is_step: bool = False):
+        """重写保存逻辑，仅保存可训练参数 (LoRA + Heads) 和训练状态。"""
+        
+        if self.monitor:
+            engine.meta[self._meta_key] = {'best_k_models': self.best_k_models}
+
+        raw_model = engine.unwrap_model()
+        full_state_dict = raw_model.state_dict()
+        lora_state_dict = {}
+        
+        # 筛选: 保存 requires_grad 的参数以及 LoRA/DoRA 相关的键
+        for name, param in raw_model.named_parameters():
+            if param.requires_grad:
+                lora_state_dict[name] = full_state_dict[name]
+        
+        # 确保 buffers (如 BN running stats) 在解冻层中也被保存
+        # 同时也保存所有名字中带 lora/dora 的 buffer (虽然通常它们没有 buffer)
+        for key, value in full_state_dict.items():
+            if 'lora_' in key or 'dora_' in key:
+                if key not in lora_state_dict:
+                    lora_state_dict[key] = value
+        
+        lora_config = {
+            'r': self.r,
+            'alpha': self.lora_alpha,
+            'target_names': self.target_names,
+            'unlock_head_keywords': self.unlock_head_keywords
+        }
+        
+        file_path = os.path.join(self.path, filename)
+        
+        try:
+            if self.use_safetensors and filename.endswith('.safetensors'):
+                # Safetensors 模式
+                metadata = {
+                    'epoch': str(engine.epoch),
+                    'global_step': str(engine.global_step),
+                    'batch_idx': str(engine.batch_idx),
+                    'is_step': str(is_step),
+                    'orbit_lora_config': json.dumps(lora_config)
+                }
+                try:
+                    metadata['meta'] = json.dumps(engine.meta)
+                except: pass
+                
+                safe_save_file(lora_state_dict, file_path, metadata=metadata)
+                
+            else:
+                # Torch 模式
+                state = {
+                    'epoch': engine.epoch,
+                    'global_step': engine.global_step,
+                    'batch_idx': engine.batch_idx,
+                    'is_step': is_step,
+                    'model_state_dict': lora_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,
+                    'orbit_lora_config': lora_config
+                }
+                torch.save(state, file_path)
+                
+            if verbose:
+                rel_path = os.path.relpath(file_path)
+                file_size = os.path.getsize(file_path) / 1024 / 1024
+                engine.print(f"Saved LoRA checkpoint: {rel_path} ({file_size:.2f} MB)", plugin='LoRA')
+        except Exception as e:
+            engine.print(f"[red]Failed to save checkpoint: {e}[/]", plugin='LoRA')
+
+    def _load(self, engine: 'Engine', file_path: str):
+        """重写加载逻辑，支持 strict=False 加载。"""
+        engine.print(f"[cyan]Loading LoRA checkpoint from: {file_path}[/]", plugin='LoRA')
+        try:
+            raw_model = engine.unwrap_model()
+            model_sd = None
+            saved_config = {}
+            
+            if file_path.endswith('.safetensors'):
+                model_sd = safe_load_file(file_path, device=str(engine.device))
+                
+                with safe_open(file_path, framework="pt", device=str(engine.device)) as f:
+                    metadata = f.metadata()
+                    if metadata and 'orbit_lora_config' in metadata:
+                        try:
+                            saved_config = json.loads(metadata['orbit_lora_config'])
+                        except: pass
+            else:
+                checkpoint = torch.load(file_path, map_location=engine.device)
+                if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+                    model_sd = checkpoint['model_state_dict']
+                    saved_config = checkpoint.get('orbit_lora_config', {})
+                else:
+                    model_sd = checkpoint if not isinstance(checkpoint, dict) else checkpoint
+
+            # 1. 加载模型参数
+            if model_sd is not None:
+                # 配置检查
+                if saved_config:
+                    if saved_config.get('r') != self.r:
+                        engine.print(f"[yellow]Warning: Loaded LoRA rank ({saved_config.get('r')}) != Current rank ({self.r})[/]", plugin='LoRA')
+                
+                missing, unexpected = raw_model.load_state_dict(model_sd, strict=False)
+                
+                # 过滤掉骨干网络的缺失警告，只关注 LoRA 部分
+                relevant_missing = [k for k in missing if 'lora_' in k or 'dora_' in k or any(h in k for h in (self.unlock_head_keywords or []))]
+                if relevant_missing:
+                    engine.print(f"[yellow]Warning: Missing relevant keys: {relevant_missing}[/]", plugin='LoRA')
+                else:
+                    engine.print("[green]LoRA weights loaded successfully.[/]", plugin='LoRA')
+
+            # 2. 恢复训练状态
+            if file_path.endswith('.safetensors'):
+                # Safetensors 仅恢复元数据
+                with safe_open(file_path, framework="pt", device=str(engine.device)) as f:
+                    metadata = f.metadata()
+                    if metadata:
+                        loaded_epoch = int(metadata.get('epoch', 0))
+                        loaded_batch_idx = int(metadata.get('batch_idx', -1))
+                        is_step = metadata.get('is_step', 'False') == 'True'
+                        engine.global_step = int(metadata.get('global_step', 0))
+                        
+                        if 'meta' in metadata:
+                            try:
+                                engine.meta.update(json.loads(metadata['meta']))
+                            except: pass
+                        
+                        if is_step:
+                            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:
+                            engine.start_epoch = loaded_epoch + 1
+                            engine.start_batch_idx = -1
+                            msg = f"Epoch {engine.start_epoch}"
+                            
+                        engine.print(f"[green]Resumed training from {msg}. Note: Optimizer state not restored (safetensors).[/]", plugin='LoRA')
+            else:
+                # Torch 恢复完整状态
+                if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+                    if engine.optimizer and 'optimizer_state_dict' in checkpoint:
+                        try:
+                            engine.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+                        except Exception as e:
+                            engine.print(f"[yellow]Warning: Failed to load optimizer state: {e}.[/]", plugin='LoRA')
+
+                    if engine.scheduler and 'scheduler_state_dict' in checkpoint:
+                        try:
+                            engine.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
+                        except: pass
+                    
+                    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:
+                        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:
+                        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]Resuming training from {msg}[/]", plugin='LoRA')
+                
+        except Exception as e:
+            engine.print(f"[red]Failed to load checkpoint: {e}[/]", plugin='LoRA')
+            import traceback
+            traceback.print_exc()

orbit/plugin/memory_estimator.py

@@ -8,8 +8,7 @@ from typing import TYPE_CHECKING, Optional, Union
 
 from orbit.callback import Callback, Event
 
-if TYPE_CHECKING:
-    from orbit.engine import Engine
+if TYPE_CHECKING: from orbit.engine import Engine
 
 class MemoryEstimator(Callback):
     """

orbit/plugin/warmup.py

@@ -1,8 +1,7 @@
 from typing import Optional, List, TYPE_CHECKING
 from orbit.callback import Callback, Event
 
-if TYPE_CHECKING:
-    from orbit.engine import Engine
+if TYPE_CHECKING: from orbit.engine import Engine
 
 class Warmup(Callback):
     """

orbit/utils/__init__.py

@@ -14,7 +14,7 @@ from .freeze import (
     set_trainable,
     freeze_layers,
     unfreeze_layers,
-    get_trainable_params
+    count_params, ParamStats,
 )
 from .seed import (
     seed_everything,
@@ -27,3 +27,26 @@ from .mask import (
     make_causal_mask,
     make_sliding_window_mask
 )
+from .layer_io import (
+    get_model_by_name,
+    save_layer, load_layer,
+    save_model, load_model
+)
+from .lora import (
+    save_lora, load_lora, inject_lora, inject_lora_file,
+    merge_lora, unmerge_lora,
+    freeze_backbone_only,
+    LoRADiagnoser
+)
+from .sft import (
+    build_sft, train_sft
+)
+from .cuda import (
+    cuda_alloc
+)
+from .image import (
+    split_to_patches, reconstruct_from_patches, pad_to_patch_size
+)
+from .moe import (
+    set_moe_training_mode
+)

orbit/utils/cuda.py

@@ -0,0 +1,10 @@
+import os
+
+def cuda_alloc(size: int = 64):
+    '''
+    设置 PyTorch CUDA 内存分配配置
+
+    Args:
+        size (int): 最大分割大小（MB）
+    '''
+    os.environ['PYTORCH_CUDA_ALLOC_CONF'] = f'max_split_size_mb:{size},expandable_segments:True'

orbit/utils/freeze.py

@@ -1,10 +1,20 @@
 import torch
 import torch.nn as nn
-from typing import Union, List, Optional, Iterable
+from typing import Union, List, Optional, Iterable, Type
+from dataclasses import dataclass
+
+@dataclass
+class ParamStats:
+    count: int
+    params: List[torch.Tensor]
+
+    def __iter__(self):
+        return iter(self.params)
 
 def set_trainable(
     model: nn.Module, 
     targets: Optional[Union[str, List[str]]] = None, 
+    target_classes: Optional[Union[Type[nn.Module], List[Type[nn.Module]]]] = None,
     trainable: bool = False
 ) -> None:
     '''设置模型参数的 requires_grad 属性，用于冻结或解冻层。
@@ -12,48 +22,82 @@ def set_trainable(
     Args:
         model (nn.Module): 目标模型。
         targets (str or List[str], optional): 要操作的层名称或参数名称模式。
-            - 如果为 None，则操作模型的所有参数。
+            - 如果为 None 且 target_classes 也为 None，则操作模型的所有参数。
             - 如果为 str，则操作名称中包含该字符串的所有参数。
             - 如果为 List[str]，则操作名称中包含列表中任意字符串的所有参数。
+        target_classes (Type[nn.Module] or List[Type[nn.Module]], optional): 要操作的模块类。
+            - 如果指定，则操作属于该类（或其子类）的所有模块的参数。
         trainable (bool): 是否可训练 (True 为解冻, False 为冻结)。
     '''
-    if targets is None:
+    if targets is None and target_classes is None:
         for param in model.parameters():
             param.requires_grad = trainable
-    else:
+        return
+
+    if targets is not None:
         if isinstance(targets, str):
             targets = [targets]
         
         for name, param in model.named_parameters():
-            # 检查参数名是否包含 targets 中的任何一个模式
-            if any(t in name for t in targets):
-                param.requires_grad = trainable
+            if any(t in name for t in targets): param.requires_grad = trainable
+
+    if target_classes is not None:
+        if not isinstance(target_classes, (list, tuple)): target_classes = [target_classes]
+        
+        target_classes = tuple(target_classes)
 
-def freeze_layers(model: nn.Module, targets: Optional[Union[str, List[str]]] = None) -> None:
+        for module in model.modules():
+            if isinstance(module, target_classes):
+                for param in module.parameters(): param.requires_grad = trainable
+
+def freeze_layers(
+    model: nn.Module, 
+    targets: Optional[Union[str, List[str]]] = None,
+    target_classes: Optional[Union[Type[nn.Module], List[Type[nn.Module]]]] = None
+) -> None:
     '''冻结模型指定层或所有层 (requires_grad=False)。
 
     Args:
         model (nn.Module): 目标模型。
-        targets (str or List[str], optional): 要冻结的层名称模式。如果不指定，则冻结整个模型。
+        targets (str or List[str], optional): 要冻结的层名称模式。
+        target_classes (Type[nn.Module] or List[Type[nn.Module]], optional): 要冻结的模块类。
+        如果不指定 targets 和 target_classes，则冻结整个模型。
     '''
-    set_trainable(model, targets, trainable=False)
+    set_trainable(model, targets, target_classes, trainable=False)
 
-def unfreeze_layers(model: nn.Module, targets: Optional[Union[str, List[str]]] = None) -> None:
+def unfreeze_layers(
+    model: nn.Module, 
+    targets: Optional[Union[str, List[str]]] = None,
+    target_classes: Optional[Union[Type[nn.Module], List[Type[nn.Module]]]] = None
+) -> None:
     '''解冻模型指定层或所有层 (requires_grad=True)。
 
     Args:
         model (nn.Module): 目标模型。
-        targets (str or List[str], optional): 要解冻的层名称模式。如果不指定，则解冻整个模型。
+        targets (str or List[str], optional): 要解冻的层名称模式。
+        target_classes (Type[nn.Module] or List[Type[nn.Module]], optional): 要解冻的模块类。
+        如果不指定 targets 和 target_classes，则解冻整个模型。
     '''
-    set_trainable(model, targets, trainable=True)
+    set_trainable(model, targets, target_classes, trainable=True)
 
-def get_trainable_params(model: nn.Module) -> Iterable[torch.Tensor]:
-    '''获取模型中 requires_grad=True 的参数，供优化器使用。
+def count_params(model: nn.Module, mode: str = 'trainable') -> ParamStats:
+    '''统计模型参数数量并获取参数列表。
 
     Args:
         model (nn.Module): 目标模型。
+        mode (str): 统计模式，可选 'trainable', 'frozen', 'all'。默认为 'trainable'。
 
     Returns:
-        Iterable[torch.Tensor]: 可训练参数的迭代器。
+        ParamStats: 包含参数总数(count)和参数列表(params)的数据类。
     '''
-    return filter(lambda p: p.requires_grad, model.parameters())
+    if mode == 'trainable':
+        params = [p for p in model.parameters() if p.requires_grad]
+    elif mode == 'frozen':
+        params = [p for p in model.parameters() if not p.requires_grad]
+    elif mode == 'all':
+        params = list(model.parameters())
+    else:
+        raise ValueError(f"Invalid mode: {mode}. Must be one of 'trainable', 'frozen', 'all'.")
+    
+    count = sum(p.numel() for p in params)
+    return ParamStats(count=count, params=params)