DOCSEARCHbyAlgolia
Submit Docs

llama-qrlhf

repository·main·Indexed 41 days ago

https://github.com/lucidrains/llama-qrlhf

A library implementing Q-learning for Large Language Models (LLMs) using an autoregressive formulation. It enables training models to maximize rewards by treating token generation as an action-selection process, featuring a Dueling Q-head architecture, conservative regularization to prevent Q-value overestimation, and support for Rotary Positional Embeddings (RoPE) and GEGLU activations.

Tokens
2.1K
Snippets
9
Records
9
Agent score
51%

What's inside llama-qrlhf

  2. Apply conservative regularization with `conservative_regularization_loss`

    main

    The conservative_regularization_loss function implements a regularization term to prevent overestimation of Q-values. It penalizes Q-values for actions not present in the dataset by pushing them towards a minimum reward value.

    Arguments:

    • q_values: Float['b n a'] - Predicted Q-values for all actions at each time step.
    • states_and_actions: Int['b n'] - The actual actions taken in the dataset.
    • action_mask: Bool['b n'] - A mask indicating which time steps should be included in the loss.
    • reward_min: float - The baseline reward value to regularize towards (defaults to 0.0).
    reg_loss = conservative_regularization_loss(
    q_values=q_values,
    states_and_actions=states_and_actions,
    action_mask=action_mask,
    reward_min=0.0
)

  3. Perform autoregressive Q-learning with `autoregressive_q_learn`

    main

    The autoregressive_q_learn function implements the core Q-learning logic adapted for language modeling. It calculates the loss based on the difference between predicted Q-values and target Q-values, using a hybrid approach: classic Bellman updates for time steps with rewards, and a max-Q target for time steps without rewards (as per the paper's contribution).

    Key features:

    • Discounting: Uses discount_gamma to encourage brevity.
    • Terminal States: If eos_id is provided, the sequence is considered terminal after the first occurrence of that token.
    • Action Masking: Only tokens following the prompt are considered valid actions for loss calculation.

    Arguments:

    • model: The current model being trained.
    • ema_model: The target model (typically an EMA version of the model).
    • states: Int['b n'] - The entire sequence (prompt + generated tokens).
    • prompt_len: Int['b n'] - The length of the prompt part of the sequence.
    • next_states: Int['b'] - The selected action (token) that becomes the next state.
    • rewards: Float['b n'] - Rewards provided at specific time steps.
    • eos_id: int | None - The token ID used to identify the end of a sequence.
    • discount_gamma: float - The reward discount factor.
    loss = autoregressive_q_learn(
    model=model,
    ema_model=ema_model,
    states=states,
    prompt_len=prompt_len,
    next_states=next_states,
    rewards=rewards,
    eos_id=tokenizer.eos_token_id,
    discount_gamma=0.998
)

  4. Initialize and use the Llama model

    main

    The Llama class is the main entry point for the architecture. It implements a transformer-based model with optional support for a Dueling Q-head, which is used for Reinforcement Learning (QRLHF).

    Key parameters:

    • num_tokens: Vocabulary size.
    • dim: Embedding dimension.
    • depth: Number of transformer layers.
    • dim_head: Dimension of each attention head.
    • heads: Number of attention heads.
    • ff_mult: Multiplier for the FeedForward layer dimension.
    • dueling_q_head: Boolean. If True, uses a DuelingHead to output Q-values instead of a standard linear layer.
    • dueling_q_head_expansion_factor: Expansion factor for the Dueling Head's hidden dimension.

    When calling forward(x, return_q_values=True), the model returns both the language modeling logits and the q_values (if dueling_q_head is enabled).

    from llama_qrlhf.llama import Llama
import torch

model = Llama(
    num_tokens = 50257,
    dim = 768,
    depth = 12,
    dueling_q_head = True
)

# x shape: (batch, seq_len)
x = torch.randint(0, 50257, (1, 128))

# Returns (logits, q_values)
logits, q_values = model(x, return_q_values=True)

  5. Reference: RotaryEmbedding parameters

    main

    The RotaryEmbedding class is used to generate frequency components for rotary positional embeddings.

    Parameters:

    • dim: The dimension of the head (dim_head).
    • theta: The base frequency (default 10000).
    from llama_qrlhf.llama import RotaryEmbedding

# dim should match the dim_head used in Attention
rotary_emb = RotaryEmbedding(dim = 64, theta = 10000)

  6. Configure Attention with Rotary Embeddings

    main

    The Attention module implements multi-head attention with support for Rotary Positional Embeddings (RoPE).

    Parameters:

    • dim: Input dimension.
    • dim_head: Dimension per head (default 64).
    • heads: Number of heads (default 8).

    It applies a causal mask during the forward pass to ensure autoregressive properties. If rotary_emb is provided, it is applied to the Query (q) and Key (k) tensors using apply_rotary_pos_emb.

    from llama_qrlhf.llama import Attention
import torch

attn = Attention(dim = 768, dim_head = 64, heads = 8)
x = torch.randn(1, 128, 768)

# rotary_emb should be generated by RotaryEmbedding(dim_head)
# and passed to the forward call
out = attn(x, rotary_emb = some_rotary_embedding)

  7. Configure FeedForward with GEGLU

    main

    The FeedForward module implements a gated feed-forward network using the GEGLU activation function.

    Parameters:

    • dim: Input dimension.
    • mult: Multiplier for the hidden dimension (default 4).

    The architecture follows: RMSNorm -> Linear (to dim_hidden * 2) -> GEGLU -> Linear (to dim).

    from llama_qrlhf.llama import FeedForward
import torch

ff = FeedForward(dim = 768, mult = 4)
x = torch.randn(1, 128, 768)

out = ff(x)

  8. Initialize the `QRLHF` training module

    main

    The QRLHF class is the main wrapper for the Q-learning training process. It manages both the active model and a target model using Exponential Moving Average (EMA) to stabilize training.

    Arguments:

    • model: The base torch.nn.Module (e.g., a Llama model).
    • dataset: The torch.utils.data.Dataset used for training.
    • accelerate_kwargs: Dictionary of arguments passed to Hugging Face accelerate.
    • ema_kwargs: Dictionary of arguments passed to the EMA wrapper (e.g., beta=0.99).
    from llama_qrlhf.llama_qrlhf import QRLHF

# Assuming 'model' and 'dataset' are already defined
trainer = QRLHF(
    model=model,
    dataset=dataset,
    ema_kwargs=dict(beta=0.99)
)

  9. Use the DuelingHead for Q-value estimation

    main

    The DuelingHead is a specialized module designed to decompose Q-values into state values and action advantages. This is a core component for the QRLHF approach.

    It uses a stem (Linear + SiLU) to project the input dimension to a hidden dimension (controlled by expansion_factor), then splits into:

    1. to_values: Predicts a scalar value for the state.
    2. to_advantages: Predicts advantages for each token in the vocabulary.

    The final Q-values are calculated as q_values = values + (advantages - advantages.mean()).

    from llama_qrlhf.llama import DuelingHead
import torch

dueling_head = DuelingHead(
    dim = 768,
    num_tokens = 50257,
    expansion_factor = 2
)

# x shape: (batch, seq_len, dim)
x = torch.randn(1, 128, 768)

q_values = dueling_head(x)
# q_values shape: (1, 128, 50257)