optimum/graphcore/models/bert/modeling_bert.py [182:268]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - if isinstance(self.cls.predictions.decoder, SerializedLinear): self.cls.predictions.decoder = self.cls.predictions.decoder.to_model() self.tie_weights() return self def _init_weights(self, module): """Initialize the weights""" def truncated_normal_(tensor, mean=0, std=1): """ Truncated Normal distribution, truncated at 2 sigma """ r = torch.tensor(truncnorm.rvs(-2, 2, loc=mean, scale=std, size=tensor.shape)) tensor.data.copy_(r) if isinstance(module, nn.Linear): truncated_normal_(module.weight, mean=0.0, std=self.config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): truncated_normal_(module.weight, mean=0.0, std=self.config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, next_sentence_label: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple[torch.Tensor], BertForPreTrainingOutput]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.bert( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output, pooled_output = outputs[:2] if labels is not None: if hasattr(self.config, "max_num_masked_tokens"): # Select only the masked tokens for the classifier labels, positions = torch.topk(labels, k=self.config.max_num_masked_tokens, dim=1) sequence_output = self.gather_indices(sequence_output, positions) prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output) total_loss = None if labels is not None and next_sentence_label is not None: masked_lm_loss = F.cross_entropy( prediction_scores.view(-1, self.config.vocab_size), labels.view(-1), ).float() next_sentence_loss = F.cross_entropy( seq_relationship_score.view(-1, 2), next_sentence_label.view(-1) ).float() total_loss = poptorch.identity_loss(masked_lm_loss + next_sentence_loss, reduction="none") # If labels are provided (training mode) only output the loss if not return_dict: output = (prediction_scores, seq_relationship_score) + outputs[2:] return (total_loss,) if total_loss is not None else output return BertForPreTrainingOutput( loss=total_loss, prediction_logits=prediction_scores if total_loss is None else None, seq_relationship_logits=seq_relationship_score if total_loss is None else None, hidden_states=outputs.hidden_states if total_loss is None else None, attentions=outputs.attentions if total_loss is None else None, ) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - optimum/graphcore/models/groupbert/modeling_groupbert.py [466:552]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - if isinstance(self.cls.predictions.decoder, SerializedLinear): self.cls.predictions.decoder = self.cls.predictions.decoder.to_model() self.tie_weights() return self def _init_weights(self, module): """Initialize the weights""" def truncated_normal_(tensor, mean=0, std=1): """ Truncated Normal distribution, truncated at 2 sigma """ r = torch.tensor(truncnorm.rvs(-2, 2, loc=mean, scale=std, size=tensor.shape)) tensor.data.copy_(r) if isinstance(module, nn.Linear): truncated_normal_(module.weight, mean=0.0, std=self.config.initializer_range) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): truncated_normal_(module.weight, mean=0.0, std=self.config.initializer_range) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, nn.LayerNorm): module.bias.data.zero_() module.weight.data.fill_(1.0) def forward( self, input_ids: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, token_type_ids: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, head_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, next_sentence_label: Optional[torch.Tensor] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple[torch.Tensor], BertForPreTrainingOutput]: return_dict = return_dict if return_dict is not None else self.config.use_return_dict outputs = self.bert( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, head_mask=head_mask, inputs_embeds=inputs_embeds, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, ) sequence_output, pooled_output = outputs[:2] if labels is not None: if hasattr(self.config, "max_num_masked_tokens"): # Select only the masked tokens for the classifier labels, positions = torch.topk(labels, k=self.config.max_num_masked_tokens, dim=1) sequence_output = self.gather_indices(sequence_output, positions) prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output) total_loss = None if labels is not None and next_sentence_label is not None: masked_lm_loss = F.cross_entropy( prediction_scores.view(-1, self.config.vocab_size), labels.view(-1), ).float() next_sentence_loss = F.cross_entropy( seq_relationship_score.view(-1, 2), next_sentence_label.view(-1) ).float() total_loss = poptorch.identity_loss(masked_lm_loss + next_sentence_loss, reduction="none") # If labels are provided (training mode) only output the loss if not return_dict: output = (prediction_scores, seq_relationship_score) + outputs[2:] return (total_loss,) if total_loss is not None else output return BertForPreTrainingOutput( loss=total_loss, prediction_logits=prediction_scores if total_loss is None else None, seq_relationship_logits=seq_relationship_score if total_loss is None else None, hidden_states=outputs.hidden_states if total_loss is None else None, attentions=outputs.attentions if total_loss is None else None, ) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -