embeddings.py

# Copyright 2024 **AUTHORS_TODO**
# License: Apache-2.0

# Copyright 2022 MosaicML Examples authors
# SPDX-License-Identifier: Apache-2.0

# Copyright 2023 MosaicML Examples authors
# SPDX-License-Identifier: Apache-2.0

# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
# Copyright (c) 2018-2021, NVIDIA CORPORATION.  All rights reserved.
# Copyright (c) 2023, Tri Dao.


import torch
import torch.nn as nn
from typing import Optional

from configuration_bert import FlexBertConfig
from normalization import get_norm_layer
from initialization import ModuleType, init_weights


class BertAlibiEmbeddings(nn.Module):
    """Construct the embeddings for words, ignoring position.

    There are no positional embeddings since we use ALiBi and token_type
    embeddings.

    This module is modeled after the Hugging Face BERT's
    :class:`~transformers.model.bert.modeling_bert.BertEmbeddings`, but is
    modified as part of Mosaic BERT's ALiBi implementation. The key change is
    that position embeddings are removed. Position information instead comes
    from attention biases that scale linearly with the position distance
    between query and key tokens.

    This module ignores the `position_ids` input to the `forward` method.
    """

    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        # ALiBi doesn't use position embeddings
        if getattr(config, "token_type_embeddings", True):
            self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
            self.use_token_type_embeddings = True
        else:
            self.use_token_type_embeddings = False

        self.LayerNorm = get_norm_layer(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        if self.use_token_type_embeddings:
            self.register_buffer(
                "token_type_ids", torch.zeros(config.max_position_embeddings, dtype=torch.long), persistent=False
            )

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        past_key_values_length: int = 0,
    ) -> torch.Tensor:
        if (input_ids is not None) == (inputs_embeds is not None):
            raise ValueError("Must specify either input_ids or input_embeds!")
        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            assert inputs_embeds is not None  # just for type checking
            input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]

        if position_ids is None:
            # great! ALiBi
            pass

        # Setting the token_type_ids to the registered buffer in constructor
        # where it is all zeros, which usually occurs when it's auto-generated;
        # registered buffer helps users when tracing the model without passing
        # token_type_ids, solves issue #5664
        if self.use_token_type_embeddings and token_type_ids is None:
            if hasattr(self, "token_type_ids"):
                buffered_token_type_ids = self.token_type_ids[:, :seq_length]
                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
                token_type_ids = buffered_token_type_ids_expanded
            else:
                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=input_ids.device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)

        if self.use_token_type_embeddings:
            token_type_embeddings = self.token_type_embeddings(token_type_ids)
            embeddings = inputs_embeds + token_type_embeddings
        else:
            embeddings = inputs_embeds

        # no position embeddings! ALiBi
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings


class FlexBertEmbeddingsBase(nn.Module):
    """A FlexBERT embeddings base class for type hints."""

    def __init__(self, config: FlexBertConfig):
        super().__init__()
        self.config = config

    def _init_weights(self, reset_params: bool = False):
        raise NotImplementedError("This is a base class and should not be used directly.")

    def reset_parameters(self):
        self._init_weights(reset_params=True)

    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
        raise NotImplementedError("This is a base class and should not be used directly.")


class FlexBertAbsoluteEmbeddings(FlexBertEmbeddingsBase):
    """Construct the embeddings with absolute positional embeddings."""

    def __init__(self, config: FlexBertConfig):
        super().__init__(config)
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)

        self.norm = get_norm_layer(config) if config.embed_norm else nn.Identity()
        self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()

        self.register_buffer(
            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
        )

    def _init_weights(self, reset_params: bool = False):
        init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)
        init_weights(self.config, self.position_embeddings, type_of_module=ModuleType.emb)

        if reset_params:
            if self.config.embed_norm:
                self.norm.reset_parameters()  # type: ignore

    def forward(
        self,
        input_ids: torch.LongTensor,
        position_ids: Optional[torch.LongTensor] = None,
    ) -> torch.Tensor:
        if position_ids is None:
            position_ids = self.position_ids[:, 0 : input_ids.shape[1]]

        embeddings = self.tok_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)

        embeddings = self.norm(embeddings + position_embeddings)
        return self.drop(embeddings)


class FlexBertCompiledSansPositionEmbeddings(FlexBertEmbeddingsBase):
    """Construct the embeddings from token embeddings without any positional embeddings."""

    def __init__(self, config: FlexBertConfig):
        super().__init__(config)
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)

        self.norm = get_norm_layer(config, compiled_norm=config.compile_model) if config.embed_norm else nn.Identity()
        self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()

    def _init_weights(self, reset_params: bool = False):
        init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)

        if reset_params:
            if self.config.embed_norm:
                self.norm.reset_parameters()  # type: ignore

    @torch.compile(dynamic=True)
    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
        return self.drop(self.norm(self.tok_embeddings(input_ids)))


class FlexBertSansPositionEmbeddings(FlexBertEmbeddingsBase):
    """Construct the embeddings from token embeddings without any positional embeddings."""

    def __init__(self, config: FlexBertConfig):
        super().__init__(config)
        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)

        self.norm = get_norm_layer(config) if config.embed_norm else nn.Identity()
        self.drop = nn.Dropout(config.embed_dropout_prob) if config.embed_dropout_prob > 0.0 else nn.Identity()

    def _init_weights(self, reset_params: bool = False):
        init_weights(self.config, self.tok_embeddings, type_of_module=ModuleType.emb)

        if reset_params:
            if self.config.embed_norm:
                self.norm.reset_parameters()  # type: ignore

    def forward(self, input_ids: torch.LongTensor, position_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
        return self.drop(self.norm(self.tok_embeddings(input_ids)))


EBB2CLS = {
    "absolute_pos": FlexBertAbsoluteEmbeddings,
    "sans_pos": FlexBertSansPositionEmbeddings,
}


def get_embedding_layer(config: FlexBertConfig) -> FlexBertEmbeddingsBase:
    try:
        if config.compile_model and config.embedding_layer == "sans_pos":
            return FlexBertCompiledSansPositionEmbeddings(config)
        elif config.compile_model:
            raise ValueError(f"{config.compile_model=} only supports sans_pos embeddings.")
        return EBB2CLS[config.embedding_layer](config)
    except KeyError:
        raise ValueError(f"Invalid embeddings layer type: {config.embedding_layer=}, must be one of {EBB2CLS.keys()}.")