#216·mae

Layer-wise learning rate decay (lrd) not applied correctly due to missing lr parameter in param groups

Author: xujialiuCreated Mar 5, 2025Updated Mar 8, 2025

In the current implementation of util.param_groups_lrd, parameter groups are assigned an lr_scale to support layer-wise learning rate decay. However, the lr_scale is not directly utilized to set the actual learning rate (lr) for each parameter group.

PyTorch optimizers like AdamW use the lr key in parameter groups to determine the learning rate. Since the current code only assigns lr_scale without applying it to the base learning rate, ​all layers effectively use the same learning rate, defeating the purpose of layer-wise decay.

**Suggested Fix:**​ In main_finetune.py

python
param_groups = lrd.param_groups_lrd(
    model_without_ddp,
    args.weight_decay,
    no_weight_decay_list=model_without_ddp.no_weight_decay(),
    layer_decay=args.layer_decay,
)
# added codes
for group in param_groups:
    group['lr'] = args.lr * group.pop('lr_scale')

optimizer = torch.optim.AdamW(param_groups, lr=args.lr)

Let me know if I can help with a PR to address this!