fix(khaosz/trainer/data_util): 修复数据集索引范围错误

khaosz/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "1.2.2"
+__version__ = "1.3.0"
 __author__ = "ViperEkura"
 
 from khaosz.model import Khaosz

khaosz/core/parameter.py

@@ -106,7 +106,7 @@ class Checkpoint(BaseModelIO):
         default_factory=TransformerConfig,
         metadata={"help": "Transformer model configuration."}
     )
-    optim_state: Dict[str, Any] = field(
+    optimizer_state: Dict[str, Any] = field(
         default=None,
         metadata={"help": "Optimizer state."}
     )
@@ -124,7 +124,7 @@ class Checkpoint(BaseModelIO):
         paths.update({
             "loss_list": paths["model"].parent / "loss.pkl",
             "loss_plot": paths["model"].parent / "loss.png",
-            "optim_state": paths["model"].parent / "optim_state.pkl",
+            "optimizer_state": paths["model"].parent / "optimizer_state.pkl",
             "sampler_state": paths["model"].parent / "sampler_state.pkl"
         })
         return paths
@@ -140,8 +140,8 @@ class Checkpoint(BaseModelIO):
             pkl.dump(self.loss_list, f)
         
         # Save optimizer state
-        with open(str(paths["optim_state"]), "wb") as f:
+        with open(str(paths["optimizer_state"]), "wb") as f:
-            pkl.dump(self.optim_state, f)
+            pkl.dump(self.optimizer_state, f)
             
         # Save sampler state
         with open(str(paths["sampler_state"]), "wb") as f:
@@ -156,9 +156,9 @@ class Checkpoint(BaseModelIO):
                 self.loss_list = pkl.load(f)
         
         # Load optimizer state
-        if paths["optim_state"].exists():
+        if paths["optimizer_state"].exists():
-            with open(str(paths["optim_state"]), "rb") as f:
+            with open(str(paths["optimizer_state"]), "rb") as f:
-                self.optim_state = pkl.load(f)
+                self.optimizer_state = pkl.load(f)
         
         # Load sampler state
         if paths["sampler_state"].exists():

khaosz/trainer/__init__.py

@@ -1,22 +1,22 @@
 from khaosz.trainer.data_util import DatasetLoader
 from khaosz.trainer.trainer import Trainer
+from khaosz.trainer.train_config import TrainConfig
 from khaosz.trainer.strategy import (
-    TrainConfig, 
     CosineScheduleConfig, 
     SgdrScheduleConfig,
     StrategyFactory,
     SchedulerFactory
 )
-from khaosz.trainer.trainer_callback import (
+from khaosz.trainer.train_callback import (
-    TrainerCallback,
+    TrainCallback,
     ProgressBarCallback,
     CheckpointCallback,
-    TrainerCallback,
+    TrainCallback,
-    SchedulerCallback
+    SchedulerCallback,
+    StepMonitorCallback
 )
 
 __all__ = [
-    # strategy
     "DatasetLoader",
     "Trainer",
     "TrainConfig",
@@ -26,9 +26,10 @@ __all__ = [
     "SchedulerFactory",
     
     # callback
-    "TrainerCallback",
+    "TrainCallback",
     "ProgressBarCallback",
     "CheckpointCallback",
-    "TrainerCallback",
+    "TrainCallback",
     "SchedulerCallback",
+    "StepMonitorCallback"
 ]

khaosz/trainer/data_util.py

@@ -110,12 +110,11 @@ class MutiSegmentFetcher:
 
 
 class BaseDataset(Dataset, ABC):
-    def __init__(self, chunk_size: int, device: str):
+    def __init__(self, chunk_size: int):
         super().__init__()
         self.segments: MutiSeg = {}
         self.chunk_size = chunk_size
         self.total_samples = 0
-        self.device = device
 
     def save(self, save_path: str):
         keys = list(self.segments.keys())
@@ -148,20 +147,20 @@ class SeqDataset(BaseDataset):
     def __init__(
         self, 
         chunk_size, 
-        device='cuda'
     ):
-        super().__init__(chunk_size, device)
+        super().__init__(chunk_size)
         self.fetcher = MutiSegmentFetcher(self.segments)
 
     def _fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
         return self.fetcher.key_fetch(begin_idx, end_idx, "sequence")
     
     def __getitem__(self, index):
-        begin_idx = index * self.chunk_size 
+        # fix the range index bug
-        end_idx = min(begin_idx + self.chunk_size, self.total_samples - 1)
+        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
+        end_idx = begin_idx + self.chunk_size
         
-        x = self._fetch_data(begin_idx, end_idx).to(device=self.device, dtype=torch.long)
+        x = self._fetch_data(begin_idx, end_idx).to(dtype=torch.long)
-        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(device=self.device, dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
         
         return {"input_ids": x, "target_ids": y}
     
@@ -175,9 +174,8 @@ class SftDataset(BaseDataset):
         eos_token_id, 
         user_token_id, 
         multi_turn=False, 
-        device='cuda'
     ):
-        super().__init__(chunk_size, device)
+        super().__init__(chunk_size)
         self.fetcher = MutiSegmentFetcher(self.segments)
         self.bos_token_id = bos_token_id
         self.eos_token_id = eos_token_id
@@ -188,11 +186,11 @@ class SftDataset(BaseDataset):
         return self.fetcher.key_fetch(begin_idx, end_idx, key)
     
     def __getitem__(self, index):
-        begin_idx = index * self.chunk_size 
+        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
-        end_idx = min(begin_idx + self.chunk_size, self.total_samples - 1)
+        end_idx = begin_idx + self.chunk_size
         
-        x = self._fetch_data(begin_idx, end_idx, "sequence").to(device=self.device, dtype=torch.long)
+        x = self._fetch_data(begin_idx, end_idx, "sequence").to(dtype=torch.long)
-        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(device=self.device, dtype=torch.long)
+        y = self._fetch_data(begin_idx + 1, end_idx + 1, "sequence").to(dtype=torch.long)
         
         # fix the eos_token_id bug(change target_ids to input_ids)
         loss_mask = build_loss_mask(x, self.bos_token_id, self.eos_token_id)
@@ -202,43 +200,41 @@ class SftDataset(BaseDataset):
 
 
 class DpoDataset(BaseDataset):
-    def __init__(self, chunk_size: int, device="cuda"):
+    def __init__(self, chunk_size: int):
-        super().__init__(chunk_size, device)
+        super().__init__(chunk_size)
         self.fetcher = MutiSegmentFetcher(self.segments)
 
     def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
         return self.fetcher.key_fetch(begin_idx, end_idx, key)
 
     def __getitem__(self, index: int):
-        start_idx = index * self.chunk_size
+        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
-        end_idx = min(start_idx + self.chunk_size, self.total_samples - 1)
+        end_idx = begin_idx + self.chunk_size
         
-        chosen = self._fetch_data(start_idx, end_idx, "chosen").to(device=self.device, dtype=torch.long)
+        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
-        rejected = self._fetch_data(start_idx, end_idx, "rejected").to(device=self.device, dtype=torch.long)
+        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
-        chosen_mask = self._fetch_data(start_idx, end_idx, "chosen_mask").to(device=self.device, dtype=torch.bool)
+        chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(dtype=torch.bool)
-        rejected_mask = self._fetch_data(start_idx, end_idx, "rejected_mask").to(device=self.device, dtype=torch.bool)
+        rejected_mask = self._fetch_data(begin_idx, end_idx, "rejected_mask").to(dtype=torch.bool)
 
         return {"chosen": chosen, "rejected": rejected, "chosen_mask": chosen_mask, "rejected_mask": rejected_mask}
 
 
 class PpoDataset(BaseDataset):
-    def __init__(self, chunk_size: int, device="cuda"):
+    def __init__(self, chunk_size: int):
-        super().__init__(chunk_size, device)
+        super().__init__(chunk_size)
         self.fetcher = MutiSegmentFetcher(self.segments)
 
     def _fetch_data(self, begin_idx: int, end_idx: int, key: str) -> Tensor:
         return self.fetcher.key_fetch(begin_idx, end_idx, key)
     
     def __getitem__(self, index: int) -> Dict[str, Tensor]:
+        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
+        end_idx = begin_idx + self.chunk_size
 
-        begin_idx = index * self.chunk_size
+        input_ids =  self._fetch_data(begin_idx, end_idx, "input_ids"),
-        end_idx = min(begin_idx + self.chunk_size, self.total_samples - 1)
+        actions = self._fetch_data(begin_idx, end_idx, "actions"),
-        
+        logprobs = self._fetch_data(begin_idx, end_idx, "logprobs"),
-
+        rewards =  self._fetch_data(begin_idx, end_idx, "rewards")
-        input_ids =  self._fetch_data(begin_idx, end_idx, "input_ids").to(self.device),
-        actions = self._fetch_data(begin_idx, end_idx, "actions").to(self.device),
-        logprobs = self._fetch_data(begin_idx, end_idx, "logprobs").to(self.device),
-        rewards =  self._fetch_data(begin_idx, end_idx, "rewards").to(self.device)
         
         return {"input_ids": input_ids, "actions": actions, "logprobs": logprobs, "rewards": rewards}
     
@@ -249,23 +245,21 @@ class DatasetLoader:
         train_type: Literal["seq", "sft", "dpo"],
         load_path: Union[str, List[str]],
         max_len: int, 
-        device: str,
         **kwargs
         ) -> BaseDataset:
         
-        dataset_router: Dict[str, Callable[[int, torch.device], BaseDataset]] = {
+        dataset_router: Dict[str, Callable[[int], BaseDataset]] = {
-            "seq": lambda m_len, device: SeqDataset(m_len, device=device),
+            "seq": lambda max_len: SeqDataset(max_len),
-            "sft": lambda m_len, device: SftDataset(
+            "sft": lambda max_len: SftDataset(
-                m_len, 
+                max_len, 
-                device=device,
                 bos_token_id=kwargs.get("bos_token_id"),
                 eos_token_id=kwargs.get("eos_token_id"),
                 user_token_id=kwargs.get("user_token_id"),
                 multi_turn=kwargs.get("multi_turn")
             ),
-            "dpo": lambda m_len, device: DpoDataset(m_len, device=device),
+            "dpo": lambda max_len: DpoDataset(max_len),
         }
-        dataset = dataset_router[train_type](max_len, device)
+        dataset = dataset_router[train_type](max_len)
         dataset.load(load_path)
         
         return dataset
@@ -297,8 +291,8 @@ class RandomSampler(Sampler[int]):
         
         start = self.current_iter % n
         for i in range(start, n):
-            yield self._indices[i]
             self.current_iter += 1
+            yield self._indices[i]
         
         self.epoch += 1
         self._indices = None

khaosz/trainer/metric_util.py

@@ -0,0 +1,65 @@
+import torch.nn as nn
+from typing import Dict
+
+def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
+    """ Compute gradient norm for each parameter in the model. """
+    norms = {}
+    for name, param in model.named_parameters():
+        norms[name] = 0.0
+        if param.grad:
+            norm = param.grad.data.norm(norm_type).item()
+            norms[name] = norm
+    return norms
+
+def grad_std(model: nn.Module) -> Dict[str, float]:
+    """ Compute standard deviation of gradients for each parameter. """
+    stds = {}
+    for name, param in model.named_parameters():
+        stds[name] = 0.0
+        if param.grad:
+            std = param.grad.data.std().item()
+            stds[name] = std
+    return stds
+
+def grad_max(model: nn.Module) -> Dict[str, float]:
+    """ Find the maximum absolute gradient value for each parameter. """
+    max_vals = {}
+    for name, param in model.named_parameters():
+        max_vals[name] = -float('inf')
+        if param.grad:
+            max_val = param.grad.data.max().item()
+            max_vals[name] = max_val
+            
+    return max_vals
+
+def grad_min(model: nn.Module) -> Dict[str, float]:
+    """ Find the minimum absolute gradient value for each parameter. """
+    min_vals = {}
+    for name, param in model.named_parameters():
+        min_vals[name] = float('inf')
+        if param.grad:
+            min_val = param.grad.data.min().item()
+            min_vals[name] = min_val
+            
+    return min_vals
+
+def grad_mean(model: nn.Module) -> Dict[str, float]:
+    """ Compute mean of gradients for each parameter. """
+    means = {}
+    for name, param in model.named_parameters():
+        means[name] = 0.0
+        if param.grad:
+            mean = param.grad.data.mean().item()
+            means[name] = mean
+            
+    return means
+
+def grad_nan_num(model: nn.Module) -> Dict[str, int]:
+    """ Count the number of NaNs in gradients for each parameter. """
+    nan_nums = {}
+    for name, param in model.named_parameters():
+        nan_nums[name] = 0
+        if param.grad:
+            nan_num = param.grad.isnan().sum().item()
+            nan_nums[name] = nan_num
+    return nan_nums

khaosz/trainer/strategy.py

@@ -5,11 +5,9 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 from torch import Tensor
-from torch.optim import Optimizer
+from typing import Any, Literal, Tuple, Callable, Dict, Union
-from torch.utils.data import Dataset
-from typing import Any, Literal, Tuple, Callable, Dict
 from abc import ABC, abstractmethod
-from dataclasses import asdict, dataclass, field
+from dataclasses import dataclass, field
 
 
 def get_logprobs(model:nn.Module, input_ids: Tensor, mask: Tensor, pad_token_id: int):
@@ -32,10 +30,14 @@ def get_logprobs(model:nn.Module, input_ids: Tensor, mask: Tensor, pad_token_id:
     
     return (token_logprobs * valid_mask).sum(dim=-1)
 
+def move_to_device(batch:Dict[str, Tensor], device: str) -> Any:
+    return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
+
 
 class BaseStrategy(ABC):
-    def __init__(self, model: nn.Module):
+    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:
@@ -46,41 +48,37 @@ class BaseStrategy(ABC):
 
 
 class SeqStrategy(BaseStrategy):
-    def __init__(self, model):
+    def __init__(self, model, device):
-        super().__init__(model)
+        super().__init__(model, device)
     
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        batch = move_to_device(batch, self.device)
         input_ids, target_ids = batch["input_ids"], batch["target_ids"]
-        B, L = input_ids.size()
+        logits = self.model(input_ids=input_ids)["logits"]
-        logits: Tensor = self.model(input_ids=input_ids)["logits"]
         
         loss = F.cross_entropy(
-            input=logits.view(B * L, -1),
+            input=logits.flatten(0, 1),
             target=target_ids.flatten()
         )
+        
         return loss
     
 
 class SftStrategy(BaseStrategy):
-    def __init__(self, model: nn.Module):
+    def __init__(self, model, device):
-        super().__init__(model)
+        super().__init__(model, device)
     
     def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
+        batch = move_to_device(batch, self.device)
         input_ids, target_ids = batch["input_ids"], batch["target_ids"]
         loss_mask, attn_mask = batch["loss_mask"], batch["attn_mask"]
         
         ignore_index = -100
-        B, L = input_ids.size()
+        logits = self.model(input_ids=input_ids, input_mask=attn_mask)["logits"]
-        
-        logits: Tensor = self.model(
-            input_ids=input_ids, 
-            input_mask=attn_mask
-        )["logits"]
-        
         target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)
         
         loss = F.cross_entropy(
-            input=logits.view(B * L, -1),
+            input=logits.flatten(0, 1),
             target=target_ids.flatten(),
             ignore_index=ignore_index
         )
@@ -89,8 +87,8 @@ class SftStrategy(BaseStrategy):
 
 
 class DpoStrategy(BaseStrategy):
-    def __init__(self, model, pad_token_id, beta):
+    def __init__(self, model, device, pad_token_id, beta):
-        super().__init__(model)
+        super().__init__(model, device)
         ref_model = copy.deepcopy(self.model)
         ref_model.requires_grad_(False)
         ref_model.eval()
@@ -100,6 +98,7 @@ class DpoStrategy(BaseStrategy):
         self.beta = beta
         
     def compute_loss(self, batch: Tuple[Tensor, ...]) -> Tensor:
+        batch = move_to_device(batch, self.device)
         good_ids, bad_ids = batch["chosen"], batch["rejected"]
         good_mask, bad_mask = batch["chosen_mask"], batch["rejected_mask"]
         
@@ -156,68 +155,20 @@ class PpoStrategy(BaseStrategy):
 
 class StrategyFactory:
     
-    def load(model, train_type, **kwargs):
+    def load(model, train_type, device, **kwargs):
         train_strategy: Dict[str, Callable[[], BaseStrategy]] = {
-            "seq": lambda: SeqStrategy(model),
+            "seq": lambda: SeqStrategy(model, device),
-            "sft": lambda: SftStrategy(model),
+            "sft": lambda: SftStrategy(model, device),
             "dpo": lambda: DpoStrategy(
                 model, 
+                device,
                 kwargs.get("pad_token_id"), 
                 kwargs.get("dpo_beta")
             )
         }
         strategy = train_strategy[train_type]()
         return strategy
-    
 
-@dataclass
-class TrainConfig:
-    
-    strategy: BaseStrategy = field(
-        default=None,
-        metadata={"help": "Training strategy."}
-    )
-    dataset: Dataset = field(
-        default=None,
-        metadata={"help": "Dataset for training."}
-    )
-    optimizer: Optimizer = field(
-        default=None,
-        metadata={"help": "Optimizer for training."}
-    )
-    checkpoint_dir: str = field(
-        default="./checkpoint",
-        metadata={"help": "Checkpoint directory."}
-    )
-    n_epoch: int = field(
-        default=1,
-        metadata={"help": "Number of epochs for training."}
-    )
-    batch_size: int = field(
-        default=4,
-        metadata={"help": "Batch size for training."}
-    )
-    checkpoint_interval: int = field(
-        default=5000,
-        metadata={"help": "Number of iterations between checkpoints."}
-    )
-    accumulation_steps: int = field(
-        default=1,
-        metadata={"help": "Number of iterations between steps."}
-    )
-    max_grad_norm: float = field(
-        default=1.0,
-        metadata={"help": "Maximum gradient norm."}
-    )
-    random_seed: int = field(
-        default=3407,
-        metadata={"help": "Random seed."}
-    )
-
-    def get_kwargs(self)-> Dict[str, Any]:
-        config_dict = asdict(self)
-        return {k: v for k, v in config_dict.items() if v is not None}
-    
 
 @dataclass
 class ScheduleConfig(ABC):

khaosz/trainer/train_callback.py

@@ -0,0 +1,230 @@
+import os
+import json
+import time
+
+from pathlib import Path
+from tqdm import tqdm
+from torch.nn.utils import clip_grad_norm_
+from torch.optim.lr_scheduler import LambdaLR
+from typing import List, Optional, Protocol, TYPE_CHECKING
+
+from khaosz.trainer.strategy import ScheduleConfig, SchedulerFactory
+from khaosz.trainer.metric_util import (
+    grad_max,
+    grad_min,
+    grad_norm,
+    grad_mean,
+    grad_std,
+    grad_nan_num
+)
+
+if TYPE_CHECKING:
+    from khaosz.trainer.trainer import Trainer
+    from khaosz.trainer.train_context import TrainContext
+
+
+class TrainCallback(Protocol):
+    """ 
+    Callback interface for trainer.
+    """
+
+    def on_train_begin(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Called at the beginning of training. """
+    
+    def on_train_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Called at the end of training. """
+
+    def on_epoch_begin(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Called at the beginning of each epoch. """
+
+    def on_epoch_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        """  Called at the end of each epoch. """
+    
+    def on_step_begin(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Called at the beginning of each step. """
+
+    def on_step_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Called at the end of each step."""
+
+    def on_batch_begin(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Called at the beginning of each batch. """
+
+    def on_batch_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Called at the end of each batch. """
+    
+    def on_error(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Called when an error occurs during training. """
+
+
+class GradientClippingCallback(TrainCallback):
+    """ 
+    Gradient clipping callback for trainer.
+    """
+    def on_step_begin(self, trainer: 'Trainer', context: 'TrainContext'):
+        _ = context
+        clip_grad_norm_(trainer.parameter.model.parameters(), trainer.train_config.max_grad_norm)
+
+
+class SchedulerCallback(TrainCallback):
+    """
+    Scheduler callback for trainer.
+    """
+    def __init__(self, schedule_config: ScheduleConfig):
+        self.schedule_config = schedule_config
+        self.scheduler: Optional[LambdaLR] = None
+    
+    def on_train_begin(self, trainer: 'Trainer', context: 'TrainContext'):
+
+        for group in trainer.train_config.optimizer.param_groups:
+            if "initial_lr" not in group:
+                group["initial_lr"] = group["lr"] 
+
+        self.schedule_config.validate()
+        lambda_scheduler_fn = SchedulerFactory.load_schedule_fn(
+            self.schedule_config
+        )
+        
+        self.scheduler = LambdaLR(
+            trainer.train_config.optimizer,
+            lambda_scheduler_fn,
+            last_epoch=context.current_iter - 1
+        )
+    
+    def on_batch_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        _ = trainer, context
+        if self.scheduler:
+            self.scheduler.step()
+
+
+class CheckpointCallback(TrainCallback):
+    """ 
+    Checkpoint callback for trainer.
+    """
+    def __init__(self, checkpoint_interval: int):
+        self.checkpoint_interval = checkpoint_interval
+        self.last_ckpt_iter = 0
+    
+    def _save_checkpoint(self, trainer: 'Trainer', context: 'TrainContext'):
+        save_path = os.path.join(trainer.train_config.checkpoint_dir, f"iter_{context.current_iter}")
+        context.checkpoint.sampler_state = context.sampler.state_dict()
+        context.checkpoint.optimizer_state = context.optimizer.state_dict()
+        context.checkpoint.save(save_path)
+        self.last_ckpt_iter = context.current_iter
+    
+    def on_batch_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        context.checkpoint.loss_list.append(context.loss)
+        
+        if context.current_iter - self.last_ckpt_iter >= self.checkpoint_interval:
+            self._save_checkpoint(trainer, context)
+            
+    def on_train_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        if context.current_iter != self.last_ckpt_iter:
+            self._save_checkpoint(trainer, context)
+
+
+class ProgressBarCallback(TrainCallback):
+    """ 
+    Progress bar callback for trainer.
+    """
+    def __init__(self):
+        self.progress_bar: tqdm = None
+    
+    def on_epoch_begin(self, trainer: 'Trainer', context: 'TrainContext'):
+        self.progress_bar = tqdm(
+            context.dataloader, 
+            desc=f"Epoch {context.epoch+1}/{trainer.train_config.n_epoch}", 
+            dynamic_ncols=True
+        )
+    
+    def on_batch_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        _ = trainer
+        self.progress_bar.set_postfix({
+            "loss": f"{context.loss:.4f}",
+            "lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}"
+        })
+        self.progress_bar.update(1)
+    
+    def on_epoch_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        _ = trainer, context
+        if self.progress_bar:
+            self.progress_bar.close()
+
+
+class StepMonitorCallback(TrainCallback):
+    """
+    Customizable logger callback for trainer.
+    
+    This callback provides flexible logging capabilities for training metrics,
+    supporting multiple log formats and custom log handlers.
+    """
+    
+    def __init__(
+        self,
+        log_dir: Optional[str] = None,
+        log_interval: int = 100,
+        metrics: Optional[List[str]] = None
+    ):
+        """
+        Args:
+            log_dir: Directory to save log files. If None, logs won't be saved to file.
+            log_interval: Log every N steps
+            metrics: List of metrics to log. Supported: ['loss', 'lr', 'grad_norm', 'grad_std', grad_max', 'grad_min', 'grad_mean', 'grad_nan_num']
+            custom_handlers: List of custom log handler functions
+            json_log: Whether to save logs in JSON format
+        """
+        
+        self.log_dir = Path(log_dir) if log_dir else Path(os.getcwd()) / "logs"
+        self.log_interval = log_interval
+        self.metrics = metrics or ['loss', 'lr']
+        self.step_num = 0
+        
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+    
+    def _handle_info(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Logs training information to console and file. """
+        
+        log_data = {
+            "timestamp": time.strftime('%Y-%m-%d %H:%M:%S'),
+            "epoch": context.epoch,
+            "iter": context.current_iter,
+            "metrics": self.metrics,
+        }
+        
+        for metric in self.metrics:
+            if metric == 'loss':
+                log_data[metric] = context.loss
+            elif metric == 'lr':
+                log_data[metric] = context.optimizer.param_groups[-1]['lr']
+            elif metric == 'grad_norm':
+                log_data[metric] = grad_norm(trainer.parameter.model)
+            elif metric == 'grad_std':
+                log_data[metric] = grad_std(trainer.parameter.model)
+            elif metric == 'grad_max':
+                log_data[metric] = grad_max(trainer.parameter.model)
+            elif metric == 'grad_min':
+                log_data[metric] = grad_min(trainer.parameter.model)
+            elif metric == 'grad_mean':
+                log_data[metric] = grad_mean(trainer.parameter.model)
+            elif metric == 'grad_nan_num':
+                log_data[metric] = grad_nan_num(trainer.parameter.model)
+            else:
+                raise ValueError(f"Invalid metric: {metric}")
+        
+        return log_data
+    
+    def _handle_log(self, trainer: 'Trainer', context: 'TrainContext'):
+        """ Logs training information to console and file. """
+        log_data = self._handle_info(trainer, context)
+        try:
+            log_file = self.log_dir / f"log_epoch_{context.epoch}_iter_{context.current_iter}.json"
+            with open(log_file, 'a') as f:
+                json.dump(log_data, f, indent=4)
+        except Exception:
+            raise
+    
+    def on_step_end(self, trainer: 'Trainer', context: 'TrainContext'):
+        if self.step_num % self.log_interval == 0:
+            self._handle_log(trainer, context)
+        
+        self.step_num += 1
+        

khaosz/trainer/train_config.py

@@ -0,0 +1,62 @@
+from dataclasses import dataclass, field
+from typing import Optional
+from torch.utils.data import Dataset
+from torch.optim import Optimizer
+from khaosz.trainer.strategy import BaseStrategy
+
+
+@dataclass
+class TrainConfig:
+    
+    strategy: BaseStrategy = field(
+        default=None,
+        metadata={"help": "Training strategy."}
+    )
+    dataset: Dataset = field(
+        default=None,
+        metadata={"help": "Dataset for training."}
+    )
+    optimizer: Optimizer = field(
+        default=None,
+        metadata={"help": "Optimizer for training."}
+    )
+    checkpoint_dir: str = field(
+        default="./checkpoint",
+        metadata={"help": "Checkpoint directory."}
+    )
+    n_epoch: int = field(
+        default=1,
+        metadata={"help": "Number of epochs for training."}
+    )
+    batch_size: int = field(
+        default=4,
+        metadata={"help": "Batch size for training."}
+    )
+    checkpoint_interval: int = field(
+        default=5000,
+        metadata={"help": "Number of iterations between checkpoints."}
+    )
+    accumulation_steps: int = field(
+        default=1,
+        metadata={"help": "Number of iterations between steps."}
+    )
+    max_grad_norm: float = field(
+        default=1.0,
+        metadata={"help": "Maximum gradient norm."}
+    )
+    random_seed: int = field(
+        default=3407,
+        metadata={"help": "Random seed."}
+    )
+    num_workers: int = field(
+        default=0,
+        metadata={"help": "Number of workers for dataloader."}
+    )
+    prefetch_factor: Optional[int] = field(
+        default=None,
+        metadata={"help": "Prefetch factor for dataloader."}
+    )
+    pin_memory: bool = field(
+        default=False,
+        metadata={"help": "Pin memory for dataloader."}
+    )

khaosz/trainer/train_context.py

@@ -0,0 +1,98 @@
+from dataclasses import dataclass, field, fields
+from typing import Optional, Self, TYPE_CHECKING
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader
+from khaosz.core.parameter import Checkpoint
+from khaosz.trainer.data_util import RandomSampler
+
+if TYPE_CHECKING:
+    from khaosz.trainer.trainer import Trainer
+
+
+@dataclass
+class TrainContext:
+    dataloader: DataLoader = field(default=None)
+    optimizer: Optimizer = field(default=None)
+    sampler: RandomSampler = field(default=None)
+    epoch: int = field(default=0)
+    current_iter: int = field(default=0)
+    loss: float = field(default=0.0)
+    checkpoint: Checkpoint = field(default=None)
+    
+    def asdict(self) -> dict:
+        return {field.name: getattr(self, field.name) 
+                for field in fields(self)}
+
+
+class TrainContextBuilder:
+    def __init__(self, trainer: 'Trainer'):
+        self.trainer = trainer
+        self._context = TrainContext(
+            dataloader=None,
+            optimizer=None,
+            sampler=None,
+            epoch=0,
+            current_iter=0,
+            loss=0.0,
+            checkpoint=None
+        )
+    
+    def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
+        if checkpoint is None:
+            checkpoint = Checkpoint(
+                model=self.trainer.parameter.model,
+                tokenizer=self.trainer.parameter.tokenizer,
+                config=self.trainer.parameter.config,
+                sampler_state=None,
+                optimizer_state=None,
+                loss_list=[]
+            )
+        self._context.checkpoint = checkpoint
+        return self
+    
+    def with_sampler(self) -> Self:
+        seed = self.trainer.train_config.random_seed
+        sampler = RandomSampler(
+            data_source=self.trainer.train_config.dataset, 
+            seed=seed
+        )
+        
+        if self._context.checkpoint and self._context.checkpoint.sampler_state:
+            sampler.load_state_dict(self._context.checkpoint.sampler_state)
+        
+        self._context.sampler = sampler
+        self._context.epoch = sampler.epoch
+        self._context.current_iter = sampler.current_iter
+        
+        if self._context.checkpoint:
+            self._context.checkpoint.sampler_state = sampler.state_dict()
+        
+        return self
+    
+    def with_optimizer(self) -> Self:
+        optimizer = self.trainer.train_config.optimizer
+        
+        if self._context.checkpoint and self._context.checkpoint.optimizer_state:
+            optimizer.load_state_dict(self._context.checkpoint.optimizer_state)
+        
+        self._context.optimizer = optimizer
+        
+        if self._context.checkpoint:
+            self._context.checkpoint.optimizer_state = optimizer.state_dict()
+        
+        return self
+    
+    def with_dataloader(self) -> Self:
+        dataloader = DataLoader(
+            self.trainer.train_config.dataset, 
+            batch_size=self.trainer.train_config.batch_size, 
+            sampler=self._context.sampler,
+            num_workers=self.trainer.train_config.num_workers,
+            pin_memory=self.trainer.train_config.pin_memory,
+            prefetch_factor=self.trainer.train_config.prefetch_factor
+        )
+        self._context.dataloader = dataloader
+        return self
+    
+    def build(self) -> TrainContext:
+        return self._context

khaosz/trainer/trainer.py

@@ -1,17 +1,17 @@
 import logging
-from typing import Optional, List, cast
+from typing import Optional, List
-from torch.utils.data import DataLoader
 
 from khaosz.core import ModelParameter, Checkpoint
-from khaosz.trainer.data_util import RandomSampler
+from khaosz.trainer.strategy import ScheduleConfig
-from khaosz.trainer.strategy import TrainConfig, ScheduleConfig
+from khaosz.trainer.train_config import TrainConfig
-from khaosz.trainer.trainer_callback import (
+from khaosz.trainer.train_callback import (
-    TrainerCallback, 
+    TrainCallback, 
     ProgressBarCallback, 
     CheckpointCallback, 
     GradientClippingCallback,
     SchedulerCallback
 )
+from khaosz.trainer.train_context import TrainContext, TrainContextBuilder
 
 logger = logging.getLogger(__name__)
 
@@ -21,120 +21,73 @@ class Trainer:
         parameter: ModelParameter,
         train_config: TrainConfig,
         schedule_config: ScheduleConfig,
-        callbacks: Optional[List[TrainerCallback]] = None
+        callbacks: Optional[List[TrainCallback]] = None
     ):
         self.parameter = parameter
         self.train_config = train_config
         self.schedule_config = schedule_config
         self.callbacks = callbacks or self._get_default_callbacks()
-        
+
-    def _get_default_callbacks(self) -> List[TrainerCallback]:
+    def _get_default_callbacks(self) -> List[TrainCallback]:
         return [
             ProgressBarCallback(),
             CheckpointCallback(self.train_config.checkpoint_interval),
             GradientClippingCallback(),
             SchedulerCallback(self.schedule_config),
         ]
-
-    def _set_train_kwargs(self, kwargs: dict):
-        seed = self.train_config.random_seed
-        sampler = RandomSampler(data_source=self.train_config.dataset, seed=seed)
-        optim = self.train_config.optimizer
-        checkpoint = cast(Checkpoint, kwargs.get('checkpoint', None))
         
-        if checkpoint is None:
+    def _build_train_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
-            checkpoint = Checkpoint(
+        return (TrainContextBuilder(self)
-                model=self.parameter.model,
+                .with_checkpoint(checkpoint)
-                tokenizer=self.parameter.tokenizer,
+                .with_sampler()
-                config=self.parameter.config,
+                .with_optimizer()
-                sampler_state=None,
+                .with_dataloader()
-                optim_state=None,
+                .build())
-                loss_list=[]
+    
-            )
+    def _call_callbacks(self, method_name: str, context: TrainContext):
-        
-        sampler_state = checkpoint.sampler_state
-        optim_state = checkpoint.optim_state
-        
-        if sampler_state: 
-            sampler.load_state_dict(sampler_state)
-        
-        if optim_state: 
-            optim.load_state_dict(optim_state)
-            
-        checkpoint.optim_state = optim.state_dict()
-        checkpoint.sampler_state = sampler.state_dict()
-
-        dataloader = DataLoader(
-            self.train_config.dataset, 
-            batch_size=self.train_config.batch_size, 
-            sampler=sampler
-        )
-        
-        kwargs["dataloader"] = dataloader
-        kwargs["optimizer"] = self.train_config.optimizer
-        kwargs["epoch"] = sampler.epoch
-        kwargs["current_iter"] = sampler.current_iter
-        kwargs["sampler"] = sampler
-        kwargs["checkpoint"] = checkpoint
-        
-    def _call_callbacks(self, method_name: str, **kwargs):
         for callback in self.callbacks:
             method = getattr(callback, method_name, None)
             if method:
-                method(self, **kwargs)
+                method(self, context)
 
-    def train(
+    def train(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
-        self,
+        context = self._build_train_context(checkpoint)
-        checkpoint: Optional[Checkpoint] = None
-    ) -> Checkpoint:
         
-        # train        
+        self._call_callbacks('on_train_begin', context)
-        train_kwargs = {
-            'checkpoint': checkpoint,
-            'dataloader': None,
-            'optimizer': None,
-            'sampler': None,
-            'epoch':  0,
-            'current_iter': 0, 
-            'loss': 0.0,
-        }
-        
-        self._set_train_kwargs(train_kwargs)
-        self._call_callbacks('on_train_begin', **train_kwargs)
-        
-        dataloader = train_kwargs['dataloader']
-        checkpoint = train_kwargs['checkpoint']
-        start_epoch = train_kwargs['epoch']
         
         try:
             self.parameter.model.train()
-            for epoch in range(start_epoch, self.train_config.n_epoch):
+            # 1.epoch
-                # epoch
+            for epoch in range(context.epoch, self.train_config.n_epoch):
-                train_kwargs["epoch"] = epoch
+                context.epoch = epoch
-                self._call_callbacks('on_epoch_begin', **train_kwargs)
+                self._call_callbacks('on_epoch_begin', context)
-                for batch in dataloader:
+                
-                    
+                for batch in context.dataloader:
-                    if train_kwargs["current_iter"] % self.train_config.accumulation_steps == 0:
+                    if context.current_iter % self.train_config.accumulation_steps == 0:
-                        # step
+                        # 2. step
-                        self._call_callbacks('on_step_begin', **train_kwargs)
+                        self._call_callbacks('on_step_begin', context)
                         self.train_config.optimizer.step()
                         self.train_config.optimizer.zero_grad()
-                        self._call_callbacks('on_step_end', **train_kwargs)
+                        self._call_callbacks('on_step_end', context)
-                        
-                    # batch
-                    self._call_callbacks('on_batch_begin', **train_kwargs)
-                    loss = self.train_config.strategy(batch)
-                    train_kwargs["loss"] = loss.item()
-                    train_kwargs["current_iter"] += 1
-                    loss.backward()
                     
-                    self._call_callbacks('on_batch_end', **train_kwargs)
+                    # 3. batch
+                    self._call_callbacks('on_batch_begin', context)
+                    loss = self.train_config.strategy(batch)
+                    context.loss = loss.item()
+                    context.current_iter += 1
+                    
+                    # to make the loss normalized by accumulation steps
+                    normalized_loss = loss / self.train_config.accumulation_steps
+                    normalized_loss.backward()
+                    
+                    self._call_callbacks('on_batch_end', context)
 
-                self._call_callbacks('on_epoch_end', **train_kwargs)
+                self._call_callbacks('on_epoch_end', context)
         
         except Exception as e:
             logger.error(f"Training failed: {str(e)}", exc_info=True)
+            self._call_callbacks('on_error', context)
             raise
         finally:
-            self._call_callbacks('on_train_end', **train_kwargs)
+            self._call_callbacks('on_train_end', context)
-            return checkpoint
+            return context.checkpoint

khaosz/trainer/trainer_callback.py

@@ -1,169 +0,0 @@
-import os
-import torch.optim as optim
-
-from tqdm import tqdm
-from torch.nn.utils import clip_grad_norm_
-from torch.optim.lr_scheduler import LambdaLR
-from typing import Optional, cast, TYPE_CHECKING
-from khaosz.core.parameter import Checkpoint
-from khaosz.trainer.data_util import RandomSampler
-from khaosz.trainer.strategy import ScheduleConfig, SchedulerFactory
-
-if TYPE_CHECKING:
-    from khaosz.trainer.trainer import Trainer
-
-
-class TrainerCallback:
-    """ 
-    Callback interface for trainer.
-    and we use '_' to ignore unused parameters.
-    """
-
-    def on_train_begin(self, trainer: 'Trainer', **kwargs):
-        """ Called at the beginning of training. """
-        _ = trainer, kwargs
-    
-    def on_train_end(self, trainer: 'Trainer', **kwargs):
-        """ Called at the end of training. """
-        _ = trainer, kwargs
-
-    def on_epoch_begin(self, trainer: 'Trainer', **kwargs):
-        """ Called at the beginning of each epoch. """
-        _ = trainer, kwargs
-
-    def on_epoch_end(self, trainer: 'Trainer', **kwargs):
-        """  Called at the end of each epoch. """
-        _ = trainer, kwargs
-
-    def on_batch_begin(self, trainer: 'Trainer', **kwargs):
-        """ Called at the beginning of each batch. """
-        _ = trainer, kwargs
-
-    def on_batch_end(self, trainer: 'Trainer', **kwargs):
-        """ Called at the end of each batch. """
-        _ = trainer, kwargs
-
-    def on_step_begin(self, trainer: 'Trainer', **kwargs):
-        """ Called at the beginning of each step. """
-        _ = trainer, kwargs
-
-    def on_step_end(self, trainer: 'Trainer', **kwargs):
-        """ Called at the end of each step."""
-        _ = trainer, kwargs
-
-
-class ProgressBarCallback(TrainerCallback):
-    """ 
-    Progress bar callback for trainer.
-    """
-    def __init__(self):
-        self.progress_bar: tqdm = None
-    
-    def on_epoch_begin(self, trainer: 'Trainer', **kwargs):
-        epoch = kwargs.get('epoch')
-        dataloader = kwargs.get('dataloader')
-        self.progress_bar = tqdm(
-            dataloader, 
-            desc=f"Epoch {epoch+1}/{trainer.train_config.n_epoch}", 
-            dynamic_ncols=True
-        )
-    
-    def on_batch_end(self, trainer: 'Trainer', **kwargs):
-        _ = trainer
-        loss = kwargs.get('loss')
-        optimizer = cast(optim.Optimizer, kwargs.get('optimizer'))
-        self.progress_bar.set_postfix({
-            "loss": f"{loss:.4f}",
-            "lr": f"{optimizer.param_groups[-1]['lr']:.2e}"
-        })
-        self.progress_bar.update(1)
-    
-    def on_epoch_end(self, trainer: 'Trainer', **kwargs):
-        _ = trainer, kwargs
-        if self.progress_bar:
-            self.progress_bar.close()
-
-
-class CheckpointCallback(TrainerCallback):
-    """ 
-    Checkpoint callback for trainer.
-    """
-    def __init__(self, checkpoint_interval: int):
-        self.checkpoint_interval = checkpoint_interval
-        self.last_ckpt_iter = 0
-    
-    @staticmethod
-    def _save_checkpoint(trainer: 'Trainer', **kwargs):
-        current_iter = kwargs.get('current_iter')
-        random_sampler = cast(RandomSampler, kwargs.get('sampler'))
-        optimizer = cast(optim.Optimizer, kwargs.get('optimizer'))
-        checkpoint = cast(Checkpoint, kwargs.get('checkpoint'))
-        
-        save_path = os.path.join(trainer.train_config.checkpoint_dir, f"iter_{current_iter}")
-        checkpoint.sampler_state = random_sampler.state_dict()
-        checkpoint.optim_state = optimizer.state_dict()
-        
-        checkpoint.save(save_path)
-    
-    def on_batch_end(self, trainer: 'Trainer', **kwargs):
-        current_iter = kwargs.get('current_iter')
-        checkpoint = cast(Checkpoint, kwargs.get('checkpoint'))
-        loss = kwargs.get('loss')
-        checkpoint.loss_list.append(loss)
-        
-        if current_iter - self.last_ckpt_iter >= self.checkpoint_interval:
-            CheckpointCallback._save_checkpoint(trainer, **kwargs)
-            self.last_ckpt_iter = current_iter
-    
-    def on_train_end(self, trainer: 'Trainer', **kwargs):
-        current_iter = kwargs.get('current_iter')
-        if current_iter != self.last_ckpt_iter:
-            CheckpointCallback._save_checkpoint(trainer, **kwargs)
-            self.last_ckpt_iter = current_iter
-
-
-class GradientClippingCallback(TrainerCallback):
-    """ 
-    Gradient clipping callback for trainer.
-    """
-    def on_step_begin(self, trainer: 'Trainer', **kwargs):
-        _ = kwargs
-        clip_grad_norm_(
-            trainer.parameter.model.parameters(),
-            trainer.train_config.max_grad_norm
-        )
-
-
-class SchedulerCallback(TrainerCallback):
-    """
-    Scheduler callback for trainer.
-    """
-    def __init__(self, schedule_config: ScheduleConfig):
-        self.schedule_config = schedule_config
-        self.scheduler: Optional[LambdaLR] = None
-        self.current_iter = 0
-    
-    def on_train_begin(self, trainer: 'Trainer', **kwargs):
-        self.current_iter = kwargs.get('current_iter')
-
-        for group in trainer.train_config.optimizer.param_groups:
-            if "initial_lr" not in group:
-                group["initial_lr"] = group["lr"] 
-
-        self.schedule_config.validate()
-        lambda_scheduler_fn = SchedulerFactory.load_schedule_fn(
-            self.schedule_config
-        )
-        
-        self.scheduler = LambdaLR(
-            trainer.train_config.optimizer,
-            lambda_scheduler_fn,
-            last_epoch=self.current_iter - 1
-        )
-    
-    def on_batch_end(self, trainer: 'Trainer', **kwargs):
-        _ =  trainer, kwargs
-        
-        if self.scheduler:
-            self.scheduler.step()
-            self.current_iter += 1

tests/conftest.py

@@ -0,0 +1,134 @@
+import os
+import json
+import numpy as np
+import tempfile
+import shutil
+import torch
+
+import pytest
+import matplotlib
+from torch.utils.data import Dataset
+
+from khaosz.core import *
+from khaosz.trainer import *
+from khaosz.trainer.data_util import *
+
+matplotlib.use("Agg")
+
+
+class RandomDataset(Dataset):
+    def __init__(self, length=None, max_length=64, vocab_size=1000):
+        self.length = length or int(np.random.randint(100, 200))
+        self.max_length = max_length
+        self.vocab_size = vocab_size
+        
+    def __len__(self):
+        return self.length
+        
+    def __getitem__(self, idx):
+        return {
+            "input_ids": torch.randint(0, self.vocab_size, (self.max_length,)),
+            "target_ids": torch.randint(0, self.vocab_size, (self.max_length,))
+        }
+
+
+class MultiTurnDataset(Dataset):
+    def __init__(self, length=None, max_length=64, vocab_size=1000):
+        self.length = length or int(np.random.randint(100, 200))
+        self.max_length = max_length
+        self.vocab_size = vocab_size
+        
+    def __len__(self):
+        return self.length
+        
+    def __getitem__(self, idx):
+        input_ids = torch.randint(0, self.vocab_size, (self.max_length,))
+        target_ids = torch.randint(0, self.vocab_size, (self.max_length,))
+        loss_mask = build_loss_mask(input_ids, 0, 1)
+        attn_mask = build_attention_mask(input_ids, 2, True)
+
+        return {
+            "input_ids": input_ids,
+            "target_ids": target_ids,
+            "loss_mask": loss_mask,
+            "attn_mask": attn_mask,
+        }
+
+
+class EarlyStoppingDataset(Dataset):
+    def __init__(self, length=10, stop_after=5):
+        self.length = length
+        self.stop_after = stop_after
+        self.count = 0
+        
+    def __len__(self):
+        return self.length
+        
+    def __getitem__(self, idx):
+        self.count += 1
+        if self.count == self.stop_after:
+            raise RuntimeError("Simulated early stopping")
+        
+        return {
+            "input_ids": torch.randint(0, 1000, (64,)),
+            "target_ids": torch.randint(0, 1000, (64,))
+        }
+
+
+@pytest.fixture
+def base_test_env(request: pytest.FixtureRequest):
+    func_name = request.function.__name__
+    test_dir = tempfile.mkdtemp(prefix=f"{func_name}_")
+    config_path = os.path.join(test_dir, "config.json")
+    
+    n_dim_choices = [8, 16, 32]
+    n_head_choices = [2, 4]
+    
+    n_dim = int(np.random.choice(n_dim_choices))
+    n_head = int(np.random.choice(n_head_choices))
+    n_kvhead = n_head // 2
+    d_ffn = n_dim * 2
+
+    config = {
+        "vocab_size": 1000,
+        "n_dim": n_dim,
+        "n_head": n_head,
+        "n_kvhead": n_kvhead,
+        "d_ffn": d_ffn,
+        "m_len": 1024,
+        "n_layer": 4,
+        "norm_eps": 1e-5
+    }
+    
+    with open(config_path, 'w') as f:
+        json.dump(config, f)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    transformer_config = TransformerConfig().load(config_path)
+    model = Transformer(transformer_config).to(device=device)
+    tokenizer = BpeTokenizer()
+    
+    yield {
+        "device": device,
+        "test_dir": test_dir,
+        "config_path": config_path,
+        "transformer_config": transformer_config,
+        "model": model,
+        "tokenizer": tokenizer,
+    }
+    
+    shutil.rmtree(test_dir)
+
+@pytest.fixture
+def random_dataset():
+    dataset = RandomDataset()
+    yield dataset
+
+@pytest.fixture
+def multi_turn_dataset():
+    dataset = MultiTurnDataset()
+    yield dataset
+    
+@pytest.fixture
+def early_stopping_dataset():
+    dataset = EarlyStoppingDataset()
+    yield dataset

tests/test_callbacks.py

@@ -0,0 +1,59 @@
+import torch
+
+from khaosz.core import *
+from khaosz.trainer import *
+from khaosz.trainer.data_util import *
+
+def test_callback_integration(base_test_env, random_dataset):
+    """Test that all callbacks are properly integrated"""
+    optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+    train_config = TrainConfig(
+        dataset=random_dataset,
+        optimizer=optimizer,
+        checkpoint_dir=base_test_env["test_dir"],
+        n_epoch=1,
+        batch_size=2,
+        checkpoint_interval=3,
+        accumulation_steps=1,
+        max_grad_norm=1.0,
+        random_seed=42
+    )
+    
+    schedule_config = CosineScheduleConfig(
+        warmup_steps=10,
+        total_steps=20
+    )
+    
+    # Create custom callbacks to track calls
+    callback_calls = []
+    
+    class TrackingCallback(TrainCallback):
+        def on_train_begin(self, trainer, context):
+            callback_calls.append('on_train_begin')
+        
+        def on_batch_end(self, trainer, context):
+            callback_calls.append('on_batch_end')
+        
+        def on_epoch_end(self, trainer, context):
+            callback_calls.append('on_epoch_end')
+    
+    train_config.strategy = StrategyFactory.load(base_test_env["model"], "seq", base_test_env["device"])
+    model_parameter = ModelParameter(
+        base_test_env["model"], 
+        base_test_env["tokenizer"], 
+        base_test_env["transformer_config"]
+    )
+    
+    trainer = Trainer(
+        model_parameter, 
+        train_config, 
+        schedule_config,
+        callbacks=[TrackingCallback(), ProgressBarCallback()]
+    )
+    
+    trainer.train()
+    
+    # Verify callbacks were called
+    assert 'on_train_begin' in callback_calls
+    assert 'on_batch_end' in callback_calls
+    assert 'on_epoch_end' in callback_calls

tests/test_dataset_loader.py

@@ -0,0 +1,67 @@
+import os
+import torch
+import pickle
+import numpy as np
+
+from khaosz.core import *
+from khaosz.trainer import *
+from khaosz.trainer.data_util import *
+
+def test_dataset_loader_random_paths(base_test_env):
+    """Test dataset loader with multiple random paths"""
+    test_dir = base_test_env["test_dir"]
+    
+    # Create multiple pkl files with random data
+    num_files = np.random.randint(2, 5)
+    pkl_paths = []
+    
+    for i in range(num_files):
+        pkl_path = os.path.join(test_dir, f"test_data_{i}.pkl")
+        seq_length = np.random.randint(50, 100)
+        dummy_data = {
+            "sequence": torch.randint(0, 1000, (seq_length,)),
+            "chosen": torch.randint(0, 1000, (seq_length,)),
+            "rejected": torch.randint(0, 1000, (seq_length,)),
+            "chosen_mask": torch.ones(seq_length, dtype=torch.bool),
+            "rejected_mask": torch.ones(seq_length, dtype=torch.bool)
+        }
+        with open(pkl_path, "wb") as f:
+            pickle.dump(dummy_data, f)
+        pkl_paths.append(pkl_path)
+    
+    # Test loading with multiple paths
+    loaded_dataset = DatasetLoader.load(
+        train_type="seq", 
+        load_path=pkl_paths, 
+        max_len=64, 
+    )
+    assert loaded_dataset is not None
+    assert len(loaded_dataset) > 0
+
+def test_dpo_strategy_with_random_data(base_test_env):
+    """Test DPO strategy with randomized preference data"""
+    test_dir = base_test_env["test_dir"]
+    
+    # Create DPO-style data
+    pkl_path = os.path.join(test_dir, "dpo_data.pkl")
+    seq_length = np.random.randint(40, 80)
+    
+    dummy_data = {
+        "chosen": torch.randint(0, 1000, (seq_length,)),
+        "rejected": torch.randint(0, 1000, (seq_length,)),
+        "chosen_mask": torch.ones(seq_length, dtype=torch.bool),
+        "rejected_mask": torch.ones(seq_length, dtype=torch.bool)
+    }
+    
+    with open(pkl_path, "wb") as f:
+        pickle.dump(dummy_data, f)
+    
+    # Load DPO dataset
+    dpo_dataset = DatasetLoader.load(
+        train_type="dpo", 
+        load_path=pkl_path, 
+        max_len=64, 
+    )
+    
+    assert dpo_dataset is not None
+    assert hasattr(dpo_dataset, 'fetcher')

tests/test_early_stopping.py

@@ -0,0 +1,42 @@
+import torch
+
+from khaosz.core import *
+from khaosz.trainer import *
+from khaosz.trainer.data_util import *
+
+def test_early_stopping_simulation(base_test_env, early_stopping_dataset):
+    """Simulate early stopping behavior"""
+    
+    optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+    train_config = TrainConfig(
+        dataset=early_stopping_dataset,
+        optimizer=optimizer,
+        checkpoint_dir=base_test_env["test_dir"],
+        n_epoch=2,
+        batch_size=2,
+        checkpoint_interval=1,
+        accumulation_steps=1,
+        max_grad_norm=1.0,
+        random_seed=42
+    )
+    
+    train_config.strategy = StrategyFactory.load(base_test_env["model"], "seq", base_test_env["device"])
+    model_parameter = ModelParameter(
+        base_test_env["model"], 
+        base_test_env["tokenizer"], 
+        base_test_env["transformer_config"]
+    )
+    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
+    trainer = Trainer(model_parameter, train_config, schedule_config)
+    
+    # Should handle early stopping gracefully
+    checkpoint = None
+    try:
+        checkpoint = trainer.train()
+        assert len(checkpoint.loss_list) == 2
+    except Exception:
+        # Handle any exceptions
+        pass
+    
+    checkpoint = trainer.train(checkpoint)
+    assert len(checkpoint.loss_list) == 10

tests/test_module.py

@@ -10,8 +10,9 @@ from khaosz.core.generator import EmbeddingEncoderCore, GeneratorCore
 from tokenizers import pre_tokenizers
 
 @pytest.fixture
-def test_env():
+def test_env(request: pytest.FixtureRequest):
-    test_dir = tempfile.mkdtemp()
+    func_name = request.function.__name__
+    test_dir = tempfile.mkdtemp(prefix=f"{func_name}_")
     config_path = os.path.join(test_dir, "config.json")
     tokenizer_path = os.path.join(test_dir, "tokenizer.json")
     model_path = os.path.join(test_dir, "model.safetensors")

tests/test_sampler.py

@@ -0,0 +1,78 @@
+from khaosz.core import *
+from khaosz.trainer import *
+from khaosz.trainer.data_util import *
+
+def test_random_sampler_consistency(random_dataset):
+    """Test RandomSampler produces consistent results with same seed"""
+    dataset = random_dataset
+    
+    # Create two samplers with same seed
+    sampler1 = RandomSampler(dataset, seed=42)
+    sampler2 = RandomSampler(dataset, seed=42)
+    
+    indices1 = list(iter(sampler1))
+    indices2 = list(iter(sampler2))
+    
+    assert indices1 == indices2
+
+def test_random_sampler_different_seeds(random_dataset):
+    """Test RandomSampler produces different results with different seeds"""
+    dataset = random_dataset
+    
+    # Create two samplers with different seeds
+    sampler1 = RandomSampler(dataset, seed=42)
+    sampler2 = RandomSampler(dataset, seed=123)
+    
+    indices1 = list(iter(sampler1))
+    indices2 = list(iter(sampler2))
+    
+    # Very high probability they should be different
+    assert indices1 != indices2
+
+def test_sampler_state_persistence(random_dataset):
+    """Test that sampler state is correctly saved and loaded"""
+    dataset = random_dataset
+    n = len(dataset)
+    
+    # Create sampler and get some indices
+    sampler = RandomSampler(dataset, seed=42)
+    iter1 = iter(sampler)
+    indices1 = [next(iter1) for _ in range(min(10, n))]
+    
+    # Save state
+    state_dict = sampler.state_dict()
+    
+    # Get more indices
+    indices2 = [next(iter1) for _ in range(min(10, n - len(indices1)))]
+    
+    # Create new sampler and load state
+    sampler2 = RandomSampler(dataset, seed=42)
+    sampler2.load_state_dict(state_dict)
+    
+    # Check that new sampler produces same sequence from saved point
+    iter2 = iter(sampler2)
+    indices3 = [next(iter2) for _ in range(min(10, n - len(indices1)))]
+    
+    assert indices2 == indices3
+
+def test_sampler_across_epochs(random_dataset):
+    """Test sampler behavior across multiple epochs"""
+    dataset = random_dataset
+    n = len(dataset)
+    
+    sampler = RandomSampler(dataset, seed=42)
+    
+    # Get indices for first epoch
+    epoch1_indices = list(iter(sampler))
+    assert len(epoch1_indices) == n
+    
+    # Get indices for second epoch
+    epoch2_indices = list(iter(sampler))
+    assert len(epoch2_indices) == n
+    
+    # Check that epochs have different order (should be random)
+    assert epoch1_indices != epoch2_indices
+    
+    # Check that all indices are present in each epoch
+    assert set(epoch1_indices) == set(range(n))
+    assert set(epoch2_indices) == set(range(n))

tests/test_train_config.py

@@ -0,0 +1,85 @@
+import torch
+import numpy as np
+
+from khaosz.core import *
+from khaosz.trainer import *
+from khaosz.trainer.data_util import *
+
+def test_different_batch_sizes(base_test_env, random_dataset):
+    """Test training with different batch sizes"""
+    batch_sizes = [1, 2, 4, 8]
+    
+    for batch_size in batch_sizes:
+        optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+        train_config = TrainConfig(
+            dataset=random_dataset,
+            optimizer=optimizer,
+            checkpoint_dir=base_test_env["test_dir"],
+            n_epoch=1,
+            batch_size=batch_size,
+            checkpoint_interval=5,
+            accumulation_steps=1,
+            max_grad_norm=1.0,
+            random_seed=np.random.randint(1000)
+        )
+        
+        assert train_config.batch_size == batch_size
+
+def test_gradient_accumulation(base_test_env, random_dataset):
+    """Test training with different gradient accumulation steps"""
+    accumulation_steps_list = [1, 2, 4]
+    
+    for accumulation_steps in accumulation_steps_list:
+        optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+        train_config = TrainConfig(
+            dataset=random_dataset,
+            optimizer=optimizer,
+            checkpoint_dir=base_test_env["test_dir"],
+            n_epoch=1,
+            batch_size=2,
+            checkpoint_interval=10,
+            accumulation_steps=accumulation_steps,
+            max_grad_norm=1.0,
+            random_seed=42
+        )
+        
+        schedule_config = CosineScheduleConfig(
+            warmup_steps=10,
+            total_steps=20
+        )
+        train_config.strategy = StrategyFactory.load(base_test_env["model"], "seq", base_test_env["device"])
+        model_parameter = ModelParameter(
+            base_test_env["model"], 
+            base_test_env["tokenizer"], 
+            base_test_env["transformer_config"]
+        )
+        
+        trainer = Trainer(model_parameter, train_config, schedule_config)
+        trainer.train()
+        
+        assert train_config.accumulation_steps == accumulation_steps
+
+def test_memory_efficient_training(base_test_env, random_dataset):
+    """Test training with memory-efficient configurations"""
+    # Test with smaller batch sizes and gradient checkpointing
+    small_batch_configs = [
+        {"batch_size": 1, "accumulation_steps": 8},
+        {"batch_size": 2, "accumulation_steps": 4},
+        {"batch_size": 4, "accumulation_steps": 2}
+    ]
+    
+    for config in small_batch_configs:
+        optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+        train_config = TrainConfig(
+            dataset=random_dataset,
+            optimizer=optimizer,
+            checkpoint_dir=base_test_env["test_dir"],
+            n_epoch=1,
+            batch_size=config["batch_size"],
+            checkpoint_interval=5,
+            accumulation_steps=config["accumulation_steps"],
+            max_grad_norm=1.0,
+            random_seed=42
+        )
+        
+        assert train_config.accumulation_steps == config["accumulation_steps"]

tests/test_train_strategy.py

@@ -0,0 +1,72 @@
+import torch
+import numpy as np
+
+from khaosz.core import *
+from khaosz.trainer import *
+from khaosz.trainer.data_util import *
+
+def test_multi_turn_training(base_test_env, multi_turn_dataset):
+    """Test training with multi-turn conversation data"""
+    optimizer = torch.optim.AdamW(base_test_env["model"].parameters())
+    train_config = TrainConfig(
+        dataset=multi_turn_dataset,
+        optimizer=optimizer,
+        checkpoint_dir=base_test_env["test_dir"],
+        n_epoch=2,
+        batch_size=2,
+        checkpoint_interval=3,
+        accumulation_steps=1,
+        max_grad_norm=1.0,
+        random_seed=int(np.random.randint(1000))
+    )
+    
+    schedule_config = CosineScheduleConfig(
+        warmup_steps=50,
+        total_steps=100
+    )
+    
+    train_config.strategy = StrategyFactory.load(
+        base_test_env["model"], 
+        "sft",
+        base_test_env["device"],
+        bos_token_id=2,
+        eos_token_id=3,
+        user_token_id=1,
+        multi_turn=True
+    )
+    
+    model_parameter = ModelParameter(
+        base_test_env["model"], 
+        base_test_env["tokenizer"], 
+        base_test_env["transformer_config"]
+    )
+    
+    trainer = Trainer(model_parameter, train_config, schedule_config)
+    checkpoint = trainer.train()
+    
+    assert len(checkpoint.loss_list) > 0
+
+def test_schedule_factory_random_configs():
+    """Test scheduler factory with random configurations"""
+    schedule_configs = [
+        CosineScheduleConfig(
+            warmup_steps=np.random.randint(50, 200),
+            total_steps=np.random.randint(1000, 5000),
+            min_rate=np.random.uniform(0.01, 0.1)
+        ),
+        SgdrScheduleConfig(
+            warmup_steps=np.random.randint(50, 200),
+            cycle_length=np.random.randint(500, 2000),
+            t_mult=np.random.randint(1, 3),
+            min_rate=np.random.uniform(0.01, 0.1)
+        )
+    ]
+    
+    for config in schedule_configs:
+        schedule_fn = SchedulerFactory.load_schedule_fn(config)
+        assert callable(schedule_fn)
+        
+        # Test the schedule function at different steps
+        for step in [0, config.warmup_steps // 2, config.warmup_steps, config.warmup_steps * 2]:
+            lr_mult = schedule_fn(step)
+            assert 0 <= lr_mult <= 1

tests/test_trainer.py

@@ -1,466 +0,0 @@
-import os
-import json
-import torch
-import shutil
-import pytest
-import pickle
-import tempfile
-import numpy as np
-
-from torch.utils.data import Dataset
-from khaosz.core import *
-from khaosz.trainer import *
-from khaosz.trainer.data_util import *
-
-import matplotlib
-matplotlib.use('Agg')
-
-
-@pytest.fixture
-def test_env():
-    """Setup test environment with randomized data"""
-    test_dir = tempfile.mkdtemp()
-    config_path = os.path.join(test_dir, "config.json")
-    
-    n_dim_choices = [8, 16, 32]
-    n_head_choices = [2, 4]
-    
-    n_dim = int(np.random.choice(n_dim_choices))
-    n_head = int(np.random.choice(n_head_choices))
-    n_kvhead =  n_head // 2
-    d_ffn = n_dim * 2
-
-    config = {
-        "vocab_size": 1000,
-        "n_dim": n_dim,
-        "n_head": n_head,
-        "n_kvhead": n_kvhead,
-        "d_ffn": d_ffn,
-        "m_len": 1024,
-        "n_layer": 4,
-        "norm_eps": 1e-5
-    }
-    
-    with open(config_path, 'w') as f:
-        json.dump(config, f)
-    
-    transformer_config = TransformerConfig().load(config_path)
-    model = Transformer(transformer_config)
-    tokenizer = BpeTokenizer()
-    
-    class RandomDataset(Dataset):
-        def __init__(self, length=None, max_length=64, vocab_size=1000):
-            self.length = length or int(np.random.randint(100, 200))
-            self.max_length = max_length
-            self.vocab_size = vocab_size
-            
-            
-        def __len__(self):
-            return self.length
-            
-        def __getitem__(self, idx):
-
-            return {
-                "input_ids": torch.randint(0, self.vocab_size, (self.max_length,)),
-                "target_ids": torch.randint(0, self.vocab_size, (self.max_length,))
-            }
-    
-    class MultiTurnDataset(Dataset):
-        def __init__(self, length=None, max_length=64, vocab_size=1000):
-            self.length = length or int(np.random.randint(100, 200))
-            self.max_length = max_length
-            self.vocab_size = vocab_size
-            
-            
-        def __len__(self):
-            return self.length
-            
-        def __getitem__(self, idx):
-            input_ids = torch.randint(0, self.vocab_size, (self.max_length,))
-            target_ids = torch.randint(0, self.vocab_size, (self.max_length,))
-            loss_mask = build_loss_mask(input_ids, 0, 1)
-            attn_mask = build_attention_mask(input_ids, 2, True)
-
-            return {
-                "input_ids": input_ids,
-                "target_ids": target_ids,
-                "loss_mask": loss_mask,
-                "attn_mask": attn_mask,
-            }
-    
-    dataset = RandomDataset()
-    multi_turn_dataset = MultiTurnDataset()
-    
-    yield {
-        "test_dir": test_dir,
-        "config_path": config_path,
-        "transformer_config": transformer_config,
-        "model": model,
-        "tokenizer": tokenizer,
-        "dataset": dataset,
-        "multi_turn_dataset": multi_turn_dataset
-    }
-    
-    shutil.rmtree(test_dir)
-
-def test_dataset_loader_random_paths(test_env):
-    """Test dataset loader with multiple random paths"""
-    test_dir = test_env["test_dir"]
-    
-    # Create multiple pkl files with random data
-    num_files = np.random.randint(2, 5)
-    pkl_paths = []
-    
-    for i in range(num_files):
-        pkl_path = os.path.join(test_dir, f"test_data_{i}.pkl")
-        seq_length = np.random.randint(50, 100)
-        dummy_data = {
-            "sequence": torch.randint(0, 1000, (seq_length,)),
-            "chosen": torch.randint(0, 1000, (seq_length,)),
-            "rejected": torch.randint(0, 1000, (seq_length,)),
-            "chosen_mask": torch.ones(seq_length, dtype=torch.bool),
-            "rejected_mask": torch.ones(seq_length, dtype=torch.bool)
-        }
-        with open(pkl_path, "wb") as f:
-            pickle.dump(dummy_data, f)
-        pkl_paths.append(pkl_path)
-    
-    # Test loading with multiple paths
-    loaded_dataset = DatasetLoader.load(
-        train_type="seq", 
-        load_path=pkl_paths, 
-        max_len=64, 
-        device="cpu"
-    )
-    assert loaded_dataset is not None
-    assert len(loaded_dataset) > 0
-
-
-def test_different_batch_sizes(test_env):
-    """Test training with different batch sizes"""
-    batch_sizes = [1, 2, 4, 8]
-    
-    for batch_size in batch_sizes:
-        optimizer = torch.optim.AdamW(test_env["model"].parameters())
-        train_config = TrainConfig(
-            dataset=test_env["dataset"],
-            optimizer=optimizer,
-            checkpoint_dir=test_env["test_dir"],
-            n_epoch=1,
-            batch_size=batch_size,
-            checkpoint_interval=5,
-            accumulation_steps=1,
-            max_grad_norm=1.0,
-            random_seed=np.random.randint(1000)
-        )
-        
-        assert train_config.batch_size == batch_size
-
-
-def test_random_sampler_consistency(test_env):
-    """Test RandomSampler produces consistent results with same seed"""
-    dataset = test_env["dataset"]
-    
-    # Create two samplers with same seed
-    sampler1 = RandomSampler(dataset, seed=42)
-    sampler2 = RandomSampler(dataset, seed=42)
-    
-    indices1 = list(iter(sampler1))
-    indices2 = list(iter(sampler2))
-    
-    assert indices1 == indices2
-
-
-def test_random_sampler_different_seeds(test_env):
-    """Test RandomSampler produces different results with different seeds"""
-    dataset = test_env["dataset"]
-    
-    # Create two samplers with different seeds
-    sampler1 = RandomSampler(dataset, seed=42)
-    sampler2 = RandomSampler(dataset, seed=123)
-    
-    indices1 = list(iter(sampler1))
-    indices2 = list(iter(sampler2))
-    
-    # Very high probability they should be different
-    assert indices1 != indices2
-
-
-def test_schedule_factory_random_configs(test_env):
-    """Test scheduler factory with random configurations"""
-    schedule_configs = [
-        CosineScheduleConfig(
-            warmup_steps=np.random.randint(50, 200),
-            total_steps=np.random.randint(1000, 5000),
-            min_rate=np.random.uniform(0.01, 0.1)
-        ),
-        SgdrScheduleConfig(
-            warmup_steps=np.random.randint(50, 200),
-            cycle_length=np.random.randint(500, 2000),
-            t_mult=np.random.randint(1, 3),
-            min_rate=np.random.uniform(0.01, 0.1)
-        )
-    ]
-    
-    for config in schedule_configs:
-        schedule_fn = SchedulerFactory.load_schedule_fn(config)
-        assert callable(schedule_fn)
-        
-        # Test the schedule function at different steps
-        for step in [0, config.warmup_steps // 2, config.warmup_steps, config.warmup_steps * 2]:
-            lr_mult = schedule_fn(step)
-            assert 0 <= lr_mult <= 1
-
-
-def test_multi_turn_training(test_env):
-    """Test training with multi-turn conversation data"""
-    optimizer = torch.optim.AdamW(test_env["model"].parameters())
-    train_config = TrainConfig(
-        dataset=test_env["multi_turn_dataset"],
-        optimizer=optimizer,
-        checkpoint_dir=test_env["test_dir"],
-        n_epoch=2,
-        batch_size=2,
-        checkpoint_interval=3,
-        accumulation_steps=1,
-        max_grad_norm=1.0,
-        random_seed=int(np.random.randint(1000))
-    )
-    
-    schedule_config = CosineScheduleConfig(
-        warmup_steps=50,
-        total_steps=100
-    )
-    
-    train_config.strategy = StrategyFactory.load(
-        test_env["model"], 
-        "sft",
-        bos_token_id=2,
-        eos_token_id=3,
-        user_token_id=1,
-        multi_turn=True
-    )
-    
-    model_parameter = ModelParameter(
-        test_env["model"], 
-        test_env["tokenizer"], 
-        test_env["transformer_config"]
-    )
-    
-    trainer = Trainer(model_parameter, train_config, schedule_config)
-    checkpoint = trainer.train()
-    
-    assert len(checkpoint.loss_list) > 0
-
-
-def test_gradient_accumulation(test_env):
-    """Test training with different gradient accumulation steps"""
-    accumulation_steps_list = [1, 2, 4]
-    
-    for accumulation_steps in accumulation_steps_list:
-        optimizer = torch.optim.AdamW(test_env["model"].parameters())
-        train_config = TrainConfig(
-            dataset=test_env["dataset"],
-            optimizer=optimizer,
-            checkpoint_dir=test_env["test_dir"],
-            n_epoch=1,
-            batch_size=2,
-            checkpoint_interval=10,
-            accumulation_steps=accumulation_steps,
-            max_grad_norm=1.0,
-            random_seed=42
-        )
-        
-        schedule_config = CosineScheduleConfig(
-            warmup_steps=10,
-            total_steps=20
-        )
-        
-        train_config.strategy = StrategyFactory.load(
-            test_env["model"], 
-            "seq"
-        )
-        
-        model_parameter = ModelParameter(
-            test_env["model"], 
-            test_env["tokenizer"], 
-            test_env["transformer_config"]
-        )
-        
-        trainer = Trainer(model_parameter, train_config, schedule_config)
-        trainer.train()
-        
-        assert train_config.accumulation_steps == accumulation_steps
-
-def test_dpo_strategy_with_random_data(test_env):
-    """Test DPO strategy with randomized preference data"""
-    test_dir = test_env["test_dir"]
-    
-    # Create DPO-style data
-    pkl_path = os.path.join(test_dir, "dpo_data.pkl")
-    seq_length = np.random.randint(40, 80)
-    
-    dummy_data = {
-        "chosen": torch.randint(0, 1000, (seq_length,)),
-        "rejected": torch.randint(0, 1000, (seq_length,)),
-        "chosen_mask": torch.ones(seq_length, dtype=torch.bool),
-        "rejected_mask": torch.ones(seq_length, dtype=torch.bool)
-    }
-    
-    with open(pkl_path, "wb") as f:
-        pickle.dump(dummy_data, f)
-    
-    # Load DPO dataset
-    dpo_dataset = DatasetLoader.load(
-        train_type="dpo", 
-        load_path=pkl_path, 
-        max_len=64, 
-        device="cpu"
-    )
-    
-    assert dpo_dataset is not None
-    assert hasattr(dpo_dataset, 'fetcher')
-
-
-def test_callback_integration(test_env):
-    """Test that all callbacks are properly integrated"""
-    optimizer = torch.optim.AdamW(test_env["model"].parameters())
-    train_config = TrainConfig(
-        dataset=test_env["dataset"],
-        optimizer=optimizer,
-        checkpoint_dir=test_env["test_dir"],
-        n_epoch=1,
-        batch_size=2,
-        checkpoint_interval=3,
-        accumulation_steps=1,
-        max_grad_norm=1.0,
-        random_seed=42
-    )
-    
-    schedule_config = CosineScheduleConfig(
-        warmup_steps=10,
-        total_steps=20
-    )
-    
-    # Create custom callbacks to track calls
-    callback_calls = []
-    
-    class TrackingCallback(TrainerCallback):
-        def on_train_begin(self, trainer, **kwargs):
-            callback_calls.append('on_train_begin')
-        
-        def on_batch_end(self, trainer, **kwargs):
-            callback_calls.append('on_batch_end')
-        
-        def on_epoch_end(self, trainer, **kwargs):
-            callback_calls.append('on_epoch_end')
-    
-    train_config.strategy = StrategyFactory.load(test_env["model"], "seq")
-    model_parameter = ModelParameter(
-        test_env["model"], 
-        test_env["tokenizer"], 
-        test_env["transformer_config"]
-    )
-    
-    trainer = Trainer(
-        model_parameter, 
-        train_config, 
-        schedule_config,
-        callbacks=[TrackingCallback(), ProgressBarCallback()]
-    )
-    
-    trainer.train()
-    
-    # Verify callbacks were called
-    assert 'on_train_begin' in callback_calls
-    assert 'on_batch_end' in callback_calls
-    assert 'on_epoch_end' in callback_calls
-
-
-def test_memory_efficient_training(test_env):
-    """Test training with memory-efficient configurations"""
-    # Test with smaller batch sizes and gradient checkpointing
-    small_batch_configs = [
-        {"batch_size": 1, "accumulation_steps": 8},
-        {"batch_size": 2, "accumulation_steps": 4},
-        {"batch_size": 4, "accumulation_steps": 2}
-    ]
-    
-    for config in small_batch_configs:
-        optimizer = torch.optim.AdamW(test_env["model"].parameters())
-        train_config = TrainConfig(
-            dataset=test_env["dataset"],
-            optimizer=optimizer,
-            checkpoint_dir=test_env["test_dir"],
-            n_epoch=1,
-            batch_size=config["batch_size"],
-            checkpoint_interval=5,
-            accumulation_steps=config["accumulation_steps"],
-            max_grad_norm=1.0,
-            random_seed=42
-        )
-        
-        assert train_config.accumulation_steps == config["accumulation_steps"]
-
-
-def test_early_stopping_simulation(test_env):
-    """Simulate early stopping behavior"""
-    class EarlyStoppingDataset(Dataset):
-        def __init__(self, length=10, stop_after=5):
-            self.length = length
-            self.stop_after = stop_after
-            self.count = 0
-            
-        def __len__(self):
-            return self.length
-            
-        def __getitem__(self, idx):
-            self.count += 1
-            if self.count == self.stop_after:
-                raise RuntimeError("Simulated early stopping")
-            
-            return {
-                "input_ids": torch.randint(0, 1000, (64,)),
-                "target_ids": torch.randint(0, 1000, (64,))
-            }
-    
-    dataset = EarlyStoppingDataset()
-    
-    optimizer = torch.optim.AdamW(test_env["model"].parameters())
-    train_config = TrainConfig(
-        dataset=dataset,
-        optimizer=optimizer,
-        checkpoint_dir=test_env["test_dir"],
-        n_epoch=2,
-        batch_size=2,
-        checkpoint_interval=1,
-        accumulation_steps=1,
-        max_grad_norm=1.0,
-        random_seed=42
-    )
-    
-    train_config.strategy = StrategyFactory.load(test_env["model"], "seq")
-    model_parameter = ModelParameter(
-        test_env["model"], 
-        test_env["tokenizer"], 
-        test_env["transformer_config"]
-    )
-    schedule_config = CosineScheduleConfig(warmup_steps=10, total_steps=20)
-    trainer = Trainer(model_parameter, train_config, schedule_config)
-    
-    # Should handle early stopping gracefully
-    checkpoint = None
-    try:
-        checkpoint = trainer.train()
-        assert len(checkpoint.loss_list) == 2
-    except Exception:
-        # Handle any exceptions
-        pass
-    
-    checkpoint = trainer.train(checkpoint)
-    assert len(checkpoint.loss_list) == 10 + 1
-
-
-if __name__ == "__main__":
-    # Run all tests
-    pytest.main([__file__, "-v"])

train.py

@@ -59,6 +59,7 @@ def train(
     strategy = StrategyFactory.load(
         model, 
         train_type,
+        device,
         **kwargs
     )
 
@@ -66,7 +67,6 @@ def train(
         train_type=train_type,
         load_path=cache_files,
         max_len=parameter.config.m_len,
-        device=device,
         **kwargs
     )