Gradient checkpoint with multiple adapters

[NguyenRichard]

I'm not sure if it can be considered as a bug since I might be using the library differently from how it's supposed to be used.

Context:

I have a PeftModel that need to be infered with 2 different inputs.
For each input I have a pretrained adapter that is frozen and a new adapter for finetuning.

My forward does:

for name, x in inputs:
   mypeft_model.base_model.set_adapter([name+'pretrain',name+'ft'])
   custom_set_pretrain_grad_false_ft_true() #Doing it because set_adapter force gradients to True cf 2759#issue-3363985341
   feature = mypeft_model(x)

(#2759 (comment))
Issue:

1. if mypeft_model contains cp.checkpoint(mymodule, x), the backpropagation will not update properly the weight of the LoRA layers in my module either because it did not 'see the set_adapter' or it did not 'see the force grad'
2. A work around I have found is to wrap the whole code inside the loop with a cp.checkpoint but it's super heavy on the memory as I have to store all in GPU until the end of the backbone (ViT-G 40 blocks transformers)

Question:
Is there anyway to 'provide' the context to the backpropagation even using gradient checkpointing when switching adapters in the forward?
I have not explored huggingface transformers.enable_gradient_checkpointing() since I'm using a custom model and I'm unsure if it fits for my problem.

[BenjaminBossan]

Thanks for the report. PEFT does not implement any specific checkpointing logic. Instead, we rely completely on gradient checkpointing provided by the modeling code (e.g. transformers).

With the information you provide, it is hard to tell if there is anything more we can do in PEFT, I'm not sure if you can provide a minimal reproducer. If you have any suggestion with regard to how PEFT could make it easier for you to add gradient checkpointing, please let us know.

[NguyenRichard]

Hi, I think it will be difficult to make a minimal reproducer (considering it needs a model, data, training loop etc) and I dont think I have enough experience in gradient checkpointing and PEFT to provide meaningful suggestions.

However, I can try to explain deeper in details what I think is the edge case for future reference. I have also found today a better work around.

Considering this fictive backbone with a common usage of cp.checkpoint to avoid the storage in memory of the attention computation in the blocks.

class ViViT(nn.Module):
		def __init__(self, with_cp):
				self.blocks = nn.ModuleList([Block(...)]) # Your transformer block with the Attention
				self.norm = nn.LayerNorm()
  				self.with_cp = with_cp
		forward(self, x):
				for blk in self.blocks:
		            if self.with_cp:
		                x = cp.checkpoint(blk, x)
		            else:
		                x = blk(x)

        		x = self.norm(x)
        		return x

This multimodal model that will contain a PeftModel of the backbone.

class MultimodalModel(nn.Module):
		def __init__(pretrained_adapters:dict[str,str], backbone_embed_dim, num_classes):
			self.backbone = ViViT(backbone_embed_dim)
			lora_config = LoraConfig(
                r=8,
                lora_alpha=32,
                target_modules=["qkv", "proj"],
                lora_dropout=0.1,
                bias="none",
                modules_to_save=[]
            )

           	self.modalities = []
			for i, (name, path) in enumerate(pretrained_adapters.items()):
				self.modalities.append(name)
				# Add pretrained adapter
	    		if i == 0:
	    			self.backbone = PeftModel.from_pretrained(
                    		provider_backbone,
                    		path,
                    		adapter_name=name,
                    		is_trainable=False
                    ft_adapter_name = f"{name}_ft"

                )
                else:
                	self.backbone.load_adapter(path, 
                			adapter_name=name, 
                			is_trainable=False)
                # Add finetuning adapter to train
                self.backbone.add_adapter(ft_adapter_name, lroa_config)
                self.head_embed_dim = len(modalities) * backbone_embed_dim
                self.num_classes = num_classes
                self.head = nn.Linear(self.head_embed_dim, self.num_classes)

        def enforce_gradient(self, backbone):
        	trainable_params = 0
        	frozen_params = 0
            for name, param in backbone.named_parameters():
                if f'_ft.' in name:
                    param.requires_grad = True
                    trainable_params += param.numel()
                else:
                    frozen_params += param.numel()
                    param.requires_grad = False

        def _adapter_switching(self, backbone, modality):
        	if isinstance(backbone, PeftModel):
        			pretrained_adapter = modality
        			ft_adapter_name = f"{modality}_ft"
        	adapters_to_activate = []
        	if pretrained_adapter in backbone.peft_config:
        		adapters_to_activate.append(pretrained_adapter)
        	if ft_adapter_name in backbone.peft_config:
        		adapters_to_activate.append(ft_adapter_name)
        	if adapters_to_activate:
				backbone.base_model.set_adapter(adapters_to_activate)
			else:
				backbone.base_model.disable_adapter_layers()

			if self.training:
				self.enforce_gradient(backbone)

		def forward(self, inputs: dict[str, torch.Tensor])
			features = []
			for (modality, input) in inputs:
				self._adapter_switching(self.backbone,modality)
				features[modality].append(self.backbone(input).mean(dim=1))
			feature_representation = torch.cat(features, dim=1)
			return self.head(feature_representation)

The bug in the backward propagation:
Training the multimodal model in this setup, the weights of the LoRA layers wont be updated (I checked the .pth after the training).
I believe the reason comes from the fact that the LoRA layers in the Attention module in the Block are encapsulated in the cp.checkpoint. When they are recalculated in the backpropagation, the backbone wont be set in the correct state by _adapter_switching like it was for the forward. I suspect here that not only it will not update the correct gradient but also that it will 'look' at the wrong weights (most likely the last set of adapters switch to)

Old work around (Allow to not save in memory every forward in the backbone, but still huge GPU usage in runtime if the backbone is heavy as all the computation are stored until we exit cp.checkpoint())

def _backbone_call(backbone, modality, input):
	self._adapter_switching(backbone,modality)
	return self.backbone(input).mean(dim=1)

remove cp.checkpoint in the backbone itself and wrap the _backbone_call instead:

def forward(self, inputs: dict[str, torch.Tensor])
	features = []
	for (modality, input) in inputs:
		if self.with_cp:
			feature = cp.checkpoint(
					self._backbone_call,
					backbone,
					modality,
					input,
					use_reentrant=False
                                        )
			features.append(feature)
		else:
			self._adapter_switching(self.backbone,modality)
			features[modality].append(self.backbone(input).mean(dim=1))
	feature_representation = torch.cat(features, dim=1)
	return self.head(feature_representation)

New work around (Replicate the cp.checkpoint inside each individual block by using chunk_size = 1)

def _backbone_blocks_call(backbone, modality, blocks_segment, input, is_first_segment):
	if is_first_segment:
		self._adapter_switching(backbone, modality)
	return self.backbone(input).mean(dim=1)
	for block in blocks_segment:
		input = block(input)
	return input

def forward(self, inputs: dict[str, torch.Tensor])
	features = []
	if self.with_cp:
		chunk_size = 1 # 1 Reproduce the cp.checkpoint inside each block but can use higher number
		for (modality, input) in inputs:
			for i in range(0, len(backbone.blocks), chunk_size):
				is_first = (i == 0)
				blocks =  backbone.blocks[i:i + chunk_size]
				feature = cp.checkpoint(
						self._backbone_blocks_call,
						backbone,
						modality,
						blocks,
						input,
						is_first,
						use_reentrant=False
						)
			features.append(backbone.norm(feature))
	else:
		self._adapter_switching(self.backbone,modality)
		features[modality].append(self.backbone(input).mean(dim=1))

	feature_representation = torch.cat(features, dim=1)
	return self.head(feature_representation)

So far it has been working (at least no visible issues and lora weights are trained again)

I dont really have any idea how to make PEFT compatible with such use case without work around like that. I'm assuming that the library was not made to be used nested (and with cp.checkpoint) like that. The main issue is that there is a conflict on the usage of set_adapter and cp.checkpoint as the later won't have the general context of the PEFT model state at the time of running the nested code.

If you have any question about this edge use case, I'm available to answer to the extent of my understanding. You're also free to close the issue as I think it is working on my end now.

[BenjaminBossan]

Thanks for investigating further and your detailed description of how you solved the issue. I hope that this can serve as inspiration for future users who encounter a similar problem.

I agree with your assessment that the conflict most likely stems from switching LoRA adapters while wrapping the calls into checkpoints (you don't show it explicitly, but I assume that the LoRA config is targeting modules inside of the Block). To understand that better, we would have to investigate the exact implementation of checkpointing in PyTorch, which I assume won't be trivial.

The only other workaround I could think of would be to only switch adapters between performing the backwards pass, i.e. use adapter 1 => loss => backwards, then switch to adapter 2 => loss => backwards, etc. However, I assume that this is not possible with your training setup, which is why you had to choose a different solution.

Also, if it's possible to avoid gradient checkpointing and choosing other options to save memory like using a quantized base model, reducing the batch size, etc. it should be considered, especially since gradient checkpointing has a runtime penalty, but sometimes, these other options have been exhausted and checkpointing is necessary.