xlm.backbones.llada
LLaDA-family backbone (modeling + LLaDAConfigBase).
Ported from the GSAI-ML/LLaDA-8B-Base Hub repo with state-dict keys kept
identical to the checkpoint (strict load, no remapping). The tokenizer is a
stock PreTrainedTokenizerFast; load it via
xlm.datamodule.load_auto_tokenizer (no custom tokenizer class needed).
ModelConfig
dataclass
LLaDA (model) configuration.
d_model = 768
class-attribute
instance-attribute
The hidden size of the model.
n_heads = 12
class-attribute
instance-attribute
The number of self-attention heads.
n_kv_heads = None
class-attribute
instance-attribute
The number of heads to use for keys and values. Defaults to n_heads.
Set this to None or n_heads for normal multi-head attention.
Set this to 1 for multi-query attention.
Set it to some in-between value for Llama2-style grouped query attention.
n_layers = 12
class-attribute
instance-attribute
The number of layers/blocks.
mlp_ratio = 4
class-attribute
instance-attribute
The ratio of the inner MLP dimensionality to d_model.
This is only used when mlp_hidden_size is not set.
mlp_hidden_size = None
class-attribute
instance-attribute
Set the exact hidden size for the MLP. Otherwise the inner MLP hidden size will be set to mlp_ratio * d_model.
activation_type = ActivationType.swiglu
class-attribute
instance-attribute
The activation function to use within the MLP layers.
block_type = BlockType.sequential
class-attribute
instance-attribute
The transformer block implementation.
block_group_size = 1
class-attribute
instance-attribute
The number of blocks to group together into a single parent block. This has no affect on the number of parameters in the model and is only used to wrap groups of blocks together with a single FSDP wrapper during training.
alibi = False
class-attribute
instance-attribute
If True, use ALiBi embeddings. Mutually exclusive with rope.
alibi_bias_max = 8.0
class-attribute
instance-attribute
Maximum absolute value of ALiBi bias.
rope = False
class-attribute
instance-attribute
Use rotary positional embeddings (RoPE). Mutually exclusive with alibi.
rope_full_precision = True
class-attribute
instance-attribute
If True, apply RoPE embeddings at full precision regardless of the input type. Otherwise,
apply RoPE at the precision of the input.
flash_attention = False
class-attribute
instance-attribute
If True, use FlashAttention.
attention_dropout = 0.1
class-attribute
instance-attribute
The dropout probability within the attention modules.
multi_query_attention = None
class-attribute
instance-attribute
Use the Multi-Query formulation of attention used in PaLM. This reduces the number of parameters and is more efficient during inference.
attention_layer_norm = False
class-attribute
instance-attribute
Apply layer norm to the keys and queries within the attention mechanism. This can help stabilize training.
residual_dropout = 0.1
class-attribute
instance-attribute
The dropout probability for the MLP and attention output within each block.
embedding_dropout = 0.1
class-attribute
instance-attribute
The dropout probability for embeddings.
input_emb_norm = False
class-attribute
instance-attribute
An input hidden_states norm implementation by gemmma.
layer_norm_type = LayerNormType.default
class-attribute
instance-attribute
The layernorm implementation to use.
layer_norm_with_affine = True
class-attribute
instance-attribute
Whether to include bias and weight parameters for the layer norms.
This only affects layer norms that are immediately followed by a linear layer in the forward pass,
so everything except QK-norms. To turn off affines for QK norms as well, set :attr:attention_layer_norm_with_affine
to False.
rms_norm_eps = 1e-05
class-attribute
instance-attribute
The rms layernorm eps param.
attention_layer_norm_with_affine = True
class-attribute
instance-attribute
Toggle affine transform for the QK norms.
max_sequence_length = 1024
class-attribute
instance-attribute
The maximum input sequence length supported by the model.
rope_theta = 10000.0
class-attribute
instance-attribute
The rope base param.
include_qkv_bias = False
class-attribute
instance-attribute
Whether or not to include bias parameters in qkv linear layers.
include_bias = False
class-attribute
instance-attribute
Whether or not to include bias parameters in linear layers. In PaLM, they got rid of all bias terms because they found that large models tend to have near 0 bias terms anyway.
bias_for_layer_norm = None
class-attribute
instance-attribute
Whether or not to include bias parameters in layer norm. This is separate from the include_bias parameter, because of a ROCm crash when biases are disabled in layer norm. When this is None (the default), it inherits the setting from include_bias.
scale_logits = False
class-attribute
instance-attribute
If True, scale the output logits by 1 / sqrt(d_model).
vocab_size = 50257
class-attribute
instance-attribute
Vocabulary size of the model.
embedding_size = 50304
class-attribute
instance-attribute
The number of embeddings, i.e. the number of tokens. If set to None it will default
to vocab_size. If vocab_size is not a multiple of 128, setting this to the
next multiple of 128 that's greater than vocab_size can improve throughput
substantially.
weight_tying = True
class-attribute
instance-attribute
Whether to tie output linear weights to the input embedding.
eos_token_id = 50256
class-attribute
instance-attribute
The ID of the end-of-sentence special token.
pad_token_id = 50256
class-attribute
instance-attribute
The ID of the token to use for padding. Defaults to the ID of the EOS token.
mask_token_id = 50256
class-attribute
instance-attribute
The ID of the token to use for mask token. Defaults to the ID of the EOS token.
init_device = None
class-attribute
instance-attribute
The torch device to use when initializing the model parameters, e.g. "cpu", "cuda:0", "meta".
init_fn = InitFnType.normal
class-attribute
instance-attribute
The weight initialization strategy.
init_std = 0.02
class-attribute
instance-attribute
The standard deviation to use when initializing weights with a "fixed distribution" init_fn, such
as "normal".
init_cutoff_factor = None
class-attribute
instance-attribute
A positive factor used to scale the cutoff values when initializing weights with a "fixed distribution" init_fn, such
as "normal". Setting this to None means values are not cutoff.
precision = None
class-attribute
instance-attribute
Precision used to train/evaluate with. You shouldn't set this directly.
See :data:TrainConfig.precision instead.
LLaDABlock
Bases: Module
A base class for transformer block implementations.
LLaDALlamaBlock
Bases: LLaDABlock
This is a transformer block where the output is computed as MLP(LN(x + Attention(LN(x))))
(plus another skip connection). This block is similar to LLaDASequentialBlock
but some operations have slightly different implementations to imitate the
behavior of Llama.
LLaDAModel
Bases: Module
forward(input_ids, input_embeddings=None, attention_mask=None, attention_bias=None, position_ids=None, past_key_values=None, use_cache=False, last_logits_only=False, output_hidden_states=None)
:param input_ids: A tensor of shape (batch_size, seq_len).
:param input_embeddings: A tensor of shape (batch_size, seq_len, d_model) with input
embeddings. When provided, it is treated as the output of the input embedding layer.
:param attention_mask: A tensor of shape (batch_size, seq_len) that indicates
which input IDs are masked. A 1 value in the mask means that
the corresponding input ID should not be ignored. A 0 means
that the corresponding input ID is masked.
This has the same meaning as the `attention_mask` in HuggingFace's `transformers`
library.
:param position_ids: An optional tensor of shape (batch_size, seq_len) with per-token
RoPE positions. When omitted, positions default to arange(seq_len). Requires
rope=True and no KV cache.
:param attention_bias: A tensor of shape (batch_size, 1, seq_len, seq_len),
(1, 1, seq_len, seq_len), or (seq_len, seq_len). This is used
to introduce causal or other biases.
If the tensor is a bool or byte tensor, a `True` or `1` at `attention_bias[:, :, i, j]`
indicates that the i-th element in the sequence is allowed to attend to the j-th
element in the sequence.
If the tensor is a float tensor, it will just be added to the attention
scores before the softmax.
The default is causal, which corresponds to a lower-diagonal byte matrix of ones.
:param past_key_values: Pre-computed keys and values for each attention block.
Can be used to speed up sequential decoding. The input_ids which have
their past given to this model should not be passed as input_ids as they have already been computed.
:param use_cache: If True, return key and value tensors for each block.
:param last_logits_only: If True, only compute the logits for the last token of each sequence.
This can speed up decoding when you only care about the next token.
LLaDAModelLM
Bases: PreTrainedModel
Extremely barebones HF model wrapper.
LLaDASequentialBlock
Bases: LLaDABlock
This is a typical transformer block where the output is computed as MLP(LN(x + Attention(LN(x))))
(plus another skip connection).