chore: 更新项目名称
This commit is contained in:
@@ -0,0 +1,35 @@
|
||||
__version__ = "1.3.2"
|
||||
__author__ = "ViperEkura"
|
||||
|
||||
from astrai.config import (
|
||||
ModelConfig,
|
||||
TrainConfig,
|
||||
)
|
||||
from astrai.model.transformer import Transformer
|
||||
from astrai.data import DatasetLoader, BpeTokenizer
|
||||
from astrai.inference.generator import (
|
||||
GenerationRequest,
|
||||
LoopGenerator,
|
||||
StreamGenerator,
|
||||
BatchGenerator,
|
||||
EmbeddingEncoder,
|
||||
GeneratorFactory,
|
||||
)
|
||||
from astrai.trainer import Trainer, StrategyFactory, SchedulerFactory
|
||||
|
||||
__all__ = [
|
||||
"Transformer",
|
||||
"ModelConfig",
|
||||
"TrainConfig",
|
||||
"DatasetLoader",
|
||||
"BpeTokenizer",
|
||||
"GenerationRequest",
|
||||
"LoopGenerator",
|
||||
"StreamGenerator",
|
||||
"BatchGenerator",
|
||||
"EmbeddingEncoder",
|
||||
"GeneratorFactory",
|
||||
"Trainer",
|
||||
"StrategyFactory",
|
||||
"SchedulerFactory",
|
||||
]
|
||||
@@ -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.")
|
||||
@@ -0,0 +1,31 @@
|
||||
from astrai.data.dataset import (
|
||||
BaseDataset,
|
||||
SEQDataset,
|
||||
DPODataset,
|
||||
SFTDataset,
|
||||
GRPODataset,
|
||||
MultiSegmentFetcher,
|
||||
DatasetLoader,
|
||||
DatasetFactory,
|
||||
)
|
||||
|
||||
from astrai.data.tokenizer import BpeTokenizer
|
||||
from astrai.data.sampler import ResumableDistributedSampler
|
||||
|
||||
__all__ = [
|
||||
# Base classes
|
||||
"BaseDataset",
|
||||
# Dataset implementations
|
||||
"SEQDataset",
|
||||
"SFTDataset",
|
||||
"DPODataset",
|
||||
"GRPODataset",
|
||||
# Fetchers
|
||||
"MultiSegmentFetcher",
|
||||
# Factory (DatasetLoader is alias for backward compatibility)
|
||||
"DatasetLoader",
|
||||
"DatasetFactory",
|
||||
# Tokenizer and sampler
|
||||
"BpeTokenizer",
|
||||
"ResumableDistributedSampler",
|
||||
]
|
||||
@@ -0,0 +1,367 @@
|
||||
"""Dataset implementations with factory pattern for training."""
|
||||
|
||||
import torch
|
||||
import bisect
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from torch import Tensor
|
||||
from torch.utils.data import Dataset
|
||||
from astrai.data.serialization import load_h5
|
||||
from typing import List, Dict, Optional, Union
|
||||
|
||||
|
||||
class BaseSegmentFetcher:
|
||||
"""Fetches data segments across multiple tensor segments.
|
||||
|
||||
Maintains cumulative lengths for efficient range queries across
|
||||
multiple discontinuous segments.
|
||||
"""
|
||||
|
||||
def __init__(self, segments: List[Tensor]):
|
||||
self.segments = segments
|
||||
self.cum_lengths = []
|
||||
|
||||
total = 0
|
||||
for seg in segments:
|
||||
total += torch.numel(seg)
|
||||
self.cum_lengths.append(total)
|
||||
|
||||
self.total_length = total
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.total_length
|
||||
|
||||
def fetch_data(self, begin_idx: int, end_idx: int) -> Tensor:
|
||||
"""Fetch data in the range [begin_idx, end_idx).
|
||||
|
||||
Args:
|
||||
begin_idx: Starting index (inclusive)
|
||||
end_idx: Ending index (exclusive)
|
||||
|
||||
Returns:
|
||||
Concatenated tensor of data in the specified range
|
||||
"""
|
||||
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)
|
||||
|
||||
# Find segment boundaries for the range
|
||||
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]))
|
||||
data = self.segments[i][start:end]
|
||||
result_segments.append(data)
|
||||
|
||||
return torch.cat(result_segments, dim=0)
|
||||
|
||||
|
||||
class MultiSegmentFetcher:
|
||||
"""Manages multiple segment fetchers for different data keys.
|
||||
|
||||
Each key corresponds to a different type of data (e.g., "sequence", "mask").
|
||||
"""
|
||||
|
||||
def __init__(self, muti_segments: Dict):
|
||||
self.muti_keys = list(muti_segments.keys())
|
||||
self.muti_fetchers = {
|
||||
key: BaseSegmentFetcher(segments) for key, segments in muti_segments.items()
|
||||
}
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Returns the minimum length across all fetchers."""
|
||||
len_list = [len(seg) for seg in self.muti_fetchers.values()]
|
||||
return min(len_list)
|
||||
|
||||
def key_fetch(
|
||||
self, begin_idx: int, end_idx: int, keys: Union[str, List[str]]
|
||||
) -> Dict:
|
||||
"""Fetch data for specific keys.
|
||||
|
||||
Args:
|
||||
begin_idx: Starting index
|
||||
end_idx: Ending index
|
||||
keys: Single key or list of keys to fetch
|
||||
|
||||
Returns:
|
||||
Dictionary of tensors if multiple keys, single tensor if one key
|
||||
"""
|
||||
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) -> Dict:
|
||||
"""Fetch all keys."""
|
||||
return self.key_fetch(begin_idx, end_idx, self.muti_keys)
|
||||
|
||||
|
||||
class BaseDataset(Dataset, ABC):
|
||||
"""Abstract base class for all dataset types.
|
||||
|
||||
Implements common functionality for window-based data fetching.
|
||||
"""
|
||||
|
||||
def __init__(self, window_size: int, stride: int):
|
||||
super().__init__()
|
||||
self.segments = {}
|
||||
self.window_size = window_size
|
||||
self.stride = stride
|
||||
self.total_samples = None
|
||||
self.fetcher: Optional[MultiSegmentFetcher] = None
|
||||
|
||||
def load(self, load_path: str):
|
||||
"""Load dataset from HDF5 file.
|
||||
|
||||
Args:
|
||||
load_path: Path to the HDF5 data file
|
||||
"""
|
||||
self.segments = load_h5(load_path)
|
||||
self.fetcher = MultiSegmentFetcher(self.segments)
|
||||
self.total_samples = len(self.fetcher)
|
||||
|
||||
def get_index(self, index: int) -> tuple:
|
||||
"""Calculate begin and end indices for a sample.
|
||||
|
||||
Args:
|
||||
index: Sample index
|
||||
|
||||
Returns:
|
||||
Tuple of (begin_idx, end_idx)
|
||||
"""
|
||||
assert self.total_samples > self.window_size
|
||||
|
||||
begin_idx = min(index * self.stride, self.total_samples - 1 - self.window_size)
|
||||
end_idx = min(begin_idx + self.window_size, self.total_samples - 1)
|
||||
|
||||
return begin_idx, end_idx
|
||||
|
||||
@abstractmethod
|
||||
def __getitem__(self, index: int) -> Dict[str, Tensor]:
|
||||
"""Get a single sample by index.
|
||||
|
||||
Must be implemented by subclasses.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __len__(self) -> int:
|
||||
assert self.total_samples is not None
|
||||
if self.total_samples <= self.window_size:
|
||||
return 0
|
||||
return (self.total_samples - 1 - self.window_size) // self.stride + 1
|
||||
|
||||
|
||||
class DatasetFactory:
|
||||
"""Factory class for creating dataset instances.
|
||||
|
||||
Supports decorator-based registration for extensible dataset types.
|
||||
All default dataset types (seq, sft, dpo, grpo) are registered automatically
|
||||
when their classes are defined with the decorator.
|
||||
|
||||
Example usage:
|
||||
@DatasetFactory.register("custom")
|
||||
class CustomDataset(BaseDataset):
|
||||
...
|
||||
|
||||
dataset = DatasetFactory.create("custom", window_size, stride)
|
||||
"""
|
||||
|
||||
SUPPORTED_TYPES = frozenset({"seq", "sft", "dpo", "grpo"})
|
||||
DATASET_MAP: Dict[str, type] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str):
|
||||
"""Decorator to register a new dataset class.
|
||||
|
||||
Args:
|
||||
name: Registration name for the dataset type
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the dataset class
|
||||
"""
|
||||
|
||||
def decorator(dataset_cls: type) -> type:
|
||||
if not issubclass(dataset_cls, BaseDataset):
|
||||
raise TypeError(f"{dataset_cls.__name__} must inherit from BaseDataset")
|
||||
cls.DATASET_MAP[name] = dataset_cls
|
||||
return dataset_cls
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def create(cls, train_type: str, window_size: int, stride: int) -> BaseDataset:
|
||||
"""Create a dataset instance.
|
||||
|
||||
Args:
|
||||
train_type: Type of training ("seq", "sft", "dpo", "grpo")
|
||||
window_size: Window size for data sampling
|
||||
stride: Stride between consecutive samples
|
||||
|
||||
Returns:
|
||||
Dataset instance
|
||||
"""
|
||||
if train_type not in cls.SUPPORTED_TYPES:
|
||||
raise ValueError(
|
||||
f"Unknown dataset type: '{train_type}'. "
|
||||
f"Supported types: {sorted(cls.SUPPORTED_TYPES)}"
|
||||
)
|
||||
|
||||
if train_type not in cls.DATASET_MAP:
|
||||
raise NotImplementedError(
|
||||
f"Dataset type '{train_type}' is supported but not yet implemented."
|
||||
)
|
||||
|
||||
dataset_cls = cls.DATASET_MAP[train_type]
|
||||
return dataset_cls(window_size, stride)
|
||||
|
||||
@classmethod
|
||||
def load(
|
||||
cls,
|
||||
train_type: str,
|
||||
load_path: str,
|
||||
window_size: int,
|
||||
stride: Optional[int] = None,
|
||||
) -> BaseDataset:
|
||||
"""Create and load a dataset in one step.
|
||||
|
||||
Args:
|
||||
train_type: Type of training dataset
|
||||
load_path: Path to the data file
|
||||
window_size: Window size for data sampling
|
||||
stride: Stride between consecutive samples (default: same as window_size)
|
||||
|
||||
Returns:
|
||||
Loaded dataset instance
|
||||
"""
|
||||
if stride is None:
|
||||
stride = window_size
|
||||
|
||||
dataset = cls.create(train_type, window_size, stride)
|
||||
dataset.load(load_path)
|
||||
|
||||
return dataset
|
||||
|
||||
@classmethod
|
||||
def available_types(cls) -> list:
|
||||
"""Return list of registered dataset type names."""
|
||||
return list(cls.DATASET_MAP.keys())
|
||||
|
||||
|
||||
# ============== Dataset Classes ==============
|
||||
# All dataset classes are registered at class definition time using the decorator
|
||||
|
||||
|
||||
@DatasetFactory.register("seq")
|
||||
class SEQDataset(BaseDataset):
|
||||
"""Dataset for sequential next-token prediction training."""
|
||||
|
||||
def __init__(self, window_size: int, stride: int):
|
||||
super().__init__(window_size, stride)
|
||||
|
||||
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, end_idx = self.get_index(index)
|
||||
|
||||
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}
|
||||
|
||||
|
||||
@DatasetFactory.register("sft")
|
||||
class SFTDataset(BaseDataset):
|
||||
"""Dataset for supervised fine-tuning with loss masking."""
|
||||
|
||||
def __init__(self, window_size: int, stride: int):
|
||||
super().__init__(window_size, stride)
|
||||
|
||||
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):
|
||||
begin_idx, end_idx = self.get_index(index)
|
||||
|
||||
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(
|
||||
dtype=torch.long
|
||||
)
|
||||
loss_mask = self._fetch_data(begin_idx + 1, end_idx + 1, "loss_mask").to(
|
||||
dtype=torch.bool
|
||||
)
|
||||
|
||||
return {"input_ids": x, "target_ids": y, "loss_mask": loss_mask}
|
||||
|
||||
|
||||
@DatasetFactory.register("dpo")
|
||||
class DPODataset(BaseDataset):
|
||||
"""Dataset for Direct Preference Optimization training."""
|
||||
|
||||
def __init__(self, window_size: int, stride: int):
|
||||
super().__init__(window_size, stride)
|
||||
|
||||
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, end_idx = self.get_index(index)
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
|
||||
@DatasetFactory.register("grpo")
|
||||
class GRPODataset(BaseDataset):
|
||||
"""Dataset for Group Relative Policy Optimization training."""
|
||||
|
||||
def __init__(self, window_size: int, stride: int):
|
||||
super().__init__(window_size, stride)
|
||||
|
||||
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, end_idx = self.get_index(index)
|
||||
|
||||
prompts = self._fetch_data(begin_idx, end_idx, "prompts")
|
||||
responses = self._fetch_data(begin_idx, end_idx, "responses")
|
||||
masks = self._fetch_data(begin_idx, end_idx, "masks")
|
||||
rewards = self._fetch_data(begin_idx, end_idx, "rewards")
|
||||
|
||||
return {
|
||||
"prompts": prompts,
|
||||
"responses": responses,
|
||||
"masks": masks,
|
||||
"rewards": rewards,
|
||||
}
|
||||
|
||||
|
||||
# Backward compatibility alias
|
||||
DatasetLoader = DatasetFactory
|
||||
@@ -0,0 +1,78 @@
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from torch.utils.data import Dataset, Sampler
|
||||
from typing import Optional
|
||||
|
||||
|
||||
class ResumableDistributedSampler(Sampler[int]):
|
||||
def __init__(
|
||||
self,
|
||||
data_source: Dataset,
|
||||
start_epoch: int = 0,
|
||||
start_iter: int = 0,
|
||||
seed: int = 42,
|
||||
drop_last: bool = False,
|
||||
shuffle: bool = True,
|
||||
process_group: Optional[dist.ProcessGroup] = None,
|
||||
):
|
||||
self.epoch = start_epoch
|
||||
self.iter = start_iter
|
||||
self.seed = seed
|
||||
self.num_samples = len(data_source)
|
||||
|
||||
if process_group is not None:
|
||||
# input process group
|
||||
self.rank = dist.get_rank(process_group)
|
||||
self.num_replicas = dist.get_world_size(process_group)
|
||||
|
||||
elif dist.is_available() and dist.is_initialized():
|
||||
# use default process group
|
||||
process_group = dist.group.WORLD
|
||||
self.rank = dist.get_rank()
|
||||
self.num_replicas = dist.get_world_size()
|
||||
|
||||
else:
|
||||
# single process
|
||||
self.rank = 0
|
||||
self.num_replicas = 1
|
||||
|
||||
self.drop_last = drop_last
|
||||
self.shuffle = shuffle
|
||||
|
||||
offset = 0 if drop_last else self.num_replicas - 1
|
||||
self.num_samples_per_replica = (self.num_samples + offset) // self.num_replicas
|
||||
self.total_size = self.num_samples_per_replica * self.num_replicas
|
||||
|
||||
self._indices = None
|
||||
|
||||
def _get_indices(self):
|
||||
if self.shuffle:
|
||||
generator = torch.Generator()
|
||||
generator.manual_seed(self.seed + self.epoch)
|
||||
indices = torch.randperm(self.num_samples, generator=generator).tolist()
|
||||
else:
|
||||
indices = torch.arange(self.num_samples).tolist()
|
||||
|
||||
if not self.drop_last and self.num_samples < self.total_size:
|
||||
padding_size = self.total_size - len(indices)
|
||||
indices += indices[:padding_size]
|
||||
|
||||
local_indices = indices[self.rank : self.total_size : self.num_replicas]
|
||||
|
||||
self.iter = self.iter % self.num_samples_per_replica
|
||||
self._indices = local_indices[self.iter :]
|
||||
|
||||
def __iter__(self):
|
||||
if self._indices is None:
|
||||
self._get_indices()
|
||||
|
||||
for i in self._indices:
|
||||
self.iter += 1
|
||||
yield i
|
||||
|
||||
self.epoch += 1
|
||||
self._indices = None
|
||||
|
||||
def __len__(self):
|
||||
return self.num_samples_per_replica
|
||||
@@ -0,0 +1,105 @@
|
||||
import os
|
||||
import h5py
|
||||
import torch
|
||||
import json
|
||||
import safetensors.torch as st
|
||||
import torch.distributed as dist
|
||||
|
||||
from pathlib import Path
|
||||
from torch import Tensor
|
||||
from typing import Any, Dict, List
|
||||
from astrai.parallel.setup import get_rank
|
||||
|
||||
|
||||
def save_h5(file_path: str, file_name: str, tensor_group: Dict[str, List[Tensor]]):
|
||||
os.makedirs(file_path, exist_ok=True)
|
||||
full_file_path = os.path.join(file_path, f"{file_name}.h5")
|
||||
with h5py.File(full_file_path, "w") as f:
|
||||
for key, tensors in tensor_group.items():
|
||||
grp = f.create_group(key)
|
||||
for idx, tensor in enumerate(tensors):
|
||||
arr = tensor.cpu().numpy()
|
||||
grp.create_dataset(f"data_{idx}", data=arr)
|
||||
|
||||
|
||||
def load_h5(file_path: str, share_memory=True) -> Dict[str, List[Tensor]]:
|
||||
tensor_group: Dict[str, List[Tensor]] = {}
|
||||
|
||||
root_path = Path(file_path)
|
||||
h5_files = list(root_path.rglob("*.h5")) + list(root_path.rglob("*.hdf5"))
|
||||
|
||||
for h5_file in h5_files:
|
||||
with h5py.File(h5_file, "r") as f:
|
||||
for key in f.keys():
|
||||
grp = f[key]
|
||||
dsets = []
|
||||
for dset_name in grp.keys():
|
||||
dset = grp[dset_name]
|
||||
tensor = torch.from_numpy(dset[:])
|
||||
if share_memory:
|
||||
tensor = tensor.share_memory_()
|
||||
dsets.append(tensor)
|
||||
|
||||
if tensor_group.get(key) is None:
|
||||
tensor_group[key] = []
|
||||
tensor_group[key].extend(dsets)
|
||||
|
||||
return tensor_group
|
||||
|
||||
|
||||
class Checkpoint:
|
||||
def __init__(
|
||||
self,
|
||||
state_dict: Dict[str, Any],
|
||||
epoch: int = 0,
|
||||
iteration: int = 0,
|
||||
):
|
||||
self.state_dict = state_dict
|
||||
self.epoch = epoch
|
||||
self.iteration = iteration
|
||||
|
||||
def save(
|
||||
self,
|
||||
save_dir: str,
|
||||
) -> None:
|
||||
|
||||
save_path = Path(save_dir)
|
||||
save_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rank = get_rank()
|
||||
if rank == 0:
|
||||
meta = {
|
||||
"epoch": self.epoch,
|
||||
"iteration": self.iteration,
|
||||
}
|
||||
with open(save_path / "meta.json", "w") as f:
|
||||
json.dump(meta, f, indent=2)
|
||||
|
||||
st.save_file(self.state_dict, save_path / "state_dict.safetensors")
|
||||
|
||||
@classmethod
|
||||
def load(
|
||||
cls,
|
||||
save_dir: str,
|
||||
) -> "Checkpoint":
|
||||
|
||||
rank = get_rank()
|
||||
save_path = Path(save_dir)
|
||||
|
||||
meta = {}
|
||||
if rank == 0:
|
||||
with open(Path(save_dir) / "meta.json", "r") as f:
|
||||
meta = json.load(f)
|
||||
|
||||
if dist.is_initialized():
|
||||
meta_list = [meta]
|
||||
dist.broadcast_object_list(meta_list, src=0)
|
||||
meta = meta_list[0]
|
||||
|
||||
state_dict = st.load_file(save_path / "state_dict.safetensors")
|
||||
|
||||
return cls(
|
||||
state_dict=state_dict,
|
||||
epoch=meta["epoch"],
|
||||
iteration=meta["iteration"],
|
||||
)
|
||||
@@ -0,0 +1,131 @@
|
||||
from tokenizers import Tokenizer, Encoding
|
||||
from tokenizers import decoders, processors, normalizers, pre_tokenizers
|
||||
from tokenizers.models import BPE
|
||||
from tokenizers.trainers import BpeTrainer
|
||||
from typing import List, Union
|
||||
|
||||
|
||||
class BpeTokenizer:
|
||||
def __init__(self, path=None):
|
||||
self._control_tokens = ["<bos>", "<eos>", "<pad>"]
|
||||
self._special_tokens = ["<|im_start|>", "<|im_end|>"]
|
||||
|
||||
model = BPE()
|
||||
self._tokenizer = Tokenizer(model)
|
||||
self._tokenizer.normalizer = normalizers.Sequence(
|
||||
[normalizers.NFC(), normalizers.Strip()]
|
||||
)
|
||||
|
||||
self._tokenizer.pre_tokenizer = pre_tokenizers.Sequence(
|
||||
[
|
||||
pre_tokenizers.UnicodeScripts(),
|
||||
pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=True),
|
||||
]
|
||||
)
|
||||
|
||||
self._tokenizer.decoder = decoders.ByteLevel()
|
||||
self._tokenizer.post_processor = processors.ByteLevel(trim_offsets=True)
|
||||
|
||||
if path is not None:
|
||||
self._tokenizer = Tokenizer.from_file(path)
|
||||
|
||||
def _prepare_trainer(
|
||||
self,
|
||||
vocab_size: int,
|
||||
min_freq: int,
|
||||
reserved_token_size: int,
|
||||
max_token_length=18,
|
||||
) -> tuple:
|
||||
assert reserved_token_size > len(self._special_tokens)
|
||||
reserved_tokens = [
|
||||
f"<|reserve{i:02d}|>"
|
||||
for i in range(reserved_token_size - len(self._special_tokens))
|
||||
]
|
||||
detail_vocab_size = vocab_size - (
|
||||
len(reserved_tokens) + len(self._special_tokens)
|
||||
)
|
||||
|
||||
alphabet = pre_tokenizers.ByteLevel.alphabet()
|
||||
min_size = len(alphabet) + len(self._control_tokens)
|
||||
assert detail_vocab_size > min_size
|
||||
|
||||
trainer = BpeTrainer(
|
||||
vocab_size=detail_vocab_size,
|
||||
min_frequency=min_freq,
|
||||
limit_alphabet=detail_vocab_size // 6,
|
||||
max_token_length=max_token_length,
|
||||
special_tokens=self._control_tokens,
|
||||
initial_alphabet=alphabet,
|
||||
show_progress=True,
|
||||
)
|
||||
|
||||
return trainer, detail_vocab_size, reserved_tokens
|
||||
|
||||
def train(self, files, vocab_size, min_freq, reserved_token_size=100):
|
||||
trainer, _, reserved_tokens = self._prepare_trainer(
|
||||
vocab_size=vocab_size,
|
||||
min_freq=min_freq,
|
||||
reserved_token_size=reserved_token_size,
|
||||
)
|
||||
self._tokenizer.train(files=files, trainer=trainer)
|
||||
self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
|
||||
|
||||
def train_from_iterator(
|
||||
self, iterator, vocab_size, min_freq, reserved_token_size=100
|
||||
):
|
||||
trainer, _, reserved_tokens = self._prepare_trainer(
|
||||
vocab_size=vocab_size,
|
||||
min_freq=min_freq,
|
||||
reserved_token_size=reserved_token_size,
|
||||
)
|
||||
self._tokenizer.train_from_iterator(iterator=iterator, trainer=trainer)
|
||||
self._tokenizer.add_special_tokens(self._special_tokens + reserved_tokens)
|
||||
|
||||
def save(self, path):
|
||||
self._tokenizer.save(path)
|
||||
|
||||
def load(self, path):
|
||||
self._tokenizer = Tokenizer.from_file(path)
|
||||
|
||||
def encode(
|
||||
self,
|
||||
tokens: Union[str, List[str]],
|
||||
out_ids: bool = True,
|
||||
add_special_tokens: bool = False,
|
||||
) -> List:
|
||||
if isinstance(tokens, str):
|
||||
encoded: Encoding = self._tokenizer.encode(
|
||||
tokens, add_special_tokens=add_special_tokens
|
||||
)
|
||||
return encoded.ids if out_ids else encoded.tokens
|
||||
elif isinstance(tokens, list):
|
||||
encoded_list: List[Encoding] = self._tokenizer.encode_batch(
|
||||
tokens, add_special_tokens=add_special_tokens
|
||||
)
|
||||
return [
|
||||
encoded.ids if out_ids else encoded.tokens for encoded in encoded_list
|
||||
]
|
||||
|
||||
def decode(self, tokens: List[int], skip_special_tokens: bool = True) -> str:
|
||||
return self._tokenizer.decode(tokens, skip_special_tokens=skip_special_tokens)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self._tokenizer.get_vocab_size()
|
||||
|
||||
@property
|
||||
def stop_ids(self) -> List[int]:
|
||||
stop_token = self._control_tokens + self._special_tokens
|
||||
stop_ids = [self._tokenizer.token_to_id(token) for token in stop_token]
|
||||
return stop_ids
|
||||
|
||||
@property
|
||||
def bos_id(self) -> int:
|
||||
return self._tokenizer.token_to_id("<bos>")
|
||||
|
||||
@property
|
||||
def eos_id(self) -> int:
|
||||
return self._tokenizer.token_to_id("<eos>")
|
||||
|
||||
@property
|
||||
def pad_id(self) -> int:
|
||||
return self._tokenizer.token_to_id("<pad>")
|
||||
@@ -0,0 +1,28 @@
|
||||
from astrai.inference.core import (
|
||||
disable_random_init,
|
||||
GeneratorCore,
|
||||
EmbeddingEncoderCore,
|
||||
KVCacheManager,
|
||||
)
|
||||
|
||||
from astrai.inference.generator import (
|
||||
GenerationRequest,
|
||||
LoopGenerator,
|
||||
StreamGenerator,
|
||||
BatchGenerator,
|
||||
EmbeddingEncoder,
|
||||
GeneratorFactory,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"disable_random_init",
|
||||
"GeneratorCore",
|
||||
"EmbeddingEncoderCore",
|
||||
"KVCacheManager",
|
||||
"GenerationRequest",
|
||||
"LoopGenerator",
|
||||
"StreamGenerator",
|
||||
"BatchGenerator",
|
||||
"EmbeddingEncoder",
|
||||
"GeneratorFactory",
|
||||
]
|
||||
@@ -0,0 +1,298 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from torch import Tensor
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Callable, List, Tuple, Union, Optional, Self
|
||||
from astrai.config import ModelParameter, ModelConfig
|
||||
|
||||
|
||||
def apply_sampling_strategies(
|
||||
logits: Tensor,
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
filter_value: float = -float("inf"),
|
||||
) -> Tensor:
|
||||
"""
|
||||
Apply sampling strategies to the logits tensor.
|
||||
|
||||
Args:
|
||||
logits (Tensor): The logits tensor.
|
||||
temperature (float): The temperature parameter.
|
||||
top_k (int): The top-k parameter.
|
||||
top_p (float): The top-p parameter.
|
||||
filter_value (float, optional): The filter value. Defaults to -float("inf").
|
||||
|
||||
Returns:
|
||||
Tensor: The sampled logits tensor.
|
||||
|
||||
"""
|
||||
|
||||
if temperature != 1.0:
|
||||
logits = logits / temperature
|
||||
|
||||
if top_k > 0:
|
||||
top_k = min(top_k, logits.size(-1))
|
||||
indices_to_remove = logits < torch.topk(logits, top_k, dim=-1)[0][..., -1, None]
|
||||
logits[indices_to_remove] = filter_value
|
||||
|
||||
if top_p < 1.0:
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
|
||||
cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
|
||||
|
||||
sorted_indices_to_remove = cumulative_probs > top_p
|
||||
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
|
||||
sorted_indices_to_remove[..., 0] = 0
|
||||
|
||||
indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
|
||||
indices_to_remove.scatter_(
|
||||
dim=1, index=sorted_indices, src=sorted_indices_to_remove
|
||||
)
|
||||
|
||||
logits[indices_to_remove] = filter_value
|
||||
|
||||
return logits
|
||||
|
||||
|
||||
@contextmanager
|
||||
def disable_random_init():
|
||||
init_functions = [
|
||||
"xavier_normal_",
|
||||
"xavier_uniform_",
|
||||
"kaiming_normal_",
|
||||
"kaiming_uniform_",
|
||||
"zeros_",
|
||||
"ones_",
|
||||
"constant_",
|
||||
"normal_",
|
||||
"uniform_",
|
||||
]
|
||||
original_funcs = {}
|
||||
for name in init_functions:
|
||||
if hasattr(nn.init, name):
|
||||
original_funcs[name] = getattr(nn.init, name)
|
||||
setattr(nn.init, name, lambda *args, **kwargs: None)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
for name, orig_func in original_funcs.items():
|
||||
setattr(nn.init, name, orig_func)
|
||||
|
||||
|
||||
class GeneratorCore:
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
self.model = parameter.model
|
||||
self.tokenizer = parameter.tokenizer
|
||||
self.config = parameter.config
|
||||
|
||||
def generate_iterator(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
attn_mask: Optional[Tensor] = None,
|
||||
kv_caches: Optional[List[Tuple[Tensor, Tensor]]] = None,
|
||||
start_pos: int = 0,
|
||||
) -> Tuple[Tensor, int]:
|
||||
|
||||
with torch.inference_mode():
|
||||
outputs = self.model(input_ids, attn_mask, kv_caches, start_pos)
|
||||
logits = outputs["logits"][:, -1, :]
|
||||
cache_increase = input_ids.size(-1)
|
||||
|
||||
logits = apply_sampling_strategies(logits, temperature, top_k, top_p)
|
||||
probs = torch.softmax(logits, dim=-1)
|
||||
next_token_id = torch.multinomial(probs, num_samples=1)
|
||||
|
||||
return next_token_id, cache_increase
|
||||
|
||||
def generate_loop(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
ids: List[int],
|
||||
temperature: float,
|
||||
top_k: int,
|
||||
top_p: float,
|
||||
attn_mask: Optional[Tensor] = None,
|
||||
kv_caches: Optional[List[Tuple[Tensor, Tensor]]] = None,
|
||||
start_pos: int = 0,
|
||||
callback: Optional[Callable[..., Any]] = None,
|
||||
) -> List[int]:
|
||||
cur_cache_pos = start_pos
|
||||
|
||||
for _ in range(len(ids), self.config.max_len):
|
||||
next_token_id, cache_increase = self.generate_iterator(
|
||||
input_ids,
|
||||
temperature,
|
||||
top_k,
|
||||
top_p,
|
||||
attn_mask,
|
||||
kv_caches,
|
||||
cur_cache_pos,
|
||||
)
|
||||
|
||||
input_ids = next_token_id
|
||||
ids.append(next_token_id.item())
|
||||
cur_cache_pos += cache_increase
|
||||
|
||||
if callback:
|
||||
callback(next_token_id.item(), ids.copy())
|
||||
|
||||
if next_token_id.item() in self.tokenizer.stop_ids:
|
||||
break
|
||||
|
||||
return ids
|
||||
|
||||
def to(self, *args, **kargs) -> Self:
|
||||
self.model.to(*args, **kargs)
|
||||
return self
|
||||
|
||||
|
||||
class EmbeddingEncoderCore:
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
self.model = parameter.model
|
||||
self.tokenizer = parameter.tokenizer
|
||||
self.config = parameter.config
|
||||
|
||||
def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
|
||||
with_batch = isinstance(sentence, list)
|
||||
ids = self.tokenizer.encode(sentence)
|
||||
batch_ids = ids if with_batch else [ids]
|
||||
max_model_len = self.config.max_len
|
||||
|
||||
all_fragments = []
|
||||
fragment_origin_idx = []
|
||||
|
||||
for i, seq in enumerate(batch_ids):
|
||||
if len(seq) > max_model_len:
|
||||
fragments = [
|
||||
seq[j : j + max_model_len]
|
||||
for j in range(0, len(seq), max_model_len)
|
||||
]
|
||||
all_fragments.extend(fragments)
|
||||
fragment_origin_idx.extend([i] * len(fragments))
|
||||
else:
|
||||
all_fragments.append(seq)
|
||||
fragment_origin_idx.append(i)
|
||||
|
||||
# if empty fragments
|
||||
if not all_fragments or not ids:
|
||||
return [] if with_batch else torch.tensor([])
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
max_len = min(max(len(seq) for seq in all_fragments), max_model_len)
|
||||
|
||||
padded_ids = []
|
||||
masks = []
|
||||
for seq in all_fragments:
|
||||
pad_len = max_len - len(seq)
|
||||
padded_seq = seq + [self.tokenizer.pad_id] * pad_len
|
||||
mask = [token_id != self.tokenizer.pad_id for token_id in padded_seq]
|
||||
padded_ids.append(padded_seq)
|
||||
masks.append(mask)
|
||||
|
||||
input_tensor = torch.tensor(padded_ids, device=device, dtype=torch.long)
|
||||
seq_mask = torch.tensor(masks, device=device, dtype=torch.bool)
|
||||
|
||||
with torch.inference_mode():
|
||||
outputs = self.model(input_tensor, seq_mask)["hidden_states"]
|
||||
# [num_fragments, seq_len, hidden_size]
|
||||
fragment_embs = torch.mul(outputs, seq_mask.unsqueeze(-1))
|
||||
|
||||
sentence_embs: List[Tensor] = []
|
||||
for i in range(len(batch_ids)):
|
||||
indices = [
|
||||
idx for idx, orig_idx in enumerate(fragment_origin_idx) if orig_idx == i
|
||||
]
|
||||
if indices:
|
||||
sum_frags = torch.sum(
|
||||
fragment_embs[indices, :, :], dim=1
|
||||
) # [frags, hidden_size]
|
||||
length = torch.sum(seq_mask[indices, :], dim=1).unsqueeze(
|
||||
1
|
||||
) # [frags, 1]
|
||||
emb = torch.sum(sum_frags / length, dim=0) # [frags, hidden_size]
|
||||
sentence_embs.append(emb.flatten())
|
||||
|
||||
if with_batch:
|
||||
return [emb.flatten() for emb in sentence_embs]
|
||||
else:
|
||||
return sentence_embs[0].flatten()
|
||||
|
||||
def to(self, *args, **kargs) -> Self:
|
||||
self.model.to(*args, **kargs)
|
||||
return self
|
||||
|
||||
|
||||
class KVCacheManager:
|
||||
def __init__(
|
||||
self,
|
||||
config: ModelConfig,
|
||||
batch_size: int,
|
||||
device: torch.device = "cuda",
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
):
|
||||
self.batch_size = batch_size
|
||||
self.device = device
|
||||
self.dtype = dtype
|
||||
self.num_layers = config.n_layers
|
||||
self.max_len = config.max_len
|
||||
self.num_heads = config.n_kv_heads
|
||||
self.head_dim = config.dim // config.n_heads
|
||||
|
||||
self._kv_cache: Tuple[Tensor, Tensor] = None
|
||||
self._seq_mask: Tensor = None
|
||||
self._initialize()
|
||||
|
||||
def _initialize(self):
|
||||
k_cache = torch.empty(
|
||||
(
|
||||
self.batch_size,
|
||||
self.max_len,
|
||||
self.num_layers,
|
||||
self.num_heads,
|
||||
self.head_dim,
|
||||
),
|
||||
device=self.device,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
v_cache = torch.empty(
|
||||
(
|
||||
self.batch_size,
|
||||
self.max_len,
|
||||
self.num_layers,
|
||||
self.num_heads,
|
||||
self.head_dim,
|
||||
),
|
||||
device=self.device,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
self._kv_cache = (k_cache, v_cache)
|
||||
self._seq_mask = torch.ones(
|
||||
(self.batch_size, self.max_len), device=self.device, dtype=torch.bool
|
||||
)
|
||||
|
||||
def update(self, active_mask: Tensor):
|
||||
k_cache, v_cache = self._kv_cache
|
||||
self._kv_cache = (k_cache[active_mask], v_cache[active_mask])
|
||||
self._seq_mask = self._seq_mask[active_mask]
|
||||
|
||||
def reset(self, full_reset=False):
|
||||
if full_reset:
|
||||
self._kv_cache = None
|
||||
self._seq_mask = None
|
||||
else:
|
||||
self._initialize()
|
||||
|
||||
def set_seq_mask(self, input_ids: Tensor, pad_id: int):
|
||||
batch_size, seq_len = input_ids.shape
|
||||
bool_mask = input_ids != pad_id
|
||||
self._seq_mask[:batch_size, :seq_len] = bool_mask
|
||||
|
||||
def get_kvcache(self) -> Tuple[Tensor, Tensor]:
|
||||
return self._kv_cache
|
||||
|
||||
def get_seq_mask(self) -> Tensor:
|
||||
return self._seq_mask
|
||||
@@ -0,0 +1,310 @@
|
||||
import torch
|
||||
from dataclasses import dataclass
|
||||
from torch import Tensor
|
||||
from typing import List, Tuple, Union, Optional, Generator
|
||||
from astrai.inference.core import GeneratorCore, EmbeddingEncoderCore, KVCacheManager
|
||||
from astrai.config.param_config import ModelParameter
|
||||
|
||||
|
||||
HistoryType = List[Tuple[str, str]]
|
||||
|
||||
|
||||
def build_prompt(
|
||||
query: str,
|
||||
system_prompt: Optional[str] = None,
|
||||
history: Optional[HistoryType] = None,
|
||||
) -> str:
|
||||
"""
|
||||
Build prompt in ChatML format for query and history.
|
||||
|
||||
Args:
|
||||
query (str): query string.
|
||||
system_prompt (Optional[str]): system prompt string.
|
||||
history (Optional[HistoryType]): history list of query and response.
|
||||
|
||||
Returns:
|
||||
str: prompt string in ChatML format.
|
||||
"""
|
||||
result = ""
|
||||
|
||||
if system_prompt:
|
||||
result += f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
|
||||
|
||||
# (convert tuple format to ChatML)
|
||||
if history:
|
||||
for user_msg, assistant_msg in history:
|
||||
result += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
|
||||
result += f"<|im_start|>assistant\n{assistant_msg}<|im_end|>\n"
|
||||
|
||||
result += f"<|im_start|>user\n{query}<|im_end|>\n"
|
||||
result += "<|im_start|>assistant\n"
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def pad_sequence(ids_list: List[List[int]], pad_id: int) -> Tuple[List[List[int]], int]:
|
||||
"""
|
||||
Pad a list of sequences to a fixed length.
|
||||
|
||||
Args:
|
||||
ids_list (List[List[int]]): A list of sequences.
|
||||
max_ids_len (int): The maximum length of sequences.
|
||||
pad_id (int): The id to pad sequences.
|
||||
|
||||
Returns:
|
||||
List[List[int]]: A list of padded sequences.
|
||||
|
||||
"""
|
||||
max_ids_len = max(len(ids) for ids in ids_list)
|
||||
new_ids_list = []
|
||||
for ids in ids_list:
|
||||
pad_len = max_ids_len - len(ids)
|
||||
padded_seq = [pad_id] * pad_len + ids
|
||||
new_ids_list.append(padded_seq)
|
||||
|
||||
return new_ids_list, max_ids_len
|
||||
|
||||
|
||||
@dataclass
|
||||
class GenerationRequest:
|
||||
"""
|
||||
Request parameters for text generation.
|
||||
|
||||
Attributes:
|
||||
top_k: Top-k sampling parameter.
|
||||
top_p: Top-p (nucleus) sampling parameter.
|
||||
temperature: Sampling temperature.
|
||||
max_len: Maximum generation length.
|
||||
query: Input query (string or list of strings for batch).
|
||||
history: Conversation history.
|
||||
system_prompt: System prompt for the conversation.
|
||||
stream: Whether to use streaming generation.
|
||||
"""
|
||||
|
||||
top_k: int
|
||||
top_p: float
|
||||
temperature: float
|
||||
max_len: int
|
||||
|
||||
query: Union[str, List[str]]
|
||||
history: Optional[Union[HistoryType, List[HistoryType]]] = None
|
||||
system_prompt: Optional[str] = None
|
||||
stream: bool = False
|
||||
|
||||
def __post_init__(self):
|
||||
if not isinstance(self.top_k, int) or self.top_k < 0:
|
||||
raise ValueError("top_k must be a non-negative integer")
|
||||
if not isinstance(self.top_p, float) or self.top_p < 0.0 or self.top_p > 1.0:
|
||||
raise ValueError("top_p must be a float between 0.0 and 1.0")
|
||||
if not isinstance(self.temperature, float) or self.temperature < 0.0:
|
||||
raise ValueError("temperature must be a non-negative float")
|
||||
|
||||
|
||||
class LoopGenerator(GeneratorCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def generate(self, request: GenerationRequest) -> str:
|
||||
device = next(self.model.parameters()).device
|
||||
cache_manager = KVCacheManager(self.config, 1, device=device)
|
||||
|
||||
prompt = build_prompt(request.query, request.history)
|
||||
ids = self.tokenizer.encode(prompt)
|
||||
input_ids = torch.tensor([ids], device=device, dtype=torch.long)
|
||||
|
||||
start_cache_pos = len(ids)
|
||||
self.model.eval()
|
||||
kv_caches = cache_manager.get_kvcache()
|
||||
|
||||
ids = self.generate_loop(
|
||||
input_ids,
|
||||
ids,
|
||||
request.temperature,
|
||||
request.top_k,
|
||||
request.top_p,
|
||||
kv_caches=kv_caches,
|
||||
)
|
||||
response = self.tokenizer.decode(ids[start_cache_pos:])
|
||||
|
||||
return response
|
||||
|
||||
|
||||
class StreamGenerator(GeneratorCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def generate(self, request: GenerationRequest) -> Generator[str, None, None]:
|
||||
device = next(self.model.parameters()).device
|
||||
cache_manager = KVCacheManager(self.config, 1, device=device)
|
||||
|
||||
prompt = build_prompt(request.query, request.history)
|
||||
ids = self.tokenizer.encode(prompt)
|
||||
input_ids = torch.tensor([ids], device=device, dtype=torch.long)
|
||||
|
||||
start_cache_pos = len(ids)
|
||||
cur_cache_pos = 0
|
||||
self.model.eval()
|
||||
kv_caches = cache_manager.get_kvcache()
|
||||
|
||||
for _ in range(len(ids), self.config.max_len):
|
||||
next_token_id, cache_increase = self.generate_iterator(
|
||||
input_ids,
|
||||
request.temperature,
|
||||
request.top_k,
|
||||
request.top_p,
|
||||
kv_caches=kv_caches,
|
||||
start_pos=cur_cache_pos,
|
||||
)
|
||||
|
||||
input_ids = next_token_id
|
||||
ids.append(next_token_id.item())
|
||||
cur_cache_pos += cache_increase
|
||||
|
||||
response = self.tokenizer.decode(ids[start_cache_pos:])
|
||||
yield response
|
||||
|
||||
if next_token_id.item() in self.tokenizer.stop_ids:
|
||||
yield response + "\n"
|
||||
break
|
||||
|
||||
|
||||
class BatchGenerator(GeneratorCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def generate(self, request: GenerationRequest) -> List[str]:
|
||||
batch_size = len(request.query)
|
||||
if request.history is None:
|
||||
request.history = [[] for _ in range(batch_size)]
|
||||
|
||||
prompts = [
|
||||
build_prompt(query, history)
|
||||
for query, history in zip(request.query, request.history)
|
||||
]
|
||||
|
||||
ids_list = [self.tokenizer.encode(prompt) for prompt in prompts]
|
||||
ids_list, max_ids_len = pad_sequence(ids_list, self.tokenizer.pad_id)
|
||||
|
||||
device = next(self.model.parameters()).device
|
||||
cache_manager = KVCacheManager(self.config, batch_size, device=device)
|
||||
|
||||
input_tensor = torch.tensor(ids_list, device=device, dtype=torch.long)
|
||||
cache_manager.set_seq_mask(input_tensor, self.tokenizer.pad_id)
|
||||
activate_task_mask = [True] * batch_size
|
||||
|
||||
start_cache_pos = max_ids_len
|
||||
cur_cache_pos = 0
|
||||
|
||||
while max_ids_len < self.config.max_len and sum(activate_task_mask) != 0:
|
||||
kv_caches = cache_manager.get_kvcache()
|
||||
attn_mask = cache_manager.get_seq_mask()
|
||||
|
||||
next_token_id, cache_increase = self.generate_iterator(
|
||||
input_tensor,
|
||||
request.temperature,
|
||||
request.top_k,
|
||||
request.top_p,
|
||||
attn_mask=attn_mask,
|
||||
kv_caches=kv_caches,
|
||||
start_pos=cur_cache_pos,
|
||||
)
|
||||
|
||||
cur_cache_pos += cache_increase
|
||||
active_mask = []
|
||||
c_ids = 0
|
||||
|
||||
for i in range(batch_size):
|
||||
if activate_task_mask[i]:
|
||||
token = next_token_id[c_ids, :].item()
|
||||
ids_list[i].append(token)
|
||||
c_ids += 1
|
||||
|
||||
is_active = token not in self.tokenizer.stop_ids
|
||||
activate_task_mask[i] = is_active
|
||||
active_mask.append(is_active)
|
||||
|
||||
active_mask = torch.tensor(active_mask, device=device, dtype=torch.bool)
|
||||
cache_manager.update(active_mask)
|
||||
input_tensor = next_token_id[active_mask, :]
|
||||
|
||||
max_ids_len += 1
|
||||
|
||||
responses = [str()] * batch_size
|
||||
for i in range(batch_size):
|
||||
responses[i] = self.tokenizer.decode(ids_list[i][start_cache_pos:])
|
||||
request.history[i].append((request.query[i], responses[i]))
|
||||
|
||||
return responses
|
||||
|
||||
|
||||
class EmbeddingEncoder(EmbeddingEncoderCore):
|
||||
def __init__(self, parameter: ModelParameter):
|
||||
super().__init__(parameter)
|
||||
|
||||
def encode(self, sentence: Union[str, List[str]]) -> Union[Tensor, List[Tensor]]:
|
||||
return super().encode(sentence)
|
||||
|
||||
|
||||
class GeneratorFactory:
|
||||
"""Factory class for creating generator instances.
|
||||
|
||||
Provides smart generator selection based on request characteristics:
|
||||
- Streaming: Use StreamGenerator for streaming output
|
||||
- Batch: Use BatchGenerator when query is a list
|
||||
- Single: Use LoopGenerator for single query non-streaming
|
||||
|
||||
Example usage:
|
||||
generator = GeneratorFactory.create_generator(parameter, request)
|
||||
result = generator.generate(request)
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def create_generator(
|
||||
parameter: ModelParameter, request: GenerationRequest
|
||||
) -> GeneratorCore:
|
||||
"""Create a generator based on request characteristics.
|
||||
|
||||
Args:
|
||||
parameter: Model parameters containing model, tokenizer, config
|
||||
request: Generation request with query, options, etc.
|
||||
|
||||
Returns:
|
||||
Appropriate GeneratorCore subclass instance
|
||||
"""
|
||||
# Streaming generation: check stream field first
|
||||
if request.stream:
|
||||
return StreamGenerator(parameter)
|
||||
|
||||
# Batch generation: query is a list of strings
|
||||
if isinstance(request.query, list):
|
||||
return BatchGenerator(parameter)
|
||||
|
||||
# Default: single query non-streaming
|
||||
return LoopGenerator(parameter)
|
||||
|
||||
@staticmethod
|
||||
def create_encoder(parameter: ModelParameter) -> EmbeddingEncoderCore:
|
||||
"""Create an embedding encoder instance.
|
||||
|
||||
Args:
|
||||
parameter: Model parameters
|
||||
|
||||
Returns:
|
||||
EmbeddingEncoderCore instance
|
||||
"""
|
||||
return EmbeddingEncoder(parameter)
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls, parameter: ModelParameter, request: GenerationRequest
|
||||
) -> GeneratorCore:
|
||||
"""Convenience method that delegates to create_generator.
|
||||
|
||||
Args:
|
||||
parameter: Model parameters
|
||||
request: Generation request
|
||||
|
||||
Returns:
|
||||
Generator instance
|
||||
"""
|
||||
return cls.create_generator(parameter, request)
|
||||
@@ -0,0 +1,10 @@
|
||||
from astrai.model.module import (
|
||||
Linear,
|
||||
RMSNorm,
|
||||
MLP,
|
||||
GQA,
|
||||
DecoderBlock,
|
||||
)
|
||||
from astrai.model.transformer import Transformer
|
||||
|
||||
__all__ = ["Linear", "RMSNorm", "MLP", "GQA", "DecoderBlock", "Transformer"]
|
||||
@@ -0,0 +1,381 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from torch import Tensor
|
||||
from typing import Optional, Tuple
|
||||
|
||||
|
||||
def repeat_kv(x: Tensor, n_rep: int) -> Tensor:
|
||||
"""
|
||||
Repeat k times along the dimension for attention heads.
|
||||
Args:
|
||||
x (Tensor): The input tensor.
|
||||
n_rep (int): The number of repetitions.
|
||||
Returns:
|
||||
Tensor: The repeated tensor.
|
||||
"""
|
||||
|
||||
bs, slen, n_heads, head_dim = x.shape
|
||||
if n_rep == 1:
|
||||
return x
|
||||
return (
|
||||
x[:, :, :, None, :]
|
||||
.expand(bs, slen, n_heads, n_rep, head_dim)
|
||||
.reshape(bs, slen, n_heads * n_rep, head_dim)
|
||||
)
|
||||
|
||||
|
||||
def get_rotary_emb(
|
||||
dim: int,
|
||||
max_len: int,
|
||||
base: float = 10000,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""
|
||||
Get the rotary embedding for the given dimension and maximum length.
|
||||
Args:
|
||||
dim (int): The dimension of the input.
|
||||
max_len (int): The maximum length of the input.
|
||||
base (float, optional): The base for the frequency. Defaults to 10000.
|
||||
Returns:
|
||||
Tensor: The rotary embedding tensor.
|
||||
"""
|
||||
|
||||
theta = base ** (-torch.arange(0, dim, 2, dtype=torch.float64) / dim)
|
||||
t = torch.arange(0, max_len, dtype=torch.float64)
|
||||
freqs = torch.outer(t, theta)
|
||||
|
||||
return torch.cos(freqs).float(), torch.sin(freqs).float()
|
||||
|
||||
|
||||
def apply_rotary_emb(x: torch.Tensor, rotary_emb: Tuple[Tensor, Tensor]) -> Tensor:
|
||||
"""
|
||||
Apply rotary embedding to the input tensor using cos/sin form.
|
||||
Args:
|
||||
x (Tensor): The input tensor (shape [..., seq_len, dim]).
|
||||
rotary_emb (Tuple[Tensor, Tensor]): The rotary embedding (shape [seq_len, dim//2]).
|
||||
Returns:
|
||||
Tensor: The output tensor (rotated, same shape as input).
|
||||
"""
|
||||
|
||||
dtype = x.dtype
|
||||
cos, sin = rotary_emb
|
||||
|
||||
cos = cos.unsqueeze(0).unsqueeze(2) # [1, seq_len, 1, dim//2]
|
||||
sin = sin.unsqueeze(0).unsqueeze(2) # [1, seq_len, 1, dim//2]
|
||||
|
||||
x_real = x[..., 0::2] # [batch, seq_len, dim//2]
|
||||
x_imag = x[..., 1::2] # [batch, seq_len, dim//2]
|
||||
|
||||
x_real_rot = x_real * cos - x_imag * sin
|
||||
x_imag_rot = x_real * sin + x_imag * cos
|
||||
|
||||
x_out = torch.stack([x_real_rot, x_imag_rot], dim=-1) # [batch, seq_len, dim//2, 2]
|
||||
x_out = x_out.view(*x_out.shape[:-2], -1) # [batch, seq_len, dim]
|
||||
|
||||
return x_out.to(dtype)
|
||||
|
||||
|
||||
class RotaryEmbedding(nn.Module):
|
||||
def __init__(self, dim: int, max_len: int, base: int = 10000):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.max_len = max_len
|
||||
self.base = base
|
||||
self.max_len_cached = None
|
||||
self._set_rotary_buffer(self.max_len)
|
||||
|
||||
def _set_rotary_buffer(self, max_len: int):
|
||||
cos_cached, sin_cached = get_rotary_emb(self.dim, max_len, self.base)
|
||||
self.register_buffer("cos_cached", cos_cached, persistent=False)
|
||||
self.register_buffer("sin_cached", sin_cached, persistent=False)
|
||||
self.max_len_cached = max_len
|
||||
|
||||
def forward(self, x: Tensor, start_pos: int = 0) -> Tuple[Tensor, Tensor]:
|
||||
seq_len = x.size(1)
|
||||
|
||||
if self.max_len_cached < seq_len + start_pos:
|
||||
self._set_rotary_buffer(seq_len + start_pos)
|
||||
|
||||
cos = self.cos_cached[start_pos : start_pos + seq_len]
|
||||
sin = self.sin_cached[start_pos : start_pos + seq_len]
|
||||
|
||||
return (cos, sin)
|
||||
|
||||
|
||||
class Linear(nn.Module):
|
||||
def __init__(self, in_dim: int, out_dim: int, bias: bool = False):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty((out_dim, in_dim)))
|
||||
self.bias = nn.Parameter(torch.zeros(out_dim)) if bias else None
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return F.linear(x, self.weight, self.bias)
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim, norm_eps):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
self.normalized_shape = (dim,)
|
||||
self.norm_eps = norm_eps
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
rms = F.rms_norm(x.float(), self.normalized_shape, self.weight, self.norm_eps)
|
||||
return rms.to(x.dtype)
|
||||
|
||||
|
||||
class MLP(nn.Module):
|
||||
def __init__(self, dim: int, dim_feed_forward: int):
|
||||
super().__init__()
|
||||
self.up = Linear(dim, dim_feed_forward)
|
||||
self.gate = Linear(dim, dim_feed_forward)
|
||||
self.down = Linear(dim_feed_forward, dim)
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
gated = self.up(x) * F.silu(self.gate(x))
|
||||
out = self.down(gated)
|
||||
return out
|
||||
|
||||
|
||||
class GQA(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
n_heads: int,
|
||||
n_kv_heads: int,
|
||||
use_qk_norm: bool,
|
||||
norm_eps: float,
|
||||
use_gated_attention: bool,
|
||||
layer_id: int,
|
||||
):
|
||||
super().__init__()
|
||||
assert dim % n_heads == 0
|
||||
assert n_heads % n_kv_heads == 0
|
||||
|
||||
self.head_dim = dim // n_heads
|
||||
self.layer_id = layer_id
|
||||
self.dim = dim
|
||||
self.n_heads = n_heads
|
||||
self.n_kv_heads = n_kv_heads
|
||||
self.n_rep = n_heads // n_kv_heads
|
||||
self.use_qk_norm = use_qk_norm
|
||||
self.use_gated_attention = use_gated_attention
|
||||
|
||||
self.q_proj = Linear(dim, n_heads * self.head_dim)
|
||||
self.k_proj = Linear(dim, n_kv_heads * self.head_dim)
|
||||
self.v_proj = Linear(dim, n_kv_heads * self.head_dim)
|
||||
self.o_proj = Linear(dim, dim)
|
||||
|
||||
if self.use_qk_norm:
|
||||
self.q_norm = RMSNorm(self.head_dim, norm_eps)
|
||||
self.k_norm = RMSNorm(self.head_dim, norm_eps)
|
||||
|
||||
if self.use_gated_attention:
|
||||
self.gate = Linear(dim, dim)
|
||||
|
||||
def _split_heads(self, x: Tensor, n_heads) -> Tensor:
|
||||
batch_size, seq_len, _ = x.shape
|
||||
x = x.reshape(batch_size, seq_len, n_heads, self.head_dim)
|
||||
return x
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
rotary_emb: Tuple[Tensor, Tensor],
|
||||
mask: Tensor = None,
|
||||
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
start_pos: int = 0,
|
||||
) -> Tensor:
|
||||
bsz, seq_len, _ = x.size()
|
||||
is_causal = mask is None
|
||||
|
||||
# x(bsz, seq_len, n_heads * head_dim) -> (bsz, seq_len, n_heads, head_dim)
|
||||
q = self._split_heads(self.q_proj(x), self.n_heads)
|
||||
k = self._split_heads(self.k_proj(x), self.n_kv_heads)
|
||||
v = self._split_heads(self.v_proj(x), self.n_kv_heads)
|
||||
q, k = apply_rotary_emb(q, rotary_emb), apply_rotary_emb(k, rotary_emb)
|
||||
|
||||
if self.use_qk_norm:
|
||||
q, k = self.q_norm(q), self.k_norm(k)
|
||||
|
||||
if kv_cache is not None:
|
||||
k_cache, v_cache = kv_cache
|
||||
|
||||
# copy to cache
|
||||
k_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = k
|
||||
v_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = v
|
||||
|
||||
# get cache
|
||||
k = k_cache[:bsz, : start_pos + seq_len, self.layer_id]
|
||||
v = v_cache[:bsz, : start_pos + seq_len, self.layer_id]
|
||||
|
||||
k, v = repeat_kv(k, self.n_rep), repeat_kv(v, self.n_rep)
|
||||
|
||||
# (bsz, seq_len, n_heads, head_dim) -> (bsz, n_heads, seq_len, head_dim)
|
||||
q, k, v = q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3)
|
||||
# (bsz, n_heads, seq_len, head_dim) - > (bsz, seq_len, n_heads*head_dim)
|
||||
sdqa_out = (
|
||||
F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
|
||||
.permute(0, 2, 1, 3)
|
||||
.contiguous()
|
||||
.flatten(2)
|
||||
)
|
||||
|
||||
if self.use_gated_attention:
|
||||
sdqa_out = sdqa_out * F.sigmoid(self.gate(x))
|
||||
|
||||
out = self.o_proj(sdqa_out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class MLA(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
n_heads: int,
|
||||
n_kv_heads: int,
|
||||
kv_lora_rank: int,
|
||||
qk_nope_head_dim: int,
|
||||
qk_rope_head_dim: int,
|
||||
norm_eps: float,
|
||||
use_gated_attention: bool,
|
||||
layer_id: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.n_heads = n_heads
|
||||
self.n_kv_heads = n_kv_heads
|
||||
self.kv_lora_rank = kv_lora_rank
|
||||
self.qk_nope_head_dim = qk_nope_head_dim
|
||||
self.qk_rope_head_dim = qk_rope_head_dim
|
||||
self.head_dim = qk_nope_head_dim + qk_rope_head_dim
|
||||
self.layer_id = layer_id
|
||||
self.n_rep = n_heads // n_kv_heads
|
||||
self.use_gated_attention = use_gated_attention
|
||||
|
||||
self.q_proj = Linear(dim, n_heads * self.head_dim, bias=False)
|
||||
self.kv_a_proj = Linear(dim, kv_lora_rank, bias=False)
|
||||
self.kv_norm = RMSNorm(kv_lora_rank, eps=norm_eps)
|
||||
|
||||
# KV (k_nope, k_rope, v)
|
||||
self.kv_b_proj = Linear(
|
||||
kv_lora_rank,
|
||||
n_kv_heads * (self.head_dim + qk_rope_head_dim + self.head_dim),
|
||||
)
|
||||
|
||||
self.o_proj = Linear(dim, dim, bias=False)
|
||||
|
||||
if use_gated_attention:
|
||||
self.gate = Linear(dim, dim, bias=False)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
rotary_emb: Tuple[Tensor, Tensor],
|
||||
mask: Tensor = None,
|
||||
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
start_pos: int = 0,
|
||||
) -> Tensor:
|
||||
bsz, seq_len, _ = x.size()
|
||||
is_causal = mask is None
|
||||
|
||||
q = self.q_proj(x)
|
||||
q = q.view(bsz, seq_len, self.n_heads, self.head_dim)
|
||||
|
||||
kv_compressed = self.kv_a_proj(x)
|
||||
kv_compressed = self.kv_norm(kv_compressed)
|
||||
|
||||
kv = self.kv_b_proj(kv_compressed)
|
||||
kv = kv.view(bsz, seq_len, self.n_kv_heads, -1)
|
||||
|
||||
k_nope, k_rope, v = torch.split(
|
||||
kv, [self.qk_nope_head_dim, self.qk_rope_head_dim, self.head_dim], dim=-1
|
||||
)
|
||||
|
||||
q_nope, q_rope = (
|
||||
q[..., : self.qk_nope_head_dim],
|
||||
q[..., self.qk_rope_head_dim :],
|
||||
)
|
||||
q_rope = apply_rotary_emb(q_rope, rotary_emb)
|
||||
k_rope = apply_rotary_emb(k_rope, rotary_emb)
|
||||
|
||||
q = torch.cat([q_nope, q_rope], dim=-1)
|
||||
k = torch.cat([k_nope, k_rope], dim=-1)
|
||||
|
||||
if kv_cache is not None:
|
||||
k_cache, v_cache = kv_cache
|
||||
k_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = k
|
||||
v_cache[:bsz, start_pos : start_pos + seq_len, self.layer_id] = v
|
||||
k = k_cache[:bsz, : start_pos + seq_len, self.layer_id]
|
||||
v = v_cache[:bsz, : start_pos + seq_len, self.layer_id]
|
||||
|
||||
q = q.permute(0, 2, 1, 3)
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
v = v.permute(0, 2, 1, 3)
|
||||
|
||||
attn_out = F.scaled_dot_product_attention(q, k, v, mask, is_causal=is_causal)
|
||||
attn_out = attn_out.permute(0, 2, 1, 3).contiguous().flatten(2)
|
||||
|
||||
if self.use_gated_attention:
|
||||
attn_out = attn_out * F.sigmoid(self.gate(x))
|
||||
|
||||
out = self.o_proj(attn_out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class DecoderBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
n_heads: int,
|
||||
dim_ffn: int,
|
||||
n_kv_heads: int,
|
||||
norm_eps: int,
|
||||
use_qk_norm: bool,
|
||||
use_gated_attention: bool,
|
||||
layer_id: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.attention = GQA(
|
||||
dim,
|
||||
n_heads,
|
||||
n_kv_heads,
|
||||
use_qk_norm,
|
||||
norm_eps,
|
||||
use_gated_attention,
|
||||
layer_id,
|
||||
)
|
||||
self.input_norm = RMSNorm(dim, norm_eps)
|
||||
self.mlp = MLP(dim, dim_ffn)
|
||||
self.post_attention_norm = RMSNorm(dim, norm_eps)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
rotary_emb: Tuple[Tensor, Tensor],
|
||||
attention_mask: Optional[Tensor] = None,
|
||||
kv_cache: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
start_pos: int = 0,
|
||||
) -> Tensor:
|
||||
# attention
|
||||
attn_output = self.attention(
|
||||
self.input_norm(x), rotary_emb, attention_mask, kv_cache, start_pos
|
||||
)
|
||||
x = attn_output + x
|
||||
|
||||
# feed forward
|
||||
x = self.mlp(self.post_attention_norm(x)) + x
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Embedding(nn.Module):
|
||||
def __init__(self, vocab_size: int, embedding_dim: int):
|
||||
super().__init__()
|
||||
self.weight = nn.Parameter(torch.empty((vocab_size, embedding_dim)))
|
||||
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return F.embedding(x, self.weight)
|
||||
@@ -0,0 +1,152 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from torch import Tensor
|
||||
from typing import Any, Mapping, Optional, Tuple
|
||||
from astrai.config.model_config import ModelConfig
|
||||
from astrai.model.module import (
|
||||
Embedding,
|
||||
DecoderBlock,
|
||||
Linear,
|
||||
RMSNorm,
|
||||
RotaryEmbedding,
|
||||
)
|
||||
|
||||
|
||||
def process_attention_mask(
|
||||
seq_mask: Tensor,
|
||||
input_tensor: Tensor,
|
||||
start_pos: int = 0,
|
||||
is_causal: bool = False,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Create attention mask for GQA
|
||||
Args:
|
||||
seq_mask (Tensor): A tensor indicating whether each position is valid or not.
|
||||
input_tensor (Tensor): The input tensor.
|
||||
start_pos (int): The starting position of the sequence.
|
||||
is_causal (bool): Whether the attention is causal or not.
|
||||
Returns:
|
||||
Tensor: The attention mask tensor.
|
||||
"""
|
||||
device = input_tensor.device
|
||||
dtype = input_tensor.dtype
|
||||
seq_len = input_tensor.size(1)
|
||||
|
||||
if seq_mask is None:
|
||||
if start_pos != 0:
|
||||
# for single prompt chat
|
||||
seq_mask = torch.ones((1, seq_len), dtype=torch.bool, device=device)
|
||||
else:
|
||||
return None
|
||||
|
||||
if seq_mask.dim() > 2:
|
||||
# shape (bsz, seq_len) or (bsz,n_heads, seq_len, seq_len + start_pos)
|
||||
# if ndim > 2, it's 4D tensor
|
||||
return seq_mask
|
||||
|
||||
batch_size = seq_mask.size(0)
|
||||
seq_mask = seq_mask[:, : start_pos + seq_len].to(device=device, dtype=torch.bool)
|
||||
# (bsz, start_pos + seq_len)
|
||||
expanded_mask = seq_mask.unsqueeze(1).expand(
|
||||
batch_size, seq_len, start_pos + seq_len
|
||||
)
|
||||
# (bsz, seq_len, start_pos + seq_len)
|
||||
|
||||
if is_causal:
|
||||
expanded_mask = torch.tril(expanded_mask, diagonal=start_pos)
|
||||
|
||||
attention_mask = torch.zeros_like(expanded_mask, dtype=dtype, device=device)
|
||||
attention_mask = attention_mask.masked_fill_(
|
||||
~expanded_mask, -torch.finfo(dtype).max / 2
|
||||
).unsqueeze(1)
|
||||
# (bsz, 1, seq_len, seq_len + start_pos)
|
||||
|
||||
return attention_mask
|
||||
|
||||
|
||||
class Transformer(nn.Module):
|
||||
def __init__(self, config: ModelConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.rotary_embeding = RotaryEmbedding(
|
||||
config.dim // config.n_heads, config.max_len
|
||||
)
|
||||
self.embed_tokens = Embedding(config.vocab_size, config.dim)
|
||||
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
DecoderBlock(
|
||||
config.dim,
|
||||
config.n_heads,
|
||||
config.dim_ffn,
|
||||
config.n_kv_heads,
|
||||
config.norm_eps,
|
||||
config.use_qk_norm,
|
||||
config.use_gated_attention,
|
||||
layer_id,
|
||||
)
|
||||
for layer_id in range(config.n_layers)
|
||||
]
|
||||
)
|
||||
|
||||
self.norm = RMSNorm(config.dim, config.norm_eps)
|
||||
self.lm_head = Linear(config.dim, config.vocab_size)
|
||||
|
||||
if self.config.tie_weight == True:
|
||||
self.lm_head.weight = self.embed_tokens.weight
|
||||
|
||||
self._init_parameters()
|
||||
|
||||
def load_state_dict(self, state_dict: Mapping[str, Any], strict=True, assign=False):
|
||||
lm_head_key = "lm_head.weight"
|
||||
embed_key = "embed_tokens.weight"
|
||||
|
||||
if self.config.tie_weight == True:
|
||||
# same tensor
|
||||
state_dict[lm_head_key] = state_dict[embed_key]
|
||||
else:
|
||||
if lm_head_key not in state_dict and embed_key in state_dict:
|
||||
# use clone to avoid sharing the same tensor
|
||||
state_dict[lm_head_key] = torch.clone(state_dict[embed_key])
|
||||
|
||||
return super().load_state_dict(state_dict, strict, assign)
|
||||
|
||||
def state_dict(self, destination=None, prefix="", keep_vars=False):
|
||||
state_dict = super().state_dict(
|
||||
destination=destination, prefix=prefix, keep_vars=keep_vars
|
||||
)
|
||||
|
||||
if self.config.tie_weight == True:
|
||||
lm_head_key = prefix + "lm_head.weight"
|
||||
if lm_head_key in state_dict:
|
||||
del state_dict[lm_head_key]
|
||||
|
||||
return state_dict
|
||||
|
||||
def _init_parameters(self):
|
||||
for param in self.parameters():
|
||||
if param.dim() > 1:
|
||||
nn.init.normal_(param, mean=0.0, std=0.006)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Tensor,
|
||||
input_mask: Optional[Tensor] = None,
|
||||
persistent_key_values: Optional[Tuple[Tensor, Tensor]] = None,
|
||||
start_pos: int = 0,
|
||||
) -> Tensor:
|
||||
assert input_ids.ndim == 2
|
||||
|
||||
x = self.embed_tokens(input_ids)
|
||||
rotary_emb = self.rotary_embeding(x, start_pos)
|
||||
|
||||
attn_mask = process_attention_mask(input_mask, x, start_pos, is_causal=True)
|
||||
|
||||
for layer in self.layers:
|
||||
x = layer(x, rotary_emb, attn_mask, persistent_key_values, start_pos)
|
||||
|
||||
hidden_states = self.norm(x)
|
||||
logits = self.lm_head(hidden_states)
|
||||
|
||||
return {"logits": logits, "hidden_states": hidden_states}
|
||||
@@ -0,0 +1,21 @@
|
||||
from astrai.parallel.setup import (
|
||||
get_world_size,
|
||||
get_rank,
|
||||
get_current_device,
|
||||
only_on_rank,
|
||||
setup_parallel,
|
||||
spawn_parallel_fn,
|
||||
)
|
||||
|
||||
from astrai.parallel.module import RowParallelLinear, ColumnParallelLinear
|
||||
|
||||
__all__ = [
|
||||
"get_world_size",
|
||||
"get_rank",
|
||||
"get_current_device",
|
||||
"only_on_rank",
|
||||
"setup_parallel",
|
||||
"spawn_parallel_fn",
|
||||
"RowParallelLinear",
|
||||
"ColumnParallelLinear",
|
||||
]
|
||||
@@ -0,0 +1,115 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
|
||||
from torch import Tensor
|
||||
from typing import Dict
|
||||
|
||||
|
||||
class ParallelModel(nn.Module):
|
||||
def __init__(self, process_group: dist.ProcessGroup):
|
||||
super().__init__()
|
||||
self.process_group = process_group
|
||||
self.rank = dist.get_rank(self.process_group)
|
||||
self.world_size = dist.get_world_size(self.process_group)
|
||||
|
||||
|
||||
class RowParallelLinear(ParallelModel):
|
||||
def __init__(
|
||||
self,
|
||||
process_group: dist.ProcessGroup,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
bias: bool = True,
|
||||
reduce_results: bool = True,
|
||||
):
|
||||
super().__init__(process_group)
|
||||
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.in_features_per_rank = in_features // self.world_size
|
||||
self.reduce_results = reduce_results
|
||||
|
||||
if in_features % self.world_size != 0:
|
||||
raise ValueError(
|
||||
f"in_features must be divisible by world_size. Got {in_features} and {self.world_size}"
|
||||
)
|
||||
|
||||
self.weight = nn.Parameter(torch.empty(out_features, self.in_features_per_rank))
|
||||
self.bias = nn.Parameter(torch.zeros(out_features)) if bias else None
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
output = F.linear(input, self.weight)
|
||||
|
||||
if self.reduce_results:
|
||||
dist.all_reduce(output, op=dist.ReduceOp.SUM, group=self.process_group)
|
||||
|
||||
if self.bias is not None:
|
||||
output += self.bias
|
||||
|
||||
return output
|
||||
|
||||
def load_state_dict(self, state_dict: Dict[str, Tensor]):
|
||||
full_weight = state_dict.get("weight")
|
||||
full_bias = state_dict.get("bias")
|
||||
|
||||
start_idx = self.rank * self.in_features_per_rank
|
||||
end_idx = start_idx + self.in_features_per_rank
|
||||
weight_slice = full_weight[:, start_idx:end_idx]
|
||||
self.weight.data.copy_(weight_slice)
|
||||
|
||||
if self.bias is not None:
|
||||
self.bias.data.copy_(full_bias)
|
||||
|
||||
|
||||
class ColumnParallelLinear(ParallelModel):
|
||||
def __init__(
|
||||
self,
|
||||
process_group: dist.ProcessGroup,
|
||||
in_features: int,
|
||||
out_features: int,
|
||||
bias: bool = True,
|
||||
gather_results: bool = True,
|
||||
):
|
||||
super().__init__(process_group)
|
||||
|
||||
self.in_features = in_features
|
||||
self.out_features = out_features
|
||||
self.out_features_per_rank = out_features // self.world_size
|
||||
self.gather_results = gather_results
|
||||
|
||||
if out_features % self.world_size != 0:
|
||||
raise ValueError(
|
||||
f"out_features must be divisible by world_size. Got {out_features} and {self.world_size}"
|
||||
)
|
||||
|
||||
self.weight = nn.Parameter(
|
||||
torch.empty(self.out_features_per_rank, self.in_features)
|
||||
)
|
||||
self.bias = (
|
||||
nn.Parameter(torch.zeros(self.out_features_per_rank)) if bias else None
|
||||
)
|
||||
|
||||
def forward(self, input: Tensor) -> Tensor:
|
||||
output = F.linear(input, self.weight, self.bias)
|
||||
|
||||
if self.gather_results:
|
||||
output_list = [torch.empty_like(output) for _ in range(self.world_size)]
|
||||
dist.all_gather(output_list, output, group=self.process_group)
|
||||
output = torch.cat(output_list, dim=-1)
|
||||
|
||||
return output
|
||||
|
||||
def load_state_dict(self, state_dict: Dict[str, Tensor]):
|
||||
full_weight = state_dict.get("weight")
|
||||
full_bias = state_dict.get("bias")
|
||||
|
||||
start_idx = self.rank * self.out_features_per_rank
|
||||
end_idx = start_idx + self.out_features_per_rank
|
||||
weight_slice = full_weight[start_idx:end_idx, :]
|
||||
self.weight.data.copy_(weight_slice)
|
||||
|
||||
if self.bias is not None:
|
||||
bias_slice = full_bias[start_idx:end_idx]
|
||||
self.bias.data.copy_(bias_slice)
|
||||
@@ -0,0 +1,170 @@
|
||||
import os
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
from functools import wraps
|
||||
from contextlib import contextmanager
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
|
||||
def get_current_device():
|
||||
return os.environ["LOCAL_DEVICE"]
|
||||
|
||||
|
||||
def get_world_size() -> int:
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
return dist.get_world_size()
|
||||
else:
|
||||
return 1
|
||||
|
||||
|
||||
def get_rank() -> int:
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
return dist.get_rank()
|
||||
else:
|
||||
return 0
|
||||
|
||||
|
||||
@contextmanager
|
||||
def setup_parallel(
|
||||
rank: int,
|
||||
world_size: int,
|
||||
backend: str = "nccl",
|
||||
master_addr: str = "localhost",
|
||||
master_port: str = "29500",
|
||||
device_type: str = "cuda",
|
||||
device_ids: Optional[List[int]] = None,
|
||||
):
|
||||
|
||||
if dist.is_available() and dist.is_initialized():
|
||||
yield dist.group.WORLD
|
||||
return
|
||||
|
||||
if world_size <= 1:
|
||||
yield None
|
||||
return
|
||||
|
||||
if device_ids is None:
|
||||
device_ids = [i for i in range(world_size)]
|
||||
|
||||
rank = device_ids[rank % len(device_ids)]
|
||||
device_id = torch.device(device_type, device_ids[rank])
|
||||
|
||||
os.environ["MASTER_ADDR"] = master_addr
|
||||
os.environ["MASTER_PORT"] = master_port
|
||||
|
||||
os.environ["LOCAL_RANK"] = str(rank)
|
||||
os.environ["WORLD_SIZE"] = str(world_size)
|
||||
os.environ["LOCAL_DEVICE"] = str(device_id)
|
||||
|
||||
dist.init_process_group(
|
||||
rank=rank, world_size=world_size, backend=backend, device_id=device_id
|
||||
)
|
||||
|
||||
try:
|
||||
if backend == "nccl" and torch.cuda.is_available():
|
||||
torch.cuda.set_device(device_id)
|
||||
elif backend == "ccl" and hasattr(torch, "xpu") and torch.xpu.is_available():
|
||||
torch.xpu.set_device(device_id)
|
||||
|
||||
yield dist.group.WORLD
|
||||
finally:
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def only_on_rank(rank, sync=False):
|
||||
"""
|
||||
decorator to run a function only on a specific rank.
|
||||
"""
|
||||
|
||||
def decorator(func):
|
||||
@wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
ret_args = None
|
||||
if get_rank() == rank:
|
||||
ret_args = func(*args, **kwargs)
|
||||
|
||||
if sync and dist.is_available() and dist.is_initialized():
|
||||
dist.barrier()
|
||||
|
||||
return ret_args
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def wrapper_spawn_func(
|
||||
rank: int,
|
||||
world_size: int,
|
||||
backend: str,
|
||||
master_addr: str,
|
||||
master_port: str,
|
||||
device_type: str,
|
||||
device_ids: List[int],
|
||||
func: Callable,
|
||||
kwargs: dict,
|
||||
):
|
||||
try:
|
||||
with setup_parallel(
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
backend=backend,
|
||||
master_addr=master_addr,
|
||||
master_port=master_port,
|
||||
device_type=device_type,
|
||||
device_ids=device_ids,
|
||||
):
|
||||
func(**kwargs)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in rank {rank}: {e}")
|
||||
raise
|
||||
|
||||
|
||||
def spawn_parallel_fn(
|
||||
func: Callable,
|
||||
world_size: int,
|
||||
backend: str = "nccl",
|
||||
master_addr: str = "localhost",
|
||||
master_port: str = "29500",
|
||||
device_type: str = "cuda",
|
||||
device_ids: Optional[List[int]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
# clear environment variables
|
||||
for key in [
|
||||
"MASTER_ADDR",
|
||||
"MASTER_PORT",
|
||||
"RANK",
|
||||
"WORLD_SIZE",
|
||||
"LOCAL_RANK",
|
||||
"LOCAL_DEVICE",
|
||||
]:
|
||||
if key in os.environ:
|
||||
del os.environ[key]
|
||||
|
||||
if world_size == 1:
|
||||
device_ids = device_ids or [0]
|
||||
device_id = torch.device(device_type, device_ids[0])
|
||||
os.environ["LOCAL_DEVICE"] = str(device_id)
|
||||
|
||||
func(**kwargs)
|
||||
return
|
||||
|
||||
wrapper_spawn_func_args = (
|
||||
world_size,
|
||||
backend,
|
||||
master_addr,
|
||||
master_port,
|
||||
device_type,
|
||||
device_ids,
|
||||
func,
|
||||
kwargs,
|
||||
)
|
||||
|
||||
mp.spawn(
|
||||
wrapper_spawn_func, nprocs=world_size, args=wrapper_spawn_func_args, join=True
|
||||
)
|
||||
@@ -0,0 +1,30 @@
|
||||
from astrai.trainer.trainer import Trainer
|
||||
from astrai.trainer.strategy import StrategyFactory, BaseStrategy
|
||||
from astrai.trainer.schedule import SchedulerFactory, BaseScheduler
|
||||
|
||||
from astrai.trainer.train_callback import (
|
||||
TrainCallback,
|
||||
GradientClippingCallback,
|
||||
SchedulerCallback,
|
||||
CheckpointCallback,
|
||||
ProgressBarCallback,
|
||||
MetricLoggerCallback,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
# Main trainer
|
||||
"Trainer",
|
||||
# Strategy factory
|
||||
"StrategyFactory",
|
||||
"BaseStrategy",
|
||||
# Scheduler factory
|
||||
"SchedulerFactory",
|
||||
"BaseScheduler",
|
||||
# Callbacks
|
||||
"TrainCallback",
|
||||
"GradientClippingCallback",
|
||||
"SchedulerCallback",
|
||||
"CheckpointCallback",
|
||||
"ProgressBarCallback",
|
||||
"MetricLoggerCallback",
|
||||
]
|
||||
@@ -0,0 +1,103 @@
|
||||
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
|
||||
|
||||
|
||||
def ctx_get_loss(ctx):
|
||||
return ctx.loss
|
||||
|
||||
|
||||
def ctx_get_lr(ctx):
|
||||
return ctx.optimizer.param_groups[-1]["lr"]
|
||||
|
||||
|
||||
def ctx_get_grad_norm(ctx):
|
||||
return grad_norm(ctx.model)
|
||||
|
||||
|
||||
def ctx_get_grad_std(ctx):
|
||||
return grad_std(ctx.model)
|
||||
|
||||
|
||||
def ctx_get_grad_max(ctx):
|
||||
return grad_max(ctx.model)
|
||||
|
||||
|
||||
def ctx_get_grad_min(ctx):
|
||||
return grad_min(ctx.model)
|
||||
|
||||
|
||||
def ctx_get_grad_mean(ctx):
|
||||
return grad_mean(ctx.model)
|
||||
|
||||
|
||||
def ctx_get_grad_nan_num(ctx):
|
||||
return grad_nan_num(ctx.model)
|
||||
@@ -0,0 +1,237 @@
|
||||
"""Learning rate scheduler implementations with factory pattern."""
|
||||
|
||||
import math
|
||||
from abc import abstractmethod, ABC
|
||||
from typing import Any, Dict, List, Type
|
||||
from torch.optim.lr_scheduler import LRScheduler
|
||||
from astrai.config.schedule_config import ScheduleConfig
|
||||
|
||||
|
||||
class BaseScheduler(LRScheduler, ABC):
|
||||
"""Base scheduler class for all other schedulers."""
|
||||
|
||||
def __init__(self, optimizer, last_epoch: int = -1):
|
||||
super().__init__(optimizer, last_epoch)
|
||||
|
||||
@abstractmethod
|
||||
def get_lr(self) -> List[float]:
|
||||
"""Calculate the current learning rate."""
|
||||
raise NotImplementedError
|
||||
|
||||
def state_dict(self) -> Dict[str, Any]:
|
||||
return super().state_dict()
|
||||
|
||||
def load_state_dict(self, state_dict: Dict[str, Any]):
|
||||
super().load_state_dict(state_dict)
|
||||
|
||||
|
||||
class SchedulerFactory:
|
||||
"""Factory class for creating learning rate schedulers.
|
||||
|
||||
Supports decorator-based registration for extensible scheduler types.
|
||||
Also supports creation from ScheduleConfig objects.
|
||||
|
||||
Example usage:
|
||||
@SchedulerFactory.register("custom")
|
||||
class CustomScheduler(BaseScheduler):
|
||||
...
|
||||
|
||||
scheduler = SchedulerFactory.create(optimizer, "custom", **kwargs)
|
||||
|
||||
# Or from config
|
||||
config = CosineScheduleConfig(total_steps=10000)
|
||||
scheduler = SchedulerFactory.load(optimizer, config)
|
||||
"""
|
||||
|
||||
SCHEDULER_MAP: Dict[str, Type[BaseScheduler]] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str):
|
||||
"""Decorator to register a new scheduler class.
|
||||
|
||||
Args:
|
||||
name: Registration name for the scheduler
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the scheduler class
|
||||
"""
|
||||
|
||||
def decorator(scheduler_cls: Type[BaseScheduler]) -> Type[BaseScheduler]:
|
||||
if not issubclass(scheduler_cls, BaseScheduler):
|
||||
raise TypeError(
|
||||
f"{scheduler_cls.__name__} must inherit from BaseScheduler"
|
||||
)
|
||||
cls.SCHEDULER_MAP[name] = scheduler_cls
|
||||
return scheduler_cls
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def create(cls, optimizer, schedule_type: str, **kwargs) -> BaseScheduler:
|
||||
"""Create a scheduler instance by type name.
|
||||
|
||||
Args:
|
||||
optimizer: PyTorch optimizer
|
||||
schedule_type: Type of scheduler ("cosine", "sgdr")
|
||||
**kwargs: Arguments passed to the scheduler constructor
|
||||
|
||||
Returns:
|
||||
Scheduler instance
|
||||
|
||||
Raises:
|
||||
ValueError: If schedule_type is not supported
|
||||
"""
|
||||
if schedule_type not in cls.SCHEDULER_MAP:
|
||||
raise ValueError(
|
||||
f"Unknown schedule type: '{schedule_type}'. "
|
||||
f"Supported types: {sorted(cls.SCHEDULER_MAP.keys())}"
|
||||
)
|
||||
|
||||
scheduler_cls = cls.SCHEDULER_MAP[schedule_type]
|
||||
return scheduler_cls(optimizer, **kwargs)
|
||||
|
||||
@staticmethod
|
||||
def load(optimizer, schedule_config: ScheduleConfig) -> BaseScheduler:
|
||||
"""Create a scheduler from a ScheduleConfig object.
|
||||
|
||||
Args:
|
||||
optimizer: PyTorch optimizer
|
||||
schedule_config: ScheduleConfig instance
|
||||
|
||||
Returns:
|
||||
Scheduler instance
|
||||
"""
|
||||
kwargs = schedule_config.get_kwargs()
|
||||
schedule_type = kwargs.pop("schedule_type")
|
||||
return SchedulerFactory.create(optimizer, schedule_type, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def available_types(cls) -> list:
|
||||
"""Return list of registered scheduler type names."""
|
||||
return list(cls.SCHEDULER_MAP.keys())
|
||||
|
||||
|
||||
# ============== Scheduler Classes ==============
|
||||
# All scheduler classes are registered at class definition time using the decorator
|
||||
|
||||
|
||||
@SchedulerFactory.register("cosine")
|
||||
class CosineScheduler(BaseScheduler):
|
||||
"""Cosine decay scheduler with warmup, implemented as PyTorch LRScheduler."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
optimizer,
|
||||
warmup_steps: int,
|
||||
lr_decay_steps: int,
|
||||
min_rate: float = 0.05,
|
||||
last_epoch: int = -1,
|
||||
):
|
||||
self.warmup_steps = warmup_steps
|
||||
self.lr_decay_steps = lr_decay_steps
|
||||
self.min_rate = min_rate
|
||||
self.total_steps = warmup_steps + lr_decay_steps
|
||||
super().__init__(optimizer, last_epoch)
|
||||
|
||||
def get_lr(self) -> List[float]:
|
||||
# warmup
|
||||
if self.last_epoch < self.warmup_steps:
|
||||
warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
|
||||
return [base_lr * warmup_factor for base_lr in self.base_lrs]
|
||||
|
||||
# cosine decay
|
||||
decay_progress = (self.last_epoch - self.warmup_steps) / self.lr_decay_steps
|
||||
decay_progress = min(decay_progress, 1.0)
|
||||
cosine_decay = 0.5 * (1.0 + math.cos(math.pi * decay_progress))
|
||||
decay_factor = max(self.min_rate, cosine_decay)
|
||||
return [base_lr * decay_factor for base_lr in self.base_lrs]
|
||||
|
||||
def state_dict(self):
|
||||
state = super().state_dict()
|
||||
state.update(
|
||||
{
|
||||
"warmup_steps": self.warmup_steps,
|
||||
"lr_decay_steps": self.lr_decay_steps,
|
||||
"min_rate": self.min_rate,
|
||||
"total_steps": self.total_steps,
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
self.warmup_steps = state_dict.pop("warmup_steps")
|
||||
self.lr_decay_steps = state_dict.pop("lr_decay_steps")
|
||||
self.min_rate = state_dict.pop("min_rate")
|
||||
self.total_steps = state_dict.pop("total_steps")
|
||||
super().load_state_dict(state_dict)
|
||||
|
||||
|
||||
@SchedulerFactory.register("sgdr")
|
||||
class SGDRScheduler(BaseScheduler):
|
||||
"""SGDR (Stochastic Gradient Descent with Warm Restarts) scheduler."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
optimizer,
|
||||
warmup_steps: int,
|
||||
cycle_length: int,
|
||||
min_rate: float = 0.05,
|
||||
t_mult: int = 2,
|
||||
last_epoch: int = -1,
|
||||
):
|
||||
self.warmup_steps = warmup_steps
|
||||
self.cycle_length = cycle_length
|
||||
self.min_rate = min_rate
|
||||
self.t_mult = t_mult
|
||||
|
||||
super().__init__(optimizer, last_epoch)
|
||||
|
||||
def get_lr(self):
|
||||
# warmup
|
||||
if self.last_epoch < self.warmup_steps:
|
||||
warmup_factor = max(self.min_rate, self.last_epoch / self.warmup_steps)
|
||||
return [base_lr * warmup_factor for base_lr in self.base_lrs]
|
||||
|
||||
# SGDR
|
||||
steps_since_warmup = self.last_epoch - self.warmup_steps
|
||||
|
||||
# 1. Calculate current cycle and position within cycle
|
||||
current_cycle_length = self.cycle_length
|
||||
total_cycles_length = 0
|
||||
cycle_num = 0
|
||||
|
||||
while total_cycles_length + current_cycle_length <= steps_since_warmup:
|
||||
total_cycles_length += current_cycle_length
|
||||
current_cycle_length *= self.t_mult
|
||||
cycle_num += 1
|
||||
|
||||
steps_in_cycle = steps_since_warmup - total_cycles_length
|
||||
|
||||
# 2. Cosine annealing within the current cycle
|
||||
cosine_factor = 0.5 * (
|
||||
1 + math.cos(math.pi * steps_in_cycle / current_cycle_length)
|
||||
)
|
||||
learning_rate_factor = self.min_rate + (1 - self.min_rate) * cosine_factor
|
||||
|
||||
return [base_lr * learning_rate_factor for base_lr in self.base_lrs]
|
||||
|
||||
def state_dict(self):
|
||||
"""Returns the state of the scheduler as a dict."""
|
||||
state = super().state_dict()
|
||||
state.update(
|
||||
{
|
||||
"warmup_steps": self.warmup_steps,
|
||||
"cycle_length": self.cycle_length,
|
||||
"min_rate": self.min_rate,
|
||||
"t_mult": self.t_mult,
|
||||
}
|
||||
)
|
||||
return state
|
||||
|
||||
def load_state_dict(self, state_dict):
|
||||
"""Loads the scheduler's state."""
|
||||
self.warmup_steps = state_dict.pop("warmup_steps")
|
||||
self.cycle_length = state_dict.pop("cycle_length")
|
||||
self.min_rate = state_dict.pop("min_rate")
|
||||
self.t_mult = state_dict.pop("t_mult")
|
||||
super().load_state_dict(state_dict)
|
||||
@@ -0,0 +1,362 @@
|
||||
"""Training strategy implementations with factory pattern."""
|
||||
|
||||
import copy
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.parallel import DistributedDataParallel as DDP
|
||||
|
||||
from torch import Tensor
|
||||
from typing import Any, Callable, Dict, Union
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
|
||||
def unwrap_model(model: nn.Module) -> nn.Module:
|
||||
"""Unwrap DDP wrapper if present to get the original model."""
|
||||
if isinstance(model, DDP):
|
||||
return model.module
|
||||
return model
|
||||
|
||||
|
||||
def create_ref_model(model: nn.Module) -> nn.Module:
|
||||
"""Create a reference model for DPO/GRPO training.
|
||||
|
||||
Handles DDP-wrapped models safely by unwrapping first,
|
||||
then creating a deep copy with frozen gradients.
|
||||
"""
|
||||
original_model = unwrap_model(model)
|
||||
ref_model = copy.deepcopy(original_model)
|
||||
ref_model.requires_grad_(False)
|
||||
ref_model.eval()
|
||||
return ref_model
|
||||
|
||||
|
||||
def move_to_device(batch: Dict[str, Tensor], device: str) -> Any:
|
||||
"""Move batch tensors to specified device with non-blocking transfer."""
|
||||
return {key: value.to(device, non_blocking=True) for key, value in batch.items()}
|
||||
|
||||
|
||||
def get_logprobs(
|
||||
model: Union[nn.Module, Callable[..., Dict[str, Tensor]]],
|
||||
input_ids: Tensor,
|
||||
mask: Tensor,
|
||||
reduction: str,
|
||||
):
|
||||
"""Compute token-wise log probabilities from model outputs.
|
||||
|
||||
Args:
|
||||
model: The language model
|
||||
input_ids: Input token IDs of shape [batch_size, seq_len]
|
||||
mask: Attention mask of shape [batch_size, seq_len]
|
||||
reduction: How to reduce over sequence dimension ("mean", "sum", "none")
|
||||
|
||||
Returns:
|
||||
Log probabilities with reduction applied over sequence dimension
|
||||
"""
|
||||
allowed_reductions = ["mean", "sum", "none"]
|
||||
if reduction not in allowed_reductions:
|
||||
raise ValueError(
|
||||
f"reduction must be one of {allowed_reductions}, got '{reduction}'"
|
||||
)
|
||||
|
||||
shifted_input_ids = input_ids[:, 1:]
|
||||
shifted_mask = mask[:, 1:]
|
||||
|
||||
logits = model(input_ids[:, :-1], mask[:, :-1])["logits"]
|
||||
log_probs = torch.log_softmax(logits.float(), dim=-1)
|
||||
|
||||
token_logprobs = torch.gather(
|
||||
log_probs, dim=-1, index=shifted_input_ids.unsqueeze(-1)
|
||||
).squeeze(-1)
|
||||
|
||||
if reduction == "mean":
|
||||
return (token_logprobs * shifted_mask).sum(dim=-1) / shifted_mask.sum(
|
||||
dim=-1
|
||||
).clamp(min=1.0)
|
||||
elif reduction == "sum":
|
||||
return (token_logprobs * shifted_mask).sum(dim=-1)
|
||||
else:
|
||||
return token_logprobs * shifted_mask
|
||||
|
||||
|
||||
class BaseStrategy(ABC):
|
||||
"""Abstract base class for training strategies."""
|
||||
|
||||
def __init__(
|
||||
self, model: Union[Callable[..., Dict[str, Tensor]]], device: str, **kwargs
|
||||
):
|
||||
self.model = model
|
||||
self.device = device
|
||||
self.extra_kwargs = kwargs
|
||||
|
||||
@abstractmethod
|
||||
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
|
||||
"""Compute loss for the given batch.
|
||||
|
||||
Args:
|
||||
batch: Dictionary containing batch tensors
|
||||
|
||||
Returns:
|
||||
Computed loss tensor
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def __call__(self, batch: Dict[str, Tensor]) -> Tensor:
|
||||
"""Allow calling strategy directly as a callable."""
|
||||
return self.compute_loss(batch)
|
||||
|
||||
|
||||
class StrategyFactory:
|
||||
"""Factory class for creating training strategy instances.
|
||||
|
||||
Supports decorator-based registration for extensible strategy types.
|
||||
All default strategies (seq, sft, dpo, grpo) are automatically registered.
|
||||
|
||||
Example usage:
|
||||
@StrategyFactory.register("custom")
|
||||
class CustomStrategy(BaseStrategy):
|
||||
...
|
||||
|
||||
strategy = StrategyFactory.create(model, "custom", device)
|
||||
"""
|
||||
|
||||
SUPPORTED_STRATEGIES = frozenset({"seq", "sft", "dpo", "grpo"})
|
||||
STRATEGY_MAP: Dict[str, type] = {}
|
||||
|
||||
@classmethod
|
||||
def register(cls, name: str):
|
||||
"""Decorator to register a new strategy class.
|
||||
|
||||
Args:
|
||||
name: Registration name for the strategy
|
||||
|
||||
Returns:
|
||||
Decorator function that registers the strategy class
|
||||
"""
|
||||
|
||||
def decorator(strategy_cls: type) -> type:
|
||||
if not issubclass(strategy_cls, BaseStrategy):
|
||||
raise TypeError(
|
||||
f"{strategy_cls.__name__} must inherit from BaseStrategy"
|
||||
)
|
||||
cls.STRATEGY_MAP[name] = strategy_cls
|
||||
return strategy_cls
|
||||
|
||||
return decorator
|
||||
|
||||
@classmethod
|
||||
def create(cls, model, train_type: str, device: str, **kwargs) -> BaseStrategy:
|
||||
"""Create a strategy instance based on training type.
|
||||
|
||||
Args:
|
||||
model: Model instance for the strategy
|
||||
train_type: Type of training ("seq", "sft", "dpo", "grpo")
|
||||
device: Device to run the strategy on
|
||||
**kwargs: Additional arguments passed to strategy constructor
|
||||
|
||||
Returns:
|
||||
Strategy instance
|
||||
|
||||
Raises:
|
||||
ValueError: If train_type is not supported
|
||||
NotImplementedError: If train_type is in supported list but not implemented
|
||||
"""
|
||||
if train_type not in cls.SUPPORTED_STRATEGIES:
|
||||
raise ValueError(
|
||||
f"Unknown training strategy: '{train_type}'. "
|
||||
f"Supported strategies: {sorted(cls.SUPPORTED_STRATEGIES)}"
|
||||
)
|
||||
|
||||
if train_type not in cls.STRATEGY_MAP:
|
||||
raise NotImplementedError(
|
||||
f"Strategy '{train_type}' is supported but not yet implemented."
|
||||
)
|
||||
|
||||
strategy_cls = cls.STRATEGY_MAP[train_type]
|
||||
return strategy_cls(model, device, **kwargs)
|
||||
|
||||
@classmethod
|
||||
def available_strategies(cls) -> list:
|
||||
"""Return list of registered strategy names."""
|
||||
return list(cls.STRATEGY_MAP.keys())
|
||||
|
||||
|
||||
# ============== Strategy Classes ==============
|
||||
# All strategies are registered at class definition time using the decorator
|
||||
|
||||
|
||||
@StrategyFactory.register("seq")
|
||||
class SEQStrategy(BaseStrategy):
|
||||
"""Standard next-token prediction training strategy.
|
||||
|
||||
Computes cross-entropy loss for next token prediction.
|
||||
"""
|
||||
|
||||
def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
|
||||
super().__init__(model, device, **kwargs)
|
||||
self.label_smoothing = label_smoothing
|
||||
|
||||
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"]
|
||||
logits = self.model(input_ids=input_ids)["logits"]
|
||||
|
||||
loss = F.cross_entropy(
|
||||
input=logits.flatten(0, 1).float(),
|
||||
target=target_ids.flatten(),
|
||||
label_smoothing=self.label_smoothing,
|
||||
)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
@StrategyFactory.register("sft")
|
||||
class SFTStrategy(BaseStrategy):
|
||||
"""Supervised Fine-tuning strategy with loss masking.
|
||||
|
||||
Applies cross-entropy loss only to tokens where loss_mask is True.
|
||||
"""
|
||||
|
||||
def __init__(self, model, device, label_smoothing: float = 0.0, **kwargs):
|
||||
super().__init__(model, device, **kwargs)
|
||||
self.label_smoothing = label_smoothing
|
||||
|
||||
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
|
||||
batch = move_to_device(batch, self.device)
|
||||
input_ids, target_ids, loss_mask = (
|
||||
batch["input_ids"],
|
||||
batch["target_ids"],
|
||||
batch["loss_mask"],
|
||||
)
|
||||
|
||||
ignore_index = -100
|
||||
logits = self.model(input_ids=input_ids)["logits"]
|
||||
target_ids = target_ids.masked_fill(loss_mask == 0, ignore_index)
|
||||
|
||||
loss = F.cross_entropy(
|
||||
input=logits.flatten(0, 1).float(),
|
||||
target=target_ids.flatten(),
|
||||
ignore_index=ignore_index,
|
||||
label_smoothing=self.label_smoothing,
|
||||
)
|
||||
|
||||
return loss
|
||||
|
||||
|
||||
@StrategyFactory.register("dpo")
|
||||
class DPOStrategy(BaseStrategy):
|
||||
"""Direct Preference Optimization strategy.
|
||||
|
||||
Implements the DPO loss from the paper "Direct Preference Optimization".
|
||||
Uses a reference model to compute KL divergence penalty.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
device: str,
|
||||
beta: float = 0.1,
|
||||
reduction: str = "mean",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(model, device, **kwargs)
|
||||
self.ref_model = create_ref_model(model)
|
||||
self.beta = beta
|
||||
self.reduction = reduction
|
||||
|
||||
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
|
||||
batch = move_to_device(batch, self.device)
|
||||
chosen_ids, rejected_ids = batch["chosen"], batch["rejected"]
|
||||
chosen_mask, rejected_mask = batch["chosen_mask"], batch["rejected_mask"]
|
||||
|
||||
contact_ids = torch.cat([chosen_ids, rejected_ids], dim=0)
|
||||
contact_mask = torch.cat([chosen_mask, rejected_mask], dim=0)
|
||||
|
||||
log_pi = get_logprobs(self.model, contact_ids, contact_mask, self.reduction)
|
||||
|
||||
with torch.no_grad():
|
||||
log_ref = get_logprobs(
|
||||
self.ref_model, contact_ids, contact_mask, self.reduction
|
||||
)
|
||||
|
||||
log_pi_chosen = log_pi[: chosen_ids.shape[0]]
|
||||
log_pi_rejected = log_pi[chosen_ids.shape[0] :]
|
||||
log_ref_chosen = log_ref[: chosen_ids.shape[0]]
|
||||
log_ref_rejected = log_ref[chosen_ids.shape[0] :]
|
||||
|
||||
pi_log_ratio = log_pi_chosen - log_pi_rejected
|
||||
ref_log_ratio = log_ref_chosen - log_ref_rejected
|
||||
|
||||
ratio_diff = pi_log_ratio - ref_log_ratio
|
||||
dpo_loss = -F.logsigmoid(self.beta * ratio_diff).mean()
|
||||
|
||||
return dpo_loss
|
||||
|
||||
|
||||
@StrategyFactory.register("grpo")
|
||||
class GRPOStrategy(BaseStrategy):
|
||||
"""Group Relative Policy Optimization strategy.
|
||||
|
||||
Implements GRPO with clipping and KL penalty.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: nn.Module,
|
||||
device: str,
|
||||
clip_eps: float = 0.2,
|
||||
kl_coef: float = 0.01,
|
||||
group_size: int = 4,
|
||||
reduction: str = "mean",
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(model, device, **kwargs)
|
||||
self.ref_model = create_ref_model(model)
|
||||
self.clip_eps = clip_eps
|
||||
self.kl_coef = kl_coef
|
||||
self.group_size = group_size
|
||||
self.reduction = reduction
|
||||
|
||||
def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor:
|
||||
batch = move_to_device(batch, self.device)
|
||||
prompts = batch["prompts"]
|
||||
responses = batch["responses"]
|
||||
masks = batch["masks"]
|
||||
rewards = batch["rewards"]
|
||||
|
||||
batch_size, group_size, response_len = responses.shape
|
||||
responses_flat = responses.view(-1, response_len)
|
||||
masks_flat = masks.view(-1, response_len)
|
||||
prompt_expanded = prompts.unsqueeze(1).repeat(1, group_size, 1).flatten(0, 1)
|
||||
|
||||
# Shape: (batch_size * group_size, seq_len + response_len)
|
||||
full_sequences = torch.cat([prompt_expanded, responses_flat], dim=-1)
|
||||
full_masks = torch.cat([torch.ones_like(prompt_expanded), masks_flat], dim=-1)
|
||||
|
||||
log_probs_policy = get_logprobs(
|
||||
self.model, full_sequences, full_masks, self.reduction
|
||||
)
|
||||
log_probs_policy = log_probs_policy.view(batch_size, group_size)
|
||||
|
||||
with torch.no_grad():
|
||||
log_probs_ref = get_logprobs(
|
||||
self.ref_model, full_sequences, full_masks, self.reduction
|
||||
)
|
||||
log_probs_ref = log_probs_ref.view(batch_size, group_size)
|
||||
|
||||
# Compute advantages from rewards with normalization
|
||||
eps = torch.finfo(log_probs_policy.dtype).eps
|
||||
mean = rewards.mean(dim=-1, keepdim=True)
|
||||
std = rewards.std(dim=-1, keepdim=True)
|
||||
advantages = (rewards - mean) / (std + eps)
|
||||
|
||||
# PPO-style clipped surrogate objective
|
||||
ratio = torch.exp(0) # Off-policy: policy_model = old_model
|
||||
surr1 = ratio * advantages
|
||||
surr2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
|
||||
|
||||
policy_loss = -torch.min(surr1, surr2).mean()
|
||||
kl_penalty = self.kl_coef * (log_probs_policy - log_probs_ref).square().mean()
|
||||
total_loss = policy_loss + kl_penalty
|
||||
|
||||
return total_loss
|
||||
@@ -0,0 +1,235 @@
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
import torch.nn as nn
|
||||
|
||||
from pathlib import Path
|
||||
from tqdm import tqdm
|
||||
from torch.nn.utils import clip_grad_norm_
|
||||
from typing import Callable, List, Optional, Protocol
|
||||
|
||||
from astrai.parallel import only_on_rank
|
||||
from astrai.trainer.metric_util import (
|
||||
ctx_get_loss,
|
||||
ctx_get_lr,
|
||||
ctx_get_grad_max,
|
||||
ctx_get_grad_min,
|
||||
ctx_get_grad_norm,
|
||||
ctx_get_grad_mean,
|
||||
ctx_get_grad_std,
|
||||
ctx_get_grad_nan_num,
|
||||
)
|
||||
from astrai.data.serialization import Checkpoint
|
||||
from astrai.trainer.train_context import TrainContext
|
||||
|
||||
|
||||
class TrainCallback(Protocol):
|
||||
"""
|
||||
Callback interface for trainer.
|
||||
"""
|
||||
|
||||
def on_train_begin(self, context: TrainContext):
|
||||
"""Called at the beginning of training."""
|
||||
|
||||
def on_train_end(self, context: TrainContext):
|
||||
"""Called at the end of training."""
|
||||
|
||||
def on_epoch_begin(self, context: TrainContext):
|
||||
"""Called at the beginning of each epoch."""
|
||||
|
||||
def on_epoch_end(self, context: TrainContext):
|
||||
"""Called at the end of each epoch."""
|
||||
|
||||
def on_step_begin(self, context: TrainContext):
|
||||
"""Called at the beginning of each step."""
|
||||
|
||||
def on_step_end(self, context: TrainContext):
|
||||
"""Called at the end of each step."""
|
||||
|
||||
def on_batch_begin(self, context: TrainContext):
|
||||
"""Called at the beginning of each batch."""
|
||||
|
||||
def on_batch_end(self, context: TrainContext):
|
||||
"""Called at the end of each batch."""
|
||||
|
||||
def on_error(self, context: TrainContext):
|
||||
"""Called when an error occurs during training."""
|
||||
|
||||
|
||||
class GradientClippingCallback(TrainCallback):
|
||||
"""
|
||||
Gradient clipping callback for trainer.
|
||||
"""
|
||||
|
||||
def __init__(self, max_grad_norm: float):
|
||||
self.max_grad_norm = max_grad_norm
|
||||
|
||||
def on_step_begin(self, context: TrainContext):
|
||||
_ = context
|
||||
clip_grad_norm_(context.model.parameters(), self.max_grad_norm)
|
||||
|
||||
|
||||
class SchedulerCallback(TrainCallback):
|
||||
"""
|
||||
Scheduler callback for trainer.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def on_train_begin(self, context: TrainContext):
|
||||
for group in context.optimizer.param_groups:
|
||||
if "initial_lr" not in group:
|
||||
group["initial_lr"] = group["lr"]
|
||||
|
||||
def on_batch_end(self, context: TrainContext):
|
||||
if context.scheduler:
|
||||
context.scheduler.step()
|
||||
|
||||
|
||||
class CheckpointCallback(TrainCallback):
|
||||
"""
|
||||
Checkpoint callback for trainer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
save_dir: str,
|
||||
interval: int,
|
||||
weight_only: bool = False,
|
||||
state_dict_fn: Optional[Callable[[nn.Module], dict]] = None,
|
||||
):
|
||||
self.save_dir = save_dir
|
||||
self.interval = interval
|
||||
self.weight_only = weight_only
|
||||
self.state_dict_fn = state_dict_fn
|
||||
self.last_ckpt_iter = 0
|
||||
|
||||
@only_on_rank(0)
|
||||
def _save_checkpoint(self, context: TrainContext):
|
||||
save_path = os.path.join(
|
||||
self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}"
|
||||
)
|
||||
state_dict = (
|
||||
self.state_dict_fn(context.model)
|
||||
if self.state_dict_fn
|
||||
else context.model.state_dict()
|
||||
)
|
||||
|
||||
context.checkpoint = Checkpoint(
|
||||
state_dict=state_dict, epoch=context.epoch, iteration=context.iteration
|
||||
)
|
||||
|
||||
context.checkpoint.save(save_path)
|
||||
self.last_ckpt_iter = context.iteration
|
||||
|
||||
def on_batch_end(self, context: TrainContext):
|
||||
if context.iteration - self.last_ckpt_iter >= self.interval:
|
||||
self._save_checkpoint(context)
|
||||
|
||||
def on_train_end(self, context: TrainContext):
|
||||
if context.iteration != self.last_ckpt_iter:
|
||||
self._save_checkpoint(context)
|
||||
|
||||
def on_error(self, context: TrainContext):
|
||||
self._save_checkpoint(context)
|
||||
|
||||
|
||||
class ProgressBarCallback(TrainCallback):
|
||||
"""
|
||||
Progress bar callback for trainer.
|
||||
"""
|
||||
|
||||
def __init__(self, num_epoch: int):
|
||||
self.num_epoch = num_epoch
|
||||
self.progress_bar: tqdm = None
|
||||
|
||||
@only_on_rank(0)
|
||||
def on_epoch_begin(self, context: TrainContext):
|
||||
self.progress_bar = tqdm(
|
||||
context.dataloader,
|
||||
desc=f"Epoch {context.epoch + 1}/{self.num_epoch}",
|
||||
dynamic_ncols=True,
|
||||
)
|
||||
|
||||
@only_on_rank(0)
|
||||
def on_batch_end(self, context: TrainContext):
|
||||
self.progress_bar.set_postfix(
|
||||
{
|
||||
"loss": f"{context.loss:.4f}",
|
||||
"lr": f"{context.optimizer.param_groups[-1]['lr']:.2e}",
|
||||
}
|
||||
)
|
||||
self.progress_bar.update(1)
|
||||
|
||||
@only_on_rank(0)
|
||||
def on_epoch_end(self, context: TrainContext):
|
||||
_ = context
|
||||
if self.progress_bar:
|
||||
self.progress_bar.close()
|
||||
|
||||
|
||||
class MetricLoggerCallback(TrainCallback):
|
||||
def __init__(
|
||||
self,
|
||||
log_dir: str,
|
||||
save_interval: int,
|
||||
log_interval: int = 10,
|
||||
metrics: List[str] = None,
|
||||
):
|
||||
self.last_log_iter = 0
|
||||
self.save_interval = save_interval
|
||||
self.log_interval = log_interval
|
||||
self.metrics = metrics or ["loss", "lr"]
|
||||
|
||||
self.log_dir = Path(log_dir) if log_dir else Path.cwd() / "logs"
|
||||
self.log_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
self.log_cache = []
|
||||
|
||||
self._metric_funcs = {
|
||||
"loss": ctx_get_loss,
|
||||
"lr": ctx_get_lr,
|
||||
"grad_norm": ctx_get_grad_norm,
|
||||
"grad_std": ctx_get_grad_std,
|
||||
"grad_max": ctx_get_grad_max,
|
||||
"grad_min": ctx_get_grad_min,
|
||||
"grad_mean": ctx_get_grad_mean,
|
||||
"grad_nan_num": ctx_get_grad_nan_num,
|
||||
}
|
||||
|
||||
def _get_log_data(self, context: TrainContext):
|
||||
return {
|
||||
"timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"epoch": context.epoch,
|
||||
"iter": context.iteration,
|
||||
**{m: self._metric_funcs[m](context) for m in self.metrics},
|
||||
}
|
||||
|
||||
@only_on_rank(0)
|
||||
def _add_log(self, log_data):
|
||||
self.log_cache.append(log_data)
|
||||
|
||||
@only_on_rank(0)
|
||||
def _save_log(self, epoch, iter):
|
||||
log_file = self.log_dir / f"epoch_{epoch}_iter_{iter}_metric.jsonl"
|
||||
|
||||
with open(log_file, "w") as f:
|
||||
for log in self.log_cache:
|
||||
f.write(json.dumps(log) + "\n")
|
||||
|
||||
def on_batch_end(self, context):
|
||||
if context.iteration % self.log_interval == 0:
|
||||
log_data = self._get_log_data(context)
|
||||
self._add_log(log_data)
|
||||
|
||||
if context.iteration - self.last_log_iter >= self.save_interval:
|
||||
self._save_log(context.epoch, context.iteration)
|
||||
self.last_log_iter = context.iteration
|
||||
|
||||
def on_train_end(self, context):
|
||||
if context.iteration != self.last_log_iter:
|
||||
self._save_log(context.epoch, context.iteration)
|
||||
|
||||
def on_error(self, context):
|
||||
self._save_log(context.epoch, context.iteration)
|
||||
@@ -0,0 +1,99 @@
|
||||
import torch.nn as nn
|
||||
from torch.optim import Optimizer
|
||||
from torch.optim.lr_scheduler import LRScheduler
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from astrai.data import ResumableDistributedSampler
|
||||
from astrai.data.serialization import Checkpoint
|
||||
from astrai.trainer.strategy import StrategyFactory, BaseStrategy
|
||||
from astrai.config.train_config import TrainConfig
|
||||
from astrai.parallel.setup import get_current_device, get_world_size, get_rank
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional, Self
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainContext:
|
||||
model: nn.Module = field(default=None)
|
||||
strategy: BaseStrategy = field(default=None)
|
||||
dataloader: DataLoader = field(default=None)
|
||||
optimizer: Optimizer = field(default=None)
|
||||
scheduler: LRScheduler = field(default=None)
|
||||
checkpoint: Checkpoint = field(default=None)
|
||||
|
||||
epoch: int = field(default=0)
|
||||
iteration: int = field(default=0)
|
||||
loss: float = field(default=0.0)
|
||||
|
||||
world_size: int = field(default=1)
|
||||
rank: int = field(default=0)
|
||||
kwargs: dict = field(default_factory=dict)
|
||||
|
||||
|
||||
class TrainContextBuilder:
|
||||
def __init__(self, config: TrainConfig):
|
||||
self.config = config
|
||||
self._context = TrainContext(
|
||||
model=config.model,
|
||||
world_size=get_world_size(),
|
||||
rank=get_rank(),
|
||||
)
|
||||
|
||||
device = get_current_device()
|
||||
self._context.model = self._context.model.to(device=device)
|
||||
|
||||
if self.config.nprocs > 1:
|
||||
fn = self.config.parallel_wrapper
|
||||
self._context.model = fn(self._context.model)
|
||||
|
||||
self._context.optimizer = self.config.optimizer_fn(self._context.model)
|
||||
self._context.scheduler = self.config.scheduler_fn(self._context.optimizer)
|
||||
|
||||
def with_checkpoint(self, checkpoint: Optional[Checkpoint]) -> Self:
|
||||
if checkpoint is None:
|
||||
checkpoint = Checkpoint(
|
||||
state_dict=self._context.model.state_dict(),
|
||||
)
|
||||
else:
|
||||
# resume from the assigned checkpoint or assigned iteration
|
||||
self._context.epoch = max(checkpoint.epoch, self.config.start_epoch)
|
||||
self._context.iteration = max(checkpoint.iteration, self.config.start_batch)
|
||||
self._context.model.load_state_dict(checkpoint.state_dict)
|
||||
|
||||
self._context.checkpoint = checkpoint
|
||||
return self
|
||||
|
||||
def with_dataloader(self) -> Self:
|
||||
# fix: change batch level iteration to sample level offset
|
||||
config = self.config
|
||||
sampler_offset = self._context.iteration * config.batch_size
|
||||
resumeable_sampler = ResumableDistributedSampler(
|
||||
data_source=config.dataset,
|
||||
start_epoch=self._context.epoch,
|
||||
start_iter=sampler_offset,
|
||||
seed=config.random_seed,
|
||||
)
|
||||
|
||||
dataloader = DataLoader(
|
||||
config.dataset,
|
||||
batch_size=config.batch_size,
|
||||
sampler=resumeable_sampler,
|
||||
num_workers=config.num_workers,
|
||||
pin_memory=config.pin_memory,
|
||||
prefetch_factor=config.prefetch_factor,
|
||||
)
|
||||
self._context.dataloader = dataloader
|
||||
return self
|
||||
|
||||
def with_strategy(self) -> Self:
|
||||
self._context.strategy = StrategyFactory.create(
|
||||
model=self._context.model,
|
||||
train_type=self.config.strategy,
|
||||
device=get_current_device(),
|
||||
**self.config.extra_kwargs,
|
||||
)
|
||||
return self
|
||||
|
||||
def build(self) -> TrainContext:
|
||||
return self._context
|
||||
@@ -0,0 +1,105 @@
|
||||
import logging
|
||||
from typing import Optional, List
|
||||
from astrai.config import TrainConfig
|
||||
from astrai.trainer.train_callback import (
|
||||
TrainCallback,
|
||||
ProgressBarCallback,
|
||||
CheckpointCallback,
|
||||
MetricLoggerCallback,
|
||||
GradientClippingCallback,
|
||||
SchedulerCallback,
|
||||
)
|
||||
from astrai.trainer.train_context import TrainContext, TrainContextBuilder
|
||||
from astrai.data.serialization import Checkpoint
|
||||
from astrai.parallel.setup import spawn_parallel_fn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class Trainer:
|
||||
def __init__(
|
||||
self, train_config: TrainConfig, callbacks: Optional[List[TrainCallback]] = None
|
||||
):
|
||||
self.train_config = train_config
|
||||
default_callbacks = self._get_default_callbacks()
|
||||
self.callbacks = (
|
||||
default_callbacks + callbacks if callbacks else default_callbacks
|
||||
)
|
||||
|
||||
def _get_default_callbacks(self) -> List[TrainCallback]:
|
||||
train_config = self.train_config
|
||||
return [
|
||||
ProgressBarCallback(train_config.n_epoch),
|
||||
CheckpointCallback(train_config.ckpt_dir, train_config.ckpt_interval),
|
||||
MetricLoggerCallback(train_config.ckpt_dir, train_config.ckpt_interval),
|
||||
GradientClippingCallback(train_config.max_grad_norm),
|
||||
SchedulerCallback(),
|
||||
]
|
||||
|
||||
def _build_context(self, checkpoint: Optional[Checkpoint]) -> TrainContext:
|
||||
return (
|
||||
TrainContextBuilder(self.train_config)
|
||||
.with_checkpoint(checkpoint)
|
||||
.with_dataloader()
|
||||
.with_strategy()
|
||||
.build()
|
||||
)
|
||||
|
||||
def _call_callbacks(self, method_name: str, context: TrainContext):
|
||||
for callback in self.callbacks:
|
||||
method = getattr(callback, method_name, None)
|
||||
if method:
|
||||
method(context)
|
||||
|
||||
def train(self, checkpoint: Optional[Checkpoint] = None):
|
||||
config = self.train_config
|
||||
spawn_parallel_fn(
|
||||
self._train_impl,
|
||||
backend=config.backend,
|
||||
world_size=config.nprocs,
|
||||
master_addr=config.master_addr,
|
||||
master_port=config.master_port,
|
||||
device_type=config.device_type,
|
||||
device_ids=config.device_ids,
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
def _train_impl(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
|
||||
context = self._build_context(checkpoint)
|
||||
self._call_callbacks("on_train_begin", context)
|
||||
|
||||
try:
|
||||
context.model.train()
|
||||
# 1.epoch
|
||||
for epoch in range(context.epoch, self.train_config.n_epoch):
|
||||
context.epoch = epoch
|
||||
self._call_callbacks("on_epoch_begin", context)
|
||||
|
||||
for batch in context.dataloader:
|
||||
if context.iteration % self.train_config.accumulation_steps == 0:
|
||||
# 2. step
|
||||
self._call_callbacks("on_step_begin", context)
|
||||
context.optimizer.step()
|
||||
context.optimizer.zero_grad()
|
||||
self._call_callbacks("on_step_end", context)
|
||||
|
||||
# 3. batch
|
||||
self._call_callbacks("on_batch_begin", context)
|
||||
loss = context.strategy(batch)
|
||||
context.loss = loss.item()
|
||||
context.iteration += 1
|
||||
|
||||
# to make the loss normalized by accumulation steps
|
||||
stand_loss = loss / self.train_config.accumulation_steps
|
||||
stand_loss.backward()
|
||||
|
||||
self._call_callbacks("on_batch_end", context)
|
||||
|
||||
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", context)
|
||||
Reference in New Issue
Block a user