feat(train): 支持分布式训练的优化器与调度器工厂配置

khaosz/config/train_config.py

@@ -1,4 +1,4 @@
-from torch import nn
+import torch.nn as nn
 from torch.utils.data import Dataset
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
@@ -30,6 +30,7 @@ class TrainConfig:
         default=None,
         metadata={"help": "Scheduler for training."}
     )
+    
     n_epoch: int = field(
         default=1,
         metadata={"help": "Number of epochs for training."}
@@ -104,7 +105,15 @@ class TrainConfig:
         default=None,
         metadata={"help": "Parallel function for training."}
     )
-    
+    optimizer_factory: Optional[Callable[[nn.Module], Optimizer]] = field(
+        default=None,
+        metadata={"help": "Optimizer factory for training."}
+    )
+    scheduler_factory: Optional[Callable[[Optimizer], LRScheduler]] = field(
+        default=None,
+        metadata={"help": "Scheduler factory for training."}
+    )
+
     # others
     extra_kwargs: dict = field(
         default_factory=dict,
@@ -115,7 +124,17 @@ class TrainConfig:
         self.validate()
     
     def validate(self):
-        required_fields = ["model", "strategy", "dataset", "optimizer", "scheduler"]
+        required_fields = ["model", "strategy", "dataset"]
+        
         for field_name in required_fields:
             if getattr(self, field_name) is None:
                 raise ValueError(f"{field_name} is required.")
+        
+        factory_case = all([self.optimizer_factory, self.scheduler_factory])
+        argument_case = all([self.optimizer, self.scheduler])
+        self.nprocs = max(self.nprocs, 1)
+        
+        if self.nprocs > 1 and not factory_case:
+            raise ValueError("Distributed training requires optimizer and scheduler factories.")
+        elif self.nprocs == 1 and not argument_case:
+            raise ValueError("Single process training requires optimizer and scheduler arguments.")

khaosz/trainer/train_context.py

@@ -80,19 +80,27 @@ class TrainContextBuilder:
         return self
     
     def with_strategy(self) -> Self:
-        device = get_current_device()
         self._context.strategy = StrategyFactory.load(
             model=self.config.model,
             train_type=self.config.strategy,
-            device=device,
+            device=get_current_device(),
             **self.config.extra_kwargs
         )
         return self
     
     def with_parallel_fn(self) -> Self:
-        fn = self.config.parallel_fn
-        if fn is not None:
+        device = get_current_device()
+        self._context.model = self._context.model.to(device=device)
+        
+        if self.config.nprocs > 1:
+            
+            fn = self.config.parallel_fn
+            optimizer_fn = self.config.optimizer_factory
+            scheduler_fn = self.config.scheduler_factory
+            
             self._context.model = fn(self._context.model)
+            self._context.optimizer = optimizer_fn(self._context.model.parameters())
+            self._context.scheduler = scheduler_fn(self._context.optimizer)
             
         return self