Config serde

[drisspg]

Stacked PRs:

• ->Config serde #1806
• Initial stab at string based config parser #1774

---

Summary: Convert base configs to use Pydantic

Related PRs

• Hugging Face: Add option for ao base configs huggingface/transformers#36526

• VLLM:
  Torchao vllm-project/vllm#14231

Overview

This PR adds serialization and deserialization capabilities to most* of our Config classes by migrating from dataclasses to Pydantic models. The implementation enables JSON serialization/deserialization and dictionary conversion for all quantization configuration classes.

Note that we have two options add this to our depedencies ( we weirdly dont have any specified currently) or I can probably write this in a way to only be enabled if we have the runtime dep>

Testing

Added a new test file test/quantization/test_config_serialization.py with tests for alot of configuration classes

Impact

W/ this we can serialize configs. However as implemented it also enforces that all params are kwarg only which may BC BREAKING, you can hack around this w/ pydantic but not sure how much of a deal this is yet.

Why do we need this?

I am currently auditing our end to end flow, and while yes the core serialization premise of tensor subclass to both define our data representation and modulate runtime behavior exists - that is not how the rest of the ecosystem is operating.

Currenlty our hugging face integration only supports 4 options:
https://huggingface.co/docs/transformers/main/en/main_classes/quantization#transformers.TorchAoConfig.quant_type

I for 1 want to have this support the full breadth of options and as well I want this to be configurable per layer. I think this gets us closer since we can specify the quant structure in a config.json on the serialized weight.

As well this config.json appears to be reused by VLLm and I am still working on the E2E flow as prototype by Jerry here:
vllm-project/vllm#13588

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 4 New Failures, 5 Cancelled Jobs, 1 Unrelated Failure

As of commit 7f411f7 with merge base ffb4350 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test-nightly (CPU Nightly, linux.4xlarge, --pre torch==2.7.0.dev20250122 --index-url https://down... / linux-job (gh)
  RuntimeError: Command docker exec -t 81cea65a3ce1279d86e49c277b6b0e85d4cd84555f98ff6586f7b867f70885ea /exec failed with exit code 139
• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch==2.7.0.dev20250122 --index-... / linux-job (gh)
  RuntimeError: Command docker exec -t f04a3e97046c910c51e45cdeefb6f189c2f4b08c8d57778d31410051c0df3a70 /exec failed with exit code 139
• Run TorchAO Experimental Tests / test-cpu-ops (macos-14) (gh)
  Process completed with exit code 2.
• Run TorchAO Experimental Tests / test-mps-ops (macos-m1-stable) (gh)
  ModuleNotFoundError: No module named 'importlib_metadata'

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• Run Regression Tests / test (CPU 2.4, linux.4xlarge, torch==2.4.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  ##[error]The operation was canceled.
• Run Regression Tests / test (CPU 2.5.1, linux.4xlarge, torch==2.5.1 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  ##[error]The operation was canceled.
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  ##[error]The operation was canceled.
• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  ##[error]The operation was canceled.
• Run Regression Tests / test (CUDA 2.5.1, linux.g5.12xlarge.nvidia.gpu, torch==2.5.1 --index-url https://download.pytorch... / linux-job (gh)
  ##[error]The operation was canceled.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Run Regression Tests / test (CPU 2.3, linux.4xlarge, torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

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