Skip to content

xlm.backbones.llada.modeling_llada

LLaDA modeling (OLMo-style bidirectional decoder).

Ported from the GSAI-ML/LLaDA-8B-Base Hub repo (modeling_llada.py, trust_remote_code module). Module/parameter names are kept identical to the checkpoint so Hub weights load with strict=True and zero remapping (model.transformer.wte, model.transformer.blocks.*, model.transformer.ln_f, model.transformer.ff_out).

Changes vs upstream (none affect state-dict keys):

  • imports from .configuration_llada use LLaDAConfigBase (renamed from the upstream LLaDAConfig; variant configs subclass it in xlm-models/llada).
  • optional position_ids threaded through LLaDAModelLM.forward -> LLaDAModel.forward -> blocks -> RotaryEmbedding so per-token RoPE positions (e.g. packed positions from an attention-mask cumsum) are supported, mirroring what the Dream backbone offers.
  • no AutoModel.register call (would clash with the remote-code module when both are imported in one process).

LayerNorm

Bases: LayerNormBase

The default :class:LayerNorm implementation which can optionally run in low precision.

RMSLayerNorm

Bases: LayerNormBase

RMS layer norm, a simplified :class:LayerNorm implementation

GemmaRMSLayerNorm

Bases: LayerNormBase

Gemma RMS layer norm, a simplified :class:LayerNorm implementation

RotaryEmbedding

Bases: Module

Rotary positional embeddings (RoPE).

LLaDABlock

Bases: Module

A base class for transformer block implementations.

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).

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.

LLaDAOutput

Bases: NamedTuple

logits instance-attribute

A tensor of shape (batch_size, seq_len, vocab_size) representing the log probabilities for the next token before normalization via (log) softmax.

attn_key_values instance-attribute

Attention keys and values from each block.

hidden_states instance-attribute

Hidden states from each block.

LLaDAGenerateOutput

Bases: NamedTuple

token_ids instance-attribute

The generated token IDs, a tensor of shape (batch_size, beam_size, max_steps). These do not include the original input IDs.

scores instance-attribute

The scores of the generated sequences, a tensor of shape (batch_size, beam_size).

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.