(self, num_tokentypes=0, parallel_output=True)
| 128 | """Pipeline version of CodeGeeX.""" |
| 129 | |
| 130 | def __init__(self, num_tokentypes=0, parallel_output=True): |
| 131 | args = get_args() |
| 132 | self.parallel_output = parallel_output |
| 133 | |
| 134 | init_method = init_method_normal(args.init_method_std) |
| 135 | |
| 136 | self.specs = [] |
| 137 | |
| 138 | # Embedding layer |
| 139 | self.specs.append( |
| 140 | TiedLayerSpec( |
| 141 | "embed", |
| 142 | EmbeddingPipe, |
| 143 | args.hidden_size, |
| 144 | args.padded_vocab_size, |
| 145 | args.max_position_embeddings, |
| 146 | args.hidden_dropout, |
| 147 | init_method=init_method, |
| 148 | num_tokentypes=num_tokentypes, |
| 149 | tied_weight_attr="word_embeddings_weight", |
| 150 | ) |
| 151 | ) |
| 152 | |
| 153 | self.specs.append(lambda x: x.transpose(0, 1).contiguous()) |
| 154 | |
| 155 | for layer_idx in range(args.num_layers): |
| 156 | self.specs.append( |
| 157 | LayerSpec( |
| 158 | ParallelTransformerLayerPipe, |
| 159 | init_method=init_method, |
| 160 | output_layer_init_method=scaled_init_method_normal( |
| 161 | args.init_method_std, args.num_layers |
| 162 | ), |
| 163 | layer_number=layer_idx, |
| 164 | self_attn_mask_type=AttnMaskType.causal, |
| 165 | ) |
| 166 | ) |
| 167 | |
| 168 | # Undo data format change |
| 169 | self.specs.append(lambda x: x.transpose(0, 1).contiguous()) |
| 170 | |
| 171 | # Final layernorm after transformer layers |
| 172 | self.specs.append( |
| 173 | LayerSpec(LayerNorm, args.hidden_size, eps=args.layernorm_epsilon) |
| 174 | ) |
| 175 | |
| 176 | def _logits_helper(embedding, lm_output): |
| 177 | """A wrapper to massage inputs/outputs from pipeline.""" |
| 178 | return parallel_lm_logits( |
| 179 | lm_output, embedding.word_embeddings_weight, self.parallel_output |
| 180 | ) |
| 181 | |
| 182 | self.specs.append( |
| 183 | TiedLayerSpec( |
| 184 | "embed", |
| 185 | EmbeddingPipe, |
| 186 | args.hidden_size, |
| 187 | args.padded_vocab_size, |
nothing calls this directly
no test coverage detected