How to train multiple LoRA at the same time, then esemble outputs to make prediction?

[quanphamnus]

Hello, I want to train k LoRA at the same time, with the same 1 base_model. For example, in my forward_logits function, I loop over self.adapter_names, then for each adapter, I use self.set_adapter(adapter) and forward through model and get corresponding output. After for loop, I concat these prediction to return a 3d tensor with shape (number of lora adapters, batchsize, dimension)

def forward_logits(self, batch, **kwargs) -> torch.Tensor:
    if self.args.dataset_type == "mcdataset":
        inputs, _, _ = batch

        inputs['input_ids'] = inputs['input_ids'].cuda()
        inputs['attention_mask'] = inputs['attention_mask'].cuda()
        # ensemble the results.
        logits_list = []
        for adapter in self.adapter_names:
            self.set_adapter(adapter)
            output = self.base_model(**inputs)
            logits = output.logits[:, -1, self.target_ids]
            logits_list.append(logits)
    return torch.stack(logits_list, dim=0)

In fit function, I average the output of forward_logits function above along dimension 0, then I want to turn on all adapters at the same time to train all of them using self.set_adapter(self.adapter_names).

    raw_logits = self.forward_logits(batch)
    pred_logits = raw_logits.mean(dim=0)
    output = torch.log_softmax(pred_logits, dim=1)
  
    nll = self.loss(output, golds, reduction="mean")
  
    # Turn on grad in all adapter
    self.set_adapter(self.adapter_names)
  
    self.accelerator.backward(nll)
    self.opt.step()
    self.opt.zero_grad()
    self.scheduler.step()

I think this idea will work correctly, but when evaluating I find that it does not work. Can anyone help me, thank you so much!

[BenjaminBossan]

I think this idea will work correctly, but when evaluating I find that it does not work. Can anyone help me, thank you so much!

Could you describe in more detail what exactly is not working?

One common issue users have with this type of training is that the optimizer is not aware of all the LoRA parameters when it is initialized. So e.g. if you do:

optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], ...)

and at this time, only LoRA adapter 0 is active, then the optimizer does not know about adapter 1, 2, etc. and will not update them.

[SahilKumar75]

There are two separate bugs in your setup BenjaminBossan touched on one of them, but there's a second one that's actually causing the evaluation failure. Let me go through both.

Bug 1: Optimizer not tracking all adapter parameters

This is the one BenjaminBossan mentioned. When you initialize AdamW, only the currently active adapter's parameters have requires_grad=True. So if adapter 0 is active at optimizer creation time, adapters 1 through k-1 are invisible to the optimizer and never get updated regardless of what you do during the forward pass. Fix this by collecting all adapter params upfront:

all_trainable = []
for adapter in self.adapter_names:
    self.base_model.set_adapter(adapter)
    all_trainable.extend([p for p in self.base_model.parameters() if p.requires_grad])

self.opt = torch.optim.AdamW(all_trainable, lr=1e-4)

Bug 2: set_adapter silently disables gradients for non-active adapters

This is the bigger problem. Under the hood, set_adapter(adapter_name) calls BaseTurnerLayer.set_adapter, which sets requires_grad=True for the activated adapter and requires_grad=False for everything else. So your forward loop is doing this:

set_adapter("adapter_0") → adapter_0: requires_grad=True, others: False → forward
set_adapter("adapter_1") → adapter_1: requires_grad=True, adapter_0: False → forward
set_adapter("adapter_2") → adapter_2: requires_grad=True, adapter_1: False → forward

By the time you call set_adapter(self.adapter_names) before backward, the computation graph for adapters 0 through k-2 was already built with requires_grad=False. Gradients are dead for most of your adapters.

The fix

After switching each adapter in the loop, manually restore requires_grad=True for that adapter before the forward pass:

def forward_logits(self, batch, **kwargs):
    inputs = {k: v.cuda() for k, v in batch[0].items()}
    logits_list = []

    for adapter in self.adapter_names:
        self.base_model.set_adapter(adapter)
        for name, param in self.base_model.named_parameters():
            if adapter in name:
                param.requires_grad_(True)
        output = self.base_model(**inputs)
        logits = output.logits[:, -1, self.target_ids]
        logits_list.append(logits)

    return torch.stack(logits_list, dim=0)

def fit_step(self, batch):
    raw_logits = self.forward_logits(batch)
    pred_logits = raw_logits.mean(dim=0)
    output = torch.log_softmax(pred_logits, dim=1)
    nll = self.loss(output, golds, reduction="mean")

    self.base_model.set_adapter(self.adapter_names)
    self.accelerator.backward(nll)
    self.opt.step()
    self.opt.zero_grad()
    self.scheduler.step()

Cleaner alternative

If you don't need separate per-adapter logits and just want all adapters to train toward the same averaged prediction, pass the full list to set_adapter all adapter deltas stack on the base model in a single forward pass:

self.base_model.set_adapter(self.adapter_names)
output = self.base_model(**inputs)

Gradient flow to all adapters works correctly here with no extra handling needed. For your explicit per-adapter logit stacking approach, stick with the loop fix above.