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
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!
Source: facebookresearch/mae