| 31 | |
| 32 | class LinearRNN(MegatronModule): |
| 33 | def __init__( |
| 34 | self, |
| 35 | config, |
| 36 | submodules: LinearRNNSubmodules, |
| 37 | layer_number=None, |
| 38 | ): |
| 39 | super().__init__(config) |
| 40 | self.config = config |
| 41 | self.la_module = config.la_module |
| 42 | self.num_heads = config.num_attention_heads |
| 43 | |
| 44 | self.head_dim = self.config.kv_channels |
| 45 | self.key_dim = int(self.config.hidden_size * self.config.expand_k) |
| 46 | self.value_dim = int(self.config.hidden_size * self.config.expand_v) |
| 47 | self.hidden_size = self.config.hidden_size |
| 48 | self.rwkv6_la_proj_low_rank_dim = self.config.rwkv6_la_proj_low_rank_dim |
| 49 | self.head_qk_dim = self.key_dim // self.num_heads |
| 50 | |
| 51 | self.la_gate_fn = self.config.la_gate_fn |
| 52 | self.time_shift = nn.ZeroPad2d((0, 0, 1, -1)) |
| 53 | |
| 54 | if config.hidden_size is not None and config.num_attention_heads is not None: |
| 55 | expand_ratio = self.hidden_size // self.num_heads |
| 56 | self.expand_ratio = expand_ratio |
| 57 | self.forget_dim = int(self.num_heads * self.expand_ratio) |
| 58 | self.input_dim = self.hidden_size |
| 59 | |
| 60 | if self.la_module == 'rwkv6': |
| 61 | self.x_proj = nn.Sequential( |
| 62 | LerpLinear(self.hidden_size, self.rwkv6_la_proj_low_rank_dim * 5), |
| 63 | nn.Tanh(), |
| 64 | nn.Linear(self.rwkv6_la_proj_low_rank_dim * 5, self.hidden_size, bias=False) |
| 65 | ) |
| 66 | |
| 67 | self.x_bias = nn.Parameter(torch.zeros(5, self.hidden_size)) |
| 68 | |
| 69 | self.r_proj = build_module( |
| 70 | submodules.r_proj, |
| 71 | self.config, |
| 72 | self.hidden_size, |
| 73 | self.key_dim, |
| 74 | ) |
| 75 | |
| 76 | self.w_proj = build_module( |
| 77 | submodules.w_proj, |
| 78 | self.config, |
| 79 | self.hidden_size, |
| 80 | self.key_dim, |
| 81 | low_rank_dim=config.rwkv6_la_gate_low_rank_dim, |
| 82 | ) |
| 83 | |
| 84 | self.k_proj = build_module( |
| 85 | submodules.k_proj, |
| 86 | self.config, |
| 87 | self.hidden_size, |
| 88 | self.key_dim, |
| 89 | ) |
| 90 | |