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
This commit is contained in:
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user