From 398e8a3ea31c9d077ce2cf56000e671722dc1de9 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 19 Aug 2026 16:17:40 +0800 Subject: [PATCH] refactor: deduplicate low-risk code paths --- astrai/config/train_config.py | 24 +++++++++++----------- astrai/dataset/dataset.py | 8 ++++---- astrai/dataset/storage.py | 2 +- astrai/inference/cache/strategy.py | 2 -- astrai/model/components/mlp.py | 8 +------- astrai/parallel/setup.py | 18 +++++++---------- astrai/serialization/checkpoint.py | 32 +++++++++++------------------- astrai/trainer/metric_util.py | 20 ++----------------- astrai/trainer/strategy.py | 17 ++-------------- astrai/trainer/train_callback.py | 21 ++++++++++---------- astrai/trainer/train_context.py | 1 - scripts/tools/train.py | 25 +++++++++++------------ 12 files changed, 64 insertions(+), 114 deletions(-) diff --git a/astrai/config/train_config.py b/astrai/config/train_config.py index 084d84a..0448333 100644 --- a/astrai/config/train_config.py +++ b/astrai/config/train_config.py @@ -11,10 +11,10 @@ from torch.utils.data import Dataset from astrai.config.base import BaseConfig from astrai.model.components.lora import LoRAConfig -_TRAIN_TYPES = frozenset({"seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"}) -_PARALLEL_MODES = frozenset({"none", "ddp", "fsdp"}) -_BACKENDS = frozenset({"nccl", "gloo"}) -_START_METHODS = frozenset({"spawn", "fork", "forkserver"}) +TRAIN_TYPES = frozenset({"seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"}) +PARALLEL_MODES = frozenset({"none", "ddp", "fsdp"}) +BACKENDS = frozenset({"nccl", "gloo"}) +START_METHODS = frozenset({"spawn", "fork", "forkserver"}) _COMPILE_MODES = frozenset({"default", "reduce-overhead", "max-autotune"}) @@ -129,31 +129,31 @@ class TrainConfig(BaseConfig): @field_validator("strategy") def _validate_strategy(cls, v: str) -> str: - if v not in _TRAIN_TYPES: + if v not in TRAIN_TYPES: raise ValueError( - f"strategy must be one of {sorted(_TRAIN_TYPES)}, got {v!r}" + f"strategy must be one of {sorted(TRAIN_TYPES)}, got {v!r}" ) return v @field_validator("parallel_mode") def _validate_parallel_mode(cls, v: str) -> str: - if v not in _PARALLEL_MODES: + if v not in PARALLEL_MODES: raise ValueError( - f"parallel_mode must be one of {sorted(_PARALLEL_MODES)}, got {v!r}" + f"parallel_mode must be one of {sorted(PARALLEL_MODES)}, got {v!r}" ) return v @field_validator("backend") def _validate_backend(cls, v: str) -> str: - if v not in _BACKENDS: - raise ValueError(f"backend must be one of {sorted(_BACKENDS)}, got {v!r}") + if v not in BACKENDS: + raise ValueError(f"backend must be one of {sorted(BACKENDS)}, got {v!r}") return v @field_validator("start_method") def _validate_start_method(cls, v: str) -> str: - if v not in _START_METHODS: + if v not in START_METHODS: raise ValueError( - f"start_method must be one of {sorted(_START_METHODS)}, got {v!r}" + f"start_method must be one of {sorted(START_METHODS)}, got {v!r}" ) return v diff --git a/astrai/dataset/dataset.py b/astrai/dataset/dataset.py index 4db6c7d..27e6bb3 100644 --- a/astrai/dataset/dataset.py +++ b/astrai/dataset/dataset.py @@ -383,10 +383,10 @@ class DatasetFactory(BaseFactory["BaseDataset"]): transform = _build_jsonl_transform(load_path, tokenizer_path) if transform is None: raise FileNotFoundError( - f"JSONL dataset config not found. Expected " - f"dataset_config.json alongside *.jsonl files, pass " - f"tokenizer_path= for the built-in messages config, or " - f"use processor= for lazy on-the-fly tokenisation." + "JSONL dataset config not found. Expected " + "dataset_config.json alongside *.jsonl files, pass " + "tokenizer_path= for the built-in messages config, or " + "use processor= for lazy on-the-fly tokenisation." ) store.load(load_path, transform=transform, **kwargs) else: diff --git a/astrai/dataset/storage.py b/astrai/dataset/storage.py index 51b7bd8..cfedab9 100644 --- a/astrai/dataset/storage.py +++ b/astrai/dataset/storage.py @@ -217,7 +217,7 @@ class Store(ABC): """ if self._window_size <= 0: raise RuntimeError("sample_window() requires window_size > 0 (stream mode)") - if self._window_size <= 0 or self._length <= self._window_size: + if self._length <= self._window_size: raise IndexError( f"Data too short for window: token_count={self._length}, " f"window_size={self._window_size}" diff --git a/astrai/inference/cache/strategy.py b/astrai/inference/cache/strategy.py index d6a584f..56355d9 100644 --- a/astrai/inference/cache/strategy.py +++ b/astrai/inference/cache/strategy.py @@ -16,8 +16,6 @@ from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import Callable, Dict, List, Optional, OrderedDict -import torch - from astrai.inference.cache.buffer import ReqToTokenPool # ---- data contract: per-task slot state ---- diff --git a/astrai/model/components/mlp.py b/astrai/model/components/mlp.py index 58c5d99..083dad3 100644 --- a/astrai/model/components/mlp.py +++ b/astrai/model/components/mlp.py @@ -29,12 +29,6 @@ class FFNOutput(TypedDict): router_stats: Optional[RouterStats] -class RoutedOutput(TypedDict): - hidden_states: Tensor - aux_loss: Optional[Tensor] - router_stats: Optional[RouterStats] - - @FFNFactory.register("mlp") class MLP(nn.Module): def __init__(self, dim: int, dim_ffn: int, down_init_std: float = 0.02): @@ -122,7 +116,7 @@ class DeepSeekMoE(nn.Module): / self.n_shared_experts ) - def _routed_forward(self, x: Tensor, include_aux_loss: bool) -> RoutedOutput: + def _routed_forward(self, x: Tensor, include_aux_loss: bool) -> FFNOutput: N, D = x.shape K = self.n_activated_experts E = self.n_routed_experts diff --git a/astrai/parallel/setup.py b/astrai/parallel/setup.py index e8d931b..e79eacb 100644 --- a/astrai/parallel/setup.py +++ b/astrai/parallel/setup.py @@ -247,18 +247,15 @@ class LocalStrategy(LaunchStrategy): ctx.join() -def _detect_launcher() -> str: - """Detect the distributed launcher from environment. - - Returns one of: "torchelastic", "torchrun", "external", "local". - """ +def _is_external_launcher() -> bool: + """Whether an external launcher (torchrun/elastic/manual env) started us.""" if dist.is_torchelastic_launched(): - return "torchelastic" + return True if "LOCAL_WORLD_SIZE" in os.environ: - return "torchrun" + return True if "RANK" in os.environ and "WORLD_SIZE" in os.environ: - return "external" - return "local" + return True + return False def spawn_parallel_fn( @@ -273,8 +270,7 @@ def spawn_parallel_fn( ): if master_port is None: master_port = find_free_port() - launcher = _detect_launcher() - if launcher in ("torchelastic", "torchrun", "external"): + if _is_external_launcher(): strategy = TorchrunStrategy( world_size, backend, master_addr, master_port, device_type, start_method ) diff --git a/astrai/serialization/checkpoint.py b/astrai/serialization/checkpoint.py index ea5a9cb..2bef9ae 100644 --- a/astrai/serialization/checkpoint.py +++ b/astrai/serialization/checkpoint.py @@ -5,7 +5,7 @@ import json import time from dataclasses import dataclass, field from pathlib import Path -from typing import Any, Dict, Optional, Union +from typing import Any, Callable, Dict, Optional, Union import safetensors.torch as st import torch @@ -22,39 +22,31 @@ def save_safetensors(state_dict: dict, path: Union[str, Path]): st.save_file(state_dict, str(path)) -def load_safetensors(path: Union[str, Path], broadcast: bool = False) -> dict: +def _broadcast_load(loader: Callable[[], dict], broadcast: bool) -> dict: + """Load on rank 0 and broadcast the object to all ranks.""" if not broadcast or not dist.is_initialized(): - return st.load_file(str(path)) - + return loader() rank = get_rank() if rank == 0: - state_dict = st.load_file(str(path)) + data = loader() else: - state_dict = {} - tmp = [state_dict] + data = {} + tmp = [data] dist.broadcast_object_list(tmp, src=0) return tmp[0] +def load_safetensors(path: Union[str, Path], broadcast: bool = False) -> dict: + return _broadcast_load(lambda: st.load_file(str(path)), broadcast) + + def save_json(data: dict, path: Union[str, Path]): with open(str(path), "w") as f: json.dump(data, f, indent=2) def load_json(path: Union[str, Path], broadcast: bool = False) -> dict: - if not broadcast or not dist.is_initialized(): - with open(str(path), "r") as f: - return json.load(f) - - rank = get_rank() - if rank == 0: - with open(str(path), "r") as f: - data = json.load(f) - else: - data = {} - tmp = [data] - dist.broadcast_object_list(tmp, src=0) - return tmp[0] + return _broadcast_load(lambda: json.loads(Path(path).read_text()), broadcast) def save_torch(obj: Any, path: Union[str, Path]): diff --git a/astrai/trainer/metric_util.py b/astrai/trainer/metric_util.py index 8a943dd..63fc2ff 100644 --- a/astrai/trainer/metric_util.py +++ b/astrai/trainer/metric_util.py @@ -94,21 +94,5 @@ def ctx_get_grad_snr(ctx): return tracker.snr -def ctx_get_moe_aux_loss(ctx): - return ctx.strategy._moe_metrics.get("aux_loss") - - -def ctx_get_router_entropy(ctx): - return ctx.strategy._moe_metrics.get("router_entropy") - - -def ctx_get_dead_expert_fraction(ctx): - return ctx.strategy._moe_metrics.get("dead_expert_fraction") - - -def ctx_get_load_imbalance_mean(ctx): - return ctx.strategy._moe_metrics.get("load_imbalance_mean") - - -def ctx_get_load_imbalance_max(ctx): - return ctx.strategy._moe_metrics.get("load_imbalance_max") +def ctx_get_moe_metric(ctx, key): + return ctx.strategy._moe_metrics.get(key) diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index 0689ae7..480284c 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -1,6 +1,6 @@ """Training strategy implementations with factory pattern.""" -from abc import ABC, abstractmethod +from abc import ABC from typing import Callable, Dict, List, Optional, TypedDict, Union import torch @@ -187,7 +187,6 @@ class BaseStrategy(ABC): self.extra_kwargs = kwargs self._rollout_runner = None - @abstractmethod def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor: """Compute loss for the given batch. @@ -197,7 +196,7 @@ class BaseStrategy(ABC): Returns: Computed loss tensor """ - raise NotImplementedError + return self.compute_loss_output(batch)["loss"] def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput: return self._normalize_output(self.compute_loss(batch)) @@ -328,9 +327,6 @@ class SEQStrategy(BaseStrategy): super().__init__(model, device, **kwargs) self.label_smoothing = label_smoothing - def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor: - return self.compute_loss_output(batch)["loss"] - def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput: batch = move_to_device(batch, self.device) input_ids, target_ids = batch["input_ids"], batch["target_ids"] @@ -369,9 +365,6 @@ class SFTStrategy(BaseStrategy): super().__init__(model, device, **kwargs) self.label_smoothing = label_smoothing - def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor: - return self.compute_loss_output(batch)["loss"] - def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput: batch = move_to_device(batch, self.device) input_ids, target_ids, position_ids, loss_mask = ( @@ -426,9 +419,6 @@ class DPOStrategy(BaseStrategy): self.beta = beta self.reduction = reduction - def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor: - return self.compute_loss_output(batch)["loss"] - def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput: batch = move_to_device(batch, self.device) chosen_ids, rejected_ids = batch["chosen"], batch["rejected"] @@ -553,9 +543,6 @@ class GRPOStrategy(BaseStrategy): if state_dict is not None: self.old_model.load_state_dict(state_dict) - def compute_loss(self, batch: Dict[str, Tensor]) -> Tensor: - return self.compute_loss_output(batch)["loss"] - def compute_loss_output(self, batch: Dict[str, Tensor]) -> LossOutput: batch = move_to_device(batch, self.device) prompts = batch["prompts"] diff --git a/astrai/trainer/train_callback.py b/astrai/trainer/train_callback.py index 7cf6842..f5fe057 100644 --- a/astrai/trainer/train_callback.py +++ b/astrai/trainer/train_callback.py @@ -3,6 +3,7 @@ import logging import os import sys import time +from functools import partial from pathlib import Path from typing import IO, Callable, List, Optional, Protocol, runtime_checkable @@ -17,15 +18,11 @@ from astrai.parallel import only_on_rank from astrai.parallel.setup import get_current_device from astrai.serialization import Checkpoint from astrai.trainer.metric_util import ( - ctx_get_dead_expert_fraction, ctx_get_grad_norm, ctx_get_grad_snr, - ctx_get_load_imbalance_max, - ctx_get_load_imbalance_mean, ctx_get_loss, ctx_get_lr, - ctx_get_moe_aux_loss, - ctx_get_router_entropy, + ctx_get_moe_metric, ctx_get_val_loss, ) from astrai.trainer.train_context import TrainContext @@ -262,11 +259,15 @@ class MetricCallback(TrainCallback): "val_loss": ctx_get_val_loss, "grad_norm": ctx_get_grad_norm, "grad_snr": ctx_get_grad_snr, - "moe_aux_loss": ctx_get_moe_aux_loss, - "router_entropy": ctx_get_router_entropy, - "dead_expert_fraction": ctx_get_dead_expert_fraction, - "load_imbalance_mean": ctx_get_load_imbalance_mean, - "load_imbalance_max": ctx_get_load_imbalance_max, + "moe_aux_loss": partial(ctx_get_moe_metric, key="aux_loss"), + "router_entropy": partial(ctx_get_moe_metric, key="router_entropy"), + "dead_expert_fraction": partial( + ctx_get_moe_metric, key="dead_expert_fraction" + ), + "load_imbalance_mean": partial( + ctx_get_moe_metric, key="load_imbalance_mean" + ), + "load_imbalance_max": partial(ctx_get_moe_metric, key="load_imbalance_max"), } def _metrics(self, context: TrainContext, names): diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 99f65f5..64bf2c4 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -204,7 +204,6 @@ class TrainContextBuilder: def _create_dataloaders( self, context: TrainContext, train_dataset, val_dataset ) -> None: - cfg = self.config sampler_offset = context.consumed_samples // context.world_size if self._resume and sampler_offset > 0: samples_per_replica = ( diff --git a/scripts/tools/train.py b/scripts/tools/train.py index 7d8bb39..9eda7c6 100644 --- a/scripts/tools/train.py +++ b/scripts/tools/train.py @@ -11,6 +11,12 @@ from click.core import ParameterSource from torch import optim from astrai.config import AutoRegressiveLMConfig, TrainConfig +from astrai.config.train_config import ( + BACKENDS, + PARALLEL_MODES, + START_METHODS, + TRAIN_TYPES, +) from astrai.dataset import DatasetFactory, dpo_collate_fn, grpo_collate_fn from astrai.model import AutoRegressiveLM from astrai.model.components.decoder_block import DecoderBlock @@ -92,12 +98,12 @@ def _merge_yaml_into_kwargs( return merged -_TRAIN_TYPE = ["seq", "sft", "dpo", "grpo", "online_grpo", "online_dpo"] -_PARALLEL = ["none", "ddp", "fsdp"] +_TRAIN_TYPE = sorted(TRAIN_TYPES) +_PARALLEL = sorted(PARALLEL_MODES) _SCHEDULES = ["cosine", "sgdr", "wsd"] _OPTIMIZERS = OptimizerFactory.list_registered() -_BACKENDS = ["nccl", "gloo"] -_START_METHODS = ["spawn", "fork", "forkserver"] +_BACKENDS = sorted(BACKENDS) +_START_METHODS = sorted(START_METHODS) @click.command( @@ -651,17 +657,10 @@ def train( decay_steps: int, **kwargs, ): - if train_type not in [ - "seq", - "sft", - "dpo", - "grpo", - "online_grpo", - "online_dpo", - ]: + if train_type not in _TRAIN_TYPE: raise ValueError( f"Invalid train_type '{train_type}'. " - f"Must be one of: seq, sft, dpo, grpo, online_grpo, online_dpo" + f"Must be one of: {', '.join(_TRAIN_TYPE)}" ) if not os.path.exists(param_path): raise FileNotFoundError(f"Model directory not found: {param_path}")