Class BitNetTransformer

bitnet/bit_transformer.py:92–148

BitNetTransformer is a transformer-based model for BitNet. Args: dim (int): The dimension of the token embeddings. depth (int): The number of transformer layers. num_tokens (int): The number of tokens in the vocabulary. heads (int, optional): The number of a

# [MAIN MODEL] BitNetTransformer
class BitNetTransformer(nn.Module):
    """
    BitNetTransformer is a transformer-based model for BitNet.

    Args:
        dim (int): The dimension of the token embeddings.
        depth (int): The number of transformer layers.
        num_tokens (int): The number of tokens in the vocabulary.
        heads (int, optional): The number of attention heads in the transformer. Defaults to 8.
        ff_mult (int, optional): The multiplier for the feed-forward layer dimension. Defaults to 4.

    Examples:
    >>> import torch
    >>> from bitnet import BitNetTransformer
    >>> x = torch.randint(0, 20000, (1, 1024))
    >>> bitnet = BitNetTransformer(
    ...     num_tokens=20000,
    ...     dim=1024,
    ...     depth=6,
    ...     heads=8,
    ...     ff_mult=4,
    ... )
    >>> logits = bitnet(x)
    >>> print(logits)
    """

    def __init__(
        self,
        dim: int,
        depth: int,
        num_tokens: int,
        heads: int = 8,
        ff_mult: int = 4,
    ):
        super().__init__()
        self.emb = nn.Embedding(num_tokens, dim)

        self.transformer = Transformer(
            dim=dim, depth=depth, heads=heads, ff_mult=ff_mult
        )

        # self.to_logits = nn.Sequential(RMSNorm(dim), nn.Linear(dim, num_tokens))
        self.to_logits = OutputHead(
            dim,
            vocab_size=num_tokens,
        )
        
        # Norm
        self.norm = nn.LayerNorm(dim)

    def forward(self, x):
        x = self.emb(x)
        # Post emb norm
        x = self.norm(x)
        
        x = self.transformer(x)
        return self.to_logits(x)