Allow per-group quantizers in QuantOptimizer, fix state_dict

[lisjin]

Changes to make torchao's version of PARQ compatible with current QAT pipeline:

1. Bug fix: When using FSDP 2, the QuantOptimizer state cannot be loaded correctly from a checkpoint because it has an extra "qat_state" key. Fix this by removing qat_state entirely and have the user set optimizer.num_steps manually after loading checkpoint.
2. New feature: Allow param_groups to be quantized using different bit-widths and/or quantization methods. If the user passes a "quant_cls" and optional "quant_kwargs" keys in a param_group, QuantOptimizer will override the default self.quantizer object with one initialized from these values.
3. New feature: Make PARQ compatible with new QAT API by relaxing the constraints in QATConfig. Before, it required that base_config or weight_config must be provided during the prepare step. Now, it accepts when only activation_config is provided. Now we can migrate from deprecated*IntXQuantizationAwareTrainingConfig classes.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit bed54ac with merge base 46ba24c ():
💚 Looks good so far! There are no failures yet. 💚

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

[andrewor14]

QAT changes looks good to me. Thanks!

[andrewor14]

Please also update the PR description to reflect the latest design.