chore: 更新项目名称

2026-03-31 09:34:11 +08:00
parent 780b9e1855
commit 2e009cf59a
51 changed files with 118 additions and 118 deletions
@@ -0,0 +1,24 @@
+from astrai.config.model_config import ModelConfig
+from astrai.config.param_config import BaseModelIO, ModelParameter
+from astrai.config.schedule_config import (
+    ScheduleConfig,
+    CosineScheduleConfig,
+    SGDRScheduleConfig,
+    ScheduleConfigFactory,
+)
+from astrai.config.train_config import TrainConfig
+
+
+__all__ = [
+    # Base I/O
+    "BaseModelIO",
+    "ModelParameter",
+    # Model configuration
+    "ModelConfig",
+    "TrainConfig",
+    # Schedule configuration
+    "ScheduleConfig",
+    "CosineScheduleConfig",
+    "SGDRScheduleConfig",
+    "ScheduleConfigFactory",
+]
@@ -0,0 +1,42 @@
+import json
+
+from dataclasses import asdict, dataclass
+from typing import Optional, Self
+
+
+@dataclass
+class ModelConfig:
+    # basic config
+    vocab_size: Optional[int] = None
+    dim: Optional[int] = None
+
+    n_layers: Optional[int] = None
+    norm_eps: Optional[float] = None
+    dim_ffn: Optional[int] = None
+    tie_weight: Optional[bool] = None
+
+    # RoPE
+    max_len: Optional[int] = None
+    rope_theta: Optional[float] = None
+
+    # GQA
+    n_heads: Optional[int] = None
+    n_kv_heads: Optional[int] = None
+    use_qk_norm: Optional[bool] = None
+    use_gated_attention: Optional[bool] = None
+
+    def load(self, config_path: str) -> Self:
+        config = {}
+        with open(config_path, "r") as f:
+            config.update(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):
+        config_dict = {k: v for k, v in asdict(self).items() if v is not None}
+        with open(config_path, "w") as f:
+            json.dump(config_dict, f, indent=4)
@@ -0,0 +1,81 @@
+import torch.nn as nn
+import safetensors.torch as st
+
+from dataclasses import dataclass, field
+from typing import Optional, Self, Union
+from pathlib import Path
+
+from astrai.data.tokenizer import BpeTokenizer
+from astrai.config.model_config import ModelConfig
+from astrai.model.transformer import Transformer
+
+
+@dataclass
+class BaseModelIO:
+    """Base class for model I/O operations."""
+
+    model: Optional[nn.Module] = field(
+        default=None, metadata={"help": "Transformer model."}
+    )
+    tokenizer: BpeTokenizer = field(
+        default_factory=BpeTokenizer, metadata={"help": "Tokenizer for the model."}
+    )
+    config: ModelConfig = field(
+        default_factory=ModelConfig,
+        metadata={"help": "Transformer model configuration."},
+    )
+
+    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 self.model is None:
+            self.model = Transformer(self.config)
+
+        if paths["model"].exists():
+            state_dict = st.load_file(str(paths["model"]))
+            self.model.load_state_dict(state_dict)
+
+        return self
+
+    def to(self, *args, **kwargs) -> "BaseModelIO":
+        """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."""
+
+    @classmethod
+    def save(cls, instance: "ModelParameter", save_dir: Union[str, Path]):
+        instance.save_components(save_dir)
+
+    @classmethod
+    def load(cls, load_dir: Union[str, Path]) -> "ModelParameter":
+        instance = cls()
+        return instance.load_components(load_dir)
@@ -0,0 +1,149 @@
+from typing import Any, Dict, Type
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+
+
+@dataclass
+class ScheduleConfig(ABC):
+    """Base configuration class for learning rate schedulers.
+
+    Provides common validation and interface for all schedule types.
+    """
+
+    schedule_type: str = field(
+        default="cosine",
+        metadata={
+            "help": "Type of learning rate schedule.",
+            "choices": ["cosine", "sgdr"],
+        },
+    )
+    warmup_steps: int = field(
+        default=1000, metadata={"help": "Number of warmup steps."}
+    )
+    min_rate: float = field(
+        default=0.05, metadata={"help": "Minimum learning rate multiplier."}
+    )
+
+    @abstractmethod
+    def get_kwargs(self) -> Dict[str, Any]:
+        """Get configuration kwargs for scheduler creation."""
+        raise NotImplementedError
+
+    def validate(self) -> None:
+        """Validate configuration parameters."""
+        if self.warmup_steps < 0:
+            raise ValueError(
+                f"warmup_steps must be non-negative, got {self.warmup_steps}"
+            )
+        if not 0 <= self.min_rate <= 1:
+            raise ValueError(f"min_rate must be between 0 and 1, got {self.min_rate}")
+
+
+@dataclass
+class CosineScheduleConfig(ScheduleConfig):
+    """Cosine annealing learning rate schedule configuration."""
+
+    total_steps: int = field(
+        default=None, metadata={"help": "Total training steps for cosine schedule."}
+    )
+
+    def __post_init__(self) -> None:
+        self.schedule_type = "cosine"
+        self.validate()
+
+    def get_kwargs(self) -> Dict[str, Any]:
+        if self.total_steps is None:
+            raise ValueError("total_steps must be specified for cosine schedule")
+
+        return {
+            "schedule_type": self.schedule_type,
+            "warmup_steps": self.warmup_steps,
+            "lr_decay_steps": self.total_steps - self.warmup_steps,
+            "min_rate": self.min_rate,
+        }
+
+    def validate(self) -> None:
+        super().validate()
+        if self.total_steps is not None and self.total_steps <= self.warmup_steps:
+            raise ValueError(
+                f"total_steps ({self.total_steps}) must be greater than warmup_steps ({self.warmup_steps})"
+            )
+
+
+@dataclass
+class SGDRScheduleConfig(ScheduleConfig):
+    """Stochastic Gradient Descent with Warm Restarts schedule configuration."""
+
+    cycle_length: int = field(
+        default=1000, metadata={"help": "Length of the first cycle in steps."}
+    )
+    t_mult: int = field(
+        default=2, metadata={"help": "Multiplier for cycle length growth."}
+    )
+
+    def __post_init__(self) -> None:
+        self.schedule_type = "sgdr"
+        self.validate()
+
+    def get_kwargs(self) -> Dict[str, Any]:
+        return {
+            "schedule_type": self.schedule_type,
+            "warmup_steps": self.warmup_steps,
+            "cycle_length": self.cycle_length,
+            "min_rate": self.min_rate,
+            "t_mult": self.t_mult,
+        }
+
+    def validate(self) -> None:
+        super().validate()
+        if self.cycle_length <= 0:
+            raise ValueError(f"cycle_length must be positive, got {self.cycle_length}")
+        if self.t_mult < 1:
+            raise ValueError(f"t_mult must be >= 1, got {self.t_mult}")
+
+
+class ScheduleConfigFactory:
+    """Factory class for creating ScheduleConfig instances.
+
+    Supports both direct instantiation and factory creation methods.
+
+    Example usage:
+        # Direct creation
+        config = CosineScheduleConfig(total_steps=10000)
+
+        # Factory method
+        config = ScheduleConfigFactory.create("cosine", total_steps=10000)
+    """
+
+    CONFIG_MAP: Dict[str, Type[ScheduleConfig]] = {
+        "cosine": CosineScheduleConfig,
+        "sgdr": SGDRScheduleConfig,
+    }
+
+    @classmethod
+    def create(cls, schedule_type: str, **kwargs) -> ScheduleConfig:
+        """Create a schedule config instance.
+
+        Args:
+            schedule_type: Type of schedule ("cosine", "sgdr")
+            **kwargs: Arguments passed to the config constructor
+
+        Returns:
+            ScheduleConfig instance
+
+        Raises:
+            ValueError: If schedule_type is not supported
+        """
+        if schedule_type not in cls.CONFIG_MAP:
+            raise ValueError(
+                f"Unknown schedule type: '{schedule_type}'. "
+                f"Supported types: {sorted(cls.CONFIG_MAP.keys())}"
+            )
+
+        config_cls = cls.CONFIG_MAP[schedule_type]
+        return config_cls(**kwargs)
+
+    @classmethod
+    def available_types(cls) -> list:
+        """Return list of available schedule type names."""
+        return list(cls.CONFIG_MAP.keys())
@@ -0,0 +1,101 @@
+import torch.nn as nn
+from torch.utils.data import Dataset
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+
+from dataclasses import dataclass, field
+from typing import Callable, List, Optional
+
+
+@dataclass
+class TrainConfig:
+    # basic setting
+    model: nn.Module = field(default=None, metadata={"help": "Model for training."})
+    strategy: str = field(default=None, metadata={"help": "Training strategy."})
+    dataset: Dataset = field(default=None, metadata={"help": "Dataset for training."})
+    optimizer_fn: Callable[[nn.Module], Optimizer] = field(
+        default=None, metadata={"help": "Optimizer factory for training."}
+    )
+    scheduler_fn: Callable[[Optimizer], LRScheduler] = field(
+        default=None, metadata={"help": "Scheduler factory for training."}
+    )
+    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."})
+    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."}
+    )
+
+    # checkpoint setting
+    start_epoch: int = field(default=0, metadata={"help": "Start epoch for training."})
+    start_batch: int = field(
+        default=0, metadata={"help": "Start batch iteration for training."}
+    )
+    ckpt_dir: str = field(
+        default="./checkpoint", metadata={"help": "Checkpoint directory."}
+    )
+    ckpt_interval: int = field(
+        default=5000, metadata={"help": "Number of iterations between checkpoints."}
+    )
+
+    # dataloader setting
+    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."}
+    )
+
+    # distributed training
+    nprocs: int = field(
+        default=1, metadata={"help": "Number of processes for distributed training."}
+    )
+    backend: str = field(
+        default="nccl", metadata={"help": "Distributed training backend."}
+    )
+    master_addr: str = field(
+        default="localhost",
+        metadata={"help": "Master address for distributed training."},
+    )
+    master_port: str = field(
+        default="29500", metadata={"help": "Master port for distributed training."}
+    )
+    parallel_wrapper: Optional[Callable] = field(
+        default=None, metadata={"help": "Parallel function for training."}
+    )
+    state_dict_fn: Optional[Callable] = field(
+        default=None, metadata={"help": "Parallel function for state  dict saving."}
+    )
+
+    # others
+    device_ids: Optional[List[int]] = field(
+        default=None, metadata={"help": "Device ids for distributed training."}
+    )
+    device_type: str = field(
+        default="cuda", metadata={"help": "Device type for distributed training."}
+    )
+    extra_kwargs: dict = field(
+        default_factory=dict, metadata={"help": "Other arguments."}
+    )
+
+    def __post_init__(self):
+        self.validate()
+
+    def validate(self):
+        required_fields = [
+            "model",
+            "strategy",
+            "dataset",
+            "optimizer_fn",
+            "scheduler_fn",
+        ]
+
+        for field_name in required_fields:
+            if getattr(self, field_name) is None:
+                raise ValueError(f"{field_name} is required.")