| 358 | |
| 359 | |
| 360 | class FTT5Encoder(nn.Module): |
| 361 | def __init__(self, encoder_weight_list, lib_path, head_num, head_size, inter_size, d_model, is_remove_padding, |
| 362 | num_layer, num_bucket=32, num_expert=0, moe_layer_index=[], max_distance=128, sparse=False, |
| 363 | q_scaling=1.0, tensor_para_size=1, pipeline_para_size=1, t5_with_bias=False, |
| 364 | position_embedding_type=0, moe_k=0, activation_type="relu", adapter_inter_size=0, |
| 365 | adapter_norm_position="pre"): |
| 366 | super().__init__() |
| 367 | |
| 368 | self.use_mpi = dist.is_mpi_available() |
| 369 | |
| 370 | if self.use_mpi: |
| 371 | try: |
| 372 | dist.init_process_group(backend='mpi') |
| 373 | except: |
| 374 | print("[INFO] WARNING: Exception occurred in dist.init_process_group(backend = 'mpi'). Maybe the process group has been initialized somewhere else.") |
| 375 | else: |
| 376 | print("[INFO] MPI is not available in this PyTorch build.") |
| 377 | assert tensor_para_size == 1, "[FATAL] MPI is required for tensor_para_size > 1." |
| 378 | assert pipeline_para_size == 1, "[FATAL] MPI is required for pipeline_para_size > 1." |
| 379 | |
| 380 | torch.classes.load_library(lib_path) |
| 381 | try: |
| 382 | self.encoder = torch.classes.FasterTransformer.T5Encoder(*encoder_weight_list, moe_layer_index, head_num, |
| 383 | head_size, inter_size, d_model, |
| 384 | is_remove_padding, num_layer, num_bucket, |
| 385 | num_expert, max_distance, sparse, q_scaling, |
| 386 | tensor_para_size, pipeline_para_size, |
| 387 | t5_with_bias, position_embedding_type, moe_k, |
| 388 | activation_type, adapter_inter_size, |
| 389 | adapter_norm_position) |
| 390 | except: |
| 391 | self.encoder = torch.classes.FasterTransformerT5Encoder(*encoder_weight_list, moe_layer_index, head_num, |
| 392 | head_size, inter_size, d_model, |
| 393 | is_remove_padding, num_layer, num_bucket, |
| 394 | num_expert, max_distance, sparse, q_scaling, |
| 395 | tensor_para_size, pipeline_para_size, |
| 396 | t5_with_bias, position_embedding_type, moe_k, |
| 397 | activation_type, adapter_inter_size, |
| 398 | adapter_norm_position) |
| 399 | |
| 400 | def forward(self, input, seq_len, inputs_embeds=None): |
| 401 | output = self.encoder.forward(input, seq_len, inputs_embeds) |
| 402 | return output |