| 21 | |
| 22 | |
| 23 | class FTT5EncoderWeight(object): |
| 24 | def __init__( |
| 25 | self, |
| 26 | config, |
| 27 | tensor_para_size, |
| 28 | pipeline_para_size, |
| 29 | *, |
| 30 | t5_with_bias=False, |
| 31 | use_gated_activation=False, |
| 32 | t5_with_moe=False, |
| 33 | position_embedding_type=0, |
| 34 | weight_data_type |
| 35 | ): |
| 36 | self.num_layer = config.num_layers |
| 37 | self.config = config |
| 38 | self.tensor_para_size = tensor_para_size |
| 39 | self.pipeline_para_size = pipeline_para_size |
| 40 | self.t5_with_bias = t5_with_bias |
| 41 | self.use_gated_activation = use_gated_activation |
| 42 | self.real_weights_num = 23 # assume all weights are allocated |
| 43 | self.t5_with_moe = t5_with_moe |
| 44 | self.position_embedding_type = position_embedding_type |
| 45 | self.weight_data_type = weight_data_type |
| 46 | self.adapter_inter_size = config.adapter_inter_size if hasattr(config, "adapter_inter_size") else 0 |
| 47 | self.w = [] |
| 48 | self.use_mpi = dist.is_mpi_available() |
| 49 | |
| 50 | if self.use_mpi: |
| 51 | try: |
| 52 | dist.init_process_group(backend='mpi') |
| 53 | except: |
| 54 | print("[INFO] WARNING: Exception occurred in dist.init_process_group(backend = 'mpi'). Maybe the process group has been initialized somewhere else.") |
| 55 | else: |
| 56 | print("[INFO] MPI is not available in this PyTorch build.") |
| 57 | assert tensor_para_size == 1, "[FATAL] MPI is required for tensor_para_size > 1." |
| 58 | assert pipeline_para_size == 1, "[FATAL] MPI is required for pipeline_para_size > 1." |
| 59 | |
| 60 | self.rank = dist.get_rank() if self.use_mpi else 0 |
| 61 | self.device_count = torch.cuda.device_count() |
| 62 | self.device = self.rank % self.device_count |
| 63 | torch.cuda.set_device(self.device) |
| 64 | |
| 65 | world_size = dist.get_world_size() if self.use_mpi else 1 |
| 66 | assert world_size == tensor_para_size * \ |
| 67 | pipeline_para_size, "[ERROR] world_size != tensor_para_size * pipeline_para_size" |
| 68 | self.tensor_para_rank = self.rank % self.tensor_para_size |
| 69 | self.pipeline_para_rank = self.rank // self.tensor_para_size |
| 70 | |
| 71 | def load_from_model(self, model): # assume this only applies to huggingface models |
| 72 | start_layer = self.pipeline_para_rank * self.num_layer // self.pipeline_para_size |
| 73 | end_layer = (self.pipeline_para_rank + 1) * self.num_layer // self.pipeline_para_size |
| 74 | |
| 75 | np_weight_dtype = self.weight_data_type |
| 76 | torch_weight_dtype = {np.float32: torch.float32, np.float16: torch.float16}[np_weight_dtype] |
| 77 | |
| 78 | encoder_weight_dict = {} |
| 79 | for name, param in model.named_parameters(): |
| 80 | if param.dim() == 2: |