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