refactor(khaosz): 重构项目结构

khaosz/config/__init__.py

@@ -0,0 +1,12 @@
+from khaosz.config.model_config import TransformerConfig
+from khaosz.config.param_config import BaseModelIO, ModelParameter, Checkpoint, ParameterLoader
+from khaosz.config.train_config import TrainConfig
+
+__all__ = [
+    "BaseModelIO",
+    "ModelParameter",
+    "Checkpoint",
+    "ParameterLoader",
+    "TransformerConfig",
+    "TrainConfig"
+]

khaosz/config/model_config.py

@@ -0,0 +1,36 @@
+import json
+
+from dataclasses import asdict, dataclass
+from typing import  Optional, Self
+
+@dataclass
+class TransformerConfig:
+    # basic config
+    vocab_size: Optional[int] = None
+    n_dim: Optional[int] = None
+    n_head: Optional[int] = None
+    n_layer: Optional[int] = None
+    m_len: Optional[int] = None
+    norm_eps: Optional[float] = None
+    d_ffn: Optional[int] = None
+    
+    # GQA
+    n_kvhead: Optional[int] = None
+    
+    
+    def load(self, config_path: str) -> Self:
+        with open(config_path, 'r') as f:
+            config: dict = json.load(f)
+        for key, value in config.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+    
+        return self
+                    
+    def save(self, config_path: str) -> None:
+        config_dict = asdict(self)
+        config_dict = {k: v for k, v in config_dict.items() if v is not None}
+        with open(config_path, 'w') as f:
+            json.dump(config_dict, f, indent=4)
+
+

khaosz/config/param_config.py

@@ -0,0 +1,238 @@
+import pickle as pkl
+import matplotlib.pyplot as plt
+import safetensors.torch as st
+import torch.nn as nn
+import torch.optim as optim
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Self, Union
+from pathlib import Path
+
+from khaosz.data.tokenizer import BpeTokenizer
+from khaosz.config.model_config import TransformerConfig
+from khaosz.model.transformer import Transformer
+
+
+class BaseModelIO:
+    """Base class for model I/O operations."""
+    
+    def __init__(
+        self, 
+        model: Optional[nn.Module] = None, 
+        tokenizer: Optional[BpeTokenizer] = None,
+        config: Optional[TransformerConfig] = None
+    ):
+        self.model = model
+        self.tokenizer = tokenizer or BpeTokenizer()
+        self.config = config or TransformerConfig()
+    
+    def _get_file_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
+        """Get standardized file paths for model components."""
+        dir_path = Path(directory)
+        return {
+            "model": dir_path / "model.safetensors",
+            "config": dir_path / "config.json", 
+            "tokenizer": dir_path / "tokenizer.json"
+        }
+    
+    def save_components(self, save_dir: Union[str, Path]):
+        """Save core model components."""
+        paths = self._get_file_paths(save_dir)
+        paths["model"].parent.mkdir(parents=True, exist_ok=True)
+        
+        if self.model is not None:
+            st.save_file(self.model.state_dict(), str(paths["model"]))
+        self.config.save(str(paths["config"]))
+        self.tokenizer.save(str(paths["tokenizer"]))
+    
+    def load_components(self, load_dir: Union[str, Path]) -> Self:
+        """Load core model components."""
+        paths = self._get_file_paths(load_dir)
+        
+        self.config.load(str(paths["config"]))
+        self.tokenizer.load(str(paths["tokenizer"]))
+        
+        if paths["model"].exists():
+            state_dict = st.load_file(str(paths["model"]))
+            if self.model is None:
+                self.model = Transformer(self.config)
+            self.model.load_state_dict(state_dict)
+        
+        return self
+    
+    def to(self, *args, **kwargs) -> Self:
+        """Move model to device."""
+        if self.model is not None:
+            self.model.to(*args, **kwargs)
+        return self
+
+
+@dataclass
+class ModelParameter(BaseModelIO):
+    """Container for model parameters with serialization capabilities."""
+    
+    model: Optional[nn.Module] = field(
+        default=None,
+        metadata={"help": "Transformer model."}
+    )
+    tokenizer: BpeTokenizer = field(
+        default_factory=BpeTokenizer,
+        metadata={"help": "Tokenizer for the model."}
+    )
+    config: TransformerConfig = field(
+        default_factory=TransformerConfig,
+        metadata={"help": "Transformer model configuration."}
+    )
+    
+    def save(self, save_dir: Union[str, Path]):
+        self.save_components(save_dir)
+    
+    def load(self, load_dir: Union[str, Path]) -> Self:
+        return self.load_components(load_dir)
+
+
+@dataclass
+class Checkpoint(BaseModelIO):
+    """Extended model parameters with training state."""
+    
+    model: Optional[nn.Module] = field(
+        default=None,
+        metadata={"help": "Transformer model."}
+    )
+    tokenizer: BpeTokenizer = field(
+        default=None,
+        metadata={"help": "Tokenizer for the model."}
+    )
+    config: TransformerConfig = field(
+        default=None,
+        metadata={"help": "Transformer model configuration."}
+    )
+    optimizer_state: Dict[str, Any] = field(
+        default=None,
+        metadata={"help": "Optimizer state."}
+    )
+    scheduler_state: Dict[str, Any] = field(
+        default=None,
+        metadata={"help": "Sampler state."}
+    )
+    loss_list: List[float] = field(
+        default_factory=list,
+        metadata={"help": "List of training losses."}
+    )
+    
+    def _get_training_paths(self, directory: Union[str, Path]) -> dict[str, Path]:
+        paths = self._get_file_paths(directory)
+        paths.update({
+            "loss_list": paths["model"].parent / "loss.pkl",
+            "loss_plot": paths["model"].parent / "loss.png",
+            "optimizer_state": paths["model"].parent / "optimizer_state.pkl",
+            "sampler_state": paths["model"].parent / "sampler_state.pkl"
+        })
+        return paths
+    
+    def save_training_state(self, save_dir: Union[str, Path]):
+        paths = self._get_training_paths(save_dir)
+        
+        # Save loss plot
+        self._plot_loss(str(paths["loss_plot"]))
+        
+        # Save loss list
+        with open(str(paths["loss_list"]), "wb") as f:
+            pkl.dump(self.loss_list, f)
+        
+        # Save optimizer state
+        with open(str(paths["optimizer_state"]), "wb") as f:
+            pkl.dump(self.optimizer_state, f)
+        
+        # Save sampler state
+        with open(str(paths["sampler_state"]), "wb") as f:
+            pkl.dump(self.scheduler_state, f)
+    
+    def load_training_state(self, load_dir: Union[str, Path]) -> Self:
+        paths = self._get_training_paths(load_dir)
+        
+        # Load loss list
+        if paths["loss_list"].exists():
+            with open(str(paths["loss_list"]), "rb") as f:
+                self.loss_list = pkl.load(f)
+        
+        # Load optimizer state
+        if paths["optimizer_state"].exists():
+            with open(str(paths["optimizer_state"]), "rb") as f:
+                self.optimizer_state = pkl.load(f)
+        
+        # Load sampler state
+        if paths["sampler_state"].exists():
+            with open(str(paths["sampler_state"]), "rb") as f:
+                self.scheduler_state = pkl.load(f)
+        
+        return self
+    
+    def _plot_loss(self, save_path: str): 
+        """Plot and save loss curve."""
+        if not self.loss_list:
+            return
+        
+        current_iter = len(self.loss_list)
+        
+        plt.figure(figsize=(10, 6))
+        plt.plot(self.loss_list)
+        plt.title(f"Training Loss - Iteration {current_iter}")
+        plt.xlabel("Batch")
+        plt.ylabel("Loss")
+        plt.grid(True)
+        plt.savefig(save_path, dpi=300, bbox_inches="tight")
+        plt.close()
+    
+    def save(self, save_dir: Union[str, Path]):
+        """Save complete checkpoint."""
+        self.save_components(save_dir)
+        self.save_training_state(save_dir)
+    
+    def load(self, load_dir: Union[str, Path]) -> Self:
+        """Load complete checkpoint."""
+        self.load_components(load_dir)
+        self.load_training_state(load_dir)
+        return self
+
+
+class ParameterLoader:
+    """Factory class for loading model parameters or checkpoints."""
+    
+    @staticmethod
+    def load(load_dir: Union[str, Path]) -> Union[ModelParameter, Checkpoint]:
+        """Load either ModelParameter or Checkpoint based on directory contents."""
+        load_dir = Path(load_dir)
+        
+        # Check for training-specific files
+        loss_file = load_dir / "loss.pkl"
+        has_training_data = loss_file.exists()
+        
+        # Create appropriate instance
+        if has_training_data:
+            checkpoint = Checkpoint()
+            checkpoint.load(str(load_dir))
+            return checkpoint
+        else:
+            params = ModelParameter()
+            params.load(str(load_dir))
+            return params
+    
+    @staticmethod
+    def create_checkpoint(
+        model: nn.Module, 
+        tokenizer: BpeTokenizer,
+        config: TransformerConfig,
+        loss_list: Optional[list[float]] = None,
+        optimizer: Optional[optim.Optimizer] = None,
+    ) -> Checkpoint:
+        """Convenience method to create a training checkpoint."""
+        return Checkpoint(
+            model=model,
+            tokenizer=tokenizer,
+            config=config,
+            loss_list=loss_list or [],
+            optimizer_state=optimizer
+        )
+
+

khaosz/config/train_config.py

@@ -0,0 +1,64 @@
+from dataclasses import dataclass, field
+from typing import Optional, TYPE_CHECKING
+from torch.utils.data import Dataset
+from torch.optim import Optimizer
+
+if TYPE_CHECKING:
+    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."}
+    )