[bug 🐛] weight decay incorrectly applied to LayerNorm and Mamba A, D parameters

[Niccolo-Ajroldi]

Description

The current implementation adds weight decay to all model parameters.

However:

1. Mamba should not have weight decay on A_log and D:
   

   mad-lab/mad/model/layers/mamba.py

   Line 116 in 69d09e2

   self.A_log._no_weight_decay = True

2. It's common practice to not have weight decay on biases and normalization layers.

Fix

I think 1. is a more crucial issue, but we should also include 2. to reflect standard practices in Language Modelling.

#6 implements a fix for both.

It creates two different param_groups for parameters with and without weight decay (see 5ab076a):

decay_params, no_decay_params = [], []
for n, p in self.model.named_parameters():
    if p.requires_grad:
        if not getattr(p, '_no_weight_decay', False) and ("bias" not in n) and ("norm" not in n):
            decay_params.append(p)
        else:
            no_decay_params.append(p)
param_groups = [
    {"params": decay_params, "weight_decay": self.mad_config.weight_decay},
    {"params": no_decay_params, "weight_decay": 0.0},
]

# optimizer:
if self.mad_config.optimizer == 'adamw':
    optimizer = torch.optim.AdamW(
        param_groups,
        lr=self.mad_config.lr
    )

To distinguish normalization layers from other modules, I had to give them a name in the model initialization.
This is achieved by replacing:

self.unembed = nn.Sequential(OrderedDict([
    ('norm', norm(layer_cfg['dim'])), 
    ('lm_head', nn.Linear(dim, vocab_size))
]))

with:

self.unembed = nn.Sequential(norm(layer_cfg['dim']), nn.Linear(dim, vocab_size))
    self.model.append(nn.Sequential(OrderedDict([
        ('norm', norm(layer_cfg['dim'])),
        ('layer', layer(**layer_cfg))
    ])))