From ee55b3a0e376a287ad8e2d4c1ea5994f46f83abb Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Mon, 26 Oct 2026 17:23:31 +0100 Subject: [PATCH 1/9] build: add quantization extra with nvidia-modelopt cosmos_framework.quantization imports ModelOpt, but no extra or dependency group declared it. Add a `quantization` extra pinning nvidia-modelopt==0.47.0 and regenerate uv.lock; it adds nvidia-modelopt and its dependency pulp, and leaves every existing locked package unchanged. Signed-off-by: Dominik Farhan --- pyproject.toml | 1 + uv.lock | 45 ++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 386421db..c4664299 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ dependencies = [ [project.optional-dependencies] automl = ["nvidia-automl-core==0.1.0"] automl-prompts = ["nvidia-automl-core==0.1.0", "gepa==0.1.2"] +quantization = ["nvidia-modelopt==0.47.0"] reasoner-eval = [ "attrs", "evaluate", diff --git a/uv.lock b/uv.lock index 733fd19c..3ee93e8d 100644 --- a/uv.lock +++ b/uv.lock @@ -2388,6 +2388,9 @@ guardrail = [ { name = "retinaface-py" }, { name = "sentencepiece" }, ] +quantization = [ + { name = "nvidia-modelopt" }, +] reasoner-eval = [ { name = "attrs" }, { name = "evaluate" }, @@ -2747,6 +2750,7 @@ requires-dist = [ { name = "nvidia-cudnn-frontend" }, { name = "nvidia-dali-cuda120", marker = "extra == 'train'" }, { name = "nvidia-ml-py" }, + { name = "nvidia-modelopt", marker = "extra == 'quantization'", specifier = "==0.47.0" }, { name = "obstore" }, { name = "omegaconf" }, { name = "open-clip-torch", marker = "extra == 'train'" }, @@ -2812,7 +2816,7 @@ requires-dist = [ { name = "xatlas", marker = "extra == 'train'" }, { name = "zarr", marker = "extra == 'train'" }, ] -provides-extras = ["automl", "automl-prompts", "reasoner-eval", "guardrail", "interactive", "serve", "train"] +provides-extras = ["automl", "automl-prompts", "quantization", "reasoner-eval", "guardrail", "interactive", "serve", "train"] [package.metadata.requires-dev] cu128 = [ @@ -9829,6 +9833,36 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fd/72/fb2af0d259a651affdce65fd6a495f0e07a685a0136baf585c5065204ee7/nvidia_ml_py-13.590.48-py3-none-any.whl", hash = "sha256:fd43d30ee9cd0b7940f5f9f9220b68d42722975e3992b6c21d14144c48760e43", size = 50680, upload-time = "2026-01-22T01:14:55.281Z" }, ] +[[package]] +name = "nvidia-modelopt" +version = "0.47.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ninja" }, + { name = "numpy" }, + { name = "nvidia-ml-py" }, + { name = "omegaconf" }, + { name = "packaging" }, + { name = "pulp" }, + { name = "pydantic" }, + { name = "pyyaml" }, + { name = "regex" }, + { name = "rich" }, + { name = "safetensors" }, + { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm')" }, + { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm')" }, + { name = "setuptools" }, + { name = "torch", version = "2.10.0", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu128') or (python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu128-train') or (python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu130') or (python_full_version < '3.13' and sys_platform == 'darwin' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-libero') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-libero') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-libero' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra != 'group-16-cosmos-framework-libero' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra != 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra != 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform == 'darwin' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra != 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero')" }, + { name = "torch", version = "2.10.0+cu128", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm')" }, + { name = "torch", version = "2.10.0+cu130", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version < '3.13' and sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128') or (python_full_version < '3.13' and sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu128-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version < '3.13' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (sys_platform != 'darwin' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-libero')" }, + { name = "torch", version = "2.13.0", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version < '3.13' and sys_platform == 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-libero') or (python_full_version >= '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-libero' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-train' and extra == 'group-16-cosmos-framework-libero') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra != 'group-16-cosmos-framework-libero' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra != 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra != 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero')" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl" }, marker = "(python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213') or (python_full_version >= '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version < '3.13' and extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-libero') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213' and extra == 'group-16-cosmos-framework-vllm') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-cu130-train') or (python_full_version < '3.13' and extra == 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-vllm') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213') or (sys_platform != 'darwin' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (sys_platform != 'darwin' and extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu128' and extra != 'group-16-cosmos-framework-cu128-train' and extra != 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu128-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-torch213-train') or (extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-cu130-train') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu128-train' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra == 'group-16-cosmos-framework-cu130' and extra == 'group-16-cosmos-framework-libero') or (extra == 'extra-16-cosmos-framework-train' and extra != 'group-16-cosmos-framework-cu130-torch213' and extra != 'group-16-cosmos-framework-cu130-torch213-train' and extra == 'group-16-cosmos-framework-libero')" }, + { name = "tqdm" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/46/f32f6183f6a26db8c6bf5c70f89f68ae2cc55c3fbb981016a772c93c4fc2/nvidia_modelopt-0.47.0-py3-none-any.whl", hash = "sha256:db13e62d7efe4a1512e8b8c3f479e28d64022306f94a194f13234e7dca1d3120", size = 2186201, upload-time = "2026-09-23T19:06:44.103Z" }, +] + [[package]] name = "nvidia-nccl-cu12" version = "2.27.5" @@ -11430,6 +11464,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" }, ] +[[package]] +name = "pulp" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/70/69be07a67621ad804d6cf347965eb4e0d7786a97330d99c31d735aaa6c5a/pulp-3.3.2.tar.gz", hash = "sha256:d0904700c207ac11e25e3b1213b70eae1d6fb25faa719d75f3f15054901258c0", size = 16305346, upload-time = "2026-05-25T09:41:26.207Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dd/6e/d674f1dde91c71e2ac19e5e5cf1ee6d5845e3aefecd9c39ed9c4b0c9a696/pulp-3.3.2-py3-none-any.whl", hash = "sha256:631b166f72086971a9597f7a0233ababa99bb8d50a01cd543f7758be5a9f86c0", size = 16391742, upload-time = "2026-05-25T09:41:22.2Z" }, +] + [[package]] name = "pure-eval" version = "0.2.3" From f364201c435cc5a907c7f8484f8b120e5713a0e5 Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Mon, 26 Oct 2026 17:35:37 +0100 Subject: [PATCH 2/9] feat(quant): quantize packed MoE experts to FP8 ModelOpt's generic fused-experts wrapper only intercepts F.linear, so on the framework's packed grouped-mm experts it quantizes nothing. - Add QuantPackedExperts, a ModelOpt wrapper for Qwen3VLMoeTextExpertsGroupedMm and MoEExperts with per-expert FP8 weight scales and one input scale per GEMM, and register_packed_experts() to install it before mtq.quantize. - Add assert_quantizers_calibrated() to fail on missing or invalid ranges. - Export packed experts in their source layout with [E,1,1] weight scales and scalar input scales, outside ModelOpt's linear-only exporter. - Skip the root Transformers modelopt_state.pth when the transformer has packed experts. Signed-off-by: Dominik Farhan --- cosmos_framework/quantization/export.py | 67 +++- cosmos_framework/quantization/moe.py | 185 ++++++++++ tests/quantization/modelopt_state_test.py | 61 ++++ tests/quantization/moe_experts_test.py | 403 ++++++++++++++++++++++ 4 files changed, 708 insertions(+), 8 deletions(-) create mode 100644 cosmos_framework/quantization/moe.py create mode 100644 tests/quantization/moe_experts_test.py diff --git a/cosmos_framework/quantization/export.py b/cosmos_framework/quantization/export.py index 7de906f8..98e3802f 100644 --- a/cosmos_framework/quantization/export.py +++ b/cosmos_framework/quantization/export.py @@ -98,6 +98,11 @@ def remap_framework_state_dict( for suffix in _SCALE_SUFFIXES: if not net_key.endswith(suffix): continue + # Packed MoE expert scales sit beside the packed parameter itself. + param_diffusers_key = net_to_diffusers.get(net_key[: -len(suffix)]) + if param_diffusers_key is not None: + diffusers_key = param_diffusers_key + suffix + break weight_key = net_key[: -len(suffix)] + ".weight" weight_diffusers_key = net_to_diffusers.get(weight_key) if weight_diffusers_key is not None: @@ -111,6 +116,22 @@ def remap_framework_state_dict( return remapped +def _remap_packed_experts( + experts: dict[str, torch.Tensor], bf16_state_dict: dict[str, torch.Tensor] +) -> dict[str, torch.Tensor]: + """:func:`remap_framework_state_dict` for packed experts, where every tensor must map. + + A dropped expert would leave its BF16 source tensor without scales under an FP8 config. + """ + remapped = remap_framework_state_dict(experts, bf16_state_dict) + if len(remapped) != len(experts): + unmapped = sorted( + key for key in experts if not remap_framework_state_dict({key: experts[key]}, bf16_state_dict) + ) + raise ValueError(f"{len(unmapped)} packed expert tensor(s) have no Diffusers key: {unmapped[:10]}") + return remapped + + # Top-level projections produced by reverse_remap use the cosmos3_vfm names; the # bf16 diffusers base provides the equivalently-named proj_in/proj_out, so we drop # ours and keep the base copies. @@ -200,6 +221,8 @@ def _remap_framework_modelopt_state( mapping used by checkpoint loading; no Diffusers or Transformers model is instantiated to reconstruct the graph. """ + from .moe import is_packed_expert_key + state = copy.deepcopy(state) modes = {name: mode_state for name, mode_state in state["modelopt_state_dict"]} if "quantize" not in modes or "real_quantize" not in modes: @@ -215,7 +238,12 @@ def _remap_framework_modelopt_state( def remap_module(name: str) -> str | None: return module_map.get(name.removeprefix("net.")) - active_modules = {key.removesuffix(".weight_scale") for key in exported_state_dict if key.endswith(".weight_scale")} + # Packed MoE experts are not linears and have no component graph state. + active_modules = { + key.removesuffix(".weight_scale") + for key in exported_state_dict + if key.endswith(".weight_scale") and not is_packed_expert_key(key) + } real_quantize = modes["real_quantize"] real_metadata = real_quantize["metadata"] @@ -390,6 +418,8 @@ def write_transformers_modelopt_state(self, component_state: dict, export_dir: P print(f"[export] wrote Transformers ModelOpt graph state for {len(quantizers)} quantizers") def overlay(self, bf16: dict, quantized: dict) -> dict: + from .moe import is_packed_expert_key + merged = dict(bf16) n_weight = n_scale = n_dropped = 0 quantized_modules: set[str] = set() @@ -401,7 +431,7 @@ def overlay(self, bf16: dict, quantized: dict) -> dict: merged[k] = v n_scale += 1 continue - if k.endswith(".weight") and v.dtype in _FP8_DTYPES: + if (k.endswith(".weight") or is_packed_expert_key(k)) and v.dtype in _FP8_DTYPES: if k not in bf16: raise KeyError( f"quantized weight {k!r} has no bf16 counterpart — layer-key " @@ -410,7 +440,9 @@ def overlay(self, bf16: dict, quantized: dict) -> dict: if bf16[k].shape != v.shape: raise ValueError(f"shape mismatch for {k}: bf16 {tuple(bf16[k].shape)} vs fp8 {tuple(v.shape)}") merged[k] = v - quantized_modules.add(k[: -len(".weight")]) + # Packed experts keep only their scales; the quantizer buffers below are per linear. + if not is_packed_expert_key(k): + quantized_modules.add(k[: -len(".weight")]) n_weight += 1 continue # Remaining bf16 tensors (norms, embed_tokens, renamed projections) are @@ -443,6 +475,7 @@ def export( modelopt_config: dict | None = None, ) -> None: from .checkpoint_io import load_sharded_safetensors, save_sharded_safetensors + from .moe import QuantPackedExperts, export_packed_experts net = mdl.net transformer_export_dir.mkdir(parents=True, exist_ok=True) @@ -451,6 +484,15 @@ def export( if n_collapsed > 0: print(f"[export] collapsed per-channel input amax to scalar on {n_collapsed} quantizers") + # Packed MoE experts are exported directly; ModelOpt's compress and unified + # exporter only handle linears, so their quantizers are disabled afterwards. + experts = export_packed_experts(net) + for module in net.modules(): + if isinstance(module, QuantPackedExperts): + for name in ("gate_up_proj", "down_proj"): + getattr(module, f"{name}_input_quantizer").disable() + getattr(module, f"{name}_weight_quantizer").disable() + # Record ModelOpt's real-quantized graph metadata before the unified # exporter unpacks those weights into plain FP8 checkpoint tensors. mtq.compress(mdl) @@ -466,6 +508,7 @@ def export( with hide_quantizers_from_state_dict(net): net_state_dict = {k: v.detach().contiguous().cpu() for k, v in net.state_dict().items()} quantized = remap_framework_state_dict(net_state_dict, bf16) + quantized.update(_remap_packed_experts(experts, bf16)) merged = self.overlay(bf16, quantized) n_shards = save_sharded_safetensors(merged, transformer_export_dir) print(f"[export] wrote {len(merged)} tensors across {n_shards} shard(s)") @@ -525,12 +568,13 @@ def _write_root_hf_quant_config(root_dir: Path, transformer_dir: Path) -> None: print(f"[assemble] wrote {root_dir / 'hf_quant_config.json'} for upstream-vLLM discovery") -def _write_root_weight_index(root_dir: Path) -> None: - """(Re)generate the root ``model.safetensors.index.json`` for whole-model loading.""" +def _write_root_weight_index(root_dir: Path) -> dict[str, str]: + """(Re)generate the root ``model.safetensors.index.json``; return its weight map.""" index = build_root_index(root_dir) dst = root_dir / "model.safetensors.index.json" index.write(dst) print(f"[assemble] wrote {dst.name} ({len(index.weight_map)} tensors, {index.metadata['total_size']} bytes)") + return index.weight_map def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer_dir: Path): @@ -540,7 +584,10 @@ def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer vision_encoder symlinked back); only ``transformer/`` is physically new (the quantized export). A fresh root weight index + a root ``hf_quant_config.json`` are written so both the diffusers and transformers loaders resolve the new weights. + The root Transformers ModelOpt state is skipped when the transformer has packed experts. """ + from .moe import is_packed_expert_key + input_dir, output_dir = Path(input_dir), Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) target_transformer = output_dir / "transformer" @@ -572,8 +619,12 @@ def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer rel_src = os.path.relpath(src, start=output_dir.resolve()) print(f"[assemble] linking {entry.name} -> {rel_src}") dst.symlink_to(rel_src) - _write_root_weight_index(output_dir) - component_state = torch.load(target_transformer / "modelopt_state.pth", map_location="cpu", weights_only=False) - Fp8DiffusersExporter().write_transformers_modelopt_state(component_state, output_dir, output_dir) + weight_map = _write_root_weight_index(output_dir) + has_packed_experts = any( + is_packed_expert_key(key) for key, file in weight_map.items() if file.startswith("transformer/") + ) + if not has_packed_experts: + component_state = torch.load(target_transformer / "modelopt_state.pth", map_location="cpu", weights_only=False) + Fp8DiffusersExporter().write_transformers_modelopt_state(component_state, output_dir, output_dir) _write_root_hf_quant_config(output_dir, target_transformer) print(f"[assemble] drop-in dir ready at {output_dir}") diff --git a/cosmos_framework/quantization/moe.py b/cosmos_framework/quantization/moe.py new file mode 100644 index 00000000..4d20054c --- /dev/null +++ b/cosmos_framework/quantization/moe.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""ModelOpt FP8 wrapper for the packed grouped-mm MoE experts. + +ModelOpt's generic ``_QuantFusedExperts`` claims any module with 3-D ``gate_up_proj`` / +``down_proj`` parameters but only intercepts ``F.linear``, so on the packed grouped-mm experts +(which run ``torch._grouped_mm``) it would quantize nothing. :class:`QuantPackedExperts` keeps +the experts packed and quantizes them in place: one per-expert weight quantizer per packed +parameter and one shared input quantizer per GEMM, matching what vLLM's FP8 MoE kernels consume. + +The quantizers act on tensors inside the expert computation, so its ``forward`` routes the tokens +and runs both grouped GEMMs itself. +""" + +from __future__ import annotations + +import torch +from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry, TensorQuantizer +from torch import nn + +from cosmos_framework.model.generator.moe.interface import MoEExperts +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import Qwen3VLMoeTextExpertsGroupedMm + +_PACKED_EXPERT_TYPES = (Qwen3VLMoeTextExpertsGroupedMm, MoEExperts) + + +class QuantPackedExperts(QuantModule): + """Quantized packed grouped-mm experts; the packed parameters stay untouched.""" + + def _setup(self): + self.gate_up_proj_input_quantizer = TensorQuantizer() + self.gate_up_proj_weight_quantizer = TensorQuantizer() + self.down_proj_input_quantizer = TensorQuantizer() + self.down_proj_weight_quantizer = TensorQuantizer() + + def forward( + self, + hidden_states: torch.Tensor, # [num_tokens,hidden_size] + topk_scores: torch.Tensor, # [num_tokens,top_k] + expert_indices: torch.Tensor, # [num_tokens,top_k] + num_tokens_per_expert: torch.Tensor, # [num_experts] + token_mask: torch.Tensor | None = None, # [num_tokens] + ) -> torch.Tensor: # [num_tokens,hidden_size] + # Group the real (token, slot) pairs by expert; ``num_tokens_per_expert`` counts exactly these. + top_k = expert_indices.shape[1] + slots = torch.arange(expert_indices.numel(), device=expert_indices.device) # [num_tokens*top_k] + if token_mask is not None: + slots = slots[token_mask.repeat_interleave(top_k)] # [num_routed] + if slots.numel() == 0: + return torch.zeros_like(hidden_states) # [num_tokens,hidden_size] + slots = slots[torch.argsort(expert_indices.reshape(-1)[slots], stable=True)] # [num_routed] + tokens = slots // top_k # [num_routed] + scores = topk_scores.reshape(-1)[slots].to(hidden_states.dtype) # [num_routed] + x = hidden_states[tokens] # [num_routed,hidden_size] + routed = self._quantized_experts(x, num_tokens_per_expert, scores) # [num_routed,hidden_size] + return torch.zeros_like(hidden_states).index_add_(0, tokens, routed) # [num_tokens,hidden_size] + + def _quantized_experts( + self, + x: torch.Tensor, # [num_routed,hidden_size] (grouped by expert) + num_tokens_per_expert: torch.Tensor, # [num_experts] + scores: torch.Tensor, # [num_routed] + ) -> torch.Tensor: # [num_routed,hidden_size] + # a2 observes ``act(gate) * up`` before the routing scores, because vLLM applies the scores + # after the down GEMM. + offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) # [num_experts] + x = self.gate_up_proj_input_quantizer(x) + h = torch._grouped_mm(x, self.gate_up_proj_weight_quantizer(self.gate_up_proj), offs=offsets) + g, u = torch.chunk(h, chunks=2, dim=-1) # 2x [num_routed,moe_intermediate_size] + a = self.down_proj_input_quantizer(self.act_fn(g) * u) # [num_routed,moe_intermediate_size] + return torch._grouped_mm( + a * scores.unsqueeze(-1), self.down_proj_weight_quantizer(self.down_proj), offs=offsets + ) # [num_routed,hidden_size] + + +def register_packed_experts() -> None: + """Register :class:`QuantPackedExperts` before ``mtq.quantize`` so ``_QuantFusedExperts`` skips it.""" + pending: dict[type[nn.Module], str] = {} + for expert_cls in _PACKED_EXPERT_TYPES: + registered = QuantModuleRegistry.get(expert_cls) + if registered is not None and issubclass(registered, QuantPackedExperts): + continue + if registered is not None: + # An earlier ``mtq.quantize`` that saw this type unregistered has globally registered + # ModelOpt's generic ``_QuantFusedExperts`` for it, which would quantize nothing here. + QuantModuleRegistry.unregister(expert_cls) + pending[expert_cls] = expert_cls.__name__ + if pending: + QuantModuleRegistry.register(pending)(QuantPackedExperts) + + +def has_packed_experts(model: nn.Module) -> bool: + return any(isinstance(module, _PACKED_EXPERT_TYPES) for module in model.modules()) + + +def assert_quantizers_calibrated(model: nn.Module) -> None: + """Fail export instead of writing a missing or zero scale for an enabled quantizer.""" + problems = [] + for name, module in model.named_modules(): + if isinstance(module, TensorQuantizer) and not module.is_enabled: + continue + if isinstance(module, TensorQuantizer): + amax = getattr(module, "amax", None) + if amax is None: + problems.append(f"{name}: no amax") + elif not torch.isfinite(amax).all(): + problems.append(f"{name}: non-finite amax") + elif name.endswith("input_quantizer") and not (amax > 0).all(): + problems.append(f"{name}: non-positive amax") + elif ( + isinstance(getattr(module, "gate_up_proj", None), nn.Parameter) + and module.gate_up_proj.ndim == 3 + and not isinstance(module, QuantPackedExperts) + ): + problems.append(f"{name}: packed experts {type(module).__name__} are not quantized") + if problems: + shown = "; ".join(problems[:10]) + raise ValueError(f"{len(problems)} quantizer problem(s), first {min(len(problems), 10)}: {shown}") + + +_FP8_MAX = 448.0 # Mirrors ``export.FP8_MAX``; export imports this module, so it is not imported back. +_PACKED_EXPERT_QUANTIZERS = ( + "gate_up_proj_input_quantizer", + "gate_up_proj_weight_quantizer", + "down_proj_input_quantizer", + "down_proj_weight_quantizer", +) +_PACKED_EXPERT_KEY_SUFFIXES = tuple( + f"experts.{name}{suffix}" + for name in ("gate_up_proj", "down_proj") + for suffix in ("", ".weight_scale", ".input_scale") +) + + +def is_packed_expert_key(key: str) -> bool: + """True for a packed expert parameter or its scale, e.g. ``...experts.gate_up_proj.weight_scale``.""" + return key.endswith(_PACKED_EXPERT_KEY_SUFFIXES) + + +def _check_per_expert_fp8(name: str, module: QuantPackedExperts) -> None: + """Refuse any layout but the one written below: FP8 E4M3, per-expert weight amax, scalar input amax.""" + for quantizer_name in _PACKED_EXPERT_QUANTIZERS: + quantizer = getattr(module, quantizer_name) + amax_shape = tuple(quantizer.amax.shape) + if quantizer.num_bits != (4, 3) or quantizer.block_sizes: + found = f"num_bits={quantizer.num_bits}, block_sizes={quantizer.block_sizes}" + elif quantizer_name.endswith("weight_quantizer"): + num_experts = getattr(module, quantizer_name.removesuffix("_weight_quantizer")).shape[0] + if amax_shape == (num_experts, 1, 1): + continue + found = f"amax shape {amax_shape}, expected ({num_experts}, 1, 1)" + elif quantizer.amax.numel() == 1: + continue + else: + found = f"amax shape {amax_shape}, expected a single element" + raise ValueError(f"{name}.{quantizer_name}: {found}; packed-expert export supports per-expert FP8 only") + + +def export_packed_experts(net: nn.Module) -> dict[str, torch.Tensor]: + """FP8 packed experts in the source layout, with per-expert weight scales and scalar input scales.""" + if not has_packed_experts(net): + return {} # Dense models export exactly as before, without the calibration check. + assert_quantizers_calibrated(net) + exported = {} + for name, module in net.named_modules(): + if not isinstance(module, QuantPackedExperts): + continue + enabled = [getattr(module, quantizer).is_enabled for quantizer in _PACKED_EXPERT_QUANTIZERS] + if not any(enabled): + continue # Skipped by the coverage filter; the experts stay BF16 from the source. + if not all(enabled): + raise ValueError(f"{name or 'experts'}: packed-expert quantizers are only partly enabled: {enabled}") + _check_per_expert_fp8(name or "experts", module) + for param_name in ("gate_up_proj", "down_proj"): + # BF16 models keep BF16 amax; the scales are FP32 on disk. + weight_amax = getattr(module, f"{param_name}_weight_quantizer").amax.detach().float() + weight_scale = torch.where(weight_amax == 0, 1.0, weight_amax / _FP8_MAX) # [num_experts,1,1] + weight = getattr(module, param_name).detach().float() + fp8 = (weight / weight_scale).clamp(-_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn) + input_amax = getattr(module, f"{param_name}_input_quantizer").amax.detach().float() + key = f"{name}.{param_name}" if name else param_name + exported[key] = fp8.contiguous().cpu() + exported[f"{key}.weight_scale"] = weight_scale.contiguous().cpu() + exported[f"{key}.input_scale"] = (input_amax / _FP8_MAX).reshape(()).contiguous().cpu() + return exported diff --git a/tests/quantization/modelopt_state_test.py b/tests/quantization/modelopt_state_test.py index 718b8830..0905318d 100644 --- a/tests/quantization/modelopt_state_test.py +++ b/tests/quantization/modelopt_state_test.py @@ -365,6 +365,67 @@ def test_assembly_keeps_distinct_states_and_source_config(export_modules, tmp_pa } +def test_assemble_skips_root_modelopt_state_for_packed_experts(export_modules, tmp_path, monkeypatch): + _, exporter = export_modules + source, output, staging = (tmp_path / name for name in ("source", "output", "staging")) + for directory in (source, output, staging): + directory.mkdir() + (staging / "config.json").write_text(json.dumps({"quantization_config": {"quant_method": "modelopt"}})) + save_file({"layers.0.mlp.experts.gate_up_proj": torch.zeros(2, 4, 4)}, str(staging / "model.safetensors")) + + def fail(*args, **kwargs): + raise AssertionError("root ModelOpt state must not be built") + + monkeypatch.setattr(exporter, "build_reasoner_modelopt_state", fail) + exporter.assemble_output_dir(source, output, staging) + assert not (output / "modelopt_state.pth").exists() + assert (output / "hf_quant_config.json").is_file() + + +def test_component_state_ignores_packed_experts(export_modules): + from cosmos_framework.inference.model import _diffusers_to_net_key + + _, exporter = export_modules + linear, experts = "layers.0.self_attn.to_q", "layers.0.mlp.experts.gate_up_proj" + net_linear = _diffusers_to_net_key(f"{linear}.weight").removesuffix(".weight") + model = torch.nn.Module() + model.add_module("net", torch.nn.Module()) + parent = model.net + *parents, leaf = net_linear.split(".") + for part in parents: + parent.add_module(part, torch.nn.Module()) + parent = parent._modules[part] + parent.add_module(leaf, torch.nn.Linear(16, 16, bias=False).bfloat16()) + model = apply_mode(model, [("quantize", copy.deepcopy(mtq.FP8_DEFAULT_CFG))]) + module = model.get_submodule(f"net.{net_linear}") + module.input_quantizer.amax = torch.tensor(3.5) + module.weight_quantizer.amax = module.weight.detach().float().abs().max() + mtq.compress(model) + bf16 = {f"{linear}.weight": torch.zeros(16, 16, dtype=torch.bfloat16), experts: torch.zeros(4, 16, 16)} + exported = { + f"{linear}.weight": torch.zeros(16, 16, dtype=torch.float8_e4m3fn), + experts: torch.zeros(4, 16, 16, dtype=torch.float8_e4m3fn), + } + for key in (linear, experts): + exported[f"{key}.weight_scale"] = torch.ones(()) + exported[f"{key}.input_scale"] = torch.ones(()) + + state = exporter._remap_framework_modelopt_state( + mto.modelopt_state(model), + bf16_state_dict=bf16, + exported_state_dict=exported, + modelopt_config=copy.deepcopy(mtq.FP8_DEFAULT_CFG), + ) + + modes = dict(state["modelopt_state_dict"]) + assert set(modes["real_quantize"]["metadata"]["q_tensor_state"]) == {linear} + names = [ + *modes["quantize"]["metadata"]["quantizer_state"], + *modes["real_quantize"]["metadata"]["real_quantizer_state"], + ] + assert names and not any("experts" in name for name in names) + + @pytest.mark.skipif(not os.environ.get("COSMOS3_REASONER_STATE_SOT"), reason="external SOT path not supplied") def test_matches_external_sot_metadata(export_modules, tmp_path): helper, _ = export_modules diff --git a/tests/quantization/moe_experts_test.py b/tests/quantization/moe_experts_test.py new file mode 100644 index 00000000..39328ba3 --- /dev/null +++ b/tests/quantization/moe_experts_test.py @@ -0,0 +1,403 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: OpenMDW-1.1 +"""ModelOpt FP8 calibration and export of the packed grouped-mm MoE experts. + +The CPU tests replace ``torch._grouped_mm`` and Triton's ``_generate_permute_indices`` with +PyTorch references. +""" + +import copy + +import modelopt.torch.quantization as mtq +import pytest +import torch +import torch.nn.functional as F +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins.huggingface import _QuantFusedExperts +from torch import nn + +import cosmos_framework.model.generator.moe.native as native_grouped_mm +from cosmos_framework.inference.model import _diffusers_to_net_key +from cosmos_framework.model.generator.moe.interface import MoEExperts +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.configuration_qwen3_vl_moe import Qwen3VLMoeTextConfig +from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import ( + Qwen3VLMoeTextExpertsGroupedMm, + Qwen3VLMoeTextExpertsNaive, +) +from cosmos_framework.quantization.export import Fp8DiffusersExporter, _remap_packed_experts, remap_framework_state_dict +from cosmos_framework.quantization.moe import ( + QuantPackedExperts, + assert_quantizers_calibrated, + export_packed_experts, + register_packed_experts, +) + +QUANTIZERS = ( + "gate_up_proj_input_quantizer", + "gate_up_proj_weight_quantizer", + "down_proj_input_quantizer", + "down_proj_weight_quantizer", +) +PER_EXPERT_WEIGHTS = [ + {"quantizer_name": "*gate_up_proj_weight_quantizer", "cfg": {"num_bits": (4, 3), "axis": 0}}, + {"quantizer_name": "*down_proj_weight_quantizer", "cfg": {"num_bits": (4, 3), "axis": 0}}, +] + + +def _reference_grouped_mm(x: torch.Tensor, w: torch.Tensor, offs: torch.Tensor) -> torch.Tensor: + # torch._grouped_mm leaves the rows past offs[-1] uninitialized; a large value stands in for that garbage. + out = x.new_full((x.shape[0], w.shape[-1]), 1e4) + start = 0 + for group, end in enumerate(offs.tolist()): + out[start:end] = x[start:end] @ w[group] + start = end + return out + + +def _reference_permute_indices( + num_tokens_per_expert: torch.Tensor, num_experts: int, max_len: int, alignment: int +) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch version of the Triton ``_generate_permute_indices``: -1 marks a padding slot.""" + sizes = num_tokens_per_expert.clamp_min(alignment).to(torch.int32) + sizes = (sizes + alignment - 1) // alignment * alignment # [num_experts] + permuted_indices = torch.full((max_len,), -1, dtype=torch.int32) + start = write = 0 + for count, size in zip(num_tokens_per_expert.tolist(), sizes.tolist(), strict=True): + permuted_indices[write : write + count] = torch.arange(start, start + count, dtype=torch.int32) + start, write = start + count, write + size + return permuted_indices, sizes + + +@pytest.fixture +def reference_routing(monkeypatch): + monkeypatch.setattr(torch, "_grouped_mm", _reference_grouped_mm) + monkeypatch.setattr(native_grouped_mm, "_generate_permute_indices", _reference_permute_indices) + + +def _config() -> Qwen3VLMoeTextConfig: + return Qwen3VLMoeTextConfig( + num_experts=4, hidden_size=16, moe_intermediate_size=8, num_experts_per_tok=2, hidden_act="silu" + ) + + +class Model(nn.Module): + def __init__(self, experts: nn.Module | None = None): + super().__init__() + self.experts = Qwen3VLMoeTextExpertsGroupedMm(_config()) if experts is None else experts + generator = torch.Generator().manual_seed(0) + for param in self.experts.parameters(): + param.data.normal_(generator=generator) + + +# The packed expert types ``register_packed_experts`` claims, built for a given top_k. +EXPERTS = { + "qwen_grouped_mm": lambda top_k: Qwen3VLMoeTextExpertsGroupedMm(_config(), top_k=top_k), + "moe_experts": lambda top_k: MoEExperts(hidden_size=16, intermediate_size=8, num_experts=4, top_k=top_k), +} + + +def _forward_inputs( + seed: int = 3, dtype: torch.dtype = torch.float32, top_k: int = 2, routing: str = "padding" +) -> tuple: + """``forward`` arguments for 6 tokens, top-``top_k`` of 4 experts, padded per ``routing``. + + Only the last token picks expert 2, so padding it leaves expert 2 without tokens. + """ + generator = torch.Generator().manual_seed(seed) + hidden_states = torch.randn(6, 16, generator=generator, dtype=dtype) + topk_scores = torch.rand(6, top_k, generator=generator, dtype=dtype) + expert_indices = torch.tensor([[0, 1, 3], [1, 3, 0], [0, 3, 1], [3, 0, 1], [1, 0, 3], [2, 1, 0]]) + expert_indices = expert_indices[:, :top_k].contiguous() + token_mask = torch.tensor([True, True, True, True, True, False]) + if routing == "no_padding": + token_mask = None + elif routing == "all_padding": + token_mask = torch.zeros(6, dtype=torch.bool) + elif routing == "empty": + hidden_states, topk_scores = hidden_states[:0], topk_scores[:0] + expert_indices, token_mask = expert_indices[:0], token_mask[:0] + else: + assert routing == "padding", routing + real = expert_indices if token_mask is None else expert_indices[token_mask] + num_tokens_per_expert = torch.bincount(real.reshape(-1), minlength=4) + return hidden_states, topk_scores, expert_indices, num_tokens_per_expert, token_mask + + +def _quantize(model: nn.Module, calibration: list[tuple] | None = None, *extra_cfg: dict) -> nn.Module: + register_packed_experts() + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["quant_cfg"] = [*config["quant_cfg"], *PER_EXPERT_WEIGHTS, *extra_cfg] + + def forward_loop(model: nn.Module) -> None: + for inputs in calibration or [_forward_inputs()]: + model.experts(*inputs) + if hasattr(model, "linear"): + model.linear(inputs[0]) + + return mtq.quantize(model, config, forward_loop) + + +@pytest.fixture +def unregistered_experts(): + if QuantModuleRegistry.get(Qwen3VLMoeTextExpertsGroupedMm) is not None: + QuantModuleRegistry.unregister(Qwen3VLMoeTextExpertsGroupedMm) + yield + register_packed_experts() + + +@pytest.mark.usefixtures("reference_routing", "unregistered_experts") +def test_registration_replaces_generic_fused_experts_wrapper(): + generic = Model() + mtq.quantize(generic, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda model: None) + assert isinstance(generic.experts, _QuantFusedExperts) + + model = _quantize(Model()) + + assert isinstance(model.experts, QuantPackedExperts) + assert not isinstance(model.experts, _QuantFusedExperts) + registered = QuantModuleRegistry.get(Qwen3VLMoeTextExpertsGroupedMm) + assert issubclass(registered, QuantPackedExperts) and not issubclass(registered, _QuantFusedExperts) + + +@pytest.mark.usefixtures("reference_routing") +@pytest.mark.parametrize("routing", ["padding", "no_padding", "all_padding", "empty"]) +@pytest.mark.parametrize("top_k", [2, 3]) +@pytest.mark.parametrize("experts", sorted(EXPERTS)) +def test_wrapper_matches_native_with_quantizers_disabled(experts, top_k, routing): + # QuantPackedExperts has its own forward; this catches it drifting from the native one. + model = Model(EXPERTS[experts](top_k)) + native = copy.deepcopy(model.experts) + inputs = _forward_inputs(seed=4, top_k=top_k, routing=routing) + expected = native(*inputs) + model = _quantize(model, [_forward_inputs(top_k=top_k)]) + assert isinstance(model.experts, QuantPackedExperts) + + for name in QUANTIZERS: + getattr(model.experts, name).disable() + # On CPU ``MoEExperts`` runs the dense ``native.naive``, so the two match up to rounding only. + torch.testing.assert_close(model.experts(*inputs), expected, rtol=1e-5, atol=1e-3) + + +@pytest.mark.usefixtures("reference_routing") +def test_each_quantizer_changes_the_wrapper_output(): + model = _quantize(Model()) + inputs = _forward_inputs(seed=4) + for name in QUANTIZERS: + getattr(model.experts, name).disable() + unquantized = model.experts(*inputs) + + for name in QUANTIZERS: + quantizer = getattr(model.experts, name) + quantizer.enable() + assert not torch.equal(model.experts(*inputs), unquantized), name + quantizer.disable() + + +@pytest.mark.usefixtures("reference_routing") +def test_forward_calibrates_per_expert_weights_and_unpadded_activations(): + calibration = [] + for seed in (1, 2): + hidden_states, _, expert_indices, counts, token_mask = _forward_inputs(seed) + hidden_states[~token_mask] = 1e3 # A masked token must not reach a1. + calibration.append((hidden_states, torch.full((6, 2), 0.25), expert_indices, counts, token_mask)) + + model = _quantize(Model(), calibration) + + assert_quantizers_calibrated(model) + experts = model.experts + for name in ("gate_up_proj", "down_proj"): + amax = getattr(experts, f"{name}_weight_quantizer").amax + assert amax.shape == (4, 1, 1) + assert torch.equal(amax, getattr(experts, name).abs().amax(dim=(1, 2), keepdim=True)) + real = [(x[mask], indices[mask]) for x, _, indices, _, mask in calibration] + assert torch.equal(experts.gate_up_proj_input_quantizer.amax, torch.stack([x.abs().max() for x, _ in real]).max()) + # a2 sees act(gate) * up on the real (token, expert) rows only, before the routing scores. + gate_up = torch.cat([torch.einsum("th,tkhf->tkf", x, experts.gate_up_proj[indices]) for x, indices in real]) + gate, up = gate_up.chunk(2, dim=-1) + expected = (F.silu(gate) * up).abs().max() + amax = experts.down_proj_input_quantizer.amax + torch.testing.assert_close(amax, expected) + assert not torch.allclose(amax, 0.25 * expected) + + +def test_assert_quantizers_calibrated_rejects_uncalibrated_and_unwrapped(): + model = nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4)) + mtq.quantize(model, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m[0](torch.randn(2, 4))) + with pytest.raises(ValueError, match=r"1\.input_quantizer"): + assert_quantizers_calibrated(model) + + model = nn.Sequential(nn.Linear(4, 4)) + model.experts = Qwen3VLMoeTextExpertsNaive(_config()) + # ModelOpt calibrates the generic fused-experts weights, so they must not be uninitialized memory. + generator = torch.Generator().manual_seed(0) + for param in model.experts.parameters(): + param.data.normal_(generator=generator) + mtq.quantize(model, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m[0](torch.randn(2, 4))) + with pytest.raises(ValueError, match="experts"): + assert_quantizers_calibrated(model) + + +@pytest.mark.usefixtures("reference_routing") +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_export_packed_experts_layout(dtype): + model = Model().to(dtype) + model.experts.gate_up_proj.data[1] = 0 + model.experts.down_proj.data[1] = 0 + model = _quantize(model, [_forward_inputs(dtype=dtype)]) + + exported = export_packed_experts(model) + + shapes = {"gate_up_proj": (4, 16, 16), "down_proj": (4, 8, 16)} + assert set(exported) == { + f"experts.{name}{suffix}" for name in shapes for suffix in ("", ".weight_scale", ".input_scale") + } + for name, shape in shapes.items(): + weight = getattr(model.experts, name).float() + fp8 = exported[f"experts.{name}"] + weight_scale = exported[f"experts.{name}.weight_scale"] + input_scale = exported[f"experts.{name}.input_scale"] + assert fp8.dtype == torch.float8_e4m3fn and fp8.shape == shape + assert weight_scale.dtype == torch.float32 and weight_scale.shape == (4, 1, 1) + assert input_scale.dtype == torch.float32 and input_scale.shape == () + weight_amax = weight.abs().amax(dim=(1, 2), keepdim=True) + torch.testing.assert_close(weight_scale, torch.where(weight_amax == 0, 1.0, weight_amax / 448)) + assert weight_scale[1].item() == 1.0 # The all-zero expert gets a unit scale, not 0. + assert torch.isfinite(fp8[1].float()).all() and (fp8[1].float() == 0).all() + amax = getattr(model.experts, f"{name}_input_quantizer").amax.float() + torch.testing.assert_close(input_scale, amax / 448) + # 32 is the FP8 E4M3 step in the top binade [256, 448]. + assert ((fp8.float() * weight_scale - weight).abs() <= 32 * weight_scale).all() + assert all( + tensor.is_contiguous() and tensor.device.type == "cpu" for tensor in (fp8, weight_scale, input_scale) + ) + + +@pytest.mark.usefixtures("reference_routing") +def test_export_skips_dense_and_disabled_experts_and_rejects_mixed(): + dense = nn.Sequential(nn.Linear(4, 4), nn.Linear(4, 4)) + mtq.quantize(dense, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m[0](torch.randn(2, 4))) + assert export_packed_experts(dense) == {} # Uncalibrated, but dense models skip the check. + + # A coverage skip disables the experts before calibration, so they never get an amax. The + # linear keeps a weight to quantize: ModelOpt (>= 0.47) refuses a config with none enabled. + skipped = Model() + skipped.linear = nn.Linear(16, 16) + skipped = _quantize(skipped, None, {"quantizer_name": "*experts*", "enable": False}) + assert isinstance(skipped.experts, QuantPackedExperts) + assert export_packed_experts(skipped) == {} + + mixed = _quantize(Model()) + mixed.experts.down_proj_input_quantizer.disable() + with pytest.raises(ValueError, match="experts"): + export_packed_experts(mixed) + + +@pytest.mark.usefixtures("reference_routing") +def test_export_rejects_non_per_expert_fp8(): + register_packed_experts() + scalar_weights = mtq.quantize(Model(), copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m.experts(*_forward_inputs())) + assert scalar_weights.experts.gate_up_proj_weight_quantizer.amax.numel() == 1 + with pytest.raises(ValueError, match="per-expert FP8") as error: + export_packed_experts(scalar_weights) + assert "experts" in str(error.value) and "gate_up_proj_weight_quantizer" in str(error.value) + + int8 = _quantize(Model(), None, {"quantizer_name": "*experts*", "cfg": {"num_bits": 8}}) + assert int8.experts.down_proj_input_quantizer.num_bits == 8 + with pytest.raises(ValueError, match="per-expert FP8") as error: + export_packed_experts(int8) + assert "num_bits=8" in str(error.value) + + +def test_remap_packed_expert_scales(): + prefix = "layers.0.mlp.experts" + bf16 = {f"{prefix}.gate_up_proj": torch.zeros(4, 16, 16, dtype=torch.bfloat16)} + net_prefix = _diffusers_to_net_key(f"{prefix}.gate_up_proj", "transformer/weights.safetensors") + exported = {net_prefix + suffix: torch.ones(1) for suffix in ("", ".weight_scale", ".input_scale")} + expected = {f"{prefix}.gate_up_proj{suffix}" for suffix in ("", ".weight_scale", ".input_scale")} + + assert set(remap_framework_state_dict(exported, bf16)) == expected + assert set(_remap_packed_experts(exported, bf16)) == expected + exported["net.unknown.experts.down_proj.weight_scale"] = torch.ones(1) + with pytest.raises(ValueError, match=r"net\.unknown\.experts\.down_proj\.weight_scale"): + _remap_packed_experts(exported, bf16) + + +@pytest.mark.usefixtures("reference_routing") +def test_overlay_accepts_packed_fp8_experts(): + model = _quantize(Model()) + bf16 = { + f"layers.0.mlp.{name}": value.bfloat16() for name, value in model.experts.named_parameters(prefix="experts") + } + quantized = {f"layers.0.mlp.{name}": value for name, value in export_packed_experts(model).items()} + + merged = Fp8DiffusersExporter().overlay(bf16, quantized) + + for key, value in quantized.items(): + assert torch.equal(merged[key].float(), value.float()) and merged[key].dtype == value.dtype + assert merged["layers.0.mlp.experts.gate_up_proj"].dtype == torch.float8_e4m3fn + assert not any("_quantizer." in key for key in merged) + quantized["layers.0.mlp.experts.down_proj"] = quantized["layers.0.mlp.experts.down_proj"][:, :4] + with pytest.raises(ValueError, match="shape mismatch"): + Fp8DiffusersExporter().overlay(bf16, quantized) + + +@pytest.mark.gpus(1) +def test_wrapper_matches_native_block_on_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for torch._grouped_mm and the Triton routing kernels") + from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock + + torch.manual_seed(0) + block = Qwen3VLMoeTextSparseMoeBlock(_config()) + for param in block.experts.parameters(): + param.data.normal_() + block = block.to(device="cuda", dtype=torch.bfloat16).eval() + native = copy.deepcopy(block) + hidden_states = torch.randn(32, 16, device="cuda", dtype=torch.bfloat16) + token_mask = torch.ones(32, dtype=torch.bool, device="cuda") + token_mask[-5:] = False + + register_packed_experts() + mtq.quantize(block, copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda model: model(hidden_states, token_mask=token_mask)) + assert isinstance(block.experts, QuantPackedExperts) + mtq.disable_quantizer(block, "*") + + with torch.no_grad(): + expected, _ = native(hidden_states, token_mask=token_mask) + actual, _ = block(hidden_states, token_mask=token_mask) + assert torch.equal(actual, expected) + + +@pytest.mark.gpus(1) +def test_a2_amax_ignores_padding_on_cuda(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required for torch._grouped_mm and the Triton routing kernels") + from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock + + torch.manual_seed(0) + block = Qwen3VLMoeTextSparseMoeBlock(_config()) + for param in block.experts.parameters(): + param.data.normal_() + block = block.to(device="cuda", dtype=torch.bfloat16).eval() + hidden_states = torch.randn(32, 16, device="cuda", dtype=torch.bfloat16) + token_mask = torch.ones(32, dtype=torch.bool, device="cuda") + token_mask[-5:] = False + routed = {} + + def record_routing(module, args, kwargs): + routed["expert_indices"] = kwargs["expert_indices"] + + block.experts.register_forward_pre_hook(record_routing, with_kwargs=True) + register_packed_experts() + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["quant_cfg"] = [*config["quant_cfg"], *PER_EXPERT_WEIGHTS] + mtq.quantize(block, config, lambda model: model(hidden_states, token_mask=token_mask)) + + # a2 on the real rows only, one row per (token, selected expert). + real = token_mask.nonzero().squeeze(-1) # [num_real] + experts = routed["expert_indices"][real] # [num_real,top_k] + gate_up = torch.einsum("th,tkhf->tkf", hidden_states[real], block.experts.gate_up_proj[experts]) + gate, up = gate_up.chunk(2, dim=-1) + expected = (F.silu(gate) * up).abs().max().float() + amax = block.experts.down_proj_input_quantizer.amax.float() + torch.testing.assert_close(amax, expected, rtol=2e-2, atol=0) From 579393f2f5bf5d80ea2605eafb3aeae29ac89b5f Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Mon, 26 Oct 2026 17:43:43 +0100 Subject: [PATCH 3/9] refactor(quant): keep the packed-expert key check out of moe.py is_packed_expert_key is a plain key-name check, but export.py imported it from quantization/moe.py, which imports the MoE model code and Triton. That broke the CPU export tests, which load export.py without the package's __init__. Move it next to the other key-name rules in export.py, and import FP8_MAX in moe.py instead of duplicating it. Signed-off-by: Dominik Farhan --- cosmos_framework/quantization/export.py | 17 +++++++++++------ cosmos_framework/quantization/moe.py | 18 ++++-------------- 2 files changed, 15 insertions(+), 20 deletions(-) diff --git a/cosmos_framework/quantization/export.py b/cosmos_framework/quantization/export.py index 98e3802f..f2b9db7f 100644 --- a/cosmos_framework/quantization/export.py +++ b/cosmos_framework/quantization/export.py @@ -139,6 +139,17 @@ def _remap_packed_experts( _FP8_DTYPES = {torch.float8_e4m3fn, getattr(torch, "float8_e5m2", torch.float8_e4m3fn)} FP8_MAX = 448.0 _SCALE_SUFFIXES = (".weight_scale", ".weight_scale_2", ".input_scale", ".input_scale_2", ".alpha", ".pre_quant_scale") +_PACKED_EXPERT_KEY_SUFFIXES = tuple( + f"experts.{name}{suffix}" + for name in ("gate_up_proj", "down_proj") + for suffix in ("", ".weight_scale", ".input_scale") +) + + +def is_packed_expert_key(key: str) -> bool: + """True for a packed expert parameter or its scale, e.g. ``...experts.gate_up_proj.weight_scale``.""" + return key.endswith(_PACKED_EXPERT_KEY_SUFFIXES) + # Linears the diffusers checkpoint keeps in bf16 (never quantized), emitted in the # ModelOpt ``ignore`` list so a reasoning-path loader (vllm Cosmos3ForConditionalGeneration) @@ -221,8 +232,6 @@ def _remap_framework_modelopt_state( mapping used by checkpoint loading; no Diffusers or Transformers model is instantiated to reconstruct the graph. """ - from .moe import is_packed_expert_key - state = copy.deepcopy(state) modes = {name: mode_state for name, mode_state in state["modelopt_state_dict"]} if "quantize" not in modes or "real_quantize" not in modes: @@ -418,8 +427,6 @@ def write_transformers_modelopt_state(self, component_state: dict, export_dir: P print(f"[export] wrote Transformers ModelOpt graph state for {len(quantizers)} quantizers") def overlay(self, bf16: dict, quantized: dict) -> dict: - from .moe import is_packed_expert_key - merged = dict(bf16) n_weight = n_scale = n_dropped = 0 quantized_modules: set[str] = set() @@ -586,8 +593,6 @@ def assemble_output_dir(input_dir: Path, output_dir: Path, quantized_transformer written so both the diffusers and transformers loaders resolve the new weights. The root Transformers ModelOpt state is skipped when the transformer has packed experts. """ - from .moe import is_packed_expert_key - input_dir, output_dir = Path(input_dir), Path(output_dir) output_dir.mkdir(parents=True, exist_ok=True) target_transformer = output_dir / "transformer" diff --git a/cosmos_framework/quantization/moe.py b/cosmos_framework/quantization/moe.py index 4d20054c..f430444b 100644 --- a/cosmos_framework/quantization/moe.py +++ b/cosmos_framework/quantization/moe.py @@ -20,6 +20,7 @@ from cosmos_framework.model.generator.moe.interface import MoEExperts from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.moe import Qwen3VLMoeTextExpertsGroupedMm +from cosmos_framework.quantization.export import FP8_MAX _PACKED_EXPERT_TYPES = (Qwen3VLMoeTextExpertsGroupedMm, MoEExperts) @@ -118,23 +119,12 @@ def assert_quantizers_calibrated(model: nn.Module) -> None: raise ValueError(f"{len(problems)} quantizer problem(s), first {min(len(problems), 10)}: {shown}") -_FP8_MAX = 448.0 # Mirrors ``export.FP8_MAX``; export imports this module, so it is not imported back. _PACKED_EXPERT_QUANTIZERS = ( "gate_up_proj_input_quantizer", "gate_up_proj_weight_quantizer", "down_proj_input_quantizer", "down_proj_weight_quantizer", ) -_PACKED_EXPERT_KEY_SUFFIXES = tuple( - f"experts.{name}{suffix}" - for name in ("gate_up_proj", "down_proj") - for suffix in ("", ".weight_scale", ".input_scale") -) - - -def is_packed_expert_key(key: str) -> bool: - """True for a packed expert parameter or its scale, e.g. ``...experts.gate_up_proj.weight_scale``.""" - return key.endswith(_PACKED_EXPERT_KEY_SUFFIXES) def _check_per_expert_fp8(name: str, module: QuantPackedExperts) -> None: @@ -174,12 +164,12 @@ def export_packed_experts(net: nn.Module) -> dict[str, torch.Tensor]: for param_name in ("gate_up_proj", "down_proj"): # BF16 models keep BF16 amax; the scales are FP32 on disk. weight_amax = getattr(module, f"{param_name}_weight_quantizer").amax.detach().float() - weight_scale = torch.where(weight_amax == 0, 1.0, weight_amax / _FP8_MAX) # [num_experts,1,1] + weight_scale = torch.where(weight_amax == 0, 1.0, weight_amax / FP8_MAX) # [num_experts,1,1] weight = getattr(module, param_name).detach().float() - fp8 = (weight / weight_scale).clamp(-_FP8_MAX, _FP8_MAX).to(torch.float8_e4m3fn) + fp8 = (weight / weight_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) input_amax = getattr(module, f"{param_name}_input_quantizer").amax.detach().float() key = f"{name}.{param_name}" if name else param_name exported[key] = fp8.contiguous().cpu() exported[f"{key}.weight_scale"] = weight_scale.contiguous().cpu() - exported[f"{key}.input_scale"] = (input_amax / _FP8_MAX).reshape(()).contiguous().cpu() + exported[f"{key}.input_scale"] = (input_amax / FP8_MAX).reshape(()).contiguous().cpu() return exported From 11d3cacd92a463a2fb2870044b2dce4cd513836f Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Mon, 26 Oct 2026 18:19:37 +0100 Subject: [PATCH 4/9] test(quant): trim redundant MoE expert tests Signed-off-by: Dominik Farhan --- tests/quantization/moe_experts_test.py | 39 +------------------------- 1 file changed, 1 insertion(+), 38 deletions(-) diff --git a/tests/quantization/moe_experts_test.py b/tests/quantization/moe_experts_test.py index 39328ba3..710f0186 100644 --- a/tests/quantization/moe_experts_test.py +++ b/tests/quantization/moe_experts_test.py @@ -154,9 +154,7 @@ def test_registration_replaces_generic_fused_experts_wrapper(): model = _quantize(Model()) assert isinstance(model.experts, QuantPackedExperts) - assert not isinstance(model.experts, _QuantFusedExperts) - registered = QuantModuleRegistry.get(Qwen3VLMoeTextExpertsGroupedMm) - assert issubclass(registered, QuantPackedExperts) and not issubclass(registered, _QuantFusedExperts) + assert issubclass(QuantModuleRegistry.get(Qwen3VLMoeTextExpertsGroupedMm), QuantPackedExperts) @pytest.mark.usefixtures("reference_routing") @@ -366,38 +364,3 @@ def test_wrapper_matches_native_block_on_cuda(): expected, _ = native(hidden_states, token_mask=token_mask) actual, _ = block(hidden_states, token_mask=token_mask) assert torch.equal(actual, expected) - - -@pytest.mark.gpus(1) -def test_a2_amax_ignores_padding_on_cuda(): - if not torch.cuda.is_available(): - pytest.skip("CUDA is required for torch._grouped_mm and the Triton routing kernels") - from cosmos_framework.model.generator.reasoner.qwen3_vl_moe.qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock - - torch.manual_seed(0) - block = Qwen3VLMoeTextSparseMoeBlock(_config()) - for param in block.experts.parameters(): - param.data.normal_() - block = block.to(device="cuda", dtype=torch.bfloat16).eval() - hidden_states = torch.randn(32, 16, device="cuda", dtype=torch.bfloat16) - token_mask = torch.ones(32, dtype=torch.bool, device="cuda") - token_mask[-5:] = False - routed = {} - - def record_routing(module, args, kwargs): - routed["expert_indices"] = kwargs["expert_indices"] - - block.experts.register_forward_pre_hook(record_routing, with_kwargs=True) - register_packed_experts() - config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - config["quant_cfg"] = [*config["quant_cfg"], *PER_EXPERT_WEIGHTS] - mtq.quantize(block, config, lambda model: model(hidden_states, token_mask=token_mask)) - - # a2 on the real rows only, one row per (token, selected expert). - real = token_mask.nonzero().squeeze(-1) # [num_real] - experts = routed["expert_indices"][real] # [num_real,top_k] - gate_up = torch.einsum("th,tkhf->tkf", hidden_states[real], block.experts.gate_up_proj[experts]) - gate, up = gate_up.chunk(2, dim=-1) - expected = (F.silu(gate) * up).abs().max().float() - amax = block.experts.down_proj_input_quantizer.amax.float() - torch.testing.assert_close(amax, expected, rtol=2e-2, atol=0) From d4caba94922d5b5c80cff23e4840171897719295 Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Mon, 26 Oct 2026 18:30:10 +0100 Subject: [PATCH 5/9] test(quant): make the root-state skip test reach its guard Signed-off-by: Dominik Farhan --- tests/quantization/modelopt_state_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/quantization/modelopt_state_test.py b/tests/quantization/modelopt_state_test.py index 0905318d..a04c31f6 100644 --- a/tests/quantization/modelopt_state_test.py +++ b/tests/quantization/modelopt_state_test.py @@ -372,6 +372,7 @@ def test_assemble_skips_root_modelopt_state_for_packed_experts(export_modules, t directory.mkdir() (staging / "config.json").write_text(json.dumps({"quantization_config": {"quant_method": "modelopt"}})) save_file({"layers.0.mlp.experts.gate_up_proj": torch.zeros(2, 4, 4)}, str(staging / "model.safetensors")) + torch.save({}, staging / "modelopt_state.pth") def fail(*args, **kwargs): raise AssertionError("root ModelOpt state must not be built") From bb8b537d41aec4ff20a3182a6d43ea79ed1b651e Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Tue, 27 Oct 2026 13:49:57 +0100 Subject: [PATCH 6/9] docs(quant): explain how to quantize packed MoE experts Document the steps the packed-expert path needs: register the wrapper before mtq.quantize or a restore, add per-expert weight quantizer entries to the config, and export with export_quantized_transformer. State that the experts are quantized as W8A8 with static input scales and that weight-only export is not supported. Signed-off-by: Dominik Farhan --- cosmos_framework/quantization/moe.py | 27 +++++++++++++++++++-------- 1 file changed, 19 insertions(+), 8 deletions(-) diff --git a/cosmos_framework/quantization/moe.py b/cosmos_framework/quantization/moe.py index f430444b..eec1af19 100644 --- a/cosmos_framework/quantization/moe.py +++ b/cosmos_framework/quantization/moe.py @@ -1,15 +1,26 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 -"""ModelOpt FP8 wrapper for the packed grouped-mm MoE experts. +"""ModelOpt FP8 quantization of the packed grouped-mm MoE experts. -ModelOpt's generic ``_QuantFusedExperts`` claims any module with 3-D ``gate_up_proj`` / -``down_proj`` parameters but only intercepts ``F.linear``, so on the packed grouped-mm experts -(which run ``torch._grouped_mm``) it would quantize nothing. :class:`QuantPackedExperts` keeps -the experts packed and quantizes them in place: one per-expert weight quantizer per packed -parameter and one shared input quantizer per GEMM, matching what vLLM's FP8 MoE kernels consume. +Quantizes both the expert weights (per-expert scales) and the expert inputs (static scales +calibrated through the wrapper's ``forward``). +This code does not support weight-only. -The quantizers act on tensors inside the expert computation, so its ``forward`` routes the tokens -and runs both grouped GEMMs itself. +How to run quantization: call :func:`register_packed_experts` before ``mtq.quantize`` (and before +restoring a saved ModelOpt state), and give the weight quantizers one scale per expert, as vLLM's +FP8 MoE layout requires; :func:`export_packed_experts` rejects ``mtq.FP8_DEFAULT_CFG`` alone:: + + config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + config["quant_cfg"] += [ + {"quantizer_name": "*gate_up_proj_weight_quantizer", "cfg": {"num_bits": (4, 3), "axis": 0}}, + {"quantizer_name": "*down_proj_weight_quantizer", "cfg": {"num_bits": (4, 3), "axis": 0}}, + ] + +Then export with :func:`~cosmos_framework.quantization.export.export_quantized_transformer`, +which writes the experts via :func:`export_packed_experts`. + +Without the registration, ModelOpt's generic ``_QuantFusedExperts`` claims these experts and, +since it only intercepts ``F.linear`` and they run ``torch._grouped_mm``, quantizes nothing. """ from __future__ import annotations From 709039a46fa219bf1ffdcfa977b0c4f88facf9f4 Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Tue, 27 Oct 2026 14:20:14 +0100 Subject: [PATCH 7/9] fix(quant): pin ModelOpt to v0.44.0 Align the quantization dependency with the ModelOpt version used for Cosmos3 FP8 checkpoints.\n\nCommit cc28c0c removed unsupported nn.Embedding quantizer metadata to make checkpoints compatible with ModelOpt v0.44.0. Pin the dependency and lockfile to that version accordingly.\n\nContext: https://github.com/NVIDIA/cosmos-framework/commit/cc28c0c36bf8f39d5d0a7afe71e237178f180155 --- pyproject.toml | 2 +- uv.lock | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c4664299..e474a650 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ dependencies = [ [project.optional-dependencies] automl = ["nvidia-automl-core==0.1.0"] automl-prompts = ["nvidia-automl-core==0.1.0", "gepa==0.1.2"] -quantization = ["nvidia-modelopt==0.47.0"] +quantization = ["nvidia-modelopt==0.44.0"] reasoner-eval = [ "attrs", "evaluate", diff --git a/uv.lock b/uv.lock index 3ee93e8d..d85f7128 100644 --- a/uv.lock +++ b/uv.lock @@ -2750,7 +2750,7 @@ requires-dist = [ { name = "nvidia-cudnn-frontend" }, { name = "nvidia-dali-cuda120", marker = "extra == 'train'" }, { name = "nvidia-ml-py" }, - { name = "nvidia-modelopt", marker = "extra == 'quantization'", specifier = "==0.47.0" }, + { name = "nvidia-modelopt", marker = "extra == 'quantization'", specifier = "==0.44.0" }, { name = "obstore" }, { name = "omegaconf" }, { name = "open-clip-torch", marker = "extra == 'train'" }, @@ -9835,7 +9835,7 @@ wheels = [ [[package]] name = "nvidia-modelopt" -version = "0.47.0" +version = "0.44.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "ninja" }, @@ -9860,7 +9860,7 @@ dependencies = [ { name = "tqdm" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/46/f32f6183f6a26db8c6bf5c70f89f68ae2cc55c3fbb981016a772c93c4fc2/nvidia_modelopt-0.47.0-py3-none-any.whl", hash = "sha256:db13e62d7efe4a1512e8b8c3f479e28d64022306f94a194f13234e7dca1d3120", size = 2186201, upload-time = "2026-09-23T19:06:44.103Z" }, + { url = "https://files.pythonhosted.org/packages/83/ab/7e12dd238638624cb9d48904e9205abe16a5b26bdd5f9b91e3357821cf90/nvidia_modelopt-0.44.0-py3-none-any.whl", hash = "sha256:9b54a853dfda161db97a0dfce4d7c24269d1d19966b5e2026094186af897f74d", size = 1604658, upload-time = "2026-05-13T20:47:04.007Z" }, ] [[package]] From a5dcfc5bcbc964cea39d52f8672bc18416423599 Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Tue, 27 Oct 2026 15:40:55 +0100 Subject: [PATCH 8/9] perf(quant): compute packed-expert weight scales on export instead of calibrating weight quantizers QuantPackedExperts ran gate_up_proj_weight_quantizer and down_proj_weight_quantizer on the full 3D expert weight tensors on every forward pass during calibration, even though the weights are static and export_packed_experts quantizes them directly from their per-expert max. Remove the two weight quantizers from QuantPackedExperts and compute the per-expert weight amax directly from gate_up_proj and down_proj inside export_packed_experts. This avoids repeatedly reducing the expert weights during calibration and removes the need for callers to append axis=0 weight-quantizer rules to mtq.FP8_DEFAULT_CFG. Signed-off-by: Dominik Farhan --- cosmos_framework/quantization/export.py | 1 - cosmos_framework/quantization/moe.py | 45 +++++++------------------ tests/quantization/moe_experts_test.py | 25 ++++---------- 3 files changed, 20 insertions(+), 51 deletions(-) diff --git a/cosmos_framework/quantization/export.py b/cosmos_framework/quantization/export.py index 160923b1..b9de291c 100644 --- a/cosmos_framework/quantization/export.py +++ b/cosmos_framework/quantization/export.py @@ -478,7 +478,6 @@ def export( if isinstance(module, QuantPackedExperts): for name in ("gate_up_proj", "down_proj"): getattr(module, f"{name}_input_quantizer").disable() - getattr(module, f"{name}_weight_quantizer").disable() # Record ModelOpt's real-quantized graph metadata before the unified # exporter unpacks those weights into plain FP8 checkpoint tensors. diff --git a/cosmos_framework/quantization/moe.py b/cosmos_framework/quantization/moe.py index eec1af19..a711c26f 100644 --- a/cosmos_framework/quantization/moe.py +++ b/cosmos_framework/quantization/moe.py @@ -2,22 +2,14 @@ # SPDX-License-Identifier: OpenMDW-1.1 """ModelOpt FP8 quantization of the packed grouped-mm MoE experts. -Quantizes both the expert weights (per-expert scales) and the expert inputs (static scales -calibrated through the wrapper's ``forward``). +Calibrates and fake-quantizes the expert inputs through the wrapper's ``forward``, and quantizes +the expert weights (with one scale per expert) directly in :func:`export_packed_experts`. This code does not support weight-only. How to run quantization: call :func:`register_packed_experts` before ``mtq.quantize`` (and before -restoring a saved ModelOpt state), and give the weight quantizers one scale per expert, as vLLM's -FP8 MoE layout requires; :func:`export_packed_experts` rejects ``mtq.FP8_DEFAULT_CFG`` alone:: - - config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - config["quant_cfg"] += [ - {"quantizer_name": "*gate_up_proj_weight_quantizer", "cfg": {"num_bits": (4, 3), "axis": 0}}, - {"quantizer_name": "*down_proj_weight_quantizer", "cfg": {"num_bits": (4, 3), "axis": 0}}, - ] - -Then export with :func:`~cosmos_framework.quantization.export.export_quantized_transformer`, -which writes the experts via :func:`export_packed_experts`. +restoring a saved ModelOpt state), then export with +:func:`~cosmos_framework.quantization.export.export_quantized_transformer`, which writes the +experts via :func:`export_packed_experts`. Without the registration, ModelOpt's generic ``_QuantFusedExperts`` claims these experts and, since it only intercepts ``F.linear`` and they run ``torch._grouped_mm``, quantizes nothing. @@ -37,13 +29,11 @@ class QuantPackedExperts(QuantModule): - """Quantized packed grouped-mm experts; the packed parameters stay untouched.""" + """Calibrates and fake-quantizes expert inputs; per-expert FP8 weights are quantized on export.""" def _setup(self): self.gate_up_proj_input_quantizer = TensorQuantizer() - self.gate_up_proj_weight_quantizer = TensorQuantizer() self.down_proj_input_quantizer = TensorQuantizer() - self.down_proj_weight_quantizer = TensorQuantizer() def forward( self, @@ -74,15 +64,13 @@ def _quantized_experts( scores: torch.Tensor, # [num_routed] ) -> torch.Tensor: # [num_routed,hidden_size] # a2 observes ``act(gate) * up`` before the routing scores, because vLLM applies the scores - # after the down GEMM. + # after the down GEMM. Per-expert weight scales are computed directly on export. offsets = torch.cumsum(num_tokens_per_expert, dim=0, dtype=torch.int32) # [num_experts] x = self.gate_up_proj_input_quantizer(x) - h = torch._grouped_mm(x, self.gate_up_proj_weight_quantizer(self.gate_up_proj), offs=offsets) + h = torch._grouped_mm(x, self.gate_up_proj, offs=offsets) g, u = torch.chunk(h, chunks=2, dim=-1) # 2x [num_routed,moe_intermediate_size] a = self.down_proj_input_quantizer(self.act_fn(g) * u) # [num_routed,moe_intermediate_size] - return torch._grouped_mm( - a * scores.unsqueeze(-1), self.down_proj_weight_quantizer(self.down_proj), offs=offsets - ) # [num_routed,hidden_size] + return torch._grouped_mm(a * scores.unsqueeze(-1), self.down_proj, offs=offsets) # [num_routed,hidden_size] def register_packed_experts() -> None: @@ -132,24 +120,17 @@ def assert_quantizers_calibrated(model: nn.Module) -> None: _PACKED_EXPERT_QUANTIZERS = ( "gate_up_proj_input_quantizer", - "gate_up_proj_weight_quantizer", "down_proj_input_quantizer", - "down_proj_weight_quantizer", ) def _check_per_expert_fp8(name: str, module: QuantPackedExperts) -> None: - """Refuse any layout but the one written below: FP8 E4M3, per-expert weight amax, scalar input amax.""" + """Refuse any input quantizer layout but FP8 E4M3 with a scalar amax.""" for quantizer_name in _PACKED_EXPERT_QUANTIZERS: quantizer = getattr(module, quantizer_name) amax_shape = tuple(quantizer.amax.shape) if quantizer.num_bits != (4, 3) or quantizer.block_sizes: found = f"num_bits={quantizer.num_bits}, block_sizes={quantizer.block_sizes}" - elif quantizer_name.endswith("weight_quantizer"): - num_experts = getattr(module, quantizer_name.removesuffix("_weight_quantizer")).shape[0] - if amax_shape == (num_experts, 1, 1): - continue - found = f"amax shape {amax_shape}, expected ({num_experts}, 1, 1)" elif quantizer.amax.numel() == 1: continue else: @@ -174,10 +155,10 @@ def export_packed_experts(net: nn.Module) -> dict[str, torch.Tensor]: _check_per_expert_fp8(name or "experts", module) for param_name in ("gate_up_proj", "down_proj"): # BF16 models keep BF16 amax; the scales are FP32 on disk. - weight_amax = getattr(module, f"{param_name}_weight_quantizer").amax.detach().float() + param = getattr(module, param_name).detach() + weight_amax = param.abs().amax(dim=(1, 2), keepdim=True).float() weight_scale = torch.where(weight_amax == 0, 1.0, weight_amax / FP8_MAX) # [num_experts,1,1] - weight = getattr(module, param_name).detach().float() - fp8 = (weight / weight_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) + fp8 = (param.float() / weight_scale).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn) input_amax = getattr(module, f"{param_name}_input_quantizer").amax.detach().float() key = f"{name}.{param_name}" if name else param_name exported[key] = fp8.contiguous().cpu() diff --git a/tests/quantization/moe_experts_test.py b/tests/quantization/moe_experts_test.py index 710f0186..73e15f30 100644 --- a/tests/quantization/moe_experts_test.py +++ b/tests/quantization/moe_experts_test.py @@ -34,14 +34,8 @@ QUANTIZERS = ( "gate_up_proj_input_quantizer", - "gate_up_proj_weight_quantizer", "down_proj_input_quantizer", - "down_proj_weight_quantizer", ) -PER_EXPERT_WEIGHTS = [ - {"quantizer_name": "*gate_up_proj_weight_quantizer", "cfg": {"num_bits": (4, 3), "axis": 0}}, - {"quantizer_name": "*down_proj_weight_quantizer", "cfg": {"num_bits": (4, 3), "axis": 0}}, -] def _reference_grouped_mm(x: torch.Tensor, w: torch.Tensor, offs: torch.Tensor) -> torch.Tensor: @@ -126,7 +120,7 @@ def _forward_inputs( def _quantize(model: nn.Module, calibration: list[tuple] | None = None, *extra_cfg: dict) -> nn.Module: register_packed_experts() config = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - config["quant_cfg"] = [*config["quant_cfg"], *PER_EXPERT_WEIGHTS, *extra_cfg] + config["quant_cfg"] = [*config["quant_cfg"], *extra_cfg] def forward_loop(model: nn.Module) -> None: for inputs in calibration or [_forward_inputs()]: @@ -177,7 +171,7 @@ def test_wrapper_matches_native_with_quantizers_disabled(experts, top_k, routing @pytest.mark.usefixtures("reference_routing") -def test_each_quantizer_changes_the_wrapper_output(): +def test_each_input_quantizer_changes_the_wrapper_output(): model = _quantize(Model()) inputs = _forward_inputs(seed=4) for name in QUANTIZERS: @@ -192,7 +186,7 @@ def test_each_quantizer_changes_the_wrapper_output(): @pytest.mark.usefixtures("reference_routing") -def test_forward_calibrates_per_expert_weights_and_unpadded_activations(): +def test_forward_calibrates_unpadded_activations(): calibration = [] for seed in (1, 2): hidden_states, _, expert_indices, counts, token_mask = _forward_inputs(seed) @@ -203,10 +197,6 @@ def test_forward_calibrates_per_expert_weights_and_unpadded_activations(): assert_quantizers_calibrated(model) experts = model.experts - for name in ("gate_up_proj", "down_proj"): - amax = getattr(experts, f"{name}_weight_quantizer").amax - assert amax.shape == (4, 1, 1) - assert torch.equal(amax, getattr(experts, name).abs().amax(dim=(1, 2), keepdim=True)) real = [(x[mask], indices[mask]) for x, _, indices, _, mask in calibration] assert torch.equal(experts.gate_up_proj_input_quantizer.amax, torch.stack([x.abs().max() for x, _ in real]).max()) # a2 sees act(gate) * up on the real (token, expert) rows only, before the routing scores. @@ -292,12 +282,11 @@ def test_export_skips_dense_and_disabled_experts_and_rejects_mixed(): @pytest.mark.usefixtures("reference_routing") def test_export_rejects_non_per_expert_fp8(): - register_packed_experts() - scalar_weights = mtq.quantize(Model(), copy.deepcopy(mtq.FP8_DEFAULT_CFG), lambda m: m.experts(*_forward_inputs())) - assert scalar_weights.experts.gate_up_proj_weight_quantizer.amax.numel() == 1 + per_channel_inputs = _quantize(Model(), None, {"quantizer_name": "*input_quantizer", "cfg": {"axis": 0}}) + assert per_channel_inputs.experts.gate_up_proj_input_quantizer.amax.numel() > 1 with pytest.raises(ValueError, match="per-expert FP8") as error: - export_packed_experts(scalar_weights) - assert "experts" in str(error.value) and "gate_up_proj_weight_quantizer" in str(error.value) + export_packed_experts(per_channel_inputs) + assert "experts" in str(error.value) and "gate_up_proj_input_quantizer" in str(error.value) int8 = _quantize(Model(), None, {"quantizer_name": "*experts*", "cfg": {"num_bits": 8}}) assert int8.experts.down_proj_input_quantizer.num_bits == 8 From 3eb4e46643145125abfdf98e3d3bfd6c64ee6418 Mon Sep 17 00:00:00 2001 From: Dominik Farhan Date: Tue, 27 Oct 2026 17:18:32 +0100 Subject: [PATCH 9/9] build: add datasets to the quantization extra cosmos_framework.quantization.calibration imports `datasets` at module level, and the package __init__ imports calibration eagerly, so `import cosmos_framework.quantization` fails with ModuleNotFoundError when only the quantization extra is installed. `datasets` was only provided by the train extra, and nvidia-modelopt does not depend on it (it imports it lazily). Left unpinned to match the train extra; uv.lock keeps it at 4.6.1. --- pyproject.toml | 2 +- uv.lock | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index e474a650..2568f0b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ dependencies = [ [project.optional-dependencies] automl = ["nvidia-automl-core==0.1.0"] automl-prompts = ["nvidia-automl-core==0.1.0", "gepa==0.1.2"] -quantization = ["nvidia-modelopt==0.44.0"] +quantization = ["datasets", "nvidia-modelopt==0.44.0"] reasoner-eval = [ "attrs", "evaluate", diff --git a/uv.lock b/uv.lock index d85f7128..5b1090b0 100644 --- a/uv.lock +++ b/uv.lock @@ -2389,6 +2389,7 @@ guardrail = [ { name = "sentencepiece" }, ] quantization = [ + { name = "datasets" }, { name = "nvidia-modelopt" }, ] reasoner-eval = [ @@ -2701,6 +2702,7 @@ requires-dist = [ { name = "boto3" }, { name = "botocore", marker = "extra == 'train'" }, { name = "cattrs" }, + { name = "datasets", marker = "extra == 'quantization'" }, { name = "datasets", marker = "extra == 'train'" }, { name = "diffusers", specifier = ">=0.39.0" }, { name = "dists-pytorch", marker = "extra == 'train'" },