Class Config

axlearn/common/decoder.py:481–518

Configures Decoder.

        dim: Required[int] = REQUIRED
        # Dropout rate applied throughout model, except for child Dropout configs with rate set
        # explicitly.
        dropout_rate: float = 0.0
        # Vector from input ids table.
        emb: BaseEmbedding.Config = TransformerTextEmbeddings.default_config()
        # Transformer model trunk.
        transformer: BaseStackedTransformerLayer.Config = StackedTransformerLayer.default_config()
        # Layer norm applied to transformer output.
        output_norm: Optional[InstantiableConfig] = LayerNorm.default_config()
        # Optional dropout rate for the transformer output.
        # If output_dropout.rate is None, it will default to cfg.dropout_rate
        output_dropout: Dropout.Config = Dropout.default_config()
        # Optional LmHead layer maps the hidden state to vocab logits (if None use emb.token_emb).
        lm_head: Optional[InstantiableConfig] = None
        pad_token_id: int = 0  # Int ID of the inputs to be masked for self-attention.
        eos_token_id: int = 1  # Int ID of the end of sequence token id.
        # Specifies how to partition the output logits of shape [batch, max_seq_len, vocab_size].
        logits_partition_spec: tuple[Union[Optional[str], tuple[Optional[str]]], ...] = (
            "data",
            None,
            "model",
        )
        # Partition spec for the pre-matmul `logits_x` activation (not the final logits
        # output, which always uses `logits_partition_spec`). Defaults to None, which reuses
        # `logits_partition_spec`; set explicitly -- it takes precedence -- when hidden-dim
        # sharding should differ, e.g. to match the embedding/lm_head weight's sharding.
        hidden_state_partition_spec: Optional[
            tuple[Union[Optional[str], tuple[Optional[str]]], ...]
        ] = None
        # Precision dtype for logits computation.
        # if None, uses the input dtype as default for logits computation.
        logits_forward_dtype: Optional[jnp.dtype] = None
        # The logit modifier to apply. If None, does not modify logits.
        output_logits_modifier: Optional[ConfigOr[logit_modifiers.LogitsToLogitsFn]] = None
        # The decoding implementation.
        decoding: DecodingLayer.Config = DecodingLayer.default_config()

    def __init__(self, cfg: Config, *, parent: Module):
        super().__init__(cfg, parent=parent)
        cfg = self.config

        # Validate logits_forward_dtype