From 0214ef90a9c396e883082f994bf8ddee1df9620c Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Wed, 30 Sep 2026 12:20:36 +0530 Subject: [PATCH] Pin qwix to 0.1.8 to fix FP8 quantization with current Flax The qwix pin (commit 408a0f48, Dec 2025) predates google/qwix@98a44ed0 (2026-01-07), which added `out_sharding` to QtProvider.conv_general_dilated. Flax >= 0.12.6 (our minimum) always passes `out_sharding` from nnx.Conv, so qwix.quantize_model fails on the first conv (Wan patch_embedding) with: TypeError: QtProvider.conv_general_dilated() got an unexpected keyword argument 'out_sharding' This breaks every qwix-quantized run (e.g. Wan 2.1 FP8 training, b/537854580). qwix 0.1.8 (PyPI, 2026-06-22) includes the fix and requires only flax>=0.12.0. Switching from a GitHub archive URL to a PyPI version also keeps direct URL references out of the package metadata. The generated requirements and deps table are edited by hand to avoid unrelated churn from re-running seed-env. --- dependencies/requirements/base_requirements/requirements.txt | 2 +- .../requirements/generated_requirements/requirements.txt | 2 +- src/install_maxdiffusion_extra_deps/extra_deps_from_github.txt | 1 - src/maxdiffusion/dependency_versions_table.py | 2 +- 4 files changed, 3 insertions(+), 4 deletions(-) diff --git a/dependencies/requirements/base_requirements/requirements.txt b/dependencies/requirements/base_requirements/requirements.txt index f500bc598..b2a5d239c 100644 --- a/dependencies/requirements/base_requirements/requirements.txt +++ b/dependencies/requirements/base_requirements/requirements.txt @@ -44,4 +44,4 @@ transformers<5.0.0 # installing GPU versions from PyPI when running seed-env torch @ https://download.pytorch.org/whl/cpu/torch-2.10.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.25.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl -qwix @ https://github.com/google/qwix/archive/408a0f48f988b6c5b180e07f0cb1d05997bf0dcc.zip +qwix==0.1.8 diff --git a/dependencies/requirements/generated_requirements/requirements.txt b/dependencies/requirements/generated_requirements/requirements.txt index 3aee20155..da356c936 100644 --- a/dependencies/requirements/generated_requirements/requirements.txt +++ b/dependencies/requirements/generated_requirements/requirements.txt @@ -147,7 +147,7 @@ pytest>=8.4.2 python-dateutil>=2.9.0.post0 pytokens>=0.4.1 pyyaml>=6.0.3 -qwix @ https://github.com/google/qwix/archive/408a0f48f988b6c5b180e07f0cb1d05997bf0dcc.zip +qwix==0.1.8 regex>=2026.2.28 requests-oauthlib>=2.0.0 requests>=2.32.5 diff --git a/src/install_maxdiffusion_extra_deps/extra_deps_from_github.txt b/src/install_maxdiffusion_extra_deps/extra_deps_from_github.txt index 130c41aed..a651439be 100644 --- a/src/install_maxdiffusion_extra_deps/extra_deps_from_github.txt +++ b/src/install_maxdiffusion_extra_deps/extra_deps_from_github.txt @@ -1,3 +1,2 @@ torch @ https://download.pytorch.org/whl/cpu/torch-2.10.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl torchvision @ https://download.pytorch.org/whl/cpu/torchvision-0.25.0%2Bcpu-cp312-cp312-manylinux_2_28_x86_64.whl -qwix @ https://github.com/google/qwix/archive/408a0f48f988b6c5b180e07f0cb1d05997bf0dcc.zip diff --git a/src/maxdiffusion/dependency_versions_table.py b/src/maxdiffusion/dependency_versions_table.py index 90dda5004..241035ac3 100644 --- a/src/maxdiffusion/dependency_versions_table.py +++ b/src/maxdiffusion/dependency_versions_table.py @@ -25,7 +25,7 @@ "pyink": "pyink", "pylint": "pylint", "pytest": "pytest", - "qwix": "qwix @ https://github.com/google/qwix/archive/408a0f48f988b6c5b180e07f0cb1d05997bf0dcc.zip", + "qwix": "qwix==0.1.8", "ruff": "ruff", "scikit-image": "scikit-image", "sentencepiece": "sentencepiece",