feat(trainer): 支持分布式训练配置与检查点加载优化

khaosz/config/train_config.py

@@ -4,7 +4,7 @@ from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Callable, Optional
 
 
 @dataclass
@@ -88,6 +88,22 @@ class TrainConfig:
         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_fn: Optional[Callable] = field(
+        default=None,
+        metadata={"help": "Parallel function for training."}
+    )
     
     # others
     extra_kwargs: dict = field(

khaosz/trainer/__init__.py

@@ -1,4 +1,5 @@
 from khaosz.trainer.trainer import Trainer
+from khaosz.trainer.checkpoint import Checkpoint
 from khaosz.trainer.strategy import StrategyFactory
 from khaosz.trainer.schedule import SchedulerFactory
 
@@ -15,6 +16,9 @@ __all__ = [
     # trainer
     "Trainer",
     
+    # checkpoint
+    "Checkpoint",
+    
     # factory
     "StrategyFactory",
     "SchedulerFactory",

khaosz/trainer/checkpoint.py

@@ -38,19 +38,23 @@ class Checkpoint:
         if save_metric_plot and self.metrics:
             self._plot_metrics()
     
-    def load(self, save_dir: str) -> "Checkpoint":
-        if not os.path.exists(save_dir):
-            raise FileNotFoundError(f"Checkpoint directory {save_dir} does not exist.")
+    @classmethod
+    def load(cls, save_dir: str) -> "Checkpoint":
+        checkpoint_path = os.path.join(save_dir, "train_state.pkl")
         
-        with open(os.path.join(save_dir, "train_state.pkl"), "rb") as f:
+        if not os.path.exists(checkpoint_path):
+            raise FileNotFoundError(f"Checkpoint file {checkpoint_path} does not exist.")
+        
+        with open(checkpoint_path, "rb") as f:
             train_state = pkl.load(f)
-            self.epoch = train_state["epoch"]
-            self.iteration = train_state["iteration"]
-            self.metrics = train_state["metrics"]
-            self.optimizer_state = train_state["optimizer_state"]
-            self.scheduler_state = train_state["scheduler_state"]
         
-        return self
+        return cls(
+            optimizer_state=train_state["optimizer_state"],
+            scheduler_state=train_state["scheduler_state"],
+            epoch=train_state["epoch"],
+            iteration=train_state["iteration"],
+            metrics=train_state["metrics"]
+        )
     
     def _plot_metrics(self):
         for metric_name, metric_value in self.metrics.items():

khaosz/trainer/train_callback.py

@@ -99,7 +99,7 @@ class CheckpointCallback(TrainCallback):
     
     @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}")
+        save_path = os.path.join(self.save_dir, f"epoch_{context.epoch}_iter_{context.iteration}")
         context.checkpoint = Checkpoint(
             context.optimizer.state_dict(), 
             context.scheduler.state_dict(), 

khaosz/trainer/train_context.py

@@ -89,5 +89,12 @@ class TrainContextBuilder:
         )
         return self
     
+    def with_parallel_fn(self) -> Self:
+        fn = self.config.parallel_fn
+        if fn is not None:
+            self._context.model = fn(self._context.model)
+            
+        return self
+    
     def build(self) -> TrainContext:
         return self._context

khaosz/trainer/trainer.py

@@ -10,6 +10,7 @@ from khaosz.trainer.train_callback import (
 )
 from khaosz.trainer.train_context import TrainContext, TrainContextBuilder
 from khaosz.trainer.checkpoint import Checkpoint
+from khaosz.parallel.setup import spawn_parallel_fn
 
 logger = logging.getLogger(__name__)
 
@@ -37,6 +38,7 @@ class Trainer:
                 .with_checkpoint(checkpoint)
                 .with_dataloader()
                 .with_strategy()
+                .with_parallel_fn()
                 .build())
     
     def _call_callbacks(self, method_name: str, context: TrainContext):
@@ -44,8 +46,19 @@ class Trainer:
             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,
+            checkpoint=checkpoint
+        )
 
-    def train(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
+    def _train_impl(self, checkpoint: Optional[Checkpoint] = None) -> Checkpoint:
         context = self._build_context(checkpoint)
         self._call_callbacks('on_train_begin', context)
         
@@ -84,5 +97,4 @@ class Trainer:
             self._call_callbacks('on_error', context)
             raise
         finally:
-            self._call_callbacks('on_train_end', context)
-            return context.checkpoint
+            self._call_callbacks('on_train_end', context)