RuntimeError When Integrating LoRA Layers

Hello community,

I’m currently working on finetuning the AnyDoor model by adding LoRA layers, inspired by a technique I found in this post. I’ve integrated LoRA layers into specific parts of the model successfully, but when I start the training process, PyTorch’s autograd throws a RuntimeError One of the differentiated Tensors does not require grad related to tensor differentiation.

Below is the relevant section of my code where I define the LoRA layers and attempt to substitute the original model layers with these:


class LoRALayer(torch.nn.Module):
    def __init__(self, in_dim, out_dim, rank, alpha):
        std_dev = 1 / torch.sqrt(torch.tensor(rank).float())
        self.W_a = torch.nn.Parameter(torch.randn(in_dim, rank) * std_dev)
        self.W_b = torch.nn.Parameter(torch.zeros(rank, out_dim))
        self.alpha = alpha

    def forward(self, x):
        x = self.alpha * (x @ self.W_a @ self.W_b)
        return x

class LinearWithLoRA(torch.nn.Module):
    def __init__(self, linear, rank, alpha):
        self.linear = linear
        self.lora = LoRALayer(linear.in_features, linear.out_features, rank, alpha)

    def forward(self, x):
        return self.linear(x) + self.lora(x)

save_memory = False
if save_memory:

# Configs
resume_path = ".ckpt/epoch=1-step=8687_ft.ckpt"
batch_size = 1
logger_freq = 1000
learning_rate = 1e-5
sd_locked = False
only_mid_control = False
n_gpus = 1
accumulate_grad_batches = 1

# First use cpu to load models. Pytorch Lightning will automatically move it to GPUs.
model = create_model("./configs/anydoor.yaml").cpu()
model.load_state_dict(load_state_dict(resume_path, location="cpu"))
model.learning_rate = learning_rate
model.sd_locked = sd_locked
model.only_mid_control = only_mid_control

for name, param in model.named_parameters():
    param.requires_grad = False

for name, param in model.named_parameters():
    if "model.diffusion_model.output_blocks" in name:
        param.requires_grad = True

lora_r = 8
lora_alpha = 16
lora_dropout = 0.05

assign_lora = partial(LinearWithLoRA, rank=lora_r, alpha=lora_alpha)

for block in model.model.diffusion_model.output_blocks:
    for layer in block:
        # Some Linear layers where I applied LoRA. Both raise the error.
        if isinstance(layer, ResBlock):
            # Access the emb_layers which is a Sequential containing Linear layers
            emb_layers = layer.emb_layers
            for i, layer in enumerate(emb_layers):
                if isinstance(layer, torch.nn.Linear):
                    # Assign LoRA or any other modifications to the Linear layer
                    emb_layers[i] = assign_lora(layer)
        if isinstance(layer, SpatialTransformer):
            layer.proj_in = assign_lora(layer.proj_in)

trainable_count = sum(p.numel() for p in model.parameters() if p.requires_grad == True)
print("trainable parameters: ", trainable_count)

with open("model_parameters.txt", "w") as file:
    for name, param in model.named_parameters():
        file.write(f"{name}: {param.requires_grad}\n")

with open("lora_model.txt", "w") as file:
    print(model, file=file)

# Datasets
DConf = OmegaConf.load("./configs/datasets.yaml")
dataset = VitonHDDataset(**DConf.Train.VitonHD)

dataloader = DataLoader(dataset, num_workers=8, batch_size=batch_size, shuffle=True)
logger = ImageLogger(batch_frequency=logger_freq)
trainer = pl.Trainer(

# Train, dataloader)

I’ve made sure to freeze the parameters of the original model and only allow gradients for the newly added LoRA layers. However, during the training initiation, I encounter the following error:

   self.precision_plugin.backward(self.lightning_module, closure_loss, *args, **kwargs)
  File "/opt/conda/envs/anydoor/lib/python3.8/site-packages/pytorch_lightning/plugins/precision/", line 91, in backward
    model.backward(closure_loss, optimizer, *args, **kwargs)
  File "/opt/conda/envs/anydoor/lib/python3.8/site-packages/pytorch_lightning/core/", line 1444, in backward
    loss.backward(*args, **kwargs)
  File "/opt/conda/envs/anydoor/lib/python3.8/site-packages/torch/", line 487, in backward
  File "/opt/conda/envs/anydoor/lib/python3.8/site-packages/torch/autograd/", line 200, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/opt/conda/envs/anydoor/lib/python3.8/site-packages/torch/autograd/", line 274, in apply
    return user_fn(self, *args)
  File "/home/ubuntu/mnt/myData/AnyDoor/ldm/modules/diffusionmodules/", line 142, in backward
    input_grads = torch.autograd.grad(
  File "/opt/conda/envs/anydoor/lib/python3.8/site-packages/torch/autograd/", line 303, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: One of the differentiated Tensors does not require grad

This error is raised when I call, dataloader) using PyTorch Lightning’s Trainer class.

I’ve already tried enabling torch.autograd.set_detect_anomaly(True) to pinpoint the issue, but the additional information provided hasn’t led me to a clear solution. The error seems to indicate a problem with tensor differentiation, possibly suggesting that a tensor involved in the computation does not have its requires_grad property set correctly. However, I’m not directly manipulating tensors’ requires_grad property except for the initial parameter freezing and subsequent modification to incorporate LoRA layers.

Has anyone encountered a similar issue or can offer insights into what might be causing this error? I’m particularly interested in understanding how to correctly integrate custom layers like LoRA into existing models without disrupting the autograd mechanism.

Any help or pointers would be greatly appreciated!

Solved it. In case can be useful here the solution: