Repository: hpcaitech/ColossalAI
File: colossalai/pipeline/policy/bert.py (Python)

Code before:
        last_hidden_state=hidden_states,
        past_key_values=next_decoder_cache,
        hidden_states=all_hidden_states,
        attentions=all_self_attentions,
        cross_attentions=all_cross_attentions,
    )


# The layer partition policy for bertmodel
class BertModelPolicy(Policy):

    def __init__(self, stage_manager: PipelineStageManager, num_layers: int, num_stages: int):
        super().__init__(stage_manager=stage_manager)
        self.stage_manager = stage_manager
        self.layers_per_stage = super().distribute_layers(num_layers, num_stages)

    def get_hold_layers(self, module: BertModel) -> List[Module]:
        """
        get pipeline layers for current stage
        """
        hold_layers = []
        if self.stage_manager.is_first_stage():
            hold_layers.append(module.embeddings)
        num_layers_per_stage_accumulated = np.cumsum(self.layers_per_stage)
        hold_layers.extend(
            module.encoder.layer[num_layers_per_stage_accumulated[self.stage_manager.stage - 1] if self.stage_manager.
                                 stage > 0 else 0:num_layers_per_stage_accumulated[self.stage_manager.stage]])

        if self.stage_manager.is_last_stage():
            hold_layers.append(module.pooler)

        return hold_layers

    def get_shared_params(self, module: BertModel) -> List[Dict[int, Tensor]]:
        '''no shared params in bertmodel'''
        pass

    def replace_forward(self, module: Module) -> None:
        module.model.forward = MethodType(partial(bert_model_forward, stage_manager=self.stage_manager), module.model)

def bert_for_pretraining_forward(
    self: BertForPreTraining,
    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,

Review comment (suggestion):
This is ugly. You can first `np.insert(num_layers_per_stage_accumulated, 0, 0)`, then
```python
start_idx = num_layers_per_stage_accumulated[self.stage_manager.stage - 1]
end_idx = num_layers_per_stage_accumulated[self.stage_manager.stage]
hold_layers.extend(module.encoder.layer[start_idx:end_idx])
```

Code after:
    )


# The layer partition policy for bertmodel
class BertModelPolicy(Policy):

    def __init__(self, stage_manager: PipelineStageManager, num_layers: int, num_stages: int):
        super().__init__(stage_manager=stage_manager)
        self.stage_manager = stage_manager
        self.layers_per_stage = self.distribute_layers(num_layers, num_stages)

    def get_hold_layers(self, module: BertModel) -> List[Module]:
        """
        get pipeline layers for current stage
        """
        hold_layers = []
        if self.stage_manager.is_first_stage():
            hold_layers.append(module.embeddings)
        start_idx, end_idx = self.get_stage_index(self.layers_per_stage, self.stage_manager.stage)
        hold_layers.extend(module.encoder.layer[start_idx:end_idx])
        if self.stage_manager.is_last_stage():
            hold_layers.append(module.pooler)

        return hold_layers

    def get_shared_params(self, module: BertModel) -> List[Dict[int, Tensor]]:
        '''no shared params in bertmodel'''
        pass

    def replace_forward(self, module: Module) -> None:
        module.model.forward = MethodType(partial(bert_model_forward, stage_manager=self.stage_manager), module.model)


def bert_for_pretraining_forward(
    self: BertForPreTraining,
    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,
    hidden_states: Optional[torch.LongTensor] = None,
    stage_manager: Optional[PipelineStageManager] = None,
) -> Union[Tuple[torch.Tensor], BertForPreTrainingOutput]:

    return_dict = return_dict if return_dict is not None else self.config.use_return_dict