Add tensor parallelism for RWKV

megatron/model/gpt2_model.py

@@ -258,6 +258,7 @@ def init_specs(self):
  LayerSpec(
  RWKVResidualLayerPipe,
  neox_args=self.neox_args,
+ init_method=self.init_method,
  layer_number=i,
  )
  )

megatron/model/rwkv/v6/rwkv.py

@@ -7,7 +7,7 @@
 import torch.nn as nn
 from torch.nn import functional as F
 from torch.utils.cpp_extension import load
-
+from megatron import mpu
 
 class WKV(torch.autograd.Function):
  """
@@ -179,7 +179,7 @@ def __init__(self, neox_args, layer_number):
 
  self.value = nn.Linear(neox_args.hidden_size, neox_args.dim_att, bias=False)
  self.output = nn.Linear(neox_args.dim_att, neox_args.hidden_size, bias=False)
- self.gate = nn.Linear(neox_args.hidden_size, neox_args.dim_att, bias=False)
+ self.gate = nn.Linear(neox_args.hidden_size, neox_args.dim_att, bias=False) # column
  self.ln_x = nn.GroupNorm(
  neox_args.num_attention_heads, neox_args.dim_att, eps=(1e-5) * (8**2)
  )
@@ -228,15 +228,19 @@ def forward(self, x):
  return self.jit_func_2(x, g)
 
 
-class RWKV_ChannelMix(nn.Module):
+class ParallelRWKV_ChannelMix(nn.Module):
  """
  Channel Mix layer. The ffn in RWKV
  """
 
- def __init__(self, neox_args, layer_number):
+ def __init__(self, neox_args, layer_number, init_method):
  super().__init__()
  self.neox_args = neox_args
  self.layer_number = layer_number
+
+ world_size = mpu.get_model_parallel_world_size()
+ self.hidden_size_per_partition = mpu.divide(neox_args.hidden_size, world_size)
+
  self.time_shift = nn.ZeroPad2d((0, 0, 1, -1))
 
  with torch.no_grad(): # fancy init of time_mix
@@ -247,29 +251,54 @@ def __init__(self, neox_args, layer_number):
  self.time_maa_k = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
  self.time_maa_r = nn.Parameter(1.0 - torch.pow(ddd, ratio_1_to_almost0))
 
- self.key = nn.Linear(neox_args.hidden_size, neox_args.dim_ffn, bias=False)
- self.receptance = nn.Linear(
- neox_args.hidden_size, neox_args.hidden_size, bias=False
- )
- self.value = nn.Linear(neox_args.dim_ffn, neox_args.hidden_size, bias=False)
-
+ #self.key = nn.Linear(neox_args.hidden_size, neox_args.dim_ffn, bias=False)
+ self.key = mpu.ColumnParallelLinear(
+ neox_args=neox_args,
+ input_size=neox_args.hidden_size,
+ output_size=neox_args.dim_ffn,
+ gather_output=False,
+ init_method=init_method,
+ bias=False,
+ )
+ #self.receptance = nn.Linear(
+ # neox_args.hidden_size, neox_args.hidden_size, bias=False
+ #)
+ self.receptance = mpu.ColumnParallelLinear(
+ neox_args=neox_args,
+ input_size=neox_args.hidden_size,
+ output_size=neox_args.hidden_size,
+ gather_output=True,
+ init_method=init_method,
+ bias=False
+ )
+ #self.value = nn.Linear(neox_args.dim_ffn, neox_args.hidden_size, bias=False)
+ self.value = mpu.RowParallelLinear(
+ neox_args=neox_args,
+ input_size=neox_args.dim_ffn,
+ output_size=neox_args.hidden_size,
+ input_is_parallel=True,
+ init_method=init_method,
+ parallel_output=False,
+ bias=False
+ )
  def forward(self, x):
  xx = self.time_shift(x) - x
  xk = x + xx * self.time_maa_k
  xr = x + xx * self.time_maa_r
 
- k = self.key(xk)
+ k, _ = self.key(xk)
  k = torch.relu(k) ** 2
- kv = self.value(k)
- return torch.sigmoid(self.receptance(xr)) * kv
+ kv, _ = self.value(k)
+ receptance, _ = self.receptance(xr)
+ return torch.sigmoid(receptance) * kv
 
 
 class RWKVResidualLayer(nn.Module):
  """
  RWKV layer definition
  """
 
- def __init__(self, neox_args, layer_number):
+ def __init__(self, neox_args, init_method, layer_number):
  super().__init__()
  self.neox_args = neox_args
  self.layer_number = layer_number
@@ -288,6 +317,7 @@ def __init__(self, neox_args, layer_number):
  self.num_attention_heads = neox_args.num_attention_heads
  assert neox_args.dim_att % self.num_attention_heads == 0
 
+ self.init_method = init_method
  if neox_args.attention_dropout > 0:
  self.drop0 = nn.Dropout(p=neox_args.attention_dropout)
 
@@ -296,7 +326,7 @@ def __init__(self, neox_args, layer_number):
 
  self.att = RWKV_TimeMix(neox_args, layer_number)
 
- self.ffn = RWKV_ChannelMix(neox_args, layer_number)
+ self.ffn = ParallelRWKV_ChannelMix(neox_args, layer_number, init_method=init_method)
 
  if neox_args.attention_dropout > 0:
  self.drop0 = nn.Dropout(p=neox_args.attention_dropout)

megatron/neox_arguments/arguments.py

@@ -1066,17 +1066,11 @@ def calculate_derived(self):
  if isinstance(self.zero_stage, int):
  assert self.zero_stage <= 2, "Zero stage 3 not compatible with Mamba"
  assert (
- self.hidden_dropout == 0.0,
+ self.hidden_dropout != 0.0,
  ), "Mamba does not yet have dropout implemented"
  if "rwkv" in self.attention_config:
- assert (
- not self.is_pipe_parallel and self.model_parallel_size == 1
- ), "RWKV not currently compatible with parallelism"
  if isinstance(self.zero_stage, int):
  assert self.zero_stage <= 2, "Zero stage 3 not compatible with RWKV"
- assert (
- self.hidden_dropout == 0.0,
- ), "RWKV does not yet have dropout implemented"
 
  # Sparsity config
  if self.sparsity_config is None: