FSDP with named_parameters for disabling weight_decay on some parameters.

[fmocking]

Hi, FSDP requires to use self.trainer.model.parameters() however to disable the weight decay on some parameters (norm and bias) using self.trainer.model.named_parameters() still triggers the error. Here is the code that I'm trying:

    def get_param_groups(self):

        bias_layers = []
        norm_layers = []
        weight_layers = []

        for parameter_name, parameter in self.trainer.model.named_parameters():
            if (
                    "LayerNorm" in parameter_name
                    or "layer_norm" in parameter_name
                    or "BatchNorm" in parameter_name
            ):
                norm_layers.append(parameter)
            elif "bias" in parameter_name:
                bias_layers.append(parameter)
            else:
                weight_layers.append(parameter)

        return [
            {"params": norm_layers, "weight_decay": 0},
            {"params": bias_layers, "weight_decay": 0},
            {"params": weight_layers, "weight_decay": self.weight_decay},
        ]

Unfortunately no luck so far. Any ideas?