LoRA: Low-Rank Adaptation of Large Language models

LLMs have been present in this world for a long time now. However when it comes to training LLMs for specific tasks, it creates bias within the LLM as it might change the logits of tokens in the LLM. To overcome this problem LoRA came into picture where it is used to freeze the actual model and then train the LoRA parameters for that specific purpose. Generally these vectors are stored in an vector db which lead to the evolution of RAGs. So that based on the requirement of the LLMs these LoRA parameters can be plugged into the LLMs and that can be used for the specific purpose needed.

Suppose we have N tasks like Task 1, Task 2... Task N. But now the LLM not trained for that so it needs to give some specific response based on the Task asked. However, now based on the query asked by the user, the embeddings of the query is compare with the embeddings of the query in the vector db and based on the cosine similarity that LoRA parameters are loaded and plugged with the LLM and then the response is generated. This lead to the origination of Agents.

LLMs have huge amount of parameters. Larger the no. of parameters more better the model. However for small and complex tasks modifying all the parameters doesn't make sense. So we keep the parameters frozen for all the parameters of the model and attach parallel layers to attention and mlp layers. This parallel connection are the LoRA adapters which are the trainable parameters for that specific tasks. We modify the rank and alpha for the LoRA adapters based on complexity of the tasks in hand.

class loRA(nn.Module):
    """LoRA module for low-rank adaptation of large language models. This module consists of two linear layers: an up projection and a down projection.
    The up projection maps the input to a lower-dimensional space defined by the rank, and the down projection maps it back to the original output space.
    The output of the down projection is scaled by alpha/rank to control the contribution of the LoRA module to the overall model output."""
    def __init__(self, in_features: int, out_features: int,
                 rank: int = 16, alpha: float = 1.0,
                 dtype: torch.dtype = torch.float32,
                 device=None):
        super().__init__()
        self.alpha = alpha
        self.rank = rank
        self.up_proj = nn.Linear(in_features, rank, bias=False, dtype=dtype, device=device)
        self.down_proj = nn.Linear(rank, out_features, bias=False, dtype=dtype, device=device)
        nn.init.zeros_(self.down_proj.weight)
        nn.init.kaiming_uniform_(self.up_proj.weight, a=math.sqrt(5))

    def forward(self, X):
        projection = self.down_proj(self.up_proj(X))
        return (self.alpha / self.rank) * projection

This is a simple LoRA layer created. Here you can see that the down_proj is initialized as zeros which means initially they have zero weights and the up_proj are kept as kaiming_uniform so that all the logits have equal probability.

Now we sum the frozen layer output with the LoRA module output

frozen_output + loRA_projection.to(frozen_output.dtype)

Now this is injected into the LLM based on the projection,

_default_projections = [
    "q_proj", "k_proj", "v_proj", "o_proj",   # attention
    "gate_proj", "up_proj", "down_proj"         # MLP
]

def create_module_dict_inject_loRA(self):
    self.module_dict = {}
    for name, module in self.model.named_modules():
        if isinstance(module, nn.Linear):
            module.requires_grad_(False)
            if any(projection in name for projection in self.projections):
                # print(f"Injecting LoRA module into {name} with shape {module.weight.shape}")
                loRA_linear = loRALinear(module, rank=self.rank, alpha=self.alpha, device=self.device)
                parent_path, attr_name = name.rsplit(".", 1)
                parent_module = self.model.get_submodule(parent_path)
                setattr(parent_module, attr_name, loRA_linear)

I have used multiple techniques for optimization.

1. 1

   KV cache

   I have used the pytorch provided KV cache to cache the K and V logits which uses cache position to identify the next position to add key and value. This is generally used in generation of tokens in a transformer. Since we have already generated the older tokens so we can store them in a KV cache and only add the logits of the newer tokens as it gets generated. This reduces the memory usage and also increases the speed of generation of tokens. I have used DynamicCache provided by pytorch which is a dynamic KV cache which can be used for any length of tokens.

2. 2

   Temperature

   I have introduced temperature based sampling which flattens the logits to some extent and gives a diverse response to the given query. This helps in generation of a more diverse set of prompts which helps in training the model better. The temperature is a hyperparameter which can be tuned based on the requirement of the model. This is implemented as follows:

   def temperature_sampling(self, temperature, next_token_logits, batch_size):
           # Returns temperature-scaled LOGITS, not probabilities: the caller feeds this into
           # topk/Categorical(logits=...), which does its own softmax. Returning softmax output
           # here made the caller re-derive pseudo-logits from an already-normalized distribution,
           # silently washing out the temperature's effect.
           temperature = max(temperature, 1e-6)
           logits = next_token_logits / temperature
           logits = logits.float()
           logits = logits - logits.max(dim=-1, keepdim=True).values
           return logits

3. 3

   Top k Sampling

   Instead of taking max logit value token, I have used categorical sampling for the top k values and generate tokens from that response. Instead of taking max logits value token, I am taking the top k logits values and then sampling from that distribution. This helps in generation of a more diverse set of prompts which helps in training the model better. This is implemented as follows:

   next_token_logits = next_token_logits - next_token_logits.max(dim=-1, keepdim=True).values
                       values, _ = torch.topk(next_token_logits, top_k)
                       min_val = values[:, -1].unsqueeze(-1)
                       next_token_logits = torch.where(next_token_logits < min_val, float('-inf'), next_token_logits)
                       next_token = torch.distributions.Categorical(logits=next_token_logits).sample()  # Sample from the distribution

4. 4

   Checkpointing

   Used Checkpointing so it takes less VRAM during the train and VRAM can be used for other purposes like increasing the GRPO batch, generation of more tokens etc. Checkpointing saves the parameters of the layers and recomputes them during the backward pass. This is by default supported in pytorch. However if you are customizing the model like in this case I am modifying the parameters then we need to enable it by:

   self.model.gradient_checkpointing_enable()  # Enable gradient checkpointing for memory efficiency

   and then save it per layer like this:

   for layer_number in range(self.model.config.num_hidden_layers):
               if self.training:
                   # gradient_checkpointing_enable() in __init__ has no effect here: it only
                   # flips a flag that the model's own forward() checks, but we bypass that
                   # entirely by driving layers through action_per_layer(). Without an explicit
                   # checkpoint call, every layer's attention/MLP activations for the full
                   # batch x seq_len are kept alive for backward, which is the dominant memory cost.
                   input = checkpoint(
                       self.action_per_layer, layer_number, input,
                       attention_mask, (cos, sin), None, None,
                       use_reentrant=False,
                   )
               else:
                   input = self.action_per_layer(layer_number, input, attention_mask=attention_mask, position_embeddings=(cos, sin))

5. 5

   Logsumexp

   Used this one of the tricks to improve numerical stability during the softmax operation. As softmax is a memory heavy operation, using logsumexp helps in reducing memory usage and avoiding numerical overflow.

6. 6

   Base model 8-bit

   Used 8-bit quantization for base model as we only use it for logits comparisons. This is a memory efficient approach. The base model is used for calculation of the KL divergence term in the loss calculation during training in a GRPO.

7. 7

   Replay Buffer

   Replay Buffers have been present in the reinforcement learning literature for a long time. They store past experiences (queries and their corresponding rewards) so that the model can learn from them multiple times, improving sample efficiency and stability during training. This I have added in my reward training so that when the model loss deviates too much from the actual path these prompts from the replay buffer can bring them back to the correct trajectory.

I did this on a Qwen3 0.6B model. I downloaded the model from hugging face and I studied through all the layers of the model. There are 28 layers in this model and each layer has an attention and mlp layer and I injected LoRA adapters into all the layers and used it for training the model on chess. A whole chess game can be written using a SAN notation. Where each previous move determines the current state of the chess board. This is what LLMs are good at. This makes it a perfect example for training chess using LLM. Here is the logs of the LLM for each layer.

Loading weights: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████| 310/310 [00:00<00:00, 359.69it/s]

model
model.embed_tokens
model.layers
model.layers.0
model.layers.0.self_attn
model.layers.0.self_attn.q_proj
model.layers.0.self_attn.k_proj
model.layers.0.self_attn.v_proj
model.layers.0.self_attn.o_proj
model.layers.0.self_attn.q_norm
model.layers.0.self_attn.k_norm
model.layers.0.mlp
model.layers.0.mlp.gate_proj
model.layers.0.mlp.up_proj
model.layers.0.mlp.down_proj
model.layers.0.mlp.act_fn
model.layers.0.input_layernorm
model.layers.0.post_attention_layernorm
model.layers.1
model.layers.1.self_attn
model.layers.1.self_attn.q_proj
model.layers.1.self_attn.k_proj
model.layers.1.self_attn.v_proj
model.layers.1.self_attn.o_proj
model.layers.1.self_attn.q_norm
model.layers.1.self_attn.k_norm
model.layers.1.mlp
model.layers.1.mlp.gate_proj
...
model.layers.26.mlp.down_proj
model.layers.26.mlp.act_fn
model.layers.26.input_layernorm
model.layers.26.post_attention_layernorm
model.layers.27
model.layers.27.self_attn
model.layers.27.self_attn.q_proj
model.layers.27.self_attn.k_proj
model.layers.27.self_attn.v_proj
model.layers.27.self_attn.o_proj
model.layers.27.self_attn.q_norm
model.layers.27.self_attn.k_norm
model.layers.27.mlp
model.layers.27.mlp.gate_proj
model.layers.27.mlp.up_proj
model.layers.27.mlp.down_proj
model.layers.27.mlp.act_fn
model.layers.27.input_layernorm
model.layers.27.post_attention_layernorm
model.norm
model.rotary_emb
lm_head

The LoRA adapters are injected into the attention and mlp layers of the transformer block. This is after addition of LoRA parameters to the model.

model
model.model
model.model.embed_tokens
model.model.layers
model.model.layers.0
model.model.layers.0.self_attn
model.model.layers.0.self_attn.q_proj
model.model.layers.0.self_attn.q_proj.frozen_linear_layer
model.model.layers.0.self_attn.q_proj.loRA_module
model.model.layers.0.self_attn.q_proj.loRA_module.up_proj
model.model.layers.0.self_attn.q_proj.loRA_module.down_proj
model.model.layers.0.self_attn.k_proj
model.model.layers.0.self_attn.k_proj.frozen_linear_layer
model.model.layers.0.self_attn.k_proj.loRA_module
model.model.layers.0.self_attn.k_proj.loRA_module.up_proj
model.model.layers.0.self_attn.k_proj.loRA_module.down_proj
model.model.layers.0.self_attn.v_proj
model.model.layers.0.self_attn.v_proj.frozen_linear_layer
model.model.layers.0.self_attn.v_proj.loRA_module
model.model.layers.0.self_attn.v_proj.loRA_module.up_proj
model.model.layers.0.self_attn.v_proj.loRA_module.down_proj
model.model.layers.0.self_attn.o_proj
model.model.layers.0.self_attn.o_proj.frozen_linear_layer
model.model.layers.0.self_attn.o_proj.loRA_module
model.model.layers.0.self_attn.o_proj.loRA_module.up_proj
model.model.layers.0.self_attn.o_proj.loRA_module.down_proj
model.model.layers.0.self_attn.q_norm
model.model.layers.0.self_attn.k_norm
model.model.layers.0.mlp
model.model.layers.0.mlp.gate_proj
model.model.layers.0.mlp.gate_proj.frozen_linear_layer
model.model.layers.0.mlp.gate_proj.loRA_module
model.model.layers.0.mlp.gate_proj.loRA_module.up_proj
model.model.layers.0.mlp.gate_proj.loRA_module.down_proj
model.model.layers.0.mlp.up_proj
model.model.layers.0.mlp.up_proj.frozen_linear_layer
model.model.layers.0.mlp.up_proj.loRA_module
model.model.layers.0.mlp.up_proj.loRA_module.up_proj
model.model.layers.0.mlp.up_proj.loRA_module.down_proj
model.model.layers.0.mlp.down_proj
model.model.layers.0.mlp.down_proj.frozen_linear_layer
model.model.layers.0.mlp.down_proj.loRA_module
model.model.layers.0.mlp.down_proj.loRA_module.up_proj
model.model.layers.0.mlp.down_proj.loRA_module.down_proj
model.model.layers.0.mlp.act_fn
model.model.layers.0.input_layernorm
model.model.layers.0.post_attention_layernorm
...
model.model.layers.27.mlp.up_proj.loRA_module.down_proj
model.model.layers.27.mlp.down_proj
model.model.layers.27.mlp.down_proj.frozen_linear_layer
model.model.layers.27.mlp.down_proj.loRA_module
model.model.layers.27.mlp.down_proj.loRA_module.up_proj
model.model.layers.27.mlp.down_proj.loRA_module.down_proj
model.model.layers.27.mlp.act_fn
model.model.layers.27.input_layernorm
model.model.layers.27.post_attention_layernorm
model.model.norm
model.model.rotary_emb
model.lm_head
loss_function

I used Supervised Training with 400k data points with 4 epochs to make the model learn the basic rules of chess. I used a batch size of 24 and a learning rate of 1e-5. The model was trained on a single GPU (RTX PRO 6000) for 4 epochs. The loss function used was Cross Entropy Loss. After the supervised training, the model was able to generate valid chess moves and the format was consistent with the SAN notation.

I used GRPO based setup where we generate multiple samples from the query provided and then give reward to each output. Some responses would be wrong, some responses would be correct and based on it we reward the output. We also added KL divergence term between the current model/policy and the base model/ base policy. This keeps the model not get diverged a lot from the base setup which may lead to instability of the model.

divergence = torch.clamp(log_probs - base_log_probs, min=-10, max=10)
ratio = torch.exp(divergence)

reward_batch = reward_batch.to(self.model_device).float()
normalized_reward = torch.tanh(reward_batch)
advantage = (normalized_reward - normalized_reward.mean()) / (normalized_reward.std() + 1e-6)  # Normalize advantages
weights = advantage * 2.0
weights = weights.unsqueeze(-1).unsqueeze(-1)

product = weights.float() * ratio.float()
product_clamped = weights.float() * torch.clamp(ratio.float(), 1.0 - self.epsilon, 1.0 + self.epsilon)
loss = -torch.min(product, product_clamped) + self.beta * divergence
loss.nan_to_num_(nan=0.0, posinf=50.0, neginf=-50.0)
loss.clamp_(min=-50.0, max=50.0)

Here you can see that the loss function is product of advantage and ratio which is added to the relative divergence between the base_log_probs and the log_probs.

I have used the SAN notation of chess for training the model. I extract the SAN notation from the generation and if the move played is correct in terms of the current board state then I reward the model. LLMs have a problem of repeating the same response again and again till the no. of tokens in the generation is exhausted. For preventing this penalize repetition of the moves response. Here is my reward function I used in the training

def board_state(self, prompt_moves, generation, play_as="white"):
    chess_board = ChessGame()
    san = r'(?:O-O-O|O-O|[KQRBN]?[a-h]?[1-8]?x?[a-h][1-8](?:=[QRBN])?[+#]?)'
    pattern = rf'(\d+)\.\s+({san})(?:\s+(?!\d+\.)({san}))?'
    san_str = re.findall(pattern, prompt_moves)
    init_moves_made = 0
    count_valid_moves = 0
    reward = 0.0
    for m in san_str:
        for san_values in m:
            for san_values in san_values.split(" "):
                if len(san_values.strip()) > 0:
                    if san_values.isdigit():
                        init_moves_made = int(san_values)
                    elif san_values[-1] == "." and san_values[:-1].isdigit():
                        init_moves_made = int(san_values[:-1])
                    else:
                        if self.is_valid_san(san_values):
                            ok, _ = chess_board.push_san(san_values)
                            if ok:
                                count_valid_moves += 1

    san = r'(?:O-O-O|O-O|[KQRBN]?[a-h]?[1-8]?x?[a-h][1-8](?:=[QRBN])?[+#]?)'
    pattern = rf'(({san}\s+)?(?:\d+\.\s+{san}(?:\s+(?!\d+\.){san})?\s*)+)'
    san_str = re.findall(pattern, generation)
    atleast_one_valid_move = False
    generation_move_no = -1
    all_correct_moves = 0
    for m in san_str:
        for san_values in m:
            for san_values in san_values.split(" "):
                if len(san_values.strip()) > 0:
                    if san_values.isdigit() and int(san_values) < init_moves_made:
                        generation_move_no = int(san_values)
                    elif san_values[-1] == "." and san_values[:-1].isdigit():
                        generation_move_no = int(san_values[:-1])
                    else:
                        if self.is_valid_san(san_values):
                            atleast_one_valid_move = True
                            if generation_move_no < init_moves_made or generation_move_no != -1:
                                continue
                            ok, _ = chess_board.push_san(san_values)
                            if ok:
                                reward += 1
                                all_correct_moves += 1
                                print(f"Valid move made: {san_values} : Current reward: {reward}")
    if generation.count("moves:") > 1:
        reward -= 1.0
    if atleast_one_valid_move:
        reward += 0.5
    if all_correct_moves > 20:
        reward += 15.0
    elif all_correct_moves > 10:
        reward += 10.0
    elif all_correct_moves > 5:
        reward += 5.0
    if reward >= 1:
        print(chess_board.board_string())
    if reward == 0 and not atleast_one_valid_move:
        reward = -1.0
    return reward