refactor: 优化参数传递，清理导入样式

astrai/__init__.py

@@ -5,17 +5,17 @@ from astrai.config import (
     ModelConfig,
     TrainConfig,
 )
-from astrai.model.transformer import Transformer
-from astrai.data import DatasetLoader, BpeTokenizer
+from astrai.data import BpeTokenizer, DatasetLoader
 from astrai.inference.generator import (
-    GenerationRequest,
-    LoopGenerator,
-    StreamGenerator,
     BatchGenerator,
     EmbeddingEncoder,
+    GenerationRequest,
     GeneratorFactory,
+    LoopGenerator,
+    StreamGenerator,
 )
-from astrai.trainer import Trainer, StrategyFactory, SchedulerFactory
+from astrai.model.transformer import Transformer
+from astrai.trainer import SchedulerFactory, StrategyFactory, Trainer
 
 __all__ = [
     "Transformer",

astrai/config/__init__.py

@@ -1,14 +1,7 @@
 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",
@@ -16,9 +9,4 @@ __all__ = [
     # Model configuration
     "ModelConfig",
     "TrainConfig",
-    # Schedule configuration
-    "ScheduleConfig",
-    "CosineScheduleConfig",
-    "SGDRScheduleConfig",
-    "ScheduleConfigFactory",
 ]

astrai/config/model_config.py

@@ -1,5 +1,4 @@
 import json
-
 from dataclasses import asdict, dataclass
 from typing import Optional, Self
 

astrai/config/param_config.py

@@ -1,13 +1,13 @@
-import torch.nn as nn
-import safetensors.torch as st
-
 from contextlib import contextmanager
 from dataclasses import dataclass, field
-from typing import Optional, Self, Union
 from pathlib import Path
+from typing import Self, Union
+
+import safetensors.torch as st
+import torch.nn as nn
 
-from astrai.data.tokenizer import BpeTokenizer
 from astrai.config.model_config import ModelConfig
+from astrai.data.tokenizer import BpeTokenizer
 from astrai.model.transformer import Transformer
 
 

astrai/config/schedule_config.py

@@ -1,149 +0,0 @@
-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())

astrai/config/train_config.py

@@ -1,11 +1,11 @@
-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
 
+import torch.nn as nn
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+from torch.utils.data import Dataset
+
 
 @dataclass
 class TrainConfig:

astrai/data/__init__.py

@@ -1,16 +1,15 @@
 from astrai.data.dataset import (
     BaseDataset,
-    SEQDataset,
+    DatasetFactory,
+    DatasetLoader,
     DPODataset,
-    SFTDataset,
     GRPODataset,
     MultiSegmentFetcher,
-    DatasetLoader,
-    DatasetFactory,
+    SEQDataset,
+    SFTDataset,
 )
-
-from astrai.data.tokenizer import BpeTokenizer
 from astrai.data.sampler import ResumableDistributedSampler
+from astrai.data.tokenizer import BpeTokenizer
 
 __all__ = [
     # Base classes

astrai/data/dataset.py

@@ -1,13 +1,14 @@
 """Dataset implementations with factory pattern for training."""
 
-import torch
 import bisect
-
 from abc import ABC, abstractmethod
+from typing import Dict, List, Optional, Union
+
+import torch
 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:

astrai/data/sampler.py

@@ -1,8 +1,8 @@
+from typing import Optional
+
 import torch
 import torch.distributed as dist
-
 from torch.utils.data import Dataset, Sampler
-from typing import Optional
 
 
 class ResumableDistributedSampler(Sampler[int]):

astrai/data/serialization.py

@@ -1,13 +1,14 @@
-import os
-import h5py
-import torch
 import json
-import safetensors.torch as st
-import torch.distributed as dist
-
+import os
 from pathlib import Path
-from torch import Tensor
 from typing import Any, Dict, List
+
+import h5py
+import safetensors.torch as st
+import torch
+import torch.distributed as dist
+from torch import Tensor
+
 from astrai.parallel.setup import get_rank
 
 

astrai/data/tokenizer.py

@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
-from tokenizers import Tokenizer, decoders, processors, normalizers, pre_tokenizers
+from typing import List, Union
+
+from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
 from tokenizers.models import BPE
 from tokenizers.trainers import BpeTrainer as BpeTrainerImpl
-from typing import List, Union
 
 
 class BaseTokenizer(ABC):

astrai/inference/__init__.py

@@ -1,16 +1,15 @@
 from astrai.inference.core import (
-    GeneratorCore,
     EmbeddingEncoderCore,
+    GeneratorCore,
     KVCacheManager,
 )
-
 from astrai.inference.generator import (
-    GenerationRequest,
-    LoopGenerator,
-    StreamGenerator,
     BatchGenerator,
     EmbeddingEncoder,
+    GenerationRequest,
     GeneratorFactory,
+    LoopGenerator,
+    StreamGenerator,
 )
 
 __all__ = [

astrai/inference/core.py

@@ -1,8 +1,9 @@
-import torch
+from typing import Any, Callable, List, Optional, Self, Tuple, Union
 
+import torch
 from torch import Tensor
-from typing import Any, Callable, List, Tuple, Union, Optional, Self
-from astrai.config import ModelParameter, ModelConfig
+
+from astrai.config import ModelConfig, ModelParameter
 
 
 def apply_sampling_strategies(

astrai/inference/generator.py

@@ -1,10 +1,11 @@
-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
+from typing import Generator, List, Optional, Tuple, Union
 
+import torch
+from torch import Tensor
+
+from astrai.config.param_config import ModelParameter
+from astrai.inference.core import EmbeddingEncoderCore, GeneratorCore, KVCacheManager
 
 HistoryType = List[Tuple[str, str]]
 

astrai/inference/server.py

@@ -1,13 +1,15 @@
-import torch
-import uvicorn
 import logging
 from pathlib import Path
-from typing import List, Optional, Dict, Any, Tuple
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import uvicorn
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, Field
+
 from astrai.config.param_config import ModelParameter
-from astrai.inference.generator import GeneratorFactory, GenerationRequest
+from astrai.inference.generator import GenerationRequest, GeneratorFactory
 
 logger = logging.getLogger(__name__)
 

astrai/model/__init__.py

@@ -1,9 +1,9 @@
 from astrai.model.module import (
+    GQA,
+    MLP,
+    DecoderBlock,
     Linear,
     RMSNorm,
-    MLP,
-    GQA,
-    DecoderBlock,
 )
 from astrai.model.transformer import Transformer
 

astrai/model/module.py

@@ -1,9 +1,9 @@
+from typing import Optional, Tuple
+
 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:

astrai/model/transformer.py

@@ -1,12 +1,13 @@
+from typing import Any, Mapping, Optional, Tuple
+
 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,
+    Embedding,
     Linear,
     RMSNorm,
     RotaryEmbedding,

astrai/parallel/__init__.py

@@ -1,14 +1,13 @@
+from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear
 from astrai.parallel.setup import (
-    get_world_size,
-    get_rank,
     get_current_device,
+    get_rank,
+    get_world_size,
     only_on_rank,
     setup_parallel,
     spawn_parallel_fn,
 )
 
-from astrai.parallel.module import RowParallelLinear, ColumnParallelLinear
-
 __all__ = [
     "get_world_size",
     "get_rank",

astrai/parallel/module.py

@@ -1,10 +1,10 @@
+from typing import Dict
+
 import torch
+import torch.distributed as dist
 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):

astrai/parallel/setup.py

@@ -1,12 +1,12 @@
 import os
+from contextlib import contextmanager
+from functools import wraps
+from typing import Callable, List, Optional
+
 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"]

astrai/trainer/__init__.py

@@ -1,15 +1,14 @@
-from astrai.trainer.trainer import Trainer
-from astrai.trainer.strategy import StrategyFactory, BaseStrategy
-from astrai.trainer.schedule import SchedulerFactory, BaseScheduler
-
+from astrai.trainer.schedule import BaseScheduler, SchedulerFactory
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 from astrai.trainer.train_callback import (
-    TrainCallback,
-    GradientClippingCallback,
-    SchedulerCallback,
     CheckpointCallback,
-    ProgressBarCallback,
+    GradientClippingCallback,
     MetricLoggerCallback,
+    ProgressBarCallback,
+    SchedulerCallback,
+    TrainCallback,
 )
+from astrai.trainer.trainer import Trainer
 
 __all__ = [
     # Main trainer

astrai/trainer/metric_util.py

@@ -1,6 +1,7 @@
-import torch.nn as nn
 from typing import Dict
 
+import torch.nn as nn
+
 
 def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
     """Compute gradient norm for each parameter in the model."""

astrai/trainer/schedule.py

@@ -1,10 +1,10 @@
 """Learning rate scheduler implementations with factory pattern."""
 
 import math
-from abc import abstractmethod, ABC
+from abc import ABC, abstractmethod
 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):
@@ -37,10 +37,6 @@ class SchedulerFactory:
             ...
 
         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]] = {}
@@ -67,7 +63,7 @@ class SchedulerFactory:
         return decorator
 
     @classmethod
-    def create(cls, optimizer, schedule_type: str, **kwargs) -> BaseScheduler:
+    def create(cls, optimizer, schedule_type: str = "none", **kwargs) -> BaseScheduler:
         """Create a scheduler instance by type name.
 
         Args:
@@ -90,29 +86,13 @@ class SchedulerFactory:
         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
+# ----------- Scheduler implementations -----------
 
 
 @SchedulerFactory.register("cosine")

astrai/trainer/strategy.py

@@ -1,14 +1,14 @@
 """Training strategy implementations with factory pattern."""
 
 import copy
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Dict, Union
+
 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
+from torch.nn.parallel import DistributedDataParallel as DDP
 
 
 def unwrap_model(model: nn.Module) -> nn.Module:

astrai/trainer/train_callback.py

@@ -1,25 +1,25 @@
-import os
 import json
+import os
 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
 
+import torch.nn as nn
+from torch.nn.utils import clip_grad_norm_
+from tqdm import tqdm
+
+from astrai.data.serialization import Checkpoint
 from astrai.parallel import only_on_rank
 from astrai.trainer.metric_util import (
+    ctx_get_grad_max,
+    ctx_get_grad_mean,
+    ctx_get_grad_min,
+    ctx_get_grad_nan_num,
+    ctx_get_grad_norm,
+    ctx_get_grad_std,
     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
 
 

astrai/trainer/train_context.py

@@ -1,16 +1,16 @@
+from dataclasses import dataclass, field
+from typing import Optional, Self
+
 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.config.train_config import TrainConfig
 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
+from astrai.parallel.setup import get_current_device, get_rank, get_world_size
+from astrai.trainer.strategy import BaseStrategy, StrategyFactory
 
 
 @dataclass

astrai/trainer/trainer.py

@@ -1,17 +1,18 @@
 import logging
-from typing import Optional, List
+from typing import List, Optional
+
 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
+from astrai.trainer.train_callback import (
+    CheckpointCallback,
+    GradientClippingCallback,
+    MetricLoggerCallback,
+    ProgressBarCallback,
+    SchedulerCallback,
+    TrainCallback,
+)
+from astrai.trainer.train_context import TrainContext, TrainContextBuilder
 
 logger = logging.getLogger(__name__)