ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

reafactor: 统一并增强项目中的工厂模式实现

Browse Source
This commit is contained in:
ViperEkura
2026-03-30 01:33:14 +08:00
parent 60f4df95bd
commit 3e33c14376
9 changed files with 550 additions and 157 deletions
Download Patch File Download Diff File Expand all files Collapse all files
khaosz/trainer/__init__.py
+18 -14
View File
@@ -1,29 +1,33 @@
from khaosz.trainer.trainer import Trainer
from khaosz.trainer.strategy import StrategyFactory
from khaosz.trainer.schedule import SchedulerFactory
from khaosz.trainer.strategy import StrategyFactory, BaseStrategy
from khaosz.trainer.schedule import SchedulerFactory, BaseScheduler
from khaosz.trainer.train_callback import (
TrainCallback,
ProgressBarCallback,
CheckpointCallback,
TrainCallback,
GradientClippingCallback,
SchedulerCallback,
MetricLoggerCallback
CheckpointCallback,
ProgressBarCallback,
MetricLoggerCallback,
)
__all__ = [
# trainer
# Main trainer
"Trainer",
# factory
# Strategy factory
"StrategyFactory",
"SchedulerFactory",
"BaseStrategy",
# callback
"TrainCallback",
"ProgressBarCallback",
"CheckpointCallback",
# Scheduler factory
"SchedulerFactory",
"BaseScheduler",
# Callbacks
"TrainCallback",
"GradientClippingCallback",
"SchedulerCallback",
"MetricLoggerCallback"
"CheckpointCallback",
"ProgressBarCallback",
"MetricLoggerCallback",
]
khaosz/trainer/schedule.py
+96 -31
View File
@@ -1,20 +1,21 @@
"""Learning rate scheduler implementations with factory pattern."""
import math
from abc import abstractmethod, ABC
from typing import Any, Dict, List
from typing import Any, Dict, List, Type
from torch.optim.lr_scheduler import LRScheduler
from khaosz.config.schedule_config import ScheduleConfig
class BaseScheduler(LRScheduler, ABC):
"""
Base scheduler class for all other schedulers.
"""
"""Base scheduler class for all other schedulers."""
def __init__(self, optimizer, last_epoch: int = -1):
super().__init__(optimizer, last_epoch)
@abstractmethod
def get_lr(self) -> List[float]:
"""Calculate the current learning rate."""
raise NotImplementedError
def state_dict(self) -> Dict[str, Any]:
@@ -24,10 +25,95 @@ class BaseScheduler(LRScheduler, ABC):
super().load_state_dict(state_dict)
class SchedulerFactory:
"""Factory class for creating learning rate schedulers.
Supports decorator-based registration for extensible scheduler types.
Also supports creation from ScheduleConfig objects.
Example usage:
@SchedulerFactory.register("custom")
class CustomScheduler(BaseScheduler):
...
scheduler = SchedulerFactory.create(optimizer, "custom", **kwargs)
# Or from config
config = CosineScheduleConfig(total_steps=10000)
scheduler = SchedulerFactory.load(optimizer, config)
"""
SCHEDULER_MAP: Dict[str, Type[BaseScheduler]] = {}
@classmethod
def register(cls, name: str):
"""Decorator to register a new scheduler class.
Args:
name: Registration name for the scheduler
Returns:
Decorator function that registers the scheduler class
"""
def decorator(scheduler_cls: Type[BaseScheduler]) -> Type[BaseScheduler]:
if not issubclass(scheduler_cls, BaseScheduler):
raise TypeError(f"{scheduler_cls.__name__} must inherit from BaseScheduler")
cls.SCHEDULER_MAP[name] = scheduler_cls
return scheduler_cls
return decorator
@classmethod
def create(cls, optimizer, schedule_type: str, **kwargs) -> BaseScheduler:
"""Create a scheduler instance by type name.
Args:
optimizer: PyTorch optimizer
schedule_type: Type of scheduler ("cosine", "sgdr")
**kwargs: Arguments passed to the scheduler constructor
Returns:
Scheduler instance
Raises:
ValueError: If schedule_type is not supported
"""
if schedule_type not in cls.SCHEDULER_MAP:
raise ValueError(
f"Unknown schedule type: '{schedule_type}'. "
f"Supported types: {sorted(cls.SCHEDULER_MAP.keys())}"
)
scheduler_cls = cls.SCHEDULER_MAP[schedule_type]
return scheduler_cls(optimizer, **kwargs)
@staticmethod
def load(optimizer, schedule_config: ScheduleConfig) -> BaseScheduler:
"""Create a scheduler from a ScheduleConfig object.
Args:
optimizer: PyTorch optimizer
schedule_config: ScheduleConfig instance
Returns:
Scheduler instance
"""
kwargs = schedule_config.get_kwargs()
schedule_type = kwargs.pop("schedule_type")
return SchedulerFactory.create(optimizer, schedule_type, **kwargs)
@classmethod
def available_types(cls) -> list:
"""Return list of registered scheduler type names."""
return list(cls.SCHEDULER_MAP.keys())
# ============== Scheduler Classes ==============
# All scheduler classes are registered at class definition time using the decorator
@SchedulerFactory.register("cosine")
class CosineScheduler(BaseScheduler):
"""
Cosine decay scheduler with warmup, implemented as PyTorch LRScheduler.
"""
"""Cosine decay scheduler with warmup, implemented as PyTorch LRScheduler."""
def __init__(
self,
@@ -75,10 +161,9 @@ class CosineScheduler(BaseScheduler):
super().load_state_dict(state_dict)
@SchedulerFactory.register("sgdr")
class SGDRScheduler(BaseScheduler):
"""
SGDR (Stochastic Gradient Descent with Warm Restarts) scheduler,
"""
"""SGDR (Stochastic Gradient Descent with Warm Restarts) scheduler."""
def __init__(
self,
@@ -141,24 +226,4 @@ class SGDRScheduler(BaseScheduler):
self.cycle_length = state_dict.pop('cycle_length')
self.min_rate = state_dict.pop('min_rate')
self.t_mult = state_dict.pop('t_mult')
super().load_state_dict(state_dict)
class SchedulerFactory:
"""
Factory class for creating learning rate schedulers.
"""
@staticmethod
def load(optimizer, schedule_config: ScheduleConfig) -> BaseScheduler:
kwargs = schedule_config.get_kwargs()
schedule_type = kwargs.pop("schedule_type")
if schedule_type == "cosine":
return CosineScheduler(optimizer, **kwargs)
elif schedule_type == "sgdr":
return SGDRScheduler(optimizer, **kwargs)
else:
raise ValueError(f"Unsupported schedule type: {schedule_type}")
super().load_state_dict(state_dict)
khaosz/trainer/strategy.py
+137 -58
View File
@@ -1,3 +1,5 @@
"""Training strategy implementations with factory pattern."""
import copy
import torch
import torch.nn as nn
@@ -17,9 +19,10 @@ def unwrap_model(model: nn.Module) -> nn.Module:
def create_ref_model(model: nn.Module) -> nn.Module:
"""
Create a reference model for DPO/GRPO training.
Handles DDP-wrapped models safely.
"""Create a reference model for DPO/GRPO training.
Handles DDP-wrapped models safely by unwrapping first,
then creating a deep copy with frozen gradients.
"""
original_model = unwrap_model(model)
ref_model = copy.deepcopy(original_model)
@@ -28,17 +31,18 @@ def create_ref_model(model: nn.Module) -> nn.Module:
return ref_model
def move_to_device(batch:Dict[str, Tensor], device: str) -> Any:
def move_to_device(batch: Dict[str, Tensor], device: str) -> Any:
"""Move batch tensors to specified device with non-blocking transfer."""
return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
def get_logprobs(
model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
input_ids: Tensor,
mask: Tensor,
reduction: str,
):
"""
Compute token-wise log probabilities from model outputs.
"""Compute token-wise log probabilities from model outputs.
Args:
model: The language model
@@ -49,7 +53,6 @@ def get_logprobs(
Returns:
Log probabilities with reduction applied over sequence dimension
"""
# reduction on seq_len dim
allowed_reductions = ["mean", "sum", "none"]
if reduction not in allowed_reductions:
raise ValueError(f"reduction must be one of {allowed_reductions}, got '{reduction}'")
@@ -60,7 +63,6 @@ def get_logprobs(
logits = model(input_ids[:, :-1], mask[:, :-1])["logits"]
log_probs = torch.log_softmax(logits.float(), dim=-1)
# [batch_size, seq_len - 1]
token_logprobs = torch.gather(
log_probs,
dim=-1,
@@ -76,20 +78,112 @@ def get_logprobs(
class BaseStrategy(ABC):
"""Abstract base class for training strategies."""
def __init__(self, model: Union[nn.Module, Callable[..., Dict[str, Tensor]]], device: str):
self.model = model
self.device = device
@abstractmethod
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
"""Compute loss for the given batch.
Args:
batch: Dictionary containing batch tensors
Returns:
Computed loss tensor
"""
raise NotImplementedError
def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
"""Allow calling strategy directly as a callable."""
return self.compute_loss(batch)
class StrategyFactory:
"""Factory class for creating training strategy instances.
Supports decorator-based registration for extensible strategy types.
All default strategies (seq, sft, dpo, grpo) are automatically registered.
Example usage:
@StrategyFactory.register("custom")
class CustomStrategy(BaseStrategy):
...
strategy = StrategyFactory.create(model, "custom", device)
"""
SUPPORTED_STRATEGIES = frozenset({"seq", "sft", "dpo", "grpo"})
STRATEGY_MAP: Dict[str, type] = {}
@classmethod
def register(cls, name: str):
"""Decorator to register a new strategy class.
Args:
name: Registration name for the strategy
Returns:
Decorator function that registers the strategy class
"""
def decorator(strategy_cls: type) -> type:
if not issubclass(strategy_cls, BaseStrategy):
raise TypeError(f"{strategy_cls.__name__} must inherit from BaseStrategy")
cls.STRATEGY_MAP[name] = strategy_cls
return strategy_cls
return decorator
@classmethod
def create(cls, model, train_type: str, device: str, **kwargs) -> BaseStrategy:
"""Create a strategy instance based on training type.
Args:
model: Model instance for the strategy
train_type: Type of training ("seq", "sft", "dpo", "grpo")
device: Device to run the strategy on
**kwargs: Additional arguments passed to strategy constructor
Returns:
Strategy instance
Raises:
ValueError: If train_type is not supported
NotImplementedError: If train_type is in supported list but not implemented
"""
if train_type not in cls.SUPPORTED_STRATEGIES:
raise ValueError(
f"Unknown training strategy: '{train_type}'. "
f"Supported strategies: {sorted(cls.SUPPORTED_STRATEGIES)}"
)
if train_type not in cls.STRATEGY_MAP:
raise NotImplementedError(
f"Strategy '{train_type}' is supported but not yet implemented."
)
strategy_cls = cls.STRATEGY_MAP[train_type]
return strategy_cls(model, device, **kwargs)
@classmethod
def available_strategies(cls) -> list:
"""Return list of registered strategy names."""
return list(cls.STRATEGY_MAP.keys())
# ============== Strategy Classes ==============
# All strategies are registered at class definition time using the decorator
@StrategyFactory.register("seq")
class SEQStrategy(BaseStrategy):
def __init__(self, model, device, label_smoothing):
"""Standard next-token prediction training strategy.
Computes cross-entropy loss for next token prediction.
"""
def __init__(self, model, device, label_smoothing: float = 0.0):
super().__init__(model, device)
self.label_smoothing = label_smoothing
@@ -99,15 +193,22 @@ class SEQStrategy(BaseStrategy):
logits = self.model(input_ids=input_ids)["logits"]
loss = F.cross_entropy(
input=logits.flatten(0, 1).float(),
target=target_ids.flatten()
input=logits.flatten(0, 1).float(),
target=target_ids.flatten(),
label_smoothing=self.label_smoothing
)
return loss
@StrategyFactory.register("sft")
class SFTStrategy(BaseStrategy):
def __init__(self, model, device, label_smoothing):
"""Supervised Fine-tuning strategy with loss masking.
Applies cross-entropy loss only to tokens where loss_mask is True.
"""
def __init__(self, model, device, label_smoothing: float = 0.0):
super().__init__(model, device)
self.label_smoothing = label_smoothing
@@ -122,19 +223,27 @@ class SFTStrategy(BaseStrategy):
loss = F.cross_entropy(
input=logits.flatten(0, 1).float(),
target=target_ids.flatten(),
ignore_index=ignore_index
ignore_index=ignore_index,
label_smoothing=self.label_smoothing
)
return loss
@StrategyFactory.register("dpo")
class DPOStrategy(BaseStrategy):
"""Direct Preference Optimization strategy.
Implements the DPO loss from the paper "Direct Preference Optimization".
Uses a reference model to compute KL divergence penalty.
"""
def __init__(
self,
model: nn.Module,
device: str,
beta: float,
reduction: str,
beta: float = 0.1,
reduction: str = "mean",
):
super().__init__(model, device)
self.ref_model = create_ref_model(model)
@@ -168,16 +277,21 @@ class DPOStrategy(BaseStrategy):
return dpo_loss
@StrategyFactory.register("grpo")
class GRPOStrategy(BaseStrategy):
"""Group Relative Policy Optimization strategy.
Implements GRPO with clipping and KL penalty.
"""
def __init__(
self,
model: nn.Module,
device: str,
clip_eps: float,
kl_coef: float,
group_size: int,
reduction: str,
clip_eps: float = 0.2,
kl_coef: float = 0.01,
group_size: int = 4,
reduction: str = "mean",
):
super().__init__(model, device)
self.ref_model = create_ref_model(model)
@@ -209,16 +323,14 @@ class GRPOStrategy(BaseStrategy):
log_probs_ref = get_logprobs(self.ref_model, full_sequences, full_masks, self.reduction)
log_probs_ref = log_probs_ref.view(batch_size, group_size)
# Compute advantages from rewards
# Compute advantages from rewards with normalization
eps = torch.finfo(log_probs_policy.dtype).eps
mean = rewards.mean(dim=-1, keepdim=True)
std = rewards.std(dim=-1, keepdim=True)
advantages = (rewards - mean) / (std + eps)
# log_ratio = log_probs_policy - log_probs_old
# ratio = torch.exp(log_ratio)
# off policy: policy_model = old_model, then ratio = 1
ratio = torch.exp(0)
# PPO-style clipped surrogate objective
ratio = torch.exp(0) # Off-policy: policy_model = old_model
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
@@ -227,36 +339,3 @@ class GRPOStrategy(BaseStrategy):
total_loss = policy_loss + kl_penalty
return total_loss
class StrategyFactory:
def load(model, train_type, device, **kwargs):
train_strategy: Dict[str, Callable[[], BaseStrategy]] = {
"seq": lambda: SEQStrategy(
model,
device,
kwargs.get("label_smoothing", 0.0)
),
"sft": lambda: SFTStrategy(
model,
device,
kwargs.get("label_smoothing", 0.0)
),
"dpo": lambda: DPOStrategy(
model,
device,
kwargs.get("dpo_beta"),
kwargs.get("reduction", "mean")
),
"grpo": lambda: GRPOStrategy(
model,
device,
kwargs.get("grpo_clip_eps"),
kwargs.get("grpo_kl_coef"),
kwargs.get("grpo_group_size"),
kwargs.get("reduction", "mean")
)
}
strategy = train_strategy[train_type]()
return strategy