refactor: inline parallel utils into executor module

- Move create_ref_model from astrai/parallel/utils.py into executor.py
- Remove unused ColumnParallelLinear/RowParallelLinear (module.py)
- Update imports in strategy.py and train_context.py
- Drop unused astrai.parallel.utils and astrai.parallel.module

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",

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)

astrai/parallel/module.py

@@ -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)

astrai/parallel/utils.py

@@ -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

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
 
 

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