FilesExpand file tree

main

/

activations.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

178 lines (137 loc) · 6.34 KB

main

/

activations.py

Copy path

Top

File metadata and controls

• Code
• Blame

178 lines (137 loc) · 6.34 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

# coding=utf-8

# Copyright 2024 HuggingFace Inc.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import torch

import torch.nn.functional as F

from torch import nn

from ..utils import deprecate

from ..utils.import_utils import is_torch_npu_available, is_torch_version

if is_torch_npu_available():

import torch_npu

ACTIVATION_FUNCTIONS = {

"swish": nn.SiLU(),

"silu": nn.SiLU(),

"mish": nn.Mish(),

"gelu": nn.GELU(),

"relu": nn.ReLU(),

}

def get_activation(act_fn: str) -> nn.Module:

"""Helper function to get activation function from string.

Args:

act_fn (str): Name of activation function.

Returns:

nn.Module: Activation function.

"""

act_fn = act_fn.lower()

if act_fn in ACTIVATION_FUNCTIONS:

return ACTIVATION_FUNCTIONS[act_fn]

else:

raise ValueError(f"Unsupported activation function: {act_fn}")

class FP32SiLU(nn.Module):

r"""

SiLU activation function with input upcasted to torch.float32.

"""

def __init__(self):

super().__init__()

def forward(self, inputs: torch.Tensor) -> torch.Tensor:

return F.silu(inputs.float(), inplace=False).to(inputs.dtype)

class GELU(nn.Module):

r"""

GELU activation function with tanh approximation support with `approximate="tanh"`.

Parameters:

dim_in (`int`): The number of channels in the input.

dim_out (`int`): The number of channels in the output.

approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.

bias (`bool`, defaults to True): Whether to use a bias in the linear layer.

"""

def __init__(self, dim_in: int, dim_out: int, approximate: str = "none", bias: bool = True):

super().__init__()

self.proj = nn.Linear(dim_in, dim_out, bias=bias)

self.approximate = approximate

def gelu(self, gate: torch.Tensor) -> torch.Tensor:

if gate.device.type == "mps" and is_torch_version("<", "2.0.0"):

# fp16 gelu not supported on mps before torch 2.0

return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(dtype=gate.dtype)

return F.gelu(gate, approximate=self.approximate)

def forward(self, hidden_states):

hidden_states = self.proj(hidden_states)

hidden_states = self.gelu(hidden_states)

return hidden_states

class GEGLU(nn.Module):

r"""

A [variant](https://arxiv.org/abs/2002.05202) of the gated linear unit activation function.

Parameters:

dim_in (`int`): The number of channels in the input.

dim_out (`int`): The number of channels in the output.

bias (`bool`, defaults to True): Whether to use a bias in the linear layer.

"""

def __init__(self, dim_in: int, dim_out: int, bias: bool = True):

super().__init__()

self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias)

def gelu(self, gate: torch.Tensor) -> torch.Tensor:

if gate.device.type == "mps" and is_torch_version("<", "2.0.0"):

# fp16 gelu not supported on mps before torch 2.0

return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)

return F.gelu(gate)

def forward(self, hidden_states, *args, **kwargs):

if len(args) > 0 or kwargs.get("scale", None) is not None:

deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."

deprecate("scale", "1.0.0", deprecation_message)

hidden_states = self.proj(hidden_states)

if is_torch_npu_available():

# using torch_npu.npu_geglu can run faster and save memory on NPU.

return torch_npu.npu_geglu(hidden_states, dim=-1, approximate=1)[0]

else:

hidden_states, gate = hidden_states.chunk(2, dim=-1)

return hidden_states * self.gelu(gate)

class SwiGLU(nn.Module):

r"""

A [variant](https://arxiv.org/abs/2002.05202) of the gated linear unit activation function. It's similar to `GEGLU`

but uses SiLU / Swish instead of GeLU.

Parameters:

dim_in (`int`): The number of channels in the input.

dim_out (`int`): The number of channels in the output.

bias (`bool`, defaults to True): Whether to use a bias in the linear layer.

"""

def __init__(self, dim_in: int, dim_out: int, bias: bool = True):

super().__init__()

self.proj = nn.Linear(dim_in, dim_out * 2, bias=bias)

self.activation = nn.SiLU()

def forward(self, hidden_states):

hidden_states = self.proj(hidden_states)

hidden_states, gate = hidden_states.chunk(2, dim=-1)

return hidden_states * self.activation(gate)

class ApproximateGELU(nn.Module):

r"""

The approximate form of the Gaussian Error Linear Unit (GELU). For more details, see section 2 of this

[paper](https://arxiv.org/abs/1606.08415).

Parameters:

dim_in (`int`): The number of channels in the input.

dim_out (`int`): The number of channels in the output.

bias (`bool`, defaults to True): Whether to use a bias in the linear layer.

"""

def __init__(self, dim_in: int, dim_out: int, bias: bool = True):

super().__init__()

self.proj = nn.Linear(dim_in, dim_out, bias=bias)

def forward(self, x: torch.Tensor) -> torch.Tensor:

x = self.proj(x)

return x * torch.sigmoid(1.702 * x)

class LinearActivation(nn.Module):

def __init__(self, dim_in: int, dim_out: int, bias: bool = True, activation: str = "silu"):

super().__init__()

self.proj = nn.Linear(dim_in, dim_out, bias=bias)

self.activation = get_activation(activation)

def forward(self, hidden_states):

hidden_states = self.proj(hidden_states)

return self.activation(hidden_states)