Merge pull request #15 from EleutherAI/geglu-patch

remove bias from geglu
EleutherAI · Feb 8, 2021 · 960a4cd · 960a4cd
2 parents 51c4ecd + 1e84d33
commit 960a4cd
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -70,10 +70,13 @@ def __init__(self):
  elif args.onnx_safe:
  self.activation_func = erf_gelu
 
- def forward(self, x, bias):
+ def forward(self, x, bias=None):
  x, gate = x.chunk(2, dim=-1)
- bias_1, bias_2 = bias.chunk(2, dim=-1)
- x = x + bias_1
+ if bias is not None:
+ bias_1, bias_2 = bias.chunk(2, dim=-1)
+ x = x + bias_1
+ else:
+ bias_1 = bias_2 = 0
  if self.bias_gelu_fusion:
  intermediate_parallel = \
  bias_gelu_impl(gate, bias_2)
@@ -140,7 +143,7 @@ def forward(self, hidden_states):
  self.activation_func(intermediate_parallel + bias_parallel)
  elif self.activation_type == "geglu":
  intermediate_parallel = \
- self.activation_func(intermediate_parallel, bias_parallel)
+ self.activation_func(intermediate_parallel)
  else:
  raise ValueError(f'Activation type {self.activation_type} not recognized')