Class BertForPretrain

examples/tutorial/sequence_parallel/model/bert.py:18–126

class BertForPretrain(nn.Module):
    def __init__(
        self,
        vocab_size,
        hidden_size,
        max_sequence_length,
        num_attention_heads,
        num_layers,
        add_binary_head,
        is_naive_fp16,
        num_tokentypes=2,
        dropout_prob=0.1,
        mlp_ratio=4,
        init_std=0.02,
        convert_fp16_to_fp32_in_softmax=False,
    ):
        super().__init__()
        self.seq_parallel_size = gpc.get_world_size(ParallelMode.SEQUENCE)
        assert (
            max_sequence_length % self.seq_parallel_size == 0
        ), "sequence length is not divisible by the sequence parallel size"
        self.sub_seq_length = max_sequence_length // self.seq_parallel_size
        self.init_std = init_std
        self.num_layers = num_layers

        if not add_binary_head:
            num_tokentypes = 0

        self.preprocessor = PreProcessor(self.sub_seq_length)
        self.embedding = Embedding(
            hidden_size=hidden_size,
            vocab_size=vocab_size,
            max_sequence_length=max_sequence_length,
            embedding_dropout_prob=dropout_prob,
            num_tokentypes=num_tokentypes,
        )
        self.bert_layers = nn.ModuleList()

        for i in range(num_layers):
            bert_layer = BertLayer(
                layer_number=i + 1,
                hidden_size=hidden_size,
                num_attention_heads=num_attention_heads,
                attention_dropout=dropout_prob,
                mlp_ratio=mlp_ratio,
                hidden_dropout=dropout_prob,
                convert_fp16_to_fp32_in_softmax=convert_fp16_to_fp32_in_softmax,
                is_naive_fp16=is_naive_fp16,
            )
            self.bert_layers.append(bert_layer)

        self.layer_norm = LayerNorm(hidden_size)
        self.head = BertDualHead(
            hidden_size, self.embedding.word_embedding_weight.size(0), add_binary_head=add_binary_head
        )
        self.reset_parameters()

    def _init_normal(self, tensor):