configuration_rost.py

"""HuggingFace config for RoST.

Ships inside the published model repository and runs on the downloader's
machine, so it must not import anything from `nanochat`.

Field names mirror `nanochat.gpt.GPTConfig` exactly rather than being renamed
to Llama's vocabulary. A rename would need a mapping table that nothing checks,
and a silently wrong mapping produces a model that loads and computes the wrong
thing -- the one failure mode this whole export has to avoid.
"""
from transformers.configuration_utils import PretrainedConfig


class RostConfig(PretrainedConfig):
    model_type = "rost"
    keys_to_ignore_at_inference = ["past_key_values"]

    def __init__(
        self,
        vocab_size=32768,
        n_layer=24,
        n_head=12,
        n_kv_head=12,
        n_embd=1536,
        sequence_len=4096,
        rope_base=100000,
        window_pattern="SSSL",
        pad_vocab_size_to=64,
        logit_softcap=15.0,
        attention_scale=1.2,
        ve_gate_channels=12,
        smear_gate_channels=24,
        bos_token_id=None,
        eos_token_id=None,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_kv_head = n_kv_head
        self.n_embd = n_embd
        self.sequence_len = sequence_len
        self.rope_base = rope_base
        self.window_pattern = window_pattern
        self.pad_vocab_size_to = pad_vocab_size_to
        # Constants in the training code, carried as config so a checkpoint
        # trained under different ones cannot be served under these.
        self.logit_softcap = logit_softcap
        self.attention_scale = attention_scale
        self.ve_gate_channels = ve_gate_channels
        self.smear_gate_channels = smear_gate_channels
        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)

    @property
    def padded_vocab_size(self):
        pad = self.pad_vocab_size_to
        return ((self.vocab_size + pad - 1) // pad) * pad

    @property
    def head_dim(self):
        return self.n_embd // self.n_head

    # Aliases so generic HuggingFace code (generation, device maps, pipelines)
    # finds what it expects without the weights being renamed.
    @property
    def hidden_size(self):
        return self.n_embd

    @property
    def num_attention_heads(self):
        return self.n_head

    @property
    def num_key_value_heads(self):
        return self.n_kv_head

    @property
    def num_hidden_layers(self):
        return self.n_layer

    @property
    def max_position_embeddings(self):
        return self.sequence_len