refactor(khaosz): 重构项目结构

khaosz/trainer/__init__.py

@@ -1,6 +1,4 @@
-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 (
     CosineScheduleConfig, 
     SgdrScheduleConfig,
@@ -17,19 +15,16 @@ from khaosz.trainer.train_callback import (
 )
 
 __all__ = [
-    "DatasetLoader",
     "Trainer",
-    "TrainConfig",
+    "StrategyFactory",
     "CosineScheduleConfig",
     "SgdrScheduleConfig",
-    "StrategyFactory",
     "SchedulerFactory",
     
     # callback
     "TrainCallback",
     "ProgressBarCallback",
     "CheckpointCallback",
-    "TrainCallback",
     "SchedulerCallback",
     "StepMonitorCallback"
 ]

khaosz/trainer/data_util.py

@@ -1,320 +0,0 @@
-import torch
-import bisect
-import pickle as pkl
-from abc import ABC, abstractmethod
-from torch import Tensor
-from torch.utils.data import Dataset, Sampler
-from typing import Callable, List, Dict, Literal, Union
-
-MutiSeg = Dict[str, List[Tensor]]
-Seg = Dict[str, Tensor]
-
-def load_pkl_files(paths: List[str]):
-    segments: MutiSeg = {}
-    total_samples = 0
-
-    for path in paths:
-        with open(path, "rb") as f:
-            pkl_file: Seg = pkl.load(f)
-        for key, value in pkl_file.items():
-            if key not in segments:
-                segments[key] = []
-            segments[key].append(value)
-        first_key = list(pkl_file.keys())[0]
-        total_samples += pkl_file[first_key].numel()
-    
-    return segments, total_samples
-
-def build_attention_mask(input_ids: Tensor, user_token_id: int, multi_turn: bool) -> Tensor:
-    seq_len = input_ids.size(0)
-    turn_id = input_ids.eq(user_token_id).cumsum(dim=-1)
-    
-    iq = turn_id.view(seq_len, 1)
-    ik = turn_id.view(1, seq_len)
-    
-    # fix the causual attention mask(iq >= ik condition)
-    seq_mask = (iq >= ik) if multi_turn else (iq == ik)
-    attention_mask = torch.tril(seq_mask)
-    
-    # fix the shape  (bsz, 1, seq_len, seq_len) unsqueeze for broadcast
-    return attention_mask.unsqueeze(0)
-
-def build_loss_mask(input_ids: Tensor, bos_token_id: int, eos_token_id: int) -> Tensor:
-    token_markers = torch.zeros_like(input_ids, dtype=torch.int8)
-    
-    is_bos_token = input_ids.eq(bos_token_id)
-    is_eos_token = input_ids.eq(eos_token_id)
-    
-    # fix the eos_token_id bug(change target_ids to input_ids)
-    token_markers[is_bos_token] = 1
-    token_markers[is_eos_token] = -1 
-    
-    cumulative_markers = torch.cumsum(token_markers, dim=-1)
-    min_cumulative = cumulative_markers.min(dim=-1, keepdim=True).values
-    loss_mask = cumulative_markers - min_cumulative
-    
-    return loss_mask.to(dtype=torch.bool)
-
-
-class BaseSegmentFetcher:
-    def __init__(self, segments: List[Tensor]):
-        self.segments = segments
-        self.cum_lengths = []
-        total = 0
-        for seg in segments:
-            total += len(seg)
-            self.cum_lengths.append(total)
-        self.total_length = total if segments else 0
-
-    def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
-        if not (0 <= begin_idx < self.total_length and 0 <= end_idx <= self.total_length):
-            raise ValueError("begin_idx or end_idx out of bounds")
-        if begin_idx >= end_idx:
-            return torch.tensor([], dtype=torch.long)
-        
-        # fix the range index bug
-        seg_start_idx = bisect.bisect_right(self.cum_lengths, begin_idx)
-        seg_end_idx = bisect.bisect_left(self.cum_lengths, end_idx)
-
-        result_segments = []
-
-        for i in range(seg_start_idx, seg_end_idx + 1):
-            prev_cum = self.cum_lengths[i - 1] if i > 0 else 0
-            start = max(begin_idx - prev_cum, 0)
-            end = min(end_idx - prev_cum, len(self.segments[i]))
-            result_segments.append(self.segments[i][start:end])
-
-        return torch.cat(result_segments, dim=0)
-    
-
-class MutiSegmentFetcher:
-    def __init__(self, muti_segments: MutiSeg):
-        self.muti_keys = list(muti_segments.keys())
-        self.muti_fetchers = {
-            key: BaseSegmentFetcher(segments)
-            for key, segments in muti_segments.items()
-        }
-        
-    def key_fetch(self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]) -> Union[Tensor, Seg]:
-        fetch_dict = {} 
-        keys = [keys] if isinstance(keys, str) else keys
-        
-        for key in keys:
-            fetcher = self.muti_fetchers[key]
-            fetch_tensor = fetcher.fetch_data(begin_idx, end_idx)
-            fetch_dict[key] = fetch_tensor
-
-        return fetch_dict if len(keys) > 1 else fetch_dict[keys[0]]
-    
-    def fetch_data(self, begin_idx: int, end_idx: int) -> Union[Tensor, Seg]:
-        return self.key_fetch(begin_idx, end_idx, self.muti_keys)
-
-
-class BaseDataset(Dataset, ABC):
-    def __init__(self, chunk_size: int):
-        super().__init__()
-        self.segments: MutiSeg = {}
-        self.chunk_size = chunk_size
-        self.total_samples = 0
-
-    def save(self, save_path: str):
-        keys = list(self.segments.keys())
-        if not keys:
-            return
-        
-        first_item = self.segments[keys[0]]
-        segment_size = len(first_item)
-        
-        for i in range(segment_size):
-            formated_segment = {key: self.segments[key][i] for key in keys}
-            pkl.dump(formated_segment, open(f"{save_path}_{i}.pkl", "wb"))
-                
-    
-    def load(self, load_path: Union[str, List[str]]):
-        paths = [load_path] if isinstance(load_path, str) else load_path
-        self.segments, self.total_samples = load_pkl_files(paths)
-        self.fetcher = MutiSegmentFetcher(self.segments)
-        
-    @abstractmethod
-    def __getitem__(self, index: int) -> Dict[str, Tensor]:
-        raise NotImplementedError
-        
-    def __len__(self) -> int:
-        assert self.total_samples // self.chunk_size > 0
-        return self.total_samples // self.chunk_size
-    
-
-class SeqDataset(BaseDataset):
-    def __init__(
-        self, 
-        chunk_size, 
-    ):
-        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):
-        # fix the range index bug
-        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(dtype=torch.long)
-        y = self._fetch_data(begin_idx + 1, end_idx + 1).to(dtype=torch.long)
-        
-        return {"input_ids": x, "target_ids": y}
-    
-    
-class SftDataset(BaseDataset):
-    def __init__(
-        self, 
-        chunk_size, 
-        bos_token_id, 
-        eos_token_id, 
-        user_token_id, 
-        multi_turn=False, 
-    ):
-        super().__init__(chunk_size)
-        self.fetcher = MutiSegmentFetcher(self.segments)
-        self.bos_token_id = bos_token_id
-        self.eos_token_id = eos_token_id
-        self.user_token_id = user_token_id
-        self.multi_turn = multi_turn
-    
-    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 = 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(dtype=torch.long)
-        y = self._fetch_data(begin_idx + 1, end_idx + 1).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)
-        attn_mask = build_attention_mask(x, self.user_token_id, self.multi_turn)
-        
-        return {"input_ids": x, "target_ids": y, "loss_mask": loss_mask, "attn_mask": attn_mask}
-
-
-class DpoDataset(BaseDataset):
-    def __init__(self, chunk_size: int):
-        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):
-        begin_idx = min(index * self.chunk_size, self.total_samples - self.chunk_size - 1)
-        end_idx = begin_idx + self.chunk_size
-        
-        chosen = self._fetch_data(begin_idx, end_idx, "chosen").to(dtype=torch.long)
-        rejected = self._fetch_data(begin_idx, end_idx, "rejected").to(dtype=torch.long)
-        chosen_mask = self._fetch_data(begin_idx, end_idx, "chosen_mask").to(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):
-        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
-
-        input_ids =  self._fetch_data(begin_idx, end_idx, "input_ids"),
-        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")
-        
-        return {"input_ids": input_ids, "actions": actions, "logprobs": logprobs, "rewards": rewards}
-    
-
-class DatasetLoader:
-    @staticmethod       
-    def load(
-        train_type: Literal["seq", "sft", "dpo"],
-        load_path: Union[str, List[str]],
-        max_len: int, 
-        **kwargs
-        ) -> BaseDataset:
-        
-        dataset_router: Dict[str, Callable[[int], BaseDataset]] = {
-            "seq": lambda max_len: SeqDataset(max_len),
-            "sft": lambda max_len: SftDataset(
-                max_len, 
-                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 max_len: DpoDataset(max_len),
-        }
-        dataset = dataset_router[train_type](max_len)
-        dataset.load(load_path)
-        
-        return dataset
-    
-
-class RandomSampler(Sampler[int]):
-    def __init__(self, data_source, generator=None, seed=42):
-        self.data_source = data_source
-        self.seed = seed
-        self.epoch = 0
-        self.current_iter = 0
-        self._indices = None
-        
-        if generator is None:
-            self.generator = torch.Generator()
-            self.generator.manual_seed(seed)
-        else:
-            self.generator = generator
-    
-    def _generate_indices(self):
-        n = len(self.data_source)
-        self._indices = torch.randperm(n, generator=self.generator).tolist()
-    
-    def __iter__(self):
-        n = len(self.data_source)
-        
-        if self._indices is None:
-            self._generate_indices()
-        
-        start = self.current_iter % n
-        for i in range(start, n):
-            self.current_iter += 1
-            yield self._indices[i]
-        
-        self.epoch += 1
-        self._indices = None
-    
-    def __len__(self):
-        return len(self.data_source)
-    
-    def state_dict(self):
-        return {
-            'epoch': self.epoch,
-            'current_iter': self.current_iter,
-            'seed': self.seed,
-            'generator_state': self.generator.get_state() if self.generator else None,
-            'indices': self._indices
-        }
-    
-    def load_state_dict(self, state_dict):
-        self.epoch = state_dict['epoch']
-        self.current_iter = state_dict['current_iter']
-        self.seed = state_dict['seed']
-        
-        if self.generator and state_dict['generator_state'] is not None:
-            self.generator.set_state(state_dict['generator_state'])
-        
-        self._indices = state_dict['indices']

khaosz/trainer/train_callback.py

@@ -106,7 +106,7 @@ class CheckpointCallback(TrainCallback):
     
     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.scheduler_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

khaosz/trainer/train_config.py

@@ -1,62 +0,0 @@
-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

@@ -1,9 +1,10 @@
 from dataclasses import dataclass, field, fields
 from typing import Optional, Self, TYPE_CHECKING
 from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
-from khaosz.core.parameter import Checkpoint
-from khaosz.trainer.data_util import RandomSampler
+from khaosz.config.param_config import Checkpoint
+from khaosz.data.data_util import ResumeableRandomSampler
 
 if TYPE_CHECKING:
     from khaosz.trainer.trainer import Trainer
@@ -13,11 +14,11 @@ if TYPE_CHECKING:
 class TrainContext:
     dataloader: DataLoader = field(default=None)
     optimizer: Optimizer = field(default=None)
-    sampler: RandomSampler = field(default=None)
+    scheduler: LRScheduler = field(default=None)
+    checkpoint: Checkpoint = 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) 
@@ -27,15 +28,7 @@ class TrainContext:
 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
-        )
+        self._context = TrainContext()
     
     def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
         if checkpoint is None:
@@ -43,32 +36,10 @@ class TrainContextBuilder:
                 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
         
@@ -82,11 +53,22 @@ class TrainContextBuilder:
         
         return self
     
+    def with_scheduler(self) -> Self:
+        return self
+        
+    
     def with_dataloader(self) -> Self:
+        resumeable_sampler = ResumeableRandomSampler(
+            data_source=self.trainer.train_config.dataset,
+            start_epoch=self._context.epoch,
+            start_iter=self._context.current_iter,
+            seed=self.trainer.train_config.random_seed
+        )
+        
         dataloader = DataLoader(
             self.trainer.train_config.dataset, 
             batch_size=self.trainer.train_config.batch_size, 
-            sampler=self._context.sampler,
+            sampler=resumeable_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

khaosz/trainer/trainer.py

@@ -1,9 +1,9 @@
 import logging
 from typing import Optional, List
 
-from khaosz.core import ModelParameter, Checkpoint
+from khaosz.config import ModelParameter, Checkpoint
 from khaosz.trainer.strategy import ScheduleConfig
-from khaosz.trainer.train_config import TrainConfig
+from khaosz.config.train_config import TrainConfig
 from khaosz.trainer.train_callback import (
     TrainCallback, 
     ProgressBarCallback, 
@@ -39,8 +39,8 @@ class Trainer:
     def _build_train_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
         return (TrainContextBuilder(self)
                 .with_checkpoint(checkpoint)
-                .with_sampler()
                 .with_optimizer()
+                .with_scheduler()
                 .with_dataloader()
                 .build())