Convert a simple custom transformer to GGUF format

[iyubondyrev]

Hello!

I would be very grateful if someone could tell me if what I want is possible, and if so, share some tips on how to achieve it.

Suppose I have a simple transformer model, similar to GPT, with several decoder layers. The model is structured like this:

class PositionalEncoding(nn.Module):
    def __init__(self,
                 emb_size: int,
                 dropout: float,
                 maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        den = torch.exp(- torch.arange(0, emb_size, 2)* torch.log(10000) / emb_size)
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        pos_embedding = torch.zeros((maxlen, emb_size))
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        pos_embedding = pos_embedding.unsqueeze(-2)

        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size

    def forward(self, tokens: Tensor):
        return self.embedding(tokens.long()) * torch.sqrt(self.emb_size)

class SimpleTransformer(nn.Module):
    def __init__(self,
                 num_layers: int,
                 emb_size: int,
                 nhead: int,
                 vocab_size: int,
                 dim_feedforward: int = 512,
                 dropout: float = 0.1):
        super(SimpleTransformer, self).__init__()
        encoder_layer = TransformerEncoderLayer(d_model=emb_size,
                                                nhead=nhead,
                                                dim_feedforward=dim_feedforward,
                                                dropout=dropout,
                                                batch_first=True
                                                )
        self.transformer: TransformerEncoder = TransformerEncoder(encoder_layer=encoder_layer,
                                              num_layers=num_layers)
        self.generator = nn.Linear(emb_size, vocab_size)
        self.tok_emb = TokenEmbedding(vocab_size, emb_size)
        self.positional_encoding = PositionalEncoding(
            emb_size, dropout=dropout)
    

    def forward(self,
                src: Tensor,
                src_mask: Tensor = None,
                src_padding_mask: Tensor = None,
                is_causal: bool = None):
        src_emb = self.positional_encoding(self.tok_emb(src))
        outs = self.transformer(src=src_emb, mask=src_mask, src_key_padding_mask=src_padding_mask, is_causal=is_causal)
        return self.generator(outs)

Is there a way to convert this model to GGUF and use it for inference? I understand that for the model to be converted to GGUF and used for inference, its architecture needs to be supported by the convert-hf-to-gguf.py script. Given that the model is quite simple, maybe it is not so hard to do.

As for the tokenizer: it is a fixed mapping from words to numbers (about 240 words in total, like a simple dict)

Thank you!

[iyubondyrev]

This PR does it for Phi-models. Is it a good starting point?