[Feature Request] Read the ._no_weight_decay attribute of parameters

[a858438680]

Feature Request

In FLA, for example, the GDN has code like this

module.A_log._no_weight_decay = True

This means that this parameter (module.A_log) should not be decayed when it is optimized.
However, it seems that this attribute (._no_weight_decay) does not affect the optimizer build function and its usage in flame.
Will this be fixed?

Motivation

Seperate decay settings are common requirements in LLM training.

Your Contribution

If my description is correct. This related to the code of torchtitan. But I think the corresponding code can be extracted and modified. If such thing is allowed, I am willing to submit a PR.