Add parq utility to create an optimizer

[metascroy]

Adds utility to create an optimizer like this:

quant_configs_and_filter_fns = [
    (QuantConfig(bitwidth=2, group_size=128), use_2bit),
    (QuantConfig(bitwidth=3, group_size=256), use_3bit),
    (QuantConfig(bitwidth=4, group_size=256), use_4bit),
    (QuantConfig(bitwidth=8), use_8bit),
 ]

optimizer = create_optimizer(model, quant_configs_and_filter_fns, base_optimizer_cls, base_optimizer_kwargs)

The filter functions are (module, fqn) -> bool, e.g.,

def use_2bit(m, fqn):
        if not fqn.endswith(".weight"):
            return False

        # Layers 0-3
        for i in range(4):
            if fqn.startswith(f"model.layers.{i}."):
                for key in ["v_proj", "gate_proj", "up_proj"]:
                    if key in fqn:
                        return True

        # Layers 4-31
        for i in range(4, 32):
            if fqn.startswith(f"model.layers.{i}."):
                for key in ["q_proj", "v_proj", "gate_proj", "up_proj", "down_proj"]:
                    if key in fqn:
                        return True

        return False

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3165

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 2 Active SEVs

There are 2 currently active SEVs. If your PR is affected, please view them below:

• Macos CI runners unavailable
• [ROCm][CI] Machines under the label linux.rocm.gpu.2 are undergoing maintenance.

✅ No Failures

As of commit 9b1df40 with merge base c96f2dd ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[lisjin]

Looks great! I just had a few nits that might not be important

[lisjin]

Standardizing the indices this way is a great idea

[lisjin]

Should we add a type hint here?

[lisjin]

Is this for debugging?

[metascroy]

After each parameter, it prints out config (bitwidth,groupsize or NONE)