enable param group configuration in llm-foundry

[vchiley]

This PR enables param group configuration in llm-foundry.

The optimizer_config defines the optimizer args.
This PR allows the user to additionally have key disable_grad which is a string or list of strings. If a string matches a parameter name, then that parameter will have requires_grad=False. This is useful for freezing parameters.
This PR additionally allows the user to specify key param_groups which is a list of dicts. In this dict, key param_str_match defines a string; if a parameter name contains this string, then it will be in this parameter group. This is useful for grouping parameters together. The dict can also contain any other key that is a valid optimizer arg.
Note: to handle name overlap conflicts, params are assigned to parameter groups and added to param_groups in the order that param_str_match appear in param_groups.

Param name comparisons are done using RegEx search.

Usage
To disable gradient for all parameters that contain the string "norm" or "bias":

    optimizer_config: {
        "name": "decoupled_lionw",
        "lr": 1e-3,
        "weight_decay": 1e-2,
        "betas": [0.9, 0.999],
        "eps": 1e-8,
        "disable_grad": ["norm", "bias"]
    }

or in the yaml as:

optimizer:
  name: decoupled_lionw
  lr: 1e-3
  weight_decay: 1e-2
  betas:
  - 0.9
  - 0.999
  eps: 1e-8
  disable_grad:
  - norm
  - bias

To create modify the optimizer parameters for all parameters that contain the string "norm" separately:

    optimizer_config: {
        "name": "decoupled_lionw",
        "lr": 1e-3,
        "weight_decay": 1e-2,
        "betas": [0.9, 0.999],
        "eps": 1e-8,
        "param_groups": [
            {
                "param_str_match": "norm",
                "lr": 1e-4,
                "weight_decay": 0.0,
            },
        ],
    }

of in yaml form:

optimizer:
  name: decoupled_lionw
  lr: 1e-3
  weight_decay: 1e-2
  betas:
  - 0.9
  - 0.999
  eps: 1e-8
  param_groups:
  - param_str_match: norm
    lr: 1e-4
    weight_decay: 0

[vchiley]

potential users: @sashaDoubov @samhavens @b-chu @bcui19 @ShashankMosaicML

[j316chuck]

Nice ✅ . Left a comment but thanks for the clean implementation.