| |
| |
|
|
| from typing import Optional |
|
|
| import torch |
| import torch.nn.functional as F |
| from flash_attn.ops.fused_dense import ColumnParallelLinear, RowParallelLinear |
| from flash_attn.utils.distributed import all_reduce, reduce_scatter |
| from torch import nn |
|
|
| from internlm.core.context import ParallelMode |
| from internlm.core.context import global_context as gpc |
| from internlm.model.utils import fused_dense_func_torch |
|
|
|
|
| class ScaleColumnParallelLinear(nn.Linear): |
| """ |
| ScaleColumnParallelLinear. |
| |
| Args: |
| in_features (int): size of each input sample |
| out_features (int): size of each output sample |
| process_group (Optional[torch.distributed.ProcessGroup]): The group of the current device for `parallel_mode`. |
| bias (bool): Whether the bias is needed for linears. True by default. But it is typically set to False |
| in the config. |
| sequence_parallel (bool): If sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: |
| we do an all_gather of x before doing the matmul. |
| If not, then the input is already gathered. |
| device (Optional[Union[str, torch.device]]): The device will be used. |
| dtype (Optional[torch.dtype]): The type of data. |
| weight_scale (int): For training stability. 1 by default. |
| """ |
|
|
| def __init__( |
| self, |
| in_features: int, |
| out_features: int, |
| process_group: Optional[torch.distributed.ProcessGroup], |
| bias: bool = True, |
| device: Optional[torch.device] = None, |
| dtype: Optional[torch.dtype] = None, |
| weight_scale: int = 1, |
| ) -> None: |
| world_size = torch.distributed.get_world_size(process_group) |
| if out_features % world_size != 0: |
| raise ValueError(f"out_features ({out_features}) must be divisible by " f"world_size ({world_size})") |
| super().__init__(in_features, out_features // world_size, bias=bias, device=device, dtype=dtype) |
| self.process_group = process_group |
| self.weight_scale = weight_scale |
|
|
| def forward(self, input): |
| |
| |
| |
| if self.weight_scale != 1: |
| weight = self.weight * self.weight_scale + (1 - self.weight_scale) * self.weight.detach() |
| else: |
| weight = self.weight |
| return fused_dense_func_torch( |
| input, |
| weight, |
| self.bias, |
| process_group=self.process_group, |
| sequence_parallel=gpc.config.parallel.sequence_parallel, |
| ) |
|
|
|
|
| class RewardModelLinear(ScaleColumnParallelLinear): |
| """ |
| RewardModelLinear. |
| Args: |
| in_features (int): size of each input sample |
| out_features (int): size of each output sample |
| process_group (Optional[torch.distributed.ProcessGroup]): The group of the current device for `parallel_mode`. |
| bias (bool): Whether the bias is needed for linears. True by default. But it is typically set to False |
| in the config. |
| sequence_parallel (bool): If sequence_parallel is True, we're doing Tensor Parallel with sequence parallelism: |
| we do an all_gather of x before doing the matmul. |
| If not, then the input is already gathered. |
| device (Optional[Union[str, torch.device]]): The device will be used. |
| dtype (Optional[torch.dtype]): The type of data. |
| weight_scale (int): For training stability. 1 by default. |
| """ |
|
|
| def __init__( |
| self, |
| in_features: int, |
| out_features: int, |
| process_group: Optional[torch.distributed.ProcessGroup], |
| bias: bool = True, |
| device: Optional[torch.device] = None, |
| dtype: Optional[torch.dtype] = None, |
| weight_scale: int = 1, |
| ) -> None: |
| super().__init__(in_features, out_features, process_group, bias, device, dtype, weight_scale) |
| torch.distributed.broadcast(self.weight, gpc.get_ranks_in_group(ParallelMode.TENSOR)[0], process_group) |
| if bias: |
| torch.distributed.broadcast(self.bias, gpc.get_ranks_in_group(ParallelMode.TENSOR)[0], process_group) |
|
|
| def forward(self, input): |
| |
| |
| |
| if self.weight_scale != 1: |
| weight = self.weight * self.weight_scale + (1 - self.weight_scale) * self.weight.detach() |
| else: |
| weight = self.weight |
| return fused_dense_func_torch( |
| input, |
| weight, |
| self.bias, |
| process_group=self.process_group, |
| sequence_parallel=gpc.config.parallel.sequence_parallel, |
| ) |
|
|
|
|
| class ColumnParallelLinearTorch(ColumnParallelLinear): |
| def forward(self, x): |
| |
| |
| |
|
|
| return fused_dense_func_torch( |
| x, self.weight, self.bias, process_group=self.process_group, sequence_parallel=self.sequence_parallel |
| ) |
|
|
|
|
| class RowParallelLinearTorch(RowParallelLinear): |
| def forward(self, x): |
| """ |
| We're doing Tensor Parallel with sequence parallelism: we do the matmul and then |
| a reduce_scatter of the result. |
| """ |
| out = fused_dense_func_torch(x, self.weight, self.bias) |
| reduce_fn = reduce_scatter if self.sequence_parallel else all_reduce |
| return reduce_fn(out, self.process_group) |
|
|
|
|
| class FeedForward(nn.Module): |
| """ |
| FeedForward. |
| |
| Args: |
| in_features (int): size of each input sample |
| hidden_features (int): size of hidden state of FFN |
| out_features (int): size of each output sample |
| process_group (Optional[torch.distributed.ProcessGroup]): The group of the current device for `parallel_mode`. |
| bias (bool): Whether the bias is needed for linears. True by default. But it is typically set to False |
| in the config. |
| device (Optional[Union[str, torch.device]]): The device will be used. |
| dtype (Optional[torch.dtype]): The type of data. |
| multiple_of (int): For efficient training. Reset the size of hidden feature. 256 by default. |
| """ |
|
|
| def __init__( |
| self, |
| in_features: int, |
| hidden_features: int, |
| out_features: int = None, |
| process_group: Optional[torch.distributed.ProcessGroup] = None, |
| bias: bool = True, |
| device: Optional[torch.device] = None, |
| dtype: Optional[torch.dtype] = None, |
| multiple_of: int = 256, |
| ): |
| super().__init__() |
|
|
| hidden_features = multiple_of * ((hidden_features + multiple_of - 1) // multiple_of) |
|
|
| self.w1 = ColumnParallelLinearTorch( |
| in_features, |
| hidden_features, |
| process_group, |
| bias, |
| sequence_parallel=gpc.config.parallel.sequence_parallel, |
| device=device, |
| dtype=dtype, |
| ) |
| self.w2 = ColumnParallelLinearTorch( |
| in_features, |
| hidden_features, |
| process_group, |
| bias, |
| sequence_parallel=gpc.config.parallel.sequence_parallel, |
| device=device, |
| dtype=dtype, |
| ) |
| self.w3 = RowParallelLinearTorch( |
| hidden_features, |
| out_features, |
| process_group, |
| bias=bias, |
| sequence_parallel=gpc.config.parallel.sequence_parallel, |
| device=device, |
| dtype=dtype, |
| ) |
|
|
| def forward(self, x): |
| out = self.w3(F.silu(self.w1(x)) * self.w2(x)) |
| return out |
|
|
