| 23 | from .checkpoint_quantization import checkpoint_quantization |
| 24 | |
| 25 | class EncoderWeights(object): |
| 26 | def __init__(self, layer_num, hidden_dim, weights=None, sparse=False, tensor_para_size=1, pipeline_para_size=1): |
| 27 | """weights need be a state_dict of bert model""" |
| 28 | self.layer_num = layer_num |
| 29 | self.int8 = False |
| 30 | self.hidden_dim = hidden_dim |
| 31 | self.weights = {} |
| 32 | self.tensor_para_size = tensor_para_size |
| 33 | self.pipeline_para_size = pipeline_para_size |
| 34 | |
| 35 | self.use_mpi = dist.is_mpi_available() |
| 36 | |
| 37 | if self.use_mpi: |
| 38 | try: |
| 39 | dist.init_process_group(backend='mpi') |
| 40 | except: |
| 41 | print("[INFO] WARNING: Exception occurred in dist.init_process_group(backend='mpi'). Maybe the process group has been initialized somewhere else.") |
| 42 | else: |
| 43 | print("[INFO] MPI is not available in this PyTorch build.") |
| 44 | assert tensor_para_size == 1, "[FATAL] MPI is required for tensor_para_size > 1." |
| 45 | assert pipeline_para_size == 1, "[FATAL] MPI is required for pipeline_para_size > 1." |
| 46 | |
| 47 | self.rank = dist.get_rank() if self.use_mpi else 0 |
| 48 | self.device_count = torch.cuda.device_count() |
| 49 | self.device = self.rank % self.device_count |
| 50 | torch.cuda.set_device(self.device) |
| 51 | |
| 52 | world_size = dist.get_world_size() if self.use_mpi else 1 |
| 53 | self.tensor_para_rank = self.rank % self.tensor_para_size |
| 54 | self.pipeline_para_rank = self.rank // self.tensor_para_size |
| 55 | start_layer = self.pipeline_para_rank * self.layer_num // self.pipeline_para_size |
| 56 | end_layer = (self.pipeline_para_rank + 1) * self.layer_num // self.pipeline_para_size |
| 57 | |
| 58 | if weights is None: |
| 59 | self._generated_weights = True |
| 60 | for i in range(layer_num): |
| 61 | pre = 'bert.encoder.layer.' + str(i) + '.' |
| 62 | self.weights[pre + 'attention.self.query.weight'] = torch.zeros(hidden_dim, hidden_dim) |
| 63 | self.weights[pre + 'attention.self.query.bias'] = torch.zeros(hidden_dim) |
| 64 | self.weights[pre + 'attention.self.key.weight'] = torch.zeros(hidden_dim, hidden_dim) |
| 65 | self.weights[pre + 'attention.self.key.bias'] = torch.zeros(hidden_dim) |
| 66 | self.weights[pre + 'attention.self.value.weight'] = torch.zeros(hidden_dim, hidden_dim) |
| 67 | self.weights[pre + 'attention.self.value.bias'] = torch.zeros(hidden_dim) |
| 68 | self.weights[pre + 'attention.output.dense.weight'] = torch.zeros(hidden_dim, hidden_dim) |
| 69 | self.weights[pre + 'attention.output.dense.bias'] = torch.zeros(hidden_dim) |
| 70 | self.weights[pre + 'attention.output.LayerNorm.weight'] = torch.zeros(hidden_dim) |
| 71 | self.weights[pre + 'attention.output.LayerNorm.bias'] = torch.zeros(hidden_dim) |
| 72 | self.weights[pre + 'intermediate.dense.weight'] = torch.zeros(4 * hidden_dim, hidden_dim) |
| 73 | self.weights[pre + 'intermediate.dense.bias'] = torch.zeros(4 * hidden_dim) |
| 74 | self.weights[pre + 'output.dense.weight'] = torch.zeros(hidden_dim, 4 * hidden_dim) |
| 75 | self.weights[pre + 'output.dense.bias'] = torch.zeros(hidden_dim) |
| 76 | self.weights[pre + 'output.LayerNorm.weight'] = torch.zeros(hidden_dim) |
| 77 | self.weights[pre + 'output.LayerNorm.bias'] = torch.zeros(hidden_dim) |
| 78 | for k, v in self.weights.items(): |
| 79 | if not k.endswith('_amax'): |
| 80 | self.weights[k] = torch.nn.init.uniform_(v, -1, 1) |
| 81 | if sparse: |
| 82 | for k, v in self.weights.items(): |
no outgoing calls
no test coverage detected