diff --git a/astrai/parallel/__init__.py b/astrai/parallel/__init__.py index e4d490c..b15f85c 100644 --- a/astrai/parallel/__init__.py +++ b/astrai/parallel/__init__.py @@ -7,8 +7,8 @@ from astrai.parallel.executor import ( FSDPExecutor, GradientState, NoneExecutor, + create_ref_model, ) -from astrai.parallel.module import ColumnParallelLinear, RowParallelLinear from astrai.parallel.setup import ( get_current_device, get_rank, @@ -17,7 +17,6 @@ from astrai.parallel.setup import ( setup_parallel, spawn_parallel_fn, ) -from astrai.parallel.utils import create_ref_model __all__ = [ "get_world_size", @@ -26,8 +25,6 @@ __all__ = [ "only_on_rank", "setup_parallel", "spawn_parallel_fn", - "RowParallelLinear", - "ColumnParallelLinear", "ExecutorFactory", "BaseExecutor", "GradientState", diff --git a/astrai/parallel/executor.py b/astrai/parallel/executor.py index a8cada3..b8523e5 100644 --- a/astrai/parallel/executor.py +++ b/astrai/parallel/executor.py @@ -4,7 +4,7 @@ import contextlib import logging import os from contextlib import contextmanager -from typing import Any, Callable, Optional, Tuple +from typing import Any, Callable, Dict, Optional, Tuple import torch import torch.distributed as dist @@ -24,6 +24,31 @@ from astrai.parallel.setup import get_rank, get_world_size logger = logging.getLogger(__name__) +def create_ref_model( + model_fn: Callable[[], nn.Module], + executor: Optional["BaseExecutor"] = None, + model: Optional[nn.Module] = None, + state_dict: Optional[Dict[str, torch.Tensor]] = None, + device: Optional[str] = None, +) -> Optional[nn.Module]: + """Create a frozen reference model from executor or state dict. + + On non-rank-0, returns None (executor.unwrap_model returns None). + """ + if state_dict is None and executor is not None and model is not None: + state_dict = executor.unwrap_model(model) + if state_dict is None: + return None + + ref_model = model_fn() + ref_model.load_state_dict(state_dict) + ref_model.requires_grad_(False) + ref_model.eval() + if device is not None: + ref_model = ref_model.to(device=device) + return ref_model + + class GradientState: def __init__(self, grad_accum_steps: int = 1): self.num_steps = max(grad_accum_steps, 1) diff --git a/astrai/parallel/module.py b/astrai/parallel/module.py deleted file mode 100644 index 8e12493..0000000 --- a/astrai/parallel/module.py +++ /dev/null @@ -1,115 +0,0 @@ -from typing import Dict - -import torch -import torch.distributed as dist -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor - - -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) diff --git a/astrai/parallel/utils.py b/astrai/parallel/utils.py deleted file mode 100644 index 225031d..0000000 --- a/astrai/parallel/utils.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Utility functions for parallel training.""" - -from typing import TYPE_CHECKING, Callable, Dict, Optional - -import torch -import torch.nn as nn - -if TYPE_CHECKING: - from astrai.parallel.executor import BaseExecutor - - -def create_ref_model( - model_fn: Callable[[], nn.Module], - executor: Optional["BaseExecutor"] = None, - model: Optional[nn.Module] = None, - state_dict: Optional[Dict[str, torch.Tensor]] = None, - device: Optional[str] = None, -) -> Optional[nn.Module]: - """Create a frozen reference model from executor or state dict. - - On non-rank-0, returns None (executor.unwrap_model returns None). - """ - if state_dict is None and executor is not None and model is not None: - state_dict = executor.unwrap_model(model) - if state_dict is None: - return None - - ref_model = model_fn() - ref_model.load_state_dict(state_dict) - ref_model.requires_grad_(False) - ref_model.eval() - if device is not None: - ref_model = ref_model.to(device=device) - return ref_model diff --git a/astrai/trainer/strategy.py b/astrai/trainer/strategy.py index cd67e57..9ddc422 100644 --- a/astrai/trainer/strategy.py +++ b/astrai/trainer/strategy.py @@ -9,7 +9,6 @@ import torch.nn.functional as F from torch import Tensor from astrai.factory import BaseFactory -from astrai.parallel.utils import create_ref_model from astrai.trainer.rollout import RolloutResult diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index 4b6251e..0ff7bfd 100644 --- a/astrai/trainer/train_context.py +++ b/astrai/trainer/train_context.py @@ -12,9 +12,8 @@ from astrai.config.train_config import TrainConfig from astrai.dataset import RDSampler from astrai.inference.core.scheduler import InferenceScheduler from astrai.model.components.lora import inject_lora -from astrai.parallel.executor import BaseExecutor, ExecutorFactory +from astrai.parallel.executor import BaseExecutor, ExecutorFactory, create_ref_model from astrai.parallel.setup import get_current_device, get_rank, get_world_size -from astrai.parallel.utils import create_ref_model from astrai.protocols import OptimizerProtocol, SchedulerProtocol from astrai.serialization import Checkpoint, load_json from astrai.tokenize import AutoTokenizer