Merge lora

Merging the LoRA Adapter into the Base Model

When training completes, the final checkpoint is written to disk. The checkpoint directory does not contain a copy of the full GaMS base model. It contains only the adapter: the two low-rank matrices \(\mathbf{A}\) and \(\mathbf{B}\) for each targeted linear layer, along with the adapter configuration and optimizer state. For a model with billions of parameters, this distinction is profound. Where a full fine-tuned checkpoint might occupy tens or even hundreds of gigabytes, a LoRA adapter trained with rank 64 occupies a small fraction of that space. The frozen base model weights — which constitute the overwhelming majority of the parameter count — were never touched during training, and therefore need not be stored again.

Two Paths to Inference

Once training is complete, there are two natural ways to deploy the result, and understanding their tradeoffs shapes how we think about the merge step.

The first approach is to leave the adapter and base model separate and combine them at runtime. The PEFT library makes this straightforward: load the base model, call PeftModel.from_pretrained with the adapter path, and the combined model is ready for inference. Moreover, libraries such as vLLM let us use multiple different LoRA adapters with the same base model during the inference. This approach is useful when we trained multiple different adapters (either during the hyperparameter search or for different tasks). We can then serve multiple "different" (same base model, different adapter) models while keeping the memory usage low. Additional benefit is that we save storage space.

The second approach is to permanently merge the adapter into the base model weights, producing a single standalone model. This merged model has no dependency on the PEFT library at inference time, behaves exactly like any other transformer, and can be loaded directly by vLLM or any other inference framework without any adapter-specific configuration. There is a modest practical speedup as well, since the matrix addition that would otherwise happen at every forward pass has already been computed once and stored. For production deployment — a single well-validated checkpoint that will serve real users — the merged form is cleaner. Additional benefit is that this approach simplifies another training runs we might perform on top of the specialized model.

Our pipeline focuses on the second approach. In this chapter we describe the merging process.

Performing the Merge

The merge script in merge_lora/merge.py takes four arguments: the path to the base model, the path to the tokenizer (which may differ from the base model path in some configurations - it differs in our case as we combined the tokenizer from the Instruct model, which has chat template, with the base model), the path to the adapter checkpoint, and a name for the output directory where the merged model will be saved.

The checkpoint path follows a consistent convention established by the training script. SFTTrainer saves intermediate and final checkpoints under a subdirectory structure of the form lora_training/experiments/<version>/checkpoint-<step_number>, relative to the training experiments directory. If, for example, training ran for three epochs and saved a checkpoint at step 1500, the adapter path would be something like r-64_ws-500_lr-2e-5_min-lr-1e-6/checkpoint-1500. This naming makes it straightforward to merge from any point in training history, not only the final checkpoint — useful when an evaluation loss suggest overfitting in later steps or we want to test the checkpoints at different steps.

With those paths in hand, the merge proceeds as a short sequence of operations:

# Load the frozen base model in bfloat16
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_path,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Attach the adapter
model = PeftModel.from_pretrained(base_model, adapter_checkpoint_path)

# Fuse adapter weights into base and discard adapter structure
model = model.merge_and_unload()

# Save the merged model and tokenizer as a self-contained artifact
model.save_pretrained(output_path)
tokenizer.save_pretrained(output_path)

The tokenizer is saved alongside the merged weights because any downstream consumer of the merged model — whether an inference server, an evaluation harness, or a colleague — expects to find both in the same directory. Keeping them co-located eliminates one class of configuration error.

To submit the merge job on the cluster, we use the provided Sbatch script, passing the four required arguments:

sbatch merge_lora/run_merge.sbatch \
    cjvt/GaMS3-12B \
    cjvt/GaMS3-12B-Instruct \
    r-64_ws-500_lr-2e-5_min-lr-1e-6/checkpoint-1500 \
    GaMS3-12B-Med

The job will request its single GPU, load the base model, attach and merge the adapter, and write the result to the output directory. Within the 30-minute wall-clock limit, even large models merge comfortably, since the computation is a single pass of matrix additions rather than any form of iterative optimization.

The Mathematics of Merging

The merge operation itself is conceptually simple. During training, each modified weight matrix W in the base model accumulated an update parameterized as the product of two low-rank matrices: a down-projection \(\mathbf{A}\) of shape \((r \times d_{in})\) and an up-projection \(\mathbf{B}\) of shape \((d_{out} \times r)\), scaled by the ratio of the LoRA alpha to the rank \(r\). The effective update to the weight matrix is therefore:

\(\mathbf{W}_{\mathrm{merged}} = \mathbf{W}_{\mathrm{base}} + \alpha r^{-1} \cdot \mathbf{B} \times \mathbf{A}\)

With alpha set to twice the rank in our configuration, the scaling factor simplifies to 2.0, but the PEFT library handles all of this bookkeeping internally. The merge_and_unload() method iterates over every adapted layer, computes this sum, replaces the weight tensor with the merged result, and then discards the adapter matrices entirely. What remains is an ordinary model with no trace of the adapter structure — just slightly different weights than the original base.

The bfloat16 precision used throughout training is preserved during the merge. Loading the base model in bfloat16 before attaching the adapter ensures that the arithmetic happens in the same numerical format as training, and the merged weights are saved in bfloat16 as well. This keeps memory requirements during the merge within reason, which is why the merge job requests only a single GPU with 64 gigabytes of RAM and a 30-minute time limit — a modest allocation compared to the multi-node training run.

Configuration Fix - End of Turn

If we proceeded directly to the inference chapter, we would find out that the trained model does not stop during the inference. The reason for that is the tokenization during the SFT training. During the pretraining, the GaMS model was trained to stop on the end of sequence (EOS) token. However, if we inspect the tokens in the first training example (we printed the first example to the training logs), we would not find token index 1, that corresponds to the EOS token (this is an actual bug in the SFTTrainer tokenization). However, we can spot the token index 106 towards the end of the example. This is the <end_of_turn> token, that GaMS uses to mark the end of each conversation turn (either system, user or assistant). Therefore, the GaMS was trained to produce the <end_of_turn> token at the end of the conversation and we will exploit that by adding this token as EOS token to the model's configuration.

To fix this we will go to the models/GaMS3-12B-Med dir, where the merged model's weights and configurations are stored. We will edit the generation_config.json file that looks like this:

{
  "_from_model_config": true,
  "bos_token_id": 2,
  "eos_token_id": 1,
  "pad_token_id": 0,
  "transformers_version": "4.56.1"
}

We can see that the only EOS token here is 1. We can add the <end_of_turn> token by converting the eos_token_id field to list. Since the <end_of_turn> index is 106, we change the configuration to:

{
  "_from_model_config": true,
  "bos_token_id": 2,
  "eos_token_id": [
    1,
    106
  ],
  "pad_token_id": 0,
  "transformers_version": "4.56.1"
}

Now the model will stop on <end_of_turn> token during the inference. If you check the configuration files on HF repositories of cjvt/GaMS3-12B and cjvt/GaMS3-12B-Instruct, you will observe the exact same difference - the first has only 1 as EOS token, while the second has 106 added as well.