xlm.backbones.llada.configuration_llada
LLaDA configuration.
Ported from the GSAI-ML/LLaDA-8B-Base Hub repo (configuration_llada.py,
trust_remote_code module). Changes vs upstream:
LLaDAConfigrenamed toLLaDAConfigBase(variant configs subclass it inxlm-models/llada), mirroringDreamConfigBase.- No
AutoConfig.registercall: registration would clash with the remote-code module when both are imported in one process (e.g. checkpoint-parity tests).
LayerNormType
Bases: StrEnum
default = 'default'
class-attribute
instance-attribute
The default LayerNorm implementation, equivalent to PyTorch's built-in version.
low_precision = 'low_precision'
class-attribute
instance-attribute
A low-precision version of the default LayerNorm.
rms = 'rms'
class-attribute
instance-attribute
An RMSNorm implementation. When using torch.compile this is
probably the fastest implementation.
gemma_rms = 'gemma_rms'
class-attribute
instance-attribute
An RMSNorm implementation by gemmma. When using torch.compile this is
probably the fastest implementation.
amd_compatible = 'amd_compatible'
class-attribute
instance-attribute
LayerNorm implemented manually to work around an issue with ROCm.
BlockType
Bases: StrEnum
llama = 'llama'
class-attribute
instance-attribute
A block similar to the sequential block with slightly different implementations of operations like attention to imitate the behavior of Llama.
InitFnType
Bases: StrEnum
mitchell = 'mitchell'
class-attribute
instance-attribute
The strategy suggested to us by Mitchell Wortsman from UW. This uses a truncated normal distribution with an adaptive standard deviation that depends on the size of the weights as well as the depth of the layer.
normal = 'normal'
class-attribute
instance-attribute
All weights are initialized from the same normal distribution.
kaiming_normal = 'kaiming_normal'
class-attribute
instance-attribute
All weights are initialized with the Kaiming method from a normal distribution. Note this currently won't work with FSDP.
fan_in = 'fan_in'
class-attribute
instance-attribute
"Fan-in variance scaling", i.e. normal with a standard deviation of 1/sqrt(d_in) where d_in
is the input dimensionality of the kernel.
full_megatron = 'full_megatron'
class-attribute
instance-attribute
This is what metaseq calls "full megatron init". It is the init used for Llama 2.
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.
ActivationCheckpointingStrategy
Bases: StrEnum
whole_layer = 'whole_layer'
class-attribute
instance-attribute
Checkpoint every transformer layer.
one_in_two = 'one_in_two'
class-attribute
instance-attribute
Checkpoint one in two transformer layers.
one_in_three = 'one_in_three'
class-attribute
instance-attribute
Checkpoint one in three transformer layers.
one_in_four = 'one_in_four'
class-attribute
instance-attribute
Checkpoint one in four transformer layers.
two_in_three = 'two_in_three'
class-attribute
instance-attribute
Checkpoint two out of every three transformer layers.
three_in_four = 'three_in_four'
class-attribute
instance-attribute
Checkpoint three out of four of every transformer layers.
four_in_five = 'four_in_five'
class-attribute
instance-attribute
Checkpoint four out of five of every transformer layers.
nine_in_ten = 'nine_in_ten'
class-attribute
instance-attribute
Checkpoint nine out of ten of every transformer layers.
fine_grained = 'fine_grained'
class-attribute
instance-attribute
Focus checkpointing on where it is cheap to recompute and saves most memory.