File size: 5,071 Bytes
b66f552 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 | # Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
from __future__ import annotations
from functools import partial
from typing import TYPE_CHECKING, Any
import torch
import torch.nn as nn
try:
from torch.distributed import DeviceMesh
except ImportError:
DeviceMesh = None
try:
from torch.distributed.tensor import Placement, Replicate, Shard, distribute_module
except ImportError:
Placement = None
Replicate = None
Shard = None
distribute_module = None
try:
from torch.distributed.tensor.parallel import ParallelStyle
except ImportError:
class ParallelStyle:
pass
from fla.modules.activations import swiglu, swiglu_linear
try:
from torch.distributed.tensor import DTensor
except (ImportError, AttributeError):
DTensor = None
if TYPE_CHECKING:
from transformers.processing_utils import Unpack
class GatedMLP(nn.Module):
def __init__(
self,
hidden_size: int,
hidden_ratio: int | None = None,
intermediate_size: int | None = None,
hidden_act: str = 'swish',
fuse_swiglu: bool = True,
) -> GatedMLP:
super().__init__()
self.hidden_size = hidden_size
# the final number of params is `hidden_ratio * hidden_size^2`
# `intermediate_size` is chosen to be a multiple of 256 closest to `2/3 * hidden_size * hidden_ratio`
if hidden_ratio is None:
hidden_ratio = 4
if intermediate_size is None:
intermediate_size = int(hidden_size * hidden_ratio * 2 / 3)
intermediate_size = 256 * ((intermediate_size + 256 - 1) // 256)
self.hidden_ratio = hidden_ratio
self.intermediate_size = intermediate_size
self.hidden_act = hidden_act
self.fuse_swiglu = fuse_swiglu
if hidden_act != 'swish':
raise ValueError(f'Unsupported hidden_act: {hidden_act}')
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
if self.fuse_swiglu:
self.swiglu_linear = SwiGLULinear()
def forward(
self,
x: torch.Tensor,
**kwargs: Unpack[Any],
) -> torch.Tensor:
gate, y = self.gate_proj(x), self.up_proj(x)
if self.fuse_swiglu:
return self.swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias)
else:
return self.down_proj(swiglu(gate, y))
class SwiGLULinear(nn.Module):
def forward(self, x, y, weight, bias):
return swiglu_linear(x, y, weight, bias)
class SwiGLULinearParallel(ParallelStyle):
def __init__(
self,
*,
input_layouts: Placement | None = None,
output_layouts: Placement | None = None,
use_local_output: bool = True,
):
super().__init__()
self.input_layouts = (input_layouts or Shard(-1),)
self.output_layouts = (output_layouts or Replicate(),)
self.desired_input_layouts = (Shard(-1),)
self.use_local_output = use_local_output
@staticmethod
def _prepare_input_fn(
input_layouts, desired_input_layouts, mod, inputs, device_mesh,
):
x, y, weight, bias = inputs
if not isinstance(x, DTensor):
x = DTensor.from_local(x, device_mesh, input_layouts, run_check=False)
if x.placements != desired_input_layouts:
x = x.redistribute(placements=desired_input_layouts, async_op=True)
if not isinstance(y, DTensor):
y = DTensor.from_local(y, device_mesh, input_layouts, run_check=False)
if y.placements != desired_input_layouts:
y = y.redistribute(placements=desired_input_layouts, async_op=True)
if not isinstance(weight, DTensor):
weight = DTensor.from_local(weight, device_mesh, (Shard(1),))
if bias is not None and not isinstance(bias, DTensor):
bias = DTensor.from_local(bias, device_mesh, (Replicate(),))
return x, y, weight, bias
@staticmethod
def _prepare_output_fn(output_layouts, use_local_output, mod, outputs, device_mesh):
# Rowwise sharding produces partial output, depending on output layouts:
# 1. to replicate -> allreduce
# 2. to shard -> reduce_scatter
if outputs.placements != output_layouts:
outputs = outputs.redistribute(placements=output_layouts, async_op=True)
# back to local tensor if use_local_output is True
return outputs.to_local() if use_local_output else outputs
def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module:
return distribute_module(
module,
device_mesh,
partition_fn=None,
input_fn=partial(self._prepare_input_fn, self.input_layouts, self.desired_input_layouts),
output_fn=partial(self._prepare_output_fn, self.output_layouts, self.use_local_output),
)
|