# This implementation was adapted from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/modules/mlp.py # Commit id: c3b219665292c61a51153d0ded4473c494296382 # Copyright (c) 2023, Tri Dao. import torch import torch.nn as nn import torch.nn.functional as F from torch.distributed import ProcessGroup try: from flash_attn.ops.activations import swiglu except ImportError: swiglu = None try: from flash_attn.ops.fused_dense import ColumnParallelLinear, RowParallelLinear except ImportError: ColumnParallelLinear, RowParallelLinear = None, None try: from flash_attn.ops.fused_dense import FusedMLP, ParallelFusedMLP except ImportError: FusedMLP, ParallelFusedMLP = None, None class Mlp(nn.Module): def __init__( self, in_features, hidden_features=None, out_features=None, activation=F.gelu, bias1=True, bias2=True, return_residual=False, device=None, dtype=None, ): factory_kwargs = {"device": device, "dtype": dtype} super().__init__() out_features = out_features if out_features is not None else in_features hidden_features = hidden_features if hidden_features is not None else in_features * 4 self.return_residual = return_residual self.fc1 = nn.Linear(in_features, hidden_features, bias=bias1, **factory_kwargs) self.activation = activation self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) def forward(self, x, cu_adapter_mask=None): if cu_adapter_mask is not None: unique_tasks = torch.unique(cu_adapter_mask).tolist() fc1_dtype = next(self.fc1.parameters()).dtype y = torch.empty(x.shape[0], self.fc1.out_features, dtype=fc1_dtype).to(x.device) for task_id in unique_tasks: task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0] task_tensor = x[task_indices] task_y = self.fc1(task_tensor, task_id=task_id) y[task_indices] = task_y else: y = self.fc1(x) y = self.activation(y) if cu_adapter_mask is not None: unique_tasks = torch.unique(cu_adapter_mask).tolist() fc2_dtype = next(self.fc2.parameters()).dtype out = torch.empty(y.shape[0], self.fc2.out_features, dtype=fc2_dtype).to(y.device) for task_id in unique_tasks: task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0] task_tensor = y[task_indices] task_out = self.fc2(task_tensor, task_id=task_id) out[task_indices] = task_out else: out = self.fc2(y) return out if not self.return_residual else (out, x) class ParallelMLP(nn.Module): def __init__( self, in_features, hidden_features=None, out_features=None, activation=F.gelu, process_group: ProcessGroup = None, sequence_parallel=True, bias1=True, bias2=True, device=None, dtype=None, ): factory_kwargs = {"device": device, "dtype": dtype} super().__init__() assert ColumnParallelLinear is not None, "Need to install fused_dense" assert RowParallelLinear is not None, "Need to install fused_dense" out_features = out_features if out_features is not None else in_features hidden_features = hidden_features if hidden_features is not None else in_features * 4 self.fc1 = ColumnParallelLinear( in_features, hidden_features, process_group, bias=bias1, sequence_parallel=sequence_parallel, **factory_kwargs, ) self.activation = activation self.fc2 = RowParallelLinear( hidden_features, out_features, process_group, bias=bias2, sequence_parallel=sequence_parallel, **factory_kwargs, ) def forward(self, x): y = self.fc1(x) y = self.activation(y) y = self.fc2(y) return y class GatedMlp(nn.Module): def __init__( self, in_features, hidden_features=None, out_features=None, activation=F.sigmoid, bias1=True, bias2=True, multiple_of=128, return_residual=False, device=None, dtype=None, ): factory_kwargs = {"device": device, "dtype": dtype} super().__init__() out_features = out_features if out_features is not None else in_features hidden_features = ( hidden_features if hidden_features is not None else int(8 * in_features / 3) ) hidden_features = (hidden_features + multiple_of - 1) // multiple_of * multiple_of self.return_residual = return_residual self.fc1 = nn.Linear(in_features, 2 * hidden_features, bias=bias1, **factory_kwargs) self.activation = activation self.fc2 = nn.Linear(hidden_features, out_features, bias=bias2, **factory_kwargs) def forward(self, x): y = self.fc1(x) if self.activation == F.sigmoid: # Special case for GLU y = F.glu(y, dim=-1) elif self.activation == F.silu and swiglu is not None: # Special case for SwiGLU y, gate = y.chunk(2, dim=-1) y = swiglu(gate, y) else: y, gate = y.chunk(2, dim=-1) y = y * self.activation(gate) y = self.fc2(y) return y if not self.return_residual else (y, x) class ParallelGatedMlp(nn.Module): """Parallel GatedMlp""" def __init__( self, in_features, process_group, hidden_features=None, out_features=None, activation=F.sigmoid, bias1=True, bias2=True, multiple_of=128, sequence_parallel=True, device=None, dtype=None, ): factory_kwargs = {"device": device, "dtype": dtype} super().__init__() out_features = out_features if out_features is not None else in_features hidden_features = ( hidden_features if hidden_features is not None else int(8 * in_features / 3) ) hidden_features = (hidden_features + multiple_of - 1) // multiple_of * multiple_of if ColumnParallelLinear is None or RowParallelLinear is None: raise ImportError("fused_dense is not installed") self.fc1 = ColumnParallelLinear( in_features, 2 * hidden_features, process_group, bias=bias1, sequence_parallel=sequence_parallel, **factory_kwargs, ) self.activation = activation self.fc2 = RowParallelLinear( hidden_features, out_features, process_group, bias=bias2, sequence_parallel=sequence_parallel, **factory_kwargs, ) def forward(self, x): y = self.fc1(x) if self.activation == F.sigmoid: # Special case for GLU y = F.glu(y, dim=-1) else: y, gate = y.chunk(2, dim=-1) y = y * self.activation(gate) y = self.fc2(y) return y