| 15 | class MOE(nn.Module): |
| 16 | |
| 17 | def __init__(self, num_experts, topk, input_dim, ffn_dim, output_dim, |
| 18 | num_heads, max_seq_len, gate_type, gate_noise): |
| 19 | super().__init__() |
| 20 | self.proj = nn.Linear(input_dim, output_dim) |
| 21 | self.activation = nn.GELU() |
| 22 | try: |
| 23 | data_group = net.create_groups_from_world(group_count=1).data_group |
| 24 | except: |
| 25 | data_group = None |
| 26 | self.model = tutel_moe.moe_layer(gate_type={ |
| 27 | 'type': gate_type, |
| 28 | 'k': topk, |
| 29 | 'fp32_gate': True, |
| 30 | 'gate_noise': gate_noise, |
| 31 | 'capacity_factor': 1.5 |
| 32 | }, |
| 33 | experts={ |
| 34 | 'type': 'ffn', |
| 35 | 'count_per_node': num_experts, |
| 36 | 'hidden_size_per_expert': ffn_dim, |
| 37 | 'activation_fn': |
| 38 | lambda x: F.gelu(x) |
| 39 | }, |
| 40 | model_dim=input_dim, |
| 41 | batch_prioritized_routing=True, |
| 42 | is_gshard_loss=False, |
| 43 | group=data_group) |
| 44 | self.embedding = nn.Parameter( |
| 45 | torch.randn(1, max_seq_len, num_heads, input_dim)) |
| 46 | |
| 47 | def forward(self, x): |
| 48 | B, T, H, D = x.shape |