Kernels:
Trusted publisher
File size: 5,460 Bytes
e19323e | 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 | # Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
from __future__ import annotations
from functools import partial
from typing import TYPE_CHECKING, Any
import torch
import torch.nn as nn
from torch.distributed import DeviceMesh
from torch.distributed.tensor import Placement, Replicate, Shard, distribute_module
from torch.distributed.tensor.parallel import ParallelStyle
from ..modules.activations import powglu, powglu_linear, 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,
powglu_power: float = 3.0,
) -> 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
self.powglu_power = powglu_power
if hidden_act not in ('swish', 'powlu'):
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 and hidden_act == 'swish':
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.hidden_act == 'powlu':
if self.fuse_swiglu:
return powglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias, self.powglu_power)
return self.down_proj(powglu(gate, y, self.powglu_power))
if self.fuse_swiglu:
return self.swiglu_linear(gate, y, self.down_proj.weight, self.down_proj.bias)
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),
)
|