From f238d93d868a09bbc395fff8867780ff7d8ac75e Mon Sep 17 00:00:00 2001 From: kinjalpatel27 <31936134+kinjalpatel27@users.noreply.github.com> Date: Thu, 16 Apr 2026 09:42:36 -0700 Subject: [PATCH 01/30] vLLM fakequant fold weight_quantizer for megatron export (#1246) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix During Megatron→vLLM fakequant export (`export_mcore_gpt_to_hf_vllm_fq`), the `weight_quantizer` is now applied as fake-quantization (quantize + dequantize) directly into the exported weight tensor, and its amax is no longer saved to `quantizer_state.pth`. On reload, if `weight_quantizer` keys are absent from the checkpoint (because they were folded at export time), the corresponding quantizer modules are disabled. This change is useful especially when amax across experts are not synced for `weight_quantizer`, this allows the `weight_quantizer` to keep them different for better accuracy. ### Usage ```python # Unchanged — export API is the same export_mcore_gpt_to_hf_vllm_fq(model, pretrained_model_name_or_path=..., export_dir=...) ``` ### Testing Step 1 — Quantize (run from Megatron-LM `examples/post_training/modelopt`): ```bash HF_MODEL_CKPT= MLM_MODEL_SAVE= \ bash quantize.sh NVFP4_DEFAULT_CFG ``` Step 2 — Export for vLLM fakequant: ```bash MLM_EXTRA_ARGS=--export-vllm-fq \ HF_MODEL_CKPT= \ MLM_MODEL_CKPT= \ EXPORT_DIR= \ bash export.sh ``` Step 3 — Serve (run from examples/vllm_serve): ```bash QUANT_CFG=NVFP4_DEFAULT_CFG \ QUANT_FILE_PATH=/quantizer_state.pth \ python3 vllm_serve_fakequant.py \ -tp 1 --served-model-name \ --host 0.0.0.0 --port 8000 \ --trust-remote-code --enforce-eager \ --disable-custom-all-reduce \ --gpu-memory-utilization 0.8 ``` ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A ### Additional Information ## Summary by CodeRabbit * **Bug Fixes** * Better handling when loading checkpoints: missing weight-quantizer entries are validated and corresponding modules are disabled to avoid load failures. * **Improvements** * Export now folds enabled weight quantizers into exported weights when present and omits internal weight-quantizer tensors from the exported state to produce cleaner exports. --------- Signed-off-by: Kinjal Patel --- examples/vllm_serve/fakequant_worker.py | 1 + examples/vllm_serve/vllm_reload_utils.py | 80 +++++++++++++++---- .../export/plugins/vllm_fakequant_megatron.py | 42 +++++++++- 3 files changed, 106 insertions(+), 17 deletions(-) diff --git a/examples/vllm_serve/fakequant_worker.py b/examples/vllm_serve/fakequant_worker.py index c56088fd7fe..b88af9c72ee 100644 --- a/examples/vllm_serve/fakequant_worker.py +++ b/examples/vllm_serve/fakequant_worker.py @@ -49,6 +49,7 @@ def _fakequant_run_prolog_worker(self) -> None: trust_remote_code = os.environ.get("TRUST_REMOTE_CODE", "false").lower() == "true" + tokenizer = AutoTokenizer.from_pretrained( self.model_runner.model_config.tokenizer, trust_remote_code=trust_remote_code ) diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index b67c92ae6ea..2b59d1be2bd 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -31,9 +31,32 @@ convert_to_quantized_model, restore_quantizer_state, ) +from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import is_quantized +def _union_quantizer_keys_across_ranks(local_quantizer_keys: list[str]) -> set[str]: + """Union of quantizer key strings from every rank (same file on all ranks → identical to local).""" + local = set(local_quantizer_keys) + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): + return local + if torch.distributed.get_world_size() <= 1: + return local + try: + world_size = torch.distributed.get_world_size() + gathered: list[list[str]] = [[] for _ in range(world_size)] + torch.distributed.all_gather_object(gathered, list(local_quantizer_keys)) + out: set[str] = set() + for g in gathered: + out.update(g) + return out + except Exception as e: + warnings.warn( + f"Could not all_gather quantizer key lists across ranks ({e}); using this rank's keys only." + ) + return local + + def _values_equal(v1: Any, v2: Any) -> bool: """Compare values, handling dicts with tensors.""" if isinstance(v1, dict) and isinstance(v2, dict): @@ -285,7 +308,7 @@ def filter_modelopt_state_quantizer_state_for_model( model: Model with quantizers (must already be converted) """ from modelopt.torch.quantization.conversion import quantizer_state - from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer + from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.utils import get_unwrapped_name model_qstate = quantizer_state(model) @@ -435,24 +458,51 @@ def load_state_dict_from_path( # Count quant keys in checkpoint and model checkpoint_quant_keys = [key for key in saved_quant_dict if "quantizer" in key] model_quant_keys = [key for key in current_state_dict if "quantizer" in key] - for key in checkpoint_quant_keys: - if key not in model_quant_keys: - print(f"Key {key} not found in model state dict, but exists in checkpoint") + ckpt_key_set = set(checkpoint_quant_keys) + global_ckpt_key_set = _union_quantizer_keys_across_ranks(checkpoint_quant_keys) + # For weight quantizers absent from the checkpoint the weights were already fake-quantized + # at export time (amax folded into weights). Disable those quantizers so that fold_weight + # is a no-op for them. Non-weight keys missing on this rank but present on another rank's + # shard are omitted from global_missing (all_gather union of key strings). + missing_wq_module_paths: set[str] = set() + global_missing_non_wq: list[str] = [] for key in model_quant_keys: - if key not in checkpoint_quant_keys: - raise ValueError(f"Key {key} not found in checkpoint state dict, but exists in model") - - checkpoint_quant_count = len(checkpoint_quant_keys) - model_quant_count = len(model_quant_keys) - - # Ensure counts match - if checkpoint_quant_count != model_quant_count: + if key in ckpt_key_set: + continue + if "weight_quantizer" in key: + # Per-rank shard: only disable using this rank's checkpoint contents. + parts = key.split(".") + weight_quantizer_index = next( + (i for i, p in enumerate(parts) if p.endswith("weight_quantizer")), + None, + ) + if weight_quantizer_index is not None: + missing_wq_module_paths.add(".".join(parts[: weight_quantizer_index + 1])) + else: + raise ValueError( + f"Missing checkpoint key {key!r} looks like a weight quantizer, but no path " + "component ends with 'weight_quantizer'; cannot map to a module to disable." + ) + elif key not in global_ckpt_key_set: + global_missing_non_wq.append(key) + + if global_missing_non_wq: + keys = sorted(global_missing_non_wq) + n = len(keys) + sample, rest = keys[:8], n - 8 warnings.warn( - f"Mismatch in quantizer state key counts: checkpoint has {checkpoint_quant_count} " - f"quant keys but model has {model_quant_count} quantizer state keys. " - f"This can happen if the model is using PP." + f"{n} quantizer key(s) missing from every rank's checkpoint (after all_gather):" + f"{sample}{' ... (+{rest} more)' if rest > 0 else ''}" ) + for name, module in model.named_modules(): + if ( + name in missing_wq_module_paths + and isinstance(module, TensorQuantizer) + and hasattr(module, "disable") + ): + module.disable() + # Update quant values saved_quant_dict = process_state_dict_for_tp(saved_quant_dict, current_state_dict) for key, value in saved_quant_dict.items(): diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 9a41ae2baff..c8e45be3650 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -117,6 +117,10 @@ def _get_quantized_state( ) -> tuple[dict[str, torch.Tensor], str, int]: """Return a state_dict, quantization format, and block_size of the module. + The weight_quantizer is folded into the weight via fake-quantization + (quantize + dequantize), and its amax is not exported. The vLLM fakequant + reload path is expected to disable the weight quantizer when the amax is absent. + Args: module: The target module to perform real quantization. dtype: The default data type. @@ -133,14 +137,48 @@ def _get_quantized_state( block_size = 0 if hasattr(module, "weight") and module.weight is not None: - weight = module.weight.to(dtype).cpu() - name_to_value["weight"] = weight + weight = module.weight.to(dtype) + # Fold the weight_quantizer into the weight by applying fake-quantization + # (quantize then dequantize). The weight_quantizer amax is not exported; + # the vLLM fakequant reload path disables the weight quantizer when absent. + weight_quantizer = getattr(module, "weight_quantizer", None) + if weight_quantizer is not None and weight_quantizer.is_enabled: + with torch.no_grad(): + # NVFP4-like kernels may need CUDA; if weights are CPU after gather, run on + # CUDA then ``weight_quantizer.to`` back (full module round-trip). + quant_device = ( + torch.device("cuda", torch.cuda.current_device()) + if weight.device.type == "cpu" and torch.cuda.is_available() + else weight.device + ) + # TensorQuantizer does not expose nn.Module.device (custom __getattr__). + param_device = next(weight_quantizer.parameters(), None) + buf_device = next(weight_quantizer.buffers(), None) + wq_dev = ( + param_device.device + if param_device is not None + else (buf_device.device if buf_device is not None else torch.device("cpu")) + ) + need_move = wq_dev != quant_device + if need_move: + weight_quantizer.to(quant_device) + try: + weight = weight_quantizer(weight.to(quant_device)).to(dtype) + finally: + if need_move: + weight_quantizer.to(wq_dev) + name_to_value["weight"] = weight.cpu() else: return name_to_value, qformat, block_size if hasattr(module, "bias") and module.bias is not None: name_to_value["bias"] = module.bias.to(dtype).cpu() + + # Only save input/output quantizer state; weight_quantizer amax is not exported + # since it has been folded into the weight above. for name, param in get_quantizer_state_dict(module).items(): + if "weight_quantizer" in name: + continue for key, value in param.items(): name_to_value[name + "." + key] = value.to(dtype).cpu() return name_to_value, qformat, block_size From 0a4908d1850023d9a1498868c1b591470d272554 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 16 Apr 2026 10:19:33 -0700 Subject: [PATCH 02/30] [chore]: weekly bump of uv.lock on main (2026-04-15) (#1266) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Automated weekly update of uv.lock file for nSpect Scanning: - `uv.lock` — upgraded all transitive dependencies to latest compatible versions Signed-off-by: github-actions[bot] Co-authored-by: github-actions[bot] --- uv.lock | 478 ++++++++++++++++++++++++++++++++++++++++++-------------- 1 file changed, 362 insertions(+), 116 deletions(-) diff --git a/uv.lock b/uv.lock index 9184371f254..0f4710f92b4 100644 --- a/uv.lock +++ b/uv.lock @@ -278,6 +278,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/06/f3/39cf3367b8107baa44f861dc802cbf16263c945b62d8265d36034fc07bea/cachetools-7.0.5-py3-none-any.whl", hash = "sha256:46bc8ebefbe485407621d0a4264b23c080cedd913921bad7ac3ed2f26c183114", size = 13918, upload-time = "2026-03-09T20:51:27.33Z" }, ] +[[package]] +name = "cbcbox" +version = "2.929" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6a/9e/c41844f3a746500b88982817233f3f2f4fa6d3a554cedb9222bac9efcbea/cbcbox-2.929-py3-none-macosx_15_0_arm64.whl", hash = "sha256:610fd250f737b19f599d56cca299682a4519472b3ff6a13044502eb71e03931e", size = 59626733, upload-time = "2026-03-23T15:28:20.758Z" }, + { url = "https://files.pythonhosted.org/packages/3f/08/172af0d618ede8668862ae9c83ffcd929bc966610c1ce357b268edb33645/cbcbox-2.929-py3-none-macosx_15_0_x86_64.whl", hash = "sha256:b369071081176fd55ba9d104c2f2330133c36f0264ebe0215c50a366f290a050", size = 87883726, upload-time = "2026-03-23T15:28:24.874Z" }, + { url = "https://files.pythonhosted.org/packages/d4/c7/a57493dd959a57dcb8dc31aeef468c0e95d066d4bbf8fc59f9dda3a1fd97/cbcbox-2.929-py3-none-manylinux2014_aarch64.whl", hash = "sha256:62841ab4ed4e1c368bcba8d62c66cd7dd03f3dfaf1ecf4a01c6385a122d1946e", size = 145740150, upload-time = "2026-03-23T15:28:29.617Z" }, + { url = "https://files.pythonhosted.org/packages/4c/86/502b81d060603e1016a8cb471064e4e65d8009c2feed60ea878b8513a728/cbcbox-2.929-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a57a403cc44fb2af0aa53cb790a7642390f1836a5ff6a7a0b31b30eb6597e3e0", size = 181208771, upload-time = "2026-03-23T15:28:35.544Z" }, + { url = "https://files.pythonhosted.org/packages/0c/85/2fef9c7c3054ebfb9b0e84fabbeac776a18514b025efa5abd768c889f560/cbcbox-2.929-py3-none-win_amd64.whl", hash = "sha256:0a91420befb965ec1763aab62d7908d3a89a51af845e2a4e05c09d99a67f1151", size = 135305901, upload-time = "2026-03-23T15:28:42.021Z" }, +] + [[package]] name = "certifi" version = "2026.2.25" @@ -287,6 +299,66 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9a/3c/c17fb3ca2d9c3acff52e30b309f538586f9f5b9c9cf454f3845fc9af4881/certifi-2026.2.25-py3-none-any.whl", hash = "sha256:027692e4402ad994f1c42e52a4997a9763c646b73e4096e4d5d6db8af1d6f0fa", size = 153684, upload-time = "2026-02-25T02:54:15.766Z" }, ] +[[package]] +name = "cffi" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pycparser", marker = "implementation_name != 'PyPy'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/93/d7/516d984057745a6cd96575eea814fe1edd6646ee6efd552fb7b0921dec83/cffi-2.0.0-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:0cf2d91ecc3fcc0625c2c530fe004f82c110405f101548512cce44322fa8ac44", size = 184283, upload-time = "2025-09-08T23:22:08.01Z" }, + { url = "https://files.pythonhosted.org/packages/9e/84/ad6a0b408daa859246f57c03efd28e5dd1b33c21737c2db84cae8c237aa5/cffi-2.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:f73b96c41e3b2adedc34a7356e64c8eb96e03a3782b535e043a986276ce12a49", size = 180504, upload-time = "2025-09-08T23:22:10.637Z" }, + { url = "https://files.pythonhosted.org/packages/50/bd/b1a6362b80628111e6653c961f987faa55262b4002fcec42308cad1db680/cffi-2.0.0-cp310-cp310-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:53f77cbe57044e88bbd5ed26ac1d0514d2acf0591dd6bb02a3ae37f76811b80c", size = 208811, upload-time = "2025-09-08T23:22:12.267Z" }, + { url = "https://files.pythonhosted.org/packages/4f/27/6933a8b2562d7bd1fb595074cf99cc81fc3789f6a6c05cdabb46284a3188/cffi-2.0.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3e837e369566884707ddaf85fc1744b47575005c0a229de3327f8f9a20f4efeb", size = 216402, upload-time = "2025-09-08T23:22:13.455Z" }, + { url = "https://files.pythonhosted.org/packages/05/eb/b86f2a2645b62adcfff53b0dd97e8dfafb5c8aa864bd0d9a2c2049a0d551/cffi-2.0.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:5eda85d6d1879e692d546a078b44251cdd08dd1cfb98dfb77b670c97cee49ea0", size = 203217, upload-time = "2025-09-08T23:22:14.596Z" }, + { url = "https://files.pythonhosted.org/packages/9f/e0/6cbe77a53acf5acc7c08cc186c9928864bd7c005f9efd0d126884858a5fe/cffi-2.0.0-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:9332088d75dc3241c702d852d4671613136d90fa6881da7d770a483fd05248b4", size = 203079, upload-time = "2025-09-08T23:22:15.769Z" }, + { url = "https://files.pythonhosted.org/packages/98/29/9b366e70e243eb3d14a5cb488dfd3a0b6b2f1fb001a203f653b93ccfac88/cffi-2.0.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc7de24befaeae77ba923797c7c87834c73648a05a4bde34b3b7e5588973a453", size = 216475, upload-time = "2025-09-08T23:22:17.427Z" }, + { url = "https://files.pythonhosted.org/packages/21/7a/13b24e70d2f90a322f2900c5d8e1f14fa7e2a6b3332b7309ba7b2ba51a5a/cffi-2.0.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cf364028c016c03078a23b503f02058f1814320a56ad535686f90565636a9495", size = 218829, upload-time = "2025-09-08T23:22:19.069Z" }, + { url = "https://files.pythonhosted.org/packages/60/99/c9dc110974c59cc981b1f5b66e1d8af8af764e00f0293266824d9c4254bc/cffi-2.0.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:e11e82b744887154b182fd3e7e8512418446501191994dbf9c9fc1f32cc8efd5", size = 211211, upload-time = "2025-09-08T23:22:20.588Z" }, + { url = "https://files.pythonhosted.org/packages/49/72/ff2d12dbf21aca1b32a40ed792ee6b40f6dc3a9cf1644bd7ef6e95e0ac5e/cffi-2.0.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8ea985900c5c95ce9db1745f7933eeef5d314f0565b27625d9a10ec9881e1bfb", size = 218036, upload-time = "2025-09-08T23:22:22.143Z" }, + { url = "https://files.pythonhosted.org/packages/e2/cc/027d7fb82e58c48ea717149b03bcadcbdc293553edb283af792bd4bcbb3f/cffi-2.0.0-cp310-cp310-win32.whl", hash = "sha256:1f72fb8906754ac8a2cc3f9f5aaa298070652a0ffae577e0ea9bd480dc3c931a", size = 172184, upload-time = "2025-09-08T23:22:23.328Z" }, + { url = "https://files.pythonhosted.org/packages/33/fa/072dd15ae27fbb4e06b437eb6e944e75b068deb09e2a2826039e49ee2045/cffi-2.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:b18a3ed7d5b3bd8d9ef7a8cb226502c6bf8308df1525e1cc676c3680e7176739", size = 182790, upload-time = "2025-09-08T23:22:24.752Z" }, + { url = "https://files.pythonhosted.org/packages/12/4a/3dfd5f7850cbf0d06dc84ba9aa00db766b52ca38d8b86e3a38314d52498c/cffi-2.0.0-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:b4c854ef3adc177950a8dfc81a86f5115d2abd545751a304c5bcf2c2c7283cfe", size = 184344, upload-time = "2025-09-08T23:22:26.456Z" }, + { url = "https://files.pythonhosted.org/packages/4f/8b/f0e4c441227ba756aafbe78f117485b25bb26b1c059d01f137fa6d14896b/cffi-2.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2de9a304e27f7596cd03d16f1b7c72219bd944e99cc52b84d0145aefb07cbd3c", size = 180560, upload-time = "2025-09-08T23:22:28.197Z" }, + { url = "https://files.pythonhosted.org/packages/b1/b7/1200d354378ef52ec227395d95c2576330fd22a869f7a70e88e1447eb234/cffi-2.0.0-cp311-cp311-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:baf5215e0ab74c16e2dd324e8ec067ef59e41125d3eade2b863d294fd5035c92", size = 209613, upload-time = "2025-09-08T23:22:29.475Z" }, + { url = "https://files.pythonhosted.org/packages/b8/56/6033f5e86e8cc9bb629f0077ba71679508bdf54a9a5e112a3c0b91870332/cffi-2.0.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:730cacb21e1bdff3ce90babf007d0a0917cc3e6492f336c2f0134101e0944f93", size = 216476, upload-time = "2025-09-08T23:22:31.063Z" }, + { url = "https://files.pythonhosted.org/packages/dc/7f/55fecd70f7ece178db2f26128ec41430d8720f2d12ca97bf8f0a628207d5/cffi-2.0.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:6824f87845e3396029f3820c206e459ccc91760e8fa24422f8b0c3d1731cbec5", size = 203374, upload-time = "2025-09-08T23:22:32.507Z" }, + { url = "https://files.pythonhosted.org/packages/84/ef/a7b77c8bdc0f77adc3b46888f1ad54be8f3b7821697a7b89126e829e676a/cffi-2.0.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:9de40a7b0323d889cf8d23d1ef214f565ab154443c42737dfe52ff82cf857664", size = 202597, upload-time = "2025-09-08T23:22:34.132Z" }, + { url = "https://files.pythonhosted.org/packages/d7/91/500d892b2bf36529a75b77958edfcd5ad8e2ce4064ce2ecfeab2125d72d1/cffi-2.0.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8941aaadaf67246224cee8c3803777eed332a19d909b47e29c9842ef1e79ac26", size = 215574, upload-time = "2025-09-08T23:22:35.443Z" }, + { url = "https://files.pythonhosted.org/packages/44/64/58f6255b62b101093d5df22dcb752596066c7e89dd725e0afaed242a61be/cffi-2.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:a05d0c237b3349096d3981b727493e22147f934b20f6f125a3eba8f994bec4a9", size = 218971, upload-time = "2025-09-08T23:22:36.805Z" }, + { url = "https://files.pythonhosted.org/packages/ab/49/fa72cebe2fd8a55fbe14956f9970fe8eb1ac59e5df042f603ef7c8ba0adc/cffi-2.0.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:94698a9c5f91f9d138526b48fe26a199609544591f859c870d477351dc7b2414", size = 211972, upload-time = "2025-09-08T23:22:38.436Z" }, + { url = "https://files.pythonhosted.org/packages/0b/28/dd0967a76aab36731b6ebfe64dec4e981aff7e0608f60c2d46b46982607d/cffi-2.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:5fed36fccc0612a53f1d4d9a816b50a36702c28a2aa880cb8a122b3466638743", size = 217078, upload-time = "2025-09-08T23:22:39.776Z" }, + { url = "https://files.pythonhosted.org/packages/2b/c0/015b25184413d7ab0a410775fdb4a50fca20f5589b5dab1dbbfa3baad8ce/cffi-2.0.0-cp311-cp311-win32.whl", hash = "sha256:c649e3a33450ec82378822b3dad03cc228b8f5963c0c12fc3b1e0ab940f768a5", size = 172076, upload-time = "2025-09-08T23:22:40.95Z" }, + { url = "https://files.pythonhosted.org/packages/ae/8f/dc5531155e7070361eb1b7e4c1a9d896d0cb21c49f807a6c03fd63fc877e/cffi-2.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:66f011380d0e49ed280c789fbd08ff0d40968ee7b665575489afa95c98196ab5", size = 182820, upload-time = "2025-09-08T23:22:42.463Z" }, + { url = "https://files.pythonhosted.org/packages/95/5c/1b493356429f9aecfd56bc171285a4c4ac8697f76e9bbbbb105e537853a1/cffi-2.0.0-cp311-cp311-win_arm64.whl", hash = "sha256:c6638687455baf640e37344fe26d37c404db8b80d037c3d29f58fe8d1c3b194d", size = 177635, upload-time = "2025-09-08T23:22:43.623Z" }, + { url = "https://files.pythonhosted.org/packages/ea/47/4f61023ea636104d4f16ab488e268b93008c3d0bb76893b1b31db1f96802/cffi-2.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d02d6655b0e54f54c4ef0b94eb6be0607b70853c45ce98bd278dc7de718be5d", size = 185271, upload-time = "2025-09-08T23:22:44.795Z" }, + { url = "https://files.pythonhosted.org/packages/df/a2/781b623f57358e360d62cdd7a8c681f074a71d445418a776eef0aadb4ab4/cffi-2.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8eca2a813c1cb7ad4fb74d368c2ffbbb4789d377ee5bb8df98373c2cc0dee76c", size = 181048, upload-time = "2025-09-08T23:22:45.938Z" }, + { url = "https://files.pythonhosted.org/packages/ff/df/a4f0fbd47331ceeba3d37c2e51e9dfc9722498becbeec2bd8bc856c9538a/cffi-2.0.0-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:21d1152871b019407d8ac3985f6775c079416c282e431a4da6afe7aefd2bccbe", size = 212529, upload-time = "2025-09-08T23:22:47.349Z" }, + { url = "https://files.pythonhosted.org/packages/d5/72/12b5f8d3865bf0f87cf1404d8c374e7487dcf097a1c91c436e72e6badd83/cffi-2.0.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b21e08af67b8a103c71a250401c78d5e0893beff75e28c53c98f4de42f774062", size = 220097, upload-time = "2025-09-08T23:22:48.677Z" }, + { url = "https://files.pythonhosted.org/packages/c2/95/7a135d52a50dfa7c882ab0ac17e8dc11cec9d55d2c18dda414c051c5e69e/cffi-2.0.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:1e3a615586f05fc4065a8b22b8152f0c1b00cdbc60596d187c2a74f9e3036e4e", size = 207983, upload-time = "2025-09-08T23:22:50.06Z" }, + { url = "https://files.pythonhosted.org/packages/3a/c8/15cb9ada8895957ea171c62dc78ff3e99159ee7adb13c0123c001a2546c1/cffi-2.0.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:81afed14892743bbe14dacb9e36d9e0e504cd204e0b165062c488942b9718037", size = 206519, upload-time = "2025-09-08T23:22:51.364Z" }, + { url = "https://files.pythonhosted.org/packages/78/2d/7fa73dfa841b5ac06c7b8855cfc18622132e365f5b81d02230333ff26e9e/cffi-2.0.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3e17ed538242334bf70832644a32a7aae3d83b57567f9fd60a26257e992b79ba", size = 219572, upload-time = "2025-09-08T23:22:52.902Z" }, + { url = "https://files.pythonhosted.org/packages/07/e0/267e57e387b4ca276b90f0434ff88b2c2241ad72b16d31836adddfd6031b/cffi-2.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3925dd22fa2b7699ed2617149842d2e6adde22b262fcbfada50e3d195e4b3a94", size = 222963, upload-time = "2025-09-08T23:22:54.518Z" }, + { url = "https://files.pythonhosted.org/packages/b6/75/1f2747525e06f53efbd878f4d03bac5b859cbc11c633d0fb81432d98a795/cffi-2.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c8f814d84194c9ea681642fd164267891702542f028a15fc97d4674b6206187", size = 221361, upload-time = "2025-09-08T23:22:55.867Z" }, + { url = "https://files.pythonhosted.org/packages/7b/2b/2b6435f76bfeb6bbf055596976da087377ede68df465419d192acf00c437/cffi-2.0.0-cp312-cp312-win32.whl", hash = "sha256:da902562c3e9c550df360bfa53c035b2f241fed6d9aef119048073680ace4a18", size = 172932, upload-time = "2025-09-08T23:22:57.188Z" }, + { url = "https://files.pythonhosted.org/packages/f8/ed/13bd4418627013bec4ed6e54283b1959cf6db888048c7cf4b4c3b5b36002/cffi-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:da68248800ad6320861f129cd9c1bf96ca849a2771a59e0344e88681905916f5", size = 183557, upload-time = "2025-09-08T23:22:58.351Z" }, + { url = "https://files.pythonhosted.org/packages/95/31/9f7f93ad2f8eff1dbc1c3656d7ca5bfd8fb52c9d786b4dcf19b2d02217fa/cffi-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:4671d9dd5ec934cb9a73e7ee9676f9362aba54f7f34910956b84d727b0d73fb6", size = 177762, upload-time = "2025-09-08T23:22:59.668Z" }, + { url = "https://files.pythonhosted.org/packages/4b/8d/a0a47a0c9e413a658623d014e91e74a50cdd2c423f7ccfd44086ef767f90/cffi-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:00bdf7acc5f795150faa6957054fbbca2439db2f775ce831222b66f192f03beb", size = 185230, upload-time = "2025-09-08T23:23:00.879Z" }, + { url = "https://files.pythonhosted.org/packages/4a/d2/a6c0296814556c68ee32009d9c2ad4f85f2707cdecfd7727951ec228005d/cffi-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45d5e886156860dc35862657e1494b9bae8dfa63bf56796f2fb56e1679fc0bca", size = 181043, upload-time = "2025-09-08T23:23:02.231Z" }, + { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" }, + { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" }, + { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" }, + { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" }, + { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" }, + { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" }, + { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" }, + { url = "https://files.pythonhosted.org/packages/eb/6d/bf9bda840d5f1dfdbf0feca87fbdb64a918a69bca42cfa0ba7b137c48cb8/cffi-2.0.0-cp313-cp313-win32.whl", hash = "sha256:74a03b9698e198d47562765773b4a8309919089150a0bb17d829ad7b44b60d27", size = 172909, upload-time = "2025-09-08T23:23:14.32Z" }, + { url = "https://files.pythonhosted.org/packages/37/18/6519e1ee6f5a1e579e04b9ddb6f1676c17368a7aba48299c3759bbc3c8b3/cffi-2.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:19f705ada2530c1167abacb171925dd886168931e0a7b78f5bffcae5c6b5be75", size = 183402, upload-time = "2025-09-08T23:23:15.535Z" }, + { url = "https://files.pythonhosted.org/packages/cb/0e/02ceeec9a7d6ee63bb596121c2c8e9b3a9e150936f4fbef6ca1943e6137c/cffi-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:256f80b80ca3853f90c21b23ee78cd008713787b1b1e93eae9f3d6a7134abd91", size = 177780, upload-time = "2025-09-08T23:23:16.761Z" }, +] + [[package]] name = "cfgv" version = "3.5.0" @@ -503,10 +575,10 @@ sdist = { url = "https://files.pythonhosted.org/packages/54/27/01d9078a77b9e31b7 [[package]] name = "cuda-pathfinder" -version = "1.5.2" +version = "1.5.3" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f2/f9/1b9b60a30fc463c14cdea7a77228131a0ccc89572e8df9cb86c9648271ab/cuda_pathfinder-1.5.2-py3-none-any.whl", hash = "sha256:0c5f160a7756c5b072723cbbd6d861e38917ef956c68150b02f0b6e9271c71fa", size = 49988, upload-time = "2026-04-06T23:01:05.17Z" }, + { url = "https://files.pythonhosted.org/packages/d3/d6/ac63065d33dd700fee7ebd7d287332401b54e31b9346e142f871e1f0b116/cuda_pathfinder-1.5.3-py3-none-any.whl", hash = "sha256:dff021123aedbb4117cc7ec81717bbfe198fb4e8b5f1ee57e0e084fec5c8577d", size = 49991, upload-time = "2026-04-14T20:09:27.037Z" }, ] [[package]] @@ -651,11 +723,23 @@ wheels = [ [[package]] name = "filelock" -version = "3.25.2" +version = "3.28.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/94/b8/00651a0f559862f3bb7d6f7477b192afe3f583cc5e26403b44e59a55ab34/filelock-3.25.2.tar.gz", hash = "sha256:b64ece2b38f4ca29dd3e810287aa8c48182bbecd1ae6e9ae126c9b35f1382694", size = 40480, upload-time = "2026-03-11T20:45:38.487Z" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/17/6e8890271880903e3538660a21d63a6c1fea969ac71d0d6b608b78727fa9/filelock-3.28.0.tar.gz", hash = "sha256:4ed1010aae813c4ee8d9c660e4792475ee60c4a0ba76073ceaf862bd317e3ca6", size = 56474, upload-time = "2026-04-14T22:54:33.625Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3b/21/2f728888c45033d34a417bfcd248ea2564c9e08ab1bfd301377cf05d5586/filelock-3.28.0-py3-none-any.whl", hash = "sha256:de9af6712788e7171df1b28b15eba2446c69721433fa427a9bee07b17820a9db", size = 39189, upload-time = "2026-04-14T22:54:32.037Z" }, +] + +[[package]] +name = "fire" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "termcolor" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/00/f8d10588d2019d6d6452653def1ee807353b21983db48550318424b5ff18/fire-0.7.1.tar.gz", hash = "sha256:3b208f05c736de98fb343310d090dcc4d8c78b2a89ea4f32b837c586270a9cbf", size = 88720, upload-time = "2025-08-16T20:20:24.175Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a4/a5/842ae8f0c08b61d6484b52f99a03510a3a72d23141942d216ebe81fefbce/filelock-3.25.2-py3-none-any.whl", hash = "sha256:ca8afb0da15f229774c9ad1b455ed96e85a81373065fb10446672f64444ddf70", size = 26759, upload-time = "2026-03-11T20:45:37.437Z" }, + { url = "https://files.pythonhosted.org/packages/e5/4c/93d0f85318da65923e4b91c1c2ff03d8a458cbefebe3bc612a6693c7906d/fire-0.7.1-py3-none-any.whl", hash = "sha256:e43fd8a5033a9001e7e2973bab96070694b9f12f2e0ecf96d4683971b5ab1882", size = 115945, upload-time = "2025-08-16T20:20:22.87Z" }, ] [[package]] @@ -841,7 +925,7 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "1.10.1" +version = "1.10.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, @@ -854,9 +938,9 @@ dependencies = [ { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/e4/28/baf5d745559503ce8d28cf5bc9551f5ac59158eafd7b6a6afff0bcdb0f50/huggingface_hub-1.10.1.tar.gz", hash = "sha256:696c53cf9c2ac9befbfb5dd41d05392a031c69fc6930d1ed9671debd405b6fff", size = 758094, upload-time = "2026-04-09T15:01:18.928Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0c/4d/00734890c7fcfe2c7ff04f1c1a167186c42b19e370a2dd8cfd8c34fc92c4/huggingface_hub-1.10.2.tar.gz", hash = "sha256:4b276f820483b709dc86a53bcb8183ea496b8d8447c9f7f88a115a12b498a95f", size = 758428, upload-time = "2026-04-14T10:42:28.498Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/83/8c/c7a33f3efaa8d6a5bc40e012e5ecc2d72c2e6124550ca9085fe0ceed9993/huggingface_hub-1.10.1-py3-none-any.whl", hash = "sha256:6b981107a62fbe68c74374418983399c632e35786dcd14642a9f2972633c8b5a", size = 642630, upload-time = "2026-04-09T15:01:17.35Z" }, + { url = "https://files.pythonhosted.org/packages/5e/c9/4c1e1216b24bcab140c83acdf8bc89a846ea17cd8a06cd18e3fd308a297f/huggingface_hub-1.10.2-py3-none-any.whl", hash = "sha256:c26c908767cc711493978dc0b4f5747ba7841602997cc98bfd628450a28cf9bc", size = 642581, upload-time = "2026-04-14T10:42:26.563Z" }, ] [[package]] @@ -871,6 +955,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f0/0f/310fb31e39e2d734ccaa2c0fb981ee41f7bd5056ce9bc29b2248bd569169/humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477", size = 86794, upload-time = "2021-09-17T21:40:39.897Z" }, ] +[[package]] +name = "hydra-core" +version = "1.3.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime" }, + { name = "omegaconf" }, + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/8e/07e42bc434a847154083b315779b0a81d567154504624e181caf2c71cd98/hydra-core-1.3.2.tar.gz", hash = "sha256:8a878ed67216997c3e9d88a8e72e7b4767e81af37afb4ea3334b269a4390a824", size = 3263494, upload-time = "2023-02-23T18:33:43.03Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547, upload-time = "2023-02-23T18:33:40.801Z" }, +] + [[package]] name = "identify" version = "2.6.18" @@ -898,6 +996,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5f/53/fb7122b71361a0d121b669dcf3d31244ef75badbbb724af388948de543e2/imagesize-2.0.0-py2.py3-none-any.whl", hash = "sha256:5667c5bbb57ab3f1fa4bc366f4fbc971db3d5ed011fd2715fd8001f782718d96", size = 9441, upload-time = "2026-03-03T14:18:27.892Z" }, ] +[[package]] +name = "immutabledict" +version = "4.3.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1d/e6/718471048fea0366c3e3d1df3acfd914ca66d571cdffcf6d37bbcd725708/immutabledict-4.3.1.tar.gz", hash = "sha256:f844a669106cfdc73f47b1a9da003782fb17dc955a54c80972e0d93d1c63c514", size = 7806, upload-time = "2026-02-15T10:32:34.668Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a3/ce/f9018bf69ae91b273b6391a095e7c93fa5e1617f25b6ba81ad4b20c9df10/immutabledict-4.3.1-py3-none-any.whl", hash = "sha256:c9facdc0ff30fdb8e35bd16532026cac472a549e182c94fa201b51b25e4bf7bf", size = 5000, upload-time = "2026-02-15T10:32:33.672Z" }, +] + [[package]] name = "importlib-metadata" version = "9.0.0" @@ -991,16 +1098,83 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bf/bf/b7802b6578ca3a6506aaac6696ac1e8de500419fee3cd288184e82a8c2aa/lief-0.17.6-cp313-cp313-win_arm64.whl", hash = "sha256:6d4eb8adce400af52cc174ac5cbe40ab10b9df5824193975d12e2d4f85b298a3", size = 3461166, upload-time = "2026-03-18T06:58:32.975Z" }, ] +[[package]] +name = "lru-dict" +version = "1.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/06/0a/dec86efe38b350314c49a8d39ef01ba7cf8bbbef1d177646320eedea7159/lru_dict-1.4.1.tar.gz", hash = "sha256:cc518ff2d38cc7a8ab56f9a6ae557f91e2e1524b57ed8e598e97f45a2bd708fc", size = 13439, upload-time = "2025-11-02T10:02:13.548Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/6c/396716746ca46fd2ac52a7a6cbd7b4cf848e5d430f431dacd209290dfa71/lru_dict-1.4.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:3766e397aa6de1ca3442729bc1fa75834ab7b0a6b017e6e197d3a66b61abde59", size = 16757, upload-time = "2025-11-02T10:00:55.767Z" }, + { url = "https://files.pythonhosted.org/packages/2d/93/c163ffb71beb18f18459461658fd16c8b8c86aed858f2dc7c7e636318f61/lru_dict-1.4.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:658e152d3a4ad0e1d75e6f53b1fa353779539920b38be99f4ea33d3bad41efdb", size = 11243, upload-time = "2025-11-02T10:00:56.715Z" }, + { url = "https://files.pythonhosted.org/packages/44/e3/fa96d54032531c67eeacf0ab6f56e10e05f25d382a29f6a381ac8ecf3814/lru_dict-1.4.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:98af7044b5c3d85a649e1afb8891829ff5210caf9143acc741b3e98ab1b66ff6", size = 11726, upload-time = "2025-11-02T10:00:57.377Z" }, + { url = "https://files.pythonhosted.org/packages/7a/23/bae4f32fb014fd2dc5512e9267a3b1ec34c3b55d16a2202a1193d9ae635d/lru_dict-1.4.1-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:906d99705b79a00b5668bdb8782ad823ccc8d26e1fc6b56327ae469a8d12e9b4", size = 29823, upload-time = "2025-11-02T10:00:58.34Z" }, + { url = "https://files.pythonhosted.org/packages/9f/3b/8c3d1e6a188ce65e0161b86dbd18f2290950baf1e9e28e4948fc123d9a67/lru_dict-1.4.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:885643fd968336d8652fddb0778184e2eeff7b7aebced6de268af6d6caef42d5", size = 30812, upload-time = "2025-11-02T10:00:59.358Z" }, + { url = "https://files.pythonhosted.org/packages/ed/11/7f061507eda944150ed959e99a3700ce6358c1241c7f697b2f1ade48646b/lru_dict-1.4.1-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:24c779334bed82f1a7eb2d1ebcba2b7aa9a1555d40a3b53e05eb6b9dfcb0609c", size = 32480, upload-time = "2025-11-02T10:01:00.141Z" }, + { url = "https://files.pythonhosted.org/packages/75/e7/94ac30d33c6f8a8eca5d7e81c0ce26fb7b79b18ea65accdcb2a652b19abc/lru_dict-1.4.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:c6099e2ecb118dfeae4a197bfcc702ea5841bfd86f19d1b340e932d0f5c47c10", size = 30199, upload-time = "2025-11-02T10:01:01.31Z" }, + { url = "https://files.pythonhosted.org/packages/4a/81/c93ee7365db67dfb497e6218aa0395b9ec878c07c732d348bfbd651bcc95/lru_dict-1.4.1-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:4e0db4f3105108598749550e639b283b07df0bb91cac3b47e86ffebcab721cc7", size = 31489, upload-time = "2025-11-02T10:01:02.363Z" }, + { url = "https://files.pythonhosted.org/packages/9f/0b/634e8b4eca2497647f802bbe1ae3f0e1e9a0de1d555cf77c022527b2682f/lru_dict-1.4.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:e21f67ba374d1945051b547e719d44a8c7880718f67a15a03e7a12e1d12ea96b", size = 29522, upload-time = "2025-11-02T10:01:03.399Z" }, + { url = "https://files.pythonhosted.org/packages/de/cc/591b959d77cc0e0ac016f11baf26d03d566bb88a53fa9b41e157bc68bc4b/lru_dict-1.4.1-cp310-cp310-win32.whl", hash = "sha256:f309b4018dd41f33bf3bd4cc0f62421da8bcca513ea044dbb22f3cd029935012", size = 13066, upload-time = "2025-11-02T10:01:04.457Z" }, + { url = "https://files.pythonhosted.org/packages/d9/bc/c14b67fdbdb5a2a81cfb907ea8a8b0c9da5aed899f34921ebf097e22a966/lru_dict-1.4.1-cp310-cp310-win_amd64.whl", hash = "sha256:e84cd1065955897de01f1fb4cbd6f87cab7706e920283bb98c27341d76dd9a8d", size = 14008, upload-time = "2025-11-02T10:01:05.421Z" }, + { url = "https://files.pythonhosted.org/packages/4c/ff/1d02bc444174f07d3ce747568989969c97dc77d0513f4c3b8b6224cb976f/lru_dict-1.4.1-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:cc74c49cf1c26d6c28d8f6988cf0354696ca38a4f6012fa63055d2800791784b", size = 16760, upload-time = "2025-11-02T10:01:06.492Z" }, + { url = "https://files.pythonhosted.org/packages/0b/d8/e2e970272ea5fe7ba6349a5e7d0bb0fd814f5d1b88a53bc72b8c2a5e034f/lru_dict-1.4.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0158db85dfb2cd2fd2ddaa47709bdb073f814e0a8a149051b70b07e59ac83231", size = 11249, upload-time = "2025-11-02T10:01:07.261Z" }, + { url = "https://files.pythonhosted.org/packages/a5/26/860b5e60f339f8038118028388926224c8b70779e8243d68772e0e0d0ab3/lru_dict-1.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:c8ac5cfd56e036bd8d7199626147044485fa64a163a5bde96bfa5a1c7fea2273", size = 11728, upload-time = "2025-11-02T10:01:08.185Z" }, + { url = "https://files.pythonhosted.org/packages/61/55/fc8f71953fd343ede33810b0a000b4130e03635ae09b28569e45735ded2f/lru_dict-1.4.1-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:2eb2058cb7b329b4b72baee4cd1bb322af1feec73de79e68edb35d333c90b698", size = 30795, upload-time = "2025-11-02T10:01:08.862Z" }, + { url = "https://files.pythonhosted.org/packages/4c/26/ad549550e6a236818a91434570d38d7a93824b0410d3db1c845a53238e1f/lru_dict-1.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6ffbb6f3c1e906e92d9129c14a88d81358be1e0b60195c1729b215a52e9670de", size = 31807, upload-time = "2025-11-02T10:01:09.581Z" }, + { url = "https://files.pythonhosted.org/packages/7c/39/72dae9ac0e95a8576a45e3bd62a6fc3e7dbb116794efa1337c7b450d4836/lru_dict-1.4.1-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:11b289d78a48a086846e46d2275707d33523f5d543475336c29c56fd5d0e65dc", size = 33437, upload-time = "2025-11-02T10:01:10.676Z" }, + { url = "https://files.pythonhosted.org/packages/a8/46/221479834703a5397fa32f07212ace38f104a31ad1af8a921cf25e053677/lru_dict-1.4.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:3fe10c1f45712e191eecb2a69604d566c64ddfe01136fd467c890ed558c3ad40", size = 31168, upload-time = "2025-11-02T10:01:11.47Z" }, + { url = "https://files.pythonhosted.org/packages/6e/13/98d36e2522fda7f6625c15332562f81f1465161a5ae021d9b3b408f8c427/lru_dict-1.4.1-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:e04820e3473bd7f55440f24c946ca4335e392d5e3e0e1e948020e94cd1954372", size = 32454, upload-time = "2025-11-02T10:01:12.522Z" }, + { url = "https://files.pythonhosted.org/packages/49/18/345ff2a98d27cddae40c84cf0466fcc329f3965cd21322bb561a94e4d332/lru_dict-1.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:edc004c88911a8f9715e716116d2520c13db89afd6c37cc0f28042ba10635163", size = 30574, upload-time = "2025-11-02T10:01:13.293Z" }, + { url = "https://files.pythonhosted.org/packages/d7/92/dfea71402a7ca46332bcb854827ee68bbc9be205e2558c3a40293eca9782/lru_dict-1.4.1-cp311-cp311-win32.whl", hash = "sha256:b0b5360264b37676c405ea0a560744d7dcb2d47adff1e7837113c15fabcc7a71", size = 13031, upload-time = "2025-11-02T10:01:13.96Z" }, + { url = "https://files.pythonhosted.org/packages/3a/7b/4c7d566d77ec3ad9128f07407494c2aec57909f8dd59f0c9910bd4c05840/lru_dict-1.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:bb4b37daad9fe4e796c462f4876cf34e52564630902bdf59a271bc482b48a361", size = 14007, upload-time = "2025-11-02T10:01:14.857Z" }, + { url = "https://files.pythonhosted.org/packages/4f/a8/89e4c26e0e751321b41b0a3007384f97d9eae7a863c49af1c68c43005ca3/lru_dict-1.4.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:7fa342c6e6bc811ee6a17eb569d37b149340d5aa5a637a53438e316a95783838", size = 16683, upload-time = "2025-11-02T10:01:15.891Z" }, + { url = "https://files.pythonhosted.org/packages/f1/34/b3c6fdd120af68b6eeb524d0de3293ff27918ec57f45eed6bef1789fd085/lru_dict-1.4.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:bd86bd202a7c1585d9dc7e5b0c3d52cf76dc56b261b4bbecfeefbbae31a5c97d", size = 11216, upload-time = "2025-11-02T10:01:16.867Z" }, + { url = "https://files.pythonhosted.org/packages/e9/7e/280267ae23f1ec1074ddaab787c5e041e090220e8e37828d51ff4e681dfd/lru_dict-1.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4617554f3e42a8f520c8494842c23b98f5b7f4d5e0410e91a4c3ad0ea5f7e094", size = 11687, upload-time = "2025-11-02T10:01:17.485Z" }, + { url = "https://files.pythonhosted.org/packages/ca/18/fec42416ceff98ae2760067ec72b0b9fc02840e729bbc18059c6a02cb01f/lru_dict-1.4.1-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:40927a6a4284d437047f547e652b15f6f0f40210deb6b9e5b77e556ff0faea0f", size = 31960, upload-time = "2025-11-02T10:01:18.158Z" }, + { url = "https://files.pythonhosted.org/packages/c2/ef/38e7ee1a5d32b9b1629d045fa5a495375383aacfb2945f4d9535b9af9630/lru_dict-1.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e2c07ecb6d42494e45d00c2541e6b0ae7659fc3cf89681521ba94b15c682d4fe", size = 32882, upload-time = "2025-11-02T10:01:18.841Z" }, + { url = "https://files.pythonhosted.org/packages/72/82/d56653ca144c291ab37bea5f23c5078ffbe64f7f5b466f91d400590b9106/lru_dict-1.4.1-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:85b28aa2de7c5f1f6c68221857accd084438df98edbd4f57595795734225770c", size = 34268, upload-time = "2025-11-02T10:01:19.564Z" }, + { url = "https://files.pythonhosted.org/packages/94/ae/382651533d60f0b598757efda56dc87cad5ac311fba8e61f86fb916bf236/lru_dict-1.4.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:cbbbb4b51e2529ccf7ee8a3c3b834052dbd54871a216cfd229dd2b1194ff293a", size = 32156, upload-time = "2025-11-02T10:01:20.22Z" }, + { url = "https://files.pythonhosted.org/packages/aa/d1/d9df7e9272ccbc96f04c477dfb9abb91fa8fabde86b7fa190cb7b3c7a024/lru_dict-1.4.1-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:e47040421a13de8bc6404557b3700c33f1f2683cbcce22fe5cacec4c938ce54b", size = 33395, upload-time = "2025-11-02T10:01:20.901Z" }, + { url = "https://files.pythonhosted.org/packages/e9/6e/dafe0f5943a7b3ab24d3429032ff85873acd626087934b8161b55340c13a/lru_dict-1.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:451f7249866cb9564bb40d73bec7ac865574dafd0a4cc91627bbf35be7e99291", size = 31591, upload-time = "2025-11-02T10:01:21.606Z" }, + { url = "https://files.pythonhosted.org/packages/a6/4d/9dd35444592bfb6805548e15971cfce821400966a51130b78dc021ee8f03/lru_dict-1.4.1-cp312-cp312-win32.whl", hash = "sha256:e8996f3f94870ecb236c55d280839390edae7f201858fee770267eac27b8b47d", size = 13119, upload-time = "2025-11-02T10:01:22.61Z" }, + { url = "https://files.pythonhosted.org/packages/8d/82/7e72e30d6c15d65466b3baca87cce15e20848ba6a488868aa54e901141a6/lru_dict-1.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:d90774db1b60c0d5c829cfa5d7fda6db96ed1519296f626575598f9f170cca37", size = 14109, upload-time = "2025-11-02T10:01:23.322Z" }, + { url = "https://files.pythonhosted.org/packages/85/95/ee171a68ae381ab988c50e3b7b136b1c598f5f683ba4a1e10c51e2480408/lru_dict-1.4.1-cp313-cp313-android_21_arm64_v8a.whl", hash = "sha256:2a5644bb1db0514abdad5e2f3d8f1beb6f7560c8cceb62079c40a4269de34b3c", size = 12248, upload-time = "2025-11-02T10:01:24.291Z" }, + { url = "https://files.pythonhosted.org/packages/a1/82/8de8e8fd96c44d46891415834ceb9f51c552840bda2d118394aca5e3153a/lru_dict-1.4.1-cp313-cp313-android_21_x86_64.whl", hash = "sha256:4209864be09ec20f6059fef8544697eb3d3729d63a983bf66457054bf3e40601", size = 12243, upload-time = "2025-11-02T10:01:25.254Z" }, + { url = "https://files.pythonhosted.org/packages/53/97/251cfb357c547a8fd06c2bc40db8a7f7eed7dbacef30d8d7e543522360e1/lru_dict-1.4.1-cp313-cp313-ios_13_0_arm64_iphoneos.whl", hash = "sha256:8fef8dd72484b4280799c502c116acfdfcf0dedf3508bc9d0d19e684a6a23267", size = 10938, upload-time = "2025-11-02T10:01:25.89Z" }, + { url = "https://files.pythonhosted.org/packages/58/14/602791d219bc87197ae80f5fa0f77ca0af8e83e9a06c7cdb89db5575839e/lru_dict-1.4.1-cp313-cp313-ios_13_0_arm64_iphonesimulator.whl", hash = "sha256:d64ddbe4c426fdc4cfc1abaea71d587d439397386a7b35d588f4fd64b695a83d", size = 11261, upload-time = "2025-11-02T10:01:26.879Z" }, + { url = "https://files.pythonhosted.org/packages/10/5d/a30a6fad150f20f084de8e243882a0488ad4929db41a2c8ce9be6cf56563/lru_dict-1.4.1-cp313-cp313-ios_13_0_x86_64_iphonesimulator.whl", hash = "sha256:000ba9a2ab4dd1ad2d91764a6d5cce75a59de51534cdda478d1ddaa3cd8d5c48", size = 10801, upload-time = "2025-11-02T10:01:27.553Z" }, + { url = "https://files.pythonhosted.org/packages/64/4d/cee327e024d42972c598b7e0cd5063a1b1d7451efba31f7de7b6ca91e7d0/lru_dict-1.4.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ffad2758ce21d8fd6f0ae2628b31330732db8429a4b5994d2e107bed0ee11e68", size = 16688, upload-time = "2025-11-02T10:01:28.17Z" }, + { url = "https://files.pythonhosted.org/packages/58/c8/2f86a1e448c5257b31424b96bf1385e7f96ec7841c2376db02811bbd395f/lru_dict-1.4.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1671e8d92fe35dfb38d3505a56338792d3e225032f8e94888b6e95b323120380", size = 11214, upload-time = "2025-11-02T10:01:28.888Z" }, + { url = "https://files.pythonhosted.org/packages/06/41/507c615cffaba67c35affd77dec25d3183bb87f404b41c8bb2b3053481ac/lru_dict-1.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:d5f01ada0cf0c1aa2bdc684e5ac0f6548be7eccc3ce8b4c0361db8445f867f04", size = 11689, upload-time = "2025-11-02T10:01:29.508Z" }, + { url = "https://files.pythonhosted.org/packages/4e/c1/35aa1359f80174016b389f8be5fd48c4a5af0a04a73afb4906e5d4279f4a/lru_dict-1.4.1-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:74204239e30b8ec7976257c5b64565d7e3e8aea0cad0dd50a9b99e171aaf3898", size = 32034, upload-time = "2025-11-02T10:01:30.164Z" }, + { url = "https://files.pythonhosted.org/packages/c2/93/46301015bddd4552a1b76982ef788a7fb2a886efff83ad2c178cc7e68349/lru_dict-1.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:a7da0e451faa4d6dcae21c0f2527c540000b2f23ed8326a0bc1d870130fd12b1", size = 32919, upload-time = "2025-11-02T10:01:30.971Z" }, + { url = "https://files.pythonhosted.org/packages/e6/cb/6d67145619d8ec3bba15fe145ff702ecf44991e33345d38c763501c1608a/lru_dict-1.4.1-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:071468a716768a9afca64659c390c1abb6d937b1897e07a0b70383f75637fce0", size = 34334, upload-time = "2025-11-02T10:01:31.663Z" }, + { url = "https://files.pythonhosted.org/packages/a5/44/50daaec6793ec2042079ed6a8b6a687b4be51b270b1d8ec5efd280116493/lru_dict-1.4.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:e77d209bcd396eb236c197bf4c95fab6848c61e0c1a5031cdde7f5c787e209f4", size = 32211, upload-time = "2025-11-02T10:01:32.468Z" }, + { url = "https://files.pythonhosted.org/packages/bd/53/355397949215e6b77b6771b973ee1dbc21fdd9f955925e47dce50d9d4727/lru_dict-1.4.1-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:b21688fd7ece56d04c0c13b42fd9f904d46fc9ff21e3de87d98f3f5a14c67f74", size = 33461, upload-time = "2025-11-02T10:01:33.2Z" }, + { url = "https://files.pythonhosted.org/packages/ff/fa/d660fa63144f38a0fd5b437a140517e3cff482d955ef6b9b4cf7651b9d85/lru_dict-1.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:989ef7352b347c82e5d5047f3b7ddf34b5a938e3f7b08775cacc9f28e97dd2a8", size = 31651, upload-time = "2025-11-02T10:01:33.874Z" }, + { url = "https://files.pythonhosted.org/packages/2e/77/0fae8d0702f7546f436efe06a684b301aad5c8a167bb2df6e42b0f821de5/lru_dict-1.4.1-cp313-cp313-win32.whl", hash = "sha256:a36e6e95b5d474ef90d04a5e3ad81ca362b473ec9534ed964222f3c0444138b8", size = 13120, upload-time = "2025-11-02T10:01:34.595Z" }, + { url = "https://files.pythonhosted.org/packages/4a/20/56a3f0d74c8fe32c01d3978387f66c9fb180c7f15bfd9fcecaa01b4e7736/lru_dict-1.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:8e73a1ec2d0f476d666ce7c91464b22854086951b319544d1850c508f5ce381f", size = 14112, upload-time = "2025-11-02T10:01:35.268Z" }, + { url = "https://files.pythonhosted.org/packages/ec/de/18ac3957e1aa6674a0a828748c819265f79b524ff30cbb0ac7f08ab786c8/lru_dict-1.4.1-pp310-pypy310_pp73-macosx_10_15_x86_64.whl", hash = "sha256:cc9dd191870555624bbf3903c8afa3f01815ca3256ed8b35cb323f0db3ce4f98", size = 10467, upload-time = "2025-11-02T10:02:05.717Z" }, + { url = "https://files.pythonhosted.org/packages/0c/53/2a0bedaa64950cc56ade72e2f5a292318473585d9a3adc797d13b38082e7/lru_dict-1.4.1-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:afdf92b332632aa6e4b8646e93723f50f41fece2a80a54d2b44e8ac67f913ceb", size = 10871, upload-time = "2025-11-02T10:02:06.353Z" }, + { url = "https://files.pythonhosted.org/packages/4e/e2/d5ea49d62ea142559fd9cafd8505d4a4f87a1d81953a9c99fa61e7ccbd6b/lru_dict-1.4.1-pp310-pypy310_pp73-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3d6770adafae25663b682420891a10a5894595f02b1e4d87766f7adc8e56e72a", size = 12969, upload-time = "2025-11-02T10:02:07.196Z" }, + { url = "https://files.pythonhosted.org/packages/a2/67/0672caac9a04dc9011f7a27fc2ec2003f0bfa008070b29940d05b4dae56a/lru_dict-1.4.1-pp310-pypy310_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:018cd3b41224ca81eb83cdf6db024409a920e5c1d3ce4e8b323cb66e24a73132", size = 13959, upload-time = "2025-11-02T10:02:08.267Z" }, + { url = "https://files.pythonhosted.org/packages/e3/7e/313385214a5011cf9fe8376928f66f70bfedc48d8f7ab424292224ed4907/lru_dict-1.4.1-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:781dbcf0c83160e525482a4ebcd7c5065851a6c7295f1cda78248a2029f23f39", size = 14084, upload-time = "2025-11-02T10:02:08.993Z" }, + { url = "https://files.pythonhosted.org/packages/8e/47/08c61cad038706b3a89b8c7587ec74ed9731c1e536329745cccb6c840916/lru_dict-1.4.1-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:9219f13e4101c064f70e1815d7c51f9be9e053983e74dfb7bcfdf92f5fcbb0e0", size = 10384, upload-time = "2025-11-02T10:02:09.656Z" }, + { url = "https://files.pythonhosted.org/packages/6b/a1/022c4d7c68c076370231488c97cf7451131fb9ca0d60d1b2785e7baa1f5b/lru_dict-1.4.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:b7e1ac7fb6e91e4d3212e153f9e2d98d163a4439b9bf9df247c22519262c26fe", size = 10822, upload-time = "2025-11-02T10:02:10.609Z" }, + { url = "https://files.pythonhosted.org/packages/65/b4/4c0a0877a77fececa9f58d804569e2aac1bfbe588e3a70e79647b5d8f7d4/lru_dict-1.4.1-pp311-pypy311_pp73-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:23424321b761c43f3021a596565f8205ecec0e175822e7a5d9b2a175578aa7de", size = 12968, upload-time = "2025-11-02T10:02:11.405Z" }, + { url = "https://files.pythonhosted.org/packages/22/06/d7e393d07dc31e656330d5a058f34e972bf590e7dc882922b426f3aec4a0/lru_dict-1.4.1-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:804ee76f98afc3d50e9a2e9c835a6820877aa6391f2add520a57f86b3f55ec3a", size = 13904, upload-time = "2025-11-02T10:02:12.144Z" }, + { url = "https://files.pythonhosted.org/packages/e8/1e/0eee8bcc16bf01b265ac83e4b870596e2f3bcc40d88aa7ec25407180fe44/lru_dict-1.4.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:3be24e24c8998302ea1c28f997505fa6843f507aad3c7d5c3a82cc01c5c11be4", size = 14062, upload-time = "2025-11-02T10:02:12.878Z" }, +] + [[package]] name = "mako" -version = "1.3.10" +version = "1.3.11" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markupsafe" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9e/38/bd5b78a920a64d708fe6bc8e0a2c075e1389d53bef8413725c63ba041535/mako-1.3.10.tar.gz", hash = "sha256:99579a6f39583fa7e5630a28c3c1f440e4e97a414b80372649c0ce338da2ea28", size = 392474, upload-time = "2025-04-10T12:44:31.16Z" } +sdist = { url = "https://files.pythonhosted.org/packages/59/8a/805404d0c0b9f3d7a326475ca008db57aea9c5c9f2e1e39ed0faa335571c/mako-1.3.11.tar.gz", hash = "sha256:071eb4ab4c5010443152255d77db7faa6ce5916f35226eb02dc34479b6858069", size = 399811, upload-time = "2026-04-14T20:19:51.493Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/87/fb/99f81ac72ae23375f22b7afdb7642aba97c00a713c217124420147681a2f/mako-1.3.10-py3-none-any.whl", hash = "sha256:baef24a52fc4fc514a0887ac600f9f1cff3d82c61d4d700a1fa84d597b88db59", size = 78509, upload-time = "2025-04-10T12:50:53.297Z" }, + { url = "https://files.pythonhosted.org/packages/68/a5/19d7aaa7e433713ffe881df33705925a196afb9532efc8475d26593921a6/mako-1.3.11-py3-none-any.whl", hash = "sha256:e372c6e333cf004aa736a15f425087ec977e1fcbd2966aae7f17c8dc1da27a77", size = 78503, upload-time = "2026-04-14T20:19:53.233Z" }, ] [[package]] @@ -1087,6 +1261,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] +[[package]] +name = "mip" +version = "1.17.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cbcbox" }, + { name = "cffi" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/96/15/7496ce00eafc5b99e98bb696588527754d7629f11ac6208fdc9355f6eec6/mip-1.17.6.tar.gz", hash = "sha256:9da8074b80bd3ef788513d5a214ef832916d82aa66487da11a49c7da9f89d270", size = 9443716, upload-time = "2026-03-23T16:20:09.009Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/1d/0eb2531e779be687a3249e7f2ba5e465dd34611072390252e5fcf89e7fef/mip-1.17.6-py3-none-any.whl", hash = "sha256:4fb7ff5d7beacbe7d007de172306a76aff8fb36760fa63a5d88a8cb3aa6f4a57", size = 88214, upload-time = "2026-03-23T16:20:07.289Z" }, +] + [[package]] name = "ml-dtypes" version = "0.5.4" @@ -1604,8 +1791,13 @@ all = [ { name = "datasets" }, { name = "deepspeed", marker = "sys_platform != 'darwin' and sys_platform != 'win32'" }, { name = "diffusers" }, + { name = "fire" }, { name = "huggingface-hub" }, + { name = "hydra-core" }, + { name = "immutabledict" }, { name = "lief" }, + { name = "lru-dict" }, + { name = "mip" }, { name = "ml-dtypes" }, { name = "nltk" }, { name = "onnx" }, @@ -1617,11 +1809,14 @@ all = [ { name = "onnxruntime-gpu", version = "1.24.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32'" }, { name = "onnxscript" }, { name = "onnxslim" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pandas", version = "3.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "peft" }, { name = "polygraphy" }, { name = "sentencepiece" }, { name = "tiktoken" }, { name = "transformers" }, + { name = "typeguard" }, { name = "wonderwords" }, ] dev = [ @@ -1634,8 +1829,13 @@ dev = [ { name = "datasets" }, { name = "deepspeed", marker = "sys_platform != 'darwin' and sys_platform != 'win32'" }, { name = "diffusers" }, + { name = "fire" }, { name = "huggingface-hub" }, + { name = "hydra-core" }, + { name = "immutabledict" }, { name = "lief" }, + { name = "lru-dict" }, + { name = "mip" }, { name = "ml-dtypes" }, { name = "mypy" }, { name = "nltk" }, @@ -1648,6 +1848,8 @@ dev = [ { name = "onnxruntime-gpu", version = "1.24.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and platform_machine != 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32'" }, { name = "onnxscript" }, { name = "onnxslim" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pandas", version = "3.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "peft" }, { name = "polygraphy" }, { name = "pre-commit" }, @@ -1673,6 +1875,7 @@ dev = [ { name = "tox" }, { name = "tox-current-env" }, { name = "transformers" }, + { name = "typeguard" }, { name = "wonderwords" }, ] dev-docs = [ @@ -1734,6 +1937,16 @@ onnx = [ { name = "onnxslim" }, { name = "polygraphy" }, ] +puzzletron = [ + { name = "fire" }, + { name = "hydra-core" }, + { name = "immutabledict" }, + { name = "lru-dict" }, + { name = "mip" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pandas", version = "3.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "typeguard" }, +] [package.metadata] requires-dist = [ @@ -1746,8 +1959,13 @@ requires-dist = [ { name = "datasets", marker = "extra == 'hf'", specifier = ">=3.0.0" }, { name = "deepspeed", marker = "sys_platform != 'darwin' and sys_platform != 'win32' and extra == 'hf'", specifier = ">=0.9.6" }, { name = "diffusers", marker = "extra == 'hf'", specifier = ">=0.32.2" }, + { name = "fire", marker = "extra == 'puzzletron'" }, { name = "huggingface-hub", marker = "extra == 'hf'", specifier = ">=0.24.0" }, + { name = "hydra-core", marker = "extra == 'puzzletron'", specifier = "==1.3.2" }, + { name = "immutabledict", marker = "extra == 'puzzletron'" }, { name = "lief", marker = "extra == 'onnx'" }, + { name = "lru-dict", marker = "extra == 'puzzletron'" }, + { name = "mip", marker = "extra == 'puzzletron'" }, { name = "ml-dtypes", marker = "extra == 'onnx'" }, { name = "mypy", marker = "extra == 'dev-lint'", specifier = "==1.17.1" }, { name = "ninja" }, @@ -1755,7 +1973,7 @@ requires-dist = [ { name = "numpy" }, { name = "nvidia-ml-py", specifier = ">=12" }, { name = "nvidia-modelopt", extras = ["all", "dev-docs", "dev-lint", "dev-test"], marker = "extra == 'dev'" }, - { name = "nvidia-modelopt", extras = ["hf", "onnx"], marker = "extra == 'all'" }, + { name = "nvidia-modelopt", extras = ["hf", "onnx", "puzzletron"], marker = "extra == 'all'" }, { name = "omegaconf", specifier = ">=2.3.0" }, { name = "onnx", marker = "extra == 'onnx'", specifier = "~=1.21.0" }, { name = "onnx-graphsurgeon", marker = "extra == 'onnx'", specifier = ">=0.6.1" }, @@ -1768,6 +1986,7 @@ requires-dist = [ { name = "onnxscript", marker = "extra == 'onnx'" }, { name = "onnxslim", marker = "extra == 'onnx'", specifier = ">=0.1.76" }, { name = "packaging" }, + { name = "pandas", marker = "extra == 'puzzletron'" }, { name = "peft", marker = "extra == 'hf'", specifier = ">=0.17.0" }, { name = "polygraphy", marker = "extra == 'onnx'", specifier = ">=0.49.22" }, { name = "pre-commit", marker = "extra == 'dev-lint'", specifier = "==4.3.0" }, @@ -1802,9 +2021,10 @@ requires-dist = [ { name = "tox-current-env", marker = "extra == 'dev-test'", specifier = ">=0.0.12" }, { name = "tqdm" }, { name = "transformers", marker = "extra == 'hf'", specifier = ">=4.56" }, + { name = "typeguard", marker = "extra == 'puzzletron'" }, { name = "wonderwords", marker = "extra == 'hf'" }, ] -provides-extras = ["onnx", "hf", "dev-lint", "dev-docs", "dev-test", "all", "dev"] +provides-extras = ["onnx", "hf", "puzzletron", "dev-lint", "dev-docs", "dev-test", "all", "dev"] [[package]] name = "omegaconf" @@ -2069,11 +2289,11 @@ wheels = [ [[package]] name = "packaging" -version = "26.0" +version = "26.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/65/ee/299d360cdc32edc7d2cf530f3accf79c4fca01e96ffc950d8a52213bd8e4/packaging-26.0.tar.gz", hash = "sha256:00243ae351a257117b6a241061796684b084ed1c516a08c48a3f7e147a9d80b4", size = 143416, upload-time = "2026-01-21T20:50:39.064Z" } +sdist = { url = "https://files.pythonhosted.org/packages/df/de/0d2b39fb4af88a0258f3bac87dfcbb48e73fbdea4a2ed0e2213f9a4c2f9a/packaging-26.1.tar.gz", hash = "sha256:f042152b681c4bfac5cae2742a55e103d27ab2ec0f3d88037136b6bfe7c9c5de", size = 215519, upload-time = "2026-04-14T21:12:49.362Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" }, + { url = "https://files.pythonhosted.org/packages/7a/c2/920ef838e2f0028c8262f16101ec09ebd5969864e5a64c4c05fad0617c56/packaging-26.1-py3-none-any.whl", hash = "sha256:5d9c0669c6285e491e0ced2eee587eaf67b670d94a19e94e3984a481aba6802f", size = 95831, upload-time = "2026-04-14T21:12:47.56Z" }, ] [[package]] @@ -2211,7 +2431,7 @@ wheels = [ [[package]] name = "peft" -version = "0.18.1" +version = "0.19.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "accelerate" }, @@ -2226,9 +2446,9 @@ dependencies = [ { name = "tqdm" }, { name = "transformers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d8/48/147b3ea999560b40a34fd78724c7777aa9d18409c2250bdcaf9c4f2db7fc/peft-0.18.1.tar.gz", hash = "sha256:2dd0d6bfce936d1850e48aaddbd250941c5c02fc8ef3237cd8fd5aac35e0bae2", size = 635030, upload-time = "2026-01-09T13:08:01.136Z" } +sdist = { url = "https://files.pythonhosted.org/packages/2a/58/2e758e0794daa49dd9a47c56da7e31ee16d66d09ec787bbe44b1486f84fd/peft-0.19.0.tar.gz", hash = "sha256:a2917070092184a462093443029bc4f9292a91b9b99880488e319309ff0a172d", size = 762553, upload-time = "2026-04-14T14:01:53.189Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b3/14/b4e3f574acf349ae6f61f9c000a77f97a3b315b4bb6ad03791e79ae4a568/peft-0.18.1-py3-none-any.whl", hash = "sha256:0bf06847a3551e3019fc58c440cffc9a6b73e6e2962c95b52e224f77bbdb50f1", size = 556960, upload-time = "2026-01-09T13:07:55.865Z" }, + { url = "https://files.pythonhosted.org/packages/1f/fd/99e9beed55de9d54f6cd038880b4db4bacb45371dd54d40ea3fda7eaff5e/peft-0.19.0-py3-none-any.whl", hash = "sha256:7feca0f07bee9101807c7fd4601353d91161ea9e1f450150ee7859b2354c7690", size = 680671, upload-time = "2026-04-14T14:01:51.279Z" }, ] [[package]] @@ -2537,9 +2757,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/87/99f21e9b20899d6dc1bf7544cfe53e5fa17acc21bb267971a540425357d3/pybind11-3.0.3-py3-none-any.whl", hash = "sha256:fb5f8e4a64946b4dcc0451c83a8c384f803bc0a62dd1ba02f199e97dbc9aad4c", size = 313717, upload-time = "2026-03-31T23:42:04.814Z" }, ] +[[package]] +name = "pycparser" +version = "3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/7d/92392ff7815c21062bea51aa7b87d45576f649f16458d78b7cf94b9ab2e6/pycparser-3.0.tar.gz", hash = "sha256:600f49d217304a5902ac3c37e1281c9fe94e4d0489de643a9504c5cdfdfc6b29", size = 103492, upload-time = "2026-01-21T14:26:51.89Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/c3/44f3fbbfa403ea2a7c779186dc20772604442dde72947e7d01069cbe98e3/pycparser-3.0-py3-none-any.whl", hash = "sha256:b727414169a36b7d524c1c3e31839a521725078d7b2ff038656844266160a992", size = 48172, upload-time = "2026-01-21T14:26:50.693Z" }, +] + [[package]] name = "pydantic" -version = "2.12.5" +version = "2.13.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-types" }, @@ -2547,99 +2776,95 @@ dependencies = [ { name = "typing-extensions" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/69/44/36f1a6e523abc58ae5f928898e4aca2e0ea509b5aa6f6f392a5d882be928/pydantic-2.12.5.tar.gz", hash = "sha256:4d351024c75c0f085a9febbb665ce8c0c6ec5d30e903bdb6394b7ede26aebb49", size = 821591, upload-time = "2025-11-26T15:11:46.471Z" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/6b/1353beb3d1cd5cf61cdec5b6f87a9872399de3bc5cae0b7ce07ff4de2ab0/pydantic-2.13.1.tar.gz", hash = "sha256:a0f829b279ddd1e39291133fe2539d2aa46cc6b150c1706a270ff0879e3774d2", size = 843746, upload-time = "2026-04-15T14:57:19.398Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5a/87/b70ad306ebb6f9b585f114d0ac2137d792b48be34d732d60e597c2f8465a/pydantic-2.12.5-py3-none-any.whl", hash = "sha256:e561593fccf61e8a20fc46dfc2dfe075b8be7d0188df33f221ad1f0139180f9d", size = 463580, upload-time = "2025-11-26T15:11:44.605Z" }, + { url = "https://files.pythonhosted.org/packages/81/5a/2225f4c176dbfed0d809e848b50ef08f70e61daa667b7fa14b0d311ae44d/pydantic-2.13.1-py3-none-any.whl", hash = "sha256:9557ecc2806faaf6037f85b1fbd963d01e30511c48085f0d573650fdeaad378a", size = 471917, upload-time = "2026-04-15T14:57:17.277Z" }, ] [[package]] name = "pydantic-core" -version = "2.41.5" +version = "2.46.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/71/70/23b021c950c2addd24ec408e9ab05d59b035b39d97cdc1130e1bce647bb6/pydantic_core-2.41.5.tar.gz", hash = "sha256:08daa51ea16ad373ffd5e7606252cc32f07bc72b28284b6bc9c6df804816476e", size = 460952, upload-time = "2025-11-04T13:43:49.098Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c6/90/32c9941e728d564b411d574d8ee0cf09b12ec978cb22b294995bae5549a5/pydantic_core-2.41.5-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:77b63866ca88d804225eaa4af3e664c5faf3568cea95360d21f4725ab6e07146", size = 2107298, upload-time = "2025-11-04T13:39:04.116Z" }, - { url = "https://files.pythonhosted.org/packages/fb/a8/61c96a77fe28993d9a6fb0f4127e05430a267b235a124545d79fea46dd65/pydantic_core-2.41.5-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:dfa8a0c812ac681395907e71e1274819dec685fec28273a28905df579ef137e2", size = 1901475, upload-time = "2025-11-04T13:39:06.055Z" }, - { url = "https://files.pythonhosted.org/packages/5d/b6/338abf60225acc18cdc08b4faef592d0310923d19a87fba1faf05af5346e/pydantic_core-2.41.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5921a4d3ca3aee735d9fd163808f5e8dd6c6972101e4adbda9a4667908849b97", size = 1918815, upload-time = "2025-11-04T13:39:10.41Z" }, - { url = "https://files.pythonhosted.org/packages/d1/1c/2ed0433e682983d8e8cba9c8d8ef274d4791ec6a6f24c58935b90e780e0a/pydantic_core-2.41.5-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e25c479382d26a2a41b7ebea1043564a937db462816ea07afa8a44c0866d52f9", size = 2065567, upload-time = "2025-11-04T13:39:12.244Z" }, - { url = "https://files.pythonhosted.org/packages/b3/24/cf84974ee7d6eae06b9e63289b7b8f6549d416b5c199ca2d7ce13bbcf619/pydantic_core-2.41.5-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f547144f2966e1e16ae626d8ce72b4cfa0caedc7fa28052001c94fb2fcaa1c52", size = 2230442, upload-time = "2025-11-04T13:39:13.962Z" }, - { url = "https://files.pythonhosted.org/packages/fd/21/4e287865504b3edc0136c89c9c09431be326168b1eb7841911cbc877a995/pydantic_core-2.41.5-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:6f52298fbd394f9ed112d56f3d11aabd0d5bd27beb3084cc3d8ad069483b8941", size = 2350956, upload-time = "2025-11-04T13:39:15.889Z" }, - { url = "https://files.pythonhosted.org/packages/a8/76/7727ef2ffa4b62fcab916686a68a0426b9b790139720e1934e8ba797e238/pydantic_core-2.41.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:100baa204bb412b74fe285fb0f3a385256dad1d1879f0a5cb1499ed2e83d132a", size = 2068253, upload-time = "2025-11-04T13:39:17.403Z" }, - { url = "https://files.pythonhosted.org/packages/d5/8c/a4abfc79604bcb4c748e18975c44f94f756f08fb04218d5cb87eb0d3a63e/pydantic_core-2.41.5-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:05a2c8852530ad2812cb7914dc61a1125dc4e06252ee98e5638a12da6cc6fb6c", size = 2177050, upload-time = "2025-11-04T13:39:19.351Z" }, - { url = "https://files.pythonhosted.org/packages/67/b1/de2e9a9a79b480f9cb0b6e8b6ba4c50b18d4e89852426364c66aa82bb7b3/pydantic_core-2.41.5-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:29452c56df2ed968d18d7e21f4ab0ac55e71dc59524872f6fc57dcf4a3249ed2", size = 2147178, upload-time = "2025-11-04T13:39:21Z" }, - { url = "https://files.pythonhosted.org/packages/16/c1/dfb33f837a47b20417500efaa0378adc6635b3c79e8369ff7a03c494b4ac/pydantic_core-2.41.5-cp310-cp310-musllinux_1_1_armv7l.whl", hash = "sha256:d5160812ea7a8a2ffbe233d8da666880cad0cbaf5d4de74ae15c313213d62556", size = 2341833, upload-time = "2025-11-04T13:39:22.606Z" }, - { url = "https://files.pythonhosted.org/packages/47/36/00f398642a0f4b815a9a558c4f1dca1b4020a7d49562807d7bc9ff279a6c/pydantic_core-2.41.5-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:df3959765b553b9440adfd3c795617c352154e497a4eaf3752555cfb5da8fc49", size = 2321156, upload-time = "2025-11-04T13:39:25.843Z" }, - { url = "https://files.pythonhosted.org/packages/7e/70/cad3acd89fde2010807354d978725ae111ddf6d0ea46d1ea1775b5c1bd0c/pydantic_core-2.41.5-cp310-cp310-win32.whl", hash = "sha256:1f8d33a7f4d5a7889e60dc39856d76d09333d8a6ed0f5f1190635cbec70ec4ba", size = 1989378, upload-time = "2025-11-04T13:39:27.92Z" }, - { url = "https://files.pythonhosted.org/packages/76/92/d338652464c6c367e5608e4488201702cd1cbb0f33f7b6a85a60fe5f3720/pydantic_core-2.41.5-cp310-cp310-win_amd64.whl", hash = "sha256:62de39db01b8d593e45871af2af9e497295db8d73b085f6bfd0b18c83c70a8f9", size = 2013622, upload-time = "2025-11-04T13:39:29.848Z" }, - { url = "https://files.pythonhosted.org/packages/e8/72/74a989dd9f2084b3d9530b0915fdda64ac48831c30dbf7c72a41a5232db8/pydantic_core-2.41.5-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:a3a52f6156e73e7ccb0f8cced536adccb7042be67cb45f9562e12b319c119da6", size = 2105873, upload-time = "2025-11-04T13:39:31.373Z" }, - { url = "https://files.pythonhosted.org/packages/12/44/37e403fd9455708b3b942949e1d7febc02167662bf1a7da5b78ee1ea2842/pydantic_core-2.41.5-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:7f3bf998340c6d4b0c9a2f02d6a400e51f123b59565d74dc60d252ce888c260b", size = 1899826, upload-time = "2025-11-04T13:39:32.897Z" }, - { url = "https://files.pythonhosted.org/packages/33/7f/1d5cab3ccf44c1935a359d51a8a2a9e1a654b744b5e7f80d41b88d501eec/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:378bec5c66998815d224c9ca994f1e14c0c21cb95d2f52b6021cc0b2a58f2a5a", size = 1917869, upload-time = "2025-11-04T13:39:34.469Z" }, - { url = "https://files.pythonhosted.org/packages/6e/6a/30d94a9674a7fe4f4744052ed6c5e083424510be1e93da5bc47569d11810/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e7b576130c69225432866fe2f4a469a85a54ade141d96fd396dffcf607b558f8", size = 2063890, upload-time = "2025-11-04T13:39:36.053Z" }, - { url = "https://files.pythonhosted.org/packages/50/be/76e5d46203fcb2750e542f32e6c371ffa9b8ad17364cf94bb0818dbfb50c/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6cb58b9c66f7e4179a2d5e0f849c48eff5c1fca560994d6eb6543abf955a149e", size = 2229740, upload-time = "2025-11-04T13:39:37.753Z" }, - { url = "https://files.pythonhosted.org/packages/d3/ee/fed784df0144793489f87db310a6bbf8118d7b630ed07aa180d6067e653a/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:88942d3a3dff3afc8288c21e565e476fc278902ae4d6d134f1eeda118cc830b1", size = 2350021, upload-time = "2025-11-04T13:39:40.94Z" }, - { url = "https://files.pythonhosted.org/packages/c8/be/8fed28dd0a180dca19e72c233cbf58efa36df055e5b9d90d64fd1740b828/pydantic_core-2.41.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f31d95a179f8d64d90f6831d71fa93290893a33148d890ba15de25642c5d075b", size = 2066378, upload-time = "2025-11-04T13:39:42.523Z" }, - { url = "https://files.pythonhosted.org/packages/b0/3b/698cf8ae1d536a010e05121b4958b1257f0b5522085e335360e53a6b1c8b/pydantic_core-2.41.5-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c1df3d34aced70add6f867a8cf413e299177e0c22660cc767218373d0779487b", size = 2175761, upload-time = "2025-11-04T13:39:44.553Z" }, - { url = "https://files.pythonhosted.org/packages/b8/ba/15d537423939553116dea94ce02f9c31be0fa9d0b806d427e0308ec17145/pydantic_core-2.41.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:4009935984bd36bd2c774e13f9a09563ce8de4abaa7226f5108262fa3e637284", size = 2146303, upload-time = "2025-11-04T13:39:46.238Z" }, - { url = "https://files.pythonhosted.org/packages/58/7f/0de669bf37d206723795f9c90c82966726a2ab06c336deba4735b55af431/pydantic_core-2.41.5-cp311-cp311-musllinux_1_1_armv7l.whl", hash = "sha256:34a64bc3441dc1213096a20fe27e8e128bd3ff89921706e83c0b1ac971276594", size = 2340355, upload-time = "2025-11-04T13:39:48.002Z" }, - { url = "https://files.pythonhosted.org/packages/e5/de/e7482c435b83d7e3c3ee5ee4451f6e8973cff0eb6007d2872ce6383f6398/pydantic_core-2.41.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:c9e19dd6e28fdcaa5a1de679aec4141f691023916427ef9bae8584f9c2fb3b0e", size = 2319875, upload-time = "2025-11-04T13:39:49.705Z" }, - { url = "https://files.pythonhosted.org/packages/fe/e6/8c9e81bb6dd7560e33b9053351c29f30c8194b72f2d6932888581f503482/pydantic_core-2.41.5-cp311-cp311-win32.whl", hash = "sha256:2c010c6ded393148374c0f6f0bf89d206bf3217f201faa0635dcd56bd1520f6b", size = 1987549, upload-time = "2025-11-04T13:39:51.842Z" }, - { url = "https://files.pythonhosted.org/packages/11/66/f14d1d978ea94d1bc21fc98fcf570f9542fe55bfcc40269d4e1a21c19bf7/pydantic_core-2.41.5-cp311-cp311-win_amd64.whl", hash = "sha256:76ee27c6e9c7f16f47db7a94157112a2f3a00e958bc626e2f4ee8bec5c328fbe", size = 2011305, upload-time = "2025-11-04T13:39:53.485Z" }, - { url = "https://files.pythonhosted.org/packages/56/d8/0e271434e8efd03186c5386671328154ee349ff0354d83c74f5caaf096ed/pydantic_core-2.41.5-cp311-cp311-win_arm64.whl", hash = "sha256:4bc36bbc0b7584de96561184ad7f012478987882ebf9f9c389b23f432ea3d90f", size = 1972902, upload-time = "2025-11-04T13:39:56.488Z" }, - { url = "https://files.pythonhosted.org/packages/5f/5d/5f6c63eebb5afee93bcaae4ce9a898f3373ca23df3ccaef086d0233a35a7/pydantic_core-2.41.5-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:f41a7489d32336dbf2199c8c0a215390a751c5b014c2c1c5366e817202e9cdf7", size = 2110990, upload-time = "2025-11-04T13:39:58.079Z" }, - { url = "https://files.pythonhosted.org/packages/aa/32/9c2e8ccb57c01111e0fd091f236c7b371c1bccea0fa85247ac55b1e2b6b6/pydantic_core-2.41.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:070259a8818988b9a84a449a2a7337c7f430a22acc0859c6b110aa7212a6d9c0", size = 1896003, upload-time = "2025-11-04T13:39:59.956Z" }, - { url = "https://files.pythonhosted.org/packages/68/b8/a01b53cb0e59139fbc9e4fda3e9724ede8de279097179be4ff31f1abb65a/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e96cea19e34778f8d59fe40775a7a574d95816eb150850a85a7a4c8f4b94ac69", size = 1919200, upload-time = "2025-11-04T13:40:02.241Z" }, - { url = "https://files.pythonhosted.org/packages/38/de/8c36b5198a29bdaade07b5985e80a233a5ac27137846f3bc2d3b40a47360/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:ed2e99c456e3fadd05c991f8f437ef902e00eedf34320ba2b0842bd1c3ca3a75", size = 2052578, upload-time = "2025-11-04T13:40:04.401Z" }, - { url = "https://files.pythonhosted.org/packages/00/b5/0e8e4b5b081eac6cb3dbb7e60a65907549a1ce035a724368c330112adfdd/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:65840751b72fbfd82c3c640cff9284545342a4f1eb1586ad0636955b261b0b05", size = 2208504, upload-time = "2025-11-04T13:40:06.072Z" }, - { url = "https://files.pythonhosted.org/packages/77/56/87a61aad59c7c5b9dc8caad5a41a5545cba3810c3e828708b3d7404f6cef/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e536c98a7626a98feb2d3eaf75944ef6f3dbee447e1f841eae16f2f0a72d8ddc", size = 2335816, upload-time = "2025-11-04T13:40:07.835Z" }, - { url = "https://files.pythonhosted.org/packages/0d/76/941cc9f73529988688a665a5c0ecff1112b3d95ab48f81db5f7606f522d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eceb81a8d74f9267ef4081e246ffd6d129da5d87e37a77c9bde550cb04870c1c", size = 2075366, upload-time = "2025-11-04T13:40:09.804Z" }, - { url = "https://files.pythonhosted.org/packages/d3/43/ebef01f69baa07a482844faaa0a591bad1ef129253ffd0cdaa9d8a7f72d3/pydantic_core-2.41.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:d38548150c39b74aeeb0ce8ee1d8e82696f4a4e16ddc6de7b1d8823f7de4b9b5", size = 2171698, upload-time = "2025-11-04T13:40:12.004Z" }, - { url = "https://files.pythonhosted.org/packages/b1/87/41f3202e4193e3bacfc2c065fab7706ebe81af46a83d3e27605029c1f5a6/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:c23e27686783f60290e36827f9c626e63154b82b116d7fe9adba1fda36da706c", size = 2132603, upload-time = "2025-11-04T13:40:13.868Z" }, - { url = "https://files.pythonhosted.org/packages/49/7d/4c00df99cb12070b6bccdef4a195255e6020a550d572768d92cc54dba91a/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:482c982f814460eabe1d3bb0adfdc583387bd4691ef00b90575ca0d2b6fe2294", size = 2329591, upload-time = "2025-11-04T13:40:15.672Z" }, - { url = "https://files.pythonhosted.org/packages/cc/6a/ebf4b1d65d458f3cda6a7335d141305dfa19bdc61140a884d165a8a1bbc7/pydantic_core-2.41.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:bfea2a5f0b4d8d43adf9d7b8bf019fb46fdd10a2e5cde477fbcb9d1fa08c68e1", size = 2319068, upload-time = "2025-11-04T13:40:17.532Z" }, - { url = "https://files.pythonhosted.org/packages/49/3b/774f2b5cd4192d5ab75870ce4381fd89cf218af999515baf07e7206753f0/pydantic_core-2.41.5-cp312-cp312-win32.whl", hash = "sha256:b74557b16e390ec12dca509bce9264c3bbd128f8a2c376eaa68003d7f327276d", size = 1985908, upload-time = "2025-11-04T13:40:19.309Z" }, - { url = "https://files.pythonhosted.org/packages/86/45/00173a033c801cacf67c190fef088789394feaf88a98a7035b0e40d53dc9/pydantic_core-2.41.5-cp312-cp312-win_amd64.whl", hash = "sha256:1962293292865bca8e54702b08a4f26da73adc83dd1fcf26fbc875b35d81c815", size = 2020145, upload-time = "2025-11-04T13:40:21.548Z" }, - { url = "https://files.pythonhosted.org/packages/f9/22/91fbc821fa6d261b376a3f73809f907cec5ca6025642c463d3488aad22fb/pydantic_core-2.41.5-cp312-cp312-win_arm64.whl", hash = "sha256:1746d4a3d9a794cacae06a5eaaccb4b8643a131d45fbc9af23e353dc0a5ba5c3", size = 1976179, upload-time = "2025-11-04T13:40:23.393Z" }, - { url = "https://files.pythonhosted.org/packages/87/06/8806241ff1f70d9939f9af039c6c35f2360cf16e93c2ca76f184e76b1564/pydantic_core-2.41.5-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:941103c9be18ac8daf7b7adca8228f8ed6bb7a1849020f643b3a14d15b1924d9", size = 2120403, upload-time = "2025-11-04T13:40:25.248Z" }, - { url = "https://files.pythonhosted.org/packages/94/02/abfa0e0bda67faa65fef1c84971c7e45928e108fe24333c81f3bfe35d5f5/pydantic_core-2.41.5-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:112e305c3314f40c93998e567879e887a3160bb8689ef3d2c04b6cc62c33ac34", size = 1896206, upload-time = "2025-11-04T13:40:27.099Z" }, - { url = "https://files.pythonhosted.org/packages/15/df/a4c740c0943e93e6500f9eb23f4ca7ec9bf71b19e608ae5b579678c8d02f/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0cbaad15cb0c90aa221d43c00e77bb33c93e8d36e0bf74760cd00e732d10a6a0", size = 1919307, upload-time = "2025-11-04T13:40:29.806Z" }, - { url = "https://files.pythonhosted.org/packages/9a/e3/6324802931ae1d123528988e0e86587c2072ac2e5394b4bc2bc34b61ff6e/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:03ca43e12fab6023fc79d28ca6b39b05f794ad08ec2feccc59a339b02f2b3d33", size = 2063258, upload-time = "2025-11-04T13:40:33.544Z" }, - { url = "https://files.pythonhosted.org/packages/c9/d4/2230d7151d4957dd79c3044ea26346c148c98fbf0ee6ebd41056f2d62ab5/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:dc799088c08fa04e43144b164feb0c13f9a0bc40503f8df3e9fde58a3c0c101e", size = 2214917, upload-time = "2025-11-04T13:40:35.479Z" }, - { url = "https://files.pythonhosted.org/packages/e6/9f/eaac5df17a3672fef0081b6c1bb0b82b33ee89aa5cec0d7b05f52fd4a1fa/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:97aeba56665b4c3235a0e52b2c2f5ae9cd071b8a8310ad27bddb3f7fb30e9aa2", size = 2332186, upload-time = "2025-11-04T13:40:37.436Z" }, - { url = "https://files.pythonhosted.org/packages/cf/4e/35a80cae583a37cf15604b44240e45c05e04e86f9cfd766623149297e971/pydantic_core-2.41.5-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:406bf18d345822d6c21366031003612b9c77b3e29ffdb0f612367352aab7d586", size = 2073164, upload-time = "2025-11-04T13:40:40.289Z" }, - { url = "https://files.pythonhosted.org/packages/bf/e3/f6e262673c6140dd3305d144d032f7bd5f7497d3871c1428521f19f9efa2/pydantic_core-2.41.5-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:b93590ae81f7010dbe380cdeab6f515902ebcbefe0b9327cc4804d74e93ae69d", size = 2179146, upload-time = "2025-11-04T13:40:42.809Z" }, - { url = "https://files.pythonhosted.org/packages/75/c7/20bd7fc05f0c6ea2056a4565c6f36f8968c0924f19b7d97bbfea55780e73/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:01a3d0ab748ee531f4ea6c3e48ad9dac84ddba4b0d82291f87248f2f9de8d740", size = 2137788, upload-time = "2025-11-04T13:40:44.752Z" }, - { url = "https://files.pythonhosted.org/packages/3a/8d/34318ef985c45196e004bc46c6eab2eda437e744c124ef0dbe1ff2c9d06b/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:6561e94ba9dacc9c61bce40e2d6bdc3bfaa0259d3ff36ace3b1e6901936d2e3e", size = 2340133, upload-time = "2025-11-04T13:40:46.66Z" }, - { url = "https://files.pythonhosted.org/packages/9c/59/013626bf8c78a5a5d9350d12e7697d3d4de951a75565496abd40ccd46bee/pydantic_core-2.41.5-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:915c3d10f81bec3a74fbd4faebe8391013ba61e5a1a8d48c4455b923bdda7858", size = 2324852, upload-time = "2025-11-04T13:40:48.575Z" }, - { url = "https://files.pythonhosted.org/packages/1a/d9/c248c103856f807ef70c18a4f986693a46a8ffe1602e5d361485da502d20/pydantic_core-2.41.5-cp313-cp313-win32.whl", hash = "sha256:650ae77860b45cfa6e2cdafc42618ceafab3a2d9a3811fcfbd3bbf8ac3c40d36", size = 1994679, upload-time = "2025-11-04T13:40:50.619Z" }, - { url = "https://files.pythonhosted.org/packages/9e/8b/341991b158ddab181cff136acd2552c9f35bd30380422a639c0671e99a91/pydantic_core-2.41.5-cp313-cp313-win_amd64.whl", hash = "sha256:79ec52ec461e99e13791ec6508c722742ad745571f234ea6255bed38c6480f11", size = 2019766, upload-time = "2025-11-04T13:40:52.631Z" }, - { url = "https://files.pythonhosted.org/packages/73/7d/f2f9db34af103bea3e09735bb40b021788a5e834c81eedb541991badf8f5/pydantic_core-2.41.5-cp313-cp313-win_arm64.whl", hash = "sha256:3f84d5c1b4ab906093bdc1ff10484838aca54ef08de4afa9de0f5f14d69639cd", size = 1981005, upload-time = "2025-11-04T13:40:54.734Z" }, - { url = "https://files.pythonhosted.org/packages/11/72/90fda5ee3b97e51c494938a4a44c3a35a9c96c19bba12372fb9c634d6f57/pydantic_core-2.41.5-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:b96d5f26b05d03cc60f11a7761a5ded1741da411e7fe0909e27a5e6a0cb7b034", size = 2115441, upload-time = "2025-11-04T13:42:39.557Z" }, - { url = "https://files.pythonhosted.org/packages/1f/53/8942f884fa33f50794f119012dc6a1a02ac43a56407adaac20463df8e98f/pydantic_core-2.41.5-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:634e8609e89ceecea15e2d61bc9ac3718caaaa71963717bf3c8f38bfde64242c", size = 1930291, upload-time = "2025-11-04T13:42:42.169Z" }, - { url = "https://files.pythonhosted.org/packages/79/c8/ecb9ed9cd942bce09fc888ee960b52654fbdbede4ba6c2d6e0d3b1d8b49c/pydantic_core-2.41.5-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:93e8740d7503eb008aa2df04d3b9735f845d43ae845e6dcd2be0b55a2da43cd2", size = 1948632, upload-time = "2025-11-04T13:42:44.564Z" }, - { url = "https://files.pythonhosted.org/packages/2e/1b/687711069de7efa6af934e74f601e2a4307365e8fdc404703afc453eab26/pydantic_core-2.41.5-graalpy311-graalpy242_311_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f15489ba13d61f670dcc96772e733aad1a6f9c429cc27574c6cdaed82d0146ad", size = 2138905, upload-time = "2025-11-04T13:42:47.156Z" }, - { url = "https://files.pythonhosted.org/packages/09/32/59b0c7e63e277fa7911c2fc70ccfb45ce4b98991e7ef37110663437005af/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:7da7087d756b19037bc2c06edc6c170eeef3c3bafcb8f532ff17d64dc427adfd", size = 2110495, upload-time = "2025-11-04T13:42:49.689Z" }, - { url = "https://files.pythonhosted.org/packages/aa/81/05e400037eaf55ad400bcd318c05bb345b57e708887f07ddb2d20e3f0e98/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:aabf5777b5c8ca26f7824cb4a120a740c9588ed58df9b2d196ce92fba42ff8dc", size = 1915388, upload-time = "2025-11-04T13:42:52.215Z" }, - { url = "https://files.pythonhosted.org/packages/6e/0d/e3549b2399f71d56476b77dbf3cf8937cec5cd70536bdc0e374a421d0599/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:c007fe8a43d43b3969e8469004e9845944f1a80e6acd47c150856bb87f230c56", size = 1942879, upload-time = "2025-11-04T13:42:56.483Z" }, - { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" }, - { url = "https://files.pythonhosted.org/packages/e6/b0/1a2aa41e3b5a4ba11420aba2d091b2d17959c8d1519ece3627c371951e73/pydantic_core-2.41.5-pp310-pypy310_pp73-macosx_10_12_x86_64.whl", hash = "sha256:b5819cd790dbf0c5eb9f82c73c16b39a65dd6dd4d1439dcdea7816ec9adddab8", size = 2103351, upload-time = "2025-11-04T13:43:02.058Z" }, - { url = "https://files.pythonhosted.org/packages/a4/ee/31b1f0020baaf6d091c87900ae05c6aeae101fa4e188e1613c80e4f1ea31/pydantic_core-2.41.5-pp310-pypy310_pp73-macosx_11_0_arm64.whl", hash = "sha256:5a4e67afbc95fa5c34cf27d9089bca7fcab4e51e57278d710320a70b956d1b9a", size = 1925363, upload-time = "2025-11-04T13:43:05.159Z" }, - { url = "https://files.pythonhosted.org/packages/e1/89/ab8e86208467e467a80deaca4e434adac37b10a9d134cd2f99b28a01e483/pydantic_core-2.41.5-pp310-pypy310_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ece5c59f0ce7d001e017643d8d24da587ea1f74f6993467d85ae8a5ef9d4f42b", size = 2135615, upload-time = "2025-11-04T13:43:08.116Z" }, - { url = "https://files.pythonhosted.org/packages/99/0a/99a53d06dd0348b2008f2f30884b34719c323f16c3be4e6cc1203b74a91d/pydantic_core-2.41.5-pp310-pypy310_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:16f80f7abe3351f8ea6858914ddc8c77e02578544a0ebc15b4c2e1a0e813b0b2", size = 2175369, upload-time = "2025-11-04T13:43:12.49Z" }, - { url = "https://files.pythonhosted.org/packages/6d/94/30ca3b73c6d485b9bb0bc66e611cff4a7138ff9736b7e66bcf0852151636/pydantic_core-2.41.5-pp310-pypy310_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:33cb885e759a705b426baada1fe68cbb0a2e68e34c5d0d0289a364cf01709093", size = 2144218, upload-time = "2025-11-04T13:43:15.431Z" }, - { url = "https://files.pythonhosted.org/packages/87/57/31b4f8e12680b739a91f472b5671294236b82586889ef764b5fbc6669238/pydantic_core-2.41.5-pp310-pypy310_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:c8d8b4eb992936023be7dee581270af5c6e0697a8559895f527f5b7105ecd36a", size = 2329951, upload-time = "2025-11-04T13:43:18.062Z" }, - { url = "https://files.pythonhosted.org/packages/7d/73/3c2c8edef77b8f7310e6fb012dbc4b8551386ed575b9eb6fb2506e28a7eb/pydantic_core-2.41.5-pp310-pypy310_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:242a206cd0318f95cd21bdacff3fcc3aab23e79bba5cac3db5a841c9ef9c6963", size = 2318428, upload-time = "2025-11-04T13:43:20.679Z" }, - { url = "https://files.pythonhosted.org/packages/2f/02/8559b1f26ee0d502c74f9cca5c0d2fd97e967e083e006bbbb4e97f3a043a/pydantic_core-2.41.5-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:d3a978c4f57a597908b7e697229d996d77a6d3c94901e9edee593adada95ce1a", size = 2147009, upload-time = "2025-11-04T13:43:23.286Z" }, - { url = "https://files.pythonhosted.org/packages/5f/9b/1b3f0e9f9305839d7e84912f9e8bfbd191ed1b1ef48083609f0dabde978c/pydantic_core-2.41.5-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:b2379fa7ed44ddecb5bfe4e48577d752db9fc10be00a6b7446e9663ba143de26", size = 2101980, upload-time = "2025-11-04T13:43:25.97Z" }, - { url = "https://files.pythonhosted.org/packages/a4/ed/d71fefcb4263df0da6a85b5d8a7508360f2f2e9b3bf5814be9c8bccdccc1/pydantic_core-2.41.5-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:266fb4cbf5e3cbd0b53669a6d1b039c45e3ce651fd5442eff4d07c2cc8d66808", size = 1923865, upload-time = "2025-11-04T13:43:28.763Z" }, - { url = "https://files.pythonhosted.org/packages/ce/3a/626b38db460d675f873e4444b4bb030453bbe7b4ba55df821d026a0493c4/pydantic_core-2.41.5-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:58133647260ea01e4d0500089a8c4f07bd7aa6ce109682b1426394988d8aaacc", size = 2134256, upload-time = "2025-11-04T13:43:31.71Z" }, - { url = "https://files.pythonhosted.org/packages/83/d9/8412d7f06f616bbc053d30cb4e5f76786af3221462ad5eee1f202021eb4e/pydantic_core-2.41.5-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:287dad91cfb551c363dc62899a80e9e14da1f0e2b6ebde82c806612ca2a13ef1", size = 2174762, upload-time = "2025-11-04T13:43:34.744Z" }, - { url = "https://files.pythonhosted.org/packages/55/4c/162d906b8e3ba3a99354e20faa1b49a85206c47de97a639510a0e673f5da/pydantic_core-2.41.5-pp311-pypy311_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:03b77d184b9eb40240ae9fd676ca364ce1085f203e1b1256f8ab9984dca80a84", size = 2143141, upload-time = "2025-11-04T13:43:37.701Z" }, - { url = "https://files.pythonhosted.org/packages/1f/f2/f11dd73284122713f5f89fc940f370d035fa8e1e078d446b3313955157fe/pydantic_core-2.41.5-pp311-pypy311_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:a668ce24de96165bb239160b3d854943128f4334822900534f2fe947930e5770", size = 2330317, upload-time = "2025-11-04T13:43:40.406Z" }, - { url = "https://files.pythonhosted.org/packages/88/9d/b06ca6acfe4abb296110fb1273a4d848a0bfb2ff65f3ee92127b3244e16b/pydantic_core-2.41.5-pp311-pypy311_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:f14f8f046c14563f8eb3f45f499cc658ab8d10072961e07225e507adb700e93f", size = 2316992, upload-time = "2025-11-04T13:43:43.602Z" }, - { url = "https://files.pythonhosted.org/packages/36/c7/cfc8e811f061c841d7990b0201912c3556bfeb99cdcb7ed24adc8d6f8704/pydantic_core-2.41.5-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:56121965f7a4dc965bff783d70b907ddf3d57f6eba29b6d2e5dabfaf07799c51", size = 2145302, upload-time = "2025-11-04T13:43:46.64Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/a1/93/f97a86a7eb28faa1d038af2fd5d6166418b4433659108a4c311b57128b2d/pydantic_core-2.46.1.tar.gz", hash = "sha256:d408153772d9f298098fb5d620f045bdf0f017af0d5cb6e309ef8c205540caa4", size = 471230, upload-time = "2026-04-15T14:49:34.52Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/a0/07f275411355b567b994e565bc5ea9dbf522978060c18e3b7edf646c0fc2/pydantic_core-2.46.1-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:84eb5414871fd0293c38d2075802f95030ff11a92cf2189942bf76fd181af77b", size = 2123782, upload-time = "2026-04-15T14:52:57.172Z" }, + { url = "https://files.pythonhosted.org/packages/ab/71/d027c7de46df5b9287ed6f0ef02346c84d61348326253a4f13695d54d66f/pydantic_core-2.46.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:5c75fb25db086bf504c55730442e471c12bc9bfae817dd359b1a36bc93049d34", size = 1948561, upload-time = "2026-04-15T14:53:12.07Z" }, + { url = "https://files.pythonhosted.org/packages/77/74/cba894bea0d51a3b2dcada9eb3af9c4cfaa271bf21123372dc82ccef029f/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:29dc09f0221425453fd9f73fd70bba15817d25b95858282702d7305a08d37306", size = 1974387, upload-time = "2026-04-15T14:50:14.048Z" }, + { url = "https://files.pythonhosted.org/packages/3b/ad/cc122887d6f20ac5d997928b0bf3016ac9c7bae07dce089333aa0c2e868b/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:139fd6722abc5e6513aa0a27b06ebeb997838c5b179cf5e83862ace45f281c56", size = 2054868, upload-time = "2026-04-15T14:49:51.912Z" }, + { url = "https://files.pythonhosted.org/packages/9f/09/22049b22d65a67253cbdced88dbce0e97162f35cc433917df37df794ede8/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ba723fd8ef6011af71f92ed54adb604e7699d172f4273e4b46f1cfb8ee8d72fd", size = 2228717, upload-time = "2026-04-15T14:49:27.384Z" }, + { url = "https://files.pythonhosted.org/packages/e6/98/b35a8a187cf977462668b5064c606e290c88c2561e053883d86193ab9c51/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:828410e082555e55da9bbb5e6c17617386fe1415c4d42765a90d372ed9cce813", size = 2298261, upload-time = "2026-04-15T14:52:20.463Z" }, + { url = "https://files.pythonhosted.org/packages/98/ae/46f8d693caefc09d8e2d3f19a6b4f2252cf6542f0b555759f2b5ec2b4ca5/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fb5cd53264c9906c163a71b489e9ac71b0ae13a2dd0241e6129f4df38ba1c814", size = 2094496, upload-time = "2026-04-15T14:49:59.711Z" }, + { url = "https://files.pythonhosted.org/packages/ee/40/7e4013639d316d2cb67dae288c768d49cc4a7a4b16ef869e486880db1a1f/pydantic_core-2.46.1-cp310-cp310-manylinux_2_31_riscv64.whl", hash = "sha256:4530a6594883d9d4a9c7ef68464ef6b4a88d839e3531c089a3942c78bffe0a66", size = 2144795, upload-time = "2026-04-15T14:52:44.731Z" }, + { url = "https://files.pythonhosted.org/packages/0d/87/c00f6450059804faf30f568009c8c98e72e6802c1ccd8b562da57953ad81/pydantic_core-2.46.1-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:ed1c71f60abbf9c9a440dc8fc6b1180c45dcab3a5e311250de99744a0166bc95", size = 2173108, upload-time = "2026-04-15T14:51:37.806Z" }, + { url = "https://files.pythonhosted.org/packages/46/15/7a8fb06c109a07dbc1f5f272b2da1290c8a25f5900a579086e433049fc1a/pydantic_core-2.46.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:254253491f1b8e3ba18c15fe924bb9b175f1a48413b74e8f0c67b8f51b6f726b", size = 2185687, upload-time = "2026-04-15T14:51:33.125Z" }, + { url = "https://files.pythonhosted.org/packages/d9/38/c52ead78febf23d32db898c7022173c674226cf3c8ee1645220ab9516931/pydantic_core-2.46.1-cp310-cp310-musllinux_1_1_armv7l.whl", hash = "sha256:dfcf6485ac38698a5b45f37467b8eb2f4f8e3edd5790e2579c5d52fdfffb2e3d", size = 2326273, upload-time = "2026-04-15T14:51:10.614Z" }, + { url = "https://files.pythonhosted.org/packages/1e/af/cb5ea2336e9938b3a0536ce4bfed4a342285caa8a6b8ff449a7bc2f179ec/pydantic_core-2.46.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:592b39150ab5b5a2cb2eb885097ee4c2e4d54e3b902f6ae32528f7e6e42c00fc", size = 2368428, upload-time = "2026-04-15T14:49:25.804Z" }, + { url = "https://files.pythonhosted.org/packages/a2/99/adcfbcbd96556120e7d795aab4fd77f5104a49051929c3805a9d736ec48f/pydantic_core-2.46.1-cp310-cp310-win32.whl", hash = "sha256:eb37b1369ad39ec046a36dc81ffd76870766bda2073f57448bbcb1fd3e4c5ad0", size = 1993405, upload-time = "2026-04-15T14:50:51.082Z" }, + { url = "https://files.pythonhosted.org/packages/c4/ff/2767be513a250293f80748740ce73b0f0677711fc791b1afab3499734dd2/pydantic_core-2.46.1-cp310-cp310-win_amd64.whl", hash = "sha256:c330dab8254d422880177436a5892ac6d9337afff9fe383fb1f8c6caedb685e1", size = 2068177, upload-time = "2026-04-15T14:52:29.899Z" }, + { url = "https://files.pythonhosted.org/packages/37/96/d83d23fc3c822326d808b8c0457d4f7afb1552e741a7c2378a974c522c63/pydantic_core-2.46.1-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:f0f84431981c6ae217ebb96c3eca8212f6f5edf116f62f62cc6c7d72971f826c", size = 2121938, upload-time = "2026-04-15T14:49:21.568Z" }, + { url = "https://files.pythonhosted.org/packages/11/44/94b1251825560f5d90e25ebcd457c4772e1f3e1a378f438c040fe2148f3e/pydantic_core-2.46.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a05f60b36549f59ab585924410187276ec17a94bae939273a213cea252c8471e", size = 1946541, upload-time = "2026-04-15T14:49:57.925Z" }, + { url = "https://files.pythonhosted.org/packages/d6/8f/79aff4c8bd6fb49001ffe4747c775c0f066add9da13dec180eb0023ada34/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b2c93fd1693afdfae7b2897f7530ed3f180d9fc92ee105df3ebdff24d5061cc8", size = 1973067, upload-time = "2026-04-15T14:51:14.765Z" }, + { url = "https://files.pythonhosted.org/packages/56/01/826ab3afb1d43cbfdc2aa592bff0f1f6f4b90f5a801478ba07bde74e706f/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0c19983759394c702a776f42f33df8d7bb7883aefaa44a69ba86356a9fd67367", size = 2053146, upload-time = "2026-04-15T14:51:48.847Z" }, + { url = "https://files.pythonhosted.org/packages/6c/32/be20ec48ccbd85cac3f8d96ca0a0f87d5c14fbf1eb438da0ac733f2546f2/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6e8debf586d7d800a718194417497db5126d4f4302885a2dff721e9df3f4851c", size = 2227393, upload-time = "2026-04-15T14:51:53.218Z" }, + { url = "https://files.pythonhosted.org/packages/b5/8e/1fae21c887f363ed1a5cf9f267027700c796b7435313c21723cd3e8aeeb3/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:54160da754d63da7780b76e5743d44f026b9daffc6b8c9696a756368c0a298c9", size = 2296193, upload-time = "2026-04-15T14:50:31.065Z" }, + { url = "https://files.pythonhosted.org/packages/0a/29/e5637b539458ffb60ba9c204fc16c52ea36828427fa667e4f9c7d83cfea9/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:74cee962c8b4df9a9b0bb63582e51986127ee2316f0c49143b2996f4b201bd9c", size = 2092156, upload-time = "2026-04-15T14:52:37.227Z" }, + { url = "https://files.pythonhosted.org/packages/bc/fa/3a453934af019c72652fb75489c504ae689de632fa2e037fec3195cd6948/pydantic_core-2.46.1-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:0ba3462872a678ebe21b15bd78eff40298b43ea50c26f230ec535c00cf93ec7e", size = 2142845, upload-time = "2026-04-15T14:51:04.847Z" }, + { url = "https://files.pythonhosted.org/packages/36/c2/71b56fa10a80b98036f4bf0fbb912833f8e9c61b15e66c236fadaf54c27c/pydantic_core-2.46.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:b718873a966d91514c5252775f568985401b54a220919ab22b19a6c4edd8c053", size = 2170756, upload-time = "2026-04-15T14:50:17.16Z" }, + { url = "https://files.pythonhosted.org/packages/e1/da/a4c761dc8d982e2c53f991c0c36d37f6fe308e149bf0a101c25b0750a893/pydantic_core-2.46.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:cb1310a9fd722da8cceec1fb59875e1c86bee37f0d8a9c667220f00ee722cc8f", size = 2183579, upload-time = "2026-04-15T14:51:20.888Z" }, + { url = "https://files.pythonhosted.org/packages/e5/d4/b0a6c00622e4afd9a807b8bb05ba8f1a0b69ca068ac138d9d36700fe767b/pydantic_core-2.46.1-cp311-cp311-musllinux_1_1_armv7l.whl", hash = "sha256:98e3ede76eb4b9db8e7b5efea07a3f3315135485794a5df91e3adf56c4d573b6", size = 2324516, upload-time = "2026-04-15T14:52:32.521Z" }, + { url = "https://files.pythonhosted.org/packages/45/f1/a4bace0c98b0774b02de99233882c48d94b399ba4394dd5e209665d05062/pydantic_core-2.46.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:780b8f24ff286e21fd010247011a68ea902c34b1eee7d775b598bc28f5f28ab6", size = 2367084, upload-time = "2026-04-15T14:50:37.832Z" }, + { url = "https://files.pythonhosted.org/packages/3a/54/ae827a3976b136d1c9a9a56c2299a8053605a69facaa0c7354ba167305eb/pydantic_core-2.46.1-cp311-cp311-win32.whl", hash = "sha256:1d452f4cad0f39a94414ca68cda7cc55ff4c3801b5ab0bc99818284a3d39f889", size = 1992061, upload-time = "2026-04-15T14:51:44.704Z" }, + { url = "https://files.pythonhosted.org/packages/55/ae/d85de69e0fdfafc0e87d88bd5d0c157a5443efaaef24eed152a8a8f8dfb6/pydantic_core-2.46.1-cp311-cp311-win_amd64.whl", hash = "sha256:f463fd6a67138d70200d2627676e9efbb0cee26d98a5d3042a35aa20f95ec129", size = 2065497, upload-time = "2026-04-15T14:51:17.077Z" }, + { url = "https://files.pythonhosted.org/packages/46/a7/9eb3b1038db630e1550924e81d1211b0dd70ac3740901fd95f30f5497990/pydantic_core-2.46.1-cp311-cp311-win_arm64.whl", hash = "sha256:155aec0a117140e86775eec113b574c1c299358bfd99467b2ea7b2ea26db2614", size = 2045914, upload-time = "2026-04-15T14:51:24.782Z" }, + { url = "https://files.pythonhosted.org/packages/ce/fb/caaa8ee23861c170f07dbd58fc2be3a2c02a32637693cbb23eef02e84808/pydantic_core-2.46.1-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:ae8c8c5eb4c796944f3166f2f0dab6c761c2c2cc5bd20e5f692128be8600b9a4", size = 2119472, upload-time = "2026-04-15T14:49:45.946Z" }, + { url = "https://files.pythonhosted.org/packages/fa/61/bcffaa52894489ff89e5e1cdde67429914bf083c0db7296bef153020f786/pydantic_core-2.46.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:daba6f5f5b986aa0682623a1a4f8d1ecb0ec00ce09cfa9ca71a3b742bc383e3a", size = 1951230, upload-time = "2026-04-15T14:52:27.646Z" }, + { url = "https://files.pythonhosted.org/packages/f8/95/80d2f43a2a1a1e3220fd329d614aa5a39e0a75d24353a3aaf226e605f1c2/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0265f3a2460539ecc97817a80c7a23c458dd84191229b655522a2674f701f14e", size = 1976394, upload-time = "2026-04-15T14:50:32.742Z" }, + { url = "https://files.pythonhosted.org/packages/8d/31/2c5b1a207926b5fc1961a2d11da940129bc3841c36cc4df03014195b2966/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bb16c0156c4b4e94aa3719138cc43c53d30ff21126b6a3af63786dcc0757b56e", size = 2068455, upload-time = "2026-04-15T14:50:01.286Z" }, + { url = "https://files.pythonhosted.org/packages/7d/36/c6aa07274359a51ac62895895325ce90107e811c6cea39d2617a99ef10d7/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1b42d80fad8e4b283e1e4138f1142f0d038c46d137aad2f9824ad9086080dd41", size = 2239049, upload-time = "2026-04-15T14:53:02.216Z" }, + { url = "https://files.pythonhosted.org/packages/0a/3f/77cdd0db8bddc714842dfd93f737c863751cf02001c993341504f6b0cd53/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9cced85896d5b795293bc36b7e2fb0347a36c828551b50cbba510510d928548c", size = 2318681, upload-time = "2026-04-15T14:50:04.539Z" }, + { url = "https://files.pythonhosted.org/packages/a1/a3/09d929a40e6727274b0b500ad06e1b3f35d4f4665ae1c8ba65acbb17e9b5/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a641cb1e74b44c418adaf9f5f450670dbec53511f030d8cde8d8accb66edc363", size = 2096527, upload-time = "2026-04-15T14:53:14.766Z" }, + { url = "https://files.pythonhosted.org/packages/89/ae/544c3a82456ebc254a9fcbe2715bab76c70acf9d291aaea24391147943e4/pydantic_core-2.46.1-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:191e7a122ab14eb12415fe3f92610fc06c7f1d2b4b9101d24d490d447ac92506", size = 2170407, upload-time = "2026-04-15T14:51:27.138Z" }, + { url = "https://files.pythonhosted.org/packages/9d/ce/0dfd881c7af4c522f47b325707bd9a2cdcf4f40e4f2fd30df0e9a3e8d393/pydantic_core-2.46.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4fe4ff660f7938b5d92f21529ce331b011aa35e481ab64b7cd03f52384e544bb", size = 2188578, upload-time = "2026-04-15T14:50:39.655Z" }, + { url = "https://files.pythonhosted.org/packages/a1/e9/980ea2a6d5114dd1a62ecc5f56feb3d34555f33bd11043f042e5f7f0724a/pydantic_core-2.46.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:18fcea085b3adc3868d8d19606da52d7a52d8bccd8e28652b0778dbe5e6a6660", size = 2188959, upload-time = "2026-04-15T14:52:42.243Z" }, + { url = "https://files.pythonhosted.org/packages/e7/f1/595e0f50f4bfc56cde2fe558f2b0978f29f2865da894c6226231e17464a5/pydantic_core-2.46.1-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:e8e589e7c9466e022d79e13c5764c2239b2e5a7993ba727822b021234f89b56b", size = 2339973, upload-time = "2026-04-15T14:52:10.642Z" }, + { url = "https://files.pythonhosted.org/packages/49/44/be9f979a6ab6b8c36865ccd92c3a38a760c66055e1f384665f35525134c4/pydantic_core-2.46.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:f78eb3d4027963bdc9baccd177f02a98bf8714bc51fe17153d8b51218918b5bc", size = 2385228, upload-time = "2026-04-15T14:51:00.77Z" }, + { url = "https://files.pythonhosted.org/packages/5b/d4/c826cd711787d240219f01d0d3ca116cb55516b8b95277820aa9c85e1882/pydantic_core-2.46.1-cp312-cp312-win32.whl", hash = "sha256:54fe30c20cab03844dc63bdc6ddca67f74a2eb8482df69c1e5f68396856241be", size = 1978828, upload-time = "2026-04-15T14:50:29.362Z" }, + { url = "https://files.pythonhosted.org/packages/22/05/8a1fcf8181be4c7a9cfc34e5fbf2d9c3866edc9dfd3c48d5401806e0a523/pydantic_core-2.46.1-cp312-cp312-win_amd64.whl", hash = "sha256:aea4e22ed4c53f2774221435e39969a54d2e783f4aee902cdd6c8011415de893", size = 2070015, upload-time = "2026-04-15T14:49:47.301Z" }, + { url = "https://files.pythonhosted.org/packages/61/d5/fea36ad2882b99c174ef4ffbc7ea6523f6abe26060fbc1f77d6441670232/pydantic_core-2.46.1-cp312-cp312-win_arm64.whl", hash = "sha256:f76fb49c34b4d66aa6e552ce9e852ea97a3a06301a9f01ae82f23e449e3a55f8", size = 2030176, upload-time = "2026-04-15T14:50:47.307Z" }, + { url = "https://files.pythonhosted.org/packages/ff/d2/bda39bad2f426cb5078e6ad28076614d3926704196efe0d7a2a19a99025d/pydantic_core-2.46.1-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:cdc8a5762a9c4b9d86e204d555444e3227507c92daba06259ee66595834de47a", size = 2119092, upload-time = "2026-04-15T14:49:50.392Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f3/69631e64d69cb3481494b2bddefe0ddd07771209f74e9106d066f9138c2a/pydantic_core-2.46.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:ba381dfe9c85692c566ecb60fa5a77a697a2a8eebe274ec5e4d6ec15fafad799", size = 1951400, upload-time = "2026-04-15T14:51:06.588Z" }, + { url = "https://files.pythonhosted.org/packages/53/1c/21cb3db6ae997df31be8e91f213081f72ffa641cb45c89b8a1986832b1f9/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1593d8de98207466dc070118322fef68307a0cc6a5625e7b386f6fdae57f9ab6", size = 1976864, upload-time = "2026-04-15T14:50:54.804Z" }, + { url = "https://files.pythonhosted.org/packages/91/9c/05c819f734318ce5a6ca24da300d93696c105af4adb90494ee571303afd8/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8262c74a1af5b0fdf795f5537f7145785a63f9fbf9e15405f547440c30017ed8", size = 2066669, upload-time = "2026-04-15T14:51:42.346Z" }, + { url = "https://files.pythonhosted.org/packages/cb/23/fadddf1c7f2f517f58731aea9b35c914e6005250f08dac9b8e53904cdbaa/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4b88949a24182e83fbbb3f7ca9b7858d0d37b735700ea91081434b7d37b3b444", size = 2238737, upload-time = "2026-04-15T14:50:45.558Z" }, + { url = "https://files.pythonhosted.org/packages/23/07/0cd4f95cb0359c8b1ec71e89c3777e7932c8dfeb9cd54740289f310aaead/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b8f3708cd55537aeaf3fd0ea55df0d68d0da51dcb07cbc8508745b34acc4c6e0", size = 2316258, upload-time = "2026-04-15T14:51:08.471Z" }, + { url = "https://files.pythonhosted.org/packages/0c/40/6fc24c3766a19c222a0d60d652b78f0283339d4cd4c173fab06b7ee76571/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f79292435fff1d4f0c18d9cfaf214025cc88e4f5104bfaed53f173621da1c743", size = 2097474, upload-time = "2026-04-15T14:49:56.543Z" }, + { url = "https://files.pythonhosted.org/packages/4b/af/f39795d1ce549e35d0841382b9c616ae211caffb88863147369a8d74fba9/pydantic_core-2.46.1-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:a2e607aeb59cf4575bb364470288db3b9a1f0e7415d053a322e3e154c1a0802e", size = 2168383, upload-time = "2026-04-15T14:51:29.269Z" }, + { url = "https://files.pythonhosted.org/packages/e6/32/0d563f74582795779df6cc270c3fc220f49f4daf7860d74a5a6cda8491ff/pydantic_core-2.46.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:ec5ca190b75878a9f6ae1fc8f5eb678497934475aef3d93204c9fa01e97370b6", size = 2186182, upload-time = "2026-04-15T14:50:19.097Z" }, + { url = "https://files.pythonhosted.org/packages/5c/07/1c10d5ce312fc4cf86d1e50bdcdbb8ef248409597b099cab1b4bb3a093f7/pydantic_core-2.46.1-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:1f80535259dcdd517d7b8ca588d5ca24b4f337228e583bebedf7a3adcdf5f721", size = 2187859, upload-time = "2026-04-15T14:49:22.974Z" }, + { url = "https://files.pythonhosted.org/packages/92/01/e1f62d4cb39f0913dbf5c95b9b119ef30ddba9493dff8c2b012f0cdd67dc/pydantic_core-2.46.1-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:24820b3c82c43df61eca30147e42853e6c127d8b868afdc0c162df829e011eb4", size = 2338372, upload-time = "2026-04-15T14:49:53.316Z" }, + { url = "https://files.pythonhosted.org/packages/44/ed/218dfeea6127fb1781a6ceca241ec6edf00e8a8933ff331af2215975a534/pydantic_core-2.46.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:f12794b1dd8ac9fb66619e0b3a0427189f5d5638e55a3de1385121a9b7bf9b39", size = 2384039, upload-time = "2026-04-15T14:53:04.929Z" }, + { url = "https://files.pythonhosted.org/packages/6c/1e/011e763cd059238249fbd5780e0f8d0b04b47f86c8925e22784f3e5fc977/pydantic_core-2.46.1-cp313-cp313-win32.whl", hash = "sha256:9bc09aed935cdf50f09e908923f9efbcca54e9244bd14a5a0e2a6c8d2c21b4e9", size = 1977943, upload-time = "2026-04-15T14:52:17.969Z" }, + { url = "https://files.pythonhosted.org/packages/8c/06/b559a490d3ed106e9b1777b8d5c8112dd8d31716243cd662616f66c1f8ea/pydantic_core-2.46.1-cp313-cp313-win_amd64.whl", hash = "sha256:fac2d6c8615b8b42bee14677861ba09d56ee076ba4a65cfb9c3c3d0cc89042f2", size = 2068729, upload-time = "2026-04-15T14:53:07.288Z" }, + { url = "https://files.pythonhosted.org/packages/9f/52/32a198946e2e19508532aa9da02a61419eb15bd2d96bab57f810f2713e31/pydantic_core-2.46.1-cp313-cp313-win_arm64.whl", hash = "sha256:f978329f12ace9f3cb814a5e44d98bbeced2e36f633132bafa06d2d71332e33e", size = 2029550, upload-time = "2026-04-15T14:52:22.707Z" }, + { url = "https://files.pythonhosted.org/packages/44/4b/1952d38a091aa7572c13460db4439d5610a524a1a533fb131e17d8eff9c2/pydantic_core-2.46.1-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:c56887c0ffa05318128a80303c95066a9d819e5e66d75ff24311d9e0a58d6930", size = 2123089, upload-time = "2026-04-15T14:50:20.658Z" }, + { url = "https://files.pythonhosted.org/packages/90/06/f3623aa98e2d7cb4ed0ae0b164c5d8a1b86e5aca01744eba980eefcd5da4/pydantic_core-2.46.1-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:614b24b875c1072631065fa85e195b40700586afecb0b27767602007920dacf8", size = 1945481, upload-time = "2026-04-15T14:50:56.945Z" }, + { url = "https://files.pythonhosted.org/packages/69/f9/a9224203b8426893e22db2cf0da27cd930ad7d76e0a611ebd707e5e6c916/pydantic_core-2.46.1-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a6382f6967c48519b6194e9e1e579e5898598b682556260eeaf05910400d827e", size = 1986294, upload-time = "2026-04-15T14:49:31.839Z" }, + { url = "https://files.pythonhosted.org/packages/96/29/954d2174db68b9f14292cef3ae8a05a25255735909adfcf45ca768023713/pydantic_core-2.46.1-graalpy311-graalpy242_311_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:93cb8aa6c93fb833bb53f3a2841fbea6b4dc077453cd5b30c0634af3dee69369", size = 2144185, upload-time = "2026-04-15T14:52:39.449Z" }, + { url = "https://files.pythonhosted.org/packages/f4/97/95de673a1356a88b2efdaa120eb6af357a81555c35f6809a7a1423ff7aef/pydantic_core-2.46.1-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:5f9107a24a4bc00293434dfa95cf8968751ad0dd703b26ea83a75a56f7326041", size = 2107564, upload-time = "2026-04-15T14:50:49.14Z" }, + { url = "https://files.pythonhosted.org/packages/00/fc/a7c16d85211ea9accddc693b7d049f20b0c06440d9264d1e1c074394ee6c/pydantic_core-2.46.1-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:2b1801ba99876984d0a03362782819238141c4d0f3f67f69093663691332fc35", size = 1939925, upload-time = "2026-04-15T14:50:36.188Z" }, + { url = "https://files.pythonhosted.org/packages/2e/23/87841169d77820ddabeb81d82002c95dcb82163846666d74f5bdeeaec750/pydantic_core-2.46.1-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b7fd82a91a20ed6d54fa8c91e7a98255b1ff45bf09b051bfe7fe04eb411e232e", size = 1995313, upload-time = "2026-04-15T14:50:22.538Z" }, + { url = "https://files.pythonhosted.org/packages/ea/96/b46609359a354fa9cd336fc5d93334f1c358b756cc81e4b397347a88fa6f/pydantic_core-2.46.1-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f135bf07c92c93def97008bc4496d16934da9efefd7204e5f22a2c92523cb1f", size = 2151197, upload-time = "2026-04-15T14:51:22.925Z" }, + { url = "https://files.pythonhosted.org/packages/f5/e7/3d1d2999ad8e78b124c752e4fc583ecd98f3bea7cc42045add2fb6e31b62/pydantic_core-2.46.1-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:b44b44537efbff2df9567cd6ba51b554d6c009260a021ab25629c81e066f1683", size = 2121103, upload-time = "2026-04-15T14:52:59.537Z" }, + { url = "https://files.pythonhosted.org/packages/de/08/50a56632994007c7a58c86f782accccbe2f3bb7ca80f462533e26424cd18/pydantic_core-2.46.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:8f9ca3af687cc6a5c89aeaa00323222fcbceb4c3cdc78efdac86f46028160c04", size = 1952464, upload-time = "2026-04-15T14:52:04.001Z" }, + { url = "https://files.pythonhosted.org/packages/75/0b/3cf631e33a55b1788add3e42ac921744bd1f39279082a027b4ef6f48bd32/pydantic_core-2.46.1-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e2678a4cbc205f00a44542dca19d15c11ccddd7440fd9df0e322e2cae55bb67a", size = 2138504, upload-time = "2026-04-15T14:52:01.812Z" }, + { url = "https://files.pythonhosted.org/packages/fa/69/f96f3dfc939450b9aeb80d3fe1943e7bc0614b14e9447d84f48d65153e0c/pydantic_core-2.46.1-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:e5a98cbb03a8a7983b0fb954e0af5e7016587f612e6332c6a4453f413f1d1851", size = 2165467, upload-time = "2026-04-15T14:52:15.455Z" }, + { url = "https://files.pythonhosted.org/packages/a8/22/bb61cccddc2ce85b179cd81a580a1746e880870060fbf4bf6024dab7e8aa/pydantic_core-2.46.1-pp311-pypy311_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:b2f098b08860bd149e090ad232f27fffb5ecf1bfd9377015445c8e17355ec2d1", size = 2183882, upload-time = "2026-04-15T14:51:50.868Z" }, + { url = "https://files.pythonhosted.org/packages/0e/01/b9039da255c5fd3a7fd85344fda8861c847ad6d8fdd115580fa4505b2022/pydantic_core-2.46.1-pp311-pypy311_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:d2623606145b55a96efdd181b015c0356804116b2f14d3c2af4832fe4f45ed5f", size = 2323011, upload-time = "2026-04-15T14:49:40.32Z" }, + { url = "https://files.pythonhosted.org/packages/24/b1/f426b20cb72d0235718ccc4de3bc6d6c0d0c2a91a3fd2f32ae11b624bcc9/pydantic_core-2.46.1-pp311-pypy311_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:420f515c42aaec607ff720867b300235bd393abd709b26b190ceacb57a9bfc17", size = 2365696, upload-time = "2026-04-15T14:49:41.936Z" }, + { url = "https://files.pythonhosted.org/packages/ef/d2/d2b0025246481aa2ce6db8ba196e29b92063343ac76e675b3a1fa478ed4d/pydantic_core-2.46.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:375cfdd2a1049910c82ba2ff24f948e93599a529e0fdb066d747975ca31fc663", size = 2190970, upload-time = "2026-04-15T14:49:33.111Z" }, ] [[package]] @@ -3493,6 +3718,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, ] +[[package]] +name = "termcolor" +version = "3.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/46/79/cf31d7a93a8fdc6aa0fbb665be84426a8c5a557d9240b6239e9e11e35fc5/termcolor-3.3.0.tar.gz", hash = "sha256:348871ca648ec6a9a983a13ab626c0acce02f515b9e1983332b17af7979521c5", size = 14434, upload-time = "2025-12-29T12:55:21.882Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/33/d1/8bb87d21e9aeb323cc03034f5eaf2c8f69841e40e4853c2627edf8111ed3/termcolor-3.3.0-py3-none-any.whl", hash = "sha256:cf642efadaf0a8ebbbf4bc7a31cec2f9b5f21a9f726f4ccbb08192c9c26f43a5", size = 7734, upload-time = "2025-12-29T12:55:20.718Z" }, +] + [[package]] name = "tiktoken" version = "0.12.0" @@ -3717,7 +3951,7 @@ wheels = [ [[package]] name = "tox" -version = "4.52.1" +version = "4.53.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cachetools" }, @@ -3733,9 +3967,9 @@ dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.11'" }, { name = "virtualenv" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bc/fb/d7d634eb513f741ffd40f4c262b7feea19d5c616882eb554045c620670a6/tox-4.52.1.tar.gz", hash = "sha256:297e71ea0ae4ef3acc45cb5fdf080b74537e6ecb5eea7d4646fa7322ca10473e", size = 273730, upload-time = "2026-04-09T16:46:45.838Z" } +sdist = { url = "https://files.pythonhosted.org/packages/04/01/d87a00063fa670ce4c48a9706b615a95ddf2c9ef5558d43af6071f166fd4/tox-4.53.0.tar.gz", hash = "sha256:62c780e42f87d34ee60f2ea20342156253794fdcbd6885fd797d98ee05009f22", size = 274048, upload-time = "2026-04-14T13:44:13.782Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3a/70/0d4fb1eefa05a24ca2f58272b4c4718090dd5ed7e38b54b9a7e757bfafc8/tox-4.52.1-py3-none-any.whl", hash = "sha256:3c4eef0a64f319df0b67dacdb7edcfeda87c8cc722581af5d98dd54f3ffdd8ef", size = 212179, upload-time = "2026-04-09T16:46:44.5Z" }, + { url = "https://files.pythonhosted.org/packages/16/03/02e2a03f3756cfb66e7e1bac41b06953f12cec75ddb961d56695d4d43dc4/tox-4.53.0-py3-none-any.whl", hash = "sha256:cc4e716d18c4889aa179d785175c438fa60c35deef20ce689ec288d8fb656096", size = 212164, upload-time = "2026-04-14T13:44:11.997Z" }, ] [[package]] @@ -3764,7 +3998,7 @@ wheels = [ [[package]] name = "transformers" -version = "5.5.3" +version = "5.5.4" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, @@ -3778,9 +4012,21 @@ dependencies = [ { name = "tqdm" }, { name = "typer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/af/35/cd5b0d1288e65d2c12db4ce84c1ec1074f7ee9bced040de6c9d69e70d620/transformers-5.5.3.tar.gz", hash = "sha256:3f60128e840b40d352655903552e1eed4f94ed49369a4d43e1bc067bd32d3f50", size = 8226047, upload-time = "2026-04-09T15:52:56.231Z" } +sdist = { url = "https://files.pythonhosted.org/packages/a5/1e/1e244ab2ab50a863e6b52cc55761910567fa532b69a6740f6e99c5fdbd98/transformers-5.5.4.tar.gz", hash = "sha256:2e67cadba81fc7608cc07c4dd54f524820bc3d95b1cabd0ef3db7733c4f8b82e", size = 8227649, upload-time = "2026-04-13T16:55:55.181Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/29/fb/162a66789c65e5afa3b051309240c26bf37fbc8fea285b4546ae747995a2/transformers-5.5.4-py3-none-any.whl", hash = "sha256:0bd6281b82966fe5a7a16f553ea517a9db1dee6284d7cb224dfd88fc0dd1c167", size = 10236696, upload-time = "2026-04-13T16:55:51.497Z" }, +] + +[[package]] +name = "typeguard" +version = "4.5.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2b/e8/66e25efcc18542d58706ce4e50415710593721aae26e794ab1dec34fb66f/typeguard-4.5.1.tar.gz", hash = "sha256:f6f8ecbbc819c9bc749983cc67c02391e16a9b43b8b27f15dc70ed7c4a007274", size = 80121, upload-time = "2026-02-19T16:09:03.392Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/a1/0b/f8524551ab2d896dfaca74ddb70a4453d515bbf4ab5451c100c7788ae155/transformers-5.5.3-py3-none-any.whl", hash = "sha256:e48f3ec31dd96505e96e66b63a1e43e1ad7a65749e108d9227caaf51051cdb02", size = 10236257, upload-time = "2026-04-09T15:52:52.866Z" }, + { url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" }, ] [[package]] @@ -3853,7 +4099,7 @@ wheels = [ [[package]] name = "virtualenv" -version = "21.2.1" +version = "21.2.4" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "distlib" }, @@ -3862,9 +4108,9 @@ dependencies = [ { name = "python-discovery" }, { name = "typing-extensions", marker = "python_full_version < '3.11'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/97/c5/aff062c66b42e2183201a7ace10c6b2e959a9a16525c8e8ca8e59410d27a/virtualenv-21.2.1.tar.gz", hash = "sha256:b66ffe81301766c0d5e2208fc3576652c59d44e7b731fc5f5ed701c9b537fa78", size = 5844770, upload-time = "2026-04-09T18:47:11.482Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0c/98/3a7e644e19cb26133488caff231be390579860bbbb3da35913c49a1d0a46/virtualenv-21.2.4.tar.gz", hash = "sha256:b294ef68192638004d72524ce7ef303e9d0cf5a44c95ce2e54a7500a6381cada", size = 5850742, upload-time = "2026-04-14T22:15:31.438Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/20/0e/f083a76cb590e60dff3868779558eefefb8dfb7c9ed020babc7aa014ccbf/virtualenv-21.2.1-py3-none-any.whl", hash = "sha256:bd16b49c53562b28cf1a3ad2f36edb805ad71301dee70ddc449e5c88a9f919a2", size = 5828326, upload-time = "2026-04-09T18:47:09.331Z" }, + { url = "https://files.pythonhosted.org/packages/27/8d/edd0bd910ff803c308ee9a6b7778621af0d10252219ad9f19ef4d4982a61/virtualenv-21.2.4-py3-none-any.whl", hash = "sha256:29d21e941795206138d0f22f4e45ff7050e5da6c6472299fb7103318763861ac", size = 5831232, upload-time = "2026-04-14T22:15:29.342Z" }, ] [[package]] @@ -4212,9 +4458,9 @@ wheels = [ [[package]] name = "zipp" -version = "3.23.0" +version = "3.23.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e3/02/0f2892c661036d50ede074e376733dca2ae7c6eb617489437771209d4180/zipp-3.23.0.tar.gz", hash = "sha256:a07157588a12518c9d4034df3fbbee09c814741a33ff63c05fa29d26a2404166", size = 25547, upload-time = "2025-06-08T17:06:39.4Z" } +sdist = { url = "https://files.pythonhosted.org/packages/30/21/093488dfc7cc8964ded15ab726fad40f25fd3d788fd741cc1c5a17d78ee8/zipp-3.23.1.tar.gz", hash = "sha256:32120e378d32cd9714ad503c1d024619063ec28aad2248dc6672ad13edfa5110", size = 25965, upload-time = "2026-04-13T23:21:46.6Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2e/54/647ade08bf0db230bfea292f893923872fd20be6ac6f53b2b936ba839d75/zipp-3.23.0-py3-none-any.whl", hash = "sha256:071652d6115ed432f5ce1d34c336c0adfd6a884660d1e9712a256d3d3bd4b14e", size = 10276, upload-time = "2025-06-08T17:06:38.034Z" }, + { url = "https://files.pythonhosted.org/packages/08/8a/0861bec20485572fbddf3dfba2910e38fe249796cb73ecdeb74e07eeb8d3/zipp-3.23.1-py3-none-any.whl", hash = "sha256:0b3596c50a5c700c9cb40ba8d86d9f2cc4807e9bedb06bcdf7fac85633e444dc", size = 10378, upload-time = "2026-04-13T23:21:45.386Z" }, ] From 6ded36bcbb79e5992b0950a46bbd6868779f61d6 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Thu, 16 Apr 2026 19:21:32 +0100 Subject: [PATCH 03/30] Add dep check for ptq and runtime check for evaluation/deployment (#1240) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: ? PTQ: model-specific dependency support - Add EXTRA_PIP_DEPS support to the launcher's `ptq.sh` so models requiring extra pip packages (e.g., `mamba-ssm` for hybrid Mamba architectures like Nemotron) can install them automatically before running PTQ. Also updates the PTQ skill with a new Step 2.5 for detecting model-specific dependencies. Container registry auth checks - Add new section 6 covering auth detection for enroot/pyxis, Docker, and Singularity/Apptainer. Includes credential locations, how to add them, and common failure modes. - Add Step 7.5 with NEL default image table, DockerHub-first strategy with NGC fallback, and build-config CLI note. - Add auth check before remote SLURM deployment. ### Usage Set EXTRA_PIP_DEPS in the launcher YAML's environment section: ``` task_0: script: common/hf/ptq.sh args: - --repo nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - --local-dir /hf-local/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 - -- - --quant nvfp4 - --tasks quant environment: - EXTRA_PIP_DEPS: "mamba-ssm causal-conv1d" ``` ### Testing Tested end-to-end: NVFP4 quantization of `NVIDIA-Nemotron-3-Nano-30B-A3B-BF16` on a B200 cluster via the launcher. Job succeeded: mamba-ssm installed automatically, calibration completed (512 samples, 84s), checkpoint exported (18 GB, 2 safetensor shards). ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A - Did you write any new necessary tests?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit ## Release Notes * **Documentation** * Added container registry authentication verification workflow for SLURM deployments, including credential checks, verification commands, common failure symptoms, and remediation guidance. * Required credential validation before SLURM job submission and added SLURM-only verification steps with image fallback recommendations. * New dependency-checking step for models that use remote/trust_remote_code, plus guidance for resolving extra package requirements and tightened build-config guidance. * Updated PTQ launcher documentation to reference the new wrapper script. * **New Features** * Support for specifying extra pip dependencies during model processing via an environment variable. --------- Signed-off-by: Kai Xu --- .claude/skills/common/slurm-setup.md | 125 ++++++++++++++++++ .claude/skills/deployment/SKILL.md | 2 + .claude/skills/evaluation/SKILL.md | 36 ++++- .claude/skills/ptq/SKILL.md | 19 +++ .../skills/ptq/references/launcher-guide.md | 6 +- tools/launcher/common/hf/ptq.sh | 8 ++ 6 files changed, 191 insertions(+), 5 deletions(-) diff --git a/.claude/skills/common/slurm-setup.md b/.claude/skills/common/slurm-setup.md index 37b9fbd56ab..d2324084fcd 100644 --- a/.claude/skills/common/slurm-setup.md +++ b/.claude/skills/common/slurm-setup.md @@ -192,3 +192,128 @@ chmod -R g+rwX /path/to/.hf_cache/ ``` Scope `chmod` to only the directories the job needs — avoid world-writable paths on shared clusters. + +--- + +## 6. Container Registry Authentication + +**Before submitting any SLURM job that pulls a container image**, check that the cluster has credentials for the image's registry. Missing auth causes jobs to fail after waiting in the queue — a costly mistake. + +### Step 1: Detect the container runtime + +Different clusters use different container runtimes. Detect which is available: + +```bash +# On the cluster (or via ssh): +which enroot 2>/dev/null && echo "RUNTIME=enroot" +which docker 2>/dev/null && echo "RUNTIME=docker" +``` + +| Runtime | Typical clusters | SLURM integration | +| --- | --- | --- | +| **enroot/pyxis** | NVIDIA internal (DGX Cloud, EOS, Selene, GCP-NRT) | `srun --container-image` | +| **Docker** | Bare-metal / on-prem with GPU | `docker run` inside job script | + +### Step 2: Check credentials for the image's registry + +Determine the registry from the image URI: + +| Image pattern | Registry | +| --- | --- | +| `nvcr.io/nvidia/...` | NGC | +| `vllm/vllm-openai:...`, `lmsysorg/sglang:...`, or no registry prefix | DockerHub | +| `ghcr.io/...` | GitHub Container Registry | +| `docker.io/...` | DockerHub (explicit) | + +Then check credentials based on the runtime: + +#### enroot/pyxis + +```bash +grep -E '^\s*machine\s+' ~/.config/enroot/.credentials 2>/dev/null +``` + +Look for `machine ` lines: +- NGC → `machine nvcr.io` +- DockerHub → `machine auth.docker.io` +- GHCR → `machine ghcr.io` + +#### Docker + +```bash +cat ~/.docker/config.json 2>/dev/null | python3 -c "import json,sys; print('\n'.join(json.load(sys.stdin).get('auths', {}).keys()))" +``` + +Look for registry keys (`https://index.docker.io/v1/`, `nvcr.io`, `ghcr.io`). + +### Step 3: If credentials are missing + +**Do not submit the job.** Instead: + +1. Tell the user which registry and runtime need authentication +2. Show the fix for their runtime: + +**enroot/pyxis:** + +```bash +mkdir -p ~/.config/enroot + +# DockerHub (get token from https://hub.docker.com/settings/security) +cat >> ~/.config/enroot/.credentials << 'EOF' +machine auth.docker.io + login + password +EOF + +# NGC (get API key from https://org.ngc.nvidia.com/setup/api-keys) +cat >> ~/.config/enroot/.credentials << 'EOF' +machine nvcr.io + login $oauthtoken + password +EOF +``` + +**Docker:** + +```bash +# DockerHub (interactive prompt) +docker login + +# NGC (use --password-stdin to avoid exposing secrets in process list) +echo "$NGC_API_KEY" | docker login nvcr.io -u '$oauthtoken' --password-stdin +``` + +3. **Suggest an alternative image** on an authenticated registry. NVIDIA clusters typically have NGC auth pre-configured, so prefer NGC-hosted images: + +| DockerHub image | NGC alternative | +| --- | --- | +| `vllm/vllm-openai:latest` | `nvcr.io/nvidia/vllm:-py3` (check [NGC catalog](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/vllm) for latest tag) | +| `nvcr.io/nvidia/tensorrt-llm/release:` | Already NGC | + +> **Note:** NGC image tags follow `YY.MM-py3` format (e.g., `26.03-py3`). Not all DockerHub images have NGC equivalents. If no NGC alternative exists and DockerHub auth is missing, the user must add DockerHub credentials or pre-cache the image as a `.sqsh` file. + +4. After the user fixes auth or switches images, verify the image is **actually pullable** before submitting (credentials alone don't guarantee the image exists): + +```bash +# enroot — test pull (aborts after manifest fetch) +enroot import --output /dev/null docker://# 2>&1 | head -10 +# Success: shows "Fetching image manifest" + layer info +# Failure: shows "401 Unauthorized" or "404 Not Found" + +# docker +docker manifest inspect 2>&1 | head -5 + +# singularity +singularity pull --dry-run docker:// 2>&1 | head -5 +``` + +> **Important**: Credentials existing for a registry does NOT mean a specific image is accessible. The image may not exist, or the credentials may lack permissions for that repository. Always verify the specific image before submitting. + +### Common failure modes + +| Symptom | Runtime | Cause | Fix | +| --- | --- | --- | --- | +| `curl: (22) ... error: 401` | enroot | No credentials for registry | Add to `~/.config/enroot/.credentials` | +| `pyxis: failed to import docker image` | enroot | Auth failed or rate limit | Check credentials; DockerHub free: 100 pulls/6h per IP | +| `unauthorized: authentication required` | docker | No `docker login` | Run `docker login [registry]` | +| Image pulls on some nodes but not others | any | Cached on one node only | Pre-cache image or ensure auth on all nodes | diff --git a/.claude/skills/deployment/SKILL.md b/.claude/skills/deployment/SKILL.md index e0189cb12e3..6f3f9b56bde 100644 --- a/.claude/skills/deployment/SKILL.md +++ b/.claude/skills/deployment/SKILL.md @@ -174,6 +174,8 @@ All checks must pass before reporting success to the user. If a cluster config exists (`~/.config/modelopt/clusters.yaml` or `.claude/clusters.yaml`), or the user mentions running on a remote machine: +0. **Check container registry auth** — before submitting any SLURM job with a container image, verify credentials exist on the cluster per `skills/common/slurm-setup.md` section 6. If credentials are missing for the image's registry, ask the user to fix auth or switch to an image on an authenticated registry (e.g., NGC). **Do not submit until auth is confirmed.** + 1. **Source remote utilities:** ```bash diff --git a/.claude/skills/evaluation/SKILL.md b/.claude/skills/evaluation/SKILL.md index f8eab5561bc..41a59c59422 100644 --- a/.claude/skills/evaluation/SKILL.md +++ b/.claude/skills/evaluation/SKILL.md @@ -28,6 +28,7 @@ Config Generation Progress: - [ ] Step 5: Confirm tasks (iterative) - [ ] Step 6: Advanced - Multi-node (Data Parallel) - [ ] Step 7: Advanced - Interceptors +- [ ] Step 7.5: Check container registry auth (SLURM only) - [ ] Step 8: Run the evaluation ``` @@ -74,9 +75,9 @@ Prompt the user with "I'll ask you 5 questions to build the base config we'll ad 4. Safety & Security (like Garak and Safety Harness) 5. Multilingual (like MMATH, Global MMLU, MMLU-Prox) -DON'T ALLOW FOR ANY OTHER OPTIONS, only the ones listed above under each category (Execution, Deployment, Auto-export, Model type, Benchmarks). YOU HAVE TO GATHER THE ANSWERS for the 5 questions before you can build the base config. +Only accept options from the categories listed above (Execution, Deployment, Auto-export, Model type, Benchmarks). YOU HAVE TO GATHER THE ANSWERS for the 5 questions before you can build the base config. -> **Note:** These categories come from NEL's `build-config` CLI. If `nel skills build-config --help` shows different options than listed above, use the CLI's current options instead. +> **Note:** These categories come from NEL's `build-config` CLI. **Always run `nel skills build-config --help` first** to get the current options — they may differ from this list (e.g., `chat_reasoning` instead of separate `chat`/`reasoning`, `general_knowledge` instead of `standard`). When the CLI's current options differ from this list, prefer the CLI's options. When you have all the answers, run the script to build the base config: @@ -181,6 +182,36 @@ If the user needs multi-node evaluation (model >120B, or more throughput), read - The docs may show incorrect parameter names for logging. Use `max_logged_requests` and `max_logged_responses` (NOT `max_saved_*` or `max_*`). +**Step 7.5: Check container registry authentication (SLURM only)** + +NEL's default deployment images by framework: + +| Framework | Default image | Registry | +| --- | --- | --- | +| vLLM | `vllm/vllm-openai:latest` | DockerHub | +| SGLang | `lmsysorg/sglang:latest` | DockerHub | +| TRT-LLM | `nvcr.io/nvidia/tensorrt-llm/release:...` | NGC | +| Evaluation tasks | `nvcr.io/nvidia/eval-factory/*:26.03` | NGC | + +Before submitting, verify the cluster has credentials for the deployment image. See `skills/common/slurm-setup.md` section 6 for the full procedure. + +```bash +ssh "grep -E '^\s*machine\s+' ~/.config/enroot/.credentials 2>/dev/null" +``` + +**Decision flow (check before submitting):** +1. Check if the cluster has credentials for the default DockerHub image (see command above) +2. If DockerHub credentials exist → use the default image and submit +3. If DockerHub credentials are missing but can be added → add them (see `slurm-setup.md` section 6), then submit +4. If DockerHub credentials cannot be added → override `deployment.image` to the NGC alternative and submit: + + ```yaml + deployment: + image: nvcr.io/nvidia/vllm:-py3 # check https://catalog.ngc.nvidia.com/orgs/nvidia/containers/vllm for latest tag + ``` + +5. **Do not retry more than once** without fixing the auth issue + **Step 8: Run the evaluation** Print the following commands to the user. Propose to execute them in order to confirm the config works as expected before the full run. @@ -303,5 +334,6 @@ Config Generation Progress: - [ ] Step 5: Confirm tasks (iterative) - [ ] Step 6: Advanced - Multi-node (Data Parallel) - [ ] Step 7: Advanced - Interceptors +- [ ] Step 7.5: Check container registry auth (SLURM only) - [ ] Step 8: Run the evaluation ``` diff --git a/.claude/skills/ptq/SKILL.md b/.claude/skills/ptq/SKILL.md index 6849f8c94d2..c4c70651a8f 100644 --- a/.claude/skills/ptq/SKILL.md +++ b/.claude/skills/ptq/SKILL.md @@ -24,6 +24,24 @@ Check the support table in `examples/llm_ptq/README.md` for verified HF models. - **Listed** → supported, use `hf_ptq.py` (step 4A/4B) - **Not listed** → read `references/unsupported-models.md` to determine if `hf_ptq.py` can still work or if a custom script is needed (step 4C) +## Step 2.5 — Check for model-specific dependencies + +If the model uses `trust_remote_code` (check `config.json` for `auto_map`), inspect its custom Python files for imports not present in the container: + +```bash +grep -h "^from \|^import " /modeling_*.py | sort -u +``` + +**Known dependency patterns:** + +| Import found | Packages to install | +| --- | --- | +| `from mamba_ssm` / `from causal_conv1d` | `mamba-ssm causal-conv1d` (Mamba/hybrid models: NemotronH, Jamba) | + +If extra deps are needed: +- **Launcher (4B)**: set `EXTRA_PIP_DEPS` in the task's `environment` section — `ptq.sh` installs them automatically +- **Manual (4A)**: `unset PIP_CONSTRAINT && pip install ` before running `hf_ptq.py` + ## Step 3 — Choose quantization format **First**, check for a model-specific recipe: @@ -128,6 +146,7 @@ Validate the exported checkpoint's quantization pattern matches the recipe. Quan ## Common Pitfalls +- **Model-specific dependencies**: Models with `trust_remote_code` may import packages not in the container (e.g., `mamba-ssm` for hybrid Mamba models). See Step 2.5. Use `EXTRA_PIP_DEPS` env var with the launcher, or install manually before running `hf_ptq.py` - **Transformers version**: New models may need a newer version of transformers than what's installed. Check `config.json` for `transformers_version`. In containers, beware of `PIP_CONSTRAINT` blocking upgrades — see `references/slurm-setup-ptq.md` for workarounds - **Gated datasets**: Some calibration datasets require HF authentication. Ensure `HF_TOKEN` is set in the job environment, or use `--dataset cnn_dailymail` as a non-gated alternative - **NFS root_squash + Docker**: See `skills/common/slurm-setup.md` section 5 diff --git a/.claude/skills/ptq/references/launcher-guide.md b/.claude/skills/ptq/references/launcher-guide.md index fb8494cf8f0..542c4ade5b1 100644 --- a/.claude/skills/ptq/references/launcher-guide.md +++ b/.claude/skills/ptq/references/launcher-guide.md @@ -12,13 +12,13 @@ uv run launch.py --yaml hf_local= --yes # Local Docker ## HF Transformers PTQ Config -The launcher provides `common/hf_ptq/hf_ptq.sh` which wraps `hf_ptq.py`. Configure via environment variables: +The launcher provides `common/hf/ptq.sh` which wraps `hf_ptq.py`. Configure via environment variables: ```yaml job_name: _ pipeline: task_0: - script: common/hf_ptq/hf_ptq.sh + script: common/hf/ptq.sh environment: - HF_MODEL: - QFORMAT: @@ -75,7 +75,7 @@ The launcher SSHes to `SLURM_HOST` via `nemo_run.SSHTunnel`. If `identity` is om ## Known Issues - **UID mapping in Docker**: May cause `getpwuid` failures. Add `USER=user` and `LOGNAME=user` to environment. -- **Megatron-LM submodule**: Only needed for `MegatronLMQuantizeTask` (Megatron models). HF PTQ via `common/hf_ptq/hf_ptq.sh` does not require it. +- **Megatron-LM submodule**: Only needed for `MegatronLMQuantizeTask` (Megatron models). HF PTQ via `common/hf/ptq.sh` does not require it. ## Dry Run diff --git a/tools/launcher/common/hf/ptq.sh b/tools/launcher/common/hf/ptq.sh index b3bc80c309e..9822362311b 100755 --- a/tools/launcher/common/hf/ptq.sh +++ b/tools/launcher/common/hf/ptq.sh @@ -25,6 +25,14 @@ set -e +# Install extra pip dependencies if specified (e.g., mamba-ssm for hybrid Mamba models). +if [ -n "$EXTRA_PIP_DEPS" ]; then + echo "Installing extra dependencies: $EXTRA_PIP_DEPS" + unset PIP_CONSTRAINT + read -r -a _deps <<< "$EXTRA_PIP_DEPS" + pip install "${_deps[@]}" +fi + REPO="" LOCAL_DIR="" PTQ_ARGS=() From 04fcf24227e3133d05b3feb7deb08f2cf261201f Mon Sep 17 00:00:00 2001 From: Chenjie Luo <108829653+cjluo-nv@users.noreply.github.com> Date: Thu, 16 Apr 2026 16:18:38 -0700 Subject: [PATCH 04/30] Fix LLM deploy test failure by defaulting expert parallelism to 1 (#1273) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes TRT-LLM DeepEP kernel failures during LLM deployment on unsupported GPUs (e.g. Blackwell SM 12.0) by defaulting expert parallelism (`ep`) to 1 instead of auto-setting it to the GPU count for MoE models. Previously, when the model config contained expert-related keys, `ep` was automatically set to `torch.cuda.device_count()`, which triggered DeepEP kernel failures on GPUs that don't support it. Now `ep` defaults to 1 while still enabling attention data parallelism for MoE models. Expert parallelism can be enabled explicitly by the caller when the environment is known to support it. ### Testing - [x] Verified that the `llm_ptq` test passes with this fix on Blackwell GPUs. - [x] 2-gpu CI test triggered: https://github.com/NVIDIA/Model-Optimizer/actions/runs/24495054531/job/71588037727 ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A Signed-off-by: Chenjie Luo --- modelopt/deploy/llm/generate.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/modelopt/deploy/llm/generate.py b/modelopt/deploy/llm/generate.py index 0df8e94c449..0f649199ec2 100644 --- a/modelopt/deploy/llm/generate.py +++ b/modelopt/deploy/llm/generate.py @@ -109,12 +109,13 @@ def _find_max_position_embeddings(cfg: dict) -> int | None: if tp < 1: tp = torch.cuda.device_count() - # Check if any key in config contains both "num" and "experts" + # Force ep=1 to avoid TRT-LLM DeepEP kernel failures on unsupported GPUs + # (e.g. Blackwell SM 12.0). Expert parallelism can be enabled explicitly + # by the caller when the environment is known to support it. ep = 1 enable_attention_dp = False for k in config: if "num" in k and "experts" in k: - ep = torch.cuda.device_count() enable_attention_dp = True break From fe8c5178c7ad1faa10f1321208123fadc1bda255 Mon Sep 17 00:00:00 2001 From: Hrishith Thadicherla <99313418+hthadicherla@users.noreply.github.com> Date: Fri, 17 Apr 2026 09:16:48 +0530 Subject: [PATCH 05/30] Removed version fixes for torch transformers in windows ptq example requirements (#1275) ### What does this PR do? Type of change: Bug fix Removed version fixes for torch and transformers ### Testing Tested quantization with a couple of models . Working as expected. ## Summary by CodeRabbit * **Chores** * Relaxed dependency specs: removed strict pin for torch to allow latest compatible installs, and constrained transformers to <5.0.0 for broader compatibility and easier updates. --------- Signed-off-by: Hrishith Thadicherla Signed-off-by: Hrishith Thadicherla <99313418+hthadicherla@users.noreply.github.com> --- examples/windows/onnx_ptq/genai_llm/requirements.txt | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/windows/onnx_ptq/genai_llm/requirements.txt b/examples/windows/onnx_ptq/genai_llm/requirements.txt index 59b54311c08..92295a0d9eb 100644 --- a/examples/windows/onnx_ptq/genai_llm/requirements.txt +++ b/examples/windows/onnx_ptq/genai_llm/requirements.txt @@ -1,3 +1,3 @@ datasets>=2.14.5 -torch==2.9.0 -transformers==4.57.3 +torch +transformers<5.0.0 From d073d8d8e6665dc63cafb67594696ec43779ab0a Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 17 Apr 2026 09:47:56 +0530 Subject: [PATCH 06/30] Update codecov.yml (#1278) Dont allow more than 1% overall project coverage drop per PR. 2% was too much for such a large codebase ## Summary by CodeRabbit * **Chores** * Updated code coverage enforcement thresholds for pull requests to maintain stricter quality standards. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .github/codecov.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/codecov.yml b/.github/codecov.yml index 3c8819a15fe..24756fdcbb2 100644 --- a/.github/codecov.yml +++ b/.github/codecov.yml @@ -9,5 +9,5 @@ coverage: project: default: target: auto - threshold: 2% # Allow atmost 2% coverage drop from main branch. + threshold: 1% # Allow atmost 1% coverage drop from main branch. patch: false From 3162ff003f259a2bc3b3eedcc1d1a72b57a9b5af Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 17 Apr 2026 11:23:19 +0530 Subject: [PATCH 07/30] Update 0.43 release date in CHANGELOG.rst (#1277) As title ## Summary by CodeRabbit * **Chores** * Updated the release date for version 0.43 in the changelog. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fdd738590aa..5c189bd28b5 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -32,7 +32,7 @@ Changelog - [Experimental] Add support for transformers>=5.0, including generic PTQ and unified HF checkpoint export for fused MoE expert modules (Mixtral, Qwen2-MoE, Qwen3-MoE, Qwen3.5-MoE, DeepSeek-V3, Jamba, OLMoE, etc.). - Improve ``megatron_preprocess_data``: add ``--reasoning_content`` support for Nemotron v3 datasets, eliminate intermediate JSONL for HuggingFace datasets, return output file prefixes from the Python API, add gzip input support (``.jsonl.gz``), add ``--strip_newlines`` flag for plain-text pretraining data, add ``--hf_streaming`` for very large datasets (only consumed rows downloaded), and auto-shuffle when ``--hf_max_samples_per_split`` is set to avoid biased sampling. -0.43 (2026-04-09) +0.43 (2026-04-16) ^^^^^^^^^^^^^^^^^ **Bug Fixes** From 7e82a5cb03b0f9081882e9ad87d363df5b112a5a Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 17 Apr 2026 11:54:35 +0530 Subject: [PATCH 08/30] [Serialization]: remove explicit weights_only default from safe_load to allow user to bypass if needed (#1279) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Remove the `kwargs.setdefault("weights_only", True)` call from `safe_load`, deferring to torch's built-in default (which is `True` for torch>=2.6) - This allows users to override via the `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1` env var when they trust a checkpoint but hit `pickle.UnpicklingError` - Add a test that verifies the default fails on unsafe objects and the env var bypass works ## Test plan - [x] `python -m pytest tests/unit/torch/utils/test_serialization.py -v` 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Serialization utility now respects PyTorch's default behavior and environment-variable configuration instead of forcibly enforcing parameter overrides, providing greater configuration flexibility. * **Tests** * Added test coverage validating environment-variable override functionality and default behavior in the serialization utility. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 --- modelopt/torch/utils/serialization.py | 6 +++-- tests/unit/torch/utils/test_serialization.py | 24 ++++++++++++++++++++ 2 files changed, 28 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/utils/serialization.py b/modelopt/torch/utils/serialization.py index da16f7514cc..dc880b86b80 100644 --- a/modelopt/torch/utils/serialization.py +++ b/modelopt/torch/utils/serialization.py @@ -54,9 +54,11 @@ def safe_save(obj: Any, f: str | os.PathLike | BinaryIO, **kwargs) -> None: def safe_load(f: str | os.PathLike | BinaryIO | bytes, **kwargs) -> Any: - """Load a checkpoint securely using weights_only=True by default.""" - kwargs.setdefault("weights_only", True) + """Load a checkpoint securely using ``weights_only=True`` by default. + NOTE: We dont set default ``weights_only`` (interpret as True for torch>=2.6) so you can override it with + ``export TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1`` if you see ``pickle.UnpicklingError`` and trust the checkpoint. + """ if isinstance(f, (bytes, bytearray)): f = BytesIO(f) diff --git a/tests/unit/torch/utils/test_serialization.py b/tests/unit/torch/utils/test_serialization.py index 32851d3a09e..cb3233739cd 100644 --- a/tests/unit/torch/utils/test_serialization.py +++ b/tests/unit/torch/utils/test_serialization.py @@ -16,7 +16,9 @@ """Tests for Modelopt's serialization utilities.""" from io import BytesIO +from pickle import UnpicklingError +import pytest import torch from modelopt.torch.opt.config import ModeloptBaseConfig @@ -70,3 +72,25 @@ def test_safe_load_with_path(tmp_path): loaded_state = safe_load(file_path) assert loaded_state["data"] == 42 + + +class _UnsafeObj: + """Not registered in torch safe globals — unpickling fails with weights_only=True.""" + + def __init__(self, v): + self.v = v + + +def test_safe_load_env_var_bypasses_weights_only(tmp_path, monkeypatch): + """Verify TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD=1 allows safe_load to load objects unsafe for weights_only.""" + file_path = tmp_path / "unsafe.pt" + torch.save({"obj": _UnsafeObj(42)}, file_path) + + # Always fails when weights_only is not set (default=True) + with pytest.raises(UnpicklingError): + safe_load(file_path) + + # With the env var, safe_load (no explicit weights_only) defers to torch's default=False + monkeypatch.setenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1") + loaded = safe_load(file_path) + assert loaded["obj"].v == 42 From 4e33368dbe0bac9093a85069b745a4bb3ebe41fa Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Fri, 17 Apr 2026 16:31:19 +0530 Subject: [PATCH 09/30] Temporarily disable latest pulp and mcore until we fix its nvidia-resiliency-ext dependency (#1285) - `megatron-core==0.17.0` released yesterday which requires nightly version of `nvidia-resiliency-ext` for an import. Pre-installed version in DLFW Pytorch container is `nvidia-resiliency-ext==0.5.0` - Temporarily pin `mcore<0.17.0` to unblock PR from merging. - Pin `pulp<4.0` as it has some breaking changes and release imminent Correct fix is to just use `nemo:26.04` container instead of PyTorch container for megatron-based tests since it always has correct combination of all packages needed for the megatron ecosystem - Done in #1286 --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- pyproject.toml | 2 +- tox.ini | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 16da6d6dc61..2993759ec10 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -43,7 +43,7 @@ dependencies = [ # modelopt.torch "PyYAML>=6.0", "omegaconf>=2.3.0", - "pulp", + "pulp<4.0", # breaking changes in upcoming 4.0 release "pydantic>=2.0", "regex", "rich", diff --git a/tox.ini b/tox.ini index 6694f7349d3..0cd99cc86f1 100644 --- a/tox.ini +++ b/tox.ini @@ -82,7 +82,8 @@ commands = [testenv:cuda13-gpu-megatron] commands_pre = # Install deps here so that it gets installed even in --current-env - pip install -U megatron-core + # Temporarily disable latest mcore until we fix its nvidia-resiliency-ext dependency + pip install 'megatron-core<0.17.0' pip install --no-build-isolation git+https://github.com/state-spaces/mamba.git pip install --no-build-isolation git+https://github.com/Dao-AILab/causal-conv1d.git pip install -e .[hf,dev-test] From e4b054bf326f72a572cb93a7d150d6ea1e9248d5 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Sat, 18 Apr 2026 00:50:54 +0530 Subject: [PATCH 10/30] Fix and Speedup megatron_mmlu by >10x via prefill scoring and global batching (#1280) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: new feature + bug fix Two improvements to Megatron inference utilities: **1. Pipeline Parallel (PP) correctness fixes** PP inference was producing garbage output (MMLU ~0.24, random chance). Two root causes: - `megatron_generate` / `megatron_prefill` used `get_forward_backward_func()` (the training pipeline scheduler), which is not designed for inference. Rewrote both functions to use explicit P2P communication via `recv_from_prev_pipeline_rank_` / `send_to_next_pipeline_rank`, matching the `run_mcore_inference` pattern. - `import_mcore_gpt_from_hf` loads HF weights into stage 0's embedding but never updates the output_layer on the last PP stage when `share_embeddings_and_output_weights=True`. At model init, `setup_embeddings_and_output_layer()` all-reduces from stage 0 to sync the output layer; after importing HF weights that all-reduce is stale. Fix: call `model.setup_embeddings_and_output_layer()` again after import. **2. `megatron_mmlu` speedup (~6x)** Replaces the `megatron_mmlu` implementation with a significantly faster approach that matches how `lm-evaluation-harness` scores multiple-choice questions. **Before:** autoregressive generation (`megatron_generate`, `osl=2`) per example, 114 separate `load_dataset` calls, batch_size=1 — 260s for 5% data. **After:** single prefill forward pass + argmax over {A,B,C,D} logits, 2 `load_dataset` calls, configurable batch_size — 18s for 5% data (~6x faster). ### Changes **PP fixes:** - `megatron_generate` / `megatron_prefill`: replace `get_forward_backward_func` with explicit P2P (`recv_from_prev_pipeline_rank_` / `send_to_next_pipeline_rank`) - `import_mcore_gpt_from_hf`: call `model.setup_embeddings_and_output_layer()` after HF weight import when PP>1 and `share_embeddings_and_output_weights=True` - `megatron_prefill`: add `skip_return_logits` param and VLM support (needed for PP non-last stages) **MMLU speedup:** - **Log-likelihood scoring**: replace `megatron_generate` with `megatron_prefill` — one forward pass per batch, no autoregressive decode loop - **Global batching**: collect all examples across all subjects, sort by descending sequence length, run in `batch_size` chunks - **2 dataset loads** instead of 114: use `load_dataset("cais/mmlu", "all")` with per-subject grouping; skip dev load when `few_shots=0` - **`percentage` → `fraction`** parameter rename for clarity - **tqdm progress bar** (rank-0 only) ### Testing - `test_megatron_generate_and_mmlu` parametrized over `tp` and `pp`. Accuracy assertion: `0.36 < score < 0.39`. Manually checked generated text is coherent. - Re-ran M-Bridge Minitron MMLU based pruning for Nano v2 9B -> 7B and all top 10 candidate's MMLU numbers are ballpark similar as before ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ❌ — `percentage` parameter renamed to `fraction`; `enable_kv_cache` removed from `megatron_mmlu` - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ — existing test updated and parametrized for TP+PP - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ 🤖 Generated with [Claude Code](https://claude.ai/claude-code) ## Summary by CodeRabbit * **Bug Fixes** * Improved pipeline-parallel generation and MMLU evaluation reliability; fixed output-layer synchronization in shared-embedding + pipeline setups. * **New Features** * MMLU scoring now uses batched prefill logit scoring for faster, batched evaluation. * **Behavior Changes** * Default MMLU sampling increased from 5% to 10%; calibration batch sizing adjusted and related CLI/help text updated. * **Tests** * Distributed tests cover tensor- and pipeline-parallel modes and tighten MMLU validation ranges. * **Documentation** * Updated pruning example and benchmark timing to reflect new sampling and speedup. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 --- CHANGELOG.rst | 1 + examples/megatron_bridge/prune_minitron.py | 18 +- examples/pruning/README.md | 4 +- .../torch/export/plugins/megatron_importer.py | 11 + .../torch/prune/plugins/mcore_minitron.py | 3 + modelopt/torch/utils/logging.py | 3 +- .../torch/utils/plugins/megatron_generate.py | 321 +++++++++--------- modelopt/torch/utils/plugins/megatron_mmlu.py | 191 ++++++----- .../utils/plugins/test_utils_megatron.py | 27 +- 9 files changed, 325 insertions(+), 254 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 5c189bd28b5..cd8e48e0857 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -22,6 +22,7 @@ Changelog **Bug Fixes** +- Fix Megatron utility functions for generation (with pipeline parallelism) and ~10x speedup in MMLU score evaluation (by batching prefill passes). - Fix Minitron pruning (``mcore_minitron``) for MoE models. Importance estimation hooks were incorrectly registered for MoE modules and NAS step was hanging before this. - Fix TRT support for remote autotuning in ONNX Autotune from 10.16+ to 10.15+ and fix TRT versioning check to the ``trtexec`` version instead of the TRT Python API when using ``trtexec`` backend. diff --git a/examples/megatron_bridge/prune_minitron.py b/examples/megatron_bridge/prune_minitron.py index 445fdea863c..0fa9a658ff2 100644 --- a/examples/megatron_bridge/prune_minitron.py +++ b/examples/megatron_bridge/prune_minitron.py @@ -18,7 +18,7 @@ while skipping pruning of num_attention_heads using following defaults: 1024 samples from nemotron-post-training-dataset-v2 for calibration, at-most 20% depth (num_layers) and 40% width is pruned per prunable hparam (hidden_size, ffn_hidden_size, ...), - top-10 candidates are evaluated for MMLU score (5% sampled data) to select the best model. + top-10 candidates are evaluated for MMLU score (10% sampled data) to select the best model. torchrun --nproc_per_node 2 prune_minitron.py \ --hf_model_name_or_path Qwen/Qwen3-8B \ @@ -140,11 +140,11 @@ def get_args() -> argparse.Namespace: parser.add_argument( "--prune_score_func", type=str, - default="mmlu_5pct", + default="mmlu_10pct", help=( "Score function to use for NAS-based pruning (--prune_target_params). Only supports MMLU at the moment. " "Format: mmlu_pct where is the percentage of MMLU data to sample per subject " - "(e.g. mmlu_5pct for 5%, mmlu_100pct for full eval)." + "(e.g. mmlu_10pct for 10%, mmlu_100pct for full eval)." ), ) parser.add_argument( @@ -299,16 +299,14 @@ def main(args: argparse.Namespace): match = re.fullmatch(r"mmlu_(\d+)pct", args.prune_score_func) if not match: raise ValueError( - f"Invalid score function: {args.prune_score_func}. " - "Expected format: mmlu_pct (e.g. mmlu_5pct)" + f"Invalid score function: {args.prune_score_func}. Expected format: mmlu_pct (e.g. mmlu_10pct)" ) - mmlu_pct = int(match.group(1)) - if not 0 < mmlu_pct <= 100: - raise ValueError("--prune_score_func percentage must be in the range [1, 100].") - _mmlu_pct = mmlu_pct / 100.0 + mmlu_frac = float(match.group(1)) / 100.0 def score_func(m): - return megatron_mmlu(m, tokenizer, percentage=_mmlu_pct) + return megatron_mmlu( + m, tokenizer, few_shots=0, fraction=mmlu_frac, batch_size=args.calib_mbs + ) pruning_config["score_func"] = score_func pruning_config["max_width_pruning"] = args.max_width_pruning diff --git a/examples/pruning/README.md b/examples/pruning/README.md index 930e9c6d259..9e84622269c 100644 --- a/examples/pruning/README.md +++ b/examples/pruning/README.md @@ -124,7 +124,7 @@ This mode can be useful when you don't know the exact dimensions you want to pru from modelopt.torch.utils.plugins.megatron_mmlu import megatron_mmlu def score_func(m): - return megatron_mmlu(m, tokenizer, percentage=0.05) # 5% sampled data for faster eval + return megatron_mmlu(m, tokenizer, fraction=0.1, batch_size=4) # 10% sampled data for faster eval # Specify target parameter count and configure the auto pruning algorithm # Save minitron scores at checkpoint so we can resume pruning without running the forward loop again @@ -147,7 +147,7 @@ mtp.prune(...) 1. **Importance Scoring**: Same as manual pruning - computes activation magnitudes for all parameters (takes ~5 minutes for an 8B model) 2. **Search Space Construction**: Generates a search space of possible architectures based search space config and other configs (`max_width_pruning`, `max_depth_pruning`, `hparams_to_skip`) -3. **Architecture Search**: Find candidate architectures that meet the parameter constraint and evaluate `top_k` (based on number of parameters) of them using `score_func` e.g. MMLU, negative validation loss, etc. (takes ~10 mins per candidate for an 8B model pruning) +3. **Architecture Search**: Find candidate architectures that meet the parameter constraint and evaluate `top_k` (based on number of parameters) of them using `score_func` e.g. MMLU, negative validation loss, etc. (takes ~5 min per candidate for an 8B model MMLU score with 10% sampled data) 4. **Best Architecture Selection**: Returns the architecture (best `export_config`) with the highest actual score from the top-K evaluated architectures 5. **Weight Slicing**: Slices the model weights according to the best pruned architecture found diff --git a/modelopt/torch/export/plugins/megatron_importer.py b/modelopt/torch/export/plugins/megatron_importer.py index a156f2cd8cc..b1d37c1ad90 100644 --- a/modelopt/torch/export/plugins/megatron_importer.py +++ b/modelopt/torch/export/plugins/megatron_importer.py @@ -747,6 +747,17 @@ def _import_state_dict(self): if hasattr(model, "output_layer") and not model.share_embeddings_and_output_weights: self.rules["output_layer"](model.output_layer) + # For PP with shared embedding/output weights, re-sync the output layer on the last + # pipeline stage from stage 0's (now HF-loaded) embedding. At model init, + # setup_embeddings_and_output_layer() zeros out the last stage's weight and all-reduces + # from stage 0. After importing HF weights into stage 0's embedding, that sync is stale, + # so we re-run it here. + if ( + model.share_embeddings_and_output_weights + and model.config.pipeline_model_parallel_size > 1 + ): + model.setup_embeddings_and_output_layer() + # MTP if hasattr(model, "mtp"): layer_pbar.set_description("Importing MTP") diff --git a/modelopt/torch/prune/plugins/mcore_minitron.py b/modelopt/torch/prune/plugins/mcore_minitron.py index 204d30603e7..e99a44a7910 100644 --- a/modelopt/torch/prune/plugins/mcore_minitron.py +++ b/modelopt/torch/prune/plugins/mcore_minitron.py @@ -171,6 +171,9 @@ class CandidateSubnet: score: float | None +torch.serialization.add_safe_globals([CandidateSubnet]) + + class MCoreMinitronSearcher(BaseSearcher): """Searcher for Minitron pruning algorithm. diff --git a/modelopt/torch/utils/logging.py b/modelopt/torch/utils/logging.py index ada1b53612b..f5aba0d1a1f 100644 --- a/modelopt/torch/utils/logging.py +++ b/modelopt/torch/utils/logging.py @@ -105,8 +105,9 @@ def no_stdout(): def print_rank_0(*args, **kwargs): """Prints only on the master process.""" + kwargs.setdefault("flush", True) if dist.is_master(): - print(*args, **kwargs, flush=True) + print(*args, **kwargs) def warn_rank_0(message, *args, **kwargs): diff --git a/modelopt/torch/utils/plugins/megatron_generate.py b/modelopt/torch/utils/plugins/megatron_generate.py index 0891f58b5bc..5625013bb44 100644 --- a/modelopt/torch/utils/plugins/megatron_generate.py +++ b/modelopt/torch/utils/plugins/megatron_generate.py @@ -17,11 +17,15 @@ import torch from megatron.core import mpu -from megatron.core.inference.communication_utils import broadcast_from_last_pipeline_stage +from megatron.core.inference.communication_utils import ( + broadcast_from_last_pipeline_stage, + recv_from_prev_pipeline_rank_, + send_to_next_pipeline_rank, +) from megatron.core.inference.contexts import StaticInferenceContext -from megatron.core.pipeline_parallel import get_forward_backward_func from megatron.core.timers import Timer from megatron.core.transformer import MegatronModule +from megatron.core.utils import get_attr_wrapped_model from tqdm import tqdm __all__ = ["megatron_generate", "megatron_prefill"] @@ -48,81 +52,105 @@ def megatron_prefill( image_sizes: torch.LongTensor | None = None, skip_return_logits: bool = False, ) -> torch.Tensor: - """A simple prefill function for Megatron Core V(LM) models.""" + """A simple prefill function for Megatron Core V(LM) models. + + Supports TP, PP, SP, and combinations thereof. For PP, activations are communicated + explicitly between pipeline stages (rather than through get_forward_backward_func) + so that the training pipeline scheduler does not interfere with inference. + """ if not isinstance(model, MegatronModule): raise ValueError("megatron_prefill only supports Megatron Core models.") model.eval() - # Create a static inference context if KV-cache is enabled. - max_batch_size = input_ids.shape[0] + batch_size = input_ids.shape[0] seq_length = input_ids.shape[-1] + device = input_ids.device - def _dummy_loss_func(output_tensor, non_loss_data=True): - """Need a dummy loss function.""" - return output_tensor - - def _forward_step_func(data, model): - """Forward step function.""" - batch_size = data["tokens"].shape[0] - seq_len = data["tokens"].shape[-1] - device = data["tokens"].device - - # ModelOpt transoformer_spec by default use arbitrary attention mask type; hence we need to - # compute the attention_mask for prefilling. Alternatively, if "causal" attention mask type - # is used, the attention_mask is not needed. During generation, the attn_mask_type is overridden - # to "no_mask" by SelfAttention.forward() if inference_context is provided. - attention_mask = ( - torch.triu(torch.ones((batch_size, seq_len, seq_len), device=device), diagonal=1) - .bool() - .view(batch_size, 1, seq_len, seq_len) - ) - - position_ids = ( - torch.arange(seq_len, dtype=torch.long, device=device) - .unsqueeze(0) - .expand(batch_size, -1) - ) - - output_tensor = model( - data["tokens"], - position_ids, - attention_mask, - runtime_gather_output=True, - ) - return output_tensor, _dummy_loss_func + pp_first = mpu.is_pipeline_first_stage() + pp_last = mpu.is_pipeline_last_stage() + is_pp = not (pp_first and pp_last) + pp_dtype = model.config.pipeline_dtype or ( + torch.bfloat16 if model.config.bf16 else torch.float32 + ) if model.config.sequence_parallel: tp = model.config.tensor_model_parallel_size - num_pad_tokens = (tp - input_ids.shape[-1] % tp) % tp + num_pad_tokens = (tp - seq_length % tp) % tp else: num_pad_tokens = 0 if num_pad_tokens > 0: - padding_shape = (input_ids.shape[0], num_pad_tokens) - padded_tokens = torch.full(padding_shape, 0, dtype=input_ids.dtype, device=input_ids.device) - tokens = torch.cat((input_ids, padded_tokens), dim=-1) + tokens = torch.cat( + [ + input_ids, + torch.zeros(batch_size, num_pad_tokens, dtype=input_ids.dtype, device=device), + ], + dim=-1, + ) else: tokens = input_ids - list_of_logits = get_forward_backward_func()( - forward_step_func=_forward_step_func, - data_iterator=[{"tokens": tokens}], - model=model, - num_microbatches=1, - seq_length=tokens.shape[-1], - micro_batch_size=max_batch_size, - decoder_seq_length=tokens.shape[-1], - forward_only=True, - collect_non_loss_data=True, + padded_seq_len = tokens.shape[-1] + + # ModelOpt transformer_spec uses arbitrary attention mask type by default; the causal mask + # must be supplied explicitly for prefill. + attention_mask = ( + torch.triu( + torch.ones((batch_size, padded_seq_len, padded_seq_len), device=device), diagonal=1 + ) + .bool() + .view(batch_size, 1, padded_seq_len, padded_seq_len) + ) + position_ids = ( + torch.arange(padded_seq_len, dtype=torch.long, device=device) + .unsqueeze(0) + .expand(batch_size, -1) ) - if skip_return_logits: - return None - if mpu.is_pipeline_last_stage(): - logits = list_of_logits[0][:, :seq_length, :].detach() + # For PP, receive activations from the previous stage before calling forward. + if is_pp and not pp_first: + pp_dtype = model.config.pipeline_dtype or ( + torch.bfloat16 if model.config.bf16 else torch.float32 + ) + recv_buffer = torch.empty( + (padded_seq_len, batch_size, model.config.hidden_size), + dtype=pp_dtype, + device=device, + ) + recv_from_prev_pipeline_rank_(recv_buffer) + get_attr_wrapped_model(model, "set_input_tensor")(recv_buffer) + + has_vision_inputs = ( + pixel_values is not None or image_grid_thw is not None or image_sizes is not None + ) + if has_vision_inputs: + forward_kwargs: dict = { + "input_ids": tokens, + "position_ids": position_ids, + "attention_mask": torch.ones( + (batch_size, padded_seq_len), dtype=torch.bool, device=device + ), + "runtime_gather_output": True, + } + if pixel_values is not None: + forward_kwargs["pixel_values"] = pixel_values + if image_grid_thw is not None: + forward_kwargs["image_grid_thw"] = image_grid_thw + if image_sizes is not None: + forward_kwargs["image_sizes"] = image_sizes + output = model(**forward_kwargs) else: - logits = None + output = model(tokens, position_ids, attention_mask, runtime_gather_output=True) + + # For PP non-last stages, forward activations to the next stage and return early. + if is_pp and not pp_last: + pp_dtype = model.config.pipeline_dtype or ( + torch.bfloat16 if model.config.bf16 else torch.float32 + ) + send_to_next_pipeline_rank(output.to(dtype=pp_dtype)) + + logits = output[:, :seq_length, :].detach() if pp_last else None if model.config.bf16: logits_dtype = torch.bfloat16 @@ -130,11 +158,12 @@ def _forward_step_func(data, model): logits_dtype = torch.float16 else: logits_dtype = torch.float32 - logits = broadcast_from_last_pipeline_stage( - [max_batch_size, seq_length, model.vocab_size], logits_dtype, logits - ) - return logits + # All PP ranks must participate in the broadcast to stay in sync. + result = broadcast_from_last_pipeline_stage( + [batch_size, seq_length, model.vocab_size], logits_dtype, logits + ) + return None if skip_return_logits else result def megatron_generate( @@ -182,6 +211,13 @@ def megatron_generate( model.eval() + pp_first = mpu.is_pipeline_first_stage() + pp_last = mpu.is_pipeline_last_stage() + is_pp = not (pp_first and pp_last) + pp_dtype = model.config.pipeline_dtype or ( + torch.bfloat16 if model.config.bf16 else torch.float32 + ) + # Create a static inference context if KV-cache is enabled. max_batch_size = input_ids.shape[0] max_seq_len = input_ids.shape[-1] + osl @@ -189,20 +225,45 @@ def megatron_generate( StaticInferenceContext(max_batch_size, max_seq_len) if enable_kv_cache else None ) - def _dummy_loss_func(output_tensor, non_loss_data=True): - """Need a dummy loss function.""" - return output_tensor + disable_tqdm = disable_tqdm or torch.distributed.get_rank() > 0 + + output_ids = torch.tensor([]) + step_pbar = tqdm(range(osl), disable=disable_tqdm, leave=False) + + time_ttft = 0 + time_remaining_outputs = 0 + timer = Timer("generate") + timer.start(barrier=True) + + for step in step_pbar: + step_pbar.set_description(get_current_memory_info()) - def _forward_step_func(data, model): - """Forward step function.""" - batch_size = data["tokens"].shape[0] - seq_len = data["tokens"].shape[-1] - device = data["tokens"].device + if model.config.sequence_parallel: + tp = model.config.tensor_model_parallel_size + num_pad_tokens = (tp - input_ids.shape[-1] % tp) % tp + else: + num_pad_tokens = 0 - # ModelOpt transoformer_spec by default use arbitrary attention mask type; hence we need to - # compute the attention_mask for prefilling. Alternatively, if "causal" attention mask type - # is used, the attention_mask is not needed. During generation, the attn_mask_type is overridden - # to "no_mask" by SelfAttention.forward() if inference_context is provided. + if inference_context is not None and step > 0: + tokens = input_ids[:, -1:] + inference_context.enable_decode_mode() + num_pad_tokens = 0 + elif num_pad_tokens > 0: + padding_shape = (input_ids.shape[0], num_pad_tokens) + padded_tokens = torch.full( + padding_shape, 0, dtype=input_ids.dtype, device=input_ids.device + ) + tokens = torch.cat((input_ids, padded_tokens), dim=-1) + else: + tokens = input_ids + + batch_size = tokens.shape[0] + seq_len = tokens.shape[-1] + device = tokens.device + + # ModelOpt transformer_spec uses arbitrary attention mask type by default; compute causal + # mask for prefill. During decode, attn_mask_type is overridden to "no_mask" by + # SelfAttention.forward() when inference_context is provided. if seq_len > 1: attention_mask = ( torch.triu(torch.ones((batch_size, seq_len, seq_len), device=device), diagonal=1) @@ -218,109 +279,57 @@ def _forward_step_func(data, model): .expand(batch_size, -1) ) - # Check if this is a VLM model (has vision inputs) - _has_pixel_values = data.get("pixel_values") is not None - _has_image_grid_thw = data.get("image_grid_thw") is not None - _has_image_sizes = data.get("image_sizes") is not None + # Check if this is a VLM model (vision inputs only passed at step 0 / prefill) + _has_pixel_values = step == 0 and pixel_values is not None + _has_image_grid_thw = step == 0 and image_grid_thw is not None + _has_image_sizes = step == 0 and image_sizes is not None has_vision_inputs = _has_pixel_values or _has_image_grid_thw or _has_image_sizes - if has_vision_inputs: - # For VLM models: - # - position_ids: [batch, seq_len] (required for RoPE with multi-modal positions) - # - attention_mask: [batch, seq_len] (simple 1D boolean mask, not 4D causal) - vlm_position_ids = ( - torch.arange(seq_len, dtype=torch.long, device=device) - .unsqueeze(0) - .expand(batch_size, -1) + # For PP, receive activations from the previous stage before calling forward. + if is_pp and not pp_first: + recv_buffer = torch.empty( + (seq_len, batch_size, model.config.hidden_size), + dtype=pp_dtype, + device=device, ) - vlm_attention_mask = torch.ones((batch_size, seq_len), dtype=torch.bool, device=device) + recv_from_prev_pipeline_rank_(recv_buffer) + get_attr_wrapped_model(model, "set_input_tensor")(recv_buffer) + if has_vision_inputs: forward_args = { - "input_ids": data["tokens"], - "position_ids": vlm_position_ids, - "attention_mask": vlm_attention_mask, + "input_ids": tokens, + "position_ids": position_ids, + "attention_mask": torch.ones( + (batch_size, seq_len), dtype=torch.bool, device=device + ), "inference_context": inference_context, "runtime_gather_output": True, } - # Add vision inputs if _has_pixel_values: - forward_args["pixel_values"] = data["pixel_values"] + forward_args["pixel_values"] = pixel_values if _has_image_grid_thw: - forward_args["image_grid_thw"] = data["image_grid_thw"] + forward_args["image_grid_thw"] = image_grid_thw if _has_image_sizes: - forward_args["image_sizes"] = data["image_sizes"] - - output_tensor = model(**forward_args) + forward_args["image_sizes"] = image_sizes + output = model(**forward_args) else: - # For text-only LLM models - output_tensor = model( - data["tokens"], + output = model( + tokens, position_ids, attention_mask, inference_context=inference_context, runtime_gather_output=True, ) - return output_tensor, _dummy_loss_func - - disable_tqdm = disable_tqdm or torch.distributed.get_rank() > 0 - - output_ids = torch.tensor([]) - step_pbar = tqdm(range(osl), disable=disable_tqdm, leave=False) - - time_ttft = 0 - time_remaining_outputs = 0 - timer = Timer("generate") - timer.start(barrier=True) - - for step in step_pbar: - step_pbar.set_description(get_current_memory_info()) - - if model.config.sequence_parallel: - tp = model.config.tensor_model_parallel_size - num_pad_tokens = (tp - input_ids.shape[-1] % tp) % tp - else: - num_pad_tokens = 0 - - if inference_context is not None and step > 0: - tokens = input_ids[:, -1:] - inference_context.enable_decode_mode() - elif num_pad_tokens > 0: - padding_shape = (input_ids.shape[0], num_pad_tokens) - padded_tokens = torch.full( - padding_shape, 0, dtype=input_ids.dtype, device=input_ids.device - ) - tokens = torch.cat((input_ids, padded_tokens), dim=-1) - else: - tokens = input_ids - - data_dict = {"tokens": tokens} - # Vision inputs should only be passed during prefill (step 0), not during decode steps - if pixel_values is not None: - data_dict["pixel_values"] = pixel_values - if image_grid_thw is not None: - data_dict["image_grid_thw"] = image_grid_thw - if image_sizes is not None: - data_dict["image_sizes"] = image_sizes - - list_of_logits = get_forward_backward_func()( - forward_step_func=_forward_step_func, - data_iterator=[data_dict], - model=model, - num_microbatches=1, - seq_length=tokens.shape[-1], - micro_batch_size=max_batch_size, - decoder_seq_length=tokens.shape[-1], - forward_only=True, - collect_non_loss_data=True, - ) if inference_context is not None: - inference_context.sequence_len_offset += tokens.shape[-1] + inference_context.sequence_len_offset += seq_len - if mpu.is_pipeline_last_stage(): - eager_ids = ( - list_of_logits[0][:, -(num_pad_tokens + 1), :].argmax(dim=-1, keepdim=True).detach() - ) + # For PP non-last stages, forward activations to the next stage. + if is_pp and not pp_last: + send_to_next_pipeline_rank(output.to(dtype=pp_dtype)) + + if pp_last: + eager_ids = output[:, -(num_pad_tokens + 1), :].argmax(dim=-1, keepdim=True).detach() else: eager_ids = None diff --git a/modelopt/torch/utils/plugins/megatron_mmlu.py b/modelopt/torch/utils/plugins/megatron_mmlu.py index fe71bc6eccc..4a07405caff 100644 --- a/modelopt/torch/utils/plugins/megatron_mmlu.py +++ b/modelopt/torch/utils/plugins/megatron_mmlu.py @@ -40,62 +40,60 @@ """A simple MMLU evaluation for Megatron LM models.""" -import requests import torch -import transformers from datasets import load_dataset +from tqdm import tqdm +from transformers import PreTrainedTokenizer -from .megatron_generate import megatron_generate +from .. import distributed as dist +from .. import print_rank_0 +from .megatron_generate import megatron_prefill __all__ = ["megatron_mmlu"] - -def _get_all_subjects(): - """All subjects (anatomy, ...) can be acquired from querying all subsets and splits.""" - response = requests.get( - "https://datasets-server.huggingface.co/splits?dataset=cais/mmlu", timeout=10 - ) - data = response.json() - all_subjects = set() - for split in data["splits"]: - all_subjects.add(split["config"]) - for name in ["all", "auxiliary_train"]: - all_subjects.discard(name) - return sorted(all_subjects) +_CHOICES = ["A", "B", "C", "D"] def megatron_mmlu( model, - tokenizer: transformers.PreTrainedTokenizer, + tokenizer: PreTrainedTokenizer, few_shots: int = 0, - percentage: float = 0.05, - enable_kv_cache: bool = False, + fraction: float = 0.05, + batch_size: int = 1, ) -> float: - """Evaluate the model on MMLU. + """Evaluate the model on MMLU using log-likelihood scoring over batched prefill passes. + + Instead of autoregressively generating tokens, a single prefill forward pass is run per + batch and the answer is selected as argmax over the four choice token logits at the last + prompt position. This is the same approach used by lm-evaluation-harness. Args: model: The model to evaluate. tokenizer: The tokenizer to use. few_shots: The number of few-shot examples to use. - percentage: The percentage of the test set to evaluate on. - enable_kv_cache: Whether to disable KV-cache. + fraction: The fraction of the test set to evaluate on. + batch_size: Number of examples to process in one forward pass. """ - all_correct = {} - all_subjects = _get_all_subjects() + print_rank_0( + f"\nMMLU ({fraction * 100}%, {few_shots}-shot, Batch Size: {batch_size}) evaluation started...\n" + "First batch may take longer to evaluate for Pipeline Parallel models." + ) + assert 0 < fraction <= 1, "Fraction must be between 0 and 1" + + # Token IDs for " A", " B", " C", " D" — the last subword handles edge cases. + choice_ids = [tokenizer.encode(f" {c}", add_special_tokens=False)[-1] for c in _CHOICES] def _format_example(example, include_answer: bool = True): - """Format an example into a multi-choices problem.""" prompt = example["question"] - for choice, answer in zip(["A", "B", "C", "D"], example["choices"]): + for choice, answer in zip(_CHOICES, example["choices"]): prompt += f"\n{choice}. {answer}" if include_answer: - prompt += "Answer: {}\n\n".format(example["answer"]) + prompt += "Answer: {}\n\n".format(_CHOICES[example["answer"]]) else: prompt += "\nAnswer:" return prompt def _generate_prompt(test_example, dev_examples, few_shots=0): - """Generating few-shot prompts.""" prompt = "The following are multiple choice questions (with answers) about {}.\n\n".format( " ".join(test_example["subject"].split("_")) ) @@ -104,51 +102,92 @@ def _generate_prompt(test_example, dev_examples, few_shots=0): prompt += _format_example(test_example, include_answer=False) return prompt - if torch.distributed.get_rank() == 0: - print(f"\nMMLU ({percentage * 100}%, {few_shots}-shot) evaluation started...\n", flush=True) - print("{:48} | (ACC) | Count/Total".format("Subject"), flush=True) - print("{:48} | {:5} | {:11}".format("-" * 48, "-" * 5, "-" * 11), flush=True) - - for subject in all_subjects: - test_data = load_dataset("cais/mmlu", subject, split="test") - dev_data = load_dataset("cais/mmlu", subject, split="dev") - - correct = [] - for idx, test_example in enumerate(test_data): - if idx > percentage * len(test_data): - break - prompt = _generate_prompt(test_example, dev_data, few_shots=few_shots) - label = ["A", "B", "C", "D"][test_example["answer"]] - tokens = tokenizer(prompt, return_tensors="pt") - generated_ids = megatron_generate( - model, - tokens.input_ids.cuda(), - osl=2, - disable_tqdm=True, - enable_kv_cache=enable_kv_cache, - ) - predict = tokenizer.batch_decode(generated_ids)[0].strip() - correct += [True] if predict.startswith(label) else [False] - all_correct[subject] = correct - - if torch.distributed.get_rank() == 0: - print( - f"{subject:48} | {sum(correct) / len(correct):.3f} | {sum(correct):5}/{len(correct):5}", - flush=True, - ) - - avg_correct = [] - - for subject, correct in all_correct.items(): - avg_correct += correct - - if torch.distributed.get_rank() == 0: - print("{:48} | {:5} | {:11}".format("-" * 48, "-" * 5, "-" * 11), flush=True) - print( - "{:48} | {:.3f} | {:5}/{:5}".format( - "average", sum(avg_correct) / len(avg_correct), sum(avg_correct), len(avg_correct) - ), - flush=True, - ) - - return sum(avg_correct) / len(avg_correct) + # Load all subjects in two dataset calls instead of 2x num_subjects calls. + # The "all" config includes a "subject" field for per-subject reporting. + test_dataset = load_dataset("cais/mmlu", "all", split="test") + dev_dataset = load_dataset("cais/mmlu", "all", split="dev") if few_shots > 0 else None + + # Group dev examples by subject for few-shot prompt construction. + dev_by_subject: dict = {} + if dev_dataset is not None: + for ex in dev_dataset: + dev_by_subject.setdefault(ex["subject"], []).append(ex) + + # Collect all examples, tracking subject membership for per-subject reporting. + all_subjects_seen: list[str] = [] + all_prompts: list[str] = [] + all_labels: list[str] = [] + + # Count test examples per subject to apply the fraction cutoff correctly. + subject_counts: dict[str, int] = {} + for ex in test_dataset: + subject_counts[ex["subject"]] = subject_counts.get(ex["subject"], 0) + 1 + + subject_idx: dict[str, int] = {} + for ex in test_dataset: + subj = ex["subject"] + idx = subject_idx.get(subj, 0) + if idx >= fraction * subject_counts[subj]: + continue + subject_idx[subj] = idx + 1 + prompt = _generate_prompt(ex, dev_by_subject.get(subj, []), few_shots=few_shots) + all_prompts.append(prompt) + all_labels.append(_CHOICES[ex["answer"]]) + all_subjects_seen.append(subj) + + # Tokenize all prompts and sort by length to minimise padding waste within batches. + encoded = [tokenizer(p, return_tensors="pt").input_ids[0] for p in all_prompts] + lengths = [e.shape[0] for e in encoded] + order = sorted(range(len(encoded)), key=lambda i: lengths[i], reverse=True) + + sorted_encoded = [encoded[i] for i in order] + sorted_lengths = [lengths[i] for i in order] + + # Run inference in global batches. + predictions: list[str] = [""] * len(encoded) + n_batches = (len(sorted_encoded) + batch_size - 1) // batch_size + pbar = tqdm( + range(0, len(sorted_encoded), batch_size), + total=n_batches, + desc="MMLU", + unit="batch", + disable=not dist.is_master(), + ) + for batch_start in pbar: + batch_enc = sorted_encoded[batch_start : batch_start + batch_size] + batch_len = sorted_lengths[batch_start : batch_start + batch_size] + max_len = max(batch_len) + + # Right-pad to max_len; causal mask means the last real token is unaffected by padding. + padded = torch.zeros(len(batch_enc), max_len, dtype=torch.long) + for i, (e, seq_len) in enumerate(zip(batch_enc, batch_len)): + padded[i, :seq_len] = e + + logits = megatron_prefill(model, padded.cuda()) # [B, max_len, vocab] + + for i, seq_len in enumerate(batch_len): + answer_logits = logits[i, seq_len - 1, choice_ids] + predictions[order[batch_start + i]] = _CHOICES[answer_logits.argmax().item()] + + examples_done = min(batch_start + batch_size, len(sorted_encoded)) + pbar.set_postfix(examples=f"{examples_done}/{len(sorted_encoded)}") + + # Compute per-subject accuracy and overall average. + subject_correct: dict[str, list[bool]] = {} + for pred, label, subj in zip(predictions, all_labels, all_subjects_seen): + subject_correct.setdefault(subj, []).append(pred == label) + + all_correct = [pred == label for pred, label in zip(predictions, all_labels)] + n_total = len(all_correct) + avg = sum(all_correct) / n_total + + print_rank_0("{:48} | (ACC) | Count/Total".format("Subject")) + print_rank_0("{:48} | {:5} | {:11}".format("-" * 48, "-" * 5, "-" * 11)) + for subj in sorted(subject_correct): + correct = subject_correct[subj] + n = len(correct) + print_rank_0(f"{subj:48} | {sum(correct) / n:.3f} | {sum(correct):5}/{n:5}") + print_rank_0("{:48} | {:5} | {:11}".format("-" * 48, "-" * 5, "-" * 11)) + print_rank_0("{:48} | {:.3f} | {:5}/{:5}".format("average", avg, sum(all_correct), n_total)) + + return avg diff --git a/tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py b/tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py index 63abe1723f7..81fca8ed961 100644 --- a/tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py +++ b/tests/gpu_megatron/torch/utils/plugins/test_utils_megatron.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. - +import pytest from _test_utils.torch.megatron.models import get_mcore_qwen3_600m from _test_utils.torch.megatron.utils import initialize_for_megatron from transformers import AutoTokenizer @@ -22,12 +22,18 @@ SEED = 1234 +# TODO: move to regression test folder -def _test_megatron_generate_and_mmlu(rank, size): - initialize_for_megatron(tensor_model_parallel_size=size, seed=SEED) - - model = get_mcore_qwen3_600m(tensor_model_parallel_size=size).cuda().eval() +def _test_megatron_generate_and_mmlu(rank, size, parallelism): + if parallelism == "tp": + initialize_for_megatron(tensor_model_parallel_size=size, seed=SEED) + model = get_mcore_qwen3_600m(tensor_model_parallel_size=size).cuda().eval() + elif parallelism == "pp": + initialize_for_megatron(pipeline_model_parallel_size=size, seed=SEED) + model = get_mcore_qwen3_600m(pipeline_model_parallel_size=size).cuda().eval() + else: + raise ValueError(f"Invalid parallelism: {parallelism}") tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") messages = [ @@ -42,10 +48,13 @@ def _test_megatron_generate_and_mmlu(rank, size): model_inputs = tokenizer([text], return_tensors="pt").to(device="cuda") output_ids = megatron_generate(model, model_inputs["input_ids"]) output_text = tokenizer.batch_decode(output_ids) - print(output_text) + print(rank, output_text) - assert megatron_mmlu(model, tokenizer) > 0.24 + assert 0.36 < megatron_mmlu(model, tokenizer, fraction=0.1, batch_size=16) < 0.39 -def test_megatron_generate_and_mmlu(dist_workers): - dist_workers.run(_test_megatron_generate_and_mmlu) +@pytest.mark.parametrize("parallelism", ["tp", "pp"]) +def test_megatron_generate_and_mmlu(dist_workers, parallelism, num_gpus): + if num_gpus == 1 and parallelism == "pp": + pytest.skip("Skipping as redundant test on 1 GPU") + dist_workers.run(_test_megatron_generate_and_mmlu, parallelism=parallelism) From dc7ad66b712855cc576312b6c5e5c6515aa61d54 Mon Sep 17 00:00:00 2001 From: sugunav14 <178320438+sugunav14@users.noreply.github.com> Date: Fri, 17 Apr 2026 15:54:39 -0700 Subject: [PATCH 11/30] GPTQ vector (#1223) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: ? ### Usage ```python # Add a code snippet demonstrating how to use this ``` ### Testing ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A - Did you write any new necessary tests?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit * **New Features** * Added backend-specific GPTQ helper registration to allow backend-tailored GPTQ behavior. * **Bug Fixes** * Prevented KV-cache state from leaking across repeated per-layer forwards during calibration. * **Tests** * Added GPU-focused tests validating GPTQ combined with vector quantization, including accuracy and end-to-end comparisons. --------- Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com> --- examples/llm_ptq/hf_ptq.py | 1 + modelopt/torch/quantization/model_calib.py | 27 +++++++++++++++++-- .../torch/quantization/utils/calib_utils.py | 15 +++++++++-- 3 files changed, 39 insertions(+), 4 deletions(-) diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index 327605406c4..37a88f97c74 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -283,6 +283,7 @@ def make_calib_dataloader( include_labels = ( args.auto_quantize_bits is not None and args.auto_quantize_method == "gradient" ) + calib_dataloader = get_dataset_dataloader( dataset_name=args.dataset, tokenizer=tokenizer, diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 35a0e931c9d..2336cb6a01b 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -49,7 +49,7 @@ reduce_amax, weight_attr_names, ) -from .utils.calib_utils import GPTQHelper +from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, GPTQHelper __all__ = [ "awq", @@ -1589,6 +1589,21 @@ def sequential_calibrate( def _layer_forward_loop(m, _inputs=layer_inputs): for args, kwargs_input in _inputs: + # Reset past_key_values to prevent the KV cache from + # accumulating across multiple forward replays (e.g. + # max_calibrate then Hessian collection in GPTQ). + # The layer doesn't need stale KV data — each replay + # should start with a fresh cache. + if ( + "past_key_values" in kwargs_input + and kwargs_input["past_key_values"] is not None + ): + kwargs_input = dict(kwargs_input) + cache = kwargs_input["past_key_values"] + if hasattr(cache, "reset"): + cache.reset() + else: + kwargs_input["past_key_values"] = None m(*args, **kwargs_input) calib_func(layer, _layer_forward_loop, **calib_kwargs) @@ -1648,7 +1663,15 @@ def gptq( print_rank_0("No quantized linear layers found, skipping GPTQ") return - gptq_handles = {name: GPTQHelper(m, name, offload_to_cpu=True) for name, m in quantized_layers} + def _make_gptq_handle(name, m): + backend = getattr(m.weight_quantizer, "backend", None) + if backend is None: + cls = GPTQHelper + else: + cls = _GPTQ_HELPER_REGISTRY.get(backend, GPTQHelper) + return cls(m, name, offload_to_cpu=True) + + gptq_handles = {name: _make_gptq_handle(name, m) for name, m in quantized_layers} for handle in gptq_handles.values(): handle.setup() diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index e52a8438d55..8c0ae1ee6ee 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -143,9 +143,7 @@ def update_weights(self, block_size, perc_damp): hessian = self.hessian.to(self.module.weight.device) self.weight = self.module.weight.data.float().clone() self._prepare_hessian_inverse(hessian, perc_damp) - self._blockwise_update(block_size) - self._print_mse_error(hessian) self.module.weight.data = self.weight.reshape(self.module.weight.shape).to( self.module.weight.data.dtype @@ -231,3 +229,16 @@ def _print_mse_error(self, hessian): mse = (delta).mm(hessian).mul(delta).mean() / (w_orig.mm(hessian).mul(w_orig).mean() + 1e-6) suffix = f", n_hessian_samples: {self.n_samples}" if self.n_samples else "" print_rank_0(f"[{self.name}] Relative MSE error: {mse.item():.2e}{suffix}") + + +_GPTQ_HELPER_REGISTRY: dict[str, type[GPTQHelper]] = {} + + +def register_gptq_helper(backend: str, factory: type[GPTQHelper]) -> None: + """Register a :class:`GPTQHelper` subclass for a quantizer backend. + + When :func:`modelopt.torch.quantization.model_calib.gptq` encounters a + module whose ``weight_quantizer.backend`` matches ``backend``, it will + construct ``factory`` instead of the default ``GPTQHelper``. + """ + _GPTQ_HELPER_REGISTRY[backend] = factory From 2d868d3f1f208f2ea9b7f8d8cdd506d9bf1cda04 Mon Sep 17 00:00:00 2001 From: realAsma <86726418+realAsma@users.noreply.github.com> Date: Fri, 17 Apr 2026 17:32:34 -0700 Subject: [PATCH 12/30] Performant layerwise calibration for large models (#1251) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Adds **performant layerwise calibration** for quantizing large models (e.g. DeepSeek-R1 671B) that don't fit entirely on GPU. ([Example commands](#example-commands)) 1. **Performant calibration for large models** — Each decoder layer is moved from CPU/disk to GPU (accelerate) or unsharded (FSDP2) **only once** and kept on GPU for the entire calibration step. Previously, every calibration batch triggered weight transfer for every layer — O(num_batches) weight movements per layer. Now it is O(1) per layer. This also means you can **increase batch size** since only one layer's weights occupy GPU at a time — e.g. DeepSeek-R1 on a single node (8×80GB) with `batch_size=16` and `gpu_max_mem_percentage=0.5`. 2. **Checkpoint save/resume** — Saves progress after each layer, so jobs that exceed cluster time limits (e.g. 4-hour Slurm windows for 100+ layer MoE models) can resume from the last completed layer. 3. **Rename** `sequential_calibrate` → `layerwise_calibrate` for clarity. ### Design details The existing layerwise state machine (skip/run/capture) already processes one layer at a time, but skip-mode layers still kept their parameters in the ModuleList — so frameworks transferred all weights every forward pass. This PR adds: - **`_SkipLayer`**: replaces fully-calibrated layers with a parameter-free dummy in the ModuleList, so framework hooks have nothing to transfer - **`persistent_materialization`**: keeps the active layer on GPU for the entire calibration step, avoiding repeated offload/reload cycles Checkpoint save is per-layer; restore is bulk — quantizer state and weights for layers 0..K-1 are restored once at the end of calibration, keeping the hot path fast. ### Example commands **Qwen3-8B** (NVFP4+GPTQ, single GPU): ```bash python hf_ptq.py \ --pyt_ckpt_path Qwen/Qwen3-8B \ --recipe nvfp4_gptq_sequential.yaml \ --calib_size 64 \ --batch_size 16 \ --dataset cnn_dailymail \ --export_path outputs/qwen3_8b_nvfp4_gptq_seq \ --gpu_max_mem_percentage 0.5 \ --use_seq_device_map \ --vllm_fakequant_export ``` **DeepSeek-R1** (NVFP4 experts-only + FP8 KV, 8×80GB): ```bash python hf_ptq.py \ --model unsloth/DeepSeek-R1-0528-BF16 \ --recipe ../../modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml \ --dataset cnn_dailymail \ --batch_size 16 \ --calib_size 64 \ --calib_seq 512 \ --gpu_max_mem_percentage 0.5 \ --use_seq_device_map \ --trust_remote_code \ --export_path output/DeepSeek-R1-BF16-nvfp4-experts-only-fp8-kv \ --vllm_fakequant_export ``` ### Example: NVFP4+GPTQ layerwise calibration on Qwen3-8B (36 layers, single GPU — 20 GB peak) **Initial run** (killed after layer 11): ``` Layerwise calibration: Found 36 transformer layers Calibrating layer 1/36 | capture: [1] Computing Hessians for 7 linear layers... GPTQ time: 51.39s Calibrating layer 2/36 | run: [1] | capture: [2] Checkpoint: saved layer 0 GPTQ time: 50.06s Calibrating layer 3/36 | skip: 1 | run: [2] | capture: [3] Checkpoint: saved layer 1 ... Calibrating layer 12/36 | skip: 10 | run: [11] | capture: [12] Checkpoint: saved layer 10 ``` **Resumed run** (picks up from layer 11, finishes all 36): ``` Layerwise calibration: Found 36 transformer layers Checkpoint: resuming layerwise calibration from layer 11/36 Calibrating layer 12 (resumed) GPTQ time: 51.45s Calibrating layer 13/36 | skip: 11 | run: [12] | capture: [13] Checkpoint: saved layer 11 ... Calibrating layer 36/36 | skip: 34 | run: [35] | capture: [36] Checkpoint: saved layer 34 GPTQ time: 50.33s Checkpoint: saved layer 35 (final) Checkpoint: restored 11 previously calibrated layers Layerwise calibration completed Quantized model exported to: outputs/qwen3_8b_nvfp4_gptq_seq GPU 0: Peak memory usage = 20.42 GB ``` ## TODO - [ ] Update CHANGELOG ## Test plan - `tests/unit/torch/quantization/test_layerwise_calibrate.py` — unit tests for skip/swap/restore - `tests/unit/torch/quantization/test_sequential_checkpoint.py` — checkpoint save/resume correctness - `tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py` — CPU-offloaded layerwise + GPTQ + checkpoint resume - `tests/gpu/torch/quantization/test_fsdp2.py` — FSDP2 layerwise calibration ### Verified - [x] Qwen3-8B: layerwise calibration + checkpoint save/restore + fakequantized checkpoint export + vLLM serve - [x] DeepSeek-R1: checkpoint resume tested - [x] DeepSeek-R1: fakequantized checkpoint export verified --------- Signed-off-by: realAsma --- CHANGELOG.rst | 1 + examples/llm_ptq/example_utils.py | 33 + examples/llm_ptq/hf_ptq.py | 17 +- .../torch/export/plugins/vllm_fakequant_hf.py | 189 +++-- modelopt/torch/quantization/config.py | 28 +- modelopt/torch/quantization/mode.py | 40 +- modelopt/torch/quantization/model_calib.py | 72 +- .../torch/quantization/plugins/accelerate.py | 82 ++- .../torch/quantization/plugins/huggingface.py | 2 +- modelopt/torch/quantization/utils/__init__.py | 2 +- .../utils/activation_collector.py | 335 --------- .../torch/quantization/utils/calib_utils.py | 2 +- .../torch/quantization/utils/core_utils.py | 107 ++- .../quantization/utils/layerwise_calib.py | 684 ++++++++++++++++++ modelopt/torch/utils/dataset_utils.py | 30 +- modelopt/torch/utils/network.py | 65 +- .../general/ptq/nvfp4_default-fp8_kv.yaml | 2 +- .../ptq/nvfp4_default-none_kv_gptq.yaml | 5 +- .../ptq/nvfp4_experts_only-fp8_kv.yaml | 7 +- .../export/test_vllm_fakequant_hf_export.py | 122 +++- .../plugins/test_accelerate_gpu.py | 592 ++++++++++++++- tests/gpu/torch/quantization/test_fsdp2.py | 133 ++++ tests/gpu/torch/quantization/test_gptq.py | 2 +- ...librate.py => test_layerwise_calibrate.py} | 38 +- .../quantization/plugins/test_huggingface.py | 2 +- tests/unit/torch/quantization/test_calib.py | 20 +- ...librate.py => test_layerwise_calibrate.py} | 250 ++++++- .../test_sequential_checkpoint.py | 185 +++++ tests/unit/torch/quantization/test_utils.py | 2 +- 29 files changed, 2467 insertions(+), 582 deletions(-) delete mode 100644 modelopt/torch/quantization/utils/activation_collector.py create mode 100644 modelopt/torch/quantization/utils/layerwise_calib.py rename tests/gpu/torch/quantization/{test_sequential_calibrate.py => test_layerwise_calibrate.py} (90%) rename tests/unit/torch/quantization/{test_sequential_calibrate.py => test_layerwise_calibrate.py} (64%) create mode 100644 tests/unit/torch/quantization/test_sequential_checkpoint.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index cd8e48e0857..80dea0e43e4 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -15,6 +15,7 @@ Changelog - Enable PTQ workflow for the Step3.5-Flash MoE model with NVFP4 W4A4 + FP8 KV cache quantization. See `modelopt_recipes/models/Step3.5-Flash/nvfp4-mlp-only.yaml `_ for more details. - Add support for vLLM fakequant reload using ModelOpt state for HF models. See `examples/vllm_serve/README.md `_ for more details. - [Early Testing] Add Claude Code PTQ skill (``.claude/skills/ptq/``) for agent-assisted post-training quantization. The skill guides the agent through environment detection, model support checking, format selection, and execution via the launcher or manual SLURM/Docker/bare GPU paths. Includes handling for unlisted models with custom module patching. This feature is in early testing — use with caution. +- Add performant layerwise calibration for large models that don't fit on GPU (e.g. DeepSeek-R1, Kimi-K2). See `modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml `_ for usage. Layerwise calibration also supports PTQ with intermediate progress saving — useful when long PTQ runs get hit with Slurm timeouts. See `modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml `_ for usage. **Backward Breaking Changes** diff --git a/examples/llm_ptq/example_utils.py b/examples/llm_ptq/example_utils.py index c2d4d4bfca8..90532efe38d 100755 --- a/examples/llm_ptq/example_utils.py +++ b/examples/llm_ptq/example_utils.py @@ -15,6 +15,7 @@ import copy import glob +import hashlib import inspect import json import logging @@ -854,3 +855,35 @@ def copy_custom_model_files(source_path: str, export_path: str, trust_remote_cod print(f"Successfully copied {len(copied_files)} custom model files to {export_path}") else: print("No custom model files found to copy") + + +def needs_checkpoint_path_update(quant_cfg: dict) -> bool: + """Check if quant_cfg has a layerwise_checkpoint_dir that should be auto-resolved to a unique subpath.""" + algorithm = quant_cfg.get("algorithm") + if not isinstance(algorithm, dict): + return False + return algorithm.get("layerwise_checkpoint_dir") is not None + + +def resolve_checkpoint_dir(quant_cfg: dict, model_path: str) -> dict: + """Append a unique ``_`` subdirectory to layerwise_checkpoint_dir. + + Allows a single recipe to be reused across models without checkpoint collisions. + Must only be called when :func:`needs_checkpoint_path_update` returns True. + """ + algorithm = quant_cfg["algorithm"] + base_dir = algorithm["layerwise_checkpoint_dir"] + + name = model_path.rstrip("/") + if "/" in name and not os.path.isabs(name): + name = name.replace("/", "--") + else: + name = Path(name).name + + config_hash = hashlib.sha256(json.dumps(quant_cfg, default=str).encode()).hexdigest()[:8] + + quant_cfg = copy.deepcopy(quant_cfg) + quant_cfg["algorithm"]["layerwise_checkpoint_dir"] = os.path.join( + base_dir, f"{name}_{config_hash}" + ) + return quant_cfg diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index 37a88f97c74..969a3d57190 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -34,6 +34,8 @@ is_enc_dec, is_nemotron_vl, load_mtp_weights, + needs_checkpoint_path_update, + resolve_checkpoint_dir, run_nemotron_vl_preview, ) from torch.utils.data import DataLoader @@ -91,8 +93,9 @@ def _set_kv_cache_constant_amax(quant_cfg: list) -> None: for i, entry in enumerate(quant_cfg): if entry.get("quantizer_name") != "*[kv]_bmm_quantizer": continue - assert isinstance(entry.get("cfg", {}), dict) - quant_cfg[i] = {**entry, "cfg": {**entry.get("cfg", {}), "use_constant_amax": True}} + cfg = entry.get("cfg") or {} + assert isinstance(cfg, dict) + quant_cfg[i] = {**entry, "cfg": {**cfg, "use_constant_amax": True}} break @@ -760,7 +763,9 @@ def export_quantized( # Load any missing weights from non-standard safetensors (handled in get_model for non-low-memory mode) # Store the MTP layer prefixes on the model for later exclusion from quantization if args.vllm_fakequant_export: - export_hf_vllm_fq_checkpoint(full_model, export_dir=export_path) + export_hf_vllm_fq_checkpoint( + full_model, export_dir=export_path, inplace_mem_efficient=True + ) else: mtp_layer_prefixes, mtp_state_dict = load_mtp_weights( full_model, args.pyt_ckpt_path @@ -1105,6 +1110,12 @@ def quantize_main( quant_cfg = copy.deepcopy(quant_cfg) _set_kv_cache_constant_amax(quant_cfg["quant_cfg"]) + if needs_checkpoint_path_update(quant_cfg): + quant_cfg = resolve_checkpoint_dir(quant_cfg, args.pyt_ckpt_path) + print( + f"Auto-resolved layerwise_checkpoint_dir: {quant_cfg['algorithm']['layerwise_checkpoint_dir']}" + ) + if args.qformat in QUANT_CFG_CHOICES: mono_quantize( args, diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index 1908354a0a7..786f9cdf593 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -24,6 +24,8 @@ from modelopt.torch.quantization.conversion import quantizer_state from modelopt.torch.quantization.nn import QuantModule, TensorQuantizer from modelopt.torch.quantization.utils import get_quantizer_state_dict +from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector from modelopt.torch.utils import get_unwrapped_name __all__ = ["export_hf_vllm_fq_checkpoint"] @@ -38,9 +40,75 @@ def disable_rotate(quantizer: TensorQuantizer): return False +def _fakequant_module_weights( + module: nn.Module, + module_name: str, + model: nn.Module, + state_dict: dict | None, + input_quantizers_folded_pqs: set, + fakequant_weights: set, + inplace: bool, +): + """Apply fake-quant to a single QuantModule's weights. + + When ``inplace=False``, reads/writes weights from/to ``state_dict``. + When ``inplace=True``, modifies the module's weight parameters directly. + """ + if not isinstance(module, QuantModule): + return + for attr_name, quantizer in module.named_children(): + if not ( + attr_name.endswith("weight_quantizer") + and isinstance(quantizer, TensorQuantizer) + and quantizer.fake_quant + and quantizer.is_enabled + ): + continue + weight_name = attr_name.removesuffix("_quantizer") + prefix = f"{module_name}." if module_name else "" + sd_key = f"{prefix}{weight_name}" + assert sd_key not in fakequant_weights, f"Weight {sd_key} has already been fakequantized" + + if inplace: + w = getattr(module, weight_name) + w_quant = quantizer(w.float()).to(w.dtype) + else: + assert state_dict is not None + if sd_key not in state_dict: + continue + w = state_dict[sd_key] + w_quant = quantizer(w.float()).to(w.dtype) + + # Fold pre_quant_scale: (x*s)@fake_quant(W) = x@(fake_quant(W)*s) + # Only valid when input_quantizer does NOT fake-quant activations. If it does + # fake_quant(x*s), the non-linearity prevents folding s into W. + inp_attr = attr_name.replace("weight_quantizer", "input_quantizer") + if hasattr(module, inp_attr): + inp_q = getattr(module, inp_attr) + if ( + hasattr(inp_q, "_pre_quant_scale") + and inp_q._pre_quant_scale is not None + and inp_q._disabled + ): + scale = inp_q._pre_quant_scale.squeeze().to(device=w_quant.device) + w_quant = (w_quant * scale[None, :]).to(w_quant.dtype) + inp_q_key = get_unwrapped_name( + f"{module_name}.{inp_attr}" if module_name else inp_attr, model + ) + input_quantizers_folded_pqs.add(inp_q_key) + + if inplace: + w.data.copy_(w_quant) + else: + assert state_dict is not None + state_dict[sd_key] = w_quant.cpu() + fakequant_weights.add(sd_key) + + def export_hf_vllm_fq_checkpoint( model: nn.Module, export_dir: Path | str, + inplace_mem_efficient: bool = False, ): """Export quantized HF weights + ``vllm_fq_modelopt_state.pth`` for vLLM fake-quant reload. @@ -53,62 +121,66 @@ def export_hf_vllm_fq_checkpoint( Args: model: In-memory quantized model. export_dir: Output dir for HF files and ``vllm_fq_modelopt_state.pth``. + inplace_mem_efficient: When True, applies fake-quant inplace one decoder layer at + a time using ``enable_weight_access_and_writeback``, avoiding full state + dict materialization. This is destructive — model weights are permanently + modified and weight quantizers are not re-enabled after export. """ export_dir = Path(export_dir) export_dir.mkdir(parents=True, exist_ok=True) # Step 1: Build the folded HF state dict. - # model.state_dict() returns detached copies of all tensors, so model - # parameters are never modified. Apply each weight quantizer's fake-quant - # to the corresponding weight tensor in the copy. - state_dict = model.state_dict() fakequant_weights = set() - input_quantizers_folded_pqs = ( - set() - ) # keys for input_quantizers where pre_quant_scale was folded + input_quantizers_folded_pqs = set() with torch.inference_mode(): - for module_name, module in model.named_modules(): - if not isinstance(module, QuantModule): - continue - for attr_name, quantizer in module.named_children(): - if not ( - attr_name.endswith("weight_quantizer") - and isinstance(quantizer, TensorQuantizer) - and quantizer.fake_quant - and quantizer.is_enabled - ): + if inplace_mem_efficient: + # Inplace path: iterate decoder layers, one offload<->onload per layer. + decoder_layers = LayerActivationCollector.get_decoder_layers(model) + assert decoder_layers is not None, ( + "inplace_mem_efficient=True requires a model with discoverable decoder layers" + ) + for name, module in model.named_modules(): + if module not in decoder_layers: continue - weight_name = attr_name.removesuffix("_quantizer") - prefix = f"{module_name}." if module_name else "" - sd_key = f"{prefix}{weight_name}" - assert sd_key not in fakequant_weights, ( - f"Weight {sd_key} has already been fakequantized" - ) - if sd_key in state_dict: - w = state_dict[sd_key] - w_quant = quantizer(w.float()).to(w.dtype).cpu() - # Fold pre_quant_scale: (x*s)@fake_quant(W) = x@(fake_quant(W)*s) - # Only valid when input_quantizer does NOT fake-quant activations. If it does - # fake_quant(x*s), the non-linearity prevents folding s into W. - inp_attr = attr_name.replace("weight_quantizer", "input_quantizer") - if hasattr(module, inp_attr): - inp_q = getattr(module, inp_attr) - if ( - hasattr(inp_q, "_pre_quant_scale") - and inp_q._pre_quant_scale is not None - and inp_q._disabled - ): - scale = inp_q._pre_quant_scale.squeeze().to(device=w_quant.device) - w_quant = (w_quant * scale[None, :]).to(w_quant.dtype) - inp_q_key = get_unwrapped_name( - f"{module_name}.{inp_attr}" if module_name else inp_attr, model - ) - input_quantizers_folded_pqs.add(inp_q_key) - state_dict[sd_key] = w_quant - fakequant_weights.add(sd_key) - - # Filter quantizer tensors out for a clean HF checkpoint. - clean_sd = {k: v for k, v in state_dict.items() if "quantizer" not in k} + with enable_weight_access_and_writeback(module, module): + for sub_name, sub_mod in module.named_modules(): + full_name = f"{name}.{sub_name}" if sub_name else name + _fakequant_module_weights( + sub_mod, + full_name, + model, + None, + input_quantizers_folded_pqs, + fakequant_weights, + inplace=True, + ) + # Meta tensors for offloaded weights (free); offload maps now have + # fakequanted values via writeback. + state_dict = model.state_dict() + else: + # Default path: full state_dict copy, fakequant into the copy. + state_dict = model.state_dict() + for module_name, module in model.named_modules(): + with enable_weight_access_and_writeback(module, model): + _fakequant_module_weights( + module, + module_name, + model, + state_dict, + input_quantizers_folded_pqs, + fakequant_weights, + inplace=False, + ) + + if inplace_mem_efficient: + # Let save_pretrained build its own state_dict so offloaded params go through + # its module_map / get_state_dict_from_offload path (modeling_utils.py:3967+). + # Passing state_dict= bypasses that path and crashes on meta tensors. + quantizer_keys = [k for k in state_dict if "quantizer" in k] + clean_sd = None + else: + clean_sd = {k: v for k, v in state_dict.items() if "quantizer" not in k} + quantizer_keys = None # Step 2: Disable weight quantizers, save modelopt state + quantizer state # dict, then re-enable. The _disabled=True flag is captured in modelopt_state @@ -161,9 +233,18 @@ def export_hf_vllm_fq_checkpoint( modelopt_state["modelopt_state_weights"] = quantizer_state_dict torch.save(modelopt_state, export_dir / "vllm_fq_modelopt_state.pth") - # Step 3: Save HF weights using the pre-built folded state dict. - model.save_pretrained(export_dir, state_dict=clean_sd, save_modelopt_state=False) - - for wq, orig_rotate in wqs_to_restore: - wq.enable() - wq._rotate = orig_rotate + # Step 3: Save HF weights. + if inplace_mem_efficient: + prev_ignore = getattr(model, "_keys_to_ignore_on_save", None) + model._keys_to_ignore_on_save = quantizer_keys + try: + model.save_pretrained(export_dir, save_modelopt_state=False) + finally: + model._keys_to_ignore_on_save = prev_ignore + else: + model.save_pretrained(export_dir, state_dict=clean_sd, save_modelopt_state=False) + + if not inplace_mem_efficient: + for wq, orig_rotate in wqs_to_restore: + wq.enable() + wq._rotate = orig_rotate diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 99c729efbcc..3f24ac09a41 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1217,16 +1217,36 @@ class QuantizeAlgorithmConfig(ModeloptBaseConfig): ), ) - use_sequential: bool = ModeloptField( + layerwise: bool = ModeloptField( default=False, - title="Enable sequential layer-by-layer calibration.", + title="Enable layerwise (layer-by-layer) calibration.", description=( - "If True, the calibration algorithm is applied sequentially to each decoder block. " - "Each layer's inputs are captured via a single forward pass that reflects the " + "If True, the calibration algorithm is applied layer by layer. " + "Each layer's inputs are captured via a forward pass that reflects the " "quantization of all preceding layers, incurring O(N) forward passes for N layers." ), ) + layerwise_checkpoint_dir: str | None = ModeloptField( + default=None, + title="Checkpoint directory for layerwise calibration.", + description=( + "If set together with layerwise=True, per-layer checkpoints are saved to this " + "directory during calibration. On restart, calibration resumes from the last " + "completed layer." + ), + ) + + @model_validator(mode="after") + def validate_layerwise_checkpoint_dir(self): + """Raise if layerwise_checkpoint_dir is set but layerwise is False.""" + if self.layerwise_checkpoint_dir is not None and not self.layerwise: + raise ValueError( + "layerwise_checkpoint_dir requires layerwise=True. " + "Set layerwise=True or remove layerwise_checkpoint_dir." + ) + return self + class MaxCalibConfig(QuantizeAlgorithmConfig): """The config for max calibration algorithm. diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index c81d5c89c74..713cdd7373c 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -60,10 +60,10 @@ from .model_calib import ( awq, gptq, + layerwise_calibrate, local_hessian_calibrate, max_calibrate, mse_calibrate, - sequential_calibrate, smoothquant, svdquant, ) @@ -213,6 +213,7 @@ def wrapped_calib_func( config: QuantizeAlgorithmConfig, forward_loop: ForwardLoop | None = None, func: Callable | None = None, + supports_layerwise: bool = True, ) -> ConvertReturnType: """Wrap the calibration function to be compatible with the ModelOpt convert entrypoint. @@ -222,7 +223,8 @@ def wrapped_calib_func( """ kwargs = config.model_dump() method = kwargs.pop("method") - sequential = kwargs.pop("use_sequential", False) + layerwise = kwargs.pop("layerwise", False) + checkpoint_dir = kwargs.pop("layerwise_checkpoint_dir", None) if method is not None and "awq" in method: # For backward compatibility kwargs["algorithm"] = method @@ -237,17 +239,24 @@ def wrapped_calib_func( module._moe_calib_experts_ratio = moe_calib_experts_ratio if func is not None: - if sequential: + if layerwise: + # All currently implemented PTQ algorithms support layerwise calibration; + # future algorithms that need full-model context must add a guard here. + if not supports_layerwise: + raise ValueError( + f"Calibration algorithm '{method}' does not support layerwise=True. " + "Set layerwise=False, or override `_supports_layerwise = True` on the " + "corresponding CalibrateModeDescriptor once the algorithm is made " + "compatible with per-layer calibration." + ) if forward_loop is None: raise ValueError("forward_loop is required for calibration but got None.") - assert method in ["max", "gptq"], ( - f"Sequential calibration currently only supports max and gptq calibration, got {method}" - ) - # Wrap with sequential processing - sequential_calibrate( + # Wrap with layerwise processing + layerwise_calibrate( model, forward_loop=forward_loop, calib_func=func, + checkpoint_dir=checkpoint_dir, **kwargs, ) else: @@ -281,6 +290,10 @@ class BaseCalibrateModeDescriptor(ModeDescriptor): _calib_func: Callable | None + # Override to False when the algorithm requires full-model context and + # cannot run per decoder layer (e.g. needs ModeloptStateManager on the root). + _supports_layerwise: bool = True + def __init__(self, *args, **kwargs): """Initialize Base calibrate mode descriptor.""" assert issubclass(self.config_class, QuantizeAlgorithmConfig), ( @@ -326,7 +339,13 @@ def convert(self) -> ConvertEntrypoint: def wrapped_func(model, config, forward_loop=None): # Access _calib_func as a class attribute to avoid binding # Check if _calib_func is defined as a class attribute - return wrapped_calib_func(model, config, forward_loop, func=self.__class__._calib_func) + return wrapped_calib_func( + model, + config, + forward_loop, + func=self.__class__._calib_func, + supports_layerwise=self.__class__._supports_layerwise, + ) return wrapped_func @@ -485,6 +504,9 @@ def config_class(self) -> type[QuantizeAlgorithmConfig]: return SVDQuantConfig _calib_func = svdquant + # create_and_replace_svdquant_linear_on_the_fly reads ModeloptStateManager from the + # root model, which is not present when layerwise_calibrate dispatches per decoder layer. + _supports_layerwise = False @property def restore(self) -> RestoreEntrypoint: diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 2336cb6a01b..b653369693d 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -28,7 +28,10 @@ from tqdm import tqdm from modelopt.torch.opt.searcher import ForwardLoop -from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector +from modelopt.torch.quantization.utils.layerwise_calib import ( + LayerActivationCollector, + _CheckpointState, +) from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method @@ -44,6 +47,7 @@ is_quantized_column_parallel_linear, is_quantized_linear, is_quantized_row_parallel_linear, + persistent_materialization, promote_nvfp4_static_quantizers, quantizer_attr_names, reduce_amax, @@ -53,9 +57,9 @@ __all__ = [ "awq", + "layerwise_calibrate", "local_hessian_calibrate", "max_calibrate", - "sequential_calibrate", "smoothquant", "svdquant", ] @@ -1552,21 +1556,27 @@ def postprocess(module, name): @torch.no_grad() -def sequential_calibrate( +def layerwise_calibrate( model: nn.Module, forward_loop: ForwardLoop, calib_func: Callable, **calib_kwargs, ): - """Sequential calibration - a sequential layer-by-layer calibration algorithm. + """Layerwise calibration - a layer-by-layer calibration algorithm. Runs the full model forward per layer but patches decoder layers with a skip / run / capture strategy so that inter-layer logic in parent modules (e.g. mask construction) executes naturally without model-specific hooks. + + If ``checkpoint_dir`` is passed (via ``calib_kwargs``), per-layer checkpoints + are saved after each layer completes. On restart, calibration resumes from + the last completed layer. """ + checkpoint_dir = calib_kwargs.pop("checkpoint_dir", None) + if forward_loop is None: raise ValueError( - "forward_loop must not be None for sequential calibration. " + "forward_loop must not be None for layerwise calibration. " "Please provide a valid forward_loop callable." ) @@ -1574,18 +1584,28 @@ def sequential_calibrate( if transformer_layers is None or len(transformer_layers) == 0: raise ValueError( "Could not find transformer layers in model. " - "Sequential calibration requires a model with identifiable transformer layers." + "Layerwise calibration requires a model with identifiable transformer layers." ) - print_rank_0(f"Sequential calibration: Found {len(transformer_layers)} transformer layers") + num_layers = len(transformer_layers) + print_rank_0(f"Layerwise calibration: Found {num_layers} transformer layers") + + ckpt = _CheckpointState.from_folder(checkpoint_dir, num_layers) + start_layer = ckpt.start_layer if ckpt else 0 input_getter = LayerActivationCollector(model) input_getter._patch_all_layers(decoder_layers=transformer_layers) + resumed_inputs = ckpt.setup_resume(transformer_layers) if ckpt and start_layer > 0 else None + try: - for layer_idx, layer in enumerate(transformer_layers): - print_rank_0(f"Calibrating layer {layer_idx + 1}/{len(transformer_layers)}") - layer_inputs = input_getter.get_input_activations(layer, forward_loop) + # Bootstrap: get first layer's inputs (or use resumed inputs). + layer_inputs = input_getter.get_first_layer_inputs( + start_layer, resumed_inputs, forward_loop + ) + + for layer_idx in range(start_layer, num_layers): + layer = transformer_layers[layer_idx] def _layer_forward_loop(m, _inputs=layer_inputs): for args, kwargs_input in _inputs: @@ -1606,14 +1626,30 @@ def _layer_forward_loop(m, _inputs=layer_inputs): kwargs_input["past_key_values"] = None m(*args, **kwargs_input) - calib_func(layer, _layer_forward_loop, **calib_kwargs) + with persistent_materialization(layer): + calib_func(layer, _layer_forward_loop, **calib_kwargs) + + # Run one more forward to get next layer's inputs and set + # output_meta on the just-calibrated layer (via "run" mode). + is_last = layer_idx + 1 >= num_layers + if not is_last: + next_inputs = input_getter.cache_outputs_for_next_layer_calib(layer, forward_loop) + else: + next_inputs = None + + if ckpt: + ckpt.save(layer_idx, layer, model, transformer_layers, next_inputs) del layer_inputs torch.cuda.empty_cache() + layer_inputs = next_inputs # noqa: F841 (used in next iteration's closure) finally: input_getter._unpatch_all_layers() - print_rank_0("Sequential calibration completed") + if ckpt: + ckpt.full_restore(transformer_layers, model) + + print_rank_0("Layerwise calibration completed") @torch.no_grad() @@ -1625,12 +1661,12 @@ def gptq( ): """GPTQ quantization. - Works in two modes depending on ``use_sequential`` in the config: + Works in two modes depending on ``layerwise`` in the config: - * **Sequential** (``use_sequential=True``): ``sequential_calibrate`` calls this + * **Layerwise** (``layerwise=True``): ``layerwise_calibrate`` calls this function once per decoder layer with updated activations, producing more accurate Hessian estimates. - * **Non-sequential** (``use_sequential=False``): called once on the full model. + * **Non-layerwise** (``layerwise=False``): called once on the full model. All layers are quantized in parallel from the original activations. Per-module steps: @@ -1643,7 +1679,7 @@ def gptq( Args: model: The module to quantize — either the full model or a single decoder - layer when invoked by ``sequential_calibrate``. + layer when invoked by ``layerwise_calibrate``. forward_loop: Callable that replays calibration inputs through *model*. perc_damp: Percentage of avg Hessian diagonal for damping (default: 0.01). block_size: Block size for GPTQ weight update. @@ -1686,8 +1722,10 @@ def _make_gptq_handle(name, m): handle.cleanup() print_rank_0("Updating weights using GPTQ algorithm...") + name_to_module = dict(model.named_modules()) for handle in gptq_handles.values(): - handle.update_weights(block_size, perc_damp) + with enable_weight_access_and_writeback(handle.module, model, name_to_module): + handle.update_weights(block_size, perc_damp) handle.free() del gptq_handles diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py index 13999df0f02..f80e2478dc6 100644 --- a/modelopt/torch/quantization/plugins/accelerate.py +++ b/modelopt/torch/quantization/plugins/accelerate.py @@ -31,51 +31,77 @@ __all__ = ["init_quantized_weights"] -def _get_cpu_offload_hook(hook): +def _get_offload_hook(hook): if isinstance(hook, AlignDevicesHook) and hook.offload and hook.weights_map is not None: - assert "weight" in hook.weights_map - if ( - isinstance(hook.weights_map, PrefixedDataset) - and hook.weights_map.prefix + "weight" not in hook.weights_map.dataset.state_dict - ): - raise NotImplementedError( - "This layer could be offloaded to disk. We don't support this yet." - ) + assert len(hook.weights_map) > 0 return hook elif isinstance(hook, SequentialHook): for h in hook.hooks: - align_hook = _get_cpu_offload_hook(h) + align_hook = _get_offload_hook(h) if align_hook is not None: return align_hook return None +def _writeback_params_to_weights_map(module, align_hook): + """Write all non-meta parameters and buffers back to the hook's CPU weights_map.""" + for name, tensor in module.state_dict(keep_vars=True).items(): + if tensor.device.type == "meta": + continue + if isinstance(align_hook.weights_map, PrefixedDataset): + key = align_hook.weights_map.prefix + name + w_map = align_hook.weights_map.dataset.state_dict + else: + w_map = align_hook.weights_map + key = name + if key in w_map: + w_map[key] = tensor.detach().to(w_map[key].device, dtype=w_map[key].dtype) + elif ( + isinstance(align_hook.weights_map, PrefixedDataset) + and hasattr(align_hook.weights_map.dataset, "index") + and key in align_hook.weights_map.dataset.index + ): + # Disk-offloaded weight: promote into state_dict so the next + # pre_forward picks up the modified tensor instead of the stale + # on-disk version. OffloadedWeightsLoader.__getitem__ gives + # state_dict priority over index, so this is sufficient. + w_map[key] = tensor.detach().cpu() + + @contextmanager def weight_access_and_writeback_context(module): - """Context manager for weight access and writeback for modules managed by accelerate.""" + """Context manager for weight access and writeback for modules managed by accelerate. + + Handles CPU-offloaded and disk-offloaded models. Iterates over the module and all + its descendants, materializing weights from any offload hook found and writing them + back on exit. ``pre_forward`` is skipped on modules whose weights are already + materialized (not on meta) to avoid overwriting them with stale CPU copies. + """ assert hasattr(module, "_hf_hook") - align_hook = _get_cpu_offload_hook(module._hf_hook) - if align_hook: - # Accelerate uses AlignDevicesHook to offload weights to CPU/Disk and then reload them in the forward pass - # The CPU/Disk offloaded weights are managed by PrefixDataset and OffloadedWeightsLoader - # See https://github.com/huggingface/accelerate/blame/f48d95c4939b281505a45b3d6e0bf554b65cc1ea/src/accelerate/utils/offload.py#L104-L141 - # TODO: Add support for disk-offloaded models if needed (they will be really slow, hence low priority) + materialized: list[tuple[torch.nn.Module, AlignDevicesHook, bool]] = [] + for mod in module.modules(): + if not hasattr(mod, "_hf_hook"): + continue + hook = _get_offload_hook(mod._hf_hook) + if hook is None: + continue + # Only call pre_forward if weights need materializing; already-materialized + # weights would be overwritten with stale CPU state_dict values. + needs_materialize = any(p.device.type == "meta" for p in mod.parameters()) + if needs_materialize: + hook.pre_forward(mod) + hook.offload = False + materialized.append((mod, hook, needs_materialize)) - # This will load the weights from CPU state_dict and move it to the GPU from meta device - align_hook.pre_forward(module) try: yield finally: - if align_hook: - # Update the weight in the CPU state_dict - if isinstance(align_hook.weights_map, PrefixedDataset): - key = align_hook.weights_map.prefix + "weight" - w_map = align_hook.weights_map.dataset.state_dict - else: - key, w_map = "weight", align_hook.weights_map - w_map[key] = module.weight.data.to(w_map[key].device, dtype=w_map[key].dtype) - align_hook.post_forward(module, None) + for mod, hook, was_materialized in materialized: + hook.offload = True + _writeback_params_to_weights_map(mod, hook) + if was_materialized: + hook.post_forward(mod, None) @contextmanager diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 82ab589934f..59bcd215bbc 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -39,7 +39,7 @@ from ..nn.modules.quant_linear import _QuantLinear from ..triton import IS_AVAILABLE as IS_TRITON_AVAILABLE from ..utils import replace_function, sync_moe_expert_amax -from ..utils.activation_collector import LayerActivationCollector +from ..utils.layerwise_calib import LayerActivationCollector from .attention import register_attention_for_kv_quant from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear, _QuantFunctionalMixin diff --git a/modelopt/torch/quantization/utils/__init__.py b/modelopt/torch/quantization/utils/__init__.py index 26603632096..dfc23c42eee 100644 --- a/modelopt/torch/quantization/utils/__init__.py +++ b/modelopt/torch/quantization/utils/__init__.py @@ -16,8 +16,8 @@ # ruff: noqa: F405 """Quantization utilities.""" -from .activation_collector import LayerActivationCollector from .core_utils import * +from .layerwise_calib import LayerActivationCollector __all__ = [ "EXPORT_MODE", diff --git a/modelopt/torch/quantization/utils/activation_collector.py b/modelopt/torch/quantization/utils/activation_collector.py deleted file mode 100644 index 5f187fdcb24..00000000000 --- a/modelopt/torch/quantization/utils/activation_collector.py +++ /dev/null @@ -1,335 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Sequential calibration layer patching and activation capture. - -This module provides :class:`LayerActivationCollector`, a stateful helper that -patches decoder layers with a skip / run / capture strategy for efficient -layer-by-layer calibration. -""" - -from collections import deque -from dataclasses import dataclass, field -from typing import Any - -import torch -import torch.nn as nn - -from modelopt.torch.opt.searcher import ForwardLoop -from modelopt.torch.utils import print_rank_0 -from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method - - -class _EarlyStopForwardError(Exception): - """Raised to halt the forward pass after capturing layer inputs.""" - - -@dataclass -class _LayerCalibState: - """Mutable per-layer state used during sequential calibration. - - Attached to each decoder layer as ``_seq_calib`` and accessed by the - patched forward to decide skip / run / capture / original behaviour. - """ - - mode: str = "original" - name: str = "" - cached_inputs: deque = field(default_factory=deque) - collected_inputs: list = field(default_factory=list) - output_meta: tuple | None = None - - -class LayerActivationCollector: - """Collects layer activations for sequential (layer-by-layer) calibration. - - Each decoder layer is patched with a unified forward whose behaviour is - governed by a per-layer :class:`_LayerCalibState`: - - * **skip** — return a zero-filled dummy whose shape and type match the - layer's real output (reconstructed from lightweight metadata). No - computation is performed. The correctly shaped dummy ensures un-patched - inter-layer operations in the parent forward (e.g. LayerNorm, tuple - unpacking) do not raise shape or type errors. - * **run** — replay previously captured inputs through the original forward, - ignoring whatever the parent passes in. Only the just-calibrated layer - uses this mode, so its output reflects updated weights. - * **capture** — record ``(args, kwargs)`` and raise - ``_EarlyStopForwardError`` to halt the forward pass early. - * **original** — call the original forward unchanged. - - Because the *run* layer discards upstream values, skip-layer outputs are - never consumed for real computation. - """ - - # Global registry of (predicate, discoverer) pairs. Populated at import time - # by plugins (e.g. huggingface.py, megatron.py). Order matters: the first - # matching entry wins, so more specific predicates (e.g. Nemotron-H) must be - # registered before generic ones (e.g. homogeneous HF models). - # - # This is intentionally a mutable class variable shared across all instances: - # plugins register once at import time, and the registry is read-only after - # that. register_decoder_layer_support() guards against duplicate entries. - _decoder_layer_support: list[tuple[Any, Any]] = [] - _LAYER_ATTR = "_seq_calib" - - def __init__(self, model: nn.Module): - """Initialize the collector for the given model.""" - self.model = model - self._decoder_layers: nn.ModuleList | None = None - self._layer_to_idx: dict[nn.Module, int] = {} - self._patched = False - - @staticmethod - def get_decoder_layers(model: nn.Module) -> nn.ModuleList | None: - """Return decoder layers supported by sequential calibration.""" - for is_supported, discoverer in LayerActivationCollector._decoder_layer_support: - if not is_supported(model): - continue - decoder_layers = discoverer(model) - if decoder_layers is not None: - return decoder_layers - return None - - @staticmethod - def is_supported(model: nn.Module) -> bool: - """Whether the model supports decoder-layer sequential calibration.""" - return LayerActivationCollector.get_decoder_layers(model) is not None - - @classmethod - def register_decoder_layer_support(cls, is_supported: Any, discoverer: Any): - """Register a (predicate, discoverer) pair for decoder-layer detection.""" - entry = (is_supported, discoverer) - if entry not in cls._decoder_layer_support: - cls._decoder_layer_support.append(entry) - - @staticmethod - def _extract_output_meta(output): - """Extract lightweight (shape, dtype, device) metadata from a layer output. - - Recursively handles tensors, tuples, lists, and non-tensor values (e.g. None). - The returned structure can be passed to ``_zeros_from_meta`` to reconstruct a - zero-filled output with identical shape and type. - """ - if isinstance(output, torch.Tensor): - return ("tensor", output.shape, output.dtype, output.device) - if isinstance(output, tuple): - return ( - "tuple", - tuple(LayerActivationCollector._extract_output_meta(o) for o in output), - ) - if isinstance(output, list): - return ("list", [LayerActivationCollector._extract_output_meta(o) for o in output]) - return ("other", output) - - @staticmethod - def _zeros_from_meta(meta): - """Reconstruct a zero-filled output from metadata produced by ``_extract_output_meta``.""" - tag = meta[0] - if tag == "tensor": - _, shape, dtype, device = meta - return torch.zeros(shape, dtype=dtype, device=device) - if tag == "tuple": - return tuple(LayerActivationCollector._zeros_from_meta(m) for m in meta[1]) - if tag == "list": - return [LayerActivationCollector._zeros_from_meta(m) for m in meta[1]] - # "other" values are expected to be lightweight non-tensors (e.g. None, small scalars). - # The value is returned directly (not copied); callers must not mutate it. - # In practice this is safe because skip-mode outputs are immediately discarded by the - # downstream run-mode layer, which replays from its own cached inputs instead. - return meta[1] - - def _patch_all_layers(self, decoder_layers: nn.ModuleList | None = None): - """Bind the unified forward to every decoder layer and the model. Called once. - - Args: - decoder_layers: Pre-resolved decoder layers. If *None*, layers are - discovered via :meth:`get_decoder_layers`. - """ - - def _patched_forward(self, *args, **kwargs): - """Unified forward bound to every decoder layer during sequential calibration. - - ``self`` here is the decoder layer module (bound via ``bind_forward_method``). - All per-layer state is accessed through ``self._seq_calib``. - """ - info: _LayerCalibState = self._seq_calib - - if info.mode == "skip": - if info.output_meta is None: - raise RuntimeError( - f"Layer {info.name} is in 'skip' mode but has no output_meta. " - "This indicates a state-machine bug: the layer should have run " - "in 'run' mode (which sets output_meta) before transitioning to 'skip'." - ) - return LayerActivationCollector._zeros_from_meta(info.output_meta) - - if info.mode == "run": - assert info.cached_inputs, ( - f"Layer {info.name} is in 'run' mode but has no cached inputs to replay." - ) - real_args, real_kwargs = info.cached_inputs.popleft() - output = self._original_forward(*real_args, **real_kwargs) - info.output_meta = LayerActivationCollector._extract_output_meta(output) - return output - - if info.mode == "capture": - info.collected_inputs.append((args, kwargs)) - raise _EarlyStopForwardError() - - return self._original_forward(*args, **kwargs) - - if decoder_layers is not None: - self._decoder_layers = decoder_layers - else: - self._decoder_layers = self.get_decoder_layers(self.model) - assert self._decoder_layers is not None - - self._layer_to_idx = {layer: i for i, layer in enumerate(self._decoder_layers)} - module_to_name = {m: name for name, m in self.model.named_modules()} - - try: - for layer in self._decoder_layers: - layer._seq_calib = _LayerCalibState( - name=module_to_name.get(layer, type(layer).__name__), - ) - bind_forward_method(layer, _patched_forward, "_original_forward") - - def _early_stop_forward(module_self, *args, **kwargs): - try: - return module_self._original_forward(*args, **kwargs) - except _EarlyStopForwardError: - return None - - bind_forward_method(self.model, _early_stop_forward, "_original_forward") - except Exception: - self._cleanup_layers() - raise - - self._patched = True - - def _cleanup_layers(self): - """Best-effort cleanup of any patched layers and model forward.""" - if hasattr(self.model, "_original_forward"): - unpatch_forward_method(self.model, "_original_forward") - - if self._decoder_layers is not None: - for layer in self._decoder_layers: - if hasattr(layer, "_original_forward"): - unpatch_forward_method(layer, "_original_forward") - if hasattr(layer, self._LAYER_ATTR): - delattr(layer, self._LAYER_ATTR) - - def _unpatch_all_layers(self): - """Restore original forwards and clean up state attributes. Called once.""" - if not self._patched: - return - self._cleanup_layers() - self._patched = False - - def _set_layer_states(self, layer_idx: int): - """Transition layer modes for the next calibration step. - - When calibrating layer *i*, three transitions happen: - - * Layer ``i - 2`` → **skip** (fully done, free its cached inputs). - * Layer ``i - 1`` → **run** (replay captured inputs with calibrated weights). - * Layer ``i`` → **capture** (record inputs, then early-stop). - """ - assert self._decoder_layers is not None - - if layer_idx > 1: - done = self._decoder_layers[layer_idx - 2]._seq_calib - # output_meta is intentionally kept: skip mode needs it to produce - # correctly shaped zero-filled outputs for the parent forward. - done.mode = "skip" - done.cached_inputs.clear() - - if layer_idx > 0: - prev = self._decoder_layers[layer_idx - 1]._seq_calib - if not prev.collected_inputs: - raise RuntimeError( - f"Layer {layer_idx - 1} ({prev.name!r}) has no collected inputs to replay. " - "Layers must be calibrated sequentially — ensure get_input_activations() " - "was called for every preceding layer in order." - ) - prev.mode = "run" - prev.cached_inputs = deque(prev.collected_inputs) - prev.collected_inputs = [] - - cur = self._decoder_layers[layer_idx]._seq_calib - cur.mode = "capture" - cur.collected_inputs = [] - - def _log_layer_summary(self, layer_idx: int): - """Log a one-line summary of layer modes for the current calibration step.""" - assert self._decoder_layers is not None - n = len(self._decoder_layers) - groups: dict[str, list[int]] = {} - for i, layer in enumerate(self._decoder_layers): - mode = layer._seq_calib.mode - if mode in ("skip", "run", "capture"): - groups.setdefault(mode, []).append(i + 1) - parts = [f"{mode}: {groups[mode]}" for mode in ("skip", "run", "capture") if mode in groups] - print_rank_0(f"Calibrating layer {layer_idx + 1}/{n} | {' | '.join(parts)}") - - # ------------------------------------------------------------------ - # Public API - # ------------------------------------------------------------------ - - @torch.no_grad() - def get_input_activations(self, layer: torch.nn.Module, forward_loop: ForwardLoop) -> list: - """Collect input activations for *layer* by running a full model forward. - - Layers before the target are skipped or re-run (if just calibrated), the - target layer captures its inputs, and an early-stop prevents unnecessary - computation beyond the target. - - :meth:`_patch_all_layers` must be called before this method. - - Note: the model forward returns ``None`` for every batch during capture - (because ``_EarlyStopForwardError`` short-circuits the forward pass). - Callers should not rely on the model's return value within *forward_loop*. - """ - if not self._patched: - raise RuntimeError( - "get_input_activations() requires _patch_all_layers() to be called first." - ) - layer_idx = self._layer_to_idx[layer] - self._set_layer_states(layer_idx) - self._log_layer_summary(layer_idx) - - info = layer._seq_calib - try: - forward_loop(self.model) - except Exception: - # Reset the current layer so subsequent calls don't see stale state. - info.mode = "original" - info.collected_inputs = [] - raise - - if not info.collected_inputs: - info.mode = "original" - raise RuntimeError( - f"Layer {info.name!r} collected no inputs during forward_loop. " - "The forward loop did not reach this layer — check that forward_loop() " - "actually calls the model and that the layer is in the forward path." - ) - - inputs = list(info.collected_inputs) - # After capture, set to original so calib_func can call the layer's - # real forward directly. The layer will transition to run → skip - # in subsequent iterations via _set_layer_states. - info.mode = "original" - return inputs diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 8c0ae1ee6ee..252f0af6fc8 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -96,7 +96,7 @@ def __init__(self, module, name, offload_to_cpu=False): self.name = name in_features = module.weight.shape[-1] device = module.weight.device - if offload_to_cpu and get_used_gpu_mem_fraction(device) > 0.65: + if device.type == "meta" or (offload_to_cpu and get_used_gpu_mem_fraction(device) > 0.65): device = "cpu" self.hessian = torch.zeros(in_features, in_features, dtype=torch.float32, device=device) self.n_samples = 0 diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 273d7564c69..29661e18f52 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -423,47 +423,70 @@ def _get_enclosing_fsdp_module( return root_model +def _set_parameter(module: nn.Module, name: str, value: nn.Parameter): + """Set a parameter on a module by dotted name (e.g. ``self_attn.q_proj.weight``).""" + parts = name.rsplit(".", 1) + if len(parts) == 2: + parent = module.get_submodule(parts[0]) + attr = parts[1] + else: + parent = module + attr = name + parent._parameters[attr] = value + + @contextmanager def fsdp2_weight_access_and_writeback_context(module: nn.Module, root_model: nn.Module): """Context manager for FSDP2 weight access and writeback. - Note this context will gather the weight across FSDP/HSDP shards. If TP is implemented with DTensor, - the weight will be a local tensor of the TP DTensor under this context. + Gathers sharded DTensor parameters across FSDP/HSDP shards so they can be + read or modified. Works for both leaf modules (single ``weight``) and + composite modules like decoder layers (all ``named_parameters``). + + If TP is implemented with DTensor, the weight will be a local tensor of the + TP DTensor under this context. """ assert isinstance(root_model, torch.distributed.fsdp.FSDPModule), "We only support FSDP2" assert not hasattr(module, "_hf_hook"), "We dont support FSDP2 with HF accelerate hooks" - assert isinstance(module.weight, torch.distributed.tensor.DTensor) fsdp_module = _get_enclosing_fsdp_module(module, root_model) assert fsdp_module is not None, "Module is not wrapped by FSDP" fsdp_device_mesh = _get_fsdp2_mesh(fsdp_module) fsdp_dim = fsdp_device_mesh.ndim - original_placements = module.weight.placements - original_device_mesh = module.weight.device_mesh - original_weight = module.weight - # Assuming the first fsdp_dim dimensions are for FSDP/HSDP, we only collect the tensor over FSDP/HSDP dimension, - # the TP will be handled by the TP reduction. - if fsdp_dim != original_device_mesh.ndim: - assert fsdp_device_mesh.mesh_dim_names == original_device_mesh.mesh_dim_names[:fsdp_dim], ( - "FSDP2 mesh should be a slice of DTesnor's device mesh." + # Collect all DTensor parameters, replacing them with local replicated copies. + originals: dict[str, tuple] = {} + for name, param in module.named_parameters(): + if not isinstance(param, torch.distributed.tensor.DTensor): + continue + original_placements = param.placements + original_device_mesh = param.device_mesh + if fsdp_dim != original_device_mesh.ndim: + assert ( + fsdp_device_mesh.mesh_dim_names == original_device_mesh.mesh_dim_names[:fsdp_dim] + ), "FSDP2 mesh should be a slice of DTensor's device mesh." + collected = param.redistribute( + placements=[Replicate()] * fsdp_dim + list(original_placements[fsdp_dim:]), + device_mesh=original_device_mesh, ) - - weight_collected = original_weight.redistribute( - placements=[Replicate()] * fsdp_dim + list(original_placements[fsdp_dim:]), - device_mesh=original_device_mesh, - ) - new_weight = nn.Parameter(weight_collected.to_local()) - module._parameters["weight"] = new_weight + originals[name] = (param, collected, original_placements, original_device_mesh) + _set_parameter(module, name, nn.Parameter(collected.to_local())) yield - original_weight.to_local().data.copy_( - weight_collected.redistribute( - placements=original_placements, device_mesh=original_device_mesh - ).to_local() - ) - module._parameters["weight"] = original_weight + # Write back and restore original DTensor parameters. + for name, ( + original_param, + collected, + original_placements, + original_device_mesh, + ) in originals.items(): + original_param.to_local().data.copy_( + collected.redistribute( + placements=original_placements, device_mesh=original_device_mesh + ).to_local() + ) + _set_parameter(module, name, original_param) @contextmanager @@ -471,7 +494,7 @@ def enable_weight_access_and_writeback(module, root_model, name_to_module: dict """Enable weight access and writeback for a module. Useful for modules with weight not intact such as Linear layer in FSDP wrapped model or - HF accelerate CPU off-loaded models. + HF accelerate offloaded models (CPU or disk). Args: module: The module to access weights for. @@ -498,6 +521,22 @@ def enable_weight_access_and_writeback(module, root_model, name_to_module: dict yield +@contextmanager +def persistent_materialization(layer): + """Keep all layer weights materialized on GPU for the duration. + + Suppresses per-forward weight transfers so that N calibration batches + pay the cost of one load/unload instead of N. + + - **FSDP2**: patches ``FSDPParamGroup.unshard/reshard`` to no-ops, then + gathers weights once via ``enable_weight_access_and_writeback``. + - **Accelerate**: materializes weights and sets ``hook.offload = False`` + so per-forward hooks skip materialization/offloading. + """ + with _disable_fsdp_unshard_reshard(layer), enable_weight_access_and_writeback(layer, layer): + yield + + def get_quantizer_state_dict(model: nn.Module): """Get the state dict of the quantizers in the model.""" # We should not call model.state_dict() here. @@ -607,6 +646,24 @@ def _init_mp_dtypes(self) -> None: ) +@contextmanager +def _disable_fsdp_unshard_reshard(layer): + """Disable FSDP2 unshard/reshard if *layer* is FSDP-wrapped.""" + if isinstance(layer, FSDPModule): + _pg_cls = torch.distributed.fsdp._fully_shard._fsdp_param_group.FSDPParamGroup + orig_unshard = _pg_cls.unshard + orig_reshard = _pg_cls.reshard + _pg_cls.unshard = lambda self, async_op=False: None + _pg_cls.reshard = lambda self: None + try: + yield + finally: + _pg_cls.unshard = orig_unshard + _pg_cls.reshard = orig_reshard + else: + yield + + def get_prefixed_param_names(parent_model, target_module): """Get parameter names for a target module prefixed with the parent model name. diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py new file mode 100644 index 00000000000..aed403ad87b --- /dev/null +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -0,0 +1,684 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Layerwise calibration layer patching, activation capture, and checkpoint save/resume. + +This module provides :class:`LayerActivationCollector`, a stateful helper that +patches decoder layers with a skip / run / capture strategy for efficient +layer-by-layer calibration, and :class:`_CheckpointState` for persisting +per-layer calibration progress to disk. +""" + +from __future__ import annotations + +import json +import os +import shutil +from collections import deque +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any + +import torch +import torch.nn as nn + +from modelopt.torch.utils import distributed as dist +from modelopt.torch.utils import print_rank_0 +from modelopt.torch.utils.network import ( + bind_forward_method, + get_module_device, + unpatch_forward_method, +) + +if TYPE_CHECKING: + from modelopt.torch.opt.searcher import ForwardLoop + + +class _EarlyStopForwardError(Exception): + """Raised to halt the forward pass after capturing layer inputs.""" + + +@dataclass +class _LayerCalibState: + """Mutable per-layer state used during layerwise calibration. + + Attached to each decoder layer as ``_layerwise_calib`` and accessed by the + patched forward to decide skip / run / capture / original behaviour. + """ + + mode: str = "original" + name: str = "" + cached_inputs: deque = field(default_factory=deque) + collected_inputs: list = field(default_factory=list) + output_meta: tuple | None = None + + +class _SkipLayer(nn.Module): + """Parameter-free stand-in for a fully calibrated decoder layer. + + Replaces the real layer in the ModuleList so that framework hooks + (accelerate, FSDP2, etc.) have no parameters to transfer. Holds a + reference to the original layer for restoration during cleanup. + """ + + def __init__(self, original: nn.Module): + super().__init__() + # Bypass nn.Module.__setattr__ to avoid registering original as a submodule. + object.__setattr__(self, "_original", original) + self._layerwise_calib = _LayerCalibState(mode="skip") + + _PROXY_BLOCKLIST = frozenset({"_hf_hook", "_old_forward"}) + + def __getattr__(self, name: str): + # Proxy non-special attribute lookups to the original layer so that + # parent-model code that accesses layer-level attributes (e.g., + # NemotronH's ``block_type``) still works when the layer is replaced + # with a _SkipLayer. Accelerate hook attrs are blocked so the + # framework does not attempt to manage this parameter-free stand-in. + try: + return super().__getattr__(name) + except AttributeError: + if name in self._PROXY_BLOCKLIST: + raise + return getattr(object.__getattribute__(self, "_original"), name) + + def forward(self, *args, **kwargs): + return LayerActivationCollector._zeros_from_meta( + self._original._layerwise_calib.output_meta + ) + + +class LayerActivationCollector: + """Collects layer activations for layerwise (layer-by-layer) calibration. + + Each decoder layer is patched with a unified forward whose behaviour is + governed by a per-layer :class:`_LayerCalibState`: + + * **skip** — return a zero-filled dummy whose shape and type match the + layer's real output (reconstructed from lightweight metadata). No + computation is performed. The correctly shaped dummy ensures un-patched + inter-layer operations in the parent forward (e.g. LayerNorm, tuple + unpacking) do not raise shape or type errors. + * **run** — replay previously captured inputs through the original forward, + ignoring whatever the parent passes in. Only the just-calibrated layer + uses this mode, so its output reflects updated weights. + * **capture** — record ``(args, kwargs)`` and raise + ``_EarlyStopForwardError`` to halt the forward pass early. + * **original** — call the original forward unchanged. + + Because the *run* layer discards upstream values, skip-layer outputs are + never consumed for real computation. + """ + + _decoder_layer_support: list[tuple[Any, Any]] = [] + _LAYER_ATTR = "_layerwise_calib" + + def __init__(self, model: nn.Module): + """Initialize the collector for the given model.""" + self.model = model + self._decoder_layers: nn.ModuleList | None = None + self._layer_to_idx: dict[nn.Module, int] = {} + self._patched = False + + def _swap_to_dummy(self, idx: int): + """Replace decoder layer *idx* with a parameter-free dummy. + + ``output_meta`` is intentionally preserved on the original layer: the + ``_SkipLayer`` reads it to produce correctly shaped zero-filled outputs + for the parent forward pass. + """ + assert self._decoder_layers is not None + layer = self._decoder_layers[idx] + layer._layerwise_calib.mode = "skip" + layer._layerwise_calib.cached_inputs.clear() + self._decoder_layers[idx] = _SkipLayer(layer) + + @staticmethod + def get_decoder_layers(model: nn.Module) -> nn.ModuleList | None: + """Return decoder layers supported by layerwise calibration.""" + for is_supported, discoverer in LayerActivationCollector._decoder_layer_support: + if not is_supported(model): + continue + decoder_layers = discoverer(model) + if decoder_layers is not None: + return decoder_layers + return None + + @staticmethod + def is_supported(model: nn.Module) -> bool: + """Whether the model supports decoder-layer layerwise calibration.""" + return LayerActivationCollector.get_decoder_layers(model) is not None + + @classmethod + def register_decoder_layer_support(cls, is_supported: Any, discoverer: Any): + """Register a (predicate, discoverer) pair for decoder-layer detection.""" + entry = (is_supported, discoverer) + if entry not in cls._decoder_layer_support: + cls._decoder_layer_support.append(entry) + + @staticmethod + def _extract_output_meta(output): + """Extract lightweight (shape, dtype, device) metadata from a layer output. + + Recursively handles tensors, tuples, lists, and non-tensor values (e.g. None). + The returned structure can be passed to ``_zeros_from_meta`` to reconstruct a + zero-filled output with identical shape and type. + """ + if isinstance(output, torch.Tensor): + return ("tensor", output.shape, output.dtype, output.device) + if isinstance(output, tuple): + return ( + "tuple", + tuple(LayerActivationCollector._extract_output_meta(o) for o in output), + ) + if isinstance(output, list): + return ("list", [LayerActivationCollector._extract_output_meta(o) for o in output]) + return ("other", output) + + @staticmethod + def _zeros_from_meta(meta): + """Reconstruct a zero-filled output from metadata produced by ``_extract_output_meta``.""" + tag = meta[0] + if tag == "tensor": + _, shape, dtype, device = meta + return torch.zeros(shape, dtype=dtype, device=device) + if tag == "tuple": + return tuple(LayerActivationCollector._zeros_from_meta(m) for m in meta[1]) + if tag == "list": + return [LayerActivationCollector._zeros_from_meta(m) for m in meta[1]] + # "other" values are lightweight non-tensors (e.g. None, small scalars). + # Returned directly (not copied); safe because skip-mode outputs are + # immediately discarded by the downstream run-mode layer. + return meta[1] + + def _patch_all_layers(self, decoder_layers: nn.ModuleList | None = None): + """Bind the unified forward to every decoder layer and the model. Called once. + + Args: + decoder_layers: Pre-resolved decoder layers. If *None*, layers are + discovered via :meth:`get_decoder_layers`. + """ + + def _patched_forward(self, *args, **kwargs): + info: _LayerCalibState = self._layerwise_calib + + if info.mode == "skip": + if info.output_meta is None: + raise RuntimeError( + f"Layer {info.name} is in 'skip' mode but has no output_meta. " + "This indicates a state-machine bug: the layer should have run " + "in 'run' mode (which sets output_meta) before transitioning to 'skip'." + ) + return LayerActivationCollector._zeros_from_meta(info.output_meta) + + if info.mode == "run": + assert info.cached_inputs, ( + f"Layer {info.name} is in 'run' mode but has no cached inputs to replay." + ) + real_args, real_kwargs = info.cached_inputs.popleft() + output = self._original_forward(*real_args, **real_kwargs) + info.output_meta = LayerActivationCollector._extract_output_meta(output) + return output + + if info.mode == "capture": + info.collected_inputs.append((args, kwargs)) + raise _EarlyStopForwardError() + + return self._original_forward(*args, **kwargs) + + if decoder_layers is not None: + self._decoder_layers = decoder_layers + else: + self._decoder_layers = self.get_decoder_layers(self.model) + assert self._decoder_layers is not None + + self._layer_to_idx = {layer: i for i, layer in enumerate(self._decoder_layers)} + module_to_name = {m: name for name, m in self.model.named_modules()} + + try: + for layer in self._decoder_layers: + layer._layerwise_calib = _LayerCalibState( + name=module_to_name.get(layer, type(layer).__name__), + ) + bind_forward_method(layer, _patched_forward, "_original_forward") + + def _early_stop_forward(module_self, *args, **kwargs): + try: + return module_self._original_forward(*args, **kwargs) + except _EarlyStopForwardError: + return None + + bind_forward_method(self.model, _early_stop_forward, "_original_forward") + except Exception: + self._cleanup_layers() + raise + + self._patched = True + + def _cleanup_layers(self): + """Best-effort cleanup of any patched layers and model forward.""" + if self._decoder_layers is not None: + for idx, layer in enumerate(self._decoder_layers): + if isinstance(layer, _SkipLayer): + self._decoder_layers[idx] = layer._original + + if hasattr(self.model, "_original_forward"): + unpatch_forward_method(self.model, "_original_forward") + + if self._decoder_layers is not None: + for layer in self._decoder_layers: + if hasattr(layer, "_original_forward"): + unpatch_forward_method(layer, "_original_forward") + if hasattr(layer, self._LAYER_ATTR): + delattr(layer, self._LAYER_ATTR) + + def _unpatch_all_layers(self): + """Restore original forwards and clean up state attributes. Called once.""" + if not self._patched: + return + self._cleanup_layers() + self._patched = False + + def _set_layer_states(self, layer_idx: int): + """Transition layer modes for the next calibration step. + + When calibrating layer *i*, three transitions happen: + + * Layer ``i - 2`` → **skip** (fully done, free its cached inputs). + * Layer ``i - 1`` → **run** (replay captured inputs with calibrated weights). + * Layer ``i`` → **capture** (record inputs, then early-stop). + """ + assert self._decoder_layers is not None + + if layer_idx > 1: + idx = layer_idx - 2 + if not isinstance(self._decoder_layers[idx], _SkipLayer): + self._swap_to_dummy(idx) + + if layer_idx > 0: + prev = self._decoder_layers[layer_idx - 1]._layerwise_calib + if not prev.collected_inputs: + raise RuntimeError( + f"Layer {layer_idx - 1} ({prev.name!r}) has no collected inputs to replay. " + "Layers must be calibrated sequentially — ensure get_input_activations() " + "was called for every preceding layer in order." + ) + prev.mode = "run" + prev.cached_inputs = deque(prev.collected_inputs) + prev.collected_inputs = [] + + cur = self._decoder_layers[layer_idx]._layerwise_calib + cur.mode = "capture" + cur.collected_inputs = [] + + def _log_layer_summary(self, layer_idx: int): + """Log a one-line summary of layer modes for the current calibration step.""" + assert self._decoder_layers is not None + n = len(self._decoder_layers) + groups: dict[str, list[int]] = {} + for i, layer in enumerate(self._decoder_layers): + mode = layer._layerwise_calib.mode + if mode in ("skip", "run", "capture"): + groups.setdefault(mode, []).append(i + 1) + + parts = [] + for mode in ("skip", "run", "capture"): + if mode not in groups: + continue + ids = groups[mode] + parts.append(f"{mode}: {len(ids)}" if mode == "skip" else f"{mode}: {ids}") + print_rank_0(f"Calibrating layer {layer_idx + 1}/{n} | {' | '.join(parts)}") + + @torch.no_grad() + def get_input_activations(self, layer: torch.nn.Module, forward_loop: ForwardLoop) -> list: + """Collect input activations for *layer* by running a full model forward. + + Layers before the target are skipped or re-run (if just calibrated), the + target layer captures its inputs, and an early-stop prevents unnecessary + computation beyond the target. + + :meth:`_patch_all_layers` must be called before this method. + + Note: the model forward returns ``None`` for every batch during capture + (because ``_EarlyStopForwardError`` short-circuits the forward pass). + Callers should not rely on the model's return value within *forward_loop*. + """ + if not self._patched: + raise RuntimeError( + "get_input_activations() requires _patch_all_layers() to be called first." + ) + layer_idx = self._layer_to_idx[layer] + self._set_layer_states(layer_idx) + self._log_layer_summary(layer_idx) + + info = layer._layerwise_calib + try: + forward_loop(self.model) + except Exception: + # Reset the current layer so subsequent calls don't see stale state. + info.mode = "original" + info.collected_inputs = [] + raise + + if not info.collected_inputs: + info.mode = "original" + raise RuntimeError( + f"Layer {info.name!r} collected no inputs during forward_loop. " + "The forward loop did not reach this layer — check that forward_loop() " + "actually calls the model and that the layer is in the forward path." + ) + + inputs = list(info.collected_inputs) + # Reset to original so calib_func can call the layer's real forward + # directly. The layer will transition to run → skip in subsequent + # iterations via _set_layer_states. + info.mode = "original" + return inputs + + def get_first_layer_inputs( + self, + start_layer: int, + resumed_inputs: list | None, + forward_loop: ForwardLoop, + ) -> list: + """Get inputs for the first layer to calibrate, handling resume. + + If *resumed_inputs* is provided, sets skip mode on layers ``0..start_layer-1`` + and seeds the start layer's ``collected_inputs`` for subsequent + ``cache_outputs_for_next_layer_calib`` calls. Otherwise, captures inputs + via a normal forward pass. + """ + assert self._decoder_layers is not None + + if resumed_inputs is not None: + print_rank_0(f"Calibrating layer {start_layer + 1} (resumed)") + for i in range(start_layer): + self._swap_to_dummy(i) + layer = self._decoder_layers[start_layer] + layer._layerwise_calib.collected_inputs = resumed_inputs + layer._layerwise_calib.mode = "original" + return resumed_inputs + + return self.get_input_activations(self._decoder_layers[start_layer], forward_loop) + + @torch.no_grad() + def cache_outputs_for_next_layer_calib( + self, layer: torch.nn.Module, forward_loop: ForwardLoop + ) -> list: + """Run a forward pass after calibrating *layer* to capture the next layer's inputs. + + This puts *layer* into "run" mode (setting its ``output_meta``) and the + next layer into "capture" mode, then runs *forward_loop*. Returns the + captured inputs for the next layer. + + Must be called only when a next layer exists (i.e. *layer* is not the + last decoder layer). + """ + assert self._decoder_layers is not None + layer_idx = self._layer_to_idx[layer] + next_idx = layer_idx + 1 + assert next_idx < len(self._decoder_layers), "No next layer to capture inputs for." + from .core_utils import persistent_materialization + + next_layer = self._decoder_layers[next_idx] + with persistent_materialization(layer): + return self.get_input_activations(next_layer, forward_loop) + + +def _move_to_device(obj: Any, device: torch.device) -> Any: + """Recursively move tensors to *device*. Non-tensors are returned as-is.""" + if isinstance(obj, torch.Tensor): + return obj.to(device) + if isinstance(obj, dict): + return {k: _move_to_device(v, device) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + moved = [_move_to_device(v, device) for v in obj] + return type(obj)(moved) + return obj + + +def _remap_output_metadata_device(meta: tuple, device: torch.device) -> tuple: + """Patch the device field inside output_meta tuples so _zeros_from_meta uses *device*.""" + tag = meta[0] + if tag == "tensor": + _, shape, dtype, _old_device = meta + return ("tensor", shape, dtype, device) + if tag == "tuple": + return ("tuple", tuple(_remap_output_metadata_device(m, device) for m in meta[1])) + if tag == "list": + return ("list", [_remap_output_metadata_device(m, device) for m in meta[1]]) + return meta + + +def _read_manifest(checkpoint_dir: str) -> dict | None: + """Read manifest.json from *checkpoint_dir*. Returns None if missing or corrupt.""" + path = os.path.join(checkpoint_dir, "manifest.json") + if not os.path.isfile(path): + return None + try: + with open(path) as f: + return json.load(f) + except (json.JSONDecodeError, OSError): + return None + + +def _write_manifest(checkpoint_dir: str, last_completed_layer: int, num_layers: int) -> None: + """Atomically write manifest.json.""" + path = os.path.join(checkpoint_dir, "manifest.json") + tmp = path + ".tmp" + with open(tmp, "w") as f: + json.dump( + {"last_completed_layer": last_completed_layer, "num_layers": num_layers}, + f, + ) + os.replace(tmp, path) + + +def _layer_dir(checkpoint_dir: str, idx: int) -> str: + return os.path.join(checkpoint_dir, f"layer_{idx:04d}") + + +def _save_layer( + checkpoint_dir: str, + idx: int, + weights: dict, + qstate: dict, + output_meta: tuple, + next_inputs: list | None, + num_layers: int, +) -> None: + """Save a single layer checkpoint and update the manifest atomically.""" + d = _layer_dir(checkpoint_dir, idx) + if os.path.isdir(d): + shutil.rmtree(d) + os.makedirs(d) + torch.save(weights, os.path.join(d, "weights.pt")) + torch.save(qstate, os.path.join(d, "quantizer_state.pt")) + torch.save(output_meta, os.path.join(d, "output_meta.pt")) + if next_inputs is not None: + torch.save(next_inputs, os.path.join(d, "next_inputs.pt")) + _write_manifest(checkpoint_dir, idx, num_layers) + + +def detect_resume_point(checkpoint_dir: str) -> tuple[int, dict] | None: + """Detect where to resume from an existing checkpoint directory. + + Returns ``(start_layer, manifest)`` if there is work to resume, + or ``None`` if the directory is empty, corrupt, or calibration was already complete. + """ + manifest = _read_manifest(checkpoint_dir) + if manifest is None: + return None + last = manifest.get("last_completed_layer") + total = manifest.get("num_layers") + if last is None or total is None: + return None + if last + 1 >= total: + return None + return (last + 1, manifest) + + +class _CheckpointState: + """Manages checkpoint save and restore for layerwise calibration. + + Handles both saving per-layer checkpoints during calibration and + restoring from a previous partial run. + + .. todo:: + Support distributed checkpoint save/restore for FSDP2: + use ``torch.distributed.checkpoint`` (or save only from rank 0 + barrier) + and broadcast restored state to all ranks during resume. + """ + + def __init__(self, checkpoint_dir: str, num_layers: int, start_layer: int = 0): + if dist.is_initialized() and dist.size() > 1: + raise RuntimeError( + "Layerwise calibration checkpointing is not supported in " + "multi-process distributed jobs (e.g. FSDP2). " + "Use single-process calibration or disable checkpointing." + ) + + self.checkpoint_dir = checkpoint_dir + self.num_layers = num_layers + self.start_layer = start_layer + + @classmethod + def from_folder(cls, checkpoint_dir: str | None, num_layers: int) -> _CheckpointState | None: + """Create from folder. Detects resume point. Returns None if no checkpoint_dir.""" + if not checkpoint_dir: + return None + os.makedirs(checkpoint_dir, exist_ok=True) + info = detect_resume_point(checkpoint_dir) + if info is not None: + manifest_num_layers = info[1].get("num_layers") + if manifest_num_layers is not None and manifest_num_layers != num_layers: + raise ValueError( + f"Checkpoint num_layers mismatch: manifest has {manifest_num_layers} " + f"but model has {num_layers}. Use a fresh checkpoint directory." + ) + start = info[0] if info else 0 + if start > 0: + print_rank_0( + f"Checkpoint: resuming layerwise calibration from layer {start}/{num_layers}" + ) + return cls(checkpoint_dir, num_layers, start_layer=start) + + def setup_resume(self, layers: nn.ModuleList) -> list | None: + """Load output_meta for skip layers 0..K-1, return next_inputs for layer K. + + Sets ``output_meta`` on each already-calibrated layer so that + skip mode can produce correctly shaped dummy outputs. + """ + if self.start_layer == 0: + return None + + last_ckpt = self.start_layer - 1 + + for i in range(self.start_layer): + d = _layer_dir(self.checkpoint_dir, i) + # weights_only=False is safe: file is internally generated by _save_layer, not user-supplied + meta = torch.load( + os.path.join(d, "output_meta.pt"), map_location="cpu", weights_only=False + ) + layer_device = get_module_device(layers[i]) + meta = _remap_output_metadata_device(meta, layer_device) + layers[i]._layerwise_calib.output_meta = meta + + d = _layer_dir(self.checkpoint_dir, last_ckpt) + next_inputs_path = os.path.join(d, "next_inputs.pt") + if not os.path.isfile(next_inputs_path): + raise FileNotFoundError(f"Cannot resume: next_inputs.pt missing for layer {last_ckpt}") + # weights_only=False is safe: file is internally generated by _save_layer, not user-supplied + next_inputs = torch.load(next_inputs_path, map_location="cpu", weights_only=False) + resume_device = get_module_device(layers[self.start_layer]) + next_inputs = _move_to_device(next_inputs, resume_device) + return next_inputs + + def full_restore(self, layers: nn.ModuleList, model: nn.Module) -> None: + """Restore weights and quantizer state for layers 0..K-1 after the calibration loop.""" + from modelopt.torch.quantization.config import QuantizeConfig + from modelopt.torch.quantization.conversion import restore_quantizer_state + from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback + + if self.start_layer == 0: + return + + dummy_config = QuantizeConfig() + name_to_module = dict(model.named_modules()) + for i in range(self.start_layer): + layer = layers[i] + d = _layer_dir(self.checkpoint_dir, i) + + # Resolve layer_device and load inside the context so params are + # materialized — otherwise get_module_device can return meta. + with enable_weight_access_and_writeback(layer, model, name_to_module): + layer_device = get_module_device(layer) + # weights_only=False is safe: files are internally generated by _save_layer + qstate = torch.load( + os.path.join(d, "quantizer_state.pt"), + map_location=layer_device, + weights_only=False, + ) + weights = torch.load( + os.path.join(d, "weights.pt"), + map_location=layer_device, + weights_only=False, + ) + restore_quantizer_state(layer, dummy_config, {"quantizer_state": qstate}) + layer.load_state_dict(weights, strict=False, assign=True) + + print_rank_0(f"Checkpoint: restored {self.start_layer} previously calibrated layers") + + def save( + self, + layer_idx: int, + layer: nn.Module, + model: nn.Module, + layers: nn.ModuleList, + next_layer_inputs: list | None = None, + ) -> None: + """Snapshot layer state and write checkpoint to disk in one step. + + Args: + layer_idx: Index of the layer just calibrated. + layer: The layer module (weights may be on GPU or managed by accelerate/FSDP2). + model: The full model (needed for ``enable_weight_access_and_writeback``). + layers: The decoder layer list (to read ``output_meta``). + next_layer_inputs: Inputs for the next layer (``None`` for the final layer). + """ + from modelopt.torch.quantization.conversion import quantizer_state + from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback + + _cpu = torch.device("cpu") + with enable_weight_access_and_writeback(layer, model): + weights = _move_to_device(layer.state_dict(), _cpu) + qstate = _move_to_device(quantizer_state(layer), _cpu) + + output_meta = getattr(layer._layerwise_calib, "output_meta", None) + if output_meta is None: + # Placeholder for the last layer: output_meta is never used for skip mode + # since there is no subsequent layer that needs a correctly shaped dummy output. + output_meta = LayerActivationCollector._extract_output_meta(torch.zeros(1)) + + _save_layer( + self.checkpoint_dir, + layer_idx, + weights, + qstate, + _move_to_device(output_meta, _cpu), + _move_to_device(next_layer_inputs, _cpu) if next_layer_inputs is not None else None, + self.num_layers, + ) + suffix = " (final)" if next_layer_inputs is None else "" + print_rank_0(f"Checkpoint: saved layer {layer_idx}{suffix}") diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 1e9a7fbbbd8..01cb3abe88f 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -601,16 +601,28 @@ def _forward_loop( dataloader: DataLoader containing the batched input data allowed_non_tensor_keys: Set of key names whose values may be non-tensor types """ - with torch.no_grad(): - is_enc_dec = model_type_is_enc_dec(model) - infer_method = model.generate if is_enc_dec else model.forward - max_working_batch_size = None # Initialize max working batch size as None + # Disable KV caching during calibration — it is unnecessary overhead and causes + # correctness issues with hybrid Mamba/attention models whose cache state is mutated + # in-place (e.g., NemotronH). + config = getattr(model, "config", None) + prev_use_cache = getattr(config, "use_cache", None) + if config is not None and prev_use_cache is not None: + config.use_cache = False - for _, data in enumerate(tqdm(dataloader)): - # Process batch and update max working batch size - max_working_batch_size = _process_batch( - data, infer_method, max_working_batch_size, allowed_non_tensor_keys - ) + try: + with torch.no_grad(): + is_enc_dec = model_type_is_enc_dec(model) + infer_method = model.generate if is_enc_dec else model.forward + max_working_batch_size = None # Initialize max working batch size as None + + for _, data in enumerate(tqdm(dataloader)): + # Process batch and update max working batch size + max_working_batch_size = _process_batch( + data, infer_method, max_working_batch_size, allowed_non_tensor_keys + ) + finally: + if config is not None and prev_use_cache is not None: + config.use_cache = prev_use_cache def create_forward_loop( diff --git a/modelopt/torch/utils/network.py b/modelopt/torch/utils/network.py index b54332375b6..440ca522d12 100644 --- a/modelopt/torch/utils/network.py +++ b/modelopt/torch/utils/network.py @@ -90,12 +90,43 @@ def is_parallel(model: nn.Module) -> bool: return isinstance(model, (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)) +def _get_execution_device_from_hook(module: nn.Module) -> torch.device | None: + """Extract the execution device from an accelerate ``_hf_hook``, if present. + + Handles both ``AlignDevicesHook`` (direct) and ``SequentialHook`` (which + may wrap one or more ``AlignDevicesHook`` instances). Returns ``None`` + when no hook is found or the hook carries no ``execution_device``. + """ + hook = getattr(module, "_hf_hook", None) + if hook is None: + return None + + dev = getattr(hook, "execution_device", None) + if dev is not None: + return torch.device("cuda", dev) if isinstance(dev, int) else torch.device(dev) + + for h in getattr(hook, "hooks", ()): + dev = getattr(h, "execution_device", None) + if dev is not None: + return torch.device("cuda", dev) if isinstance(dev, int) else torch.device(dev) + + return None + + def get_module_device(module: nn.Module) -> torch.device: - """Get the device of a PyTorch module.""" + """Get the device of a PyTorch module. + + For modules managed by accelerate (``_hf_hook``), returns the hook's + ``execution_device`` which is the authoritative device even when + parameters are offloaded to CPU/meta between forward calls. + """ + hook_device = _get_execution_device_from_hook(module) + if hook_device is not None: + return hook_device + try: return next(module.parameters()).device except StopIteration: - # For modules without parameters return torch.device("cpu") @@ -590,21 +621,29 @@ def get_unwrapped_name(name: str, model: nn.Module | None = None) -> str: @contextmanager def temporarily_remove_accelerate_hook(module): - """Context manager to temporarily remove accelerate hook from a module.""" - accelerate_hook = None - if hasattr(module, "_hf_hook"): - # A module with forward method patched by accelerate - from accelerate.hooks import add_hook_to_module, remove_hook_from_module + """Context manager to temporarily bypass the accelerate hook on a module. + + Swaps ``module.forward`` with the pre-hook forward (``_old_forward``) so + that code inside the context sees the un-hooked forward. On exit the + hook-wrapped forward is restored and ``_old_forward`` is updated to + reflect any changes made inside the context. - accelerate_hook = module._hf_hook - remove_hook_from_module(module) + This avoids ``remove_hook_from_module`` / ``add_hook_to_module`` entirely, + sidestepping ``init_hook`` which would call ``set_module_tensor_to_device`` + and fail when newly-added quantizer modules have weights on the meta device. + """ + hooked_forward = None + cached_old_forward = None + if hasattr(module, "_hf_hook"): + hooked_forward = module.forward + cached_old_forward = module._old_forward + module.forward = cached_old_forward try: yield finally: - if accelerate_hook is not None: - from accelerate.hooks import add_hook_to_module - - add_hook_to_module(module, accelerate_hook) + if hooked_forward is not None: + module._old_forward = module.forward + module.forward = hooked_forward def bind_forward_method( diff --git a/modelopt_recipes/general/ptq/nvfp4_default-fp8_kv.yaml b/modelopt_recipes/general/ptq/nvfp4_default-fp8_kv.yaml index 6fe4a8c3d12..862929ef34c 100644 --- a/modelopt_recipes/general/ptq/nvfp4_default-fp8_kv.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_default-fp8_kv.yaml @@ -15,7 +15,7 @@ metadata: recipe_type: ptq - description: NVFP4 MLP/MoE weight only (W4A16), FP8 KV cache, max calibration. + description: NVFP4 W4A4, FP8 KV cache, max calibration. quantize: algorithm: max quant_cfg: diff --git a/modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml b/modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml index a62051b659a..99098c9d6d0 100644 --- a/modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml @@ -15,11 +15,12 @@ metadata: recipe_type: ptq - description: NVFP4 weight and activation (W4A4), gptq sequential calibration. + description: NVFP4 weight and activation (W4A4), gptq layerwise calibration. quantize: algorithm: method: gptq - use_sequential: true + layerwise: true + layerwise_checkpoint_dir: output/layerwise_ckpts/ quant_cfg: - quantizer_name: '*' enable: false diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml index cc332733a02..220d0622327 100644 --- a/modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml @@ -15,9 +15,12 @@ metadata: recipe_type: ptq - description: NVFP4 static weight and dynamic activation for expert layers only (W4A4), FP8 KV cache, max calibration. + description: NVFP4 static weight and dynamic activation for expert layers only (W4A4), FP8 KV cache, max layerwise calibration. quantize: - algorithm: max + algorithm: + method: max + # Max calibration is fast and does not typically need checkpointing. + layerwise: true quant_cfg: - quantizer_name: '*' enable: false diff --git a/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py b/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py index 8ee71ed453e..df5610ca3df 100644 --- a/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py +++ b/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py @@ -12,16 +12,19 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import copy from copy import deepcopy import pytest import torch from _test_utils.torch.transformers_models import create_tiny_llama_dir -from transformers import AutoModelForCausalLM +from accelerate import init_empty_weights, load_checkpoint_and_dispatch +from transformers import AutoConfig, AutoModelForCausalLM import modelopt.torch.quantization as mtq from modelopt.torch.export import export_hf_vllm_fq_checkpoint from modelopt.torch.quantization.model_quant import fold_weight +from modelopt.torch.quantization.utils import enable_weight_access_and_writeback from modelopt.torch.utils import safe_load @@ -111,3 +114,120 @@ def forward_loop(model): "_amax" in k for k in quantizer_state_dict_before[name] ): assert any("_amax" in k for k in state), f"input quantizer {name} should preserve _amax" + + +def _make_cpu_offloaded_model(tmp_path, num_hidden_layers=3): + """Create a tiny LLaMA model with layer 0 offloaded to CPU via accelerate.""" + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_hidden_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + + device_map = { + n: 0 + for n, m in model.named_modules() + if "layers" not in n or n.split("layers.")[-1].isdigit() + } + device_map["model.layers.0"] = "cpu" + + model = load_checkpoint_and_dispatch(model, tiny_llama_dir, device_map=device_map) + return model, config, tiny_llama_dir + + +def _make_layerwise_cfg(base_cfg): + """Add layerwise=True to a quant config's algorithm field.""" + cfg = copy.deepcopy(base_cfg) + algo = cfg.get("algorithm", "max") + if isinstance(algo, str): + cfg["algorithm"] = {"method": algo, "layerwise": True} + else: + algo["layerwise"] = True + return cfg + + +@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG]) +def test_hf_vllm_export_offload(tmp_path, quant_cfg): + """Verifies the inplace_mem_efficient=True path mutates offloaded weights in place + and produces folded values matching deepcopy+fold_weight reference. Does NOT + exercise save_pretrained -- transformers' load_offloaded_parameter doesn't unwrap + SequentialHook, a pre-existing limitation unrelated to this PR's new code. + """ + num_hidden_layers = 3 + + model, _config, _tiny_llama_dir = _make_cpu_offloaded_model( + tmp_path / "offloaded", num_hidden_layers=num_hidden_layers + ) + model.eval() + + seq_cfg = _make_layerwise_cfg(quant_cfg) + + def forward_loop(model): + input_ids = torch.randint(0, model.config.vocab_size, (1, 128)).cuda() + with torch.no_grad(): + model(input_ids) + + model = mtq.quantize(model, seq_cfg, forward_loop) + quantizer_state_dict_before = mtq.utils.get_quantizer_state_dict(model) + + folded_model = deepcopy(model) + with enable_weight_access_and_writeback(folded_model.model.layers[0], folded_model): + fold_weight(folded_model) + expected_weights = { + k: v.detach().clone() + for k, v in folded_model.state_dict().items() + if "quantizer" not in k + } + del folded_model + + export_dir = tmp_path / "vllm_export_offload" + export_dir.mkdir(exist_ok=True) + + # Snapshot the offloaded layer's weight before/after export to verify the + # inplace_mem_efficient path actually mutates offloaded weights (would otherwise + # be unfalsifiable if the function silently took the copy path). + with enable_weight_access_and_writeback(model.model.layers[0], model): + weight_before = model.model.layers[0].self_attn.q_proj.weight.data.clone() + + # Skip save_pretrained: transformers' load_offloaded_parameter doesn't unwrap + # SequentialHook, a pre-existing upstream limitation unrelated to this PR. The + # delta under test is inplace fake-quant + weight writeback, which runs before + # save_pretrained. + original_save_pretrained = model.save_pretrained + model.save_pretrained = lambda *args, **kwargs: None + try: + export_hf_vllm_fq_checkpoint(model, export_dir=export_dir, inplace_mem_efficient=True) + finally: + model.save_pretrained = original_save_pretrained + + with enable_weight_access_and_writeback(model.model.layers[0], model): + weight_after = model.model.layers[0].self_attn.q_proj.weight.data.clone() + assert not torch.equal(weight_before, weight_after), ( + "inplace path must mutate offloaded weights" + ) + + with enable_weight_access_and_writeback(model.model.layers[0], model): + actual_weights = { + k: v.detach().clone() for k, v in model.state_dict().items() if "quantizer" not in k + } + for key, expected in expected_weights.items(): + actual = actual_weights.get(key) + assert actual is not None, f"missing {key} after export" + assert torch.allclose(actual, expected, atol=1e-6), f"mismatch at {key}" + + modelopt_state_file = export_dir / "vllm_fq_modelopt_state.pth" + assert modelopt_state_file.exists(), ( + f"vllm_fq_modelopt_state.pth file should be created in {export_dir}" + ) + + quantizer_state_dict = safe_load(modelopt_state_file)["modelopt_state_weights"] + assert len(quantizer_state_dict) > 0, ( + f"modelopt_state_weights should not be empty in {modelopt_state_file}" + ) + for name, state in quantizer_state_dict.items(): + if "weight_quantizer" in name: + assert state == {}, f"weight quantizer {name} should have empty state after fold" + elif "input_quantizer" in name and any( + "_amax" in k for k in quantizer_state_dict_before[name] + ): + assert any("_amax" in k for k in state), f"input quantizer {name} should preserve _amax" diff --git a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py index 8ed1039e59d..49e74e5851f 100644 --- a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py +++ b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py @@ -13,9 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy +import json +import os +import shutil + import pytest import torch -from _test_utils.torch.quantization.quantize_common import INT4_AWQ_CLIP_CFG from _test_utils.torch.transformers_models import create_tiny_llama_dir from accelerate import init_empty_weights, load_checkpoint_and_dispatch from transformers import AutoConfig, AutoModelForCausalLM @@ -25,19 +29,11 @@ enable_weight_access_and_writeback, is_quantized_linear, ) +from modelopt.torch.quantization.utils.layerwise_calib import _layer_dir -@pytest.mark.parametrize( - "quant_cfg", - [ - mtq.INT4_AWQ_CFG, - mtq.INT8_SMOOTHQUANT_CFG, - INT4_AWQ_CLIP_CFG, - mtq.NVFP4_SVDQUANT_DEFAULT_CFG, - mtq.INT8_DEFAULT_CFG, - ], -) -def test_cpu_offloaded_tinyllama(tmp_path, quant_cfg): +def test_cpu_offloaded_tinyllama(tmp_path): + quant_cfg = mtq.NVFP4_AWQ_LITE_CFG tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=2) config = AutoConfig.from_pretrained(tiny_llama_dir) @@ -73,3 +69,575 @@ def test_cpu_offloaded_tinyllama(tmp_path, quant_cfg): assert torch.allclose(module.weight, model_ref.get_submodule(name).weight) assert torch.allclose(output_ref.logits, output_test.logits) + + +def _make_cpu_offloaded_model(tmp_path, num_hidden_layers=3): + """Create a tiny LLaMA model with layer 0 offloaded to CPU via accelerate.""" + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_hidden_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + + device_map = { + n: 0 + for n, m in model.named_modules() + if "layers" not in n or n.split("layers.")[-1].isdigit() + } + device_map["model.layers.0"] = "cpu" + + model = load_checkpoint_and_dispatch(model, tiny_llama_dir, device_map=device_map) + inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda() + return model, config, tiny_llama_dir, inputs + + +def _make_layerwise_cfg(base_cfg): + """Add layerwise=True to a quant config's algorithm field.""" + cfg = copy.deepcopy(base_cfg) + algo = cfg.get("algorithm", "max") + if isinstance(algo, str): + cfg["algorithm"] = {"method": algo, "layerwise": True} + else: + algo["layerwise"] = True + return cfg + + +def _make_layerwise_checkpoint_cfg(base_cfg, checkpoint_dir): + """Add layerwise=True and layerwise_checkpoint_dir to a quant config's algorithm field.""" + cfg = _make_layerwise_cfg(base_cfg) + cfg["algorithm"]["layerwise_checkpoint_dir"] = checkpoint_dir + return cfg + + +@pytest.mark.parametrize("use_checkpoint", [False, True], ids=["no_ckpt", "ckpt"]) +def test_layerwise_calibrate_cpu_offloaded(tmp_path, use_checkpoint): + """Layerwise calibration on CPU-offloaded model matches GPU-only reference.""" + quant_cfg = mtq.NVFP4_AWQ_LITE_CFG + num_layers = 3 + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda() + + if use_checkpoint: + ckpt_dir = str(tmp_path / "seq_ckpt") + seq_cfg = _make_layerwise_checkpoint_cfg(quant_cfg, ckpt_dir) + else: + seq_cfg = _make_layerwise_cfg(quant_cfg) + + # Reference: GPU-only model with layerwise calibration + ref_cfg = _make_layerwise_cfg(quant_cfg) + model_ref = AutoModelForCausalLM.from_pretrained( + tiny_llama_dir, torch_dtype=config.torch_dtype + ).cuda() + mtq.quantize(model_ref, ref_cfg, lambda model: model(inputs)) + output_ref = model_ref(inputs) + + # Test: CPU-offloaded model + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + device_map = { + n: 0 + for n, m in model.named_modules() + if "layers" not in n or n.split("layers.")[-1].isdigit() + } + device_map["model.layers.0"] = "cpu" + model = load_checkpoint_and_dispatch(model, tiny_llama_dir, device_map=device_map) + + mtq.quantize(model, seq_cfg, lambda model: model(inputs)) + output_test = model(inputs) + + for name, module in model.named_modules(): + if is_quantized_linear(module): + with enable_weight_access_and_writeback(module, model): + assert torch.allclose(module.weight, model_ref.get_submodule(name).weight), ( + f"Weight mismatch at {name}" + ) + + assert torch.allclose(output_ref.logits, output_test.logits) + + if use_checkpoint: + manifest_path = os.path.join(ckpt_dir, "manifest.json") + assert os.path.isfile(manifest_path) + with open(manifest_path) as f: + manifest = json.load(f) + assert manifest["last_completed_layer"] == num_layers - 1 + assert manifest["num_layers"] == num_layers + + +def test_sequential_checkpoint_resume_cpu_offloaded(tmp_path): + """Resume from a partial checkpoint on a CPU-offloaded model matches a full run.""" + quant_cfg = mtq.NVFP4_AWQ_LITE_CFG + num_layers = 3 + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda() + + ckpt_dir = str(tmp_path / "seq_ckpt") + seq_ckpt_cfg = _make_layerwise_checkpoint_cfg(quant_cfg, ckpt_dir) + + # Full reference run with checkpointing + with init_empty_weights(): + model_ref = AutoModelForCausalLM.from_config(config) + device_map = { + n: 0 + for n, m in model_ref.named_modules() + if "layers" not in n or n.split("layers.")[-1].isdigit() + } + device_map["model.layers.0"] = "cpu" + model_ref = load_checkpoint_and_dispatch(model_ref, tiny_llama_dir, device_map=device_map) + mtq.quantize(model_ref, seq_ckpt_cfg, lambda model: model(inputs)) + output_ref = model_ref(inputs) + + # Simulate crash after layer 0 by truncating the manifest and removing later layers + last_completed_layer = 0 + manifest_path = os.path.join(ckpt_dir, "manifest.json") + with open(manifest_path, "w") as f: + json.dump({"last_completed_layer": last_completed_layer, "num_layers": num_layers}, f) + for i in range(last_completed_layer + 1, num_layers): + d = _layer_dir(ckpt_dir, i) + if os.path.isdir(d): + shutil.rmtree(d) + + # Resume from a fresh CPU-offloaded model + with init_empty_weights(): + model_resumed = AutoModelForCausalLM.from_config(config) + device_map = { + n: 0 + for n, m in model_resumed.named_modules() + if "layers" not in n or n.split("layers.")[-1].isdigit() + } + device_map["model.layers.0"] = "cpu" + model_resumed = load_checkpoint_and_dispatch( + model_resumed, tiny_llama_dir, device_map=device_map + ) + mtq.quantize(model_resumed, seq_ckpt_cfg, lambda model: model(inputs)) + output_resumed = model_resumed(inputs) + + assert torch.allclose(output_ref.logits, output_resumed.logits), ( + "Resumed checkpoint should produce identical output to full run" + ) + + +def test_sequential_checkpoint_resume_multi_offload(tmp_path): + """Resume with multiple layers offloaded exercises per-layer device resolution.""" + num_layers = 3 + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda() + + ckpt_dir = str(tmp_path / "seq_ckpt") + seq_ckpt_cfg = _make_layerwise_checkpoint_cfg(mtq.INT4_AWQ_CFG, ckpt_dir) + + def _make_multi_offload_model(): + with init_empty_weights(): + m = AutoModelForCausalLM.from_config(config) + dmap = { + n: 0 + for n, mod in m.named_modules() + if "layers" not in n or n.split("layers.")[-1].isdigit() + } + dmap["model.layers.0"] = "cpu" + dmap["model.layers.1"] = "cpu" + return load_checkpoint_and_dispatch(m, tiny_llama_dir, device_map=dmap) + + # Full reference run + model_ref = _make_multi_offload_model() + mtq.quantize(model_ref, seq_ckpt_cfg, lambda model: model(inputs)) + output_ref = model_ref(inputs) + + # Simulate crash after layer 0 + last_completed_layer = 0 + manifest_path = os.path.join(ckpt_dir, "manifest.json") + with open(manifest_path, "w") as f: + json.dump({"last_completed_layer": last_completed_layer, "num_layers": num_layers}, f) + for i in range(last_completed_layer + 1, num_layers): + d = _layer_dir(ckpt_dir, i) + if os.path.isdir(d): + shutil.rmtree(d) + + # Resume from fresh model with same offload layout + model_resumed = _make_multi_offload_model() + mtq.quantize(model_resumed, seq_ckpt_cfg, lambda model: model(inputs)) + output_resumed = model_resumed(inputs) + + assert torch.allclose(output_ref.logits, output_resumed.logits), ( + "Resumed checkpoint with multi-offload should match full run" + ) + + +def _make_gptq_sequential_cfg(base_cfg): + """Create a sequential GPTQ config from a base quantization config.""" + cfg = copy.deepcopy(base_cfg) + cfg["algorithm"] = {"method": "gptq", "layerwise": True} + return cfg + + +def _make_gptq_sequential_checkpoint_cfg(base_cfg, checkpoint_dir): + """Create a sequential GPTQ config with checkpoint dir.""" + cfg = _make_gptq_sequential_cfg(base_cfg) + cfg["algorithm"]["layerwise_checkpoint_dir"] = checkpoint_dir + return cfg + + +@pytest.mark.parametrize("use_checkpoint", [False, True], ids=["no_ckpt", "ckpt"]) +def test_sequential_gptq_cpu_offloaded(tmp_path, use_checkpoint): + """Sequential GPTQ (weight-modifying) on CPU-offloaded model matches GPU-only reference.""" + num_layers = 3 + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda() + + if use_checkpoint: + ckpt_dir = str(tmp_path / "gptq_ckpt") + seq_cfg = _make_gptq_sequential_checkpoint_cfg(mtq.NVFP4_AWQ_LITE_CFG, ckpt_dir) + else: + seq_cfg = _make_gptq_sequential_cfg(mtq.NVFP4_AWQ_LITE_CFG) + + # Reference: GPU-only model + ref_cfg = _make_gptq_sequential_cfg(mtq.NVFP4_AWQ_LITE_CFG) + model_ref = AutoModelForCausalLM.from_pretrained( + tiny_llama_dir, torch_dtype=config.torch_dtype + ).cuda() + mtq.quantize(model_ref, ref_cfg, lambda model: model(inputs)) + output_ref = model_ref(inputs) + + # Test: CPU-offloaded model + model, _, _, _ = _make_cpu_offloaded_model(tmp_path / "offloaded", num_hidden_layers=num_layers) + mtq.quantize(model, seq_cfg, lambda model: model(inputs)) + output_test = model(inputs) + + for name, module in model.named_modules(): + if is_quantized_linear(module): + with enable_weight_access_and_writeback(module, model): + assert torch.allclose(module.weight, model_ref.get_submodule(name).weight), ( + f"Weight mismatch at {name}" + ) + + assert torch.allclose(output_ref.logits, output_test.logits) + + +def test_sequential_gptq_checkpoint_resume_cpu_offloaded(tmp_path): + """GPTQ checkpoint resume with CPU offloading restores modified weights correctly.""" + num_layers = 3 + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda() + + ckpt_dir = str(tmp_path / "gptq_ckpt") + seq_ckpt_cfg = _make_gptq_sequential_checkpoint_cfg(mtq.NVFP4_AWQ_LITE_CFG, ckpt_dir) + + # Full reference run with checkpointing + with init_empty_weights(): + model_ref = AutoModelForCausalLM.from_config(config) + device_map = { + n: 0 + for n, m in model_ref.named_modules() + if "layers" not in n or n.split("layers.")[-1].isdigit() + } + device_map["model.layers.0"] = "cpu" + model_ref = load_checkpoint_and_dispatch(model_ref, tiny_llama_dir, device_map=device_map) + mtq.quantize(model_ref, seq_ckpt_cfg, lambda model: model(inputs)) + output_ref = model_ref(inputs) + + # Simulate crash after layer 0 + last_completed_layer = 0 + manifest_path = os.path.join(ckpt_dir, "manifest.json") + with open(manifest_path, "w") as f: + json.dump({"last_completed_layer": last_completed_layer, "num_layers": num_layers}, f) + for i in range(last_completed_layer + 1, num_layers): + d = _layer_dir(ckpt_dir, i) + if os.path.isdir(d): + shutil.rmtree(d) + + # Resume from fresh CPU-offloaded model + with init_empty_weights(): + model_resumed = AutoModelForCausalLM.from_config(config) + device_map = { + n: 0 + for n, m in model_resumed.named_modules() + if "layers" not in n or n.split("layers.")[-1].isdigit() + } + device_map["model.layers.0"] = "cpu" + model_resumed = load_checkpoint_and_dispatch( + model_resumed, tiny_llama_dir, device_map=device_map + ) + mtq.quantize(model_resumed, seq_ckpt_cfg, lambda model: model(inputs)) + output_resumed = model_resumed(inputs) + + assert torch.allclose(output_ref.logits, output_resumed.logits), ( + "GPTQ resumed checkpoint should produce identical output to full run" + ) + + +class _TupleReturningBlock(torch.nn.Module): + """Decoder layer that returns a tuple, mimicking HuggingFace decoder layers.""" + + def __init__(self, dim=16): + super().__init__() + self.linear = torch.nn.Linear(dim, dim, bias=False) + + def forward(self, x, **kwargs): + return (self.linear(x), None) + + +class _TupleUnpackingModel(torch.nn.Module): + """Parent model that unpacks layer outputs as tuples.""" + + def __init__(self, n_layers=4, dim=16): + super().__init__() + self.layers = torch.nn.ModuleList([_TupleReturningBlock(dim) for _ in range(n_layers)]) + + def forward(self, x): + for layer in self.layers: + x, _ = layer(x) + return x + + +def test_skip_dummy_has_no_hf_hook(monkeypatch): + """Dummies must not carry _hf_hook from the original layer.""" + from accelerate.hooks import AlignDevicesHook, add_hook_to_module + + from modelopt.torch.quantization.utils.layerwise_calib import ( + LayerActivationCollector, + _SkipLayer, + ) + + monkeypatch.setattr( + LayerActivationCollector, + "_decoder_layer_support", + [(lambda m: hasattr(m, "layers"), lambda m: m.layers)], + ) + + model = _TupleUnpackingModel(n_layers=4, dim=16) + data = [torch.randn(2, 16)] + + for layer in model.layers: + hook = AlignDevicesHook(execution_device=torch.device("cpu")) + add_hook_to_module(layer, hook) + + def forward_loop(m): + for d in data: + m(d) + + collector = LayerActivationCollector(model) + collector._patch_all_layers() + try: + for layer in list(model.layers): + collector.get_input_activations(layer, forward_loop) + + for i in range(2): + dummy = model.layers[i] + assert isinstance(dummy, _SkipLayer) + assert not hasattr(dummy, "_hf_hook"), f"Dummy at {i} should not have _hf_hook" + finally: + collector._unpatch_all_layers() + + +def test_persistent_materialization_cpu_offloaded(tmp_path): + """persistent_materialization keeps CPU-offloaded weights on GPU and writes back modifications.""" + import torch.nn as nn + from accelerate.hooks import AlignDevicesHook + + from modelopt.torch.quantization.utils import persistent_materialization + + model, config, _, inputs = _make_cpu_offloaded_model(tmp_path) + offloaded_layer = model.model.layers[0] + + # Verify offloaded (meta device) + assert all(p.device.type == "meta" for p in offloaded_layer.parameters()) + + # Save reference weight + linear = None + with enable_weight_access_and_writeback(offloaded_layer, model): + linear = next(m for m in offloaded_layer.modules() if isinstance(m, nn.Linear)) + ref_weight = linear.weight.clone() + + with persistent_materialization(offloaded_layer): + # Params materialized on GPU + assert all( + p.device.type == "cuda" for p in offloaded_layer.parameters() if p.device.type != "meta" + ) + + # Run multiple forward passes (hooks don't re-offload) + for _ in range(3): + model(inputs) + + # Modify a weight + linear.weight.data.add_(1.0) + + # Verify hooks have offload=False during context + for mod in offloaded_layer.modules(): + if hasattr(mod, "_hf_hook"): + hook = mod._hf_hook + if isinstance(hook, AlignDevicesHook): + assert not hook.offload + + # After context: back to meta device (offloaded) + assert all(p.device.type == "meta" for p in offloaded_layer.parameters()) + + # Verify weight modification persisted through writeback + with enable_weight_access_and_writeback(offloaded_layer, model): + assert torch.allclose(linear.weight, ref_weight + 1.0) + + +def _make_disk_offload_device_map(model): + """Build a device_map with layer 0 on disk, everything else on GPU 0. + + Ancestor modules (``""`` and ``"model"``) are excluded so that + ``dispatch_model`` does not attach a ``place_submodules=True`` hook that + would try to move disk-offloaded meta tensors to GPU (which fails because + no ``value`` is available — unlike CPU offload where weights are on CPU and + can be moved directly). + """ + device_map = { + n: 0 + for n, m in model.named_modules() + if n not in ("", "model") and ("layers" not in n or n.split("layers.")[-1].isdigit()) + } + device_map["model.layers.0"] = "disk" + return device_map + + +def _make_disk_offloaded_model(tmp_path, num_hidden_layers=3): + """Create a tiny LLaMA model with layer 0 offloaded to disk via accelerate.""" + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_hidden_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + + device_map = _make_disk_offload_device_map(model) + offload_dir = str(tmp_path / "offload") + model = load_checkpoint_and_dispatch( + model, tiny_llama_dir, device_map=device_map, offload_folder=offload_dir + ) + inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda() + return model, config, tiny_llama_dir, inputs + + +def test_disk_offloaded_tinyllama(tmp_path): + quant_cfg = mtq.NVFP4_AWQ_LITE_CFG + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=2) + + config = AutoConfig.from_pretrained(tiny_llama_dir) + + model_ref = AutoModelForCausalLM.from_pretrained( + tiny_llama_dir, torch_dtype=config.torch_dtype + ).cuda() + inputs = torch.randint(0, model_ref.config.vocab_size, (1, 4)).cuda() + + mtq.quantize(model_ref, quant_cfg, lambda model: model(inputs)) + output_ref = model_ref(inputs) + + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + + device_map = _make_disk_offload_device_map(model) + offload_dir = str(tmp_path / "offload") + model = load_checkpoint_and_dispatch( + model, tiny_llama_dir, device_map=device_map, offload_folder=offload_dir + ) + + assert all(p.device == torch.device("meta") for p in model.model.layers[0].parameters()) + + mtq.quantize(model, quant_cfg, lambda model: model(inputs)) + output_test = model(inputs) + + for name, module in model.named_modules(): + if is_quantized_linear(module): + with enable_weight_access_and_writeback(module, model): + assert torch.allclose(module.weight, model_ref.get_submodule(name).weight) + + assert torch.allclose(output_ref.logits, output_test.logits) + + +def test_persistent_materialization_disk_offloaded(tmp_path): + """persistent_materialization keeps disk-offloaded weights on GPU and writes back modifications.""" + import torch.nn as nn + from accelerate.hooks import AlignDevicesHook + + from modelopt.torch.quantization.utils import persistent_materialization + + model, config, _, inputs = _make_disk_offloaded_model(tmp_path) + offloaded_layer = model.model.layers[0] + + # Verify offloaded (meta device) + assert all(p.device.type == "meta" for p in offloaded_layer.parameters()) + + # Save reference weight + linear = None + with enable_weight_access_and_writeback(offloaded_layer, model): + linear = next(m for m in offloaded_layer.modules() if isinstance(m, nn.Linear)) + ref_weight = linear.weight.clone() + + with persistent_materialization(offloaded_layer): + # Params materialized on GPU + assert all( + p.device.type == "cuda" for p in offloaded_layer.parameters() if p.device.type != "meta" + ) + + # Run multiple forward passes (hooks don't re-offload) + for _ in range(3): + model(inputs) + + # Modify a weight + linear.weight.data.add_(1.0) + + # Verify hooks have offload=False during context + for mod in offloaded_layer.modules(): + if hasattr(mod, "_hf_hook"): + hook = mod._hf_hook + if isinstance(hook, AlignDevicesHook): + assert not hook.offload + + # After context: back to meta device (offloaded) + assert all(p.device.type == "meta" for p in offloaded_layer.parameters()) + + # Verify weight modification persisted through writeback + with enable_weight_access_and_writeback(offloaded_layer, model): + assert torch.allclose(linear.weight, ref_weight + 1.0) + + +@pytest.mark.parametrize("use_checkpoint", [False, True], ids=["no_ckpt", "ckpt"]) +def test_layerwise_calibrate_disk_offloaded(tmp_path, use_checkpoint): + """Layerwise calibration on disk-offloaded model matches GPU-only reference.""" + quant_cfg = mtq.NVFP4_AWQ_LITE_CFG + num_layers = 3 + tiny_llama_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=num_layers) + config = AutoConfig.from_pretrained(tiny_llama_dir) + inputs = torch.randint(0, config.vocab_size, (1, 4)).cuda() + + if use_checkpoint: + ckpt_dir = str(tmp_path / "seq_ckpt") + seq_cfg = _make_layerwise_checkpoint_cfg(quant_cfg, ckpt_dir) + else: + seq_cfg = _make_layerwise_cfg(quant_cfg) + + # Reference: GPU-only model with layerwise calibration + ref_cfg = _make_layerwise_cfg(quant_cfg) + model_ref = AutoModelForCausalLM.from_pretrained( + tiny_llama_dir, torch_dtype=config.torch_dtype + ).cuda() + mtq.quantize(model_ref, ref_cfg, lambda model: model(inputs)) + output_ref = model_ref(inputs) + + # Test: disk-offloaded model + with init_empty_weights(): + model = AutoModelForCausalLM.from_config(config) + device_map = _make_disk_offload_device_map(model) + offload_dir = str(tmp_path / "offload") + model = load_checkpoint_and_dispatch( + model, tiny_llama_dir, device_map=device_map, offload_folder=offload_dir + ) + + mtq.quantize(model, seq_cfg, lambda model: model(inputs)) + output_test = model(inputs) + + for name, module in model.named_modules(): + if is_quantized_linear(module): + with enable_weight_access_and_writeback(module, model): + assert torch.allclose(module.weight, model_ref.get_submodule(name).weight), ( + f"Weight mismatch at {name}" + ) + + assert torch.allclose(output_ref.logits, output_test.logits) diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index 4889b6dc8cc..c5584ece5cf 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -128,3 +128,136 @@ def test_fsdp_simple_linear(dist_workers): ) def test_nested_fsdp2_backward(quant_cfg, dist_workers): dist_workers.run(partial(_test_nested_fsdp2_backward, quant_cfg=quant_cfg)) + + +class _DecoderBlock(nn.Module): + """Minimal decoder block for FSDP2 sequential tests.""" + + def __init__(self, dim=32): + super().__init__() + self.attn = nn.Linear(dim, dim, bias=False) + self.ffn = nn.Sequential( + nn.Linear(dim, dim, bias=False), nn.ReLU(), nn.Linear(dim, dim, bias=False) + ) + self.norm = nn.LayerNorm(dim) + + def forward(self, x): + x = x + self.attn(self.norm(x)) + x = x + self.ffn(x) + return x + + +class _SimpleTransformerModel(nn.Module): + """Model with ``model.layers`` for layerwise calibration discovery.""" + + def __init__(self, n_layers=3, dim=32): + super().__init__() + self.layers = nn.ModuleList([_DecoderBlock(dim) for _ in range(n_layers)]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +def _test_layerwise_calibrate_fsdp2(rank, size): + """Layerwise calibration on FSDP2-wrapped model matches non-FSDP reference.""" + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector + + dim = 32 + torch.manual_seed(1) + model = _SimpleTransformerModel(n_layers=3, dim=dim).cuda() + inputs = torch.randn(2, 2, dim).cuda() + synchronize_state_dict(model) + + # Register discoverer for our simple model + old_support = LayerActivationCollector._decoder_layer_support[:] + LayerActivationCollector._decoder_layer_support = [ + ( + lambda m: hasattr(m, "layers") and isinstance(m.layers, nn.ModuleList), + lambda m: m.layers, + ), + *old_support, + ] + + try: + # Reference: non-FSDP layerwise calibration + ref_model = copy.deepcopy(model) + seq_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) + seq_cfg["algorithm"] = {"method": "max", "layerwise": True} + mtq.quantize(ref_model, seq_cfg, lambda m: m(inputs)) + output_ref = ref_model(inputs) + + # Test: FSDP2-wrapped layerwise calibration + for layer in model.layers: + fully_shard(layer) + model = fully_shard(model) + mtq.quantize(model, seq_cfg, lambda m: m(inputs)) + output_test = model(inputs) + + assert torch.allclose(output_ref, output_test) + finally: + LayerActivationCollector._decoder_layer_support = old_support + + +def test_layerwise_calibrate_fsdp2(dist_workers): + dist_workers.run(_test_layerwise_calibrate_fsdp2) + + +def _test_persistent_materialization(rank, size): + """persistent_materialization keeps weights accessible and writes back modifications.""" + from torch.distributed.tensor import DTensor + + from modelopt.torch.quantization.utils import ( + enable_weight_access_and_writeback, + persistent_materialization, + ) + + dim = 32 + torch.manual_seed(1) + model = nn.Sequential( + nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, dim)), + nn.Sequential(nn.Linear(dim, dim), nn.Linear(dim, dim)), + ).cuda(rank) + synchronize_state_dict(model) + + fully_shard(model[0]) + fully_shard(model[1]) + model = fully_shard(model) + + layer = model[0] + inputs = torch.randn(2, dim).cuda(rank) + + # Warmup forward to trigger FSDP2's lazy_init (mirrors real usage where + # layerwise_calibrate always runs get_first_layer_inputs first). + model(inputs) + + # Save reference weight (gathered) + with enable_weight_access_and_writeback(layer[0], model): + ref_weight = layer[0].weight.clone() + + # Verify sharded before context + assert isinstance(next(iter(layer.parameters())), DTensor) + + with persistent_materialization(layer): + # Params are local tensors (not DTensors) + assert not isinstance(layer[0].weight, DTensor) + assert layer[0].weight.device.type == "cuda" + + # Run multiple forward passes (FSDP hooks fire, unshard/reshard are no-ops) + for _ in range(3): + layer(inputs) + + # Modify a weight + layer[0].weight.data.add_(1.0) + + # After context: params restored to DTensors (sharded) + assert isinstance(next(iter(layer.parameters())), DTensor) + + # Verify modification persisted + with enable_weight_access_and_writeback(layer[0], model): + assert torch.allclose(layer[0].weight, ref_weight + 1.0) + + +def test_persistent_materialization(dist_workers): + dist_workers.run(_test_persistent_materialization) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index d183855abbd..2d5f9d6d707 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -219,7 +219,7 @@ def test_gptq_e2e_flow(quant_cfg): model.eval() quant_cfg = copy.deepcopy(quant_cfg) - quant_cfg["algorithm"] = {"method": "gptq", "use_sequential": True} + quant_cfg["algorithm"] = {"method": "gptq", "layerwise": True} calib_dataloader = get_dataset_dataloader( dataset_name="cnn_dailymail", tokenizer=tokenizer, diff --git a/tests/gpu/torch/quantization/test_sequential_calibrate.py b/tests/gpu/torch/quantization/test_layerwise_calibrate.py similarity index 90% rename from tests/gpu/torch/quantization/test_sequential_calibrate.py rename to tests/gpu/torch/quantization/test_layerwise_calibrate.py index ba71e896c72..d38b82f46fb 100644 --- a/tests/gpu/torch/quantization/test_sequential_calibrate.py +++ b/tests/gpu/torch/quantization/test_layerwise_calibrate.py @@ -13,13 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Integration tests for sequential_calibrate and LayerActivationCollector.""" +"""Integration tests for layerwise_calibrate and LayerActivationCollector.""" import torch import torch.nn as nn -from modelopt.torch.quantization.model_calib import sequential_calibrate -from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector +from modelopt.torch.quantization.model_calib import layerwise_calibrate +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector class _DecoderBlock(nn.Module): @@ -101,7 +101,7 @@ def _register_test_discoverer(monkeypatch): ) -def test_seq_calib_func_called_per_layer(monkeypatch): +def test_layerwise_calib_func_called_per_layer(monkeypatch): _register_test_discoverer(monkeypatch) model, data = _make_model_and_data(n_layers=4) call_count = [0] @@ -109,7 +109,7 @@ def test_seq_calib_func_called_per_layer(monkeypatch): def counting_calib(layer, forward_loop, **kwargs): call_count[0] += 1 - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: _run_forward(m, data), calib_func=counting_calib, @@ -118,7 +118,7 @@ def counting_calib(layer, forward_loop, **kwargs): assert call_count[0] == 4 -def test_seq_calib_func_receives_correct_layer(monkeypatch): +def test_layerwise_calib_func_receives_correct_layer(monkeypatch): _register_test_discoverer(monkeypatch) model, data = _make_model_and_data(n_layers=3) called_layers = [] @@ -126,7 +126,7 @@ def test_seq_calib_func_receives_correct_layer(monkeypatch): def track_layers(layer, forward_loop, **kwargs): called_layers.append(layer) - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: _run_forward(m, data), calib_func=track_layers, @@ -136,7 +136,7 @@ def track_layers(layer, forward_loop, **kwargs): assert called_layers[i] is layer -def test_seq_calib_kwargs_forwarded(monkeypatch): +def test_layerwise_calib_kwargs_forwarded(monkeypatch): _register_test_discoverer(monkeypatch) model, data = _make_model_and_data(n_layers=2) received_kwargs = [] @@ -144,7 +144,7 @@ def test_seq_calib_kwargs_forwarded(monkeypatch): def capture_kwargs(layer, forward_loop, **kwargs): received_kwargs.append(kwargs) - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: _run_forward(m, data), calib_func=capture_kwargs, @@ -158,7 +158,7 @@ def capture_kwargs(layer, forward_loop, **kwargs): assert kw["method"] == "max" -def test_seq_calib_layer_forward_loop_runs_all_batches(monkeypatch): +def test_layerwise_calib_layer_forward_loop_runs_all_batches(monkeypatch): """The per-layer forward loop passed to calib_func should replay all batches.""" _register_test_discoverer(monkeypatch) n_batches = 5 @@ -178,7 +178,7 @@ def counting_forward(*args, **kw): layer.forward = orig_forward batch_counts.append(counter["n"]) - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: _run_forward(m, data), calib_func=count_batches, @@ -188,13 +188,13 @@ def counting_forward(*args, **kw): assert count == n_batches -def test_seq_calib_does_not_alter_weights(monkeypatch): - """sequential_calibrate itself should not modify model weights.""" +def test_layerwise_calib_does_not_alter_weights(monkeypatch): + """layerwise_calibrate itself should not modify model weights.""" _register_test_discoverer(monkeypatch) model, data = _make_model_and_data(n_layers=3) weights_before = {n: p.clone() for n, p in model.named_parameters()} - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: _run_forward(m, data), calib_func=lambda layer, forward_loop, **kw: None, @@ -204,7 +204,7 @@ def test_seq_calib_does_not_alter_weights(monkeypatch): assert torch.equal(p, weights_before[n]), f"Weight {n} was modified" -def test_seq_calib_activations_update_across_layers(monkeypatch): +def test_layerwise_calib_activations_update_across_layers(monkeypatch): """Subsequent layers should see activations transformed by prior layers.""" _register_test_discoverer(monkeypatch) torch.manual_seed(0) @@ -228,7 +228,7 @@ def capture_forward(*args, **kw): layer_idx = list(model.layers).index(layer) layer_inputs_record[layer_idx] = activations - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: [m(t) for t in tokens], calib_func=record_inputs, @@ -240,7 +240,7 @@ def capture_forward(*args, **kw): def test_mode_transitions_across_calibration_steps(monkeypatch): - """Verify layer modes after each sequential calibration step. + """Verify layer modes after each layerwise calibration step. After get_input_activations(layers[i]) returns, the current layer is reset to 'original'. Layers further back are left in 'run' (just calibrated) or @@ -259,7 +259,7 @@ def forward_loop(m): try: def modes(): - return [model.layers[i]._seq_calib.mode for i in range(5)] + return [model.layers[i]._layerwise_calib.mode for i in range(5)] collector.get_input_activations(model.layers[0], forward_loop) assert modes() == ["original", "original", "original", "original", "original"] @@ -316,7 +316,7 @@ def weight_doubling_calib(layer, layer_forward_loop, **kwargs): layer.weight.mul_(2.0) layer_forward_loop(layer) - sequential_calibrate( + layerwise_calibrate( model, forward_loop=forward_loop, calib_func=weight_doubling_calib, diff --git a/tests/unit/torch/quantization/plugins/test_huggingface.py b/tests/unit/torch/quantization/plugins/test_huggingface.py index 692ab07d4aa..ae638c42ee2 100644 --- a/tests/unit/torch/quantization/plugins/test_huggingface.py +++ b/tests/unit/torch/quantization/plugins/test_huggingface.py @@ -35,7 +35,7 @@ get_homogeneous_hf_decoder_layers, is_homogeneous_hf_model, ) -from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector pytest.importorskip("transformers") diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index b3c372eb335..d2e6fdd03e8 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -27,7 +27,7 @@ from modelopt.torch.quantization.model_calib import ( apply_pre_quant_scale_and_smooth, disable_pre_quant_scale_and_resmooth, - sequential_calibrate, + layerwise_calibrate, ) from modelopt.torch.quantization.nn import TensorQuantizer @@ -379,7 +379,7 @@ def test_svdquant_lora_weights(): assert lora_residual.shape == module.weight.shape -def test_sequential_calibrate_support_gate(): +def test_layerwise_calibrate_support_gate(): class _UnsupportedModel(nn.Module): def __init__(self): super().__init__() @@ -392,17 +392,17 @@ def forward(self, x): with ( torch.no_grad(), - pytest.raises(ValueError, match="Sequential calibration requires a model"), + pytest.raises(ValueError, match="Layerwise calibration requires a model"), ): - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: m(torch.randn(2, 4)), calib_func=lambda layer, loop: loop(layer), ) -def test_sequential_calibrate_propagates_inputs_without_replaying_full_model(monkeypatch): - from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector +def test_layerwise_calibrate_propagates_inputs_without_replaying_full_model(monkeypatch): + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector class _ToyLayer(nn.Module): def __init__(self, scale: float, bias: float): @@ -463,7 +463,7 @@ def _pre_hook(_module, args): handle.remove() observed_layer_inputs.append(captured) - sequential_calibrate(model, _forward_loop, _calib_func) + layerwise_calibrate(model, _forward_loop, _calib_func) assert forward_loop_calls == len(model.layers) assert len(observed_layer_inputs) == len(model.layers) @@ -482,9 +482,9 @@ def _pre_hook(_module, args): assert torch.allclose(observed, expected) -def test_sequential_calibrate_handles_inter_layer_logic(monkeypatch): +def test_layerwise_calibrate_handles_inter_layer_logic(monkeypatch): """Verify that parent-level inter-layer logic (e.g. mask selection) works correctly.""" - from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector + from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector class _ToyLayer(nn.Module): def __init__(self, scale: float): @@ -537,7 +537,7 @@ def _pre_hook(_module, args): handle.remove() observed_layer_inputs.append(captured) - sequential_calibrate(model, _forward_loop, _calib_func) + layerwise_calibrate(model, _forward_loop, _calib_func) assert len(observed_layer_inputs) == 3 # Layer 0 gets raw batch diff --git a/tests/unit/torch/quantization/test_sequential_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py similarity index 64% rename from tests/unit/torch/quantization/test_sequential_calibrate.py rename to tests/unit/torch/quantization/test_layerwise_calibrate.py index 14c1903de2e..3739feff969 100644 --- a/tests/unit/torch/quantization/test_sequential_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -13,16 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Unit tests for sequential_calibrate and LayerActivationCollector.""" +"""Unit tests for layerwise_calibrate and LayerActivationCollector.""" +import copy from collections import deque import pytest import torch import torch.nn as nn -from modelopt.torch.quantization.model_calib import sequential_calibrate -from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.model_calib import layerwise_calibrate +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector, _SkipLayer class _DecoderBlock(nn.Module): @@ -60,7 +63,7 @@ def forward(self, x, **kwargs): class _FlatMLP(nn.Module): - """No decoder-layer structure -- should be rejected by sequential_calibrate.""" + """No decoder-layer structure -- should be rejected by layerwise_calibrate.""" def __init__(self, dim=16): super().__init__() @@ -180,7 +183,7 @@ def forward_loop(m): collector._unpatch_all_layers() assert not hasattr(model, "_original_forward") - assert not hasattr(model.layers[0], "_seq_calib") + assert not hasattr(model.layers[0], "_layerwise_calib") assert not hasattr(model.layers[0], "_original_forward") @@ -201,38 +204,38 @@ def bad_forward_loop(m): collector._unpatch_all_layers() assert not hasattr(model, "_original_forward") - assert not hasattr(model.layers[0], "_seq_calib") + assert not hasattr(model.layers[0], "_layerwise_calib") -# sequential_calibrate tests -def test_seq_calib_raises_on_none_forward_loop(monkeypatch): +# layerwise_calibrate tests +def test_layerwise_calib_raises_on_none_forward_loop(monkeypatch): _register_test_discoverer(monkeypatch) model, data = _make_model_and_data(n_layers=2) with pytest.raises(ValueError, match="forward_loop must not be None"): - sequential_calibrate( + layerwise_calibrate( model, forward_loop=None, calib_func=lambda *a, **kw: None, ) -def test_seq_calib_raises_on_unrecognized_model(): +def test_layerwise_calib_raises_on_unrecognized_model(): model = _FlatMLP() with pytest.raises(ValueError, match="Could not find transformer layers"): - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: m(torch.randn(2, 16)), calib_func=lambda *a, **kw: None, ) -def test_seq_calib_empty_forward_loop_raises(monkeypatch): - """If forward_loop feeds no data, sequential_calibrate raises RuntimeError.""" +def test_layerwise_calib_empty_forward_loop_raises(monkeypatch): + """If forward_loop feeds no data, layerwise_calibrate raises RuntimeError.""" _register_test_discoverer(monkeypatch) model = _SimpleTransformerModel(n_layers=2, dim=16) with pytest.raises(RuntimeError, match="collected no inputs during forward_loop"): - sequential_calibrate( + layerwise_calibrate( model, forward_loop=lambda m: None, calib_func=lambda *a, **kw: None, @@ -344,11 +347,11 @@ def forward_loop(m): try: # Layer 0 starts as capture — no output_meta yet collector.get_input_activations(model.layers[0], forward_loop) - assert model.layers[0]._seq_calib.output_meta is None + assert model.layers[0]._layerwise_calib.output_meta is None # Calibrating layer 1 puts layer 0 into run, which sets output_meta collector.get_input_activations(model.layers[1], forward_loop) - meta = model.layers[0]._seq_calib.output_meta + meta = model.layers[0]._layerwise_calib.output_meta assert meta is not None assert meta[0] == "tuple", "Tuple-returning layer should produce tuple metadata" finally: @@ -375,11 +378,11 @@ def forward_loop(m): # Before calibrating layer 2, layer 1 transitions to run. # Its cached_inputs should be populated from collected_inputs. collector._set_layer_states(2) - assert len(model.layers[1]._seq_calib.cached_inputs) == n_batches + assert len(model.layers[1]._layerwise_calib.cached_inputs) == n_batches # After the forward loop, all cached inputs should be consumed forward_loop(model) - assert len(model.layers[1]._seq_calib.cached_inputs) == 0 + assert len(model.layers[1]._layerwise_calib.cached_inputs) == 0 finally: collector._unpatch_all_layers() @@ -399,24 +402,24 @@ def test_set_layer_states_transitions(monkeypatch): try: def modes(): - return [model.layers[i]._seq_calib.mode for i in range(5)] + return [model.layers[i]._layerwise_calib.mode for i in range(5)] collector._set_layer_states(0) assert modes() == ["capture", "original", "original", "original", "original"] - model.layers[0]._seq_calib.collected_inputs = [fake_inp] + model.layers[0]._layerwise_calib.collected_inputs = [fake_inp] collector._set_layer_states(1) assert modes() == ["run", "capture", "original", "original", "original"] - model.layers[1]._seq_calib.collected_inputs = [fake_inp] + model.layers[1]._layerwise_calib.collected_inputs = [fake_inp] collector._set_layer_states(2) assert modes() == ["skip", "run", "capture", "original", "original"] - model.layers[2]._seq_calib.collected_inputs = [fake_inp] + model.layers[2]._layerwise_calib.collected_inputs = [fake_inp] collector._set_layer_states(3) assert modes() == ["skip", "skip", "run", "capture", "original"] - model.layers[3]._seq_calib.collected_inputs = [fake_inp] + model.layers[3]._layerwise_calib.collected_inputs = [fake_inp] collector._set_layer_states(4) assert modes() == ["skip", "skip", "skip", "run", "capture"] finally: @@ -446,8 +449,8 @@ def test_run_asserts_on_empty_cached_inputs(monkeypatch): collector = LayerActivationCollector(model) collector._patch_all_layers() try: - model.layers[0]._seq_calib.mode = "run" - model.layers[0]._seq_calib.cached_inputs = deque() + model.layers[0]._layerwise_calib.mode = "run" + model.layers[0]._layerwise_calib.cached_inputs = deque() with pytest.raises(AssertionError, match="no cached inputs to replay"): model(torch.randn(2, 16)) @@ -455,8 +458,8 @@ def test_run_asserts_on_empty_cached_inputs(monkeypatch): collector._unpatch_all_layers() -def test_cleanup_removes_seq_calib_attr(monkeypatch): - """After unpatch, no layer should have the _seq_calib attribute.""" +def test_cleanup_removes_layerwise_calib_attr(monkeypatch): + """After unpatch, no layer should have the _layerwise_calib attribute.""" _register_test_discoverer(monkeypatch) model = _TupleUnpackingModel(n_layers=3, dim=16) data = [torch.randn(2, 16)] @@ -472,7 +475,9 @@ def forward_loop(m): collector._unpatch_all_layers() for i, layer in enumerate(model.layers): - assert not hasattr(layer, "_seq_calib"), f"Layer {i} still has _seq_calib after cleanup" + assert not hasattr(layer, "_layerwise_calib"), ( + f"Layer {i} still has _layerwise_calib after cleanup" + ) assert not hasattr(layer, "_original_forward"), ( f"Layer {i} still has _original_forward after cleanup" ) @@ -517,15 +522,17 @@ def forward_loop(m): for d in data: m(d) + originals = list(model.layers) collector = LayerActivationCollector(model) collector._patch_all_layers() try: - for layer in model.layers: + for layer in originals: collector.get_input_activations(layer, forward_loop) - # After full calibration, layers 0 and 1 have been through 'run' and have output_meta - meta_0 = model.layers[0]._seq_calib.output_meta - meta_1 = model.layers[1]._seq_calib.output_meta + # After full calibration, layers 0 and 1 have been through 'run' and have output_meta. + # Access via originals since skip-position entries are now _SkipLayer dummies. + meta_0 = originals[0]._layerwise_calib.output_meta + meta_1 = originals[1]._layerwise_calib.output_meta assert meta_0 is not None assert meta_1 is not None # SmallBlock returns 3-element tuple, BigBlock returns 1-element tuple @@ -533,3 +540,182 @@ def forward_loop(m): assert len(meta_1[1]) == 1 finally: collector._unpatch_all_layers() + + +# --------------------------------------------------------------------------- +# _SkipLayer swap / restore tests +# --------------------------------------------------------------------------- + + +def test_skip_layers_replaced_with_dummy(monkeypatch): + """After calibrating enough layers, skip-position entries must be _SkipLayer with no params.""" + _register_test_discoverer(monkeypatch) + model = _TupleUnpackingModel(n_layers=5, dim=16) + data = [torch.randn(2, 16) for _ in range(2)] + + def forward_loop(m): + for d in data: + m(d) + + collector = LayerActivationCollector(model) + collector._patch_all_layers() + try: + for layer in list(model.layers): + collector.get_input_activations(layer, forward_loop) + + # Layers 0..2 should be dummies (swapped when calibrating layers 2..4) + for i in range(3): + assert isinstance(model.layers[i], _SkipLayer), f"Layer {i} should be _SkipLayer" + assert list(model.layers[i].parameters()) == [], ( + f"Layer {i} dummy should have no params" + ) + # Layers 3 (run) and 4 (original) remain real + for i in range(3, 5): + assert not isinstance(model.layers[i], _SkipLayer), f"Layer {i} should still be real" + finally: + collector._unpatch_all_layers() + + +def test_cleanup_restores_original_layers(monkeypatch): + """After _unpatch_all_layers, all ModuleList entries must be the original modules.""" + _register_test_discoverer(monkeypatch) + model = _TupleUnpackingModel(n_layers=5, dim=16) + originals = list(model.layers) + data = [torch.randn(2, 16)] + + def forward_loop(m): + for d in data: + m(d) + + collector = LayerActivationCollector(model) + collector._patch_all_layers() + for layer in originals: + collector.get_input_activations(layer, forward_loop) + collector._unpatch_all_layers() + + for i, orig in enumerate(originals): + assert model.layers[i] is orig, f"Layer {i} not restored to original after cleanup" + assert not hasattr(orig, "_layerwise_calib"), f"Layer {i} still has _layerwise_calib" + + +def _int8_layerwise_config(algorithm: dict) -> dict: + """Start from the shipped INT8 config and enable layerwise in the algorithm block. + + Using a real shipped config guarantees the same include/exclude rules + production PTQ relies on, so algorithm dispatch matches real usage. + """ + cfg = copy.deepcopy(mtq.INT8_SMOOTHQUANT_CFG) + cfg["algorithm"] = algorithm + return cfg + + +def _awq_layerwise_config() -> dict: + """INT4 weight-only AWQ config sized for the _DecoderBlock test model.""" + cfg = copy.deepcopy(mtq.INT4_AWQ_CFG) + # Resize AWQ block to fit dim=16 hidden. + for entry in cfg["quant_cfg"]: + if entry.get("quantizer_name") == "*weight_quantizer": + entry.setdefault("cfg", {})["block_sizes"] = {-1: 8, "type": "static"} + cfg["algorithm"] = {"method": "awq_lite", "alpha_step": 0.5, "layerwise": True} + return cfg + + +def _svdquant_layerwise_config() -> dict: + """SVDQuant config sized for the _DecoderBlock test model.""" + cfg = copy.deepcopy(mtq.INT4_AWQ_CFG) + for entry in cfg["quant_cfg"]: + if entry.get("quantizer_name") == "*weight_quantizer": + entry.setdefault("cfg", {})["block_sizes"] = {-1: 8, "type": "static"} + cfg["algorithm"] = {"method": "svdquant", "lowrank": 4, "layerwise": True} + return cfg + + +def test_mtq_quantize_layerwise_e2e_max(monkeypatch): + """End-to-end: mtq.quantize with layerwise=True produces populated amax values. + + ``max`` is the representative algorithm for the layerwise happy path because + every other algorithm seeds amax via max_calibrate first — if max works, the + shared skip/run/capture machinery is sound. Other algorithms are covered by + the dispatch-only test below to avoid hardware requirements (e.g. gptq needs + CUDA) or unnecessary duplication. + """ + _register_test_discoverer(monkeypatch) + config = _int8_layerwise_config({"method": "max", "layerwise": True}) + + torch.manual_seed(0) + model = _SimpleTransformerModel(n_layers=3, dim=16) + calib_data = [torch.randint(0, 32, (2, 8)) for _ in range(2)] + + def forward_loop(m): + for batch in calib_data: + m(batch) + + model = mtq.quantize(model, config, forward_loop=forward_loop) + + for i, layer in enumerate(model.layers): + assert not isinstance(layer, _SkipLayer), f"layer {i} left as _SkipLayer" + assert not hasattr(layer, "_layerwise_calib"), f"layer {i} leaked _layerwise_calib" + + amax_count = sum( + 1 + for layer in model.layers + for module in layer.modules() + if ( + isinstance(module, TensorQuantizer) + and module.is_enabled + and getattr(module, "_amax", None) is not None + ) + ) + assert amax_count > 0, "no TensorQuantizer in decoder layers had _amax populated" + + with torch.no_grad(): + model(calib_data[0]) + + +@pytest.mark.parametrize( + "algorithm", + ["gptq", "awq_lite", "smoothquant", "mse"], +) +def test_mtq_quantize_layerwise_dispatches_for_algorithm(monkeypatch, algorithm): + """Every layerwise-supporting algorithm must route through layerwise_calibrate. + + Stubs layerwise_calibrate to a spy so the dispatch contract is checked without + running the algorithm's full calibration — lets ``gptq`` (CUDA-only at runtime) + and other expensive algorithms participate in CPU unit tests. + """ + spy: dict = {} + + def stub(model, forward_loop, calib_func, **kwargs): + spy["calib_func"] = calib_func + spy["kwargs"] = kwargs + + monkeypatch.setattr("modelopt.torch.quantization.mode.layerwise_calibrate", stub) + + if algorithm == "awq_lite": + config = _awq_layerwise_config() + else: + config = _int8_layerwise_config({"method": algorithm, "layerwise": True}) + + torch.manual_seed(0) + model = _SimpleTransformerModel(n_layers=2, dim=16) + mtq.quantize( + model, + config, + forward_loop=lambda m: m(torch.randint(0, 32, (2, 8))), + ) + + assert "calib_func" in spy, f"{algorithm} did not dispatch through layerwise_calibrate" + assert callable(spy["calib_func"]) + + +def test_mtq_quantize_layerwise_raises_for_unsupported_algorithm(): + """Modes with ``_supports_layerwise = False`` must raise a clear ValueError.""" + config = _svdquant_layerwise_config() + torch.manual_seed(0) + model = _SimpleTransformerModel(n_layers=2, dim=16) + with pytest.raises(ValueError, match="does not support layerwise=True"): + mtq.quantize( + model, + config, + forward_loop=lambda m: m(torch.randint(0, 32, (2, 8))), + ) diff --git a/tests/unit/torch/quantization/test_sequential_checkpoint.py b/tests/unit/torch/quantization/test_sequential_checkpoint.py new file mode 100644 index 00000000000..0e592a68c75 --- /dev/null +++ b/tests/unit/torch/quantization/test_sequential_checkpoint.py @@ -0,0 +1,185 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for layerwise calibration checkpoint save/resume.""" + +import json +import os +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from modelopt.torch.quantization.model_calib import layerwise_calibrate +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils.network import get_module_device + + +class _DecoderBlock(nn.Module): + def __init__(self, dim=16): + super().__init__() + self.linear = nn.Linear(dim, dim, bias=False) + + def forward(self, x, **kwargs): + return self.linear(x) + + +class _SimpleTransformerModel(nn.Module): + def __init__(self, n_layers=3, dim=16): + super().__init__() + self.layers = nn.ModuleList([_DecoderBlock(dim) for _ in range(n_layers)]) + self.embed = nn.Embedding(32, dim) + + def forward(self, x, **kwargs): + x = self.embed(x) + for layer in self.layers: + x = layer(x) + return x + + +def _register_test_discoverer(monkeypatch): + monkeypatch.setattr( + LayerActivationCollector, + "_decoder_layer_support", + [(lambda m: hasattr(m, "layers"), lambda m: m.layers)], + ) + + +def _dummy_calib_func(layer, forward_loop, **kwargs): + """Scale all weights by 0.5 to produce a visible, deterministic change.""" + forward_loop(layer) + with torch.no_grad(): + for p in layer.parameters(): + p.mul_(0.5) + + +def _make_model_and_forward(n_layers=3, dim=16, seed=42): + torch.manual_seed(seed) + model = _SimpleTransformerModel(n_layers=n_layers, dim=dim) + tokens = [torch.randint(0, 32, (2, 8)) for _ in range(2)] + + def forward_loop(m): + for t in tokens: + m(t) + + return model, forward_loop + + +def test_full_run_creates_checkpoints(monkeypatch, tmp_path): + """layerwise_calibrate with checkpoint_dir creates correct layer dirs and manifest.""" + _register_test_discoverer(monkeypatch) + model, forward_loop = _make_model_and_forward(n_layers=3) + ckpt_dir = str(tmp_path / "ckpt") + + layerwise_calibrate(model, forward_loop, _dummy_calib_func, checkpoint_dir=ckpt_dir) + + manifest_path = os.path.join(ckpt_dir, "manifest.json") + assert os.path.isfile(manifest_path) + with open(manifest_path) as f: + manifest = json.load(f) + assert manifest["last_completed_layer"] == 2 + assert manifest["num_layers"] == 3 + + for i in range(3): + layer_dir = os.path.join(ckpt_dir, f"layer_{i:04d}") + assert os.path.isdir(layer_dir) + assert os.path.isfile(os.path.join(layer_dir, "weights.pt")) + assert os.path.isfile(os.path.join(layer_dir, "quantizer_state.pt")) + assert os.path.isfile(os.path.join(layer_dir, "output_meta.pt")) + # All layers except the last should have next_inputs + assert os.path.isfile(os.path.join(ckpt_dir, "layer_0000", "next_inputs.pt")) + assert os.path.isfile(os.path.join(ckpt_dir, "layer_0001", "next_inputs.pt")) + assert not os.path.isfile(os.path.join(ckpt_dir, "layer_0002", "next_inputs.pt")) + + +def test_resume_matches_full_run(monkeypatch, tmp_path): + """Resume from a truncated checkpoint produces the same final weights as a full run.""" + _register_test_discoverer(monkeypatch) + ckpt_dir = str(tmp_path / "ckpt") + + # Full reference run + ref_model, forward_loop = _make_model_and_forward(n_layers=3) + layerwise_calibrate(ref_model, forward_loop, _dummy_calib_func, checkpoint_dir=ckpt_dir) + ref_weights = {n: p.clone() for n, p in ref_model.named_parameters()} + + # Simulate crash after layer 0: truncate manifest + manifest_path = os.path.join(ckpt_dir, "manifest.json") + with open(manifest_path, "w") as f: + json.dump({"last_completed_layer": 0, "num_layers": 3}, f) + + # Resume from a fresh model + resumed_model, forward_loop = _make_model_and_forward(n_layers=3) + layerwise_calibrate(resumed_model, forward_loop, _dummy_calib_func, checkpoint_dir=ckpt_dir) + + for name, ref_param in ref_weights.items(): + resumed_param = dict(resumed_model.named_parameters())[name] + assert torch.allclose(ref_param, resumed_param, atol=1e-6), ( + f"Parameter {name} diverged after resume" + ) + + +def test_no_checkpoint_unchanged(monkeypatch): + """Without checkpoint_dir, calibration still works and modifies parameters.""" + _register_test_discoverer(monkeypatch) + model, forward_loop = _make_model_and_forward(n_layers=3) + original_weights = {n: p.clone() for n, p in model.named_parameters()} + + layerwise_calibrate(model, forward_loop, _dummy_calib_func) + + changed = False + for name, param in model.named_parameters(): + if not torch.allclose(original_weights[name], param): + changed = True + break + assert changed, "Expected calibration to modify at least one parameter" + + +# --------------------------------------------------------------------------- +# get_module_device tests +# --------------------------------------------------------------------------- + + +def test_get_module_device_no_hook(): + """Falls back to parameter device when no _hf_hook is present.""" + layer = nn.Linear(4, 4) + assert get_module_device(layer) == torch.device("cpu") + + +def test_get_module_device_with_direct_hook(): + """Returns execution_device from a direct AlignDevicesHook-style hook.""" + layer = nn.Linear(4, 4) + layer._hf_hook = SimpleNamespace(execution_device=torch.device("cuda:0")) + assert get_module_device(layer) == torch.device("cuda:0") + + +def test_get_module_device_with_sequential_hook(): + """Returns execution_device from an AlignDevicesHook wrapped in SequentialHook.""" + layer = nn.Linear(4, 4) + inner_hook = SimpleNamespace(execution_device=torch.device("cuda:1")) + layer._hf_hook = SimpleNamespace(hooks=[inner_hook]) + assert get_module_device(layer) == torch.device("cuda:1") + + +def test_get_module_device_hook_without_execution_device(): + """Falls back to parameters when hook has no execution_device.""" + layer = nn.Linear(4, 4) + layer._hf_hook = SimpleNamespace() + assert get_module_device(layer) == torch.device("cpu") + + +def test_get_module_device_parameterless_module(): + """Returns cpu for a module with no parameters and no hook.""" + module = nn.Module() + assert get_module_device(module) == torch.device("cpu") diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index 92fe1345f94..73d3423ba55 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -20,7 +20,7 @@ convert_quantization_axis_to_reduce_axis, reduce_block_amax, ) -from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector +from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector @pytest.mark.parametrize( From 76b6fd51a54ab5b79bdd0d76facec0307473a237 Mon Sep 17 00:00:00 2001 From: "Chenhan D. Yu" <5185878+ChenhanYu@users.noreply.github.com> Date: Fri, 17 Apr 2026 23:11:17 -0700 Subject: [PATCH 13/30] fix: DFlash regression tests and vLLM server liveness (#1288) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - **hf_online_dflash.yaml**: Add 100K-sample training config with regression baselines (B200 loss curve), `MAX_FINAL_LOSS`/`MIN_FINAL_ACC`/`MIN_ACCEPTANCE_LENGTH` thresholds, vLLM nightly container for DFlash support - **vllm_smoke_test.sh**: Parse acceptance length from vLLM server log for regression check; `pip install pandas` workaround for broken nightly container; capture server output to temp file - **query.sh**: Detect vLLM server death during startup (PID liveness check) + 600s timeout to prevent infinite polling that wastes GPU hours; `pip install pandas` workaround - Fix empty `environment:` key in DFlash YAML causing nemo_run `ListParseError` ## Test plan - [x] E2E pipeline passed on 8x B200 (training + vLLM smoke test + AR eval) - [x] Training regression: final loss 3.82 < 5.0, acc 0.20 > 0.15 - [x] vLLM acceptance length: 1.79 >= 1.4 threshold - [x] AR evaluation: 2.02 overall on MT-Bench (8 categories) - [x] Server liveness check prevents GPU waste on vLLM crash 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **New Features** * Added optional regression validation for vLLM acceptance metrics * Introduced configurable vLLM server startup timeout (default 600 seconds) * **Improvements** * Enhanced logging for vLLM server startup with progress tracking and waited time reporting * Faster detection of vLLM server process failures during initialization * **Configuration Updates** * Increased training dataset size and logging granularity * Scaled tensor parallelism from 4 to 8 across multiple pipelines * Expanded PTQ quantization to multi-step pipeline * Added configurable training metric thresholds --------- Signed-off-by: Chenhan Yu Co-authored-by: Claude Opus 4.6 (1M context) --- .../common/megatron_lm/quantize/quantize.sh | 16 +- .../common/megatron_lm/quantize/task.py | 13 +- .../common/specdec/vllm_smoke_test.sh | 35 +- tools/launcher/common/tensorrt_llm/eval.sh | 59 ++++ .../tensorrt_llm/extra_llm_api_options.yaml | 52 +++ tools/launcher/common/vllm/query.sh | 20 +- .../Qwen/Qwen3-30B-A3B/megatron_lm_ptq.yaml | 53 +++ .../Qwen/Qwen3-8B/hf_offline_eagle3.yaml | 22 +- .../Qwen/Qwen3-8B/hf_online_dflash.yaml | 24 +- .../Qwen/Qwen3-8B/megatron_lm_ptq.yaml | 46 ++- uv.lock | 310 +++++++++++++++--- 11 files changed, 577 insertions(+), 73 deletions(-) create mode 100644 tools/launcher/common/tensorrt_llm/eval.sh create mode 100644 tools/launcher/common/tensorrt_llm/extra_llm_api_options.yaml create mode 100644 tools/launcher/examples/Qwen/Qwen3-30B-A3B/megatron_lm_ptq.yaml diff --git a/tools/launcher/common/megatron_lm/quantize/quantize.sh b/tools/launcher/common/megatron_lm/quantize/quantize.sh index 6e4d21b9945..1bb0d60e80d 100755 --- a/tools/launcher/common/megatron_lm/quantize/quantize.sh +++ b/tools/launcher/common/megatron_lm/quantize/quantize.sh @@ -36,10 +36,18 @@ CONVERT_EXE="bash modules/Megatron-LM/examples/post_training/modelopt/convert.sh EXPORT_EXE="bash modules/Megatron-LM/examples/post_training/modelopt/export.sh" export MLM_EXTRA_ARGS=${@} -${QUANTIZE_EXE} ${MLM_MODEL_CFG} ${QUANT_CFG} - -export MLM_EXTRA_ARGS="--mmlu-dataset ${MMLU_DATASET:-/hf-local/cais/mmlu} --fraction 0.01 --lower-bound 0.38 --disable-tqdm" -MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${MMLU_EXE} ${MLM_MODEL_CFG} +TP=${TP:-1} PP=${PP:-1} EP=${EP:-1} ETP=${ETP:-1} ${QUANTIZE_EXE} ${MLM_MODEL_CFG} ${QUANT_CFG} + +export MLM_EXTRA_ARGS="--mmlu-dataset ${MMLU_DATASET:-/hf-local/cais/mmlu} --fraction 0.01 --lower-bound ${MMLU_LOWER_BOUND:-0.38} --disable-tqdm" +TP=${TP:-1} PP=${PP:-1} EP=${EP:-1} ETP=${ETP:-1} MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${MMLU_EXE} ${MLM_MODEL_CFG} + +# Export quantized checkpoint to HF format (PP=all GPUs) +TOTAL_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo ${NUM_GPUS:-1}) +echo "=== Exporting ${MLM_MODEL_CFG} ${QUANT_CFG} (PP=${TOTAL_GPUS}) ===" +export MLM_EXTRA_ARGS= +TP=1 PP=${TOTAL_GPUS} EP=1 ETP=1 MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${EXPORT_EXE} ${MLM_MODEL_CFG} +ls ${EXPORT_DIR} +cat ${EXPORT_DIR}/hf_quant_config.json ################################################################################################### diff --git a/tools/launcher/common/megatron_lm/quantize/task.py b/tools/launcher/common/megatron_lm/quantize/task.py index 4ba10030e61..95833fe3960 100644 --- a/tools/launcher/common/megatron_lm/quantize/task.py +++ b/tools/launcher/common/megatron_lm/quantize/task.py @@ -66,6 +66,10 @@ class MegatronLMQuantizeConfig: model: str = "Qwen/Qwen3-8B" quant_cfg: str = "NVFP4_DEFAULT_CFG" tp: int = 4 + pp: int = 1 + ep: int = 1 + etp: int = 1 + extra_args: str = "" calib_dataset: str = "abisee/cnn_dailymail" calib_size: int = 32 mmlu_dataset: str = "cais/mmlu" @@ -92,14 +96,21 @@ def __post_init__(self): if self.config is not None: c = self.config self.script = self.script or "common/megatron_lm/quantize/quantize.sh" - self.args = [ + args = [ f"--calib-dataset-path-or-name {c.hf_local}{c.calib_dataset}", f"--calib-size {c.calib_size}", ] + if c.extra_args: + args.append(c.extra_args) + self.args = args self.environment = [ {"MLM_MODEL_CFG": c.model}, {"QUANT_CFG": c.quant_cfg}, {"HF_MODEL_CKPT": f"{c.hf_local}{c.model}"}, {"MMLU_DATASET": f"{c.hf_local}{c.mmlu_dataset}"}, {"TP": str(c.tp)}, + {"PP": str(c.pp)}, + {"EP": str(c.ep)}, + {"ETP": str(c.etp)}, + {"MMLU_LOWER_BOUND": str(c.mmlu_lower_bound)}, ] diff --git a/tools/launcher/common/specdec/vllm_smoke_test.sh b/tools/launcher/common/specdec/vllm_smoke_test.sh index 1e508cb1169..4b9d5a63b4f 100644 --- a/tools/launcher/common/specdec/vllm_smoke_test.sh +++ b/tools/launcher/common/specdec/vllm_smoke_test.sh @@ -32,7 +32,10 @@ SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" source ${SCRIPT_DIR}/../service_utils.sh 2>/dev/null || true -cleanup() { kill $SERVER_PID 2>/dev/null; sleep 2; kill -9 $SERVER_PID 2>/dev/null; } +# Ensure pandas is available (missing in some vLLM nightly builds) +pip install pandas 2>/dev/null || true + +cleanup() { kill $SERVER_PID 2>/dev/null; sleep 2; kill -9 $SERVER_PID 2>/dev/null; rm -f "${VLLM_LOG:-}" 2>/dev/null; } trap cleanup EXIT MODEL=${HF_MODEL_CKPT} @@ -72,7 +75,8 @@ if [ "${DISABLE_PREFIX_CACHING:-}" = "1" ]; then OPTIONAL_ARGS="${OPTIONAL_ARGS} --no-enable-prefix-caching" fi -# Start vLLM server +# Start vLLM server (capture output for regression check parsing) +VLLM_LOG=$(mktemp /tmp/vllm_server_XXXXXX.log) if [ -n "$SPEC_CONFIG" ]; then vllm serve ${MODEL} \ --speculative-config "${SPEC_CONFIG}" \ @@ -80,14 +84,14 @@ if [ -n "$SPEC_CONFIG" ]; then --tensor-parallel-size ${TP} \ --port ${PORT} \ ${OPTIONAL_ARGS} \ - & + > >(tee -a "$VLLM_LOG") 2>&1 & else vllm serve ${MODEL} \ --max-num-batched-tokens 32768 \ --tensor-parallel-size ${TP} \ --port ${PORT} \ ${OPTIONAL_ARGS} \ - & + > >(tee -a "$VLLM_LOG") 2>&1 & fi SERVER_PID=$! @@ -168,4 +172,27 @@ if [ $FAIL -gt 0 ]; then exit 1 fi +# Regression check: minimum acceptance length for speculative decoding +if [ -n "${MIN_ACCEPTANCE_LENGTH:-}" ]; then + # Parse mean acceptance length from vLLM's SpecDecoding metrics log. + # vLLM logs: "SpecDecoding metrics: Mean acceptance length: X.XX, ..." + # Take the last reported value (most accurate, covers all prompts). + AVG_ACCEPT=$(grep -oP 'Mean acceptance length: \K[0-9.]+' "$VLLM_LOG" 2>/dev/null | tail -1 || true) + if [ -n "$AVG_ACCEPT" ]; then + echo "" + echo "=== Acceptance Length Regression Check ===" + echo " Mean acceptance length: ${AVG_ACCEPT}" + echo " Threshold: ${MIN_ACCEPTANCE_LENGTH}" + PASS_CHECK=$(python3 -c "print('yes' if float('${AVG_ACCEPT}') >= float('${MIN_ACCEPTANCE_LENGTH}') else 'no')") + if [ "$PASS_CHECK" = "yes" ]; then + echo " PASS: ${AVG_ACCEPT} >= ${MIN_ACCEPTANCE_LENGTH}" + else + echo " REGRESSION: ${AVG_ACCEPT} < ${MIN_ACCEPTANCE_LENGTH}" + exit 1 + fi + else + echo "WARNING: Could not parse acceptance length from vLLM log, skipping regression check" + fi +fi + echo "Done" diff --git a/tools/launcher/common/tensorrt_llm/eval.sh b/tools/launcher/common/tensorrt_llm/eval.sh new file mode 100644 index 00000000000..3e3f2d1b768 --- /dev/null +++ b/tools/launcher/common/tensorrt_llm/eval.sh @@ -0,0 +1,59 @@ +#!/bin/bash + +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +SCRIPT_DIR="$(dirname "$(readlink -f "$0")")" + +################################################################################################### + +if [[ -z ${HF_MODEL_CKPT} ]]; then + export HF_MODEL_CKPT=/scratchspace/export +fi + +if [[ -z ${TP} ]]; then + TP=4 +fi + +if [[ -z ${EP} ]]; then + EP=4 +fi + +if [[ -z ${EXTRA_LLM_API_OPTIONS} ]]; then + EXTRA_LLM_API_OPTIONS=common/tensorrt_llm/extra_llm_api_options.yaml +fi + + +TARGET_FILENAME="config.json" + + +# Find all files matching the target filename, print their paths null-terminated +find "${HF_MODEL_CKPT}" -type f -name "$TARGET_FILENAME" -print0 | while IFS= read -r -d '' filepath; do + # Extract the directory path from the full file path + dir_path=$(dirname "$filepath") + + echo "Processing model: $dir_path" + # Place your commands here to run within or on the $dir_path + # Example: cd "$dir_path" && some_command + + trtllm-llmapi-launch trtllm-eval \ + --model ${dir_path} \ + --disable_kv_cache_reuse \ + --tp_size ${TP} \ + --ep_size ${EP} \ + --trust_remote_code \ + --extra_llm_api_options ${EXTRA_LLM_API_OPTIONS} \ + mmlu +done diff --git a/tools/launcher/common/tensorrt_llm/extra_llm_api_options.yaml b/tools/launcher/common/tensorrt_llm/extra_llm_api_options.yaml new file mode 100644 index 00000000000..f5f43140861 --- /dev/null +++ b/tools/launcher/common/tensorrt_llm/extra_llm_api_options.yaml @@ -0,0 +1,52 @@ +context_parallel_size: 1 + + # backend: _autodeploy + # reasoning_parser: nano-v3 + # tool_parser: qwen3_coder + # + # runtime: trtllm + # compile_backend: torch-cudagraph + # max_batch_size: 64 + # max_seq_len: 16384 + # enable_chunked_prefill: true + # attn_backend: flashinfer + # model_factory: AutoModelForCausalLM + # skip_loading_weights: false + # free_mem_ratio: 0.65 + # cuda_graph_batch_sizes: [1, 2, 4, 8, 16, 24, 32, 64, 128, 256, 320, 384] + # kv_cache_config: + # # disable kv_cache reuse since not supported for hybrid/ssm models + # enable_block_reuse: false + # transforms: + # detect_sharding: + # sharding_dims: ['ep', 'bmm'] + # allreduce_strategy: 'AUTO' + # manual_config: + # head_dim: 128 + # tp_plan: + # # mamba SSM layer + # "in_proj": "mamba" + # "out_proj": "rowwise" + # # attention layer + # "q_proj": "colwise" + # "k_proj": "colwise" + # "v_proj": "colwise" + # "o_proj": "rowwise" + # # NOTE: consider not sharding shared experts and/or + # # latent projections at all, keeping them replicated. + # # To do so, comment out the corresponding entries. + # # moe layer: SHARED experts + # "up_proj": "colwise" + # "down_proj": "rowwise" + # # MoLE: latent projections: simple shard + # "fc1_latent_proj": "gather" + # "fc2_latent_proj": "gather" + # multi_stream_moe: + # stage: compile + # enabled: true + # insert_cached_ssm_attention: + # cache_config: + # mamba_dtype: float32 + # fuse_mamba_a_log: + # stage: post_load_fusion + # enabled: true diff --git a/tools/launcher/common/vllm/query.sh b/tools/launcher/common/vllm/query.sh index 4ce6ded1965..d1513623c34 100755 --- a/tools/launcher/common/vllm/query.sh +++ b/tools/launcher/common/vllm/query.sh @@ -58,6 +58,9 @@ source ${SCRIPT_DIR}/../service_utils.sh # gpus_per_node: 4 ################################################################################################### +# Ensure pandas is available (missing in some vLLM nightly builds) +pip install pandas 2>/dev/null || true + export OPENAI_API_KEY="token-abc123" if [ -z ${SLURM_ARRAY_TASK_ID} ]; then @@ -108,13 +111,26 @@ SERVER_PID=$! # Wait for server to start up by polling the health endpoint echo "Waiting for server to start..." +MAX_WAIT=${VLLM_STARTUP_TIMEOUT:-600} +WAITED=0 while true; do + if ! kill -0 $SERVER_PID 2>/dev/null; then + echo "ERROR: vLLM server process died during startup" + wait $SERVER_PID 2>/dev/null + exit 1 + fi response=$(curl -s -o /dev/null -w "%{http_code}" "http://$(hostname -f):8000/health" || true) if [ "$response" -eq 200 ]; then - echo "Server is up!" + echo "Server is up! (waited ${WAITED}s)" break fi - echo "Server not ready yet, retrying in 10 seconds..." + WAITED=$((WAITED + 10)) + if [ $WAITED -ge $MAX_WAIT ]; then + echo "ERROR: vLLM server failed to start within ${MAX_WAIT}s" + kill $SERVER_PID 2>/dev/null + exit 1 + fi + echo "Server not ready yet (${WAITED}/${MAX_WAIT}s), retrying in 10 seconds..." sleep 10 done diff --git a/tools/launcher/examples/Qwen/Qwen3-30B-A3B/megatron_lm_ptq.yaml b/tools/launcher/examples/Qwen/Qwen3-30B-A3B/megatron_lm_ptq.yaml new file mode 100644 index 00000000000..0eeca6531c9 --- /dev/null +++ b/tools/launcher/examples/Qwen/Qwen3-30B-A3B/megatron_lm_ptq.yaml @@ -0,0 +1,53 @@ +# Qwen3-30B-A3B PTQ quantization (8 GPUs, MoE model). +# +# 2-step pipeline: NVFP4 then FP8, each followed by MMLU evaluation. +# MMLU uses EP for expert parallelism. +# +# Usage: +# uv run launch.py --yaml examples/Qwen/Qwen3-30B-A3B/megatron_lm_ptq.yaml --yes + +job_name: Qwen3-30B-A3B_PTQ +pipeline: + skip: false + allow_to_fail: false + note: + + task_0: + _target_: common.megatron_lm.quantize.task.MegatronLMQuantizeTask + config: + model: Qwen/Qwen3-30B-A3B + quant_cfg: NVFP4_DEFAULT_CFG + tp: 1 + pp: 1 + ep: 8 + etp: 1 + calib_dataset: abisee/cnn_dailymail + calib_size: 32 + mmlu_dataset: cais/mmlu + mmlu_lower_bound: 0.75 + hf_local: /hf-local/ + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 8 + gpus_per_node: 8 + + task_1: + _target_: common.megatron_lm.quantize.task.MegatronLMQuantizeTask + config: + model: Qwen/Qwen3-30B-A3B + quant_cfg: FP8_DEFAULT_CFG + tp: 1 + pp: 1 + ep: 8 + etp: 1 + calib_dataset: abisee/cnn_dailymail + calib_size: 32 + mmlu_dataset: cais/mmlu + mmlu_lower_bound: 0.75 + hf_local: /hf-local/ + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 8 + gpus_per_node: 8 diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml index ae2c1e957cf..24068c4bb3e 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_offline_eagle3.yaml @@ -27,8 +27,8 @@ pipeline: script: common/tensorrt_llm/query.sh args: - --model <> - - --tp_size 4 - - --ep_size 4 + - --tp_size 8 + - --ep_size 8 - --max_num_tokens 32000 - --port 8000 - --host 0.0.0.0 @@ -41,8 +41,8 @@ pipeline: slurm_config: _factory_: "slurm_factory" nodes: 1 - ntasks_per_node: 4 - gpus_per_node: 4 + ntasks_per_node: 8 + gpus_per_node: 8 container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0 # Step 2: Dump hidden states from target model @@ -52,15 +52,15 @@ pipeline: - --input-data /scratchspace/data - --output-dir /scratchspace/offline_hidden_states - --max-seq-len 8192 - - --tp 4 - - --moe-ep 4 + - --tp 8 + - --moe-ep 8 environment: - HF_MODEL_CKPT: <> slurm_config: _factory_: "slurm_factory" nodes: 1 - ntasks_per_node: 4 - gpus_per_node: 4 + ntasks_per_node: 8 + gpus_per_node: 8 container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0 # Step 3: Train EAGLE3 draft head (offline, single task) @@ -78,7 +78,7 @@ pipeline: _factory_: "slurm_factory" nodes: 1 ntasks_per_node: 1 - gpus_per_node: 4 + gpus_per_node: 8 container: nvcr.io/nvidia/tensorrt-llm/release:1.2.0 # Step 4: Benchmark speculative decoding (VLLM backend) @@ -89,7 +89,7 @@ pipeline: - --draft_length 3 - --output_length 4096 - --engine VLLM - - --tp_size 4 + - --tp_size 8 - --ep_size 1 - --speculative_algorithm EAGLE3 - --mtbench /hf-local/HuggingFaceH4/mt_bench_prompts/raw/question.jsonl @@ -100,5 +100,5 @@ pipeline: _factory_: "slurm_factory" nodes: 1 ntasks_per_node: 1 - gpus_per_node: 4 + gpus_per_node: 8 container: vllm/vllm-openai:latest diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash.yaml index e255f51dc32..7c7f2a959dc 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/hf_online_dflash.yaml @@ -5,6 +5,21 @@ # task_1: vLLM smoke test with DFlash speculative decoding # task_2: MT-Bench per-category HF AR evaluation (1 GPU) # +# Convergence baseline (8x B200, batch_size=1, seq_len=4096, 5-layer draft, block_size=16): +# 100K samples, 1 epoch (~12,500 steps) +# Step 100 (epoch 0.01): loss=8.900 acc=0.029 +# Step 1000 (epoch 0.08): loss=5.845 acc=0.096 +# Step 2500 (epoch 0.20): loss=4.981 acc=0.138 +# Step 5000 (epoch 0.40): loss=4.383 acc=0.176 +# Step 7500 (epoch 0.60): loss=4.040 acc=0.196 +# Step 10000 (epoch 0.80): loss=3.900 acc=0.210 +# Step 12500 (epoch 1.00): loss=3.821 acc=0.200 +# Average train_loss=4.493, training time=5094s +# +# Regression criteria (set via environment): +# MAX_FINAL_LOSS: final loss must be below this (default: 5.0) +# MIN_FINAL_ACC: final accuracy must be above this (default: 0.15) +# # Reference: "DFlash: Block Diffusion for Flash Speculative Decoding" (arXiv:2602.06036) # # Usage: @@ -22,13 +37,14 @@ pipeline: args: - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dflash.yaml - model.model_name_or_path=<> - - data.data_path=/hf-local/modelopt/Speculative-Decoding-Dataset-v1-Qwen3-8B/sample-1K-openai.jsonl + - data.data_path=/hf-local/modelopt/Speculative-Decoding-Dataset-v1-Qwen3-8B/sample-100K-openai.jsonl - data.chat_template=examples/Qwen/Qwen3-8B/chat_template_train.jinja - training.output_dir=/scratchspace/dflash_bs16 + - training.per_device_train_batch_size=1 - training.num_train_epochs=1 - training.training_seq_len=4096 - training.save_steps=5000 - - training.logging_steps=1000 + - training.logging_steps=100 - training.disable_tqdm=true - training.answer_only_loss=true - dflash.dflash_block_size=16 @@ -37,6 +53,8 @@ pipeline: - dflash.dflash_mask_token_id=151669 - dflash.dflash_architecture_config.num_hidden_layers=5 environment: + - MAX_FINAL_LOSS: "5.0" + - MIN_FINAL_ACC: "0.15" slurm_config: _factory_: "slurm_factory" nodes: 1 @@ -51,8 +69,10 @@ pipeline: - DRAFT_CKPT_DIR: /scratchspace/dflash_bs16 - SPEC_METHOD: "dflash" - NUM_SPEC_TOKENS: "7" + - MIN_ACCEPTANCE_LENGTH: "1.4" slurm_config: _factory_: "slurm_factory" + container: "vllm/vllm-openai:nightly" nodes: 1 ntasks_per_node: 1 gpus_per_node: 1 diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml index a67d0bceb07..33b9da18e66 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml @@ -1,4 +1,9 @@ -# Qwen3-8B NVFP4 quantization (4 GPUs, for Slurm clusters). +# Qwen3-8B PTQ quantization (8 GPUs, for Slurm clusters). +# +# 3-step pipeline: +# task_0: NVFP4 quantize → MMLU → export +# task_1: FP8 quantize → MMLU → export +# task_2: TRT-LLM eval MMLU on all exported checkpoints # # Uses MegatronLMQuantizeTask with typed config — see common/megatron_lm/quantize/task.py # for all available fields. @@ -8,7 +13,7 @@ # # For single-GPU local Docker, use megatron_lm_ptq_local.yaml instead. -job_name: Qwen3-8B_NVFP4_DEFAULT_CFG +job_name: Qwen3-8B_PTQ pipeline: skip: false allow_to_fail: false @@ -19,13 +24,44 @@ pipeline: config: model: Qwen/Qwen3-8B quant_cfg: NVFP4_DEFAULT_CFG - tp: 4 + tp: 8 + calib_dataset: abisee/cnn_dailymail + calib_size: 32 + mmlu_dataset: cais/mmlu + mmlu_lower_bound: 0.68 + hf_local: /hf-local/ + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 8 + gpus_per_node: 8 + + task_1: + _target_: common.megatron_lm.quantize.task.MegatronLMQuantizeTask + config: + model: Qwen/Qwen3-8B + quant_cfg: FP8_DEFAULT_CFG + tp: 8 calib_dataset: abisee/cnn_dailymail calib_size: 32 mmlu_dataset: cais/mmlu + mmlu_lower_bound: 0.68 hf_local: /hf-local/ slurm_config: _factory_: "slurm_factory" nodes: 1 - ntasks_per_node: 4 - gpus_per_node: 4 + ntasks_per_node: 8 + gpus_per_node: 8 + + # Step 3: TRT-LLM eval MMLU on all exported checkpoints + task_2: + script: common/tensorrt_llm/eval.sh + environment: + - HF_MODEL_CKPT: /scratchspace/export + - TP: "8" + - EP: "1" + slurm_config: + _factory_: "slurm_factory" + nodes: 1 + ntasks_per_node: 1 + gpus_per_node: 8 diff --git a/uv.lock b/uv.lock index 0f4710f92b4..cfa742f1081 100644 --- a/uv.lock +++ b/uv.lock @@ -32,9 +32,6 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", ] -[manifest] -overrides = [{ name = "torch", marker = "sys_platform == 'never'" }] - [[package]] name = "accelerate" version = "1.13.0" @@ -47,7 +44,7 @@ dependencies = [ { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch", marker = "sys_platform == 'never'" }, + { name = "torch" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ca/14/787e5498cd062640f0f3d92ef4ae4063174f76f9afd29d13fc52a319daae/accelerate-1.13.0.tar.gz", hash = "sha256:d631b4e0f5b3de4aff2d7e9e6857d164810dfc3237d54d017f075122d057b236", size = 402835, upload-time = "2026-03-04T19:34:12.359Z" } wheels = [ @@ -573,6 +570,21 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/54/27/01d9078a77b9e31b79b9716e66ca4db74f4744c5232bcb3e8769395c4280/cppimport-22.8.2.tar.gz", hash = "sha256:bbb4957102db41bc99ad72c233bce92f9d1fd91be352fc07878c4361033a401f", size = 26635, upload-time = "2022-08-02T16:50:36.872Z" } +[[package]] +name = "cuda-bindings" +version = "12.9.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cuda-pathfinder", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/d8/b546104b8da3f562c1ff8ab36d130c8fe1dd6a045ced80b4f6ad74f7d4e1/cuda_bindings-12.9.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d3c842c2a4303b2a580fe955018e31aea30278be19795ae05226235268032e5", size = 12148218, upload-time = "2025-10-21T14:51:28.855Z" }, + { url = "https://files.pythonhosted.org/packages/45/e7/b47792cc2d01c7e1d37c32402182524774dadd2d26339bd224e0e913832e/cuda_bindings-12.9.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c912a3d9e6b6651853eed8eed96d6800d69c08e94052c292fec3f282c5a817c9", size = 12210593, upload-time = "2025-10-21T14:51:36.574Z" }, + { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" }, + { url = "https://files.pythonhosted.org/packages/63/56/e465c31dc9111be3441a9ba7df1941fe98f4aa6e71e8788a3fb4534ce24d/cuda_bindings-12.9.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:32bdc5a76906be4c61eb98f546a6786c5773a881f3b166486449b5d141e4a39f", size = 11906628, upload-time = "2025-10-21T14:51:49.905Z" }, + { url = "https://files.pythonhosted.org/packages/a3/84/1e6be415e37478070aeeee5884c2022713c1ecc735e6d82d744de0252eee/cuda_bindings-12.9.4-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56e0043c457a99ac473ddc926fe0dc4046694d99caef633e92601ab52cbe17eb", size = 11925991, upload-time = "2025-10-21T14:51:56.535Z" }, +] + [[package]] name = "cuda-pathfinder" version = "1.5.3" @@ -637,18 +649,18 @@ name = "deepspeed" version = "0.18.9" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "einops", marker = "sys_platform != 'win32'" }, - { name = "hjson", marker = "sys_platform != 'win32'" }, - { name = "msgpack", marker = "sys_platform != 'win32'" }, - { name = "ninja", marker = "sys_platform != 'win32'" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' and sys_platform != 'win32'" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and sys_platform != 'win32'" }, - { name = "packaging", marker = "sys_platform != 'win32'" }, - { name = "psutil", marker = "sys_platform != 'win32'" }, - { name = "py-cpuinfo", marker = "sys_platform != 'win32'" }, - { name = "pydantic", marker = "sys_platform != 'win32'" }, - { name = "torch", marker = "sys_platform == 'never'" }, - { name = "tqdm", marker = "sys_platform != 'win32'" }, + { name = "einops", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "hjson", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "msgpack", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "ninja", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "packaging", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "psutil", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "py-cpuinfo", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "pydantic", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "torch", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "tqdm", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/24/61/5ea1c63b139fe7530b196b68ce0bdffa9cde79882e527dcecae58bd6c770/deepspeed-0.18.9.tar.gz", hash = "sha256:ee4818dcf342794f74f429a0aeebef90291ec808fa82609c5140c23e665c4011", size = 1663466, upload-time = "2026-03-30T16:43:16.566Z" } @@ -1520,7 +1532,11 @@ name = "networkx" version = "3.4.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ + "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'win32'", "(python_full_version < '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')", + "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform == 'darwin'", + "python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform == 'win32'", "python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", ] @@ -1534,15 +1550,27 @@ name = "networkx" version = "3.6.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ + "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'win32'", "(python_full_version >= '3.13' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')", + "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform == 'darwin'", "(python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'darwin'", "(python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'darwin'", "python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", + "python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } wheels = [ @@ -1751,6 +1779,108 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/74/f4c001f4714c3ad9ce037e18cf2b9c64871a84951eaa0baf683a9ca9301c/numpy-2.4.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:f2cf083b324a467e1ab358c105f6cad5ea950f50524668a80c486ff1db24e119", size = 12509075, upload-time = "2026-03-29T13:21:57.644Z" }, ] +[[package]] +name = "nvidia-cublas-cu12" +version = "12.8.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.10.2.21" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.3.3.83" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, +] + +[[package]] +name = "nvidia-cufile-cu12" +version = "1.13.1.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.9.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.3.90" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.5.8.93" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, +] + +[[package]] +name = "nvidia-cusparselt-cu12" +version = "0.7.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, +] + [[package]] name = "nvidia-ml-py" version = "13.595.45" @@ -1779,7 +1909,7 @@ dependencies = [ { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "setuptools" }, - { name = "torch", marker = "sys_platform == 'never'" }, + { name = "torch" }, { name = "tqdm" }, ] @@ -2026,6 +2156,38 @@ requires-dist = [ ] provides-extras = ["onnx", "hf", "puzzletron", "dev-lint", "dev-docs", "dev-test", "all", "dev"] +[[package]] +name = "nvidia-nccl-cu12" +version = "2.27.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.8.93" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, +] + +[[package]] +name = "nvidia-nvshmem-cu12" +version = "3.4.5" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.8.90" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, +] + [[package]] name = "omegaconf" version = "2.3.0" @@ -2442,7 +2604,7 @@ dependencies = [ { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch", marker = "sys_platform == 'never'" }, + { name = "torch" }, { name = "tqdm" }, { name = "transformers" }, ] @@ -3782,7 +3944,7 @@ dependencies = [ { name = "huggingface-hub" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch", marker = "sys_platform == 'never'" }, + { name = "torch" }, { name = "torchvision" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7b/1e/e924b3b2326a856aaf68586f9c52a5fc81ef45715eca408393b68c597e0e/timm-1.0.26.tar.gz", hash = "sha256:f66f082f2f381cf68431c22714c8b70f723837fa2a185b155961eab90f2d5b10", size = 2419859, upload-time = "2026-03-23T18:12:10.272Z" } @@ -3870,15 +4032,63 @@ name = "torch" version = "2.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ + { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "filelock" }, { name = "fsspec" }, { name = "jinja2" }, { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "setuptools", marker = "python_full_version >= '3.12'" }, { name = "sympy" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "typing-extensions" }, ] +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/30/bfebdd8ec77db9a79775121789992d6b3b75ee5494971294d7b4b7c999bc/torch-2.10.0-2-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:2b980edd8d7c0a68c4e951ee1856334a43193f98730d97408fbd148c1a933313", size = 79411457, upload-time = "2026-02-10T21:44:59.189Z" }, + { url = "https://files.pythonhosted.org/packages/0f/8b/4b61d6e13f7108f36910df9ab4b58fd389cc2520d54d81b88660804aad99/torch-2.10.0-2-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:418997cb02d0a0f1497cf6a09f63166f9f5df9f3e16c8a716ab76a72127c714f", size = 79423467, upload-time = "2026-02-10T21:44:48.711Z" }, + { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" }, + { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" }, + { url = "https://files.pythonhosted.org/packages/16/ee/efbd56687be60ef9af0c9c0ebe106964c07400eade5b0af8902a1d8cd58c/torch-2.10.0-3-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:a1ff626b884f8c4e897c4c33782bdacdff842a165fee79817b1dd549fdda1321", size = 915510070, upload-time = "2026-03-11T14:16:39.386Z" }, + { url = "https://files.pythonhosted.org/packages/36/ab/7b562f1808d3f65414cd80a4f7d4bb00979d9355616c034c171249e1a303/torch-2.10.0-3-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ac5bdcbb074384c66fa160c15b1ead77839e3fe7ed117d667249afce0acabfac", size = 915518691, upload-time = "2026-03-11T14:15:43.147Z" }, + { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" }, + { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" }, + { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" }, + { url = "https://files.pythonhosted.org/packages/0c/1a/c61f36cfd446170ec27b3a4984f072fd06dab6b5d7ce27e11adb35d6c838/torch-2.10.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:5276fa790a666ee8becaffff8acb711922252521b28fbce5db7db5cf9cb2026d", size = 145992962, upload-time = "2026-01-21T16:24:14.04Z" }, + { url = "https://files.pythonhosted.org/packages/b5/60/6662535354191e2d1555296045b63e4279e5a9dbad49acf55a5d38655a39/torch-2.10.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:aaf663927bcd490ae971469a624c322202a2a1e68936eb952535ca4cd3b90444", size = 915599237, upload-time = "2026-01-21T16:23:25.497Z" }, + { url = "https://files.pythonhosted.org/packages/40/b8/66bbe96f0d79be2b5c697b2e0b187ed792a15c6c4b8904613454651db848/torch-2.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:a4be6a2a190b32ff5c8002a0977a25ea60e64f7ba46b1be37093c141d9c49aeb", size = 113720931, upload-time = "2026-01-21T16:24:23.743Z" }, + { url = "https://files.pythonhosted.org/packages/76/bb/d820f90e69cda6c8169b32a0c6a3ab7b17bf7990b8f2c680077c24a3c14c/torch-2.10.0-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:35e407430795c8d3edb07a1d711c41cc1f9eaddc8b2f1cc0a165a6767a8fb73d", size = 79411450, upload-time = "2026-01-21T16:25:30.692Z" }, + { url = "https://files.pythonhosted.org/packages/78/89/f5554b13ebd71e05c0b002f95148033e730d3f7067f67423026cc9c69410/torch-2.10.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:3282d9febd1e4e476630a099692b44fdc214ee9bf8ee5377732d9d9dfe5712e4", size = 145992610, upload-time = "2026-01-21T16:25:26.327Z" }, + { url = "https://files.pythonhosted.org/packages/ae/30/a3a2120621bf9c17779b169fc17e3dc29b230c29d0f8222f499f5e159aa8/torch-2.10.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a2f9edd8dbc99f62bc4dfb78af7bf89499bca3d753423ac1b4e06592e467b763", size = 915607863, upload-time = "2026-01-21T16:25:06.696Z" }, + { url = "https://files.pythonhosted.org/packages/6f/3d/c87b33c5f260a2a8ad68da7147e105f05868c281c63d65ed85aa4da98c66/torch-2.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:29b7009dba4b7a1c960260fc8ac85022c784250af43af9fb0ebafc9883782ebd", size = 113723116, upload-time = "2026-01-21T16:25:21.916Z" }, + { url = "https://files.pythonhosted.org/packages/61/d8/15b9d9d3a6b0c01b883787bd056acbe5cc321090d4b216d3ea89a8fcfdf3/torch-2.10.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:b7bd80f3477b830dd166c707c5b0b82a898e7b16f59a7d9d42778dd058272e8b", size = 79423461, upload-time = "2026-01-21T16:24:50.266Z" }, + { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" }, + { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" }, + { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" }, + { url = "https://files.pythonhosted.org/packages/c9/5c/dee910b87c4d5c0fcb41b50839ae04df87c1cfc663cf1b5fca7ea565eeaa/torch-2.10.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6d3707a61863d1c4d6ebba7be4ca320f42b869ee657e9b2c21c736bf17000294", size = 79498198, upload-time = "2026-01-21T16:24:34.704Z" }, + { url = "https://files.pythonhosted.org/packages/c9/6f/f2e91e34e3fcba2e3fc8d8f74e7d6c22e74e480bbd1db7bc8900fdf3e95c/torch-2.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5c4d217b14741e40776dd7074d9006fd28b8a97ef5654db959d8635b2fe5f29b", size = 146004247, upload-time = "2026-01-21T16:24:29.335Z" }, + { url = "https://files.pythonhosted.org/packages/98/fb/5160261aeb5e1ee12ee95fe599d0541f7c976c3701d607d8fc29e623229f/torch-2.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6b71486353fce0f9714ca0c9ef1c850a2ae766b409808acd58e9678a3edb7738", size = 915716445, upload-time = "2026-01-21T16:22:45.353Z" }, + { url = "https://files.pythonhosted.org/packages/6a/16/502fb1b41e6d868e8deb5b0e3ae926bbb36dab8ceb0d1b769b266ad7b0c3/torch-2.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2ee399c644dc92ef7bc0d4f7e74b5360c37cdbe7c5ba11318dda49ffac2bc57", size = 113757050, upload-time = "2026-01-21T16:24:19.204Z" }, + { url = "https://files.pythonhosted.org/packages/1a/0b/39929b148f4824bc3ad6f9f72a29d4ad865bcf7ebfc2fa67584773e083d2/torch-2.10.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:3202429f58309b9fa96a614885eace4b7995729f44beb54d3e4a47773649d382", size = 79851305, upload-time = "2026-01-21T16:24:09.209Z" }, + { url = "https://files.pythonhosted.org/packages/d8/14/21fbce63bc452381ba5f74a2c0a959fdf5ad5803ccc0c654e752e0dbe91a/torch-2.10.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:aae1b29cd68e50a9397f5ee897b9c24742e9e306f88a807a27d617f07adb3bd8", size = 146005472, upload-time = "2026-01-21T16:22:29.022Z" }, + { url = "https://files.pythonhosted.org/packages/54/fd/b207d1c525cb570ef47f3e9f836b154685011fce11a2f444ba8a4084d042/torch-2.10.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6021db85958db2f07ec94e1bc77212721ba4920c12a18dc552d2ae36a3eb163f", size = 915612644, upload-time = "2026-01-21T16:21:47.019Z" }, + { url = "https://files.pythonhosted.org/packages/36/53/0197f868c75f1050b199fe58f9bf3bf3aecac9b4e85cc9c964383d745403/torch-2.10.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff43db38af76fda183156153983c9a096fc4c78d0cd1e07b14a2314c7f01c2c8", size = 113997015, upload-time = "2026-01-21T16:23:00.767Z" }, + { url = "https://files.pythonhosted.org/packages/0e/13/e76b4d9c160e89fff48bf16b449ea324bda84745d2ab30294c37c2434c0d/torch-2.10.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:cdf2a523d699b70d613243211ecaac14fe9c5df8a0b0a9c02add60fb2a413e0f", size = 79498248, upload-time = "2026-01-21T16:23:09.315Z" }, +] [[package]] name = "torch-geometric" @@ -3908,7 +4118,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "torch", marker = "sys_platform == 'never'" }, + { name = "torch" }, { name = "torchvision" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6f/36/574c0c46e818533b78b3c09505211162918188325ab4165ef11a3f295755/torchprofile-0.0.4.tar.gz", hash = "sha256:96b6da17d752a06b02977e078aea95614893b31d4117dd5dcd081f30ce65611b", size = 4557, upload-time = "2021-06-22T04:58:03.592Z" } @@ -3918,35 +4128,35 @@ wheels = [ [[package]] name = "torchvision" -version = "0.26.0" +version = "0.25.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pillow" }, - { name = "torch", marker = "sys_platform == 'never'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/74/b4/cdfee31e0402ea035135462cb0ab496e974d56fab6b4e7a1f0cbccb8cd28/torchvision-0.26.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a06d4772a8e13e772906ed736cc53ec6639e5e60554f8e5fa6ca165aabebc464", size = 1863503, upload-time = "2026-03-23T18:13:01.384Z" }, - { url = "https://files.pythonhosted.org/packages/e4/74/11fee109841e80ad14e5ca2d80bff6b10eb11b7838ff06f35bfeaa9f7251/torchvision-0.26.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:2adfbe438473236191ff077a4a9a0c767436879c89628aa97137e959b0c11a94", size = 7766423, upload-time = "2026-03-23T18:12:56.049Z" }, - { url = "https://files.pythonhosted.org/packages/5e/00/24d8c7845c3f270153fb81395a5135b2778e2538e81d14c6aea5106c689c/torchvision-0.26.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:b6f9ad1ecc0eab52647298b379ee9426845f8903703e6127973f8f3d049a798b", size = 7518249, upload-time = "2026-03-23T18:12:51.743Z" }, - { url = "https://files.pythonhosted.org/packages/d7/ed/e53cd7c0da7ae002e5e929c1796ebbe7ec0c700c29f7a0a6696497fb3d8b/torchvision-0.26.0-cp310-cp310-win_amd64.whl", hash = "sha256:f13f12b3791a266de2d599cb8162925261622a037d87fc03132848343cf68f75", size = 3669784, upload-time = "2026-03-23T18:12:49.949Z" }, - { url = "https://files.pythonhosted.org/packages/b4/bd/d552a2521bade3295b2c6e7a4a0d1022261cab7ca7011f4e2a330dbb3caa/torchvision-0.26.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:55bd6ad4ae77be01ba67a410b05b51f53b0d0ee45f146eb6a0dfb9007e70ab3c", size = 1863499, upload-time = "2026-03-23T18:12:58.696Z" }, - { url = "https://files.pythonhosted.org/packages/33/bf/21b899792b08cae7a298551c68398a79e333697479ed311b3b067aab4bdc/torchvision-0.26.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:1c55dc8affbcc0eb2060fbabbe996ae9e5839b24bb6419777f17848945a411b1", size = 7767527, upload-time = "2026-03-23T18:12:44.348Z" }, - { url = "https://files.pythonhosted.org/packages/9a/45/57bbf9e216850d065e66dd31a50f57424b607f1d878ab8956e56a1f4e36b/torchvision-0.26.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:fd10b5f994c210f4f6d6761cf686f82d748554adf486cb0979770c3252868c8f", size = 7519925, upload-time = "2026-03-23T18:12:53.283Z" }, - { url = "https://files.pythonhosted.org/packages/10/58/ed8f7754299f3e91d6414b6dc09f62b3fa7c6e5d63dfe48d69ab81498a37/torchvision-0.26.0-cp311-cp311-win_amd64.whl", hash = "sha256:de6424b12887ad884f39a0ee446994ae3cd3b6a00a9cafe1bead85a031132af0", size = 3983834, upload-time = "2026-03-23T18:13:00.224Z" }, - { url = "https://files.pythonhosted.org/packages/ae/e7/56b47cc3b132aea90ccce22bcb8975dec688b002150012acc842846039d0/torchvision-0.26.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c409e1c3fdebec7a3834465086dbda8bf7680eff79abf7fd2f10c6b59520a7a4", size = 1863502, upload-time = "2026-03-23T18:12:57.326Z" }, - { url = "https://files.pythonhosted.org/packages/f4/ec/5c31c92c08b65662fe9604a4067ae8232582805949f11ddc042cebe818ed/torchvision-0.26.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:406557718e62fdf10f5706e88d8a5ec000f872da913bf629aab9297622585547", size = 7767944, upload-time = "2026-03-23T18:12:42.805Z" }, - { url = "https://files.pythonhosted.org/packages/f5/d8/cb6ccda1a1f35a6597645818641701207b3e8e13553e75fce5d86bac74b2/torchvision-0.26.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:d61a5abb6b42a0c0c311996c2ac4b83a94418a97182c83b055a2a4ae985e05aa", size = 7522205, upload-time = "2026-03-23T18:12:54.654Z" }, - { url = "https://files.pythonhosted.org/packages/1c/a9/c272623a0f735c35f0f6cd6dc74784d4f970e800cf063bb76687895a2ab9/torchvision-0.26.0-cp312-cp312-win_amd64.whl", hash = "sha256:7993c01648e7c61d191b018e84d38fe0825c8fcb2720cd0f37caf7ba14404aa1", size = 4255155, upload-time = "2026-03-23T18:12:32.652Z" }, - { url = "https://files.pythonhosted.org/packages/da/80/0762f77f53605d10c9477be39bb47722cc8e383bbbc2531471ce0e396c07/torchvision-0.26.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5d63dd43162691258b1b3529b9041bac7d54caa37eae0925f997108268cbf7c4", size = 1860809, upload-time = "2026-03-23T18:12:47.629Z" }, - { url = "https://files.pythonhosted.org/packages/e6/81/0b3e58d1478c660a5af4268713486b2df7203f35abd9195fea87348a5178/torchvision-0.26.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:a39c7a26538c41fda453f9a9692b5ff9b35a5437db1d94f3027f6f509c160eac", size = 7727494, upload-time = "2026-03-23T18:12:46.062Z" }, - { url = "https://files.pythonhosted.org/packages/b6/dc/d9ab5d29115aa05e12e30f1397a3eeae1d88a511241dc3bce48dc4342675/torchvision-0.26.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:b7e6213620bbf97742e5f79832f9e9d769e6cf0f744c5b53dad80b76db633691", size = 7521747, upload-time = "2026-03-23T18:12:36.815Z" }, - { url = "https://files.pythonhosted.org/packages/a9/1b/f1bc86a918c5f6feab1eeff11982e2060f4704332e96185463d27855bdf5/torchvision-0.26.0-cp313-cp313-win_amd64.whl", hash = "sha256:4280c35ec8cba1fcc8294fb87e136924708726864c379e4c54494797d86bc474", size = 4319880, upload-time = "2026-03-23T18:12:38.168Z" }, - { url = "https://files.pythonhosted.org/packages/66/28/b4ad0a723ed95b003454caffcc41894b34bd8379df340848cae2c33871de/torchvision-0.26.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:358fc4726d0c08615b6d83b3149854f11efb2a564ed1acb6fce882e151412d23", size = 1951973, upload-time = "2026-03-23T18:12:48.781Z" }, - { url = "https://files.pythonhosted.org/packages/71/e2/7a89096e6cf2f3336353b5338ba925e0addf9d8601920340e6bdf47e8eb3/torchvision-0.26.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:3daf9cc149cf3cdcbd4df9c59dae69ffca86c6823250442c3bbfd63fc2e26c61", size = 7728679, upload-time = "2026-03-23T18:12:26.196Z" }, - { url = "https://files.pythonhosted.org/packages/69/1d/4e1eebc17d18ce080a11dcf3df3f8f717f0efdfa00983f06e8ba79259f61/torchvision-0.26.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:82c3965eca27e86a316e31e4c3e5a16d353e0bcbe0ef8efa2e66502c54493c4b", size = 7609138, upload-time = "2026-03-23T18:12:35.327Z" }, - { url = "https://files.pythonhosted.org/packages/f3/a4/f1155e943ae5b32400d7000adc81c79bb0392b16ceb33bcf13e02e48cced/torchvision-0.26.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ebc043cc5a4f0bf22e7680806dbba37ffb19e70f6953bbb44ed1a90aeb5c9bea", size = 4248202, upload-time = "2026-03-23T18:12:41.423Z" }, + { name = "torch" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/50/ae/cbf727421eb73f1cf907fbe5788326a08f111b3f6b6ddca15426b53fec9a/torchvision-0.25.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a95c47abb817d4e90ea1a8e57bd0d728e3e6b533b3495ae77d84d883c4d11f56", size = 1874919, upload-time = "2026-01-21T16:27:47.617Z" }, + { url = "https://files.pythonhosted.org/packages/64/68/dc7a224f606d53ea09f9a85196a3921ec3a801b0b1d17e84c73392f0c029/torchvision-0.25.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:acc339aba4a858192998c2b91f635827e40d9c469d9cf1455bafdda6e4c28ea4", size = 2343220, upload-time = "2026-01-21T16:27:44.26Z" }, + { url = "https://files.pythonhosted.org/packages/f9/fa/8cce5ca7ffd4da95193232493703d20aa06303f37b119fd23a65df4f239a/torchvision-0.25.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:0d9a3f925a081dd2ebb0b791249b687c2ef2c2717d027946654607494b9b64b6", size = 8068106, upload-time = "2026-01-21T16:27:37.805Z" }, + { url = "https://files.pythonhosted.org/packages/8b/b9/a53bcf8f78f2cd89215e9ded70041765d50ef13bf301f9884ec6041a9421/torchvision-0.25.0-cp310-cp310-win_amd64.whl", hash = "sha256:b57430fbe9e9b697418a395041bb615124d9c007710a2712fda6e35fb310f264", size = 3697295, upload-time = "2026-01-21T16:27:36.574Z" }, + { url = "https://files.pythonhosted.org/packages/3e/be/c704bceaf11c4f6b19d64337a34a877fcdfe3bd68160a8c9ae9bea4a35a3/torchvision-0.25.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:db74a551946b75d19f9996c419a799ffdf6a223ecf17c656f90da011f1d75b20", size = 1874923, upload-time = "2026-01-21T16:27:46.574Z" }, + { url = "https://files.pythonhosted.org/packages/ae/e9/f143cd71232430de1f547ceab840f68c55e127d72558b1061a71d0b193cd/torchvision-0.25.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:f49964f96644dbac2506dffe1a0a7ec0f2bf8cf7a588c3319fed26e6329ffdf3", size = 2344808, upload-time = "2026-01-21T16:27:43.191Z" }, + { url = "https://files.pythonhosted.org/packages/43/ae/ad5d6165797de234c9658752acb4fce65b78a6a18d82efdf8367c940d8da/torchvision-0.25.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:153c0d2cbc34b7cf2da19d73450f24ba36d2b75ec9211b9962b5022fb9e4ecee", size = 8070752, upload-time = "2026-01-21T16:27:33.748Z" }, + { url = "https://files.pythonhosted.org/packages/23/19/55b28aecdc7f38df57b8eb55eb0b14a62b470ed8efeb22cdc74224df1d6a/torchvision-0.25.0-cp311-cp311-win_amd64.whl", hash = "sha256:ea580ffd6094cc01914ad32f8c8118174f18974629af905cea08cb6d5d48c7b7", size = 4038722, upload-time = "2026-01-21T16:27:41.355Z" }, + { url = "https://files.pythonhosted.org/packages/56/3a/6ea0d73f49a9bef38a1b3a92e8dd455cea58470985d25635beab93841748/torchvision-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c2abe430c90b1d5e552680037d68da4eb80a5852ebb1c811b2b89d299b10573b", size = 1874920, upload-time = "2026-01-21T16:27:45.348Z" }, + { url = "https://files.pythonhosted.org/packages/51/f8/c0e1ef27c66e15406fece94930e7d6feee4cb6374bbc02d945a630d6426e/torchvision-0.25.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b75deafa2dfea3e2c2a525559b04783515e3463f6e830cb71de0fb7ea36fe233", size = 2344556, upload-time = "2026-01-21T16:27:40.125Z" }, + { url = "https://files.pythonhosted.org/packages/68/2f/f24b039169db474e8688f649377de082a965fbf85daf4e46c44412f1d15a/torchvision-0.25.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:f25aa9e380865b11ea6e9d99d84df86b9cc959f1a007cd966fc6f1ab2ed0e248", size = 8072351, upload-time = "2026-01-21T16:27:21.074Z" }, + { url = "https://files.pythonhosted.org/packages/ad/16/8f650c2e288977cf0f8f85184b90ee56ed170a4919347fc74ee99286ed6f/torchvision-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:f9c55ae8d673ab493325d1267cbd285bb94d56f99626c00ac4644de32a59ede3", size = 4303059, upload-time = "2026-01-21T16:27:11.08Z" }, + { url = "https://files.pythonhosted.org/packages/f5/5b/1562a04a6a5a4cf8cf40016a0cdeda91ede75d6962cff7f809a85ae966a5/torchvision-0.25.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:24e11199e4d84ba9c5ee7825ebdf1cd37ce8deec225117f10243cae984ced3ec", size = 1874918, upload-time = "2026-01-21T16:27:39.02Z" }, + { url = "https://files.pythonhosted.org/packages/36/b1/3d6c42f62c272ce34fcce609bb8939bdf873dab5f1b798fd4e880255f129/torchvision-0.25.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5f271136d2d2c0b7a24c5671795c6e4fd8da4e0ea98aeb1041f62bc04c4370ef", size = 2309106, upload-time = "2026-01-21T16:27:30.624Z" }, + { url = "https://files.pythonhosted.org/packages/c7/60/59bb9c8b67cce356daeed4cb96a717caa4f69c9822f72e223a0eae7a9bd9/torchvision-0.25.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:855c0dc6d37f462482da7531c6788518baedca1e0847f3df42a911713acdfe52", size = 8071522, upload-time = "2026-01-21T16:27:29.392Z" }, + { url = "https://files.pythonhosted.org/packages/32/a5/9a9b1de0720f884ea50dbf9acb22cbe5312e51d7b8c4ac6ba9b51efd9bba/torchvision-0.25.0-cp313-cp313-win_amd64.whl", hash = "sha256:cef0196be31be421f6f462d1e9da1101be7332d91984caa6f8022e6c78a5877f", size = 4321911, upload-time = "2026-01-21T16:27:35.195Z" }, + { url = "https://files.pythonhosted.org/packages/52/99/dca81ed21ebaeff2b67cc9f815a20fdaa418b69f5f9ea4c6ed71721470db/torchvision-0.25.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a8f8061284395ce31bcd460f2169013382ccf411148ceb2ee38e718e9860f5a7", size = 1896209, upload-time = "2026-01-21T16:27:32.159Z" }, + { url = "https://files.pythonhosted.org/packages/28/cc/2103149761fdb4eaed58a53e8437b2d716d48f05174fab1d9fcf1e2a2244/torchvision-0.25.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:146d02c9876858420adf41f3189fe90e3d6a409cbfa65454c09f25fb33bf7266", size = 2310735, upload-time = "2026-01-21T16:27:22.327Z" }, + { url = "https://files.pythonhosted.org/packages/76/ad/f4c985ad52ddd3b22711c588501be1b330adaeaf6850317f66751711b78c/torchvision-0.25.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:c4d395cb2c4a2712f6eb93a34476cdf7aae74bb6ea2ea1917f858e96344b00aa", size = 8089557, upload-time = "2026-01-21T16:27:27.666Z" }, + { url = "https://files.pythonhosted.org/packages/63/cc/0ea68b5802e5e3c31f44b307e74947bad5a38cc655231d845534ed50ddb8/torchvision-0.25.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5e6b449e9fa7d642142c0e27c41e5a43b508d57ed8e79b7c0a0c28652da8678c", size = 4344260, upload-time = "2026-01-21T16:27:17.018Z" }, ] [[package]] @@ -4029,6 +4239,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" }, ] +[[package]] +name = "triton" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/f7/f1c9d3424ab199ac53c2da567b859bcddbb9c9e7154805119f8bd95ec36f/triton-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a6550fae429e0667e397e5de64b332d1e5695b73650ee75a6146e2e902770bea", size = 188105201, upload-time = "2026-01-20T16:00:29.272Z" }, + { url = "https://files.pythonhosted.org/packages/e0/12/b05ba554d2c623bffa59922b94b0775673de251f468a9609bc9e45de95e9/triton-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8e323d608e3a9bfcc2d9efcc90ceefb764a82b99dea12a86d643c72539ad5d3", size = 188214640, upload-time = "2026-01-20T16:00:35.869Z" }, + { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, + { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, + { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, +] + [[package]] name = "typer" version = "0.24.1" From feec81ad2b049d9b5f6c0dd3d5f318fbef150aa4 Mon Sep 17 00:00:00 2001 From: jingyu-ml <108295447+jingyu-ml@users.noreply.github.com> Date: Fri, 17 Apr 2026 23:54:19 -0700 Subject: [PATCH 14/30] Add the Skip softmax for diffusion (#1166) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: new feature, new example ## Summary - Add skip-softmax sparse attention (BLASST) for diffusion models via dedicated Triton kernels — an inference kernel with tile skipping and a calibration kernel with vectorized multi-threshold sparsity measurement - Add `triton_skip_softmax` method with exponential model calibration (`scale_factor = a * exp(b * sparsity)`) and log-space fitting for diffusion models - Add Triton kernel backends for diffusers and LTX attention dispatch - Fix calibration to skip RULER dataset generation when user provides their own `forward_loop` (required for non-LLM models) ## Changes ### Triton kernels (`modelopt/torch/kernels/triton_fa.py`) - **`_attn_fwd`**: Forward kernel with optional tile skipping — tiles whose max attention score is far below the running softmax max are skipped entirely (no V load, no softmax, no accumulation). Runtime sparsity measurement via atomic counters. - **`_attn_fwd_calibrate`**: Calibration kernel that computes full attention while measuring how many tiles would be skipped at each of N thresholds simultaneously. Uses per-program output buffers (zero atomic contention) and vectorized multi-threshold comparison. - **`attention()`** / **`attention_calibrate()`**: Python wrappers for inference and calibration kernels. ### Kernel backends (`modelopt/torch/sparsity/attention_sparsity/kernels/`) - **`diffusers_triton_attention.py`**: Registers `modelopt_triton` backend in diffusers' attention dispatch. Handles [B, S, H, D] → varlen layout conversion, calibration/inference mode switching, thread-local configuration, and counter accumulation. - **`ltx_triton_attention.py`**: Patches `ltx_core.Attention` modules for Triton dispatch with the same calibration/inference modes. ### Method (`modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py`) - `TritonSkipSoftmaxMethod`: Context managers for calibration (→ calibration kernel) and inference (→ forward kernel with tile skipping). Three threshold priority levels: raw threshold > calibrated scale_factor > static threshold. ### Calibration (`modelopt/torch/sparsity/attention_sparsity/calibration/`) - **`calibrator.py`**: `DynamicThresholdCalibrator` with `fit_logspace` option — fits exponential model in log space (minimizes relative error) for diffusion models where scale_factors span many orders of magnitude. Records observed sparsity range for extrapolation warnings. - **`calibrate.py`**: Skips RULER dataset when `forward_loop` is provided; passes `fit_logspace` through from config. ### Config & conversion - **`config.py`**: `CalibrationConfig.fit_logspace` field (default False, recommended True for diffusion models). `skip_softmax_raw_threshold` field for direct threshold mode. - **`conversion.py`**: Auto-registers diffusers/LTX Triton backends on `sparsify()`. Updated summary display. ### Example - **`wan22_skip_softmax.py`**: End-to-end example for WAN 2.2 5B/14B with baseline, raw-threshold, and calibrated modes. Supports runtime sparsity reporting. ## Threshold modes | Mode | How it works | Use case | |------|-------------|----------| | **Raw threshold** (`--raw-threshold -0.7`) | Passed directly to kernel as `skip_threshold_log2` | Quick testing, sweeps | | **Calibrated** (`--calibrate --target-sparsity 0.5`) | `scale_factor = a * exp(b * target)`, then `threshold = scale_factor / seq_k` at runtime | Production use with seqlen adaptation | | **Static** (default `skip_softmax_threshold=0.1`) | `log2(lambda) * sm_scale` | Fallback | ## Usage ```bash # Fixed raw threshold (no calibration) python examples/diffusers/sparsity/wan22_skip_softmax.py \ --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ --raw-threshold -0.7 \ --prompt "A cat playing piano" --output out.mp4 # With calibration (log-space fit for diffusion models) python examples/diffusers/sparsity/wan22_skip_softmax.py \ --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ --calibrate --target-sparsity 0.5 \ --prompt "A cat playing piano" --output out.mp4 # Dense baseline for comparison python examples/diffusers/sparsity/wan22_skip_softmax.py \ --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ --baseline \ --prompt "A cat playing piano" --output baseline.mp4 ``` ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ ### Additional Information ## Summary by CodeRabbit ## Release Notes * **New Features** * Added skip-softmax sparse attention support for Diffusers models, enabling efficient video generation * Added support for both eager and Triton attention backends for sparse attention * Added new example script for Wan 2.2 text-to-video generation with sparse attention optimization * **Documentation** * Updated documentation with sparse attention configuration guide and usage examples * **Tests** * Added comprehensive unit tests for kernel backend registration and skip-softmax functionality --------- Signed-off-by: Jingyu Xin --- .github/codecov.yml | 12 + .github/workflows/example_tests.yml | 2 +- examples/diffusers/README.md | 54 ++ examples/diffusers/sparsity/README.md | 76 +++ .../diffusers/sparsity/wan22_skip_softmax.py | 485 ++++++++++++++++++ modelopt/torch/export/diffusers_utils.py | 11 + modelopt/torch/kernels/__init__.py | 4 + modelopt/torch/kernels/triton_fa.py | 333 +++++++++++- .../src/conv/bench_implicit_gemm.py | 208 ++++++++ .../calibration/calibrate.py | 43 +- .../calibration/calibrator.py | 108 +++- .../sparsity/attention_sparsity/config.py | 20 + .../sparsity/attention_sparsity/conversion.py | 36 ++ .../attention_sparsity/kernels/__init__.py | 38 +- .../kernels/diffusers_triton_attention.py | 251 +++++++++ .../kernels/ltx_triton_attention.py | 190 +++++++ .../methods/flash_skip_softmax.py | 16 +- .../attention_sparsity/methods/registry.py | 4 + .../methods/triton_skip_softmax.py | 250 ++++++++- .../attention_sparsity/plugins/huggingface.py | 12 +- .../attention_sparsity/stats_manager.py | 30 +- tests/_test_utils/torch/diffusers_models.py | 93 ++++ .../diffusers_sparsity/test_sparsity.py | 104 ++++ .../test_diffusers_triton_attention.py | 181 +++++++ .../test_triton_fa_calibrate.py | 287 +++++++++++ .../test_wan22_skip_softmax.py | 280 ++++++++++ tests/unit/torch/kernels/test_triton_fa.py | 39 ++ .../test_calibrator_fitting.py | 183 +++++++ .../test_flash_skip_softmax.py | 124 +++++ .../test_kernel_backends.py | 131 +++++ .../test_ltx_triton_attention.py | 126 +++++ .../test_sparse_attention_calibration.py | 173 +++++++ .../test_sparse_attention_conversion.py | 128 +++++ .../test_triton_skip_softmax.py | 222 ++++++++ tox.ini | 2 + 35 files changed, 4180 insertions(+), 76 deletions(-) create mode 100644 examples/diffusers/sparsity/README.md create mode 100644 examples/diffusers/sparsity/wan22_skip_softmax.py create mode 100644 modelopt/torch/quantization/src/conv/bench_implicit_gemm.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/kernels/diffusers_triton_attention.py create mode 100644 modelopt/torch/sparsity/attention_sparsity/kernels/ltx_triton_attention.py create mode 100644 tests/examples/diffusers_sparsity/test_sparsity.py create mode 100644 tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py create mode 100644 tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.py create mode 100644 tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py create mode 100644 tests/unit/torch/kernels/test_triton_fa.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_calibrator_fitting.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_kernel_backends.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_ltx_triton_attention.py create mode 100644 tests/unit/torch/sparsity/attention_sparsity/test_triton_skip_softmax.py diff --git a/.github/codecov.yml b/.github/codecov.yml index 24756fdcbb2..b4ac8367690 100644 --- a/.github/codecov.yml +++ b/.github/codecov.yml @@ -11,3 +11,15 @@ coverage: target: auto threshold: 1% # Allow atmost 1% coverage drop from main branch. patch: false + +# Exclude GPU-only Triton kernel files from ALL codecov calculations (project +# and patch checks, all flags). Rationale: these files are dominated by +# @triton.jit kernel bodies that CPU unit tests cannot exercise. GPU tests +# cover them end-to-end (see tests/gpu/torch/sparsity/attention_sparsity/) but +# the `gpu`-flag upload may race with the PR status check, so relying on flag +# combination alone leaves the project check flaky. Dropping these files here +# makes the check deterministic — local `pytest --cov` and GPU runs still +# measure them; only the codecov PR status ignores them. +ignore: + - "modelopt/torch/kernels/triton_fa.py" + - "modelopt/torch/kernels/hf_triton_attention.py" diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index e316618852a..73f5c1fa9b7 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -63,7 +63,7 @@ jobs: strategy: &torch_strategy fail-fast: false matrix: - example: [llm_distill, llm_qat, llm_sparsity] + example: [llm_distill, llm_qat, llm_sparsity, diffusers_sparsity] include: - example: speculative_decoding docker_image: "26.01" diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 6e7259ee52d..84f248bfb15 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -13,6 +13,7 @@ Cache Diffusion is a technique that reuses cached outputs from previous diffusio | Pre-Requisites | Required & optional packages to use this technique | \[[Link](#pre-requisites)\] | | | Getting Started | Learn how to optimize your models using quantization/cache diffusion to reduce precision and improve inference efficiency | \[[Link](#getting-started)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | | Support Matrix | View the support matrix to see quantization/cahce diffusion compatibility and feature availability across different models | \[[Link](#support-matrix)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | +| Sparse Attention (Skip-Softmax) | Skip-softmax sparse attention for diffusion models | \[[Link](#sparse-attention-skip-softmax)\] | | | Cache Diffusion | Caching technique to accelerate inference without compromising quality | \[[Link](#cache-diffusion)\] | | | Post Training Quantization (PTQ) | Example scripts on how to run PTQ on diffusion models | \[[Link](#post-training-quantization-ptq)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | | Quantization Aware Training (QAT) | Example scripts on how to run QAT on diffusion models | \[[Link](#quantization-aware-training-qat)\] | \[[docs](https://nvidia.github.io/Model-Optimizer/guides/1_quantization.html)\] | @@ -290,6 +291,59 @@ mto.restore(pipe.unet, your_quantized_ckpt) By following these steps, your PEFT LoRA model should be efficiently quantized using ModelOpt, ready for deployment while maximizing performance. +## Sparse Attention (Skip-Softmax) + +Skip-softmax sparse attention skips KV tiles whose attention scores are negligible during the softmax computation, reducing FLOPs without retraining. An exponential model (`scale_factor = a * exp(b * target_sparsity)`) is calibrated once, then the target sparsity can be adjusted at runtime without recalibration. + +### Getting Started + +```python +import modelopt.torch.sparsity.attention_sparsity as mtsa + +# 1. Define config with calibration +config = { + "sparse_cfg": { + "calibration": { + "target_sparse_ratio": {"prefill": 0.5}, + }, + "*.attn1": { + "method": "triton_skip_softmax", + "backend": "triton", + "is_causal": False, + "collect_stats": True, + "enable": True, + }, + "*.attn2": {"enable": False}, + "default": {"enable": False}, + }, +} + +# 2. Provide a calibration forward loop +def forward_loop(model): + pipeline(prompt="a cat", num_frames=81, num_inference_steps=40, ...) + +# 3. Sparsify + calibrate +mtsa.sparsify(transformer, config, forward_loop=forward_loop) + +# 4. Generate as usual — sparsity is applied automatically +output = pipeline(prompt="a dog on the beach", ...) +``` + +### Example Scripts + +#### Wan 2.2 [Script](./sparsity/wan22_skip_softmax.py) + +The 14B model automatically sparsifies both `transformer` and `transformer_2`. + +```bash + +# 5B/14B model +python sparsity/wan22_skip_softmax.py \ + --model-path Wan-AI/Wan2.2-T2V-A14B-Diffusers|Wan-AI/Wan2.2-TI2V-5B-Diffusers \ + --calibrate --target-sparsity 0.5 --calib-size 4 \ + --prompt "A sunset over mountains" --output out.mp4 +``` + ## Cache Diffusion Cache Diffusion methods, such as [DeepCache](https://arxiv.org/abs/2312.00858), [Block Caching](https://arxiv.org/abs/2312.03209) and [T-Gate](https://arxiv.org/abs/2404.02747), optimize performance by reusing cached outputs from previous steps instead of recalculating them. This **training-free** caching approach is compatible with a variety of models, like **DiT** and **UNet**, enabling considerable acceleration without compromising quality. diff --git a/examples/diffusers/sparsity/README.md b/examples/diffusers/sparsity/README.md new file mode 100644 index 00000000000..dc44fbcd173 --- /dev/null +++ b/examples/diffusers/sparsity/README.md @@ -0,0 +1,76 @@ +# Skip-Softmax Sparse Attention for Diffusion Models + +> [!WARNING] +> **Third-Party License Notice — LTX-2** +> +> LTX-2 packages (`ltx-core`, `ltx-pipelines`, `ltx-trainer`) are third-party dependencies +> developed and provided by [Lightricks](https://github.com/Lightricks/LTX-2). They are +> **NOT** covered by the Apache 2.0 license governing NVIDIA Model Optimizer. +> +> You **MUST** comply with the +> [LTX Community License Agreement](https://github.com/Lightricks/LTX-2/blob/main/LICENSE) +> when installing and using LTX-2 with NVIDIA Model Optimizer. Any derivative models or +> fine-tuned weights produced from LTX-2 (including quantized, distilled, or sparsified +> checkpoints) remain subject to the LTX Community License Agreement, not Apache 2.0. + +Skip-softmax sparse attention (BLASST, ) skips KV +tiles whose attention scores are negligible during the FlashAttention computation, +reducing FLOPs without retraining. + +Two modes are supported: +- **Fixed raw threshold** — pass a log2-space threshold directly to the Triton + kernel. No calibration needed. Good for quick testing and sweeps. +- **Calibrated threshold** — an exponential model + (`scale_factor = a * exp(b * target_sparsity)`) is calibrated once via the + Triton calibration kernel, then the target sparsity can be adjusted at runtime + without recalibration. Log-space fitting (`fit_logspace=True`) is recommended + for diffusion models where scale_factors span many orders of magnitude. + +## Supported Models + +| Model | Script | Notes | +|-------|--------|-------| +| WAN 2.2 5B | `wan22_skip_softmax.py` | Single transformer, self-attention only | +| WAN 2.2 14B | `wan22_skip_softmax.py` | Dual transformer (auto-detected) | +| LTX-2 | (coming soon) | Via `ltx_triton_attention.py` backend | + +## Quick Start + +```bash +# Fixed raw threshold (no calibration, fast) +python wan22_skip_softmax.py \ + --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ + --raw-threshold -0.7 \ + --prompt "A cat playing piano" --output out.mp4 + +# With calibration +python wan22_skip_softmax.py \ + --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ + --calibrate --target-sparsity 0.5 \ + --prompt "A cat playing piano" --output out.mp4 + +# Dense baseline (no sparsity, for comparison) +python wan22_skip_softmax.py \ + --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ + --baseline \ + --prompt "A cat playing piano" --output baseline.mp4 + +# Report runtime sparsity (per-layer tile skip ratios) +python wan22_skip_softmax.py \ + --model-path /path/to/Wan2.2-T2V-A14B-Diffusers \ + --raw-threshold -0.7 --report-avg-sparsity \ + --prompt "A cat playing piano" --output out.mp4 +``` + +## Threshold Modes + +| Mode | How threshold reaches the kernel | Use case | +|------|----------------------------------|----------| +| **Raw threshold** (`--raw-threshold -0.7`) | Passed directly as `skip_threshold_log2` — no conversion | Quick testing, sweeps | +| **Calibrated** (`--calibrate --target-sparsity 0.5`) | `scale_factor = a * exp(b * target)`, then backend computes `threshold = scale_factor / seq_k`, then kernel converts `log2(threshold) * sm_scale` | Production use with automatic seqlen adaptation | +| **Static lambda** (default `skip_softmax_threshold=0.1`) | `log2(lambda) * sm_scale` | Fallback when neither raw nor calibrated | + +## Known Issues + +- **14B dual transformer calibration**: Transformers are calibrated sequentially — transformer_2's calibration runs while transformer_1 is already sparsified, introducing asymmetric calibration conditions. +- **Minimum achievable sparsity**: Even the strictest threshold may yield 30-40% sparsity on diffusion models (many tiles are inherently negligible). Targets below this floor cause extrapolation; an inference-time warning is emitted. diff --git a/examples/diffusers/sparsity/wan22_skip_softmax.py b/examples/diffusers/sparsity/wan22_skip_softmax.py new file mode 100644 index 00000000000..e335451e2b5 --- /dev/null +++ b/examples/diffusers/sparsity/wan22_skip_softmax.py @@ -0,0 +1,485 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Wan 2.2 inference with skip-softmax sparse attention. + +This example applies skip-softmax sparse attention to the Wan 2.2 video +generation model (text-to-video). Four modes are supported: + +1. **Baseline** — pass ``--baseline`` for dense inference (default diffusers backend). +2. **Triton baseline** — pass ``--triton-baseline`` for dense Triton FA kernel + (no skip-softmax, same kernel as sparse runs for apples-to-apples comparison). +3. **Fixed raw threshold** — pass ``--raw-threshold`` to supply a log2-space + threshold directly to the Triton kernel. No calibration data is needed. +4. **Calibrated threshold** — pass ``--calibrate`` to run exponential-model + calibration (``scale_factor = a * exp(b * target_sparsity)``). + +During calibration, ``triton_skip_softmax`` with the Triton calibration kernel +collects sparsity statistics across multiple threshold trials. The fitted +exponential model then allows runtime control of the target sparsity ratio +without recalibration. + +The Wan 2.2 5B model has 40 transformer blocks with self-attention (attn1) +and cross-attention (attn2). Only self-attention is sparsified. + +Usage:: + + # Baseline (dense, no sparsity) + python wan22_skip_softmax.py --baseline --prompt "A cat playing piano" \\ + --output baseline.mp4 + + # Fixed raw threshold (no calibration needed) + python wan22_skip_softmax.py --raw-threshold -5.0 --report-avg-sparsity \\ + --prompt "A cat playing piano" --output out.mp4 + + # With calibration + python wan22_skip_softmax.py --calibrate --target-sparsity 0.25 \\ + --report-avg-sparsity --prompt "A cat playing piano" --output out.mp4 +""" + +import argparse +import gc +import os + +import torch +from datasets import load_dataset +from diffusers import AutoencoderKLWan, WanPipeline +from diffusers.utils import export_to_video + +import modelopt.torch.sparsity.attention_sparsity as mtsa +from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule + +DEFAULT_MODEL_PATH = os.environ.get("WAN22_MODEL_PATH", "Wan-AI/Wan2.2-TI2V-5B-Diffusers") + +# fmt: off +# ruff: noqa: RUF001 +DEFAULT_NEGATIVE_PROMPT = ( # Official Wan 2.2 negative prompt (Chinese) + "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，" + "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，" + "画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，" + "杂乱的背景，三条腿，背景人很多，倒着走" +) +# fmt: on + +# Default threshold trials for calibration +DEFAULT_THRESHOLD_TRIALS = [ + 1e-12, + 1e-10, + 1e-8, + 1e-6, + 5e-6, + 1e-5, + 5e-5, + 1e-4, + 5e-4, + 1e-3, + 5e-3, + 1e-2, + 2e-2, + 5e-2, + 1e-1, + 2e-1, + 3e-1, + 5e-1, + 7e-1, + 8e-1, + 9e-1, + 9.9e-1, +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Wan 2.2 video generation with skip-softmax sparse attention" + ) + parser.add_argument( + "--prompt", + type=str, + default=None, + help="Text prompt for generation (optional, skips generation if not set)", + ) + parser.add_argument("--output", type=str, default="output.mp4", help="Output video path") + parser.add_argument( + "--model-path", type=str, default=DEFAULT_MODEL_PATH, help="Wan 2.2 model path or HF ID" + ) + parser.add_argument( + "--num-frames", type=int, default=81, help="Number of frames (must be 4k+1)" + ) + parser.add_argument("--height", type=int, default=480, help="Video height") + parser.add_argument("--width", type=int, default=832, help="Video width") + parser.add_argument("--num-steps", type=int, default=40, help="Number of inference steps") + parser.add_argument( + "--guidance-scale", type=float, default=4.0, help="Classifier-free guidance scale" + ) + parser.add_argument( + "--guidance-scale-2", + type=float, + default=3.0, + help="Second guidance scale for 14B dual-transformer model (ignored by 5B)", + ) + parser.add_argument( + "--negative-prompt", + type=str, + default=DEFAULT_NEGATIVE_PROMPT, + help="Negative prompt", + ) + parser.add_argument("--seed", type=int, default=42, help="Random seed") + + # Sparse attention options + parser.add_argument( + "--baseline", + action="store_true", + help="Run dense inference with default diffusers backend (no sparsity)", + ) + parser.add_argument( + "--triton-baseline", + action="store_true", + help="Run dense inference with Triton FA kernel (no skip-softmax, " + "apples-to-apples comparison with sparse runs)", + ) + parser.add_argument( + "--raw-threshold", + type=float, + default=None, + help="Raw skip_threshold_log2 value passed directly to the Triton kernel. " + "Negative values (e.g., -5.0 means tile must be within 5 units of running max). " + "Bypasses calibration and lambda conversion. Typical range: -1 to -30.", + ) + parser.add_argument( + "--skip-first-last", + type=int, + default=2, + help="Number of first/last transformer layers to keep dense (default: 2)", + ) + parser.add_argument( + "--report-avg-sparsity", + action="store_true", + help="Report per-layer and overall average tile sparsity after generation", + ) + + # Calibration options + parser.add_argument( + "--calibrate", + action="store_true", + help="Calibrate threshold via exponential model (recommended)", + ) + parser.add_argument( + "--target-sparsity", + type=float, + default=0.5, + help="Target sparsity ratio for calibration (0.0-1.0)", + ) + parser.add_argument( + "--calib-steps", + type=int, + default=40, + help="Inference steps for calibration", + ) + parser.add_argument( + "--calib-frames", + type=int, + default=151, + help="Number of frames for calibration", + ) + parser.add_argument( + "--calib-size", + type=int, + default=4, + help="Number of calibration prompts from OpenVid-1M dataset", + ) + return parser.parse_args() + + +def build_pipeline(model_path: str) -> WanPipeline: + """Build the Wan 2.2 text-to-video pipeline.""" + vae = AutoencoderKLWan.from_pretrained(model_path, subfolder="vae", torch_dtype=torch.float32) + pipe = WanPipeline.from_pretrained(model_path, vae=vae, torch_dtype=torch.bfloat16) + pipe.to("cuda") + return pipe + + +def build_sparse_config(args: argparse.Namespace, num_blocks: int) -> dict: + """Build sparse attention config from CLI args. + + Two modes: + - **Raw threshold**: ``--raw-threshold`` sets ``skip_softmax_raw_threshold`` + directly on the Triton kernel — no calibration needed. + - **Calibrated**: ``--calibrate`` collects multi-threshold sparsity statistics + via the Triton calibration kernel, then fits an exponential model: + ``scale_factor = a * exp(b * sparsity)``. + """ + attn_cfg: dict = { + "method": "triton_skip_softmax", + "skip_softmax_threshold": 0.0 if args.triton_baseline else 0.1, + "backend": "triton", + "is_causal": False, # Diffusion = bidirectional attention + "collect_stats": True, + "enable": True, + } + + # Raw threshold bypasses calibration and lambda conversion + if args.raw_threshold is not None: + attn_cfg["skip_softmax_raw_threshold"] = args.raw_threshold + + sparse_cfg: dict = { + "*.attn1*": attn_cfg, # Self-attention only + "*.attn2*": {"enable": False}, # Text cross-attention + "default": {"enable": False}, + } + + # Keep first/last N layers dense for quality + for i in range(args.skip_first_last): + sparse_cfg[f"*blocks.{i}.attn*"] = {"enable": False} + sparse_cfg[f"*blocks.{num_blocks - 1 - i}.attn*"] = {"enable": False} + + config: dict = {"sparse_cfg": sparse_cfg} + + # Add calibration config only when calibrating (not with raw threshold) + if args.calibrate and args.raw_threshold is None: + sparse_cfg["calibration"] = { + "target_sparse_ratio": {"prefill": args.target_sparsity}, + "threshold_trials": DEFAULT_THRESHOLD_TRIALS, + "fit_logspace": True, + } + + return config + + +def load_calib_prompts(calib_size: int) -> list[str]: + """Load calibration prompts from OpenVid-1M dataset.""" + dataset = load_dataset("nkp37/OpenVid-1M", split="train") + prompts = list(dataset["caption"][:calib_size]) + print(f"Loaded {len(prompts)} calibration prompts from OpenVid-1M") + return prompts + + +def build_calibration_forward_loop( + pipe: WanPipeline, + calib_size: int = 4, + num_steps: int = 40, + num_frames: int = 151, + height: int = 480, + width: int = 832, + seed: int = 42, + guidance_scale: float = 4.0, + guidance_scale_2: float | None = 3.0, + negative_prompt: str = "", +): + """Build a forward loop for exponential model calibration. + + Uses prompts from OpenVid-1M dataset (same as quantization examples). + Each prompt is run individually (batch_size=1). + """ + calib_prompts = load_calib_prompts(calib_size) + + def forward_loop(model): + for i, prompt in enumerate(calib_prompts): + print(f"Calibration [{i + 1}/{len(calib_prompts)}]: {prompt[:60]}...") + kw: dict = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "num_frames": num_frames, + "height": height, + "width": width, + "num_inference_steps": num_steps, + "guidance_scale": guidance_scale, + "generator": torch.Generator(device="cuda").manual_seed(seed), + } + if guidance_scale_2 is not None: + kw["guidance_scale_2"] = guidance_scale_2 + pipe(**kw) + + return forward_loop + + +def enable_sparsity_measurement(model: torch.nn.Module) -> None: + """Enable runtime sparsity measurement on all sparse attention modules.""" + for _name, module in model.named_modules(): + if isinstance(module, SparseAttentionModule) and module.is_enabled: + method = module._sparse_method_instance + if hasattr(method, "enable_measure_sparsity"): + method.reset_sparsity_counters() + method.enable_measure_sparsity(True) + + +def print_sparsity_summary(model: torch.nn.Module) -> None: + """Print per-module sparsity statistics including runtime kernel counters.""" + enabled, disabled = [], [] + for name, module in model.named_modules(): + if isinstance(module, SparseAttentionModule): + if module.is_enabled: + enabled.append((name, module)) + else: + disabled.append(name) + + print(f"\nSparse attention: {len(enabled)} enabled, {len(disabled)} disabled") + for name, module in enabled: + info = module.get_threshold_info() + print(f" {name}: {info}") + + +def print_runtime_sparsity(model: torch.nn.Module) -> None: + """Print runtime tile sparsity measured via kernel atomic counters.""" + total_all = 0 + skipped_all = 0 + per_module: list[tuple[str, int, int]] = [] + + for name, module in model.named_modules(): + if isinstance(module, SparseAttentionModule) and module.is_enabled: + method = module._sparse_method_instance + if hasattr(method, "get_sparsity_counters"): + total, skipped = method.get_sparsity_counters() + if total > 0: + per_module.append((name, total, skipped)) + total_all += total + skipped_all += skipped + + if total_all == 0: + print("\nNo runtime sparsity data collected.") + return + + print("\n" + "=" * 70) + print("Runtime tile sparsity (measured via kernel atomic counters)") + print("=" * 70) + for name, total, skipped in per_module: + ratio = skipped / total + print(f" {name}: {skipped:,}/{total:,} tiles skipped ({ratio:.1%})") + ratio_all = skipped_all / total_all + print("-" * 70) + print(f" Overall: {skipped_all:,}/{total_all:,} tiles skipped ({ratio_all:.1%})") + print("=" * 70) + + +def _get_num_blocks(transformer: torch.nn.Module) -> int: + """Count transformer blocks by looking for *.blocks.N.* submodules.""" + max_idx = -1 + for name, _ in transformer.named_modules(): + parts = name.split(".") + for i, part in enumerate(parts): + if part == "blocks" and i + 1 < len(parts) and parts[i + 1].isdigit(): + max_idx = max(max_idx, int(parts[i + 1])) + if max_idx < 0: + raise ValueError( + "Could not detect transformer blocks (expected submodules matching *.blocks.N.*). " + "Check that the model architecture uses 'blocks' as the layer container name." + ) + return max_idx + 1 + + +def main() -> None: + args = parse_args() + + # ---- Build pipeline ---- + print(f"Loading Wan 2.2 from {args.model_path}...") + pipe = build_pipeline(args.model_path) + + # ---- Collect transformers ---- + # Wan 2.2 5B has one transformer; 14B has two (transformer + transformer_2) + transformers = [] + if pipe.transformer is not None: + transformers.append(("transformer", pipe.transformer)) + if getattr(pipe, "transformer_2", None) is not None: + transformers.append(("transformer_2", pipe.transformer_2)) + is_14b = len(transformers) > 1 + + # ---- Sparsify (unless baseline) ---- + if args.baseline: + print("Baseline mode: running dense inference (default diffusers backend)") + elif args.triton_baseline: + print("Triton baseline: dense Triton FA kernel (no skip-softmax)") + for name, transformer in transformers: + num_blocks = _get_num_blocks(transformer) + print(f"Applying Triton backend to {name} ({num_blocks} blocks)...") + config = build_sparse_config(args, num_blocks=num_blocks) + mtsa.sparsify(transformer, config, forward_loop=None) + else: + # Build calibration forward loop if needed + forward_loop = None + if args.raw_threshold is not None: + print(f"Using fixed raw threshold: {args.raw_threshold} (skipping calibration)") + if args.calibrate: + print("Warning: --calibrate is ignored when --raw-threshold is set") + elif args.calibrate: + forward_loop = build_calibration_forward_loop( + pipe, + calib_size=args.calib_size, + num_steps=args.calib_steps, + num_frames=args.calib_frames, + height=args.height, + width=args.width, + seed=args.seed, + guidance_scale=args.guidance_scale, + guidance_scale_2=args.guidance_scale_2 if is_14b else None, + negative_prompt=args.negative_prompt, + ) + else: + print( + "Warning: neither --baseline, --raw-threshold, nor --calibrate specified; " + "using default static threshold" + ) + + for name, transformer in transformers: + num_blocks = _get_num_blocks(transformer) + print(f"Applying skip-softmax to {name} ({num_blocks} blocks)...") + config = build_sparse_config(args, num_blocks=num_blocks) + mtsa.sparsify(transformer, config, forward_loop=forward_loop) + + # ---- Free calibration memory before inference ---- + if not args.baseline and not args.triton_baseline and forward_loop is not None: + gc.collect() + torch.cuda.empty_cache() + print("Cleared CUDA cache after calibration") + + # ---- Generate (optional) ---- + if args.prompt: + # Enable runtime sparsity measurement before generation + if args.report_avg_sparsity and not args.baseline: + for _name, transformer in transformers: + enable_sparsity_measurement(transformer) + + print(f"Generating: {args.prompt[:80]}...") + pipe_kwargs: dict = { + "prompt": args.prompt, + "negative_prompt": args.negative_prompt, + "num_frames": args.num_frames, + "height": args.height, + "width": args.width, + "num_inference_steps": args.num_steps, + "guidance_scale": args.guidance_scale, + "generator": torch.Generator(device="cuda").manual_seed(args.seed), + } + if is_14b and args.guidance_scale_2 is not None: + pipe_kwargs["guidance_scale_2"] = args.guidance_scale_2 + output = pipe(**pipe_kwargs) + + try: + export_to_video(output.frames[0], args.output, fps=16) + print(f"Saved to {args.output}") + except ImportError as exc: + # Minimal CI envs may lack opencv/imageio — skip export silently, + # the inference itself already ran successfully. + print(f"Video export skipped (no opencv/imageio backend): {exc}") + + # ---- Print stats ---- + if not args.baseline: + for name, transformer in transformers: + print(f"\n{name}:") + print_sparsity_summary(transformer) + if args.report_avg_sparsity: + print_runtime_sparsity(transformer) + + +if __name__ == "__main__": + main() diff --git a/modelopt/torch/export/diffusers_utils.py b/modelopt/torch/export/diffusers_utils.py index 2c99db79835..9c4bdd06ff4 100644 --- a/modelopt/torch/export/diffusers_utils.py +++ b/modelopt/torch/export/diffusers_utils.py @@ -46,6 +46,17 @@ try: # optional for LTX-2 export paths from ltx_pipelines.ti2vid_two_stages import TI2VidTwoStagesPipeline as _TI2VidTwoStagesPipeline + warnings.warn( + "LTX-2 packages (ltx-core, ltx-pipelines, ltx-trainer) are provided by Lightricks " + "and are NOT covered by the Apache 2.0 license governing NVIDIA Model Optimizer. " + "You MUST comply with the LTX Community License Agreement when installing and using " + "LTX-2 with NVIDIA Model Optimizer. Any derivative models or fine-tuned weights from " + "LTX-2 (including quantized or distilled checkpoints) remain subject to the LTX " + "Community License Agreement, not Apache 2.0. " + "See: https://github.com/Lightricks/LTX-2/blob/main/LICENSE", + UserWarning, + stacklevel=2, + ) TI2VidTwoStagesPipeline = _TI2VidTwoStagesPipeline except Exception: # pragma: no cover TI2VidTwoStagesPipeline = None diff --git a/modelopt/torch/kernels/__init__.py b/modelopt/torch/kernels/__init__.py index 24d27a1ba2f..fa07b06e20c 100644 --- a/modelopt/torch/kernels/__init__.py +++ b/modelopt/torch/kernels/__init__.py @@ -21,6 +21,7 @@ IS_AVAILABLE = False attention = None +attention_calibrate = None register_triton_attention = None if torch.cuda.is_available(): @@ -32,8 +33,10 @@ ), ): from .triton_fa import attention as _attention + from .triton_fa import attention_calibrate as _attention_calibrate attention = _attention + attention_calibrate = _attention_calibrate IS_AVAILABLE = True from .hf_triton_attention import register_triton_attention as _register_triton_attention @@ -42,5 +45,6 @@ __all__ = [ "IS_AVAILABLE", "attention", + "attention_calibrate", "register_triton_attention", ] diff --git a/modelopt/torch/kernels/triton_fa.py b/modelopt/torch/kernels/triton_fa.py index 8d3b11f1af5..8044383889f 100644 --- a/modelopt/torch/kernels/triton_fa.py +++ b/modelopt/torch/kernels/triton_fa.py @@ -252,6 +252,9 @@ def _attn_fwd( DENSE_WINDOW_SIZE: tl.constexpr = 64, # Tokens near diagonal kept dense (absolute, BLOCK_N-independent) APPLY_SKIP_SOFTMAX: tl.constexpr = False, # Skip KV tiles with negligible scores SKIP_THRESHOLD_LOG2: tl.constexpr = 0.0, # log2(lambda) * sm_scale, pre-scaled for comparison on scaled scores + Sparsity_total=None, # Optional int64 scalar for counting total tiles (atomic) + Sparsity_skipped=None, # Optional int64 scalar for counting skipped tiles (atomic) + MEASURE_SPARSITY: tl.constexpr = False, # When True, count total/skipped tiles via atomic adds ): # --- Grid: (batch, num_q_heads, num_q_tiles) --- # Example: batch=2, num_q_heads=32, seq_len=256, BLOCK_M=128 @@ -347,6 +350,12 @@ def _attn_fwd( # Per-tile: skip entire tile only if ALL rows are negligible skip_tile = tl.min(can_skip.to(tl.int32)) == 1 + # Optional runtime sparsity measurement via atomic counters + if MEASURE_SPARSITY: + tl.atomic_add(Sparsity_total, 1) # count every tile + if skip_tile: + tl.atomic_add(Sparsity_skipped, 1) # count skipped tiles + if not skip_tile: m_new = tl.maximum(row_max, tile_row_max) p = tl.math.exp2(scores - m_new[:, None]) @@ -385,7 +394,9 @@ def _attn_fwd( row_max = m_new # --- Final normalization: output = acc / row_sum --- - acc = acc / row_sum[:, None] + # Clamp denominator to avoid 0/0 NaN when skip-softmax skips all KV tiles. + # Safe because acc is also 0 in that case (never accumulated), so 0/eps = 0. + acc = acc / tl.maximum(row_sum[:, None], 1e-6) # Save LSE for backward pass (log2-space: lse = max + log2(sum)) if STORE_LSE: @@ -768,6 +779,8 @@ def forward( num_sink_tokens, dense_window_size, skip_softmax_threshold, + skip_softmax_raw_threshold, + measure_sparsity, ): HEAD_DIM = q.shape[2] num_q_heads = q.shape[1] @@ -788,20 +801,36 @@ def forward( # Triton tiles must be powers of 2; pad head dim BLOCK_D = triton.next_power_of_2(HEAD_DIM) - # Skip-softmax: convert threshold to scaled log2 space for the kernel. - # The BLASST reference (https://arxiv.org/pdf/2512.12087) checks - # ln(lambda) on unscaled scores. Our kernel works in log2-scaled space - # (scores pre-multiplied by qk_scale = sm_scale * LOG2E), so we - # pre-scale: threshold_scaled = log2(lambda) * sm_scale. - apply_skip = skip_softmax_threshold is not None and skip_softmax_threshold > 0.0 - if apply_skip: + # Skip-softmax threshold in scaled log2 space for the kernel. + # Two modes: + # 1. raw_threshold: passed directly as skip_threshold_log2 (for testing) + # 2. lambda threshold: converted via log2(lambda) * sm_scale + if skip_softmax_raw_threshold is not None: + apply_skip = True + skip_threshold_log2 = skip_softmax_raw_threshold + elif skip_softmax_threshold is not None and skip_softmax_threshold > 0.0: + apply_skip = True + # The BLASST reference (https://arxiv.org/pdf/2512.12087) checks + # ln(lambda) on unscaled scores. Our kernel works in log2-scaled space + # (scores pre-multiplied by qk_scale = sm_scale * LOG2E), so we + # pre-scale: threshold_scaled = log2(lambda) * sm_scale. skip_threshold_log2 = math.log2(skip_softmax_threshold) * sm_scale else: + apply_skip = False skip_threshold_log2 = 0.0 o = torch.empty_like(q) lse = torch.empty(q.shape[0], num_q_heads, device=q.device, dtype=torch.float32) + # Optional runtime sparsity counters (single int64 scalars for atomic adds) + do_measure = measure_sparsity and apply_skip + if do_measure: + sparsity_total = torch.zeros(1, dtype=torch.int64, device=q.device) + sparsity_skipped = torch.zeros(1, dtype=torch.int64, device=q.device) + else: + sparsity_total = None + sparsity_skipped = None + # Grid: (batch, q_heads, q_tiles). Uses a function because BLOCK_M is autotuned. def grid(META): return (batch, num_q_heads, triton.cdiv(max_input_len, META["BLOCK_M"])) @@ -839,9 +868,17 @@ def grid(META): DENSE_WINDOW_SIZE=dense_window_size, APPLY_SKIP_SOFTMAX=apply_skip, SKIP_THRESHOLD_LOG2=skip_threshold_log2, + Sparsity_total=sparsity_total, + Sparsity_skipped=sparsity_skipped, + MEASURE_SPARSITY=do_measure, # BLOCK_M, BLOCK_N, num_warps, num_stages chosen by autotune ) + # Store sparsity counters on the output tensor for retrieval by callers + if do_measure: + o._sparsity_total = sparsity_total.item() + o._sparsity_skipped = sparsity_skipped.item() + ctx.save_for_backward(q, k, v, o, lse, b_start_loc, b_seq_len, b_start_loc_k, b_seq_len_k) ctx.max_input_len = max_input_len ctx.max_input_len_k = max_input_len_k @@ -985,6 +1022,8 @@ def backward(ctx, grad_output): None, None, None, + None, + None, ) @@ -1006,6 +1045,8 @@ def attention( num_sink_tokens: int = 0, dense_window_size: int = 64, skip_softmax_threshold: float | None = None, + skip_softmax_raw_threshold: float | None = None, + measure_sparsity: bool = False, ) -> torch.Tensor: """Variable-length flash attention with GQA, autograd, and optional N:M sparse softmax and skip-softmax. @@ -1037,6 +1078,16 @@ def attention( softmax contribution is negligible. Tiles are skipped entirely (no softmax, V load, or BMM2). The threshold is applied on unscaled scores. Set to ``None`` or ``0`` to disable. + skip_softmax_raw_threshold: Raw ``skip_threshold_log2`` value passed + directly to the kernel without conversion. The kernel skips tiles + where ``tile_row_max < row_max + raw_threshold``. Typical values + are negative (e.g., ``-5.0`` means tiles must be within 5 units of + the running max in the kernel's scaled score space). Takes + precedence over ``skip_softmax_threshold`` when both are set. + measure_sparsity: When True and skip-softmax is active, count total + and skipped tiles via atomic counters. The counts are stored as + ``_sparsity_total`` and ``_sparsity_skipped`` attributes on the + returned output tensor. Returns: Output tensor [total_q_tokens, num_q_heads, head_dim]. @@ -1059,7 +1110,271 @@ def attention( num_sink_tokens, dense_window_size, skip_softmax_threshold, + skip_softmax_raw_threshold, + measure_sparsity, + ) + + +# --------------------------------------------------------------------------- +# Calibration kernel: collect multi-threshold skip-softmax sparsity stats +# --------------------------------------------------------------------------- +@triton.jit +def _attn_fwd_calibrate( + Q, + K, + V, + qk_scale, + b_start_loc, + b_seq_len, + b_start_loc_k, + b_seq_len_k, + Out, + stride_qbs, + stride_qh, + stride_kbs, + stride_kh, + stride_vbs, + stride_vh, + stride_obs, + stride_oh, + Threshold_trials, # [NUM_THRESHOLDS] float32 — pre-scaled to log2 space + Per_program_totals, # [num_programs * NUM_THRESHOLDS] int32 — per-program tile counts + Per_program_skipped, # [num_programs * NUM_THRESHOLDS] int32 — per-program skip counts + kv_group_num: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_D: tl.constexpr, + BLOCK_N: tl.constexpr, + IS_CAUSAL: tl.constexpr, + HEAD_DIM: tl.constexpr, + NUM_THRESHOLDS: tl.constexpr, + PADDED_THRESHOLDS: tl.constexpr, # next_power_of_2(NUM_THRESHOLDS) for tl.arange +): + """Forward kernel with multi-threshold sparsity measurement. + + Computes full attention (no skipping) while counting how many KV tiles + would be skipped at each threshold. Each program writes its local counts + to ``Per_program_totals`` and ``Per_program_skipped``; the Python wrapper + sums across programs afterward. This avoids global atomic contention. + """ + batch_idx = tl.program_id(0) + head_idx = tl.program_id(1) + tile_q = tl.program_id(2) + kv_head_idx = head_idx // kv_group_num + + seq_len_q = tl.load(b_seq_len + batch_idx) + seq_len_kv = tl.load(b_seq_len_k + batch_idx) + q_offset = tl.load(b_start_loc + batch_idx) + kv_offset = tl.load(b_start_loc_k + batch_idx) + + if tile_q * BLOCK_M >= seq_len_q: + return + + q_pos = tile_q * BLOCK_M + tl.arange(0, BLOCK_M) + kv_pos = tl.arange(0, BLOCK_N) + dim_pos = tl.arange(0, BLOCK_D) + d_mask = dim_pos < HEAD_DIM + + q_ptrs = (q_offset + q_pos[:, None]) * stride_qbs + head_idx * stride_qh + dim_pos[None, :] + q = tl.load(Q + q_ptrs, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :], other=0.0) + + k_base = K + kv_head_idx * stride_kh + v_base = V + kv_head_idx * stride_vh + + row_max = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + row_sum = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, BLOCK_D], dtype=tl.float32) + + # Pre-load all thresholds once (vectorized, stays in registers). + # tl.arange requires power-of-2 size, so use PADDED_THRESHOLDS with masking. + thresh_offs = tl.arange(0, PADDED_THRESHOLDS) + thresh_mask = thresh_offs < NUM_THRESHOLDS + thresholds = tl.load(Threshold_trials + thresh_offs, mask=thresh_mask, other=float("inf")) + + # Per-program local counters: avoid global atomic contention in inner loop. + # Each program accumulates locally, then writes once to Per_program buffers. + local_skipped = tl.zeros([PADDED_THRESHOLDS], dtype=tl.int32) + num_tiles = 0 + + kv_bound = seq_len_kv if not IS_CAUSAL else tl.minimum((tile_q + 1) * BLOCK_M, seq_len_kv) + + for kv_start in range(0, kv_bound, BLOCK_N): + kv_start = tl.multiple_of(kv_start, BLOCK_N) + + k_offs = (kv_offset + kv_start + kv_pos[None, :]) * stride_kbs + dim_pos[:, None] + k = tl.load( + k_base + k_offs, + mask=((kv_start + kv_pos[None, :]) < seq_len_kv) & d_mask[:, None], + other=0.0, + ) + + scores = tl.dot(q, k) * qk_scale + scores = _apply_mask(scores, q_pos, kv_pos, seq_len_q, seq_len_kv, kv_start, IS_CAUSAL) + + tile_row_max = tl.max(scores, 1) + + # --- Vectorized multi-threshold sparsity measurement --- + # A tile is skipped iff ALL Q rows satisfy: tile_row_max < row_max + thresh. + # Equivalently: max(tile_row_max - row_max) < thresh (worst-case row + # must still be below threshold for the tile to be skippable). + max_gap = tl.max(tile_row_max - row_max) # scalar + skip_mask = (max_gap < thresholds).to(tl.int32) # [PADDED_THRESHOLDS] + local_skipped += skip_mask + num_tiles += 1 + + # --- Always compute full attention (no skipping) --- + m_new = tl.maximum(row_max, tile_row_max) + p = tl.math.exp2(scores - m_new[:, None]) + l_new = tl.sum(p, 1) + correction = tl.math.exp2(row_max - m_new) + row_sum = row_sum * correction + l_new + acc = acc * correction[:, None] + + v_offs = (kv_offset + kv_start + kv_pos[:, None]) * stride_vbs + dim_pos[None, :] + v = tl.load( + v_base + v_offs, + mask=((kv_start + kv_pos[:, None]) < seq_len_kv) & d_mask[None, :], + other=0.0, + ) + acc = tl.dot(p.to(v.dtype), v, acc) + row_max = m_new + + # --- Write per-program counters (no atomics, just stores) --- + # Compute unique flat program index for this (batch, head, q_tile) + num_q_tiles = tl.cdiv(tl.load(b_seq_len + 0), BLOCK_M) # conservative upper bound + num_heads = tl.num_programs(1) + prog_idx = batch_idx * num_heads * num_q_tiles + head_idx * num_q_tiles + tile_q + base = prog_idx * NUM_THRESHOLDS + tl.store( + Per_program_totals + base + thresh_offs, + tl.full([PADDED_THRESHOLDS], num_tiles, dtype=tl.int32), + mask=thresh_mask, ) + tl.store( + Per_program_skipped + base + thresh_offs, + local_skipped, + mask=thresh_mask, + ) + + acc = acc / tl.maximum(row_sum[:, None], 1e-6) + o_ptrs = (q_offset + q_pos[:, None]) * stride_obs + head_idx * stride_oh + dim_pos[None, :] + tl.store(Out + o_ptrs, acc, mask=(q_pos[:, None] < seq_len_q) & d_mask[None, :]) + + +def attention_calibrate( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b_start_loc: torch.Tensor, + b_seq_len: torch.Tensor, + max_input_len: int, + is_causal: bool = True, + softmax_scale: float | None = None, + b_start_loc_k: torch.Tensor | None = None, + b_seq_len_k: torch.Tensor | None = None, + max_input_len_k: int | None = None, + *, + threshold_trials: list[float] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Flash attention with multi-threshold skip-softmax sparsity measurement. + + Computes full attention (identical output to dense attention) while + measuring how many KV tiles would be skipped at each threshold in + ``threshold_trials``. No autograd — forward only. + + Args: + q, k, v, b_start_loc, b_seq_len, max_input_len, is_causal, + softmax_scale, b_start_loc_k, b_seq_len_k, max_input_len_k: + Same as :func:`attention`. + threshold_trials: List of threshold values to measure sparsity for. + Each value is converted to log2-scaled space for the kernel. + + Returns: + Tuple of (output, sparsity_counters): + - output: ``[total_q_tokens, num_q_heads, head_dim]`` + - sparsity_counters: ``[num_thresholds, 2]`` int64 tensor where + ``[:, 0]`` = total tile evaluations, ``[:, 1]`` = skipped tiles. + Sparsity per threshold = ``counters[:, 1] / counters[:, 0]``. + """ + if threshold_trials is None or len(threshold_trials) == 0: + raise ValueError("threshold_trials must be a non-empty list") + + HEAD_DIM = q.shape[2] + num_q_heads = q.shape[1] + num_kv_heads = k.shape[1] + kv_group_num = num_q_heads // num_kv_heads + batch = b_seq_len.shape[0] + sm_scale = 1.0 / (HEAD_DIM**0.5) if softmax_scale is None else softmax_scale + qk_scale = sm_scale * LOG2E + BLOCK_D = triton.next_power_of_2(HEAD_DIM) + BLOCK_M = 128 + BLOCK_N = 64 + + if b_seq_len_k is None: + b_seq_len_k = b_seq_len + b_start_loc_k = b_start_loc + + num_thresholds = len(threshold_trials) + + # Convert thresholds to log2-scaled space: log2(lambda) * sm_scale + threshold_tensor = torch.tensor( + [math.log2(t) * sm_scale for t in threshold_trials], + dtype=torch.float32, + device=q.device, + ) + + o = torch.empty_like(q) + + num_q_tiles = triton.cdiv(max_input_len, BLOCK_M) + grid = (batch, num_q_heads, num_q_tiles) + num_programs = batch * num_q_heads * num_q_tiles + + # Per-program output buffers (no atomics needed — each program writes its own row) + per_program_totals = torch.zeros( + num_programs * num_thresholds, dtype=torch.int32, device=q.device + ) + per_program_skipped = torch.zeros( + num_programs * num_thresholds, dtype=torch.int32, device=q.device + ) + + _attn_fwd_calibrate[grid]( + q, + k, + v, + qk_scale, + b_start_loc, + b_seq_len, + b_start_loc_k, + b_seq_len_k, + o, + q.stride(0), + q.stride(1), + k.stride(0), + k.stride(1), + v.stride(0), + v.stride(1), + o.stride(0), + o.stride(1), + threshold_tensor, + per_program_totals, + per_program_skipped, + kv_group_num=kv_group_num, + BLOCK_M=BLOCK_M, + BLOCK_D=BLOCK_D, + BLOCK_N=BLOCK_N, + IS_CAUSAL=is_causal, + HEAD_DIM=HEAD_DIM, + NUM_THRESHOLDS=num_thresholds, + PADDED_THRESHOLDS=triton.next_power_of_2(num_thresholds), + num_warps=4, + num_stages=1, + ) + + # Reduce across programs: sum per-program counts → [num_thresholds] + totals = per_program_totals.view(num_programs, num_thresholds).sum(dim=0).to(torch.int64) + skipped = per_program_skipped.view(num_programs, num_thresholds).sum(dim=0).to(torch.int64) + sparsity_counters = torch.stack([totals, skipped], dim=1) # [num_thresholds, 2] + + return o, sparsity_counters -__all__ = ["attention"] +__all__ = ["attention", "attention_calibrate"] diff --git a/modelopt/torch/quantization/src/conv/bench_implicit_gemm.py b/modelopt/torch/quantization/src/conv/bench_implicit_gemm.py new file mode 100644 index 00000000000..807ce178387 --- /dev/null +++ b/modelopt/torch/quantization/src/conv/bench_implicit_gemm.py @@ -0,0 +1,208 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Latency benchmark: implicit GEMM (quant / non-quant) vs cuDNN conv3d. + +Usage: + python -m experimental.conv.bench_implicit_gemm + python -m experimental.conv.bench_implicit_gemm --shapes wan22 + python -m experimental.conv.bench_implicit_gemm --shapes all --warmup 20 --iters 100 +""" + +import argparse + +import torch +import torch.nn.functional as F + +# --------------------------------------------------------------------------- +# Benchmark shapes +# --------------------------------------------------------------------------- + +# (name, N, Cin, D, H, W, Cout, kD, kH, kW, stride, padding, dilation) +SHAPES = { + "small": [ + ("small_16x32_3x3x3", 1, 16, 8, 8, 8, 32, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), + ], + "medium": [ + ("med_64x128_3x3x3", 1, 64, 16, 32, 32, 128, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), + ("med_128x256_3x3x3", 1, 128, 8, 16, 16, 256, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), + ("med_128x128_1x3x3", 1, 128, 16, 32, 32, 128, 1, 3, 3, (1, 1, 1), (0, 1, 1), (1, 1, 1)), + ], + "wan22": [ + ("wan22_128x512", 1, 128, 21, 60, 106, 512, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), + ("wan22_512x512", 1, 512, 21, 60, 106, 512, 1, 1, 1, (1, 1, 1), (0, 0, 0), (1, 1, 1)), + ("wan22_512x128", 1, 512, 21, 60, 106, 128, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), + ], + "stride": [ + ("stride2_64x128", 1, 64, 16, 32, 32, 128, 3, 3, 3, (2, 2, 2), (1, 1, 1), (1, 1, 1)), + ("stride2_128x256", 1, 128, 16, 32, 32, 256, 3, 3, 3, (2, 2, 2), (1, 1, 1), (1, 1, 1)), + ], +} + + +def get_shapes(name: str): + """Return list of benchmark shapes by name or all shapes.""" + if name == "all": + result = [] + for v in SHAPES.values(): + result.extend(v) + return result + return SHAPES[name] + + +# --------------------------------------------------------------------------- +# Timing utility +# --------------------------------------------------------------------------- + + +def bench_fn(fn, warmup: int, iters: int) -> float: + """Benchmark a callable, return median time in ms.""" + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + times = [] + for _ in range(iters): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + fn() + end.record() + torch.cuda.synchronize() + times.append(start.elapsed_time(end)) + + times.sort() + return times[len(times) // 2] # median + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def run_benchmark(shapes_name: str, warmup: int, iters: int, fp4_block_size: int): + """Run latency benchmark for the given shapes.""" + from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda + + shapes = get_shapes(shapes_name) + + # Header + print(f"\n{'=' * 100}") + print( + f"Conv3D Latency Benchmark | warmup={warmup} iters={iters} fp4_block_size={fp4_block_size}" + ) + print(f"GPU: {torch.cuda.get_device_name()}") + print(f"{'=' * 100}") + print( + f"{'Shape':<25} {'M':>10} {'K':>8} {'N':>6} " + f"{'cuDNN':>9} {'GEMM':>9} {'GEMM+FP4':>9} " + f"{'GEMM/cuDNN':>11} {'FP4/cuDNN':>10}" + ) + print("-" * 100) + + for name, n, cin, d, h, w, cout, kd, kh, kw, stride, padding, dilation in shapes: + torch.manual_seed(42) + x = torch.randn(n, cin, d, h, w, device="cuda", dtype=torch.float32) + weight = torch.randn(cout, cin, kd, kh, kw, device="cuda", dtype=torch.float32) + act_amax = x.abs().max().unsqueeze(0) + + # Compute GEMM dimensions for display + sd, sh, sw = stride + dd, dh, dw = dilation + pd, ph, pw = padding + od = (d + 2 * pd - dd * (kd - 1) - 1) // sd + 1 + oh = (h + 2 * ph - dh * (kh - 1) - 1) // sh + 1 + ow = (w + 2 * pw - dw * (kw - 1) - 1) // sw + 1 + gemm_m = n * od * oh * ow + gemm_k = cin * kd * kh * kw + gemm_n = cout + + # cuDNN (torch.nn.functional.conv3d) + t_cudnn = bench_fn( + lambda: F.conv3d(x, weight, stride=stride, padding=padding, dilation=dilation), + warmup, + iters, + ) + + # Implicit GEMM (non-quantized) + t_gemm = bench_fn( + lambda: conv3d_implicit_gemm_cuda( + x, + weight, + stride=stride, + padding=padding, + dilation=dilation, + quant_act=False, + fp4_block_size=fp4_block_size, + ), + warmup, + iters, + ) + + # Implicit GEMM (FP4 quantized) + t_fp4 = bench_fn( + lambda: conv3d_implicit_gemm_cuda( + x, + weight, + stride=stride, + padding=padding, + dilation=dilation, + act_amax=act_amax, + quant_act=True, + fp4_block_size=fp4_block_size, + ), + warmup, + iters, + ) + + ratio_gemm = t_gemm / t_cudnn + ratio_fp4 = t_fp4 / t_cudnn + + print( + f"{name:<25} {gemm_m:>10,} {gemm_k:>8,} {gemm_n:>6,} " + f"{t_cudnn:>8.3f}ms {t_gemm:>8.3f}ms {t_fp4:>8.3f}ms " + f"{ratio_gemm:>10.2f}x {ratio_fp4:>9.2f}x" + ) + + print(f"{'=' * 100}") + print("Ratios > 1.0x mean slower than cuDNN; < 1.0x mean faster.") + print() + + +def main(): + """Entry point for the benchmark CLI.""" + parser = argparse.ArgumentParser(description="Conv3D latency benchmark") + parser.add_argument( + "--shapes", + default="all", + choices=[*list(SHAPES.keys()), "all"], + help="Which shape set to benchmark (default: all)", + ) + parser.add_argument("--warmup", type=int, default=20, help="Warmup iterations") + parser.add_argument("--iters", type=int, default=100, help="Benchmark iterations") + parser.add_argument( + "--fp4-block-size", + type=int, + default=128, + choices=[128, 256], + help="FP4 block size (default: 128)", + ) + args = parser.parse_args() + + run_benchmark(args.shapes, args.warmup, args.iters, args.fp4_block_size) + + +if __name__ == "__main__": + main() diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py index dbc4d5bc274..f63feac69ed 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrate.py @@ -21,7 +21,6 @@ import torch import torch.nn as nn -from transformers import AutoTokenizer from modelopt.torch.utils import get_module_device @@ -32,8 +31,10 @@ from .ruler_dataset import RulerDatasetBuilder -def _load_tokenizer(tokenizer_name_or_path: str) -> "AutoTokenizer": +def _load_tokenizer(tokenizer_name_or_path: str): """Load tokenizer and ensure pad_token is set.""" + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path) if not tokenizer.pad_token: tokenizer.pad_token = tokenizer.eos_token @@ -255,11 +256,14 @@ def calibrate_sparse_attention( print(f"Calibrating {len(sparse_modules)} sparse attention modules together...") - # Extract tokenizer and build calibration data if needed - tokenizer = _extract_tokenizer_from_model(model) + # Extract tokenizer and build calibration data only if no forward_loop is provided. + # When the user supplies their own forward_loop (e.g. for diffusion models), + # RULER dataset generation is skipped entirely. + tokenizer = None calibration_data = None - if calibrate_prefill or calibrate_decode: + if forward_loop is None and (calibrate_prefill or calibrate_decode): + tokenizer = _extract_tokenizer_from_model(model) builder = RulerDatasetBuilder( samples=calib_config.samples, max_seqlen=calib_config.max_seqlen, @@ -280,14 +284,19 @@ def calibrate_sparse_attention( print("PREFILL PHASE CALIBRATION") print("=" * 60) - if calibration_data is None: + if forward_loop is None and calibration_data is None: raise RuntimeError("calibration_data must be built before prefill") - prefill_forward_loop = forward_loop or create_calibration_forward_loop( - calibration_data, tokenizer, chunk_size=calib_config.chunk_size - ) + if forward_loop is not None: + prefill_forward_loop = forward_loop + else: + assert calibration_data is not None and tokenizer is not None + prefill_forward_loop = create_calibration_forward_loop( + calibration_data, tokenizer, chunk_size=calib_config.chunk_size + ) prefill_calibrator = DynamicThresholdCalibrator( threshold_trials=calib_config.threshold_trials, + fit_logspace=calib_config.fit_logspace, ) prefill_result = prefill_calibrator.calibrate(model, prefill_forward_loop, phase="prefill") @@ -302,14 +311,15 @@ def calibrate_sparse_attention( print("DECODE PHASE CALIBRATION") print("=" * 60) - if calibration_data is None: - raise RuntimeError("calibration_data must be built before decode") + if calibration_data is None or tokenizer is None: + raise RuntimeError("calibration_data and tokenizer must be built before decode") decode_forward_loop = create_decode_calibration_forward_loop( calibration_data, tokenizer, num_decode_tokens=calib_config.num_decode_tokens ) decode_calibrator = DynamicThresholdCalibrator( threshold_trials=calib_config.threshold_trials, + fit_logspace=calib_config.fit_logspace, ) decode_result = decode_calibrator.calibrate(model, decode_forward_loop, phase="decode") @@ -323,15 +333,20 @@ def calibrate_sparse_attention( warnings.warn("No calibration produced valid results") return {} - # Extract a and b for each phase + # Extract a, b, and observed sparsity range for each phase calibration_params: dict[str, dict[str, float]] = {} for phase in ["prefill", "decode"]: if phase in calibration_results: result = calibration_results[phase] - calibration_params[phase] = { + params: dict[str, float] = { "a": result["a"], "b": result["b"], } + if "min_observed_sparsity" in result: + params["min_observed_sparsity"] = result["min_observed_sparsity"] + if "max_observed_sparsity" in result: + params["max_observed_sparsity"] = result["max_observed_sparsity"] + calibration_params[phase] = params # Apply calibration params to all modules print("\n" + "=" * 60) @@ -341,7 +356,7 @@ def calibrate_sparse_attention( for phase, params in calibration_params.items(): result = calibration_results[phase] print(f" {phase}:") - print(f" Model: scale_factor = {params['a']:.6f} * exp({params['b']:.4f} * sparsity)") + print(f" Model: scale_factor = {params['a']:.6e} * exp({params['b']:.4f} * sparsity)") print(f" R-squared: {result['r_squared']:.6f}") for module_name, module in sparse_modules: diff --git a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py index 68212069378..d3ed3303256 100644 --- a/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py +++ b/modelopt/torch/sparsity/attention_sparsity/calibration/calibrator.py @@ -55,12 +55,16 @@ class DynamicThresholdCalibrator: def __init__( self, threshold_trials: list[float] | None = None, + fit_logspace: bool = False, ): """Initialize dynamic threshold calibrator. Args: threshold_trials: List of thresholds to try during calibration. Should span a range that achieves sparsities from ~10% to ~95%. + fit_logspace: If True, fit the exponential model in log space + (minimizes relative error). Recommended for diffusion models + where scale_factors span many orders of magnitude. """ # Default threshold trials if not provided self.threshold_trials = threshold_trials or [ @@ -85,6 +89,7 @@ def __init__( 9.5e-1, 9.9e-1, ] + self.fit_logspace = fit_logspace def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dict[str, Any]: """Calibrate a and b parameters for Exponential model. @@ -167,6 +172,8 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic # Filter out extreme sparsities (must be in (10%, 90%)) # Extreme values are unreliable for fitting valid_mask = (sparsities >= 0.10) & (sparsities <= 0.90) + if self.fit_logspace: + valid_mask &= scale_factors > 0 # log requires positive values scale_factors = scale_factors[valid_mask] sparsities = sparsities[valid_mask] @@ -176,47 +183,81 @@ def calibrate(self, model: nn.Module, forward_loop: Callable, phase: str) -> dic ) return {} - # Define Exponential model: sf = a * exp(b * S) - def exponential(sparsity, a, b): - return a * np.exp(b * sparsity) + # Record observed sparsity range for feasibility checks at inference + min_observed_sparsity = float(np.min(sparsities)) + max_observed_sparsity = float(np.max(sparsities)) - # Fit the model try: - popt, pcov = curve_fit( - exponential, - sparsities, - scale_factors, - p0=[1.0, 5.0], # Initial guess - bounds=([0.0, 0.0], [np.inf, 20.0]), # Bounds for a and b - maxfev=10000, - ) - a, b = popt + if self.fit_logspace: + # Log-space fit: minimizes relative error. Recommended for + # diffusion models where scale_factors span many orders of + # magnitude (e.g. 0.06 to 57,000) — a linear-space fit would + # be dominated by the largest values. + log_scale_factors = np.log(scale_factors) + + def log_exponential(sparsity, log_a, b): + return log_a + b * sparsity + + popt, pcov = curve_fit( + log_exponential, + sparsities, + log_scale_factors, + p0=[0.0, 10.0], + maxfev=10000, + ) + log_a, b = popt + a = np.exp(log_a) + + # R-squared in log space (where the fit was performed) + pred = log_exponential(sparsities, log_a, b) + ss_res = np.sum((log_scale_factors - pred) ** 2) + ss_tot = np.sum((log_scale_factors - np.mean(log_scale_factors)) ** 2) + else: + # Linear-space fit (default): minimizes absolute error. + + def exponential(sparsity, a, b): + return a * np.exp(b * sparsity) + + popt, pcov = curve_fit( + exponential, + sparsities, + scale_factors, + p0=[1.0, 5.0], + bounds=([0.0, 0.0], [np.inf, 20.0]), + maxfev=10000, + ) + a, b = popt + + pred = exponential(sparsities, a, b) + ss_res = np.sum((scale_factors - pred) ** 2) + ss_tot = np.sum((scale_factors - np.mean(scale_factors)) ** 2) except Exception as e: warnings.warn(f"Curve fitting failed: {e}") return {} - # Calculate R-squared and RMSE - pred_scale_factors = exponential(sparsities, a, b) - ss_res = np.sum((scale_factors - pred_scale_factors) ** 2) - ss_tot = np.sum((scale_factors - np.mean(scale_factors)) ** 2) r_squared = 1 - (ss_res / ss_tot) if ss_tot > 0 else 0 - rmse = np.sqrt(np.mean((scale_factors - pred_scale_factors) ** 2)) - print(f"\n{phase.capitalize()} Calibration Results (Exponential Model):") + fit_label = "log-space" if self.fit_logspace else "linear-space" + print(f"\n{phase.capitalize()} Calibration Results (Exponential Model, {fit_label} fit):") print(" Model: scale_factor = a * exp(b * sparsity)") - print(f" Fitted a: {a:.6f}") + print(f" Fitted a: {a:.6e}") print(f" Fitted b: {b:.4f}") print(f" R-squared: {r_squared:.6f}") - print(f" RMSE: {rmse:.2f}") + print( + f" Observed sparsity range: [{min_observed_sparsity:.1%}, {max_observed_sparsity:.1%}]" + ) print(f" Data points used: {int(np.sum(valid_mask))} / {len(all_data_points)}") # Show scale_factor for various target sparsities print("\nScale factors for different target sparsities:") - print(f" {'Target':<10} {'Scale Factor':<15}") - print(f" {'-' * 10} {'-' * 15}") - for target in [0.5, 0.7, 0.8, 0.9, 0.95]: + print(f" {'Target':<10} {'Scale Factor':<15} {'Note':<20}") + print(f" {'-' * 10} {'-' * 15} {'-' * 20}") + for target in [0.3, 0.4, 0.5, 0.6, 0.7, 0.8]: sf = a * np.exp(b * target) - print(f" {target:<10.0%} {sf:<15.2f}") + note = "" + if target < min_observed_sparsity or target > max_observed_sparsity: + note = "(extrapolation)" + print(f" {target:<10.0%} {sf:<15.4f} {note:<20}") # Print calibration data summary by threshold print("\nCalibration data summary (per threshold):") @@ -239,10 +280,11 @@ def exponential(sparsity, a, b): "a": float(a), "b": float(b), "r_squared": float(r_squared), - "rmse": float(rmse), "num_data_points": int(np.sum(valid_mask)), "total_samples": len(all_data_points), "calibration_type": "exponential", + "min_observed_sparsity": min_observed_sparsity, + "max_observed_sparsity": max_observed_sparsity, } def _enable_calibration_mode(self, modules: list[nn.Module]): @@ -333,6 +375,16 @@ def _extract_calibration_stats( return aggregated_stats def _set_thresholds(self, modules: list[nn.Module], thresholds: list[float]): - """Set thresholds list on sparse attention modules.""" + """Set thresholds list on sparse attention modules. + + Supports both flash_skip_softmax (sets ``thresholds`` attribute) and + triton_skip_softmax (sets ``_threshold_trials`` attribute). + """ for module in modules: - module._sparse_method_instance.thresholds = thresholds + method = module._sparse_method_instance + if hasattr(method, "_threshold_trials"): + # triton_skip_softmax: calibration uses Triton calibration kernel + method._threshold_trials = thresholds + else: + # flash_skip_softmax: calibration uses F.softmax patching + method.thresholds = thresholds diff --git a/modelopt/torch/sparsity/attention_sparsity/config.py b/modelopt/torch/sparsity/attention_sparsity/config.py index fa415b322bf..eed50b87af1 100644 --- a/modelopt/torch/sparsity/attention_sparsity/config.py +++ b/modelopt/torch/sparsity/attention_sparsity/config.py @@ -139,6 +139,17 @@ class SparseAttentionAttributeConfig(ModeloptBaseConfig): ), ) + skip_softmax_raw_threshold: float | None = ModeloptField( + default=None, + title="Raw skip-softmax threshold (skip_threshold_log2).", + description=( + "Raw value passed directly to the Triton kernel as skip_threshold_log2. " + "The kernel skips tiles where tile_row_max < row_max + raw_threshold. " + "Typical values are negative (e.g., -5.0). Takes precedence over " + "skip_softmax_threshold and calibration when set." + ), + ) + @field_validator("method") @classmethod def validate_method(cls, v): @@ -326,6 +337,15 @@ class CalibrationConfig(ModeloptBaseConfig): ), ) + fit_logspace: bool = ModeloptField( + default=False, + title="Fit in log space", + description=( + "If True, fit the exponential model in log space (minimizes relative error). " + "Recommended for diffusion models where scale_factors span many orders of magnitude." + ), + ) + cache_dir: str | None = ModeloptField( default=None, title="Cache directory", diff --git a/modelopt/torch/sparsity/attention_sparsity/conversion.py b/modelopt/torch/sparsity/attention_sparsity/conversion.py index 6ba238e77c0..cc928198509 100644 --- a/modelopt/torch/sparsity/attention_sparsity/conversion.py +++ b/modelopt/torch/sparsity/attention_sparsity/conversion.py @@ -115,6 +115,37 @@ def is_attn_sparsified(model: nn.Module) -> bool: return any(isinstance(module, SparseAttentionModule) for module in model.modules()) +def _register_diffusers_backends_if_needed(model: nn.Module) -> None: + """Register diffusers/LTX Triton attention backends if the model needs them. + + Called before plugin registration so that the backends are available + when ``SparseAttentionModule.forward()`` activates the skip-softmax context. + """ + import contextlib + + # Register the diffusers Triton backend if the model is a diffusers ModelMixin + try: + from diffusers.models.modeling_utils import ModelMixin + + if isinstance(model, ModelMixin): + from .kernels import register_diffusers_triton_attention + + if register_diffusers_triton_attention is not None: + register_diffusers_triton_attention() + except (ImportError, Exception): + pass + + # Patch ltx_core Attention modules if present (independent of diffusers) + try: + from .kernels import register_ltx_triton_attention + except (ImportError, RuntimeError): + return + + if register_ltx_triton_attention is not None: + with contextlib.suppress(Exception): + register_ltx_triton_attention(model) + + def convert_to_sparse_attention_model( model: ModelLikeModule, config: SparseAttentionConfig ) -> ConvertReturnType: @@ -130,6 +161,9 @@ def convert_to_sparse_attention_model( # Initialize the true module if necessary model = model.init_modellike() if isinstance(model, ModelLikeModule) else model + # Register diffusers backends for diffusion models + _register_diffusers_backends_if_needed(model) + # Set the correct attn_implementation for the chosen backend _set_attn_implementation(model, config) @@ -484,6 +518,8 @@ def print_sparse_attention_summary(model: nn.Module): # Group by (method, threshold) groups: dict[tuple[str, str], int] = {} for _, module in sparse_modules: + if not module.is_enabled: + continue method = getattr(module, "_method", "unknown") threshold = _format_threshold(module.get_threshold_info()) groups[(method, threshold)] = groups.get((method, threshold), 0) + 1 diff --git a/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py b/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py index dee1bc472a2..0cc4a202f57 100644 --- a/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py +++ b/modelopt/torch/sparsity/attention_sparsity/kernels/__init__.py @@ -13,12 +13,48 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Re-exports from modelopt.torch.kernels for backward compatibility.""" +"""Kernel integrations for sparse attention: Triton FA and diffusers/LTX backends.""" + +import contextlib +import threading from modelopt.torch.kernels import IS_AVAILABLE, attention, register_triton_attention +# --------------------------------------------------------------------------- +# Optional backend registrations (depend on diffusers / ltx_core) +# --------------------------------------------------------------------------- +register_diffusers_triton_attention = None +register_ltx_triton_attention = None + +# Suppress ImportError (missing package) and RuntimeError (triton without GPU driver) +with contextlib.suppress(ImportError, RuntimeError): + from .diffusers_triton_attention import register_diffusers_triton_attention + +with contextlib.suppress(ImportError, RuntimeError): + from .ltx_triton_attention import register_ltx_triton_attention + +# --------------------------------------------------------------------------- +# Thread-local flag for flash_skip_softmax's eager-attention context +# --------------------------------------------------------------------------- +_thread_local = threading.local() + + +def set_skip_softmax_context(active: bool) -> None: + """Set whether skip-softmax softmax patching is active (thread-local).""" + _thread_local.skip_softmax_active = active + + +def get_skip_softmax_context() -> bool: + """Return whether skip-softmax softmax patching is active.""" + return getattr(_thread_local, "skip_softmax_active", False) + + __all__ = [ "IS_AVAILABLE", "attention", + "get_skip_softmax_context", + "register_diffusers_triton_attention", + "register_ltx_triton_attention", "register_triton_attention", + "set_skip_softmax_context", ] diff --git a/modelopt/torch/sparsity/attention_sparsity/kernels/diffusers_triton_attention.py b/modelopt/torch/sparsity/attention_sparsity/kernels/diffusers_triton_attention.py new file mode 100644 index 00000000000..2923447cf02 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/kernels/diffusers_triton_attention.py @@ -0,0 +1,251 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Triton flash attention backend for diffusers models. + +Registers a ``modelopt_triton`` backend in diffusers' ``_AttentionBackendRegistry`` +that converts the diffusers [B, S, H, D] layout to the Triton FA kernel's varlen +[total_tokens, H, D] format. + +Two modes: +- **Inference**: Calls ``attention()`` with skip-softmax tile skipping. +- **Calibration**: Calls ``attention_calibrate()`` to collect multi-threshold + sparsity statistics without skipping any tiles. +""" + +import inspect +import math +import threading + +import torch +from diffusers.models.attention_dispatch import ( + AttentionBackendName, + _AttentionBackendRegistry, + attention_backend, +) + +from modelopt.torch.kernels import attention, attention_calibrate + +_BACKEND_NAME = "modelopt_triton" +_BACKEND_REGISTERED = False + +# Thread-local storage for per-forward skip-softmax configuration. +_thread_local = threading.local() + + +def set_triton_skip_softmax_config( + threshold: float | None = None, + calibration_mode: bool = False, + threshold_trials: list[float] | None = None, + scale_factor: float | None = None, + raw_threshold: float | None = None, + measure_sparsity: bool = False, +) -> None: + """Set thread-local skip-softmax config for the next Triton attention call. + + Args: + threshold: Skip-softmax threshold for inference mode (static). + calibration_mode: If True, use the calibration kernel to collect + multi-threshold sparsity stats instead of skipping tiles. + threshold_trials: List of thresholds to measure sparsity for + (only used when calibration_mode=True). + scale_factor: Calibrated scale factor for dynamic threshold computation. + When set, the actual threshold is computed as ``scale_factor / seq_k`` + at attention call time, adapting to the actual sequence length. + raw_threshold: Raw ``skip_threshold_log2`` value passed directly to the + kernel without conversion. Takes precedence over other thresholds. + measure_sparsity: If True, count total and skipped tiles during + inference via atomic counters in the forward kernel. + """ + _thread_local.skip_threshold = threshold + _thread_local.calibration_mode = calibration_mode + _thread_local.threshold_trials = threshold_trials + _thread_local.scale_factor = scale_factor + _thread_local.raw_threshold = raw_threshold + _thread_local.measure_sparsity = measure_sparsity + # Accumulated counters across all attention calls in one forward pass + _thread_local.calibration_counters = None + _thread_local.calibration_seq_k = None + # Accumulated runtime sparsity counters (total_tiles, skipped_tiles) + _thread_local.sparsity_total = 0 + _thread_local.sparsity_skipped = 0 + + +def clear_triton_skip_softmax_config() -> None: + """Clear thread-local skip-softmax config.""" + _thread_local.skip_threshold = None + _thread_local.calibration_mode = False + _thread_local.threshold_trials = None + _thread_local.scale_factor = None + _thread_local.raw_threshold = None + _thread_local.measure_sparsity = False + _thread_local.calibration_counters = None + _thread_local.calibration_seq_k = None + _thread_local.sparsity_total = 0 + _thread_local.sparsity_skipped = 0 + + +def get_calibration_counters() -> "torch.Tensor | None": + """Return accumulated calibration counters ``[num_thresholds, 2]`` or None.""" + return getattr(_thread_local, "calibration_counters", None) + + +def get_calibration_seq_k() -> int | None: + """Return KV sequence length observed during calibration, or None.""" + return getattr(_thread_local, "calibration_seq_k", None) + + +def get_sparsity_counters() -> tuple[int, int]: + """Return accumulated runtime sparsity counters ``(total_tiles, skipped_tiles)``.""" + return ( + getattr(_thread_local, "sparsity_total", 0), + getattr(_thread_local, "sparsity_skipped", 0), + ) + + +# --------------------------------------------------------------------------- +# Triton attention implementation for diffusers layout +# --------------------------------------------------------------------------- + + +def _diffusers_triton_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, +) -> torch.Tensor: + """Compute attention via Triton FA kernel on diffusers layout ``[B, S, H, D]``.""" + batch, seq_q, num_heads_q, head_dim = query.shape + seq_k = key.shape[1] + device = query.device + + # Reshape from diffusers [B, S, H, D] -> flat [B*S, H, D] + q = query.reshape(batch * seq_q, num_heads_q, head_dim).contiguous() + k = key.reshape(batch * seq_k, key.shape[2], head_dim).contiguous() + v = value.reshape(batch * seq_k, value.shape[2], head_dim).contiguous() + + # Build varlen metadata + b_start_loc_q = torch.arange(batch, device=device, dtype=torch.int32) * seq_q + b_seq_len_q = torch.full((batch,), seq_q, device=device, dtype=torch.int32) + + if scale is None: + scale = 1.0 / math.sqrt(head_dim) + + kw: dict = { + "b_start_loc": b_start_loc_q, + "b_seq_len": b_seq_len_q, + "max_input_len": seq_q, + "is_causal": is_causal, + "softmax_scale": scale, + } + + if seq_q != seq_k: + b_start_loc_k = torch.arange(batch, device=device, dtype=torch.int32) * seq_k + b_seq_len_k = torch.full((batch,), seq_k, device=device, dtype=torch.int32) + kw["b_start_loc_k"] = b_start_loc_k + kw["b_seq_len_k"] = b_seq_len_k + kw["max_input_len_k"] = seq_k + + # --- Calibration mode: collect multi-threshold stats --- + calib_mode = getattr(_thread_local, "calibration_mode", False) + if calib_mode: + trials = getattr(_thread_local, "threshold_trials", None) + if trials and attention_calibrate is not None: + o, counters = attention_calibrate(q, k, v, **kw, threshold_trials=trials) + + # Accumulate counters across all attention calls in this forward pass + prev = getattr(_thread_local, "calibration_counters", None) + if prev is None: + _thread_local.calibration_counters = counters + else: + _thread_local.calibration_counters = prev + counters + + # Store actual KV sequence length for calibration stats + _thread_local.calibration_seq_k = seq_k + + return o.view(batch, seq_q, num_heads_q, head_dim) + + # --- Inference mode: skip-softmax with raw, dynamic, or static threshold --- + raw_thresh = getattr(_thread_local, "raw_threshold", None) + if raw_thresh is not None: + # Raw threshold: passed directly to kernel as skip_threshold_log2 + kw["skip_softmax_raw_threshold"] = raw_thresh + else: + scale_factor = getattr(_thread_local, "scale_factor", None) + if scale_factor is not None and scale_factor > 0.0: + # Dynamic threshold: adapt to actual sequence length + kw["skip_softmax_threshold"] = scale_factor / seq_k + else: + threshold = getattr(_thread_local, "skip_threshold", None) + if threshold is not None and threshold > 0.0: + kw["skip_softmax_threshold"] = threshold + + assert attention is not None, "Triton attention kernel not available (requires CUDA + triton)" + do_measure = getattr(_thread_local, "measure_sparsity", False) + if do_measure: + kw["measure_sparsity"] = True + o = attention(q, k, v, **kw) + + # Accumulate runtime sparsity counters from the kernel output + if do_measure and hasattr(o, "_sparsity_total"): + prev_total = getattr(_thread_local, "sparsity_total", 0) + prev_skipped = getattr(_thread_local, "sparsity_skipped", 0) + _thread_local.sparsity_total = prev_total + o._sparsity_total + _thread_local.sparsity_skipped = prev_skipped + o._sparsity_skipped + + return o.view(batch, seq_q, num_heads_q, head_dim) + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + + +def register_diffusers_triton_attention() -> None: + """Register ``modelopt_triton`` backend in diffusers. + + Safe to call multiple times; registration happens only once. + """ + global _BACKEND_REGISTERED + if _BACKEND_REGISTERED: + return + + new_member = str.__new__(AttentionBackendName, _BACKEND_NAME) + new_member._name_ = "MODELOPT_TRITON" + new_member._value_ = _BACKEND_NAME + AttentionBackendName._member_map_["MODELOPT_TRITON"] = new_member + AttentionBackendName._value2member_map_[_BACKEND_NAME] = new_member + + _AttentionBackendRegistry._backends[new_member] = _diffusers_triton_attention + _AttentionBackendRegistry._constraints[new_member] = [] + _AttentionBackendRegistry._supported_arg_names[new_member] = set( + inspect.signature(_diffusers_triton_attention).parameters.keys() + ) + + _BACKEND_REGISTERED = True + + +def get_triton_attention_backend(): + """Return a context manager that activates the modelopt_triton backend.""" + if not _BACKEND_REGISTERED: + raise RuntimeError( + "modelopt_triton backend not registered. " + "Call register_diffusers_triton_attention() first." + ) + return attention_backend(_BACKEND_NAME) diff --git a/modelopt/torch/sparsity/attention_sparsity/kernels/ltx_triton_attention.py b/modelopt/torch/sparsity/attention_sparsity/kernels/ltx_triton_attention.py new file mode 100644 index 00000000000..fd53e7f9f49 --- /dev/null +++ b/modelopt/torch/sparsity/attention_sparsity/kernels/ltx_triton_attention.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Triton flash attention wrapper for LTX-2 (ltx_core) skip-softmax sparse attention. + +Two modes: +- **Inference**: ``attention()`` with skip-softmax tile skipping. +- **Calibration**: ``attention_calibrate()`` to collect multi-threshold stats. +""" + +import math +import threading + +import torch + +from modelopt.torch.kernels import attention, attention_calibrate +from modelopt.torch.utils.logging import warn_rank_0 + +# Thread-local storage for skip-softmax configuration +_thread_local = threading.local() + + +def set_ltx_triton_context( + active: bool, + threshold: float | None = None, + calibration_mode: bool = False, + threshold_trials: list[float] | None = None, + scale_factor: float | None = None, + raw_threshold: float | None = None, + **kwargs, +) -> None: + """Set thread-local Triton config for LTX-2 attention.""" + _thread_local.active = active + _thread_local.threshold = threshold + _thread_local.calibration_mode = calibration_mode + _thread_local.threshold_trials = threshold_trials + _thread_local.scale_factor = scale_factor + _thread_local.raw_threshold = raw_threshold + if not calibration_mode: + _thread_local.calibration_counters = None + _thread_local.calibration_seq_k = None + + +def clear_ltx_triton_context() -> None: + """Clear thread-local Triton config.""" + _thread_local.active = False + _thread_local.threshold = None + _thread_local.calibration_mode = False + _thread_local.threshold_trials = None + _thread_local.scale_factor = None + _thread_local.raw_threshold = None + _thread_local.calibration_counters = None + _thread_local.calibration_seq_k = None + + +def _get_ltx_triton_context() -> tuple[bool, float | None, float | None]: + """Return (active, threshold, scale_factor).""" + return ( + getattr(_thread_local, "active", False), + getattr(_thread_local, "threshold", None), + getattr(_thread_local, "scale_factor", None), + ) + + +def get_calibration_counters() -> "torch.Tensor | None": + """Return accumulated calibration counters ``[num_thresholds, 2]`` or None.""" + return getattr(_thread_local, "calibration_counters", None) + + +def get_calibration_seq_k() -> int | None: + """Return KV sequence length observed during calibration, or None.""" + return getattr(_thread_local, "calibration_seq_k", None) + + +def _ltx_triton_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + heads: int, + mask: torch.Tensor | None = None, + threshold: float | None = None, +) -> torch.Tensor: + """Triton FA attention on LTX-2 layout ``[B, T, H*D]``.""" + b, seq_q, dim_total = q.shape + dim_head = dim_total // heads + seq_k = k.shape[1] + device = q.device + + q_flat = q.view(b, seq_q, heads, dim_head).reshape(b * seq_q, heads, dim_head).contiguous() + k_flat = k.view(b, seq_k, heads, dim_head).reshape(b * seq_k, heads, dim_head).contiguous() + v_flat = v.view(b, seq_k, heads, dim_head).reshape(b * seq_k, heads, dim_head).contiguous() + + b_start_loc_q = torch.arange(b, device=device, dtype=torch.int32) * seq_q + b_seq_len_q = torch.full((b,), seq_q, device=device, dtype=torch.int32) + + scale = 1.0 / math.sqrt(dim_head) + + kw: dict = { + "b_start_loc": b_start_loc_q, + "b_seq_len": b_seq_len_q, + "max_input_len": seq_q, + "is_causal": False, + "softmax_scale": scale, + } + + if seq_q != seq_k: + b_start_loc_k = torch.arange(b, device=device, dtype=torch.int32) * seq_k + b_seq_len_k = torch.full((b,), seq_k, device=device, dtype=torch.int32) + kw["b_start_loc_k"] = b_start_loc_k + kw["b_seq_len_k"] = b_seq_len_k + kw["max_input_len_k"] = seq_k + + # --- Calibration mode --- + calib_mode = getattr(_thread_local, "calibration_mode", False) + if calib_mode: + trials = getattr(_thread_local, "threshold_trials", None) + if trials and attention_calibrate is not None: + o, counters = attention_calibrate(q_flat, k_flat, v_flat, **kw, threshold_trials=trials) + + prev = getattr(_thread_local, "calibration_counters", None) + if prev is None: + _thread_local.calibration_counters = counters + else: + _thread_local.calibration_counters = prev + counters + + # Store actual KV sequence length for calibration stats + _thread_local.calibration_seq_k = seq_k + + return o.view(b, seq_q, heads * dim_head) + + # --- Inference mode: raw, dynamic, or static threshold --- + raw_thresh = getattr(_thread_local, "raw_threshold", None) + scale_factor = getattr(_thread_local, "scale_factor", None) + if raw_thresh is not None: + kw["skip_softmax_raw_threshold"] = raw_thresh + elif scale_factor is not None and scale_factor > 0.0: + kw["skip_softmax_threshold"] = scale_factor / seq_k + elif threshold is not None and threshold > 0.0: + kw["skip_softmax_threshold"] = threshold + + assert attention is not None, "Triton attention kernel not available (requires CUDA + triton)" + o = attention(q_flat, k_flat, v_flat, **kw) + return o.view(b, seq_q, heads * dim_head) + + +class _TritonLTXAttentionWrapper: + """Wraps ltx_core attention_function for Triton dispatch.""" + + def __init__(self, original_fn): + self._original_fn = original_fn + + def __call__(self, q, k, v, heads, mask=None): + active, threshold, _scale_factor = _get_ltx_triton_context() + if active: + return _ltx_triton_attention(q, k, v, heads, mask, threshold) + return self._original_fn(q, k, v, heads, mask) + + +def register_ltx_triton_attention(model: torch.nn.Module) -> None: + """Patch all ``ltx_core.Attention`` modules for Triton dispatch.""" + from ltx_core.model.transformer.attention import Attention + + for module in model.modules(): + if isinstance(module, Attention): + warn_rank_0( + "LTX-2 packages (ltx-core, ltx-pipelines, ltx-trainer) are provided by " + "Lightricks and are NOT covered by the Apache 2.0 license governing NVIDIA " + "Model Optimizer. You MUST comply with the LTX Community License Agreement " + "when installing and using LTX-2 with NVIDIA Model Optimizer. Any derivative " + "models or fine-tuned weights from LTX-2 (including quantized or distilled " + "checkpoints) remain subject to the LTX Community License Agreement, not " + "Apache 2.0. See: https://github.com/Lightricks/LTX-2/blob/main/LICENSE", + UserWarning, + stacklevel=2, + ) + fn = module.attention_function + if not isinstance(fn, _TritonLTXAttentionWrapper): + module.attention_function = _TritonLTXAttentionWrapper(fn) diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py index 2501b58f659..117e337809f 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/flash_skip_softmax.py @@ -20,6 +20,7 @@ """ import math +from contextlib import ExitStack from typing import Any import numpy as np @@ -369,7 +370,11 @@ def get_threshold_info(self) -> dict[str, Any]: } def get_sparse_context(self, module: torch.nn.Module): - """Return a context manager that patches F.softmax with sparse masking.""" + """Return a context manager that patches F.softmax with sparse masking. + + Also registers the diffusers eager backend so that diffusion models + (which don't call F.softmax directly) route through the patched path. + """ original_softmax = F.softmax def sparse_softmax(input, dim=-1, *args, **kwargs): @@ -379,7 +384,14 @@ def sparse_softmax(input, dim=-1, *args, **kwargs): input = self.apply_sparsity(input, sparse_mask) return original_softmax(input, dim, *args, **kwargs) - return replace_function(torch.nn.functional, "softmax", sparse_softmax) + from ..kernels import set_skip_softmax_context + + stack = ExitStack() + set_skip_softmax_context(True) + stack.callback(set_skip_softmax_context, False) + + stack.enter_context(replace_function(torch.nn.functional, "softmax", sparse_softmax)) + return stack @property def name(self) -> str: diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/registry.py b/modelopt/torch/sparsity/attention_sparsity/methods/registry.py index 80371466430..3cb4f9010e1 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/registry.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/registry.py @@ -40,6 +40,10 @@ def __init__(self): # Video shape for VSA (T, H, W). None for non-VSA methods. self.video_shape: tuple[int, int, int] | None = None + def set_calibration_mode(self, enabled: bool) -> None: + """Enable or disable calibration mode (called by DynamicThresholdCalibrator).""" + self._calibration_mode = enabled + def forward_attention( self, query: torch.Tensor, diff --git a/modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py b/modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py index 4db51e894e7..1e2f3905e7a 100644 --- a/modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py +++ b/modelopt/torch/sparsity/attention_sparsity/methods/triton_skip_softmax.py @@ -13,10 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Skip-softmax method for attention via Triton kernel tile skipping.""" +"""Skip-softmax method for attention via Triton kernel tile skipping. + +Supports two modes: +- **Inference**: KV tiles with negligible scores are skipped in-kernel. +- **Calibration**: The Triton calibration kernel collects multi-threshold + sparsity statistics without skipping any tiles. +""" from contextlib import contextmanager +import torch + from .registry import SparseAttentionMethod, register_sparse_method @@ -39,21 +47,251 @@ def __init__(self, method_config=None): super().__init__() method_config = method_config or {} self.skip_softmax_threshold = method_config.get("skip_softmax_threshold", 0.1) + self.skip_softmax_raw_threshold: float | None = method_config.get( + "skip_softmax_raw_threshold", None + ) + # Calibration state + self._threshold_trials: list[float] | None = None + # Runtime sparsity measurement + self._measure_sparsity: bool = False + self._sparsity_total: int = 0 + self._sparsity_skipped: int = 0 @property def name(self) -> str: """Method name identifier.""" return "triton_skip_softmax" + def calculate_sparsity(self, attention_scores): + """Return a no-op mask (skip decision is made inside the Triton kernel).""" + mask = torch.ones_like(attention_scores, dtype=torch.bool) + return mask, {} + + def apply_sparsity(self, attention_scores, sparse_mask=None): + """Not supported — tile skipping is fused into the Triton kernel.""" + raise NotImplementedError( + "triton_skip_softmax applies tile skipping inside the Triton kernel. " + "Use backend='triton', not backend='pytorch'." + ) + def get_sparse_context(self, module): - """Return context manager that activates skip-softmax during forward.""" + """Return context manager that activates skip-softmax during forward. + + In calibration mode, configures the Triton backend to use the + calibration kernel which collects multi-threshold sparsity stats. + In inference mode, sets the skip threshold for tile skipping. + """ + if self._calibration_mode and self._threshold_trials: + return self._triton_calibration_context(module) + return self._triton_inference_context(module) + + @contextmanager + def _triton_inference_context(self, module): + """Inference: activate skip-softmax with calibrated or fixed threshold.""" + module._apply_skip_softmax = True + + backend_kwargs: dict = {} + if self._measure_sparsity: + backend_kwargs["measure_sparsity"] = True + + # Priority: raw_threshold > scale_factor (calibrated) > static threshold + if self.skip_softmax_raw_threshold is not None: + self._set_triton_backends( + raw_threshold=self.skip_softmax_raw_threshold, **backend_kwargs + ) + else: + scale_factor = self._get_scale_factor() + if scale_factor is not None: + self._set_triton_backends(scale_factor=scale_factor, **backend_kwargs) + else: + self._set_triton_backends(threshold=self.skip_softmax_threshold, **backend_kwargs) + with self._get_diffusers_backend_context(): + try: + yield + finally: + # Collect accumulated runtime sparsity counters before clearing + if self._measure_sparsity: + self._collect_sparsity_counters() + module._apply_skip_softmax = False + self._clear_triton_backends() - @contextmanager - def _skip_softmax_context(): - module._apply_skip_softmax = True + @contextmanager + def _triton_calibration_context(self, module): + """Calibration: collect multi-threshold sparsity stats via Triton kernel.""" + module._apply_skip_softmax = True + self._set_triton_backends(calibration_mode=True, threshold_trials=self._threshold_trials) + with self._get_diffusers_backend_context(): try: yield + # After forward pass, extract counters and build stats + self._collect_calibration_stats(module) finally: module._apply_skip_softmax = False + self._clear_triton_backends() + + def _get_scale_factor(self) -> float | None: + """Compute scale_factor from calibration params, or None if uncalibrated. + + The scale_factor is sequence-length-independent. Backends divide by the + actual ``seq_k`` at call time: ``threshold = scale_factor / seq_k``. + """ + if self.calibration_params and self.target_sparse_ratio: + import math + import warnings + + params = self.calibration_params.get("prefill", {}) + a = params.get("a", 0) + b = params.get("b", 0) + target = self.target_sparse_ratio.get("prefill", 0.5) + if a > 0 and b > 0: + # Warn if target is outside the calibrated range + min_s = params.get("min_observed_sparsity") + max_s = params.get("max_observed_sparsity") + if min_s is not None and target < min_s: + warnings.warn( + f"Target sparsity {target:.1%} is below the minimum observed " + f"during calibration ({min_s:.1%}). The model is extrapolating " + f"and runtime sparsity will likely be higher than the target.", + stacklevel=2, + ) + elif max_s is not None and target > max_s: + warnings.warn( + f"Target sparsity {target:.1%} is above the maximum observed " + f"during calibration ({max_s:.1%}). The model is extrapolating.", + stacklevel=2, + ) + return a * math.exp(b * target) + return None + + @staticmethod + @contextmanager + def _get_diffusers_backend_context(): + """Activate the modelopt_triton diffusers backend if registered.""" + try: + from ..kernels.diffusers_triton_attention import get_triton_attention_backend + + with get_triton_attention_backend(): + yield + except (ImportError, RuntimeError): + yield + + def _set_triton_backends(self, **kwargs): + """Set config on both diffusers and LTX Triton backends.""" + try: + from ..kernels.diffusers_triton_attention import set_triton_skip_softmax_config + + set_triton_skip_softmax_config(**kwargs) + except ImportError: + pass + try: + from ..kernels.ltx_triton_attention import set_ltx_triton_context + + set_ltx_triton_context(active=True, **kwargs) + except ImportError: + pass + + def _clear_triton_backends(self): + """Clear config on both Triton backends.""" + try: + from ..kernels.diffusers_triton_attention import clear_triton_skip_softmax_config + + clear_triton_skip_softmax_config() + except ImportError: + pass + try: + from ..kernels.ltx_triton_attention import clear_ltx_triton_context + + clear_ltx_triton_context() + except ImportError: + pass + + def _collect_calibration_stats(self, module): + """Read Triton calibration counters and store as stats on the module.""" + counters = None + seq_k = None + + try: + from ..kernels.diffusers_triton_attention import ( + get_calibration_counters, + get_calibration_seq_k, + ) + + counters = get_calibration_counters() + seq_k = get_calibration_seq_k() + except ImportError: + pass + + if counters is None: + try: + from ..kernels.ltx_triton_attention import ( + get_calibration_counters, + get_calibration_seq_k, + ) + + counters = get_calibration_counters() + seq_k = get_calibration_seq_k() + except ImportError: + pass + + if counters is None or self._threshold_trials is None: + return + + # counters: [num_thresholds, 2] — [:, 0]=total, [:, 1]=skipped + total = counters[:, 0].float() + skipped = counters[:, 1].float() + sparsity_list = (skipped / total.clamp(min=1)).tolist() + + # Use actual KV sequence length from backend for the exponential model fit. + # The calibrator uses: scale_factor = threshold * sample_length, so this + # must be the real sequence length, not the total tile count. + sample_length = seq_k if seq_k is not None else 0 + + module._last_stats = { + "sparsity": sparsity_list, + "sample_length": sample_length, + "phase": "prefill", + } + + def get_threshold_info(self) -> dict: + """Get threshold information for debugging/display.""" + scale_factor = self._get_scale_factor() + if scale_factor is not None: + return { + "type": "dynamic_calibrated", + "formula": "threshold = scale_factor / seq_k (computed at runtime)", + "scale_factor": scale_factor, + "calibration_params": self.calibration_params, + "target_sparse_ratio": self.target_sparse_ratio, + } + return { + "type": "static", + "value": self.skip_softmax_threshold, + } + + # ------------------------------------------------------------------ + # Runtime sparsity measurement + # ------------------------------------------------------------------ + + def enable_measure_sparsity(self, enabled: bool = True) -> None: + """Enable or disable runtime sparsity measurement.""" + self._measure_sparsity = enabled + + def reset_sparsity_counters(self) -> None: + """Reset accumulated sparsity counters to zero.""" + self._sparsity_total = 0 + self._sparsity_skipped = 0 + + def get_sparsity_counters(self) -> tuple[int, int]: + """Return accumulated ``(total_tiles, skipped_tiles)``.""" + return self._sparsity_total, self._sparsity_skipped + + def _collect_sparsity_counters(self) -> None: + """Read runtime sparsity counters from the backend and accumulate.""" + try: + from ..kernels.diffusers_triton_attention import get_sparsity_counters - return _skip_softmax_context() + total, skipped = get_sparsity_counters() + self._sparsity_total += total + self._sparsity_skipped += skipped + except ImportError: + pass diff --git a/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py b/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py index 599832943dc..d26b73f0b4e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py +++ b/modelopt/torch/sparsity/attention_sparsity/plugins/huggingface.py @@ -16,7 +16,6 @@ """Dynamic sparse attention registration for HuggingFace models.""" import torch.nn as nn -import transformers from modelopt.torch.opt.dynamic import DynamicModule @@ -112,11 +111,22 @@ def _is_supported_model(model: nn.Module) -> bool: """ # Check for HuggingFace PreTrainedModel try: + import transformers + if isinstance(model, transformers.PreTrainedModel): return True except ImportError: pass + # Check for diffusers ModelMixin + try: + from diffusers.models.modeling_utils import ModelMixin + + if isinstance(model, ModelMixin): + return True + except ImportError: + pass + # Support any PyTorch model with attention modules return isinstance(model, nn.Module) diff --git a/modelopt/torch/sparsity/attention_sparsity/stats_manager.py b/modelopt/torch/sparsity/attention_sparsity/stats_manager.py index 1eabdfe3586..3b8d9e2b92e 100644 --- a/modelopt/torch/sparsity/attention_sparsity/stats_manager.py +++ b/modelopt/torch/sparsity/attention_sparsity/stats_manager.py @@ -66,12 +66,13 @@ def collect(self, stats: dict): self.aggregated_stats["total_calls"] += 1 self.aggregated_stats["total_blocks"] += stats.get("total_blocks", 0) - incoming = stats["sparse_blocks"] - if "sparse_blocks" not in self.aggregated_stats: - self.aggregated_stats["sparse_blocks"] = list(incoming) - else: - for i, val in enumerate(incoming): - self.aggregated_stats["sparse_blocks"][i] += val + incoming = stats.get("sparse_blocks") + if incoming is not None: + if "sparse_blocks" not in self.aggregated_stats: + self.aggregated_stats["sparse_blocks"] = list(incoming) + else: + for i, val in enumerate(incoming): + self.aggregated_stats["sparse_blocks"][i] += val phase = stats.get("phase", "unknown") if phase in self.aggregated_stats["phase_counts"]: @@ -79,14 +80,15 @@ def collect(self, stats: dict): # In calibration mode, store per-sample stats if self.calibration_mode: - self.per_sample_stats.append( - { - "module": self.module_name, - "sparsity": stats.get("sparsity", 0.0), - "sample_length": stats.get("sample_length", 0), - "phase": phase, - } - ) + sample_stat = { + "module": self.module_name, + "sparsity": stats.get("sparsity", 0.0), + "sample_length": stats.get("sample_length", 0), + "phase": phase, + } + if "normalized_gaps" in stats: + sample_stat["normalized_gaps"] = stats["normalized_gaps"] + self.per_sample_stats.append(sample_stat) def get_summary(self) -> dict: """Get aggregated statistics summary. diff --git a/tests/_test_utils/torch/diffusers_models.py b/tests/_test_utils/torch/diffusers_models.py index c2f9d9b3d79..352cc60d793 100644 --- a/tests/_test_utils/torch/diffusers_models.py +++ b/tests/_test_utils/torch/diffusers_models.py @@ -32,6 +32,16 @@ except Exception: # pragma: no cover - optional diffusers models Flux2Transformer2DModel = None +try: + from diffusers.models.transformers import WanTransformer3DModel +except Exception: # pragma: no cover - optional diffusers models + WanTransformer3DModel = None + +try: + from diffusers.models.autoencoders import AutoencoderKLWan +except Exception: # pragma: no cover - optional diffusers models + AutoencoderKLWan = None + import modelopt.torch.opt as mto @@ -157,3 +167,86 @@ def df_modelopt_state_and_output_tester(model_ref, model_test): assert model_ref_state == model_test_state df_output_tester(model_ref, model_test) + + +def get_tiny_wan22_transformer(**config_kwargs): + """Create a tiny WanTransformer3DModel for testing.""" + if WanTransformer3DModel is None: + pytest.skip("WanTransformer3DModel is not available in this diffusers version.") + + kwargs = { + "patch_size": (1, 2, 2), + "num_attention_heads": 2, + "attention_head_dim": 12, + "in_channels": 16, + "out_channels": 16, + "text_dim": 32, + "freq_dim": 256, + "ffn_dim": 32, + "num_layers": 2, + "cross_attn_norm": True, + "qk_norm": "rms_norm_across_heads", + "rope_max_seq_len": 32, + } + kwargs.update(**config_kwargs) + return WanTransformer3DModel(**kwargs) + + +def get_tiny_wan22_vae(**config_kwargs): + """Create a tiny AutoencoderKLWan for testing.""" + if AutoencoderKLWan is None: + pytest.skip("AutoencoderKLWan is not available in this diffusers version.") + + kwargs = { + "base_dim": 3, + "z_dim": 16, + "dim_mult": [1, 1, 1, 1], + "num_res_blocks": 1, + "temperal_downsample": [False, True, True], + } + kwargs.update(**config_kwargs) + return AutoencoderKLWan(**kwargs) + + +def create_tiny_wan22_pipeline_dir(tmp_path: Path) -> Path: + """Create and save a tiny Wan 2.2 (14B-style) pipeline to a directory. + + Uses the same tiny config as diffusers' own Wan 2.2 tests: + - Transformer: 2 heads, 12 head_dim, 2 layers (hidden_dim=24) + - VAE: base_dim=3, z_dim=16 + - Text encoder: hf-internal-testing/tiny-random-t5 (hidden_size=32) + - Dual transformer (14B style) with boundary_ratio=0.875 + + The saved directory can be loaded with ``WanPipeline.from_pretrained(path)``. + """ + from diffusers import UniPCMultistepScheduler, WanPipeline + from transformers import AutoTokenizer, T5EncoderModel + + torch.manual_seed(0) + vae = get_tiny_wan22_vae() + + torch.manual_seed(0) + transformer = get_tiny_wan22_transformer() + + torch.manual_seed(0) + transformer_2 = get_tiny_wan22_transformer() + + scheduler = UniPCMultistepScheduler( + prediction_type="flow_prediction", use_flow_sigmas=True, flow_shift=3.0 + ) + text_encoder = T5EncoderModel.from_pretrained("hf-internal-testing/tiny-random-t5") + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-t5") + + pipe = WanPipeline( + transformer=transformer, + transformer_2=transformer_2, + vae=vae, + scheduler=scheduler, + text_encoder=text_encoder, + tokenizer=tokenizer, + boundary_ratio=0.875, + ) + + save_dir = tmp_path / "tiny_wan22" + pipe.save_pretrained(save_dir) + return save_dir diff --git a/tests/examples/diffusers_sparsity/test_sparsity.py b/tests/examples/diffusers_sparsity/test_sparsity.py new file mode 100644 index 00000000000..d33be1df68a --- /dev/null +++ b/tests/examples/diffusers_sparsity/test_sparsity.py @@ -0,0 +1,104 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for skip-softmax sparse attention on Wan 2.2 (examples/diffusers/sparsity/). + +Uses a tiny Wan 2.2 model (dual transformer, 2 layers, hidden_dim=24) created +from scratch. Tests run the wan22_skip_softmax.py example script in baseline, +triton-baseline, and raw-threshold modes. +""" + +import pytest +from _test_utils.examples.run_command import run_example_command +from _test_utils.torch.diffusers_models import create_tiny_wan22_pipeline_dir + +EXAMPLE_PATH = "diffusers/sparsity" + +# Tiny inference settings — fast but exercises all code paths +_TINY_ARGS = [ + "--num-frames", + "5", + "--height", + "16", + "--width", + "16", + "--num-steps", + "2", + "--guidance-scale", + "1.0", + "--skip-first-last", + "0", + "--negative-prompt", + "", +] + + +@pytest.fixture(scope="session") +def tiny_wan22_path(tmp_path_factory): + """Create a tiny Wan 2.2 pipeline saved to disk (session-scoped).""" + return str(create_tiny_wan22_pipeline_dir(tmp_path_factory.mktemp("tiny_wan22"))) + + +def test_wan22_baseline(tiny_wan22_path, tmp_path): + """Dense baseline — no sparsity, default diffusers attention backend.""" + cmd = [ + "python", + "wan22_skip_softmax.py", + "--model-path", + tiny_wan22_path, + "--baseline", + "--prompt", + "test", + "--output", + str(tmp_path / "baseline.mp4"), + *_TINY_ARGS, + ] + run_example_command(cmd, EXAMPLE_PATH) + + +def test_wan22_triton_baseline(tiny_wan22_path, tmp_path): + """Triton kernel without skip-softmax (threshold=0, apples-to-apples).""" + cmd = [ + "python", + "wan22_skip_softmax.py", + "--model-path", + tiny_wan22_path, + "--triton-baseline", + "--prompt", + "test", + "--output", + str(tmp_path / "triton_baseline.mp4"), + *_TINY_ARGS, + ] + run_example_command(cmd, EXAMPLE_PATH) + + +def test_wan22_raw_threshold(tiny_wan22_path, tmp_path): + """Skip-softmax with a fixed raw threshold — no calibration needed.""" + cmd = [ + "python", + "wan22_skip_softmax.py", + "--model-path", + tiny_wan22_path, + "--raw-threshold", + "-5.0", + "--report-avg-sparsity", + "--prompt", + "test", + "--output", + str(tmp_path / "raw_threshold.mp4"), + *_TINY_ARGS, + ] + run_example_command(cmd, EXAMPLE_PATH) diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py b/tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py new file mode 100644 index 00000000000..f479b8883f9 --- /dev/null +++ b/tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU tests for diffusers and LTX Triton attention wrappers.""" + +import pytest +import torch + +pytestmark = [ + pytest.mark.filterwarnings("ignore::UserWarning"), + pytest.mark.filterwarnings("ignore::RuntimeWarning"), + pytest.mark.filterwarnings("ignore::DeprecationWarning"), +] + +diffusers = pytest.importorskip("diffusers") + +from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE +from modelopt.torch.sparsity.attention_sparsity.kernels import ( + diffusers_triton_attention as diffusers_mod, +) +from modelopt.torch.sparsity.attention_sparsity.kernels import ltx_triton_attention as ltx_mod + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestDiffusersTritonAttention: + """Exercise _diffusers_triton_attention on a real device.""" + + @pytest.fixture(autouse=True) + def _reset_thread_local(self): + diffusers_mod.clear_triton_skip_softmax_config() + yield + diffusers_mod.clear_triton_skip_softmax_config() + + def _make_qkv(self, b=1, seq_q=128, seq_k=128, h=4, d=64, dtype=torch.float16): + q = torch.randn(b, seq_q, h, d, device="cuda", dtype=dtype) + k = torch.randn(b, seq_k, h, d, device="cuda", dtype=dtype) + v = torch.randn(b, seq_k, h, d, device="cuda", dtype=dtype) + return q, k, v + + def test_basic_forward(self): + q, k, v = self._make_qkv() + out = diffusers_mod._diffusers_triton_attention(q, k, v) + assert out.shape == q.shape + assert not torch.isnan(out).any() + + def test_skip_softmax_threshold_path(self): + diffusers_mod.set_triton_skip_softmax_config(threshold=0.01) + q, k, v = self._make_qkv() + out = diffusers_mod._diffusers_triton_attention(q, k, v) + assert out.shape == q.shape + + def test_raw_threshold_path(self): + diffusers_mod.set_triton_skip_softmax_config(raw_threshold=-10.0) + q, k, v = self._make_qkv() + out = diffusers_mod._diffusers_triton_attention(q, k, v) + assert out.shape == q.shape + + def test_scale_factor_path(self): + diffusers_mod.set_triton_skip_softmax_config(scale_factor=2.0) + q, k, v = self._make_qkv() + out = diffusers_mod._diffusers_triton_attention(q, k, v) + assert out.shape == q.shape + + def test_measure_sparsity_returns_counts(self): + diffusers_mod.set_triton_skip_softmax_config(threshold=0.1, measure_sparsity=True) + q, k, v = self._make_qkv(seq_q=512, seq_k=512) + diffusers_mod._diffusers_triton_attention(q, k, v) + total, _skipped = diffusers_mod.get_sparsity_counters() + assert total > 0 + + def test_calibration_mode(self): + diffusers_mod.set_triton_skip_softmax_config( + calibration_mode=True, + threshold_trials=[1e-3, 1e-2, 1e-1], + ) + q, k, v = self._make_qkv(seq_q=128, seq_k=128) + out = diffusers_mod._diffusers_triton_attention(q, k, v) + assert out.shape == q.shape + counters = diffusers_mod.get_calibration_counters() + assert counters is not None + assert counters.shape == (3, 2) + assert diffusers_mod.get_calibration_seq_k() == 128 + + def test_cross_attention_different_seq_lengths(self): + q, k, v = self._make_qkv(seq_q=128, seq_k=256) + out = diffusers_mod._diffusers_triton_attention(q, k, v) + assert out.shape == q.shape + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestLTXTritonAttention: + """Exercise _ltx_triton_attention (LTX layout [B, T, H*D]).""" + + @pytest.fixture(autouse=True) + def _reset_thread_local(self): + ltx_mod.clear_ltx_triton_context() + yield + ltx_mod.clear_ltx_triton_context() + + def _make_qkv(self, b=1, seq_q=128, seq_k=128, heads=4, dim_head=64, dtype=torch.float16): + dim = heads * dim_head + q = torch.randn(b, seq_q, dim, device="cuda", dtype=dtype) + k = torch.randn(b, seq_k, dim, device="cuda", dtype=dtype) + v = torch.randn(b, seq_k, dim, device="cuda", dtype=dtype) + return q, k, v + + def test_inference_basic(self): + q, k, v = self._make_qkv() + out = ltx_mod._ltx_triton_attention(q, k, v, heads=4) + assert out.shape == q.shape + assert not torch.isnan(out).any() + + def test_inference_with_raw_threshold(self): + ltx_mod.set_ltx_triton_context(active=True, raw_threshold=-10.0) + q, k, v = self._make_qkv() + out = ltx_mod._ltx_triton_attention(q, k, v, heads=4) + assert out.shape == q.shape + + def test_inference_with_scale_factor(self): + ltx_mod.set_ltx_triton_context(active=True, scale_factor=5.0) + q, k, v = self._make_qkv() + out = ltx_mod._ltx_triton_attention(q, k, v, heads=4) + assert out.shape == q.shape + + def test_inference_with_static_threshold(self): + q, k, v = self._make_qkv() + out = ltx_mod._ltx_triton_attention(q, k, v, heads=4, threshold=0.1) + assert out.shape == q.shape + + def test_cross_attention_different_seq(self): + q, k, v = self._make_qkv(seq_q=128, seq_k=256) + out = ltx_mod._ltx_triton_attention(q, k, v, heads=4) + assert out.shape == q.shape + + def test_calibration_mode(self): + ltx_mod.set_ltx_triton_context( + active=True, + calibration_mode=True, + threshold_trials=[1e-3, 1e-2, 1e-1], + ) + q, k, v = self._make_qkv(seq_q=128, seq_k=128) + out = ltx_mod._ltx_triton_attention(q, k, v, heads=4) + assert out.shape == q.shape + counters = ltx_mod.get_calibration_counters() + assert counters is not None + assert counters.shape == (3, 2) + assert ltx_mod.get_calibration_seq_k() == 128 + + def test_wrapper_dispatch(self): + """TritonLTXAttentionWrapper dispatches based on thread-local flag.""" + called = {"original": 0} + + def original_fn(q, k, v, heads, mask=None): + called["original"] += 1 + return q + + wrapper = ltx_mod._TritonLTXAttentionWrapper(original_fn) + q, k, v = self._make_qkv() + + # When inactive, original_fn is called + ltx_mod.clear_ltx_triton_context() + wrapper(q, k, v, heads=4) + assert called["original"] == 1 + + # When active, triton path runs + ltx_mod.set_ltx_triton_context(active=True, threshold=0.1) + out = wrapper(q, k, v, heads=4) + assert called["original"] == 1 # unchanged + assert out.shape == q.shape diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.py b/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.py new file mode 100644 index 00000000000..37c4da9969c --- /dev/null +++ b/tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa_calibrate.py @@ -0,0 +1,287 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""GPU tests for the Triton flash attention calibration kernel. + +Exercises ``attention_calibrate`` which computes full attention while counting +how many KV tiles would be skipped at each threshold in ``threshold_trials``. +""" + +import pytest +import torch +from conftest import make_qkv, make_varlen_meta + +pytestmark = [ + pytest.mark.filterwarnings("ignore::UserWarning"), + pytest.mark.filterwarnings("ignore::RuntimeWarning"), + pytest.mark.filterwarnings("ignore::DeprecationWarning"), +] + +from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + from modelopt.torch.kernels import attention, attention_calibrate + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestAttentionCalibrate: + """Multi-threshold sparsity measurement kernel.""" + + def _make_inputs(self, batch=1, seq_len=256, num_heads=4, head_dim=64): + total = batch * seq_len + torch.manual_seed(42) + q = torch.randn(total, num_heads, head_dim, device="cuda", dtype=torch.float16) + k = torch.randn(total, num_heads, head_dim, device="cuda", dtype=torch.float16) + v = torch.randn(total, num_heads, head_dim, device="cuda", dtype=torch.float16) + locs, lens = make_varlen_meta([seq_len] * batch) + return q, k, v, locs, lens + + def test_output_matches_dense(self): + """Calibration kernel computes full attention — output should match dense.""" + q, k, v, locs, lens = self._make_inputs() + scale = 1.0 / (64**0.5) + out_dense = attention(q, k, v, locs, lens, 256, softmax_scale=scale, is_causal=False) + out_calib, counters = attention_calibrate( + q, + k, + v, + locs, + lens, + 256, + softmax_scale=scale, + is_causal=False, + threshold_trials=[1e-3, 1e-2, 1e-1], + ) + assert out_calib.shape == q.shape + # Online softmax differences between dense and calibrate kernel are within a small tol + torch.testing.assert_close(out_calib, out_dense, rtol=5e-3, atol=5e-3) + + def test_counter_shape_and_values(self): + """Counters have shape [num_thresholds, 2] and sane values.""" + q, k, v, locs, lens = self._make_inputs() + scale = 1.0 / (64**0.5) + trials = [1e-4, 1e-2, 1e-1, 5e-1] + _, counters = attention_calibrate( + q, + k, + v, + locs, + lens, + 256, + softmax_scale=scale, + is_causal=False, + threshold_trials=trials, + ) + assert counters.shape == (len(trials), 2) + totals = counters[:, 0] + skipped = counters[:, 1] + # Totals are equal across thresholds (every tile evaluated for every threshold) + assert (totals == totals[0]).all() + # Skipped counts monotonically increase with threshold + skipped_list = skipped.tolist() + assert all(skipped_list[i] <= skipped_list[i + 1] for i in range(len(skipped_list) - 1)) + # No tile can be skipped more than total + assert (skipped <= totals).all() + + def test_different_seq_q_seq_k(self): + """Cross-attention varlen with separate Q and K/V metadata.""" + batch = 1 + seq_q, seq_k = 128, 256 + num_heads, head_dim = 4, 64 + scale = 1.0 / (head_dim**0.5) + + torch.manual_seed(11) + q = torch.randn(seq_q * batch, num_heads, head_dim, device="cuda", dtype=torch.float16) + k = torch.randn(seq_k * batch, num_heads, head_dim, device="cuda", dtype=torch.float16) + v = torch.randn(seq_k * batch, num_heads, head_dim, device="cuda", dtype=torch.float16) + b_start_loc = torch.arange(batch, device="cuda", dtype=torch.int32) * seq_q + b_seq_len = torch.full((batch,), seq_q, device="cuda", dtype=torch.int32) + b_start_loc_k = torch.arange(batch, device="cuda", dtype=torch.int32) * seq_k + b_seq_len_k = torch.full((batch,), seq_k, device="cuda", dtype=torch.int32) + + out, counters = attention_calibrate( + q, + k, + v, + b_start_loc, + b_seq_len, + seq_q, + softmax_scale=scale, + is_causal=False, + b_start_loc_k=b_start_loc_k, + b_seq_len_k=b_seq_len_k, + max_input_len_k=seq_k, + threshold_trials=[1e-2, 1e-1], + ) + assert out.shape == q.shape + assert counters.shape == (2, 2) + + def test_threshold_order_doesnt_affect_counts(self): + """Skipped counts at the same threshold are independent of trial ordering.""" + q, k, v, locs, lens = self._make_inputs() + scale = 1.0 / (64**0.5) + _, c1 = attention_calibrate( + q, + k, + v, + locs, + lens, + 256, + softmax_scale=scale, + is_causal=False, + threshold_trials=[1e-3, 1e-1], + ) + _, c2 = attention_calibrate( + q, + k, + v, + locs, + lens, + 256, + softmax_scale=scale, + is_causal=False, + threshold_trials=[1e-1, 1e-3], + ) + # Both runs measure the same two thresholds — the skipped counts should match + # after permuting back to the same order. + assert c1[0, 1].item() == c2[1, 1].item() + assert c1[1, 1].item() == c2[0, 1].item() + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestMeasureSparsity: + """Runtime sparsity counters during inference.""" + + def test_measure_sparsity_returns_counts(self): + """measure_sparsity=True attaches _sparsity_total/_sparsity_skipped to output.""" + torch.manual_seed(99) + batch, seq_len, num_heads, head_dim = 1, 1024, 4, 64 + total = batch * seq_len + q = torch.randn(total, num_heads, head_dim, device="cuda", dtype=torch.float16) + k = torch.randn(total, num_heads, head_dim, device="cuda", dtype=torch.float16) + v = torch.randn(total, num_heads, head_dim, device="cuda", dtype=torch.float16) + locs, lens = make_varlen_meta([seq_len] * batch) + scale = 1.0 / (head_dim**0.5) + + out = attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + is_causal=False, + skip_softmax_threshold=0.5, + measure_sparsity=True, + ) + assert hasattr(out, "_sparsity_total") + assert hasattr(out, "_sparsity_skipped") + assert out._sparsity_total > 0 + assert out._sparsity_skipped <= out._sparsity_total + + def test_measure_sparsity_without_skip_is_noop(self): + """Without skip-softmax, measure_sparsity doesn't attach counters.""" + q, k, v = make_qkv(256, 4, 4, 64, dtype=torch.float16) + locs, lens = make_varlen_meta([256]) + scale = 1.0 / (64**0.5) + + out = attention( + q, k, v, locs, lens, 256, softmax_scale=scale, is_causal=False, measure_sparsity=True + ) + # No skip-softmax active => counters should not be attached + assert not hasattr(out, "_sparsity_total") + + def test_raw_threshold_path(self): + """Raw threshold is passed directly to the kernel without conversion.""" + q, k, v = make_qkv(256, 4, 4, 64, dtype=torch.float16) + locs, lens = make_varlen_meta([256]) + scale = 1.0 / (64**0.5) + out_raw = attention( + q, + k, + v, + locs, + lens, + 256, + softmax_scale=scale, + is_causal=False, + skip_softmax_raw_threshold=-20.0, + ) + # With a very negative raw threshold, almost no tiles are skipped + # Output should be close to dense + out_dense = attention(q, k, v, locs, lens, 256, softmax_scale=scale, is_causal=False) + torch.testing.assert_close(out_raw, out_dense, rtol=1e-2, atol=1e-2) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestBackwardWithSparsity: + """Backward pass with skip-softmax (covers _attn_bwd_dq / _attn_bwd_dkdv).""" + + def test_backward_with_skip_softmax(self): + """Backward pass runs without error when skip-softmax is active.""" + seq_len, num_heads, head_dim = 128, 4, 64 + scale = 1.0 / (head_dim**0.5) + torch.manual_seed(7) + q, k, v = make_qkv(seq_len, num_heads, num_heads, head_dim, dtype=torch.float32) + q.requires_grad_(True) + k.requires_grad_(True) + v.requires_grad_(True) + locs, lens = make_varlen_meta([seq_len]) + + out = attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + is_causal=True, + skip_softmax_threshold=1e-3, + ) + out.sum().backward() + assert q.grad is not None + assert k.grad is not None + assert v.grad is not None + assert not torch.isnan(q.grad).any() + assert not torch.isnan(k.grad).any() + assert not torch.isnan(v.grad).any() + + def test_backward_with_sparsity_nm(self): + """Backward pass with 2:4 N:M sparsity runs without error.""" + seq_len, num_heads, head_dim = 128, 4, 64 + scale = 1.0 / (head_dim**0.5) + torch.manual_seed(13) + q, k, v = make_qkv(seq_len, num_heads, num_heads, head_dim, dtype=torch.float32) + q.requires_grad_(True) + k.requires_grad_(True) + v.requires_grad_(True) + locs, lens = make_varlen_meta([seq_len]) + + out = attention( + q, + k, + v, + locs, + lens, + seq_len, + softmax_scale=scale, + is_causal=True, + sparsity_n=2, + sparsity_m=4, + ) + out.sum().backward() + assert q.grad is not None diff --git a/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py b/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py new file mode 100644 index 00000000000..72b20df9329 --- /dev/null +++ b/tests/gpu/torch/sparsity/attention_sparsity/test_wan22_skip_softmax.py @@ -0,0 +1,280 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""End-to-end tests for skip-softmax sparse attention on a tiny Wan 2.2 pipeline. + +Uses ``create_tiny_wan22_pipeline_dir`` (matches diffusers' own tiny Wan 2.2 +test config): dual 2-layer transformer, tiny VAE, tiny T5 text encoder. +The full ``WanPipeline`` is loaded, sparsified with ``mtsa.sparsify``, and +run end-to-end with a 2-step denoising loop — asserting no NaN/Inf in the +pipeline output. +""" + +import pytest +import torch + +pytestmark = [ + pytest.mark.filterwarnings("ignore::UserWarning"), + pytest.mark.filterwarnings("ignore::RuntimeWarning"), + pytest.mark.filterwarnings("ignore::DeprecationWarning"), +] + +diffusers = pytest.importorskip("diffusers") + +from modelopt.torch.kernels import IS_AVAILABLE as TRITON_KERNEL_AVAILABLE + +if TRITON_KERNEL_AVAILABLE: + import modelopt.torch.sparsity.attention_sparsity as mtsa + from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule + + +# --------------------------------------------------------------------------- +# Tiny Wan 2.2 pipeline fixture — shared across tests (pipeline load is costly) +# --------------------------------------------------------------------------- + + +@pytest.fixture(scope="module") +def tiny_wan22_path(tmp_path_factory): + """Create and save a tiny Wan 2.2 pipeline to disk once per module.""" + from _test_utils.torch.diffusers_models import create_tiny_wan22_pipeline_dir + + return str(create_tiny_wan22_pipeline_dir(tmp_path_factory.mktemp("tiny_wan22"))) + + +@pytest.fixture +def tiny_wan22_pipe(tiny_wan22_path): + """Load a fresh copy of the tiny Wan 2.2 pipeline on CUDA (per test).""" + from diffusers import WanPipeline + + pipe = WanPipeline.from_pretrained(tiny_wan22_path, torch_dtype=torch.bfloat16) + pipe.to("cuda") + return pipe + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +_TINY_PIPE_KWARGS = { + "prompt": "test", + "negative_prompt": "", + "num_frames": 5, + "height": 16, + "width": 16, + "num_inference_steps": 2, + "guidance_scale": 1.0, +} + + +def _skip_softmax_cfg(raw_threshold=-5.0): + """Sparse config targeting Wan 2.2 self-attention (attn1) only.""" + return { + "sparse_cfg": { + "*attn1*": { + "method": "triton_skip_softmax", + "backend": "triton", + "skip_softmax_raw_threshold": raw_threshold, + "enable": True, + }, + "default": {"enable": False}, + }, + } + + +def _sparsify_both_transformers(pipe, cfg): + """Apply sparsify to both transformer and transformer_2 (14B-style dual).""" + mtsa.sparsify(pipe.transformer, cfg) + mtsa.sparsify(pipe.transformer_2, cfg) + + +def _run_pipe(pipe, seed=0): + """Run the pipeline with a fixed seed; return output frames tensor.""" + generator = torch.Generator(device="cuda").manual_seed(seed) + with torch.no_grad(): + output = pipe(generator=generator, **_TINY_PIPE_KWARGS) + # output.frames[0] is a list of PIL images; for assertion we just need shape+health + return output + + +def _count_sparse_modules(module): + return sum(isinstance(m, SparseAttentionModule) for m in module.modules()) + + +# --------------------------------------------------------------------------- +# E2E tests +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestWan22PipelineE2E: + """End-to-end skip-softmax flow on a tiny Wan 2.2 pipeline.""" + + def test_baseline_pipeline_runs(self, tiny_wan22_pipe): + """Dense baseline: pipeline produces finite frames without sparsification.""" + output = _run_pipe(tiny_wan22_pipe) + assert output.frames is not None + assert len(output.frames) == 1 + # Frames are PIL Images; ensure list isn't empty + assert len(output.frames[0]) > 0 + + def test_sparsify_inserts_modules_in_both_transformers(self, tiny_wan22_pipe): + """Both transformer and transformer_2 get SparseAttentionModule instances.""" + _sparsify_both_transformers(tiny_wan22_pipe, _skip_softmax_cfg()) + assert _count_sparse_modules(tiny_wan22_pipe.transformer) > 0 + assert _count_sparse_modules(tiny_wan22_pipe.transformer_2) > 0 + + def test_skip_softmax_pipeline_runs_e2e(self, tiny_wan22_pipe): + """Sparsified pipeline runs end-to-end producing finite frames.""" + _sparsify_both_transformers(tiny_wan22_pipe, _skip_softmax_cfg(raw_threshold=-5.0)) + output = _run_pipe(tiny_wan22_pipe) + assert output.frames is not None + assert len(output.frames[0]) > 0 + + def test_tight_threshold_matches_dense_within_tolerance(self, tiny_wan22_pipe, tiny_wan22_path): + """raw_threshold=-50 (effectively dense) → output close to unsparsified run.""" + from diffusers import WanPipeline + + # Dense run: fresh pipe, no sparsification + dense_pipe = WanPipeline.from_pretrained(tiny_wan22_path, torch_dtype=torch.bfloat16) + dense_pipe.to("cuda") + dense_frame0 = _run_pipe(dense_pipe).frames[0][0] + + # Sparse run: same seed, raw_threshold=-50 (≈ no tiles skipped) + _sparsify_both_transformers(tiny_wan22_pipe, _skip_softmax_cfg(raw_threshold=-50.0)) + sparse_frame0 = _run_pipe(tiny_wan22_pipe).frames[0][0] + + # Both are PIL images — convert to tensor and compare + import numpy as np + + d = np.asarray(dense_frame0, dtype=np.float32) + s = np.asarray(sparse_frame0, dtype=np.float32) + # Pixel-wise MAE should be small for tight threshold (but not bit-exact due to + # different code paths in the online softmax accumulation). + mae = np.abs(d - s).mean() + assert mae < 20.0, f"MAE between dense and tight-sparse frames was {mae:.2f}" + + def test_measure_sparsity_counts_accumulate(self, tiny_wan22_pipe): + """measure_sparsity=True + a permissive threshold → nonzero sparsity counters.""" + from modelopt.torch.sparsity.attention_sparsity.methods.triton_skip_softmax import ( + TritonSkipSoftmaxMethod, + ) + + _sparsify_both_transformers(tiny_wan22_pipe, _skip_softmax_cfg(raw_threshold=-2.0)) + + # Enable measurement + reset counters on every sparse module + for module in (tiny_wan22_pipe.transformer, tiny_wan22_pipe.transformer_2): + for m in module.modules(): + if isinstance(m, SparseAttentionModule): + method = m._sparse_method_instance + if isinstance(method, TritonSkipSoftmaxMethod): + method.enable_measure_sparsity(True) + method.reset_sparsity_counters() + + _run_pipe(tiny_wan22_pipe) + + # Sum counters across all sparse modules + total_sum = 0 + for module in (tiny_wan22_pipe.transformer, tiny_wan22_pipe.transformer_2): + for m in module.modules(): + if isinstance(m, SparseAttentionModule): + method = m._sparse_method_instance + if isinstance(method, TritonSkipSoftmaxMethod): + total, skipped = method.get_sparsity_counters() + assert skipped <= total + total_sum += total + + assert total_sum > 0, "Expected nonzero sparsity counters after pipeline run" + + def test_save_restore_roundtrip(self, tiny_wan22_pipe): + """Sparsified transformer saves & restores via modelopt_state, preserving + per-module method choice. The ``*attn1*`` pattern maps to triton_skip_softmax; + ``attn2`` modules keep the default method. The restored model must show the + identical (module_name → method) mapping. + """ + from _test_utils.torch.diffusers_models import get_tiny_wan22_transformer + + import modelopt.torch.opt as mto + + _sparsify_both_transformers(tiny_wan22_pipe, _skip_softmax_cfg()) + state = mto.modelopt_state(tiny_wan22_pipe.transformer) + + # Restore into a fresh transformer of the same shape + torch.manual_seed(0) + restored = get_tiny_wan22_transformer().to("cuda", dtype=torch.bfloat16).eval() + mto.restore_from_modelopt_state(restored, state) + + def _method_map(module): + return { + name: m._method + for name, m in module.named_modules() + if isinstance(m, SparseAttentionModule) + } + + orig_map = _method_map(tiny_wan22_pipe.transformer) + restored_map = _method_map(restored) + assert orig_map == restored_map, ( + f"Restored method map differs from original:\n" + f" orig: {orig_map}\n" + f" restored: {restored_map}" + ) + # Sanity: the ``*attn1*`` pattern should have produced at least one + # triton_skip_softmax module in the restored model. + triton_attn1 = [ + name + for name, method in restored_map.items() + if method == "triton_skip_softmax" and "attn1" in name + ] + assert triton_attn1, ( + f"Expected at least one attn1 module with triton_skip_softmax after restore, " + f"got {restored_map}" + ) + + +@pytest.mark.skipif(not TRITON_KERNEL_AVAILABLE, reason="Need CUDA + triton") +class TestWan22Calibration: + """Multi-threshold calibration path on a tiny Wan 2.2 transformer.""" + + def test_calibration_collects_stats_per_module(self, tiny_wan22_pipe): + """A forward pass under calibration_mode populates per-module _last_stats.""" + from modelopt.torch.sparsity.attention_sparsity.methods.triton_skip_softmax import ( + TritonSkipSoftmaxMethod, + ) + + _sparsify_both_transformers(tiny_wan22_pipe, _skip_softmax_cfg()) + + threshold_trials = [1e-3, 1e-2, 1e-1] + for module in (tiny_wan22_pipe.transformer, tiny_wan22_pipe.transformer_2): + for m in module.modules(): + if isinstance(m, SparseAttentionModule): + method = m._sparse_method_instance + if isinstance(method, TritonSkipSoftmaxMethod): + method._calibration_mode = True + method._threshold_trials = threshold_trials + + _run_pipe(tiny_wan22_pipe) + + # At least one sparse module should report stats of the correct shape + found_stats = False + for module in (tiny_wan22_pipe.transformer, tiny_wan22_pipe.transformer_2): + for m in module.modules(): + if isinstance(m, SparseAttentionModule) and m._last_stats is not None: + stats = m._last_stats + assert len(stats["sparsity"]) == len(threshold_trials) + found_stats = True + break + if found_stats: + break + assert found_stats, "No sparse module reported calibration stats" diff --git a/tests/unit/torch/kernels/test_triton_fa.py b/tests/unit/torch/kernels/test_triton_fa.py new file mode 100644 index 00000000000..ac054e10e6b --- /dev/null +++ b/tests/unit/torch/kernels/test_triton_fa.py @@ -0,0 +1,39 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU smoke tests for the Triton flash attention module. + +The ``@triton.jit`` kernels and the ``attention`` / ``attention_calibrate`` +Python wrappers require a GPU and are fully exercised in +``tests/gpu/torch/sparsity/attention_sparsity/test_triton_fa*.py``. + +This file only verifies that the module is importable on CPU-only CI runners, +so upstream code paths that conditionally import it don't break. +""" + +import pytest + + +def test_triton_fa_importable_on_cpu(): + """Module imports cleanly without CUDA; exports the public API names.""" + try: + import triton # noqa: F401 + except ImportError: + pytest.skip("triton is not installed") + + from modelopt.torch.kernels import triton_fa + + assert "attention" in triton_fa.__all__ + assert "attention_calibrate" in triton_fa.__all__ diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_calibrator_fitting.py b/tests/unit/torch/sparsity/attention_sparsity/test_calibrator_fitting.py new file mode 100644 index 00000000000..c7c6f56928b --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_calibrator_fitting.py @@ -0,0 +1,183 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for DynamicThresholdCalibrator exponential fitting (no GPU required). + +Tests the calibration math (curve_fit, filtering, aggregation) using synthetic +data injected via mock forward loops and mock sparse attention modules. +""" + +import numpy as np +import pytest + +pytest.importorskip("transformers") + + +from _test_utils.torch.sparsity.sparse_attention_common import SimpleAttentionModel + +from modelopt.torch.sparsity.attention_sparsity import sparsify +from modelopt.torch.sparsity.attention_sparsity.calibration.calibrator import ( + DynamicThresholdCalibrator, +) +from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule + + +class TestDynamicThresholdCalibratorInit: + def test_default_thresholds(self): + cal = DynamicThresholdCalibrator() + assert len(cal.threshold_trials) == 20 + assert cal.threshold_trials[0] == 1e-6 + assert cal.threshold_trials[-1] == 9.9e-1 + + def test_custom_thresholds(self): + trials = [0.01, 0.1, 0.5] + cal = DynamicThresholdCalibrator(threshold_trials=trials) + assert cal.threshold_trials == trials + + def test_fit_logspace(self): + cal = DynamicThresholdCalibrator(fit_logspace=True) + assert cal.fit_logspace is True + + +class TestExponentialFitting: + """Test the calibration pipeline with synthetic stats injected via mock modules.""" + + def _make_sparse_model(self): + """Create a model with flash_skip_softmax sparse attention applied.""" + model = SimpleAttentionModel(hidden_size=64, num_heads=4) + config = { + "sparse_cfg": { + "*attention*": { + "method": "flash_skip_softmax", + "thresholds": {"prefill": [0.1], "decode": [0.1]}, + "br": 64, + "bc": 64, + "enable": True, + } + }, + } + return sparsify(model, config) + + def _inject_synthetic_stats(self, model, threshold_trials, sample_length=4096): + """Inject synthetic calibration stats into the sparse modules. + + Generates stats matching the exponential model: scale_factor = a * exp(b * sparsity), + with a=0.1, b=5.0. For each threshold, the expected sparsity is: + sparsity = ln(threshold * sample_length / a) / b + """ + a_true, b_true = 0.1, 5.0 + + for module in model.modules(): + if isinstance(module, SparseAttentionModule): + sparsity_list = [] + for t in threshold_trials: + sf = t * sample_length + # Invert the model: sparsity = ln(sf / a) / b + s = np.log(sf / a_true) / b_true + s = max(0.0, min(1.0, s)) + sparsity_list.append(s) + module._last_stats = { + "sparsity": sparsity_list, + "sample_length": sample_length, + "phase": "prefill", + } + + def test_calibrate_with_synthetic_linear(self): + """Test linear-space fitting recovers known parameters.""" + model = self._make_sparse_model() + trials = [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 0.1, 0.3, 0.5] + cal = DynamicThresholdCalibrator(threshold_trials=trials, fit_logspace=False) + + def forward_loop(m): + # Simulate 3 samples with different lengths + for length in [2048, 4096, 8192]: + self._inject_synthetic_stats(m, trials, sample_length=length) + for module in m.modules(): + if isinstance(module, SparseAttentionModule) and module._stats_manager: + stats = module._last_stats + if stats: + module._stats_manager.collect(stats) + + result = cal.calibrate(model, forward_loop, "prefill") + # Should produce valid result with a and b close to ground truth + if result: + assert "a" in result + assert "b" in result + assert result["r_squared"] > 0.8 + + def test_calibrate_with_synthetic_logspace(self): + """Test log-space fitting recovers known parameters.""" + model = self._make_sparse_model() + trials = [1e-4, 5e-4, 1e-3, 5e-3, 1e-2, 5e-2, 0.1, 0.3, 0.5] + cal = DynamicThresholdCalibrator(threshold_trials=trials, fit_logspace=True) + + def forward_loop(m): + for length in [2048, 4096, 8192]: + self._inject_synthetic_stats(m, trials, sample_length=length) + for module in m.modules(): + if isinstance(module, SparseAttentionModule) and module._stats_manager: + stats = module._last_stats + if stats: + module._stats_manager.collect(stats) + + result = cal.calibrate(model, forward_loop, "prefill") + if result: + assert "a" in result + assert "b" in result + assert result["r_squared"] > 0.9 # Log-space should fit better + + def test_calibrate_no_modules_raises(self): + """Test error when no sparse attention modules exist.""" + model = SimpleAttentionModel(hidden_size=64, num_heads=4) + cal = DynamicThresholdCalibrator(threshold_trials=[0.01]) + + with pytest.raises(ValueError, match="No sparse attention modules"): + cal.calibrate(model, lambda m: None, "prefill") + + def test_calibrate_empty_stats_returns_empty(self): + """Test empty dict returned when forward loop produces no stats.""" + model = self._make_sparse_model() + cal = DynamicThresholdCalibrator(threshold_trials=[0.01]) + + result = cal.calibrate(model, lambda m: None, "prefill") + assert result == {} + + +class TestSetThresholds: + """Test _set_thresholds for both method types.""" + + def test_set_thresholds_flash_method(self): + model = SimpleAttentionModel(hidden_size=64, num_heads=4) + config = { + "sparse_cfg": { + "*attention*": { + "method": "flash_skip_softmax", + "thresholds": {"prefill": [0.1], "decode": [0.1]}, + "br": 64, + "bc": 64, + "enable": True, + } + }, + } + model = sparsify(model, config) + cal = DynamicThresholdCalibrator() + + modules = [m for m in model.modules() if isinstance(m, SparseAttentionModule)] + trials = [0.001, 0.01, 0.1] + cal._set_thresholds(modules, trials) + + for module in modules: + method = module._sparse_method_instance + assert method.thresholds == trials diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_flash_skip_softmax.py b/tests/unit/torch/sparsity/attention_sparsity/test_flash_skip_softmax.py index 5c8b7d98446..0724fa9ac03 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_flash_skip_softmax.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_flash_skip_softmax.py @@ -272,3 +272,127 @@ def test_apply_sparsity_without_mask(self): # Verify output shape matches input assert sparse_attn.shape == attn.shape + + def test_calibrated_path_prefill(self): + """Dynamic calibrated threshold path is exercised when params/targets are set.""" + method = FlashSkipSoftmax( + { + "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, + "br": 128, + "bc": 128, + "backend": "pytorch", + "is_causal": False, + } + ) + method.calibration_params = {"prefill": {"a": 1.0, "b": 5.0}} + method.target_sparse_ratio = {"prefill": 0.5} + + attn = torch.randn(1, 2, 128, 256) + mask, stats = method.calc_correction_factor_and_p(attn, "prefill") + # calibrated single-threshold path yields one sparsity entry + assert len(stats["sparsity"]) == 1 + assert mask is not None + + def test_calibrated_path_decode(self): + """Decode with calibrated params.""" + method = FlashSkipSoftmax( + { + "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, + "br": 128, + "bc": 128, + "backend": "pytorch", + "is_causal": False, + } + ) + method.calibration_params = {"decode": {"a": 0.5, "b": 4.0}} + method.target_sparse_ratio = {"decode": 0.6} + + attn = torch.randn(1, 2, 1, 256) + mask, stats = method.calc_correction_factor_and_p(attn, "decode") + assert stats["phase"] == "decode" + assert len(stats["sparsity"]) == 1 + + def test_get_threshold_info_calibrated(self): + """get_threshold_info returns dynamic_calibrated type when calibrated.""" + method = FlashSkipSoftmax( + { + "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, + "br": 128, + "bc": 128, + "backend": "pytorch", + "is_causal": True, + } + ) + method.calibration_params = {"prefill": {"a": 1.0, "b": 5.0}} + method.target_sparse_ratio = {"prefill": 0.5} + info = method.get_threshold_info() + assert info["type"] == "dynamic_calibrated" + assert "phases" in info + assert "prefill" in info["phases"] + assert "scale_factor" in info["phases"]["prefill"] + + def test_get_threshold_info_static(self): + """get_threshold_info returns static type when no calibration.""" + method = FlashSkipSoftmax( + { + "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, + "br": 128, + "bc": 128, + "backend": "pytorch", + "is_causal": True, + } + ) + info = method.get_threshold_info() + assert info["type"] == "static" + assert "value" in info + + def test_get_sparse_context_patches_softmax(self): + """get_sparse_context returns an ExitStack that patches F.softmax.""" + import torch.nn.functional as F + + method = FlashSkipSoftmax( + { + "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, + "br": 64, + "bc": 64, + "backend": "pytorch", + "is_causal": True, + } + ) + + module = type("M", (), {"_last_stats": None})() + original_softmax = F.softmax + stack = method.get_sparse_context(module) + with stack: + # Inside the context, softmax should be patched + assert F.softmax is not original_softmax + # Call it once to exercise the sparse_softmax wrapper + scores = torch.randn(1, 1, 64, 64) + F.softmax(scores, dim=-1) + assert module._last_stats is not None + + # After the context, softmax is restored + assert F.softmax is original_softmax + + def test_calibration_mode_skips_apply(self): + """In calibration mode, sparse_softmax wrapper does not apply mask.""" + import torch.nn.functional as F + + method = FlashSkipSoftmax( + { + "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, + "br": 64, + "bc": 64, + "backend": "pytorch", + "is_causal": True, + } + ) + method.set_calibration_mode(True) + module = type("M", (), {"_last_stats": None})() + + with method.get_sparse_context(module): + scores = torch.randn(1, 1, 64, 64) + # Should not apply sparsity — output is regular softmax + out = F.softmax(scores, dim=-1) + assert torch.allclose(out.sum(dim=-1), torch.ones_like(out.sum(dim=-1))) + method.set_calibration_mode(False) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_kernel_backends.py b/tests/unit/torch/sparsity/attention_sparsity/test_kernel_backends.py new file mode 100644 index 00000000000..775723e66c4 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_kernel_backends.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for diffusers kernel backend registration and thread-local context. + +The forward pass of ``_diffusers_triton_attention`` requires a GPU and is +exercised in ``tests/gpu/torch/sparsity/attention_sparsity/ +test_diffusers_triton_attention.py``. These CPU tests cover backend +registration, thread-local config, and the conversion-time plumbing. +""" + +import pytest +import torch.nn as nn + +pytest.importorskip("diffusers") + + +# --------------------------------------------------------------------------- +# Thread-local skip-softmax context +# --------------------------------------------------------------------------- + + +class TestSkipSoftmaxContext: + def test_default_is_false(self): + from modelopt.torch.sparsity.attention_sparsity.kernels import get_skip_softmax_context + + assert get_skip_softmax_context() is False + + def test_set_and_get(self): + from modelopt.torch.sparsity.attention_sparsity.kernels import ( + get_skip_softmax_context, + set_skip_softmax_context, + ) + + set_skip_softmax_context(True) + assert get_skip_softmax_context() is True + set_skip_softmax_context(False) + assert get_skip_softmax_context() is False + + +# --------------------------------------------------------------------------- +# Diffusers triton attention backend registration and config +# --------------------------------------------------------------------------- + + +class TestDiffusersTritonBackend: + """Backend registration, thread-local config — no kernel execution.""" + + @pytest.fixture(autouse=True) + def _reset(self): + from modelopt.torch.sparsity.attention_sparsity.kernels import ( + diffusers_triton_attention as mod, + ) + + mod._BACKEND_REGISTERED = False + mod.clear_triton_skip_softmax_config() + yield + mod.clear_triton_skip_softmax_config() + + def test_set_clear_config(self): + from modelopt.torch.sparsity.attention_sparsity.kernels.diffusers_triton_attention import ( + clear_triton_skip_softmax_config, + set_triton_skip_softmax_config, + ) + + set_triton_skip_softmax_config(threshold=0.1) + clear_triton_skip_softmax_config() + + def test_register_idempotent(self): + from modelopt.torch.sparsity.attention_sparsity.kernels.diffusers_triton_attention import ( + register_diffusers_triton_attention, + ) + + register_diffusers_triton_attention() + register_diffusers_triton_attention() # Should be a no-op + + def test_get_backend_before_register_raises(self): + from modelopt.torch.sparsity.attention_sparsity.kernels.diffusers_triton_attention import ( + get_triton_attention_backend, + ) + + with pytest.raises(RuntimeError, match="not registered"): + get_triton_attention_backend() + + +# --------------------------------------------------------------------------- +# conversion._register_diffusers_backends_if_needed +# --------------------------------------------------------------------------- + + +class TestRegisterDiffusersBackends: + def test_no_diffusers_model_no_error(self): + """Non-ModelMixin models pass through without registering.""" + from modelopt.torch.sparsity.attention_sparsity.conversion import ( + _register_diffusers_backends_if_needed, + ) + + _register_diffusers_backends_if_needed(nn.Linear(10, 10)) + + def test_with_diffusers_model(self): + """A ModelMixin subclass triggers diffusers backend registration.""" + from diffusers.models.modeling_utils import ModelMixin + + from modelopt.torch.sparsity.attention_sparsity.conversion import ( + _register_diffusers_backends_if_needed, + ) + from modelopt.torch.sparsity.attention_sparsity.kernels import ( + diffusers_triton_attention as mod, + ) + + mod._BACKEND_REGISTERED = False + + class _TinyMixinModel(ModelMixin): + def __init__(self): + super().__init__() + self.lin = nn.Linear(4, 4) + + _register_diffusers_backends_if_needed(_TinyMixinModel()) + assert mod._BACKEND_REGISTERED is True diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_ltx_triton_attention.py b/tests/unit/torch/sparsity/attention_sparsity/test_ltx_triton_attention.py new file mode 100644 index 00000000000..4751fbae353 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_ltx_triton_attention.py @@ -0,0 +1,126 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for LTX Triton attention thread-local context and registration. + +The ``_ltx_triton_attention`` wrapper forward pass is exercised end-to-end in +``tests/gpu/torch/sparsity/attention_sparsity/test_diffusers_triton_attention.py``. +These CPU tests only cover the thread-local context helpers and registration. +""" + +import contextlib +import sys +import types +from unittest.mock import patch + +import pytest +import torch + + +@pytest.fixture +def ltx_mod(): + """Import ltx_triton_attention and ensure thread-local state is reset.""" + from modelopt.torch.sparsity.attention_sparsity.kernels import ltx_triton_attention as mod + + mod.clear_ltx_triton_context() + try: + yield mod + finally: + mod.clear_ltx_triton_context() + + +class TestThreadLocalContext: + """Test set/clear/get thread-local context functions.""" + + def test_set_context_populates_fields(self, ltx_mod): + ltx_mod.set_ltx_triton_context( + active=True, + threshold=0.1, + calibration_mode=False, + threshold_trials=[0.01, 0.1], + scale_factor=2.0, + raw_threshold=-5.0, + ) + active, threshold, scale_factor = ltx_mod._get_ltx_triton_context() + assert active is True + assert threshold == 0.1 + assert scale_factor == 2.0 + + def test_set_context_without_calibration_mode_clears_counters(self, ltx_mod): + """Setting non-calibration mode resets calibration_counters to None.""" + ltx_mod._thread_local.calibration_counters = torch.tensor([[1, 2]]) + ltx_mod.set_ltx_triton_context(active=True, calibration_mode=False) + assert ltx_mod._thread_local.calibration_counters is None + + def test_set_context_in_calibration_mode_preserves_counters(self, ltx_mod): + """Setting calibration mode does NOT clear the existing counters.""" + existing = torch.tensor([[5, 3]]) + ltx_mod._thread_local.calibration_counters = existing + ltx_mod.set_ltx_triton_context(active=True, calibration_mode=True) + assert ltx_mod._thread_local.calibration_counters is existing + + def test_clear_context_resets_all(self, ltx_mod): + ltx_mod.set_ltx_triton_context(active=True, threshold=0.1, scale_factor=2.0) + ltx_mod.clear_ltx_triton_context() + active, threshold, scale_factor = ltx_mod._get_ltx_triton_context() + assert active is False + assert threshold is None + assert scale_factor is None + + def test_get_calibration_counters_returns_none_initially(self, ltx_mod): + assert ltx_mod.get_calibration_counters() is None + assert ltx_mod.get_calibration_seq_k() is None + + +class TestRegisterLTXTritonAttention: + """Test register_ltx_triton_attention patches ltx_core Attention modules.""" + + def test_no_ltx_core_no_error(self, ltx_mod): + """If ltx_core is absent, the patch attempt raises ImportError cleanly.""" + with contextlib.suppress(ImportError, ModuleNotFoundError): + ltx_mod.register_ltx_triton_attention(torch.nn.Linear(4, 4)) + + def test_patches_ltx_attention_modules(self, ltx_mod): + """When ltx_core.Attention exists, modules get wrapped.""" + + class FakeAttention(torch.nn.Module): + def __init__(self): + super().__init__() + self.attention_function = lambda q, k, v, heads, mask=None: q + + fake_attn_mod = types.ModuleType("ltx_core.model.transformer.attention") + fake_attn_mod.Attention = FakeAttention + + patched = { + "ltx_core": types.ModuleType("ltx_core"), + "ltx_core.model": types.ModuleType("ltx_core.model"), + "ltx_core.model.transformer": types.ModuleType("ltx_core.model.transformer"), + "ltx_core.model.transformer.attention": fake_attn_mod, + } + with patch.dict(sys.modules, patched): + + class Parent(torch.nn.Module): + def __init__(self): + super().__init__() + self.attn1 = FakeAttention() + + parent = Parent() + ltx_mod.register_ltx_triton_attention(parent) + + from modelopt.torch.sparsity.attention_sparsity.kernels.ltx_triton_attention import ( + _TritonLTXAttentionWrapper, + ) + + assert isinstance(parent.attn1.attention_function, _TritonLTXAttentionWrapper) diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py index 0785cdf2275..9519856a875 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_calibration.py @@ -19,6 +19,9 @@ pytest.importorskip("transformers") +# Imports added for new tests +from unittest.mock import MagicMock, patch + import numpy as np from _test_utils.torch.sparsity.sparse_attention_common import SimpleAttentionModel from pydantic import ValidationError @@ -30,7 +33,10 @@ ) from modelopt.torch.sparsity.attention_sparsity.calibration.calibrate import ( _extract_calibration_config, + _extract_tokenizer_from_model, calibrate_sparse_attention, + create_calibration_forward_loop, + create_decode_calibration_forward_loop, ) from modelopt.torch.sparsity.attention_sparsity.calibration.ruler_dataset import ( _generate_target_lengths, @@ -416,3 +422,170 @@ def test_extract_calibration_config_none(self): calib_config = _extract_calibration_config(config) assert calib_config is None + + def test_extract_calibration_config_invalid_type(self): + """_extract_calibration_config raises when calibration is not a dict.""" + config = { + "sparse_cfg": { + "calibration": "not-a-dict", + "*attn*": {"method": "flash_skip_softmax"}, + }, + } + with pytest.raises(ValueError, match="must be a dict"): + _extract_calibration_config(config) + + def test_calibrate_no_sparse_modules(self): + """Config with calibration but no sparse modules -> empty dict.""" + model = SimpleAttentionModel(hidden_size=64, num_heads=4) + + def noop_forward(m): + pass + + config = { + "sparse_cfg": { + "calibration": { + "target_sparse_ratio": {"prefill": 0.5, "decode": 0.0}, + "samples": 4, + "max_seqlen": 1024, + }, + }, + } + # Without sparse_cfg for modules, no sparse modules => returns empty + result = calibrate_sparse_attention(model, config, forward_loop=noop_forward) + assert result == {} + + def test_calibrate_both_phases_zero(self): + """Config with both prefill/decode=0.0 returns empty immediately.""" + model = SimpleAttentionModel(hidden_size=64, num_heads=4) + config = { + "sparse_cfg": { + "calibration": { + "target_sparse_ratio": {"prefill": 0.0, "decode": 0.0}, + "samples": 4, + "max_seqlen": 1024, + }, + }, + } + # Both phases disabled => returns empty without attempting calibration + result = calibrate_sparse_attention(model, config, forward_loop=lambda m: None) + assert result == {} + + def test_calibrate_with_user_forward_loop(self): + """User-provided forward_loop skips RULER dataset building.""" + import numpy as np + + model = SimpleAttentionModel(hidden_size=64, num_heads=4) + # Sparsify first WITHOUT calibration, so we can call calibrate_sparse_attention + # ourselves with a user-supplied forward_loop. + base_config = { + "sparse_cfg": { + "*attention*": { + "method": "flash_skip_softmax", + "thresholds": {"prefill": [1e-3], "decode": [1e-4]}, + "br": 64, + "bc": 64, + "enable": True, + }, + }, + } + sparse_model = sparsify(model, base_config) + + # Now build a calibration-enabled config and call calibrate directly + calib_config = { + "sparse_cfg": { + "calibration": { + "target_sparse_ratio": {"prefill": 0.5, "decode": 0.0}, + "samples": 4, + "max_seqlen": 1024, + }, + }, + } + + # Seed synthetic stats in each module so calibrator converges. + # Use a known threshold list (matches the DynamicThresholdCalibrator default). + a_true, b_true = 0.1, 5.0 + threshold_trials = [ + 1e-6, + 5e-6, + 1e-5, + 5e-5, + 1e-4, + 5e-4, + 1e-3, + 5e-3, + 1e-2, + 2e-2, + 5e-2, + 1e-1, + 2e-1, + 3e-1, + 5e-1, + 7e-1, + 8e-1, + 9e-1, + 9.5e-1, + 9.9e-1, + ] + + def fake_forward(m): + for module in m.modules(): + if isinstance(module, SparseAttentionModule): + sparsity_list = [] + for t in threshold_trials: + sf = t * 4096 + s = np.log(sf / a_true) / b_true + sparsity_list.append(max(0.0, min(1.0, s))) + module._last_stats = { + "sparsity": sparsity_list, + "sample_length": 4096, + "phase": "prefill", + } + if module._stats_manager is not None: + module._stats_manager.collect(module._last_stats) + + # This exercises the main calibrate_sparse_attention pipeline end-to-end + result = calibrate_sparse_attention(sparse_model, calib_config, forward_loop=fake_forward) + # If regression fit succeeded, we have results; otherwise we expect a warning + # and an empty dict. Either is acceptable — we just want the code path covered. + assert isinstance(result, dict) + + +class TestCreateCalibrationForwardLoop: + """Test create_calibration_forward_loop factory (no GPU required).""" + + def test_factory_returns_callable(self): + """Factory returns a callable forward_loop closure.""" + fake_tok = MagicMock() + fake_tok.pad_token = None + fake_tok.eos_token = "" + with patch( + "transformers.AutoTokenizer.from_pretrained", + MagicMock(return_value=fake_tok), + ): + fn = create_calibration_forward_loop([], "gpt2") + assert callable(fn) + + def test_decode_factory_returns_callable(self): + """create_decode_calibration_forward_loop returns a callable.""" + fake_tok = MagicMock() + fake_tok.pad_token = None + fake_tok.eos_token = "" + with patch( + "transformers.AutoTokenizer.from_pretrained", + MagicMock(return_value=fake_tok), + ): + fn = create_decode_calibration_forward_loop([], "gpt2") + assert callable(fn) + + +class TestExtractTokenizer: + """Test _extract_tokenizer_from_model logic.""" + + def test_raises_when_no_config(self): + model = type("M", (), {})() # no config attribute + with pytest.raises(ValueError, match="Could not load tokenizer"): + _extract_tokenizer_from_model(model) + + def test_returns_tokenizer_path_from_config(self): + model = type("M", (), {"config": type("C", (), {"_name_or_path": "gpt2"})()})() + assert _extract_tokenizer_from_model(model) == "gpt2" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py index eab020022b7..93389a46105 100644 --- a/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py +++ b/tests/unit/torch/sparsity/attention_sparsity/test_sparse_attention_conversion.py @@ -19,6 +19,8 @@ pytest.importorskip("transformers") +from unittest.mock import MagicMock, patch + import torch.nn as nn from _test_utils.torch.sparsity.sparse_attention_common import ( FLASH_SKIP_SOFTMAX_DEFAULT_CFG, @@ -29,8 +31,10 @@ import modelopt.torch.opt as mto import modelopt.torch.sparsity.attention_sparsity as sparse_attn from modelopt.torch.sparsity.attention_sparsity.conversion import ( + _set_attn_implementation, disable_sparse_attention, enable_sparse_attention, + export_sparse_attention_config, print_sparse_attention_summary, ) from modelopt.torch.sparsity.attention_sparsity.sparse_attention import SparseAttentionModule @@ -249,3 +253,127 @@ def test_get_stats_without_stats_manager(self): stats = module.get_stats() assert stats == {} break + + +class TestSetAttnImplementation: + """Cover the _set_attn_implementation logic in conversion.py.""" + + def test_triton_backend_sets_attn_impl(self): + """triton backend sets _attn_implementation=modelopt_triton on model.config.""" + model = type( + "M", + (), + {"config": type("C", (), {"_attn_implementation": "eager"})()}, + )() + config = type("Cfg", (), {"sparse_cfg": {}})() + config.sparse_cfg = { + "*": {"method": "triton_skip_softmax", "backend": "triton"}, + } + with patch( + "modelopt.torch.sparsity.attention_sparsity.kernels.register_triton_attention", + MagicMock(return_value=True), + ): + _set_attn_implementation(model, config) + assert model.config._attn_implementation == "modelopt_triton" + + def test_triton_backend_register_failure_raises(self): + """When register_triton_attention returns False, a RuntimeError is raised.""" + model = type( + "M", + (), + {"config": type("C", (), {"_attn_implementation": "eager"})()}, + )() + config = type("Cfg", (), {"sparse_cfg": {}})() + config.sparse_cfg = {"*": {"method": "triton_skip_softmax", "backend": "triton"}} + with ( + patch( + "modelopt.torch.sparsity.attention_sparsity.kernels.register_triton_attention", + MagicMock(return_value=False), + ), + pytest.raises(RuntimeError, match="Failed to register"), + ): + _set_attn_implementation(model, config) + + def test_triton_backend_no_triton_raises(self): + """When register_triton_attention is None, ImportError is raised.""" + model = type( + "M", + (), + {"config": type("C", (), {"_attn_implementation": "eager"})()}, + )() + config = type("Cfg", (), {"sparse_cfg": {}})() + config.sparse_cfg = {"*": {"method": "triton_skip_softmax", "backend": "triton"}} + with ( + patch( + "modelopt.torch.sparsity.attention_sparsity.kernels.register_triton_attention", + None, + ), + pytest.raises(ImportError, match="Triton backend requires"), + ): + _set_attn_implementation(model, config) + + def test_mixed_backends_raises(self): + """Mixing pytorch and triton backends is not supported.""" + model = type("M", (), {"config": None})() + config = type("Cfg", (), {"sparse_cfg": {}})() + config.sparse_cfg = { + "layer1": {"method": "triton_skip_softmax", "backend": "triton"}, + "layer2": {"method": "flash_skip_softmax", "backend": "pytorch"}, + } + with pytest.raises(ValueError, match="Mixed backends"): + _set_attn_implementation(model, config) + + def test_vsa_only_is_noop(self): + """VSA-only configs do not change _attn_implementation.""" + model = type( + "M", + (), + {"config": type("C", (), {"_attn_implementation": "eager"})()}, + )() + config = type("Cfg", (), {"sparse_cfg": {}})() + config.sparse_cfg = {"*": {"method": "vsa"}} + _set_attn_implementation(model, config) + # Should remain eager — VSA patches SDPA directly + assert model.config._attn_implementation == "eager" + + def test_mixed_vsa_and_non_vsa_raises(self): + """VSA + non-VSA methods are rejected.""" + model = type("M", (), {"config": None})() + config = type("Cfg", (), {"sparse_cfg": {}})() + config.sparse_cfg = { + "layer1": {"method": "vsa"}, + "layer2": {"method": "flash_skip_softmax", "backend": "pytorch"}, + } + with pytest.raises(ValueError, match="Cannot mix VSA"): + _set_attn_implementation(model, config) + + +class TestExportSparseAttentionConfig: + """Cover export_sparse_attention_config branches.""" + + def test_returns_none_without_calibration(self): + """When no module has calibration params, returns None.""" + model = SimpleAttentionModel() + model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG) + out = export_sparse_attention_config(model) + assert out is None + + def test_exports_when_calibration_present(self): + """Calibration params are reflected in the exported config.""" + model = SimpleAttentionModel() + model = sparse_attn.sparsify(model, FLASH_SKIP_SOFTMAX_DEFAULT_CFG) + + for module in model.modules(): + if isinstance(module, SparseAttentionModule): + module._sparse_method_instance.calibration_params = { + "prefill": {"a": 3.14, "b": 7.5}, + "decode": {"a": 0.5, "b": 9.0}, + } + + out = export_sparse_attention_config(model) + assert out is not None + assert "config_groups" in out + tsf = out["threshold_scale_factor"] + assert tsf["prefill"] == {"a": 3.14, "b": 7.5} + assert tsf["decode"] == {"a": 0.5, "b": 9.0} + assert out["producer"]["name"] == "modelopt" diff --git a/tests/unit/torch/sparsity/attention_sparsity/test_triton_skip_softmax.py b/tests/unit/torch/sparsity/attention_sparsity/test_triton_skip_softmax.py new file mode 100644 index 00000000000..8f7ef9f1271 --- /dev/null +++ b/tests/unit/torch/sparsity/attention_sparsity/test_triton_skip_softmax.py @@ -0,0 +1,222 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for TritonSkipSoftmaxMethod (no GPU required).""" + +import math +import warnings + +import pytest +import torch + +from modelopt.torch.sparsity.attention_sparsity.methods.triton_skip_softmax import ( + TritonSkipSoftmaxMethod, +) + + +class TestInit: + def test_default_config(self): + m = TritonSkipSoftmaxMethod() + assert m.skip_softmax_threshold == 0.1 + assert m.skip_softmax_raw_threshold is None + assert m._threshold_trials is None + assert m._measure_sparsity is False + + def test_custom_config(self): + m = TritonSkipSoftmaxMethod( + {"skip_softmax_threshold": 0.05, "skip_softmax_raw_threshold": -3.0} + ) + assert m.skip_softmax_threshold == 0.05 + assert m.skip_softmax_raw_threshold == -3.0 + + def test_name(self): + assert TritonSkipSoftmaxMethod().name == "triton_skip_softmax" + + +class TestCalculateSparsity: + def test_returns_all_ones_mask(self): + m = TritonSkipSoftmaxMethod() + scores = torch.randn(2, 4, 8, 8) + mask, stats = m.calculate_sparsity(scores) + assert mask.shape == scores.shape + assert mask.all() + assert stats == {} + + +class TestApplySparsity: + def test_raises_not_implemented(self): + m = TritonSkipSoftmaxMethod() + with pytest.raises(NotImplementedError, match="Triton kernel"): + m.apply_sparsity(torch.randn(2, 2)) + + +class TestGetScaleFactor: + def test_uncalibrated_returns_none(self): + m = TritonSkipSoftmaxMethod() + assert m._get_scale_factor() is None + + def test_no_target_returns_none(self): + m = TritonSkipSoftmaxMethod() + m.calibration_params = {"prefill": {"a": 1.0, "b": 5.0}} + m.target_sparse_ratio = None + assert m._get_scale_factor() is None + + def test_calibrated_computation(self): + m = TritonSkipSoftmaxMethod() + m.calibration_params = {"prefill": {"a": 2.0, "b": 3.0}} + m.target_sparse_ratio = {"prefill": 0.5} + expected = 2.0 * math.exp(3.0 * 0.5) + assert m._get_scale_factor() == pytest.approx(expected) + + def test_zero_a_returns_none(self): + m = TritonSkipSoftmaxMethod() + m.calibration_params = {"prefill": {"a": 0, "b": 5.0}} + m.target_sparse_ratio = {"prefill": 0.5} + assert m._get_scale_factor() is None + + def test_zero_b_returns_none(self): + m = TritonSkipSoftmaxMethod() + m.calibration_params = {"prefill": {"a": 1.0, "b": 0}} + m.target_sparse_ratio = {"prefill": 0.5} + assert m._get_scale_factor() is None + + def test_warns_below_min_observed(self): + m = TritonSkipSoftmaxMethod() + m.calibration_params = { + "prefill": { + "a": 1.0, + "b": 5.0, + "min_observed_sparsity": 0.3, + "max_observed_sparsity": 0.8, + } + } + m.target_sparse_ratio = {"prefill": 0.1} + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + result = m._get_scale_factor() + assert result is not None + assert len(w) == 1 + assert "below the minimum" in str(w[0].message) + + def test_warns_above_max_observed(self): + m = TritonSkipSoftmaxMethod() + m.calibration_params = { + "prefill": { + "a": 1.0, + "b": 5.0, + "min_observed_sparsity": 0.3, + "max_observed_sparsity": 0.8, + } + } + m.target_sparse_ratio = {"prefill": 0.95} + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + result = m._get_scale_factor() + assert result is not None + assert len(w) == 1 + assert "above the maximum" in str(w[0].message) + + +class TestGetThresholdInfo: + def test_static_threshold(self): + m = TritonSkipSoftmaxMethod({"skip_softmax_threshold": 0.05}) + info = m.get_threshold_info() + assert info["type"] == "static" + assert info["value"] == 0.05 + + def test_calibrated_threshold(self): + m = TritonSkipSoftmaxMethod() + m.calibration_params = {"prefill": {"a": 2.0, "b": 3.0}} + m.target_sparse_ratio = {"prefill": 0.5} + info = m.get_threshold_info() + assert info["type"] == "dynamic_calibrated" + assert "scale_factor" in info + + +class TestSparsityMeasurement: + def test_enable_disable(self): + m = TritonSkipSoftmaxMethod() + assert m._measure_sparsity is False + m.enable_measure_sparsity(True) + assert m._measure_sparsity is True + m.enable_measure_sparsity(False) + assert m._measure_sparsity is False + + def test_reset_counters(self): + m = TritonSkipSoftmaxMethod() + m._sparsity_total = 100 + m._sparsity_skipped = 50 + m.reset_sparsity_counters() + assert m._sparsity_total == 0 + assert m._sparsity_skipped == 0 + + def test_get_counters(self): + m = TritonSkipSoftmaxMethod() + m._sparsity_total = 200 + m._sparsity_skipped = 80 + total, skipped = m.get_sparsity_counters() + assert total == 200 + assert skipped == 80 + + +class TestGetSparseContext: + def test_inference_mode_selected(self): + m = TritonSkipSoftmaxMethod() + m._calibration_mode = False + module = type("M", (), {"_apply_skip_softmax": False})() + ctx = m.get_sparse_context(module) + # Should return the inference context (a generator-based context manager) + assert hasattr(ctx, "__enter__") + + def test_calibration_mode_selected(self): + m = TritonSkipSoftmaxMethod() + m._calibration_mode = True + m._threshold_trials = [0.01, 0.1] + module = type("M", (), {"_apply_skip_softmax": False, "_last_stats": None})() + ctx = m.get_sparse_context(module) + assert hasattr(ctx, "__enter__") + + def test_calibration_mode_without_trials_falls_back_to_inference(self): + m = TritonSkipSoftmaxMethod() + m._calibration_mode = True + m._threshold_trials = None # No trials = falls back to inference + module = type("M", (), {"_apply_skip_softmax": False})() + ctx = m.get_sparse_context(module) + assert hasattr(ctx, "__enter__") + + +class TestCollectCalibrationStats: + """Defensive null-guards in _collect_calibration_stats (no GPU required). + + The happy-path (calibration context populating ``module._last_stats`` from + real counters) is exercised by GPU tests in + ``tests/gpu/torch/sparsity/attention_sparsity/``. + """ + + def test_no_counters_is_noop(self): + """Skips writing stats when neither backend has counters.""" + m = TritonSkipSoftmaxMethod() + m._threshold_trials = [0.01] + module = type("M", (), {"_last_stats": None})() + m._collect_calibration_stats(module) + assert module._last_stats is None + + def test_no_threshold_trials_is_noop(self): + """Skips writing stats when threshold_trials was never set.""" + m = TritonSkipSoftmaxMethod() + m._threshold_trials = None + module = type("M", (), {"_last_stats": None})() + m._collect_calibration_stats(module) + assert module._last_stats is None diff --git a/tox.ini b/tox.ini index 0cd99cc86f1..7bfa1e41e57 100644 --- a/tox.ini +++ b/tox.ini @@ -45,6 +45,8 @@ deps = # Install megatron-core to test torch-only install can still import plugins torch: megatron-core + # diffusers is needed for unit tests of the sparse-attention/quantization diffusers backend + torch: diffusers torch: .[dev-test] torch_deploy: .[onnx,torch,dev-test] From 92622a9aa60737ad84b62b5bdb324e6318ee4cfd Mon Sep 17 00:00:00 2001 From: bkartal-dev Date: Sat, 18 Apr 2026 00:28:49 -0700 Subject: [PATCH 15/30] Add nvfp4_mse and nvfp4_local_hessian options to the ptq script. (#1113) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bugfix Add newly added quant configs to the example PTQ script. ### Testing I have locally run auto_quantize with these two quant_configs, and obtained successfully exported HF artifacts. ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A - Did you write any new necessary tests?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit * **New Features** * Added support for three new quantization formats: nvfp4_mse, nvfp4_local_hessian, and nvfp4_experts_only, expanding available export options when using auto-quantize. * **Bug Fixes / UX** * Updated the invalid-quantization error message to include the newly accepted format identifiers. Signed-off-by: Bilal Kartal Signed-off-by: bkartal-dev --- examples/llm_ptq/hf_ptq.py | 1 + examples/llm_ptq/scripts/huggingface_example.sh | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/llm_ptq/hf_ptq.py b/examples/llm_ptq/hf_ptq.py index 969a3d57190..831d230a672 100755 --- a/examples/llm_ptq/hf_ptq.py +++ b/examples/llm_ptq/hf_ptq.py @@ -337,6 +337,7 @@ def auto_quantize( "nvfp4_mlp_only", "nvfp4_experts_only", "nvfp4_omlp_only", + "nvfp4_local_hessian", "mxfp8", ] for qformat in qformat_list diff --git a/examples/llm_ptq/scripts/huggingface_example.sh b/examples/llm_ptq/scripts/huggingface_example.sh index f511abc289e..d9c4ff8a7a0 100755 --- a/examples/llm_ptq/scripts/huggingface_example.sh +++ b/examples/llm_ptq/scripts/huggingface_example.sh @@ -53,9 +53,9 @@ esac IFS="," for qformat in $QFORMAT; do case $qformat in - fp8 | fp8_pc_pt | fp8_pb_wo | int8_wo | int8_sq | int4_awq | w4a8_awq | fp16 | bf16 | nvfp4 | nvfp4_awq | w4a8_nvfp4_fp8 | w4a8_mxfp4_fp8 | nvfp4_mlp_only | nvfp4_experts_only | nvfp4_omlp_only | nvfp4_svdquant | mxfp8) ;; + fp8 | fp8_pc_pt | fp8_pb_wo | int8_wo | int8_sq | int4_awq | w4a8_awq | fp16 | bf16 | nvfp4 | nvfp4_awq | nvfp4_mse | w4a8_nvfp4_fp8 | w4a8_mxfp4_fp8 | nvfp4_experts_only | nvfp4_mlp_only | nvfp4_omlp_only | nvfp4_svdquant | mxfp8 | nvfp4_local_hessian) ;; *) - echo "Unknown quant argument: Expected one of: [fp8, fp8_pc_pt, fp8_pb_wo, int8_wo, int8_sq, int4_awq, w4a8_awq, fp16, bf16, nvfp4, nvfp4_awq, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8, nvfp4_mlp_only, nvfp4_experts_only, nvfp4_omlp_only, nvfp4_svdquant, mxfp8]" >&2 + echo "Unknown quant argument: Expected one of: [fp8, fp8_pc_pt, fp8_pb_wo, int8_wo, int8_sq, int4_awq, w4a8_awq, fp16, bf16, nvfp4, nvfp4_awq, nvfp4_mse, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8, nvfp4_experts_only, nvfp4_mlp_only, nvfp4_omlp_only, nvfp4_svdquant, mxfp8, nvfp4_local_hessian]" >&2 exit 1 ;; esac From 760c9807279eceec120deaae3a941dd7e86c3379 Mon Sep 17 00:00:00 2001 From: Ajinkya Rasane <131806219+ajrasane@users.noreply.github.com> Date: Sat, 18 Apr 2026 20:00:45 +0530 Subject: [PATCH 16/30] Add ResNet50 support for torch_onnx quantization workflow (#1263) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - Add end-to-end ResNet50 support in the torch_onnx quantization → ONNX export → TRT engine pipeline - Fix multiple Conv2d-related export issues that blocked Conv2d-heavy models from working with FP8/INT8/MXFP8/NVFP4/auto quantization modes - Fix `configure_linear_module_onnx_quantizers` to handle all modules with block quantization (not just `nn.Linear`), fixing NVFP4/MXFP8 export for models with quantized non-Linear modules - Add `--trt_build` flag to `torch_quant_to_onnx.py` and simplify test infrastructure ### Files Changed - `modelopt/torch/_deploy/utils/torch_onnx.py` — Disable FP8 Conv2d weight quantizers and autocast during ONNX export - `modelopt/torch/quantization/export_onnx.py` — Fix `configure_linear_module_onnx_quantizers` for all module types with block quantization - `examples/torch_onnx/torch_quant_to_onnx.py` — Add `--trt_build` flag, calibration for FP8 override quantizers, Conv2d→FP8 override for auto mode, filter_func updates - `examples/torch_onnx/README.md` — Add ResNet50 to supported models table - `tests/examples/torch_onnx/test_torch_quant_to_onnx.py` — Add ResNet50 test entry, simplify using `--trt_build` - `tests/_test_utils/torch/vision_models.py` — Add ResNet50 to timm model registry ### Quantization modes passing - ✅ FP8, INT8, MXFP8, NVFP4, Auto (all 5 modes pass export + TRT build) - INT4_AWQ excluded (pre-existing limitation for all models) ## Test plan - [x] All 5 resnet50 test modes pass: `pytest tests/examples/torch_onnx/test_torch_quant_to_onnx.py -k resnet50` (5/5 passed) - [x] Full regression: 18 passed, 2 failed (pre-existing swinv2_tiny fp8/int8 failures) 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **New Features** * Added ResNet50 to supported ONNX export vision models with FP8, INT8, MXFP8, and NVFP4 support. * Optional TensorRT engine build after export via a new CLI flag. * **Improvements** * Enhanced quantization calibration and export flows for FP8/INT8 models, including broader block-quantization support across module types and safer export handling. * Tests updated to include ResNet50 in the model matrix. --------- Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> Signed-off-by: ajrasane Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- examples/torch_onnx/README.md | 1 + examples/torch_onnx/torch_quant_to_onnx.py | 152 +++++++++++--- modelopt/onnx/export/fp8_exporter.py | 85 +++++++- modelopt/onnx/utils.py | 186 ++++++++++++++++++ modelopt/torch/_deploy/utils/torch_onnx.py | 52 ++++- modelopt/torch/quantization/export_onnx.py | 37 +++- tests/_test_utils/torch/vision_models.py | 1 + .../torch_onnx/test_torch_quant_to_onnx.py | 36 +--- 8 files changed, 480 insertions(+), 70 deletions(-) diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index d540770116f..cfd4dc380ce 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -307,6 +307,7 @@ python torch_quant_to_onnx.py \ | [vit_base_patch16_224](https://huggingface.co/timm/vit_base_patch16_224.augreg_in21k_ft_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [swin_tiny_patch4_window7_224](https://huggingface.co/timm/swin_tiny_patch4_window7_224.ms_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | | [swinv2_tiny_window8_256](https://huggingface.co/timm/swinv2_tiny_window8_256.ms_in1k) | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | +| [resnet50](https://huggingface.co/timm/resnet50.a1_in1k) | ✅ | ✅ | ✅ | ✅ | | ✅ | ## Resources diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index 7f74e617e86..98daa2c13d4 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -17,6 +17,7 @@ import copy import json import re +import subprocess import sys import warnings from pathlib import Path @@ -35,13 +36,19 @@ import modelopt.torch.quantization as mtq """ -This script is used to quantize a timm model using dynamic quantization like MXFP8 or NVFP4, -or using auto quantization for optimal per-layer quantization. +Quantize a timm vision model and export to ONNX for TensorRT deployment. + +Supports FP8, INT8, MXFP8, NVFP4, and AUTO (mixed-precision) quantization modes end-to-end +(quantize + ONNX export + TRT build). INT4_AWQ is quantize/export-only; it is not compatible +with ``--trt_build``. The script will: -1. Given the model name, create a timm torch model. -2. Quantize the torch model in MXFP8, NVFP4, INT4_AWQ, or AUTO mode. -3. Export the quantized torch model to ONNX format. +1. Load a pretrained timm model (e.g., ViT, Swin, ResNet). +2. Quantize the model using the specified mode. For models with Conv2d layers, + Conv2d quantization is automatically overridden for TensorRT compatibility + (FP8 for MXFP8/NVFP4, INT8 for INT4_AWQ). +3. Export the quantized model to ONNX with FP16 weights. +4. Optionally evaluate accuracy on ImageNet-1k before and after quantization. """ @@ -81,6 +88,15 @@ }, ] +# Auto-quantize format configs that use block quantization and need Conv2d overrides for TRT. +# TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors. +_NEEDS_FP8_CONV_OVERRIDE: set[str] = { + "NVFP4_AWQ_LITE_CFG", + "NVFP4_DEFAULT_CFG", + "MXFP8_DEFAULT_CFG", +} +_NEEDS_INT8_CONV_OVERRIDE: set[str] = {"INT4_AWQ_CFG"} + def get_quant_config(quantize_mode): """Get quantization config, overriding Conv2d for TRT compatibility. @@ -106,19 +122,26 @@ def get_quant_config(quantize_mode): def filter_func(name): - """Filter function to exclude certain layers from quantization.""" + """Filter function to exclude certain layers from quantization. + + ``downsample.reduction`` (Swin/SwinV2) is excluded because it operates on 4D tensors + and TRT's DynamicQuantize layer (used for MXFP8/NVFP4) requires 2D/3D input. + """ pattern = re.compile( r".*(time_emb_proj|time_embedding|conv_in|conv_out|conv_shortcut|add_embedding|" - r"pos_embed|time_text_embed|context_embedder|norm_out|x_embedder|patch_embed|cpb_mlp|downsample).*" + r"pos_embed|time_text_embed|context_embedder|norm_out|x_embedder|patch_embed|cpb_mlp|" + r"maxpool|global_pool|downsample\.reduction).*" ) return pattern.match(name) is not None -def load_calibration_data(model_name, data_size, batch_size, device, with_labels=False): +def load_calibration_data(model, data_size, batch_size, device, with_labels=False): """Load and prepare calibration data. Args: - model_name: Name of the timm model + model: The timm model being quantized; used to derive the calibration transforms so the + data pipeline matches the exact model config (respects --no_pretrained and + --model_kwargs). data_size: Number of samples to load batch_size: Batch size for data loader device: Device to load data to @@ -126,7 +149,6 @@ def load_calibration_data(model_name, data_size, batch_size, device, with_labels If False, return just the images (for standard quantize) """ dataset = load_dataset("zh-plus/tiny-imagenet") - model = timm.create_model(model_name, pretrained=True, num_classes=1000) data_config = timm.data.resolve_model_data_config(model) transforms = timm.data.create_transform(**data_config, is_training=False) @@ -147,6 +169,36 @@ def load_calibration_data(model_name, data_size, batch_size, device, with_labels ) +def _calibrate_uncalibrated_quantizers(model, data_loader): + """Calibrate FP8 quantizers that weren't calibrated by mtq.quantize(). + + When MXFP8/NVFP4 modes override Conv2d to FP8, the FP8 quantizers may not + be calibrated because the MXFP8/NVFP4 quantization pipeline skips standard + calibration. This function explicitly calibrates those uncalibrated quantizers. + """ + uncalibrated = [] + for _, module in model.named_modules(): + for attr_name in ("input_quantizer", "weight_quantizer"): + if not hasattr(module, attr_name): + continue + quantizer = getattr(module, attr_name) + if quantizer.is_enabled and not quantizer.block_sizes and quantizer.amax is None: + quantizer.enable_calib() + uncalibrated.append(quantizer) + + if not uncalibrated: + return + + model.eval() + with torch.no_grad(): + for batch in data_loader: + model(batch) + + for quantizer in uncalibrated: + quantizer.disable_calib() + quantizer.load_calib_amax(strict=False) + + def quantize_model(model, config, data_loader=None): """Quantize the model using the given config and calibration data.""" if data_loader is not None: @@ -159,7 +211,14 @@ def forward_loop(model): else: quantized_model = mtq.quantize(model, config) + # Disable filtered quantizers BEFORE calibrating override quantizers so we don't + # waste time calibrating quantizers that are about to be turned off. mtq.disable_quantizer(quantized_model, filter_func) + + # Calibrate any FP8 override quantizers that weren't calibrated by mtq.quantize(). + if data_loader is not None: + _calibrate_uncalibrated_quantizers(quantized_model, data_loader) + return quantized_model @@ -209,11 +268,19 @@ def auto_quantize_model( _disable_inplace_relu(model) constraints = {"effective_bits": effective_bits} - # Convert string format names to actual config objects + # Convert string format names to config objects, incorporating Conv2d TRT overrides. + # TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors. + # By including the overrides in the format configs, the auto_quantize search + # correctly accounts for Conv2d being FP8/INT8 in the effective_bits budget. format_configs = [] for fmt in quantization_formats: if isinstance(fmt, str): - format_configs.append(getattr(mtq, fmt)) + config = copy.deepcopy(getattr(mtq, fmt)) + if fmt in _NEEDS_FP8_CONV_OVERRIDE: + config["quant_cfg"].extend(_FP8_CONV_OVERRIDE) + elif fmt in _NEEDS_INT8_CONV_OVERRIDE: + config["quant_cfg"].extend(_INT8_CONV_OVERRIDE) + format_configs.append(config) else: format_configs.append(fmt) @@ -248,7 +315,10 @@ def get_model_input_shape(model): def main(): parser = argparse.ArgumentParser( - description="Quantize timm models to FP8, MXFP8, INT8, NVFP4, INT4_AWQ, or use AUTO quantization" + description=( + "Quantize timm models to FP8, MXFP8, INT8, NVFP4, or use AUTO quantization. " + "INT4_AWQ is supported for quantize/export only and is not compatible with --trt_build." + ) ) # Model hyperparameters @@ -320,6 +390,11 @@ def main(): default=128, help="Number of scoring steps for auto quantization. Default is 128.", ) + parser.add_argument( + "--trt_build", + action="store_true", + help="Build a TensorRT engine from the exported ONNX model using trtexec.", + ) parser.add_argument( "--no_pretrained", action="store_true", @@ -362,7 +437,7 @@ def main(): if args.quantize_mode == "auto": # Auto quantization requires labels for loss computation data_loader = load_calibration_data( - args.timm_model_name, + model, args.calibration_data_size, args.batch_size, device, @@ -378,18 +453,18 @@ def main(): args.num_score_steps, ) else: - # Standard quantization - only load calibration data if needed + # Standard quantization - load calibration data + # Note: MXFP8 is dynamic and does not need calibration itself, but when + # Conv2d layers are overridden to FP8 (for TRT compatibility), those FP8 + # quantizers require calibration data. config = get_quant_config(args.quantize_mode) - if args.quantize_mode == "mxfp8": - data_loader = None - else: - data_loader = load_calibration_data( - args.timm_model_name, - args.calibration_data_size, - args.batch_size, - device, - with_labels=False, - ) + data_loader = load_calibration_data( + model, + args.calibration_data_size, + args.batch_size, + device, + with_labels=False, + ) quantized_model = quantize_model(model, config, data_loader) @@ -421,6 +496,33 @@ def main(): print(f"Quantized ONNX model is saved to {args.onnx_save_path}") + if args.trt_build: + build_trt_engine(args.onnx_save_path) + + +def build_trt_engine(onnx_path): + """Build a TensorRT engine from the exported ONNX model using trtexec.""" + cmd = [ + "trtexec", + f"--onnx={onnx_path}", + "--stronglyTyped", + "--builderOptimizationLevel=4", + ] + print(f"\nBuilding TensorRT engine: {' '.join(cmd)}") + try: + result = subprocess.run(cmd, capture_output=True, text=True, timeout=600) + except FileNotFoundError as e: + raise RuntimeError( + "trtexec not found on PATH; install TensorRT or drop --trt_build." + ) from e + except subprocess.TimeoutExpired as e: + raise RuntimeError(f"trtexec timed out building {onnx_path} after 600s.") from e + if result.returncode != 0: + raise RuntimeError( + f"TensorRT engine build failed for {onnx_path}:\n{result.stdout}\n{result.stderr}" + ) + print("TensorRT engine build succeeded.") + if __name__ == "__main__": main() diff --git a/modelopt/onnx/export/fp8_exporter.py b/modelopt/onnx/export/fp8_exporter.py index ffcbd894238..dcae618dd0a 100644 --- a/modelopt/onnx/export/fp8_exporter.py +++ b/modelopt/onnx/export/fp8_exporter.py @@ -101,13 +101,87 @@ def compress_weights(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return gs.export_onnx(graph) + @staticmethod + def _quantize_conv_weights_to_fp8(graph: gs.Graph) -> int: + """Add FP8 weight DequantizeLinear for Conv layers with unquantized weights. + + Conv weight quantizers are disabled during TorchScript ONNX export because the + TRT_FP8DequantizeLinear custom op produces outputs with unknown shapes, causing + the _convolution symbolic to fail. This method restores FP8 weight quantization + by inserting DQ nodes in the ONNX graph, mirroring the compress_weights logic. + + For each Conv node with an unquantized constant weight: + 1. Compute per-tensor scale = max(abs(weight)) / 448.0 + 2. Quantize weights to FP8E4M3FN + 3. Insert a DequantizeLinear(fp8_weights, scale) before the Conv weight input + + Args: + graph: The onnx-graphsurgeon graph to modify in-place. + + Returns: + Number of Conv weight DQ nodes inserted. + """ + fp8_max = 448.0 + count = 0 + + for node in list(graph.nodes): + if node.op != "Conv": + continue + if len(node.inputs) < 2: + continue + + weight_input = node.inputs[1] + if not isinstance(weight_input, gs.Constant): + continue + + # Skip if weight already has a DQ producer + if any(out.op == "DequantizeLinear" for out in weight_input.outputs): + continue + + torch_weights = torch.from_numpy(weight_input.values.copy()) + amax = torch_weights.abs().max().float() + if amax == 0: + continue + scale_val = (amax / fp8_max).item() + + # Quantize weights to FP8 (WAR: numpy doesn't support fp8) + fp8_data = (torch_weights / scale_val).to(torch.float8_e4m3fn).view(torch.uint8).numpy() + fp8_tensor = onnx.TensorProto() + fp8_tensor.data_type = onnx.TensorProto.FLOAT8E4M3FN + fp8_tensor.dims.extend(fp8_data.shape) + fp8_tensor.raw_data = fp8_data.tobytes() + fp8_constant = gs.Constant( + node.name + "/weight_quantizer/fp8_weights", LazyValues(fp8_tensor) + ) + + # Scale in FP16 — DQ output type matches scale dtype, must match activation type + import numpy as np + + scale_constant = gs.Constant( + node.name + "/weight_quantizer/scale", + np.array(scale_val, dtype=np.float16), + ) + + dq_output = gs.Variable(node.name + "/weight_quantizer/dq_output") + dq_node = gs.Node( + op="DequantizeLinear", + name=node.name + "/weight_quantizer/DequantizeLinear", + inputs=[fp8_constant, scale_constant], + outputs=[dq_output], + ) + graph.nodes.append(dq_node) + node.inputs[1] = dq_output + count += 1 + + return count + @staticmethod def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Post-processes the ONNX model for FP8 quantization. - Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear: - - TRT_FP8QuantizeLinear -> QuantizeLinear with FP8E4M3FN zero_point and saturate=1 - - TRT_FP8DequantizeLinear -> DequantizeLinear + Converts TRT_FP8 QDQ ops to native ONNX QuantizeLinear/DequantizeLinear and + adds FP8 weight DQ for Conv layers whose weight quantizers were disabled during + TorchScript export. Args: onnx_model: The ONNX model containing TRT_FP8 quantization nodes. @@ -144,5 +218,10 @@ def post_process(onnx_model: onnx.ModelProto) -> onnx.ModelProto: f"Converted {node.name} from TRT_FP8DequantizeLinear to DequantizeLinear" ) + # Add FP8 weight DQ for Conv layers that had weight quantizers disabled during export + count = FP8QuantExporter._quantize_conv_weights_to_fp8(graph) + if count > 0: + logger.info(f"Inserted FP8 weight DequantizeLinear for {count} Conv nodes") + graph.cleanup().toposort() return gs.export_onnx(graph) diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index 54efd0a1110..ac93bc2a26c 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -1504,6 +1504,192 @@ def remove_redundant_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: return onnx_model +def fold_dq_fp32_to_fp16_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + """Remove Cast(FP32->FP16) nodes after DequantizeLinear by setting DQ output to FP16. + + When convert_float_to_float16 blocks DequantizeLinear, it inserts Cast nodes to bridge + the FP32 DQ output to the FP16 graph. This function removes those Cast nodes by: + 1. Converting the DQ scale initializer from FP32 to FP16 + 2. Updating the DQ output type to FP16 in value_info + 3. Bypassing and removing the Cast node + + NVFP4 uses a nested DQ chain (scale is itself a DQ output). When the outer DQ's scale + is produced by another DQ, recursively retype the inner DQ's chain so the whole + chain produces FP16 tensors under strongly-typed TRT parsing. + + Args: + onnx_model: The ONNX model with DQ -> Cast(FP32->FP16) patterns. + + Returns: + The ONNX model with Cast nodes removed and DQ outputs set to FP16. + """ + import numpy as np + + dq_ops = {"DequantizeLinear", "TRT_FP8DequantizeLinear"} + + # Build a map of tensor name -> producer node + producer_map: dict[str, onnx.NodeProto] = {} + for node in onnx_model.graph.node: + for out in node.output: + producer_map[out] = node + + # Build initializer lookup + initializer_map: dict[str, onnx.TensorProto] = { + init.name: init for init in onnx_model.graph.initializer + } + + value_info_map: dict[str, onnx.ValueInfoProto] = { + vi.name: vi for vi in onnx_model.graph.value_info + } + + retyped_dq_outputs: set[str] = set() + + def _convert_fp32_init_to_fp16(init: onnx.TensorProto) -> None: + scale_data = np.frombuffer(init.raw_data, dtype=np.float32) + if not scale_data.size: + scale_data = np.array(init.float_data, dtype=np.float32) + init.data_type = onnx.TensorProto.FLOAT16 + init.raw_data = scale_data.astype(np.float16).tobytes() + del init.float_data[:] + + def _retype_dq_chain(dq_node: onnx.NodeProto, depth: int = 0) -> None: + """Propagate FP16 output type down through a DQ's scale chain.""" + if depth > 4 or len(dq_node.input) < 2: + return + scale_name = dq_node.input[1] + scale_init = initializer_map.get(scale_name) + if scale_init is not None: + if scale_init.data_type == onnx.TensorProto.FLOAT: + _convert_fp32_init_to_fp16(scale_init) + return + scale_producer = producer_map.get(scale_name) + if scale_producer is None or scale_producer.op_type not in dq_ops: + return + _retype_dq_chain(scale_producer, depth + 1) + if scale_name in value_info_map: + value_info_map[scale_name].type.tensor_type.elem_type = onnx.TensorProto.FLOAT16 + retyped_dq_outputs.add(scale_name) + + nodes_to_remove = [] + for node in onnx_model.graph.node: + if node.op_type != "Cast": + continue + + cast_to = None + for attr in node.attribute: + if attr.name == "to": + cast_to = attr.i + if cast_to != onnx.TensorProto.FLOAT16: + continue + + producer = producer_map.get(node.input[0]) + if producer is None or producer.op_type not in dq_ops: + continue + + _retype_dq_chain(producer) + + _bypass_cast_node(onnx_model, node) + nodes_to_remove.append(node) + + dq_output_name = producer.output[0] + retyped_dq_outputs.add(dq_output_name) + + for name in retyped_dq_outputs: + vi = value_info_map.get(name) + if vi is not None: + vi.type.tensor_type.elem_type = onnx.TensorProto.FLOAT16 + + logger.debug(f"Folded {len(nodes_to_remove)} DQ -> Cast(FP32->FP16) patterns") + for node in nodes_to_remove: + onnx_model.graph.node.remove(node) + + return onnx_model + + +def fold_qdq_scale_fp16_to_fp32_casts(onnx_model: onnx.ModelProto) -> onnx.ModelProto: + """Remove Cast(FP16->FP32) nodes feeding into Q/DQ scale inputs. + + When convert_float_to_float16 blocks QuantizeLinear/DequantizeLinear, it inserts + Cast(FP16->FP32) nodes before every scale input. In opset >=20 Q/DQ natively accept + FP16 scales, and leaving the cast in place forces DQ outputs to FP32, breaking + downstream FP16 matmul/add operations under strongly-typed TRT parsing. + + This function bypasses each such Cast and, when the upstream Constant is FP16, + wires the DQ output to FP16 in value_info so shape inference stays consistent. + + Args: + onnx_model: The ONNX model with Cast(FP16->FP32) -> Q/DQ.scale patterns. + + Returns: + The ONNX model with redundant scale-path casts removed. + """ + qdq_ops = { + "QuantizeLinear", + "DequantizeLinear", + "TRT_FP8QuantizeLinear", + "TRT_FP8DequantizeLinear", + } + + producer_map: dict[str, onnx.NodeProto] = {} + consumer_map: dict[str, list[tuple[onnx.NodeProto, int]]] = {} + for node in onnx_model.graph.node: + for out in node.output: + producer_map[out] = node + for idx, inp in enumerate(node.input): + if inp: + consumer_map.setdefault(inp, []).append((node, idx)) + + type_map = _build_tensor_type_map(onnx_model) + + nodes_to_remove: list[onnx.NodeProto] = [] + dq_outputs_retyped: set[str] = set() + visited_casts: set[int] = set() + for node in onnx_model.graph.node: + if node.op_type not in qdq_ops or len(node.input) < 2: + continue + + scale_name = node.input[1] + cast_node = producer_map.get(scale_name) + if cast_node is None or cast_node.op_type != "Cast": + continue + if id(cast_node) in visited_casts: + # Already handled (e.g. shared scale Cast across paired Q/DQ). + if node.op_type.endswith("DequantizeLinear"): + dq_outputs_retyped.add(node.output[0]) + continue + if get_cast_to_type(cast_node) != onnx.TensorProto.FLOAT: + continue + if type_map.get(cast_node.input[0]) != onnx.TensorProto.FLOAT16: + continue + + # Only bypass when every consumer of this Cast is a Q/DQ scale input; otherwise + # other ops would silently receive FP16 instead of the FP32 they requested. + cast_output = cast_node.output[0] + consumers = consumer_map.get(cast_output, []) + if not consumers or not all(c.op_type in qdq_ops and i == 1 for c, i in consumers): + continue + + # Bypass the cast so the scale stays FP16 + _bypass_cast_node(onnx_model, cast_node) + nodes_to_remove.append(cast_node) + visited_casts.add(id(cast_node)) + + # For DQ nodes, the output type follows the scale type — update value_info. + if node.op_type.endswith("DequantizeLinear"): + dq_outputs_retyped.add(node.output[0]) + + for vi in onnx_model.graph.value_info: + if vi.name in dq_outputs_retyped: + vi.type.tensor_type.elem_type = onnx.TensorProto.FLOAT16 + + logger.debug(f"Folded {len(nodes_to_remove)} Cast(FP16->FP32) -> Q/DQ.scale patterns") + for cast_node in nodes_to_remove: + if cast_node in onnx_model.graph.node: + onnx_model.graph.node.remove(cast_node) + + return onnx_model + + def remove_node_training_mode(onnx_model: onnx.ModelProto, node_op_type: str) -> onnx.ModelProto: """Remove `training_mode` attribute and extra training outputs from nodes of a given op type. diff --git a/modelopt/torch/_deploy/utils/torch_onnx.py b/modelopt/torch/_deploy/utils/torch_onnx.py index 8cb741dbc7a..9ec110b7887 100644 --- a/modelopt/torch/_deploy/utils/torch_onnx.py +++ b/modelopt/torch/_deploy/utils/torch_onnx.py @@ -16,6 +16,7 @@ """Utility functions related to Onnx.""" import base64 +import contextlib import inspect import json import logging @@ -46,6 +47,8 @@ from modelopt.onnx.utils import ( change_casts_to_fp16, check_model_uses_external_data, + fold_dq_fp32_to_fp16_casts, + fold_qdq_scale_fp16_to_fp32_casts, get_input_names, get_input_shapes, get_node_names, @@ -402,6 +405,29 @@ def is_fp8_quantized(model: nn.Module) -> bool: return False +@contextlib.contextmanager +def _disable_fp8_conv_weight_quantizers(model: nn.Module): + """Temporarily disable FP8 weight quantizers on Conv layers during ONNX export. + + The TorchScript ONNX exporter requires static kernel shapes for Conv operations, + but the TRT_FP8DequantizeLinear custom op produces outputs with unknown shapes in + the TorchScript IR, causing the _convolution symbolic to fail. Disabling Conv weight + quantizers during export allows the Conv to export with static-shape FP16/FP32 weights. + FP8 weight quantization is restored as a post-processing step in FP8QuantExporter. + """ + disabled = [] + for _, module in model.named_modules(): + if isinstance(module, (nn.Conv1d, nn.Conv2d, nn.Conv3d)): + if hasattr(module, "weight_quantizer") and module.weight_quantizer.is_enabled: + module.weight_quantizer.disable() + disabled.append(module) + try: + yield + finally: + for module in disabled: + module.weight_quantizer.enable() + + def quantize_weights(model: nn.Module, onnx_model: onnx.ModelProto) -> onnx.ModelProto: """Real quantizes the weights in the onnx model. @@ -522,7 +548,11 @@ def get_onnx_bytes_and_metadata( input_none_names = list(set(tree_spec_input.names) - set(input_names)) use_torch_autocast = not ( - is_fp4_quantized(model) or is_mxfp8_quantized(model) or weights_dtype == "fp32" + is_fp4_quantized(model) + or is_mxfp8_quantized(model) + or is_fp8_quantized(model) + or is_int8_quantized(model) + or weights_dtype == "fp32" ) autocast = torch.autocast("cuda") if use_torch_autocast else nullcontext() @@ -556,7 +586,13 @@ def get_onnx_bytes_and_metadata( if is_fp4_quantized(model) or is_mxfp8_quantized(model) else nullcontext() ) - with torch.inference_mode(), autocast, quantizer_context: + # Disable FP8 Conv weight quantizers: TorchScript custom ops produce outputs with + # unknown shapes, causing _convolution symbolic to fail. Conv weights are quantized + # to FP8 in post-processing by FP8QuantExporter instead. + conv_wq_context = ( + _disable_fp8_conv_weight_quantizers(model) if is_fp8_quantized(model) else nullcontext() + ) + with torch.inference_mode(), autocast, quantizer_context, conv_wq_context: additional_kwargs = {} if not dynamo_export: additional_kwargs["dynamic_axes"] = dynamic_axes @@ -598,7 +634,12 @@ def get_onnx_bytes_and_metadata( onnx_opt_graph = qdq_to_dq(onnx_opt_graph) if weights_dtype in ["fp16", "bf16"]: - if is_int4_quantized(model) or is_mxfp8_quantized(model) or is_fp8_quantized(model): + if ( + is_int4_quantized(model) + or is_mxfp8_quantized(model) + or is_fp8_quantized(model) + or is_int8_quantized(model) + ): assert weights_dtype == "fp16", "BF16 + MXFP8/INT4 mixed precision is not supported yet" onnx_opt_graph = convert_float_to_float16( onnx_opt_graph, @@ -610,6 +651,11 @@ def get_onnx_bytes_and_metadata( # Change FP32 cast nodes feeding into Concat/Add to FP16 op_list = ["Concat", "Add", "Sqrt", "LayerNormalization", "Clip", "Mul", "Exp"] onnx_opt_graph = change_casts_to_fp16(onnx_opt_graph, op_list) + # Remove Cast(FP32->FP16) nodes after DQ by setting DQ output to FP16 directly + onnx_opt_graph = fold_dq_fp32_to_fp16_casts(onnx_opt_graph) + # Remove Cast(FP16->FP32) feeding Q/DQ scales so DQ stays FP16 for downstream + # MatMul/Add layers under strongly-typed TRT parsing. + onnx_opt_graph = fold_qdq_scale_fp16_to_fp32_casts(onnx_opt_graph) else: onnx_opt_graph = convert_to_f16( onnx_opt_graph, low_precision_type=weights_dtype, keep_io_types=False diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index fe4fb8d70ad..05efe48842f 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -656,9 +656,36 @@ def export_fp4( @contextlib.contextmanager def configure_linear_module_onnx_quantizers(model): - """Sets the onnx export attributes for the given model.""" + """Sets the onnx export attributes for the given model. + + For modules with block quantization (NVFP4/MXFP8): + - Weight quantizers use "static" export (TRT_FP4QDQ for NVFP4, DQ-only for MXFP8) + - Input/activation quantizers use "dynamic" export (TRT_FP4DynamicQuantize, etc.) + + This must be set for ALL modules with block quantization, not just nn.Linear, + because models like ResNet have non-Linear modules (e.g., MaxPool2d) with NVFP4/MXFP8 + input quantizers that would otherwise default to the static path and produce + TRT_FP4QDQ nodes on activations (which the NVFP4 exporter cannot handle). + """ + sentinel = object() + originals: list[tuple] = [] for _, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - module.input_quantizer._onnx_quantizer_type = "dynamic" - module.weight_quantizer._onnx_quantizer_type = "static" - yield + for attr_name, new_value in ( + ("input_quantizer", "dynamic"), + ("weight_quantizer", "static"), + ): + quantizer = getattr(module, attr_name, None) + if quantizer is None or not quantizer.block_sizes: + continue + original = getattr(quantizer, "_onnx_quantizer_type", sentinel) + originals.append((quantizer, original)) + quantizer._onnx_quantizer_type = new_value + try: + yield + finally: + for quantizer, original in originals: + if original is sentinel: + if hasattr(quantizer, "_onnx_quantizer_type"): + delattr(quantizer, "_onnx_quantizer_type") + else: + quantizer._onnx_quantizer_type = original diff --git a/tests/_test_utils/torch/vision_models.py b/tests/_test_utils/torch/vision_models.py index 5fed1d20c11..942167d7639 100644 --- a/tests/_test_utils/torch/vision_models.py +++ b/tests/_test_utils/torch/vision_models.py @@ -117,6 +117,7 @@ def get_model_and_input(on_gpu: bool = False): # "dm_nfnet_f0", "efficientnet_b0", "swin_tiny_patch4_window7_224", + "resnet50", ], _create_timm_fn, ), diff --git a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py index 7c2692c1d99..6d6e0d9de57 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -14,9 +14,6 @@ # limitations under the License. -import os -import subprocess - import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command @@ -28,34 +25,9 @@ "vit_tiny": ("vit_tiny_patch16_224", '{"depth": 1}'), "swin_tiny": ("swin_tiny_patch4_window7_224", '{"depths": [1, 1, 1, 1]}'), "swinv2_tiny": ("swinv2_tiny_window8_256", '{"depths": [1, 1, 1, 1]}'), + "resnet50": ("resnet50", None), } -# Builder optimization level: 4 for low-bit modes, 3 otherwise -_LOW_BIT_MODES = {"fp8", "int8", "nvfp4"} - - -def _verify_trt_engine_build(onnx_save_path, quantize_mode): - """Verify the exported ONNX model can be compiled into a TensorRT engine.""" - example_dir = os.path.join( - os.path.dirname(__file__), "..", "..", "..", "examples", "torch_onnx" - ) - onnx_path = os.path.join(example_dir, onnx_save_path) - assert os.path.exists(onnx_path), f"ONNX file not found: {onnx_path}" - - opt_level = "4" if quantize_mode in _LOW_BIT_MODES else "3" - cmd = [ - "trtexec", - f"--onnx={onnx_path}", - "--stronglyTyped", - f"--builderOptimizationLevel={opt_level}", - ] - - result = subprocess.run(cmd, capture_output=True, text=True, timeout=600) - assert result.returncode == 0, ( - f"TensorRT engine build failed for {onnx_save_path} " - f"(mode={quantize_mode}):\n{result.stdout}\n{result.stderr}" - ) - @pytest.mark.parametrize("quantize_mode", _QUANT_MODES) @pytest.mark.parametrize("model_key", list(_MODELS)) @@ -63,7 +35,6 @@ def test_torch_onnx(model_key, quantize_mode): timm_model_name, model_kwargs = _MODELS[model_key] onnx_save_path = f"{model_key}.{quantize_mode}.onnx" - # Step 1: Quantize and export to ONNX cmd_parts = extend_cmd_parts( ["python", "torch_quant_to_onnx.py"], timm_model_name=timm_model_name, @@ -73,8 +44,5 @@ def test_torch_onnx(model_key, quantize_mode): calibration_data_size="1", num_score_steps="1", ) - cmd_parts.append("--no_pretrained") + cmd_parts.extend(["--no_pretrained", "--trt_build"]) run_example_command(cmd_parts, "torch_onnx") - - # Step 2: Verify the exported ONNX model builds a TensorRT engine - _verify_trt_engine_build(onnx_save_path, quantize_mode) From 2004779a6751407be75ecab9caab25e1e596e5f4 Mon Sep 17 00:00:00 2001 From: Farid Adilazuarda <42537562+faridlazuarda@users.noreply.github.com> Date: Sat, 18 Apr 2026 15:55:42 +0100 Subject: [PATCH 17/30] Update README.md for DMS (fix cd experimental/DMS to cd Model-Optimizer/experimental/DMS) (#879) ## What does this PR do? **Type of change:** ? **Overview:** ? ## Usage ```python # Add a code snippet demonstrating how to use this ``` ## Testing ## Before your PR is "*Ready for review*" - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes/No - **Did you write any new necessary tests?**: Yes/No - **Did you add or update any necessary documentation?**: Yes/No - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: Yes/No ## Additional Information ## Summary by CodeRabbit * **Documentation** * Updated DMS installation instructions to reflect the repository structure and correct directory navigation during setup. * Clarified the setup steps so users follow the accurate directory change before running installation commands. * Small wording improvements to reduce confusion during the installation process. Signed-off-by: Farid Adilazuarda <42537562+faridlazuarda@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- experimental/dms/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/experimental/dms/README.md b/experimental/dms/README.md index 5e49f011afa..2107c80079a 100644 --- a/experimental/dms/README.md +++ b/experimental/dms/README.md @@ -40,7 +40,7 @@ Clone and install: ```bash git clone https://github.com/NVIDIA/Model-Optimizer -cd experimental/dms +cd Model-Optimizer/experimental/dms pip install -e . ``` From 2b315eda47996b00f8bae8d9461135b28cda7605 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Sat, 18 Apr 2026 21:19:36 +0530 Subject: [PATCH 18/30] Replace mip package with pulp (#663) Replace mip package with more popular pulp package for puzzle mip solving. Both use the CBC solver under the hood ## Testing - Results very close for Qwen3-8B and Nemotron-Nano-12B-v2 ## Summary by CodeRabbit * **Chores** * Simplified GPU test environment setup by removing unnecessary system dependency installation * Updated internal optimization solver dependencies in the puzzletron module Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .../mip/mip_with_multi_layer_replacements.py | 40 +++++++++---------- pyproject.toml | 1 - 2 files changed, 19 insertions(+), 22 deletions(-) diff --git a/modelopt/torch/puzzletron/mip/mip_with_multi_layer_replacements.py b/modelopt/torch/puzzletron/mip/mip_with_multi_layer_replacements.py index a906f886364..6fba86a873a 100644 --- a/modelopt/torch/puzzletron/mip/mip_with_multi_layer_replacements.py +++ b/modelopt/torch/puzzletron/mip/mip_with_multi_layer_replacements.py @@ -24,7 +24,7 @@ from random import random from typing import Any, TypeAlias -from mip import BINARY, Model, maximize, minimize, xsum +import pulp from .utils import consecutive_ngrams, get_nested_key, sort_replacements @@ -55,16 +55,15 @@ def run_mip( ) print("\n\n\n") - if not replacements: - return [], 0.0, {} - - mip_model = Model() + # Create pulp problem with appropriate sense (minimize or maximize) + sense = pulp.LpMaximize if bigger_is_better else pulp.LpMinimize + problem = pulp.LpProblem(name="multi_layer_replacement", sense=sense) objective_vars = [] constraint_vars = {constraint_key: [] for constraint_key in constraints} choice_indicators_by_layer = defaultdict(list) - for replacement_id, replacement in replacements.items(): - is_chosen = mip_model.add_var(var_type=BINARY) + for i, (replacement_id, replacement) in enumerate(replacements.items()): + is_chosen = pulp.LpVariable(f"choice_{i}", cat=pulp.LpBinary) replacement["is_chosen"] = is_chosen for parent_layer_idx in replacement["parent_layer_indices"]: @@ -79,7 +78,7 @@ def run_mip( # MIP constraints: each parent layer must come from exactly one chosen replacement for parent_layer_idx, curr_choice_indicators in choice_indicators_by_layer.items(): - mip_model += xsum(curr_choice_indicators) == 1 + problem += pulp.lpSum(curr_choice_indicators) == 1 # MIP constraints: the sum of chosen replacement costs must be lower than the max cost for constraint_key, max_cost in constraints.items(): @@ -87,22 +86,21 @@ def run_mip( if isinstance(max_cost, Iterable): min_cost, max_cost = max_cost - if max_cost is not None: - mip_model += xsum(constraint_vars[constraint_key]) <= max_cost - if min_cost is not None: - mip_model += xsum(constraint_vars[constraint_key]) >= min_cost + # PuLP is stricter than mip - it doesn't allow NaN/inf in constraints + if max_cost is not None and math.isfinite(max_cost): + problem += pulp.lpSum(constraint_vars[constraint_key]) <= max_cost + if min_cost is not None and math.isfinite(min_cost): + problem += pulp.lpSum(constraint_vars[constraint_key]) >= min_cost # MIP objective - mip_model.objective = ( - maximize(xsum(objective_vars)) if bigger_is_better else minimize(xsum(objective_vars)) - ) - - if max_seconds_per_solution is not None: - mip_model.max_seconds = max_seconds_per_solution + problem += (pulp.lpSum(objective_vars), "objective") - mip_model.optimize() + # Configure and run solver + solver = pulp.PULP_CBC_CMD(msg=True, timeLimit=max_seconds_per_solution) + problem.solve(solver) - if is_chosen.x is None: + # Check if solution is feasible + if problem.status != pulp.LpStatusOptimal: return [] # raise InfeasibleError() @@ -112,7 +110,7 @@ def run_mip( chosen_replacements: ChosenReplacements = [] chosen_layers = [] for replacement_id, replacement in replacements.items(): - is_chosen = replacement["is_chosen"].x >= 0.99 + is_chosen = replacement["is_chosen"].varValue >= 0.99 if is_chosen: assert replacement not in chosen_replacements chosen_replacements.append(replacement) diff --git a/pyproject.toml b/pyproject.toml index 2993759ec10..25cb6338aa5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -91,7 +91,6 @@ puzzletron = [ # Dependedencies for modelopt.torch.puzzletron subpackage "hydra-core==1.3.2", "immutabledict", "lru-dict", - "mip", "pandas", "typeguard", ] From 3d0f0db49ed77bfe58d2911944271054fa32d11c Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Sat, 18 Apr 2026 22:26:02 +0530 Subject: [PATCH 19/30] [CI] Replace tox with nox, use nemo:26.04 for megatron tests, and simplify CI workflows (#1286) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: New feature / infrastructure improvement Follow-up to #1285 for correct CI test environment for megatron based tests Replaces `tox` + `tox-current-env` with `nox` for all test, lint, docs, and wheel build sessions. The primary motivation was that `tox-current-env` is incompatible with uv venvs in NGC containers (e.g. NeMo's `/opt/venv`) — it picks the system Python via `sys._base_executable` instead of the container's venv Python which has megatron packages pre-installed. Key changes: - **`noxfile.py`** replaces `tox.ini` with GPU, CPU unit, partial-install, pre-commit, docs, and wheel sessions - **GPU sessions** use `venv_backend="none"` (run directly in container env) and `python -m pip/pytest` to avoid PATH mismatches - **uv** is set as the default venv backend (if available) for CPU sessions (faster installs) Also includes CI workflow simplifications: - **`_pr_gate.yml`** new reusable workflow centralizing file-change detection + linux-check wait logic (was duplicated across 3 workflow files) - **Collapsed pr/non-pr job pairs** into single jobs with conditional `runs-on` in `gpu_tests.yml`, `example_tests.yml`, `regression_tests.yml` - **Collapsed `multi-py` / `multi-torch` / `multi-transformers`** into a single `multi-version` matrix job in `unit_tests.yml` - **PR path filtering** for unit test secondary jobs (multi-version, launcher, partial-install) — skipped if no relevant files changed - **Fixed schedule/workflow_dispatch skipping** — jobs with `needs: [pr-gate]` were incorrectly skipped when all pr-gate internal jobs were skipped; fixed by making the gate job always run - **multi-version, launcher, partial-install** now also run on `schedule` / `workflow_dispatch` ### Usage ```bash python -m pip install nox uv # install nox and uv (once) nox -l # list all sessions nox -s gpu_megatron # run a GPU session (inside container) nox -s "unit-3.12(torch_211, tf_latest)" # run a specific unit test combination nox -s "unit-3.12(torch_211, tf_latest)" -R # force-recreate venv (e.g. after dep changes) COVERAGE_PROCESS_START=pyproject.toml nox -s "unit-3.12(torch_211, tf_latest)" # with coverage ``` ### Testing - Ran `nox -l` to verify all session names - Ran `gpu_megatron` session locally inside NeMo container — confirmed it uses `/opt/venv/bin/python` correctly - Manually triggered nightly-runs: - Unit: https://github.com/NVIDIA/Model-Optimizer/actions/runs/24608013657 - GPU: https://github.com/NVIDIA/Model-Optimizer/actions/runs/24608018763 - Examples: https://github.com/NVIDIA/Model-Optimizer/actions/runs/24608017322 ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: N/A — CI infrastructure only - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ (added `nox` and `uv` to `dev-test`, both Apache-2.0) - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A — no user-facing changes ### Additional Information Supersedes the tox-current-env workaround in the parent branch. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 --- .github/CODEOWNERS | 2 +- .github/workflows/_example_tests_runner.yml | 1 - .github/workflows/_pr_gate.yml | 65 +++ .github/workflows/bump_uv_lock.yml | 3 +- .github/workflows/code_quality.yml | 8 +- .github/workflows/example_tests.yml | 144 ++---- .github/workflows/gpu_tests.yml | 99 ++-- .github/workflows/pages.yml | 14 +- .github/workflows/regression_tests.yml | 77 ++++ .github/workflows/release.yml | 6 +- .github/workflows/unit_tests.yml | 122 ++--- .vscode/settings.json | 1 - CLAUDE.md | 8 +- CONTRIBUTING.md | 2 +- .../torch/export/unified_export_megatron.py | 6 +- noxfile.py | 195 ++++++++ pyproject.toml | 14 +- .../export/test_unified_export_megatron.py | 3 +- .../torch/peft/plugins/test_megatron_peft.py | 21 +- .../torch/speculative/test_dflash.py | 0 tests/unit/torch/speculative/conftest.py | 21 + tools/launcher/tests/conftest.py | 4 +- tox.ini | 149 ------ uv.lock | 426 +++++------------- 24 files changed, 647 insertions(+), 744 deletions(-) create mode 100644 .github/workflows/_pr_gate.yml create mode 100644 .github/workflows/regression_tests.yml create mode 100644 noxfile.py rename tests/{gpu_regression => regression}/torch/speculative/test_dflash.py (100%) create mode 100644 tests/unit/torch/speculative/conftest.py delete mode 100644 tox.ini diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 9bac2756b8b..19188206ea0 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -9,7 +9,7 @@ LICENSE @NVIDIA/modelopt-setup-codeowners LICENSE_HEADER @NVIDIA/modelopt-setup-codeowners pyproject.toml @NVIDIA/modelopt-setup-codeowners SECURITY.md @NVIDIA/modelopt-setup-codeowners -tox.ini @NVIDIA/modelopt-setup-codeowners +noxfile.py @NVIDIA/modelopt-setup-codeowners uv.lock @NVIDIA/modelopt-setup-codeowners # Library diff --git a/.github/workflows/_example_tests_runner.yml b/.github/workflows/_example_tests_runner.yml index b34aa87e0de..8adadbac7af 100644 --- a/.github/workflows/_example_tests_runner.yml +++ b/.github/workflows/_example_tests_runner.yml @@ -48,7 +48,6 @@ jobs: - name: Install dependencies run: | # use `python -m pip` instead of `pip` to avoid conflicts with system pip for nemo containers - pip uninstall -y nvidia-modelopt python -m pip install ".${{ inputs.pip_install_extras }}" if [[ "${{ inputs.example }}" == *"diffusers"* ]]; then diff --git a/.github/workflows/_pr_gate.yml b/.github/workflows/_pr_gate.yml new file mode 100644 index 00000000000..d1e6dad0f9f --- /dev/null +++ b/.github/workflows/_pr_gate.yml @@ -0,0 +1,65 @@ +name: PR Gate + +on: + workflow_call: + inputs: + files: + description: "Newline-separated list of file patterns to watch for changes" + required: true + type: string + outputs: + any_changed: + description: "Whether any relevant files changed" + value: ${{ jobs.check-file-changes.outputs.any_changed }} + +jobs: + check-file-changes: + runs-on: ubuntu-latest + outputs: + any_changed: ${{ steps.changed-tests.outputs.any_changed || steps.non-pr.outputs.any_changed }} + steps: + # For non-PR triggers (schedule, workflow_dispatch), always run tests + - id: non-pr + if: "!startsWith(github.ref, 'refs/heads/pull-request/')" + run: echo "any_changed=true" >> $GITHUB_OUTPUT + - if: startsWith(github.ref, 'refs/heads/pull-request/') + uses: actions/checkout@v6 + with: + fetch-depth: 0 + - if: startsWith(github.ref, 'refs/heads/pull-request/') + id: get-pr-info + uses: nv-gha-runners/get-pr-info@main + # Extract SHAs from pr-info JSON via shell to avoid fromJSON on potentially-empty outputs + - if: startsWith(github.ref, 'refs/heads/pull-request/') + id: pr-shas + env: + PR_INFO: ${{ steps.get-pr-info.outputs.pr-info }} + run: | + echo "head_sha=$(echo "$PR_INFO" | jq -r '.head.sha')" >> $GITHUB_OUTPUT + echo "base_sha=$(echo "$PR_INFO" | jq -r '.base.sha')" >> $GITHUB_OUTPUT + # Get commit from main branch that is present in the PR to use as base for changed files + - if: startsWith(github.ref, 'refs/heads/pull-request/') + id: calculate-merge-base + run: | + (echo -n "merge-base="; git merge-base "${{ steps.pr-shas.outputs.base_sha }}" "${{ steps.pr-shas.outputs.head_sha }}") | tee --append "${GITHUB_OUTPUT}" + - if: startsWith(github.ref, 'refs/heads/pull-request/') + name: Check for changes in test-relevant directories + id: changed-tests + uses: step-security/changed-files@v46.0.5 + with: + base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }} + sha: ${{ steps.pr-shas.outputs.head_sha }} + files: ${{ inputs.files }} + fail_on_initial_diff_error: true + wait-checks: + needs: [check-file-changes] + if: >- + startsWith(github.ref, 'refs/heads/pull-request/') && + needs.check-file-changes.outputs.any_changed == 'true' + uses: ./.github/workflows/_wait_for_checks.yml + permissions: + checks: read + secrets: inherit + with: + match_pattern: "^linux$" # Wait for Unit tests / linux (DCO is a prerequisite of linux) + delay: 300s diff --git a/.github/workflows/bump_uv_lock.yml b/.github/workflows/bump_uv_lock.yml index 51f84e5c277..e0933418b6e 100644 --- a/.github/workflows/bump_uv_lock.yml +++ b/.github/workflows/bump_uv_lock.yml @@ -3,7 +3,8 @@ name: Bump uv.lock on: schedule: - cron: "0 9 * * 1" # Every Monday at 9:00 UTC - workflow_dispatch: # On-demand + workflow_dispatch: + # On-demand permissions: contents: write diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml index 723f3af1209..3330d87332a 100644 --- a/.github/workflows/code_quality.yml +++ b/.github/workflows/code_quality.yml @@ -5,10 +5,12 @@ on: branches: [main, release/*, feature/*] schedule: - cron: "0 0 * * *" # Nightly - workflow_dispatch: # On-demand + workflow_dispatch: + # On-demand + -# Cancel previous runs if new commit is pushed to the same PR concurrency: + # Cancel previous runs if new commit is pushed to the same PR group: ${{ github.workflow }}-${{ github.event.pull_request.number }} cancel-in-progress: true @@ -24,4 +26,4 @@ jobs: with: extra_args: --results=verified,unknown - name: Run code quality checks - run: pip install tox && tox -e pre-commit-all + run: pip install nox uv && nox -s pre_commit_all diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index 73f5c1fa9b7..e8d307ccef9 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -6,61 +6,34 @@ on: # NOTE: paths cannot be used since push happens to copied PR and only latest commit to PR is used schedule: - cron: "0 0 * * *" # Nightly - workflow_dispatch: # On-demand + workflow_dispatch: + # On-demand + -# Cancel previous runs if new commit is pushed to the same PR concurrency: + # Cancel previous runs if new commit is pushed to the same PR group: ${{ github.workflow }}-${{ startsWith(github.ref, 'refs/heads/pull-request/') && github.ref || github.sha }} cancel-in-progress: true jobs: - check-file-changes: - if: startsWith(github.ref, 'refs/heads/pull-request/') - runs-on: ubuntu-latest - outputs: - any_changed: ${{ steps.changed-tests.outputs.any_changed }} - steps: - - uses: actions/checkout@v6 - with: - fetch-depth: 0 - - id: get-pr-info - uses: nv-gha-runners/get-pr-info@main - # Get commit from main branch that is present in the PR to use as base for changed files - - id: calculate-merge-base - env: - PR_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }} - BASE_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.sha }} - run: | - (echo -n "merge-base="; git merge-base "$BASE_SHA" "$PR_SHA") | tee --append "${GITHUB_OUTPUT}" - - name: Check for changes in test-relevant directories - id: changed-tests - uses: step-security/changed-files@v46.0.5 - with: - base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }} - sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }} - files: | - .github/workflows/example_tests.yml - examples/** - modelopt/** - pyproject.toml - tests/examples/** - fail_on_initial_diff_error: true - wait-checks: - needs: [check-file-changes] - if: needs.check-file-changes.outputs.any_changed == 'true' - uses: ./.github/workflows/_wait_for_checks.yml + pr-gate: + uses: ./.github/workflows/_pr_gate.yml permissions: checks: read secrets: inherit with: - match_pattern: "^DCO$|^linux$" # Wait for DCO and Unit tests / linux to pass - delay: 300s + files: | + .github/workflows/example_tests.yml + examples/** + modelopt/** + pyproject.toml + tests/examples/** ##### PyTorch Example Tests (speculative_decoding requires 26.01 image) ##### - torch-pr: - needs: [check-file-changes, wait-checks] - if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true' - strategy: &torch_strategy + torch: + needs: [pr-gate] + if: needs.pr-gate.outputs.any_changed == 'true' + strategy: fail-fast: false matrix: example: [llm_distill, llm_qat, llm_sparsity, diffusers_sparsity] @@ -74,24 +47,12 @@ jobs: example: ${{ matrix.example }} timeout_minutes: 30 pip_install_extras: "[hf,dev-test]" - runner: linux-amd64-gpu-rtxpro6000-latest-1 - - torch-non-pr: - if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }} - strategy: *torch_strategy - uses: ./.github/workflows/_example_tests_runner.yml - secrets: inherit - with: - docker_image: "nvcr.io/nvidia/pytorch:${{ matrix.docker_image || '26.03' }}-py3" - example: ${{ matrix.example }} - timeout_minutes: 30 - pip_install_extras: "[hf,dev-test]" - runner: linux-amd64-gpu-rtxpro6000-latest-2 + runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} - ##### TensorRT-LLM Example Tests ##### + ##### TensorRT-LLM Example Tests (pr/non-pr split: non-pr runs extra autodeploy+eval examples) ##### trtllm-pr: - needs: [check-file-changes, wait-checks] - if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true' + needs: [pr-gate] + if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.pr-gate.outputs.any_changed == 'true' strategy: fail-fast: false matrix: @@ -118,40 +79,24 @@ jobs: pip_install_extras: "[hf,dev-test]" runner: linux-amd64-gpu-rtxpro6000-latest-2 - ##### NeMo Example Tests ##### - nemo-pr: - needs: [check-file-changes, wait-checks] - if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true' - strategy: &nemo_strategy - fail-fast: false - matrix: - example: [megatron_bridge] + ##### Megatron Example Tests ##### + megatron: + needs: [pr-gate] + if: needs.pr-gate.outputs.any_changed == 'true' uses: ./.github/workflows/_example_tests_runner.yml secrets: inherit with: docker_image: "nvcr.io/nvidia/nemo:26.02" - example: ${{ matrix.example }} + example: megatron_bridge timeout_minutes: 30 pip_install_extras: "[hf,puzzletron,dev-test]" - runner: linux-amd64-gpu-rtxpro6000-latest-1 - - nemo-non-pr: - if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }} - strategy: *nemo_strategy - uses: ./.github/workflows/_example_tests_runner.yml - secrets: inherit - with: - docker_image: "nvcr.io/nvidia/nemo:26.02" - example: ${{ matrix.example }} - timeout_minutes: 30 - pip_install_extras: "[hf,puzzletron,dev-test]" - runner: linux-amd64-gpu-rtxpro6000-latest-2 + runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} ##### ONNX/TensorRT Example Tests ##### - onnx-pr: - needs: [check-file-changes, wait-checks] - if: startsWith(github.ref, 'refs/heads/pull-request/') && needs.check-file-changes.outputs.any_changed == 'true' - strategy: &onnx_strategy + onnx: + needs: [pr-gate] + if: needs.pr-gate.outputs.any_changed == 'true' + strategy: fail-fast: false matrix: example: [diffusers, torch_onnx] @@ -160,34 +105,23 @@ jobs: with: docker_image: "nvcr.io/nvidia/tensorrt:26.02-py3" example: ${{ matrix.example }} - pip_install_extras: "[all,dev-test]" - runner: linux-amd64-gpu-rtxpro6000-latest-1 - - onnx-non-pr: - if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }} - strategy: *onnx_strategy - uses: ./.github/workflows/_example_tests_runner.yml - secrets: inherit - with: - docker_image: "nvcr.io/nvidia/tensorrt:26.02-py3" - example: ${{ matrix.example }} - pip_install_extras: "[all,dev-test]" - runner: linux-amd64-gpu-rtxpro6000-latest-2 + pip_install_extras: "[onnx,hf,dev-test]" + runner: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} ##### Required Check for PR ##### example-pr-required-check: # Run even if example tests are skipped if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} - needs: [check-file-changes, torch-pr, trtllm-pr, nemo-pr, onnx-pr] + needs: [pr-gate, torch, trtllm-pr, megatron, onnx] runs-on: ubuntu-latest steps: - - name: Required GPU tests did not succeed + - name: Required example tests did not succeed if: | - needs.check-file-changes.result != 'success' || - (needs.check-file-changes.outputs.any_changed == 'true' && ( - needs.torch-pr.result != 'success' || + needs.pr-gate.result != 'success' || + (needs.pr-gate.outputs.any_changed == 'true' && ( + needs.torch.result != 'success' || needs.trtllm-pr.result != 'success' || - needs.nemo-pr.result != 'success' || - needs.onnx-pr.result != 'success' + needs.megatron.result != 'success' || + needs.onnx.result != 'success' )) run: exit 1 diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index 30b47bb2160..628aead7ee8 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -6,64 +6,35 @@ on: # NOTE: paths cannot be used since push happens to copied PR and only latest commit to PR is used schedule: - cron: "0 0 * * *" # Nightly - workflow_dispatch: # On-demand + workflow_dispatch: + # On-demand + -# Cancel previous runs if new commit is pushed to the same PR concurrency: + # Cancel previous runs if new commit is pushed to the same PR group: ${{ github.workflow }}-${{ startsWith(github.ref, 'refs/heads/pull-request/') && github.ref || github.sha }} cancel-in-progress: true jobs: - check-file-changes: - if: startsWith(github.ref, 'refs/heads/pull-request/') - runs-on: ubuntu-latest - outputs: - any_changed: ${{ steps.changed-tests.outputs.any_changed }} - steps: - - uses: actions/checkout@v6 - with: - fetch-depth: 0 - - id: get-pr-info - uses: nv-gha-runners/get-pr-info@main - # Get commit from main branch that is present in the PR to use as base for changed files - - id: calculate-merge-base - env: - PR_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }} - BASE_SHA: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).base.sha }} - run: | - (echo -n "merge-base="; git merge-base "$BASE_SHA" "$PR_SHA") | tee --append "${GITHUB_OUTPUT}" - - name: Check for changes in test-relevant directories - id: changed-tests - uses: step-security/changed-files@v46.0.5 - with: - base_sha: ${{ steps.calculate-merge-base.outputs.merge-base }} - sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info).head.sha }} - files: | - .github/workflows/gpu_tests.yml - modelopt/** - tests/gpu/** - tests/gpu_regression/** - examples/speculative_decoding/** - examples/dataset/** - modelopt_recipes/general/speculative_decoding/** - tools/launcher/** - pyproject.toml - tox.ini - fail_on_initial_diff_error: true - wait-checks: - needs: [check-file-changes] - if: needs.check-file-changes.outputs.any_changed == 'true' - uses: ./.github/workflows/_wait_for_checks.yml + pr-gate: + uses: ./.github/workflows/_pr_gate.yml permissions: checks: read secrets: inherit with: - match_pattern: "^DCO$|^linux$" # Wait for DCO and Unit tests / linux to pass - delay: 300s - gpu-tests-pr: - needs: [check-file-changes, wait-checks] - if: needs.check-file-changes.outputs.any_changed == 'true' - strategy: &gpu_strategy + files: | + .github/workflows/gpu_tests.yml + modelopt/** + noxfile.py + pyproject.toml + tests/gpu/** + tests/gpu_megatron/** + tests/gpu_trtllm/** + + gpu-tests: + needs: [pr-gate] + if: needs.pr-gate.outputs.any_changed == 'true' + strategy: fail-fast: false matrix: include: @@ -71,24 +42,21 @@ jobs: timeout: 60 container_image: pytorch:26.01-py3 # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers until https://github.com/tbenthompson/cppimport/pull/98 - - example: gpu-regression - timeout: 15 - container_image: pytorch:26.01-py3 - - example: gpu-megatron + - example: gpu_megatron timeout: 45 - container_image: pytorch:26.01-py3 - - example: gpu-trtllm + container_image: nemo:26.04 + - example: gpu_trtllm timeout: 30 container_image: tensorrt-llm/release:1.3.0rc10 - runs-on: linux-amd64-gpu-rtxpro6000-latest-1 + runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} timeout-minutes: ${{ matrix.timeout }} - container: &gpu_container + container: image: nvcr.io/nvidia/${{ matrix.container_image }} env: GIT_DEPTH: 1000 # For correct version PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages HF_TOKEN: ${{ secrets.HF_TOKEN }} - steps: &gpu_steps + steps: - uses: actions/checkout@v6 - uses: nv-gha-runners/setup-proxy-cache@main - name: Setup environment variables @@ -99,8 +67,7 @@ jobs: COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml COVERAGE_FILE: ${{ github.workspace }}/.coverage run: | - pip install tox-current-env - COV_ARGS="--cov" tox -e cuda13-${{ matrix.example }} --current-env + python -m pip install nox && nox -s ${{ matrix.example }} - name: Upload GPU coverage to Codecov uses: codecov/codecov-action@v5 with: @@ -109,19 +76,13 @@ jobs: flags: gpu fail_ci_if_error: false # test may be skipped if relevant file changes are not detected verbose: true - gpu-tests-non-pr: - if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }} - strategy: *gpu_strategy - runs-on: linux-amd64-gpu-rtxpro6000-latest-2 - timeout-minutes: ${{ matrix.timeout }} - container: *gpu_container - steps: *gpu_steps + gpu-pr-required-check: - # Run even if gpu-tests-pr is skipped + # Run even if gpu-tests is skipped if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} - needs: [check-file-changes, gpu-tests-pr] + needs: [pr-gate, gpu-tests] runs-on: ubuntu-latest steps: - name: Required GPU tests did not succeed - if: ${{ needs.check-file-changes.result != 'success' || (needs.check-file-changes.outputs.any_changed == 'true' && needs.gpu-tests-pr.result != 'success') }} + if: ${{ needs.pr-gate.result != 'success' || (needs.pr-gate.outputs.any_changed == 'true' && needs.gpu-tests.result != 'success') }} run: exit 1 diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index 6789142f97f..43e2b5cc78d 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -8,16 +8,18 @@ on: branches: [main] schedule: - cron: "0 0 * * *" # Nightly - workflow_dispatch: # On-demand + workflow_dispatch: + # On-demand + -# Cancel previous runs if new commit is pushed concurrency: + # Cancel previous runs if new commit is pushed group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} cancel-in-progress: true permissions: - contents: write # push to gh-pages branch - pull-requests: write # post/update preview URL comment on PRs + contents: write # push to gh-pages branch + pull-requests: write # post/update preview URL comment on PRs jobs: build-docs: @@ -27,7 +29,7 @@ jobs: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup - name: Build docs - run: pip install tox && tox -e build-docs + run: pip install nox uv && nox -s docs - name: Upload docs artifact if: github.event_name == 'push' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' uses: actions/upload-artifact@v4 @@ -47,7 +49,7 @@ jobs: - uses: ./.github/actions/ubuntu-setup - name: Build docs if: github.event.action != 'closed' - run: pip install tox && tox -e build-docs + run: pip install nox uv && nox -s docs - name: Deploy / remove PR preview uses: rossjrw/pr-preview-action@v1 with: diff --git a/.github/workflows/regression_tests.yml b/.github/workflows/regression_tests.yml new file mode 100644 index 00000000000..3e0fd6aba6f --- /dev/null +++ b/.github/workflows/regression_tests.yml @@ -0,0 +1,77 @@ +name: Regression tests + +on: + push: + branches: ["pull-request/[0-9]+"] + # NOTE: paths cannot be used since push happens to copied PR and only latest commit to PR is used + schedule: + - cron: "0 0 * * *" # Nightly + workflow_dispatch: + # On-demand + + +concurrency: + # Cancel previous runs if new commit is pushed to the same PR + group: ${{ github.workflow }}-${{ startsWith(github.ref, 'refs/heads/pull-request/') && github.ref || github.sha }} + cancel-in-progress: true + +jobs: + pr-gate: + uses: ./.github/workflows/_pr_gate.yml + permissions: + checks: read + secrets: inherit + with: + files: | + .github/workflows/regression_tests.yml + modelopt/torch/** + noxfile.py + pyproject.toml + tests/regression/** + examples/speculative_decoding/** + examples/dataset/** + modelopt_recipes/general/speculative_decoding/** + tools/launcher/** + + regression-tests: + needs: [pr-gate] + if: needs.pr-gate.outputs.any_changed == 'true' + runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} + timeout-minutes: 15 + container: + image: nvcr.io/nvidia/pytorch:26.01-py3 + env: + GIT_DEPTH: 1000 # For correct version + PIP_CONSTRAINT: "" # Disable pip constraint for upgrading packages + HF_TOKEN: ${{ secrets.HF_TOKEN }} + steps: + - uses: actions/checkout@v6 + - uses: nv-gha-runners/setup-proxy-cache@main + - name: Setup environment variables + run: | + echo "LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/include:/usr/lib/x86_64-linux-gnu" >> $GITHUB_ENV + - name: Run regression tests + env: + COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml + COVERAGE_FILE: ${{ github.workspace }}/.coverage + run: python -m pip install nox && nox -s regression + - name: Upload regression coverage to Codecov + uses: codecov/codecov-action@v5 + with: + token: ${{ secrets.CODECOV_TOKEN }} + files: coverage.xml + flags: regression + fail_ci_if_error: false # test may be skipped if relevant file changes are not detected + verbose: true + + regression-pr-required-check: + # Run even if regression-tests is skipped + if: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && always() }} + needs: [pr-gate, regression-tests] + runs-on: ubuntu-latest + steps: + - name: Required regression tests did not succeed + if: | + needs.pr-gate.result != 'success' || + (needs.pr-gate.outputs.any_changed == 'true' && needs.regression-tests.result != 'success') + run: exit 1 diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index e3e62d78def..35edb9e53ef 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -13,13 +13,11 @@ jobs: steps: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup - - name: Install dependencies - run: pip install tox - name: Run basic unit tests - run: tox -e py312-torch29-tf_latest-unit + run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_latest)" - name: Build Wheel run: | - tox -e build-wheel + nox -s build_wheel echo "WHEEL_PATH=$(find dist -name "*.whl" | head -n 1)" >> $GITHUB_ENV - name: Upload GitHub Release Artifact env: diff --git a/.github/workflows/unit_tests.yml b/.github/workflows/unit_tests.yml index afa57fc3a86..9832f0cc605 100644 --- a/.github/workflows/unit_tests.yml +++ b/.github/workflows/unit_tests.yml @@ -9,16 +9,18 @@ on: paths: - ".github/workflows/unit_tests.yml" - "modelopt/**" - - "tests/unit/**" + - "noxfile.py" - "pyproject.toml" - - "tox.ini" + - "tests/unit/**" - "tools/launcher/**" schedule: - cron: "0 0 * * *" # Nightly - workflow_dispatch: # On-demand + workflow_dispatch: + # On-demand + -# Cancel previous runs if new commit is pushed concurrency: + # Cancel previous runs if new commit is pushed group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }} cancel-in-progress: true @@ -30,6 +32,30 @@ jobs: secrets: inherit with: match_pattern: "^DCO$" + check-file-changes: + runs-on: ubuntu-latest + outputs: + any_changed: ${{ steps.changed.outputs.any_changed || steps.non-pr.outputs.any_changed }} + steps: + - id: non-pr + if: github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' + run: echo "any_changed=true" >> $GITHUB_OUTPUT + - if: github.event_name == 'pull_request' + uses: actions/checkout@v6 + with: + fetch-depth: 0 + - if: github.event_name == 'pull_request' + name: Check for changes in test-relevant paths + id: changed + uses: step-security/changed-files@v46.0.5 + with: + files: | + .github/workflows/unit_tests.yml + modelopt/** + noxfile.py + pyproject.toml + tests/unit/** + tools/launcher/** linux: needs: [check-dco] runs-on: ubuntu-latest @@ -38,7 +64,10 @@ jobs: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup - name: Run unit tests - run: pip install tox && COV_ARGS="--cov" tox -e py312-torch211-tf_latest-unit + env: + COVERAGE_PROCESS_START: ${{ github.workspace }}/pyproject.toml + COVERAGE_FILE: ${{ github.workspace }}/.coverage + run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_latest)" - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v5 with: @@ -47,8 +76,8 @@ jobs: fail_ci_if_error: true verbose: true windows: - if: github.event_name == 'pull_request' - needs: [linux] + if: needs.check-file-changes.outputs.any_changed == 'true' + needs: [linux, check-file-changes] runs-on: windows-latest timeout-minutes: 30 steps: @@ -57,55 +86,47 @@ jobs: with: python-version: "3.12" - name: Run unit tests (without coverage) - # Some issues with torch 2.10 on Windows, so using 2.9 for now - run: pip install tox && tox -e py312-torch29-tf_latest-unit - multi-py: - if: github.event_name == 'pull_request' - needs: [linux] + run: pip install nox uv && nox -s "unit-3.12(torch_211, tf_latest)" + multi-version: + if: needs.check-file-changes.outputs.any_changed == 'true' + needs: [linux, check-file-changes] runs-on: ubuntu-latest timeout-minutes: 30 strategy: fail-fast: false matrix: - py: [10, 11, 13] + include: + - {nox_session: "unit-3.10(torch_211, tf_latest)", python_version: "3.10"} + - {nox_session: "unit-3.11(torch_211, tf_latest)", python_version: "3.11"} + - {nox_session: "unit-3.13(torch_211, tf_latest)", python_version: "3.13"} + - {nox_session: "unit-3.12(torch_28, tf_latest)", python_version: "3.12"} + - {nox_session: "unit-3.12(torch_29, tf_latest)", python_version: "3.12"} + - {nox_session: "unit-3.12(torch_210, tf_latest)", python_version: "3.12"} + - {nox_session: "unit-3.12(torch_211, tf_min)", python_version: "3.12"} steps: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup with: - python-version: "3.${{ matrix.py }}" + python-version: ${{ matrix.python_version }} - name: Run unit tests - run: pip install tox && tox -e py3${{ matrix.py }}-torch211-tf_latest-unit - multi-torch: - if: github.event_name == 'pull_request' - needs: [linux] - runs-on: ubuntu-latest - timeout-minutes: 30 - strategy: - fail-fast: false - matrix: - torch: [28, 29, 210] - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/ubuntu-setup - - name: Run unit tests - run: pip install tox && tox -e py312-torch${{ matrix.torch }}-tf_latest-unit - multi-transformers: - if: github.event_name == 'pull_request' - needs: [linux] + run: pip install nox uv && nox -s "${{ matrix.nox_session }}" + partial-install: + if: needs.check-file-changes.outputs.any_changed == 'true' + needs: [linux, check-file-changes] runs-on: ubuntu-latest timeout-minutes: 30 strategy: fail-fast: false matrix: - tf: [min] + test-env: [onnx, torch] steps: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup - name: Run unit tests - run: pip install tox && tox -e py312-torch211-tf_${{ matrix.tf }}-unit + run: pip install nox uv && nox -s "partial_unit(subset='${{ matrix.test-env }}')" launcher: - if: github.event_name == 'pull_request' - needs: [linux] + if: needs.check-file-changes.outputs.any_changed == 'true' + needs: [linux, check-file-changes] runs-on: ubuntu-latest timeout-minutes: 15 steps: @@ -120,33 +141,18 @@ jobs: uv venv .venv uv pip install -e . pytest uv run python3 -m pytest -v - partial-install: - if: github.event_name == 'pull_request' - needs: [linux] - runs-on: ubuntu-latest - timeout-minutes: 30 - strategy: - fail-fast: false - matrix: - test-env: [onnx, torch] - steps: - - uses: actions/checkout@v6 - - uses: ./.github/actions/ubuntu-setup - - name: Run unit tests - run: pip install tox && tox -e py312-partial-unit-${{ matrix.test-env }} unit-pr-required-check: # Run even if some jobs are skipped if: ${{ github.event_name == 'pull_request' && always() }} - needs: [linux, windows, multi-py, multi-torch, multi-transformers, partial-install, launcher] + needs: [check-file-changes, linux, windows, multi-version, partial-install, launcher] runs-on: ubuntu-latest steps: - name: Required unit tests did not succeed if: >- - ${{ needs.linux.result != 'success' || - needs.windows.result != 'success' || - needs.multi-py.result != 'success' || - needs.multi-torch.result != 'success' || - needs.multi-transformers.result != 'success' || - needs.partial-install.result != 'success' || - needs.launcher.result != 'success' }} + ${{ needs.linux.result != 'success' || (needs.check-file-changes.outputs.any_changed == 'true' && ( + needs.windows.result != 'success' || + needs.multi-version.result != 'success' || + needs.partial-install.result != 'success' || + needs.launcher.result != 'success' + )) }} run: exit 1 diff --git a/.vscode/settings.json b/.vscode/settings.json index 0e8465ad38d..5a1c278b0f7 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -19,7 +19,6 @@ ".mypy_cache": true, ".pytest_cache": true, ".ruff_cache": true, - ".tox": true, "**/__pycache__/**": true, "**/*.pyc": true, "**/runs": true, diff --git a/CLAUDE.md b/CLAUDE.md index d07775f3a85..4af38586788 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -37,9 +37,9 @@ Primarily Python codebase with optional C++/CUDA extensions supporting PyTorch, | Pattern match | `pytest tests/unit -k "test_quantize"` | | Lint + format (all files) | `pre-commit run --all-files` | | Lint (diff only) | `pre-commit run --from-ref origin/main --to-ref HEAD` | -| Run via tox (CPU unit) | `tox -e py312-torch210-tf_latest-unit` | -| Build docs | `tox -e build-docs` | -| Build wheel | `tox -e build-wheel` | +| Run via nox (CPU unit) | `nox -s "unit-3.12(torch_211, tf_latest)"` | +| Build docs | `nox -s docs` | +| Build wheel | `nox -s build_wheel` | ## Architecture @@ -104,7 +104,7 @@ A **recipe** is a declarative YAML specification of an optimization configuratio | `modelopt_recipes/general/ptq/` | Built-in PTQ recipe YAML files | | `pyproject.toml` | Optional dependency groups (`[onnx]`, `[hf]`, `[all]`, `[dev]`); ruff, mypy, pytest, bandit, and coverage config | | `.pre-commit-config.yaml` | Pre-commit hooks (ruff, mypy, clang-format, license headers) | -| `tox.ini` | Test environment definitions | +| `noxfile.py` | Test session definitions | ## Design Patterns diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 23f07850a6b..53e879a7bf3 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -89,7 +89,7 @@ We use [pytest](https://docs.pytest.org/) for all tests. For any new features / - `tests/gpu_trtllm`: Fast GPU-based unit tests for the core ModelOpt library for TensorRT-LLM features. In most cases, they should not take more than a few seconds to run. - `tests/examples`: Integration tests for ModelOpt examples. They should not take more than a few minutes to run. Please refer to [example test README](./tests/examples/README.md) for more details. -Please refer to [tox.ini](./tox.ini) for more details on how to run the tests and their dependencies. +Please refer to [noxfile.py](./noxfile.py) for more details on how to run the tests and their dependencies. ## ✍️ Signing your work diff --git a/modelopt/torch/export/unified_export_megatron.py b/modelopt/torch/export/unified_export_megatron.py index c901a2159d8..89b718623da 100644 --- a/modelopt/torch/export/unified_export_megatron.py +++ b/modelopt/torch/export/unified_export_megatron.py @@ -76,6 +76,7 @@ from megatron.core.parallel_state import ( get_pipeline_model_parallel_rank, get_pipeline_model_parallel_world_size, + get_tensor_model_parallel_rank, ) from megatron.core.ssm.mamba_layer import MambaLayer from megatron.core.transformer.identity_op import IdentityOp @@ -258,13 +259,14 @@ def save_pretrained( """ pp_rank = get_pipeline_model_parallel_rank() pp_size = get_pipeline_model_parallel_world_size() + tp_rank = get_tensor_model_parallel_rank() # We use the 1st PP rank to handle VLM because vision_models # and vision_proj only exist in the first stage. - is_first_stage_main_rank = pp_rank == 0 + is_first_stage_main_rank = pp_rank == 0 and tp_rank == 0 # We use the last PP rank to write the config because # medusa_heads and eagle_module only exist in the last stage. - is_last_stage_main_rank = pp_rank == pp_size - 1 + is_last_stage_main_rank = pp_rank == pp_size - 1 and tp_rank == 0 # Main export process layer_state_dicts = self.layer_state_dicts diff --git a/noxfile.py b/noxfile.py new file mode 100644 index 00000000000..fcef3d30875 --- /dev/null +++ b/noxfile.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Nox session definitions for testing, linting, docs, and wheel builds. + +Usage: + python -m pip install nox uv # install nox and uv (once) + nox -l # list all sessions + nox -s gpu_megatron # run a GPU session (inside container) + nox -s "unit-3.12(torch_211, tf_latest)" # run a specific unit test combination + nox -s "unit-3.12(torch_211, tf_latest)" -R # force-recreate venv (e.g. after dep changes) + COVERAGE_PROCESS_START=pyproject.toml nox -s "unit-3.12(torch_211, tf_latest)" # with coverage +""" + +import glob +import os +import shutil + +import nox + +nox.options.default_venv_backend = "uv" if shutil.which("uv") else "virtualenv" +nox.options.envdir = "/tmp/.nox" +nox.options.reuse_existing_virtualenvs = True + +TORCH_VERSIONS = { + "torch_28": "torchvision~=0.23.0", + "torch_29": "torchvision~=0.24.0", + "torch_210": "torchvision~=0.25.0", + "torch_211": "torchvision~=0.26.0", +} + +TRANSFORMERS_VERSIONS = { + "tf_latest": None, + "tf_min": "transformers~=4.56.0", +} + + +def _cov_args(): + """Return --cov when COVERAGE_PROCESS_START is set (CI only).""" + return ["--cov"] if os.environ.get("COVERAGE_PROCESS_START") else [] + + +# ─── CPU unit tests ─────────────────────────────────────────────────────────── +@nox.session(python=["3.10", "3.11", "3.12", "3.13"]) +@nox.parametrize("tf_ver", [nox.param(k, id=k) for k in TRANSFORMERS_VERSIONS]) +@nox.parametrize("torch_ver", [nox.param(k, id=k) for k in TORCH_VERSIONS]) +def unit(session, torch_ver, tf_ver): + """Unit tests — parametrized over torch and transformers versions.""" + session.install(TORCH_VERSIONS[torch_ver], "-e", ".[all,dev-test]") + tf_pin = TRANSFORMERS_VERSIONS[tf_ver] + if tf_pin: + session.install(tf_pin) + session.run("python", "-m", "pytest", "tests/unit", *_cov_args()) + + +@nox.session(python="3.12") +@nox.parametrize("subset", ["onnx", "torch", "torch_deploy"]) +def partial_unit(session, subset): + """Unit tests with partial installs.""" + if subset == "onnx": + session.install("torchvision~=0.26.0", ".[onnx,dev-test]") + session.run("python", "-m", "pytest", "tests/unit/onnx") + elif subset == "torch": + session.install("megatron-core", ".[dev-test]") + session.run( + "python", + "-m", + "pytest", + "tests/unit/torch", + "--ignore=tests/unit/torch/deploy", + "--ignore=tests/unit/torch/puzzletron", + ) + else: # torch_deploy + session.install(".[onnx,dev-test]") + session.run("python", "-m", "pytest", "tests/unit/torch/deploy") + + +# ─── GPU sessions (run inside containers — no new venv) ────────────────────── +# `venv_backend="none"` skips creating a new venv so the session runs directly in the container's +# existing Python environment (e.g. /opt/venv in NeMo) instead of an isolated one. +# Use `python -m pip/pytest` to ensure the container's active venv Python is used, +# not a stale PATH entry (e.g. NeMo container has pip → /usr/local/bin/pip but python → /opt/venv/bin/python). +# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later +@nox.session(venv_backend="none") +def gpu(session): + # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers + # until https://github.com/tbenthompson/cppimport/pull/98 + session.run( + "python", + "-m", + "pip", + "install", + "--no-build-isolation", + "git+https://github.com/Dao-AILab/fast-hadamard-transform.git", + ) + session.run("python", "-m", "pip", "install", "-e", ".[all,dev-test]") + session.run("python", "-m", "pip", "uninstall", "-y", "cupy-cuda12x") + session.run("python", "-m", "pip", "install", "cupy-cuda13x") + session.run( + "python", + "-m", + "pip", + "install", + "--no-build-isolation", + "git+https://github.com/state-spaces/mamba.git", + "git+https://github.com/Dao-AILab/causal-conv1d.git", + ) + session.run("python", "-m", "pytest", "tests/gpu", *_cov_args()) + + +# Container: nvcr.io/nvidia/nemo:26.04 or later +@nox.session(venv_backend="none") +def gpu_megatron(session): + # nemo:26.04 has transformers 5.x but system-wide installed trtllm 1.2.0 which does not support it causing import errors + session.run("pip", "uninstall", "-y", "tensorrt_llm") + session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") + session.run("python", "-m", "pytest", "tests/gpu_megatron", *_cov_args()) + + +# Container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10 or later +@nox.session(venv_backend="none") +def gpu_trtllm(session): + session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") + session.run("python", "-m", "pytest", "tests/gpu_trtllm", *_cov_args()) + + +# Container: nvcr.io/nvidia/pytorch:26.01-py3 or later +@nox.session(venv_backend="none") +def regression(session): + session.run("python", "-m", "pip", "install", "-e", ".[hf,dev-test]") + session.run("python", "-m", "pytest", "tests/regression", *_cov_args()) + + +# ─── Code quality ───────────────────────────────────────────────────────────── +@nox.session +def pre_commit_all(session): + session.install("-e", ".[all,dev-lint]") + session.run("pre-commit", "run", "--all-files", "--show-diff-on-failure") + + +@nox.session +def pre_commit_diff(session): + session.install("-e", ".[all,dev-lint]") + session.run("pre-commit", "run", "--from-ref", "origin/main", "--to-ref", "HEAD") + + +# ─── Docs ───────────────────────────────────────────────────────────────────── +@nox.session +def docs(session): + session.install("-e", ".[all,dev-docs]") + shutil.rmtree("docs/build", ignore_errors=True) + shutil.rmtree("docs/source/reference/generated", ignore_errors=True) + with session.chdir("docs"): + session.run( + "sphinx-build", + "source", + "build/html", + "--fail-on-warning", + "--show-traceback", + "--keep-going", + ) + + +@nox.session +def docs_debug(session): + session.install("-e", ".[all,dev-docs]") + shutil.rmtree("docs/build", ignore_errors=True) + shutil.rmtree("docs/source/reference/generated", ignore_errors=True) + with session.chdir("docs"): + session.run("sphinx-autobuild", "source", "build/html", "--host", "0.0.0.0") + + +# ─── Wheel build ────────────────────────────────────────────────────────────── +@nox.session +def build_wheel(session): + shutil.rmtree("build", ignore_errors=True) + session.install("twine") + session.run("pip", "wheel", "--no-deps", "--wheel-dir=dist", ".") + wheels = glob.glob("dist/*.whl") + session.run("twine", "check", *wheels) + (modelopt_wheel,) = glob.glob("dist/nvidia_modelopt-*.whl") + session.install(modelopt_wheel, "-f", "dist") + with session.chdir("dist"): + session.run("python", "-c", "import modelopt; print(modelopt.__version__)") diff --git a/pyproject.toml b/pyproject.toml index 25cb6338aa5..fdd60b5193a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -82,7 +82,7 @@ hf = [ "peft>=0.17.0", "sentencepiece>=0.2.1", # Also implicitly used in test_unified_export_megatron, test_vllm_fakequant_megatron_export "tiktoken", - "transformers>=4.56", # Should match modelopt/torch/__init__.py and tox.ini + "transformers>=4.56", # Should match modelopt/torch/__init__.py and noxfile.py "wonderwords", ] @@ -111,17 +111,18 @@ dev-docs = [ "sphinx-togglebutton>=0.3.2", ] dev-test = [ - "coverage[toml]>=7.13.0", # a1_coverage.pth for subprocess tracking requires this + "coverage[toml]>=7.13.0", # a1_coverage.pth for subprocess tracking requires this + "nox", "pytest", "pytest-cov", "pytest-instafail", "pytest-timeout", + "uv", + # test-specific dependencies "timm", - "torchprofile==0.0.4", # optional dependency for modelopt.torch + "torchprofile==0.0.4", # optional dependency for modelopt.torch "torchvision", "torch-geometric", - "tox>4.18", - "tox-current-env>=0.0.12", ] # Compound extras via self-references all = ["nvidia-modelopt[onnx,hf,puzzletron]"] @@ -208,6 +209,7 @@ extend-ignore = [ [tool.ruff.lint.per-file-ignores] "__init__.py" = ["F401", "F403"] "examples/*" = ["D"] +"noxfile.py" = ["D", "E501"] "tests/*" = ["B017", "D", "E402", "PT012"] "*/_[a-zA-Z]*" = ["D"] # Private packages (_abc/*.py) or modules (_xyz.py) "*.ipynb" = [ @@ -279,7 +281,7 @@ module = ["examples.*"] disable_error_code = ["attr-defined"] [tool.bandit] -exclude_dirs = [".github/", "examples/", "tests/"] +exclude_dirs = [".github/", "examples/", "noxfile.py", "tests/"] # Do not change `skips`. It should be consistent with NVIDIA's Wheel-CI-CD bandit.yml config. # Use of `# nosec BXXX` requires special approval skips = [ diff --git a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py index 925c0e3c51c..8ccb3d42906 100644 --- a/tests/gpu_megatron/torch/export/test_unified_export_megatron.py +++ b/tests/gpu_megatron/torch/export/test_unified_export_megatron.py @@ -24,6 +24,7 @@ from _test_utils.torch.megatron.models import get_mcore_gpt_model from _test_utils.torch.megatron.utils import get_forward from _test_utils.torch.transformers_models import create_tiny_llama_dir +from safetensors import safe_open from safetensors.torch import save_file import modelopt.torch.quantization as mtq @@ -275,8 +276,6 @@ def _test_qkv_slicing_gqa_tp2(tmp_path, rank, size): # Verify Q/K/V projections were exported (collect keys from all shard files) if rank == 0: - from safetensors import safe_open - safetensors_files = list(export_dir.glob("*.safetensors")) assert safetensors_files, "no safetensors files found in export dir" keys = [] diff --git a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py index 4f76b3f3ae8..50e3f43ab7d 100644 --- a/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py +++ b/tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py @@ -20,8 +20,11 @@ import torch import torch.nn.init as init from _test_utils.torch.megatron.models import get_mcore_gpt_model -from _test_utils.torch.megatron.utils import initialize_for_megatron -from megatron.core import dist_checkpointing +from _test_utils.torch.megatron.utils import ( + initialize_for_megatron, + load_distributed_checkpoint, + save_distributed_checkpoint, +) import modelopt.torch.peft as mtpeft import modelopt.torch.quantization as mtq @@ -148,20 +151,6 @@ } -def save_distributed_checkpoint(checkpoint_path, gpt_model): - sharded_state_dict = gpt_model.sharded_state_dict(prefix="") - dist_checkpointing.save(sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path) - - -def load_distributed_checkpoint(checkpoint_path, gpt_model): - sharded_state_dict = gpt_model.sharded_state_dict(prefix="") - checkpoint = dist_checkpointing.load( - sharded_state_dict=sharded_state_dict, checkpoint_dir=checkpoint_path - ) - gpt_model.load_state_dict(checkpoint) - return gpt_model - - def _gpt_model_provider(tp_size: int, hidden_size=256, vocab_size=64, meta_device=False): """Build the model.""" diff --git a/tests/gpu_regression/torch/speculative/test_dflash.py b/tests/regression/torch/speculative/test_dflash.py similarity index 100% rename from tests/gpu_regression/torch/speculative/test_dflash.py rename to tests/regression/torch/speculative/test_dflash.py diff --git a/tests/unit/torch/speculative/conftest.py b/tests/unit/torch/speculative/conftest.py new file mode 100644 index 00000000000..083769483a3 --- /dev/null +++ b/tests/unit/torch/speculative/conftest.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import platform + + +# `Compiler: cl is not found` on Windows +def pytest_ignore_collect(collection_path, config): + return platform.system() == "Windows" diff --git a/tools/launcher/tests/conftest.py b/tools/launcher/tests/conftest.py index da98e8ad6cb..072518cc795 100644 --- a/tools/launcher/tests/conftest.py +++ b/tools/launcher/tests/conftest.py @@ -20,8 +20,8 @@ uv pip install pytest uv run python3 -m pytest tests/ -v -Or via tox from Model-Optimizer root: - tox -e py312-launcher +Or via nox from Model-Optimizer root: + nox -s "unit-3.12(torch_211, tf_latest)" """ import os diff --git a/tox.ini b/tox.ini deleted file mode 100644 index 7bfa1e41e57..00000000000 --- a/tox.ini +++ /dev/null @@ -1,149 +0,0 @@ -[tox] -envlist= - pre-commit-all - py312-torch210-tf_latest-unit - cuda13-gpu - cuda13-gpu-regression - cuda13-gpu-megatron -skipsdist = True -toxworkdir = /tmp/{env:USER}-modelopt-tox -passenv = - SETUPTOOLS_SCM_PRETEND_VERSION - -############################ -# CPU Unit test environments -############################ -[testenv:{py310,py311,py312,py313}-torch{28,29,210,211}-tf_{min,latest}-unit] -deps = - # torch version auto-selected based on torchvision version - torch28: torchvision~=0.23.0 - torch29: torchvision~=0.24.0 - torch210: torchvision~=0.25.0 - torch211: torchvision~=0.26.0 - - -e .[all,dev-test] - - # Should match pyproject.toml - tf_min: transformers~=4.56.0 -commands = - python -m pytest tests/unit {env:COV_ARGS:} - - -##################################################################### -# Environment to run unit tests with subset of dependencies installed -##################################################################### -[testenv:{py310,py311,py312,py313}-partial-unit-{onnx,torch,torch_deploy}] -allowlist_externals = - bash, rm -deps = - # Make sure torch 2.10 is used - torchvision~=0.26.0 - - # ONNX unit tests heavily rely on torch / torchvision - onnx: .[onnx,dev-test] - onnx: torchvision - - # Install megatron-core to test torch-only install can still import plugins - torch: megatron-core - # diffusers is needed for unit tests of the sparse-attention/quantization diffusers backend - torch: diffusers - torch: .[dev-test] - - torch_deploy: .[onnx,torch,dev-test] -commands = - onnx: python -m pytest tests/unit/onnx - torch: python -m pytest tests/unit/torch --ignore tests/unit/torch/deploy --ignore tests/unit/torch/puzzletron - torch_deploy: python -m pytest tests/unit/torch/deploy - - -########################################################### -# GPU test environments (Should be used with --current-env) -########################################################### -[testenv:cuda13-gpu] -commands_pre = - # Install deps here so that it gets installed even in --current-env - pip install --no-build-isolation git+https://github.com/Dao-AILab/fast-hadamard-transform.git - pip install -e .[all,dev-test] - - # Install cupy-cuda13x for INT4 ONNX quantization (default is cupy-cuda12x) - pip uninstall -y cupy-cuda12x - pip install cupy-cuda13x - - # Install mamba and causal-conv1d for Nemotron tests - pip install --no-build-isolation git+https://github.com/state-spaces/mamba.git - pip install --no-build-isolation git+https://github.com/Dao-AILab/causal-conv1d.git -commands = - python -m pytest tests/gpu {env:COV_ARGS:} - -[testenv:cuda13-gpu-regression] -commands_pre = - pip install -e .[hf,dev-test] -commands = - python -m pytest tests/gpu_regression {env:COV_ARGS:} - -[testenv:cuda13-gpu-megatron] -commands_pre = - # Install deps here so that it gets installed even in --current-env - # Temporarily disable latest mcore until we fix its nvidia-resiliency-ext dependency - pip install 'megatron-core<0.17.0' - pip install --no-build-isolation git+https://github.com/state-spaces/mamba.git - pip install --no-build-isolation git+https://github.com/Dao-AILab/causal-conv1d.git - pip install -e .[hf,dev-test] -commands = - python -m pytest tests/gpu_megatron {env:COV_ARGS:} - -[testenv:cuda13-gpu-trtllm] -# Expected to be run in TRT-LLM container -commands_pre = - # Install deps here so that it gets installed even in --current-env - pip install -e .[hf,dev-test] -commands = - python -m pytest tests/gpu_trtllm {env:COV_ARGS:} - -############################################# -# Code quality checks on all files or on diff -############################################# -[testenv:{pre-commit}-{all,diff}] -deps = - -e .[all,dev-lint] -commands = - all: pre-commit run --all-files --show-diff-on-failure {posargs} - diff: pre-commit run --from-ref origin/main --to-ref HEAD {posargs} - - -######################### -# Run documentation build -######################### -[testenv:{build,debug}-docs] -allowlist_externals = - rm -deps = - -e .[all,dev-docs] -changedir = docs -commands_pre = - rm -rf build - rm -rf source/reference/generated -commands = - sphinx-build source build/html --fail-on-warning --show-traceback --keep-going - debug: sphinx-autobuild source build/html --host 0.0.0.0 - - -################# -# Run wheel build -################# -[testenv:build-wheel] -allowlist_externals = - bash, cd, rm -deps = - twine -commands = - # Clean build directory to avoid any stale files getting into the wheel - rm -rf build - - # Build and check wheel - pip wheel --no-deps --wheel-dir=dist . - twine check dist/* - - # Install and test the wheel - bash -c "find dist -name 'nvidia_modelopt-*.whl' | xargs pip install -f dist" - bash -c "cd dist; python -c 'import modelopt; print(modelopt.__version__);'" diff --git a/uv.lock b/uv.lock index cfa742f1081..e223ce17eae 100644 --- a/uv.lock +++ b/uv.lock @@ -32,6 +32,9 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", ] +[manifest] +overrides = [{ name = "torch", marker = "sys_platform == 'never'" }] + [[package]] name = "accelerate" version = "1.13.0" @@ -44,7 +47,7 @@ dependencies = [ { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch" }, + { name = "torch", marker = "sys_platform == 'never'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/ca/14/787e5498cd062640f0f3d92ef4ae4063174f76f9afd29d13fc52a319daae/accelerate-1.13.0.tar.gz", hash = "sha256:d631b4e0f5b3de4aff2d7e9e6857d164810dfc3237d54d017f075122d057b236", size = 402835, upload-time = "2026-03-04T19:34:12.359Z" } wheels = [ @@ -206,6 +209,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, ] +[[package]] +name = "argcomplete" +version = "3.6.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/38/61/0b9ae6399dd4a58d8c1b1dc5a27d6f2808023d0b5dd3104bb99f45a33ff6/argcomplete-3.6.3.tar.gz", hash = "sha256:62e8ed4fd6a45864acc8235409461b72c9a28ee785a2011cc5eb78318786c89c", size = 73754, upload-time = "2025-10-20T03:33:34.741Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/74/f5/9373290775639cb67a2fce7f629a1c240dce9f12fe927bc32b2736e16dfc/argcomplete-3.6.3-py3-none-any.whl", hash = "sha256:f5007b3a600ccac5d25bbce33089211dfd49eab4a7718da3f10e3082525a92ce", size = 43846, upload-time = "2025-10-20T03:33:33.021Z" }, +] + [[package]] name = "async-timeout" version = "5.0.1" @@ -266,15 +278,6 @@ toml = [ { name = "tomli", marker = "python_full_version < '3.11'" }, ] -[[package]] -name = "cachetools" -version = "7.0.5" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/af/dd/57fe3fdb6e65b25a5987fd2cdc7e22db0aef508b91634d2e57d22928d41b/cachetools-7.0.5.tar.gz", hash = "sha256:0cd042c24377200c1dcd225f8b7b12b0ca53cc2c961b43757e774ebe190fd990", size = 37367, upload-time = "2026-03-09T20:51:29.451Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/06/f3/39cf3367b8107baa44f861dc802cbf16263c945b62d8265d36034fc07bea/cachetools-7.0.5-py3-none-any.whl", hash = "sha256:46bc8ebefbe485407621d0a4264b23c080cedd913921bad7ac3ed2f26c183114", size = 13918, upload-time = "2026-03-09T20:51:27.33Z" }, -] - [[package]] name = "cbcbox" version = "2.929" @@ -471,6 +474,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a7/06/3d6badcf13db419e25b07041d9c7b4a2c331d3f4e7134445ec5df57714cd/coloredlogs-15.0.1-py2.py3-none-any.whl", hash = "sha256:612ee75c546f53e92e70049c9dbfcc18c935a2b9a53b66085ce9ef6a6e5c0934", size = 46018, upload-time = "2021-06-11T10:22:42.561Z" }, ] +[[package]] +name = "colorlog" +version = "6.10.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a2/61/f083b5ac52e505dfc1c624eafbf8c7589a0d7f32daa398d2e7590efa5fda/colorlog-6.10.1.tar.gz", hash = "sha256:eb4ae5cb65fe7fec7773c2306061a8e63e02efc2c72eba9d27b0fa23c94f1321", size = 17162, upload-time = "2025-10-16T16:14:11.978Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6d/c1/e419ef3723a074172b68aaa89c9f3de486ed4c2399e2dbd8113a4fdcaf9e/colorlog-6.10.1-py3-none-any.whl", hash = "sha256:2d7e8348291948af66122cff006c9f8da6255d224e7cf8e37d8de2df3bad8c9c", size = 11743, upload-time = "2025-10-16T16:14:10.512Z" }, +] + [[package]] name = "coverage" version = "7.13.5" @@ -570,21 +585,6 @@ dependencies = [ ] sdist = { url = "https://files.pythonhosted.org/packages/54/27/01d9078a77b9e31b79b9716e66ca4db74f4744c5232bcb3e8769395c4280/cppimport-22.8.2.tar.gz", hash = "sha256:bbb4957102db41bc99ad72c233bce92f9d1fd91be352fc07878c4361033a401f", size = 26635, upload-time = "2022-08-02T16:50:36.872Z" } -[[package]] -name = "cuda-bindings" -version = "12.9.4" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cuda-pathfinder", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/7a/d8/b546104b8da3f562c1ff8ab36d130c8fe1dd6a045ced80b4f6ad74f7d4e1/cuda_bindings-12.9.4-cp310-cp310-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d3c842c2a4303b2a580fe955018e31aea30278be19795ae05226235268032e5", size = 12148218, upload-time = "2025-10-21T14:51:28.855Z" }, - { url = "https://files.pythonhosted.org/packages/45/e7/b47792cc2d01c7e1d37c32402182524774dadd2d26339bd224e0e913832e/cuda_bindings-12.9.4-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c912a3d9e6b6651853eed8eed96d6800d69c08e94052c292fec3f282c5a817c9", size = 12210593, upload-time = "2025-10-21T14:51:36.574Z" }, - { url = "https://files.pythonhosted.org/packages/a9/c1/dabe88f52c3e3760d861401bb994df08f672ec893b8f7592dc91626adcf3/cuda_bindings-12.9.4-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:fda147a344e8eaeca0c6ff113d2851ffca8f7dfc0a6c932374ee5c47caa649c8", size = 12151019, upload-time = "2025-10-21T14:51:43.167Z" }, - { url = "https://files.pythonhosted.org/packages/63/56/e465c31dc9111be3441a9ba7df1941fe98f4aa6e71e8788a3fb4534ce24d/cuda_bindings-12.9.4-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:32bdc5a76906be4c61eb98f546a6786c5773a881f3b166486449b5d141e4a39f", size = 11906628, upload-time = "2025-10-21T14:51:49.905Z" }, - { url = "https://files.pythonhosted.org/packages/a3/84/1e6be415e37478070aeeee5884c2022713c1ecc735e6d82d744de0252eee/cuda_bindings-12.9.4-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:56e0043c457a99ac473ddc926fe0dc4046694d99caef633e92601ab52cbe17eb", size = 11925991, upload-time = "2025-10-21T14:51:56.535Z" }, -] - [[package]] name = "cuda-pathfinder" version = "1.5.3" @@ -659,11 +659,24 @@ dependencies = [ { name = "psutil", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, { name = "py-cpuinfo", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, { name = "pydantic", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "torch", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "torch", marker = "sys_platform == 'never'" }, { name = "tqdm", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/24/61/5ea1c63b139fe7530b196b68ce0bdffa9cde79882e527dcecae58bd6c770/deepspeed-0.18.9.tar.gz", hash = "sha256:ee4818dcf342794f74f429a0aeebef90291ec808fa82609c5140c23e665c4011", size = 1663466, upload-time = "2026-03-30T16:43:16.566Z" } +[[package]] +name = "dependency-groups" +version = "1.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/62/55/f054de99871e7beb81935dea8a10b90cd5ce42122b1c3081d5282fdb3621/dependency_groups-1.3.1.tar.gz", hash = "sha256:78078301090517fd938c19f64a53ce98c32834dfe0dee6b88004a569a6adfefd", size = 10093, upload-time = "2025-05-02T00:34:29.452Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/c7/d1ec24fb280caa5a79b6b950db565dab30210a66259d17d5bb2b3a9f878d/dependency_groups-1.3.1-py3-none-any.whl", hash = "sha256:51aeaa0dfad72430fcfb7bcdbefbd75f3792e5919563077f30bc0d73f4493030", size = 8664, upload-time = "2025-05-02T00:34:27.085Z" }, +] + [[package]] name = "diffusers" version = "0.37.1" @@ -967,6 +980,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f0/0f/310fb31e39e2d734ccaa2c0fb981ee41f7bd5056ce9bc29b2248bd569169/humanfriendly-10.0-py2.py3-none-any.whl", hash = "sha256:1697e1a8a8f550fd43c2865cd84542fc175a61dcb779b6fee18cf6b6ccba1477", size = 86794, upload-time = "2021-09-17T21:40:39.897Z" }, ] +[[package]] +name = "humanize" +version = "4.15.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/66/a3921783d54be8a6870ac4ccffcd15c4dc0dd7fcce51c6d63b8c63935276/humanize-4.15.0.tar.gz", hash = "sha256:1dd098483eb1c7ee8e32eb2e99ad1910baefa4b75c3aff3a82f4d78688993b10", size = 83599, upload-time = "2025-12-20T20:16:13.19Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c5/7b/bca5613a0c3b542420cf92bd5e5fb8ebd5435ce1011a091f66bb7693285e/humanize-4.15.0-py3-none-any.whl", hash = "sha256:b1186eb9f5a9749cd9cb8565aee77919dd7c8d076161cf44d70e59e3301e1769", size = 132203, upload-time = "2025-12-20T20:16:11.67Z" }, +] + [[package]] name = "hydra-core" version = "1.3.2" @@ -1532,11 +1554,7 @@ name = "networkx" version = "3.4.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform == 'win32'", "(python_full_version < '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version < '3.11' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')", - "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform == 'darwin'", - "python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform == 'win32'", "python_full_version < '3.11' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version < '3.11' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", ] @@ -1550,27 +1568,15 @@ name = "networkx" version = "3.6.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform == 'win32'", "(python_full_version >= '3.13' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.13' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')", - "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform == 'darwin'", "(python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version == '3.12.*' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')", - "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'darwin'", "(python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version == '3.11.*' and platform_machine == 'aarch64' and sys_platform != 'darwin' and sys_platform != 'win32')", - "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'darwin'", "python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'", - "python_full_version >= '3.13' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version >= '3.13' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32'", ] sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } wheels = [ @@ -1627,6 +1633,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/b2/d0896bdcdc8d28a7fc5717c305f1a861c26e18c05047949fb371034d98bd/nodeenv-1.10.0-py2.py3-none-any.whl", hash = "sha256:5bb13e3eed2923615535339b3c620e76779af4cb4c6a90deccc9e36b274d3827", size = 23438, upload-time = "2025-12-20T14:08:52.782Z" }, ] +[[package]] +name = "nox" +version = "2026.4.10" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "argcomplete" }, + { name = "attrs" }, + { name = "colorlog" }, + { name = "dependency-groups" }, + { name = "humanize" }, + { name = "packaging" }, + { name = "tomli", marker = "python_full_version < '3.11'" }, + { name = "virtualenv" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e7/6b/e672c862a43cfca704d32359221fa3780226daa1e5db5dfc401bcc8be9c9/nox-2026.4.10.tar.gz", hash = "sha256:2d0af5374f3f37a295428c927d1b04a8182aa01762897d172446dda2f1ce9692", size = 4034839, upload-time = "2026-04-10T17:42:42.209Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/95/4df134a100b5a9a12378d5301b934366686ef6fbdaffcd21211d5654970e/nox-2026.4.10-py3-none-any.whl", hash = "sha256:082c117627590d9b90aa21f86df89b310b07c5842539524203bcb3c719f116c1", size = 75536, upload-time = "2026-04-10T17:42:40.664Z" }, +] + [[package]] name = "numpy" version = "2.2.6" @@ -1779,108 +1804,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/74/f4c001f4714c3ad9ce037e18cf2b9c64871a84951eaa0baf683a9ca9301c/numpy-2.4.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:f2cf083b324a467e1ab358c105f6cad5ea950f50524668a80c486ff1db24e119", size = 12509075, upload-time = "2026-03-29T13:21:57.644Z" }, ] -[[package]] -name = "nvidia-cublas-cu12" -version = "12.8.4.1" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/61/e24b560ab2e2eaeb3c839129175fb330dfcfc29e5203196e5541a4c44682/nvidia_cublas_cu12-12.8.4.1-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:8ac4e771d5a348c551b2a426eda6193c19aa630236b418086020df5ba9667142", size = 594346921, upload-time = "2025-03-07T01:44:31.254Z" }, -] - -[[package]] -name = "nvidia-cuda-cupti-cu12" -version = "12.8.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/02/2adcaa145158bf1a8295d83591d22e4103dbfd821bcaf6f3f53151ca4ffa/nvidia_cuda_cupti_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ea0cb07ebda26bb9b29ba82cda34849e73c166c18162d3913575b0c9db9a6182", size = 10248621, upload-time = "2025-03-07T01:40:21.213Z" }, -] - -[[package]] -name = "nvidia-cuda-nvrtc-cu12" -version = "12.8.93" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/05/6b/32f747947df2da6994e999492ab306a903659555dddc0fbdeb9d71f75e52/nvidia_cuda_nvrtc_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:a7756528852ef889772a84c6cd89d41dfa74667e24cca16bb31f8f061e3e9994", size = 88040029, upload-time = "2025-03-07T01:42:13.562Z" }, -] - -[[package]] -name = "nvidia-cuda-runtime-cu12" -version = "12.8.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0d/9b/a997b638fcd068ad6e4d53b8551a7d30fe8b404d6f1804abf1df69838932/nvidia_cuda_runtime_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:adade8dcbd0edf427b7204d480d6066d33902cab2a4707dcfc48a2d0fd44ab90", size = 954765, upload-time = "2025-03-07T01:40:01.615Z" }, -] - -[[package]] -name = "nvidia-cudnn-cu12" -version = "9.10.2.21" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/ba/51/e123d997aa098c61d029f76663dedbfb9bc8dcf8c60cbd6adbe42f76d049/nvidia_cudnn_cu12-9.10.2.21-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:949452be657fa16687d0930933f032835951ef0892b37d2d53824d1a84dc97a8", size = 706758467, upload-time = "2025-06-06T21:54:08.597Z" }, -] - -[[package]] -name = "nvidia-cufft-cu12" -version = "11.3.3.83" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/13/ee4e00f30e676b66ae65b4f08cb5bcbb8392c03f54f2d5413ea99a5d1c80/nvidia_cufft_cu12-11.3.3.83-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:4d2dd21ec0b88cf61b62e6b43564355e5222e4a3fb394cac0db101f2dd0d4f74", size = 193118695, upload-time = "2025-03-07T01:45:27.821Z" }, -] - -[[package]] -name = "nvidia-cufile-cu12" -version = "1.13.1.3" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/bb/fe/1bcba1dfbfb8d01be8d93f07bfc502c93fa23afa6fd5ab3fc7c1df71038a/nvidia_cufile_cu12-1.13.1.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1d069003be650e131b21c932ec3d8969c1715379251f8d23a1860554b1cb24fc", size = 1197834, upload-time = "2025-03-07T01:45:50.723Z" }, -] - -[[package]] -name = "nvidia-curand-cu12" -version = "10.3.9.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/aa/6584b56dc84ebe9cf93226a5cde4d99080c8e90ab40f0c27bda7a0f29aa1/nvidia_curand_cu12-10.3.9.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:b32331d4f4df5d6eefa0554c565b626c7216f87a06a4f56fab27c3b68a830ec9", size = 63619976, upload-time = "2025-03-07T01:46:23.323Z" }, -] - -[[package]] -name = "nvidia-cusolver-cu12" -version = "11.7.3.90" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-cublas-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, - { name = "nvidia-cusparse-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/85/48/9a13d2975803e8cf2777d5ed57b87a0b6ca2cc795f9a4f59796a910bfb80/nvidia_cusolver_cu12-11.7.3.90-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:4376c11ad263152bd50ea295c05370360776f8c3427b30991df774f9fb26c450", size = 267506905, upload-time = "2025-03-07T01:47:16.273Z" }, -] - -[[package]] -name = "nvidia-cusparse-cu12" -version = "12.5.8.93" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-nvjitlink-cu12", marker = "platform_machine != 'aarch64' and platform_machine != 's390x' and sys_platform != 'darwin' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/c2/f5/e1854cb2f2bcd4280c44736c93550cc300ff4b8c95ebe370d0aa7d2b473d/nvidia_cusparse_cu12-12.5.8.93-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1ec05d76bbbd8b61b06a80e1eaf8cf4959c3d4ce8e711b65ebd0443bb0ebb13b", size = 288216466, upload-time = "2025-03-07T01:48:13.779Z" }, -] - -[[package]] -name = "nvidia-cusparselt-cu12" -version = "0.7.1" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/56/79/12978b96bd44274fe38b5dde5cfb660b1d114f70a65ef962bcbbed99b549/nvidia_cusparselt_cu12-0.7.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:f1bb701d6b930d5a7cea44c19ceb973311500847f81b634d802b7b539dc55623", size = 287193691, upload-time = "2025-02-26T00:15:44.104Z" }, -] - [[package]] name = "nvidia-ml-py" version = "13.595.45" @@ -1909,7 +1832,7 @@ dependencies = [ { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "scipy", version = "1.17.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "setuptools" }, - { name = "torch" }, + { name = "torch", marker = "sys_platform == 'never'" }, { name = "tqdm" }, ] @@ -1969,6 +1892,7 @@ dev = [ { name = "ml-dtypes" }, { name = "mypy" }, { name = "nltk" }, + { name = "nox" }, { name = "onnx" }, { name = "onnx-graphsurgeon" }, { name = "onnxconverter-common" }, @@ -2002,10 +1926,9 @@ dev = [ { name = "torch-geometric" }, { name = "torchprofile" }, { name = "torchvision" }, - { name = "tox" }, - { name = "tox-current-env" }, { name = "transformers" }, { name = "typeguard" }, + { name = "uv" }, { name = "wonderwords" }, ] dev-docs = [ @@ -2027,6 +1950,7 @@ dev-lint = [ ] dev-test = [ { name = "coverage", extra = ["toml"] }, + { name = "nox" }, { name = "pytest" }, { name = "pytest-cov" }, { name = "pytest-instafail" }, @@ -2035,8 +1959,7 @@ dev-test = [ { name = "torch-geometric" }, { name = "torchprofile" }, { name = "torchvision" }, - { name = "tox" }, - { name = "tox-current-env" }, + { name = "uv" }, ] hf = [ { name = "accelerate" }, @@ -2100,6 +2023,7 @@ requires-dist = [ { name = "mypy", marker = "extra == 'dev-lint'", specifier = "==1.17.1" }, { name = "ninja" }, { name = "nltk", marker = "extra == 'hf'" }, + { name = "nox", marker = "extra == 'dev-test'" }, { name = "numpy" }, { name = "nvidia-ml-py", specifier = ">=12" }, { name = "nvidia-modelopt", extras = ["all", "dev-docs", "dev-lint", "dev-test"], marker = "extra == 'dev'" }, @@ -2120,7 +2044,7 @@ requires-dist = [ { name = "peft", marker = "extra == 'hf'", specifier = ">=0.17.0" }, { name = "polygraphy", marker = "extra == 'onnx'", specifier = ">=0.49.22" }, { name = "pre-commit", marker = "extra == 'dev-lint'", specifier = "==4.3.0" }, - { name = "pulp" }, + { name = "pulp", specifier = "<4.0" }, { name = "pydantic", specifier = ">=2.0" }, { name = "pytest", marker = "extra == 'dev-test'" }, { name = "pytest-cov", marker = "extra == 'dev-test'" }, @@ -2147,47 +2071,14 @@ requires-dist = [ { name = "torch-geometric", marker = "extra == 'dev-test'" }, { name = "torchprofile", marker = "extra == 'dev-test'", specifier = "==0.0.4" }, { name = "torchvision", marker = "extra == 'dev-test'" }, - { name = "tox", marker = "extra == 'dev-test'", specifier = ">4.18" }, - { name = "tox-current-env", marker = "extra == 'dev-test'", specifier = ">=0.0.12" }, { name = "tqdm" }, { name = "transformers", marker = "extra == 'hf'", specifier = ">=4.56" }, { name = "typeguard", marker = "extra == 'puzzletron'" }, + { name = "uv", marker = "extra == 'dev-test'" }, { name = "wonderwords", marker = "extra == 'hf'" }, ] provides-extras = ["onnx", "hf", "puzzletron", "dev-lint", "dev-docs", "dev-test", "all", "dev"] -[[package]] -name = "nvidia-nccl-cu12" -version = "2.27.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6e/89/f7a07dc961b60645dbbf42e80f2bc85ade7feb9a491b11a1e973aa00071f/nvidia_nccl_cu12-2.27.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ad730cf15cb5d25fe849c6e6ca9eb5b76db16a80f13f425ac68d8e2e55624457", size = 322348229, upload-time = "2025-06-26T04:11:28.385Z" }, -] - -[[package]] -name = "nvidia-nvjitlink-cu12" -version = "12.8.93" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/f6/74/86a07f1d0f42998ca31312f998bd3b9a7eff7f52378f4f270c8679c77fb9/nvidia_nvjitlink_cu12-12.8.93-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:81ff63371a7ebd6e6451970684f916be2eab07321b73c9d244dc2b4da7f73b88", size = 39254836, upload-time = "2025-03-07T01:49:55.661Z" }, -] - -[[package]] -name = "nvidia-nvshmem-cu12" -version = "3.4.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b5/09/6ea3ea725f82e1e76684f0708bbedd871fc96da89945adeba65c3835a64c/nvidia_nvshmem_cu12-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:042f2500f24c021db8a06c5eec2539027d57460e1c1a762055a6554f72c369bd", size = 139103095, upload-time = "2025-09-06T00:32:31.266Z" }, -] - -[[package]] -name = "nvidia-nvtx-cu12" -version = "12.8.90" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/eb/86626c1bbc2edb86323022371c39aa48df6fd8b0a1647bc274577f72e90b/nvidia_nvtx_cu12-12.8.90-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5b17e2001cc0d751a5bc2c6ec6d26ad95913324a4adb86788c944f8ce9ba441f", size = 89954, upload-time = "2025-03-07T01:42:44.131Z" }, -] - [[package]] name = "omegaconf" version = "2.3.0" @@ -2604,7 +2495,7 @@ dependencies = [ { name = "psutil" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch" }, + { name = "torch", marker = "sys_platform == 'never'" }, { name = "tqdm" }, { name = "transformers" }, ] @@ -3061,19 +2952,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, ] -[[package]] -name = "pyproject-api" -version = "1.10.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "packaging" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/45/7b/c0e1333b61d41c69e59e5366e727b18c4992688caf0de1be10b3e5265f6b/pyproject_api-1.10.0.tar.gz", hash = "sha256:40c6f2d82eebdc4afee61c773ed208c04c19db4c4a60d97f8d7be3ebc0bbb330", size = 22785, upload-time = "2025-10-09T19:12:27.21Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/54/cc/cecf97be298bee2b2a37dd360618c819a2a7fd95251d8e480c1f0eb88f3b/pyproject_api-1.10.0-py3-none-any.whl", hash = "sha256:8757c41a79c0f4ab71b99abed52b97ecf66bd20b04fa59da43b5840bac105a09", size = 13218, upload-time = "2025-10-09T19:12:24.428Z" }, -] - [[package]] name = "pyreadline3" version = "3.5.4" @@ -3944,7 +3822,7 @@ dependencies = [ { name = "huggingface-hub" }, { name = "pyyaml" }, { name = "safetensors" }, - { name = "torch" }, + { name = "torch", marker = "sys_platform == 'never'" }, { name = "torchvision" }, ] sdist = { url = "https://files.pythonhosted.org/packages/7b/1e/e924b3b2326a856aaf68586f9c52a5fc81ef45715eca408393b68c597e0e/timm-1.0.26.tar.gz", hash = "sha256:f66f082f2f381cf68431c22714c8b70f723837fa2a185b155961eab90f2d5b10", size = 2419859, upload-time = "2026-03-23T18:12:10.272Z" } @@ -4018,76 +3896,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/7b/61/cceae43728b7de99d9b847560c262873a1f6c98202171fd5ed62640b494b/tomli-2.4.1-py3-none-any.whl", hash = "sha256:0d85819802132122da43cb86656f8d1f8c6587d54ae7dcaf30e90533028b49fe", size = 14583, upload-time = "2026-03-25T20:22:03.012Z" }, ] -[[package]] -name = "tomli-w" -version = "1.2.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/19/75/241269d1da26b624c0d5e110e8149093c759b7a286138f4efd61a60e75fe/tomli_w-1.2.0.tar.gz", hash = "sha256:2dd14fac5a47c27be9cd4c976af5a12d87fb1f0b4512f81d69cce3b35ae25021", size = 7184, upload-time = "2025-01-15T12:07:24.262Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c7/18/c86eb8e0202e32dd3df50d43d7ff9854f8e0603945ff398974c1d91ac1ef/tomli_w-1.2.0-py3-none-any.whl", hash = "sha256:188306098d013b691fcadc011abd66727d3c414c571bb01b1a174ba8c983cf90", size = 6675, upload-time = "2025-01-15T12:07:22.074Z" }, -] - [[package]] name = "torch" version = "2.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "filelock" }, - { name = "fsspec" }, - { name = "jinja2" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "setuptools", marker = "python_full_version >= '3.12'" }, - { name = "sympy" }, - { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, - { name = "typing-extensions" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/5b/30/bfebdd8ec77db9a79775121789992d6b3b75ee5494971294d7b4b7c999bc/torch-2.10.0-2-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:2b980edd8d7c0a68c4e951ee1856334a43193f98730d97408fbd148c1a933313", size = 79411457, upload-time = "2026-02-10T21:44:59.189Z" }, - { url = "https://files.pythonhosted.org/packages/0f/8b/4b61d6e13f7108f36910df9ab4b58fd389cc2520d54d81b88660804aad99/torch-2.10.0-2-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:418997cb02d0a0f1497cf6a09f63166f9f5df9f3e16c8a716ab76a72127c714f", size = 79423467, upload-time = "2026-02-10T21:44:48.711Z" }, - { url = "https://files.pythonhosted.org/packages/d3/54/a2ba279afcca44bbd320d4e73675b282fcee3d81400ea1b53934efca6462/torch-2.10.0-2-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:13ec4add8c3faaed8d13e0574f5cd4a323c11655546f91fbe6afa77b57423574", size = 79498202, upload-time = "2026-02-10T21:44:52.603Z" }, - { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" }, - { url = "https://files.pythonhosted.org/packages/16/ee/efbd56687be60ef9af0c9c0ebe106964c07400eade5b0af8902a1d8cd58c/torch-2.10.0-3-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:a1ff626b884f8c4e897c4c33782bdacdff842a165fee79817b1dd549fdda1321", size = 915510070, upload-time = "2026-03-11T14:16:39.386Z" }, - { url = "https://files.pythonhosted.org/packages/36/ab/7b562f1808d3f65414cd80a4f7d4bb00979d9355616c034c171249e1a303/torch-2.10.0-3-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ac5bdcbb074384c66fa160c15b1ead77839e3fe7ed117d667249afce0acabfac", size = 915518691, upload-time = "2026-03-11T14:15:43.147Z" }, - { url = "https://files.pythonhosted.org/packages/b3/7a/abada41517ce0011775f0f4eacc79659bc9bc6c361e6bfe6f7052a6b9363/torch-2.10.0-3-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:98c01b8bb5e3240426dcde1446eed6f40c778091c8544767ef1168fc663a05a6", size = 915622781, upload-time = "2026-03-11T14:17:11.354Z" }, - { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" }, - { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" }, - { url = "https://files.pythonhosted.org/packages/0c/1a/c61f36cfd446170ec27b3a4984f072fd06dab6b5d7ce27e11adb35d6c838/torch-2.10.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:5276fa790a666ee8becaffff8acb711922252521b28fbce5db7db5cf9cb2026d", size = 145992962, upload-time = "2026-01-21T16:24:14.04Z" }, - { url = "https://files.pythonhosted.org/packages/b5/60/6662535354191e2d1555296045b63e4279e5a9dbad49acf55a5d38655a39/torch-2.10.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:aaf663927bcd490ae971469a624c322202a2a1e68936eb952535ca4cd3b90444", size = 915599237, upload-time = "2026-01-21T16:23:25.497Z" }, - { url = "https://files.pythonhosted.org/packages/40/b8/66bbe96f0d79be2b5c697b2e0b187ed792a15c6c4b8904613454651db848/torch-2.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:a4be6a2a190b32ff5c8002a0977a25ea60e64f7ba46b1be37093c141d9c49aeb", size = 113720931, upload-time = "2026-01-21T16:24:23.743Z" }, - { url = "https://files.pythonhosted.org/packages/76/bb/d820f90e69cda6c8169b32a0c6a3ab7b17bf7990b8f2c680077c24a3c14c/torch-2.10.0-cp310-none-macosx_11_0_arm64.whl", hash = "sha256:35e407430795c8d3edb07a1d711c41cc1f9eaddc8b2f1cc0a165a6767a8fb73d", size = 79411450, upload-time = "2026-01-21T16:25:30.692Z" }, - { url = "https://files.pythonhosted.org/packages/78/89/f5554b13ebd71e05c0b002f95148033e730d3f7067f67423026cc9c69410/torch-2.10.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:3282d9febd1e4e476630a099692b44fdc214ee9bf8ee5377732d9d9dfe5712e4", size = 145992610, upload-time = "2026-01-21T16:25:26.327Z" }, - { url = "https://files.pythonhosted.org/packages/ae/30/a3a2120621bf9c17779b169fc17e3dc29b230c29d0f8222f499f5e159aa8/torch-2.10.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:a2f9edd8dbc99f62bc4dfb78af7bf89499bca3d753423ac1b4e06592e467b763", size = 915607863, upload-time = "2026-01-21T16:25:06.696Z" }, - { url = "https://files.pythonhosted.org/packages/6f/3d/c87b33c5f260a2a8ad68da7147e105f05868c281c63d65ed85aa4da98c66/torch-2.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:29b7009dba4b7a1c960260fc8ac85022c784250af43af9fb0ebafc9883782ebd", size = 113723116, upload-time = "2026-01-21T16:25:21.916Z" }, - { url = "https://files.pythonhosted.org/packages/61/d8/15b9d9d3a6b0c01b883787bd056acbe5cc321090d4b216d3ea89a8fcfdf3/torch-2.10.0-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:b7bd80f3477b830dd166c707c5b0b82a898e7b16f59a7d9d42778dd058272e8b", size = 79423461, upload-time = "2026-01-21T16:24:50.266Z" }, - { url = "https://files.pythonhosted.org/packages/cc/af/758e242e9102e9988969b5e621d41f36b8f258bb4a099109b7a4b4b50ea4/torch-2.10.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:5fd4117d89ffd47e3dcc71e71a22efac24828ad781c7e46aaaf56bf7f2796acf", size = 145996088, upload-time = "2026-01-21T16:24:44.171Z" }, - { url = "https://files.pythonhosted.org/packages/23/8e/3c74db5e53bff7ed9e34c8123e6a8bfef718b2450c35eefab85bb4a7e270/torch-2.10.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:787124e7db3b379d4f1ed54dd12ae7c741c16a4d29b49c0226a89bea50923ffb", size = 915711952, upload-time = "2026-01-21T16:23:53.503Z" }, - { url = "https://files.pythonhosted.org/packages/6e/01/624c4324ca01f66ae4c7cd1b74eb16fb52596dce66dbe51eff95ef9e7a4c/torch-2.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:2c66c61f44c5f903046cc696d088e21062644cbe541c7f1c4eaae88b2ad23547", size = 113757972, upload-time = "2026-01-21T16:24:39.516Z" }, - { url = "https://files.pythonhosted.org/packages/c9/5c/dee910b87c4d5c0fcb41b50839ae04df87c1cfc663cf1b5fca7ea565eeaa/torch-2.10.0-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:6d3707a61863d1c4d6ebba7be4ca320f42b869ee657e9b2c21c736bf17000294", size = 79498198, upload-time = "2026-01-21T16:24:34.704Z" }, - { url = "https://files.pythonhosted.org/packages/c9/6f/f2e91e34e3fcba2e3fc8d8f74e7d6c22e74e480bbd1db7bc8900fdf3e95c/torch-2.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5c4d217b14741e40776dd7074d9006fd28b8a97ef5654db959d8635b2fe5f29b", size = 146004247, upload-time = "2026-01-21T16:24:29.335Z" }, - { url = "https://files.pythonhosted.org/packages/98/fb/5160261aeb5e1ee12ee95fe599d0541f7c976c3701d607d8fc29e623229f/torch-2.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6b71486353fce0f9714ca0c9ef1c850a2ae766b409808acd58e9678a3edb7738", size = 915716445, upload-time = "2026-01-21T16:22:45.353Z" }, - { url = "https://files.pythonhosted.org/packages/6a/16/502fb1b41e6d868e8deb5b0e3ae926bbb36dab8ceb0d1b769b266ad7b0c3/torch-2.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2ee399c644dc92ef7bc0d4f7e74b5360c37cdbe7c5ba11318dda49ffac2bc57", size = 113757050, upload-time = "2026-01-21T16:24:19.204Z" }, - { url = "https://files.pythonhosted.org/packages/1a/0b/39929b148f4824bc3ad6f9f72a29d4ad865bcf7ebfc2fa67584773e083d2/torch-2.10.0-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:3202429f58309b9fa96a614885eace4b7995729f44beb54d3e4a47773649d382", size = 79851305, upload-time = "2026-01-21T16:24:09.209Z" }, - { url = "https://files.pythonhosted.org/packages/d8/14/21fbce63bc452381ba5f74a2c0a959fdf5ad5803ccc0c654e752e0dbe91a/torch-2.10.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:aae1b29cd68e50a9397f5ee897b9c24742e9e306f88a807a27d617f07adb3bd8", size = 146005472, upload-time = "2026-01-21T16:22:29.022Z" }, - { url = "https://files.pythonhosted.org/packages/54/fd/b207d1c525cb570ef47f3e9f836b154685011fce11a2f444ba8a4084d042/torch-2.10.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:6021db85958db2f07ec94e1bc77212721ba4920c12a18dc552d2ae36a3eb163f", size = 915612644, upload-time = "2026-01-21T16:21:47.019Z" }, - { url = "https://files.pythonhosted.org/packages/36/53/0197f868c75f1050b199fe58f9bf3bf3aecac9b4e85cc9c964383d745403/torch-2.10.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ff43db38af76fda183156153983c9a096fc4c78d0cd1e07b14a2314c7f01c2c8", size = 113997015, upload-time = "2026-01-21T16:23:00.767Z" }, - { url = "https://files.pythonhosted.org/packages/0e/13/e76b4d9c160e89fff48bf16b449ea324bda84745d2ab30294c37c2434c0d/torch-2.10.0-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:cdf2a523d699b70d613243211ecaac14fe9c5df8a0b0a9c02add60fb2a413e0f", size = 79498248, upload-time = "2026-01-21T16:23:09.315Z" }, + { name = "filelock", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "fsspec", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "jinja2", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "setuptools", marker = "(python_full_version >= '3.12' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.12' and sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "sympy", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "typing-extensions", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, ] [[package]] @@ -4118,7 +3939,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "torch" }, + { name = "torch", marker = "sys_platform == 'never'" }, { name = "torchvision" }, ] sdist = { url = "https://files.pythonhosted.org/packages/6f/36/574c0c46e818533b78b3c09505211162918188325ab4165ef11a3f295755/torchprofile-0.0.4.tar.gz", hash = "sha256:96b6da17d752a06b02977e078aea95614893b31d4117dd5dcd081f30ce65611b", size = 4557, upload-time = "2021-06-22T04:58:03.592Z" } @@ -4134,7 +3955,7 @@ dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "pillow" }, - { name = "torch" }, + { name = "torch", marker = "sys_platform == 'never'" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/50/ae/cbf727421eb73f1cf907fbe5788326a08f111b3f6b6ddca15426b53fec9a/torchvision-0.25.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a95c47abb817d4e90ea1a8e57bd0d728e3e6b533b3495ae77d84d883c4d11f56", size = 1874919, upload-time = "2026-01-21T16:27:47.617Z" }, @@ -4159,41 +3980,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/63/cc/0ea68b5802e5e3c31f44b307e74947bad5a38cc655231d845534ed50ddb8/torchvision-0.25.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5e6b449e9fa7d642142c0e27c41e5a43b508d57ed8e79b7c0a0c28652da8678c", size = 4344260, upload-time = "2026-01-21T16:27:17.018Z" }, ] -[[package]] -name = "tox" -version = "4.53.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cachetools" }, - { name = "colorama" }, - { name = "filelock" }, - { name = "packaging" }, - { name = "platformdirs" }, - { name = "pluggy" }, - { name = "pyproject-api" }, - { name = "python-discovery" }, - { name = "tomli", marker = "python_full_version < '3.11'" }, - { name = "tomli-w" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, - { name = "virtualenv" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/04/01/d87a00063fa670ce4c48a9706b615a95ddf2c9ef5558d43af6071f166fd4/tox-4.53.0.tar.gz", hash = "sha256:62c780e42f87d34ee60f2ea20342156253794fdcbd6885fd797d98ee05009f22", size = 274048, upload-time = "2026-04-14T13:44:13.782Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/16/03/02e2a03f3756cfb66e7e1bac41b06953f12cec75ddb961d56695d4d43dc4/tox-4.53.0-py3-none-any.whl", hash = "sha256:cc4e716d18c4889aa179d785175c438fa60c35deef20ce689ec288d8fb656096", size = 212164, upload-time = "2026-04-14T13:44:11.997Z" }, -] - -[[package]] -name = "tox-current-env" -version = "0.0.16" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "tox" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/35/f3/504ab508410e0660d989eb3ea175f6245271c5582dc4a93268fa9feceeda/tox_current_env-0.0.16.tar.gz", hash = "sha256:2e453c3e82e837d35846004a678db4504e24e5c0419d6e42aa07ca8294fad1bd", size = 25129, upload-time = "2025-03-12T15:52:14.845Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/05/38/b485aed12bb714c53aedbed4919d7ffda7a0f21760229c894feedcc05a61/tox_current_env-0.0.16-py3-none-any.whl", hash = "sha256:88db23525a6514dee6350df01f5b591654b72e7709df3c57e9e1edb1bbc22737", size = 13543, upload-time = "2025-03-12T15:52:13.226Z" }, -] - [[package]] name = "tqdm" version = "4.67.3" @@ -4239,18 +4025,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/91/88/b55b3117287a8540b76dbdd87733808d4d01c8067a3b339408c250bb3600/typeguard-4.5.1-py3-none-any.whl", hash = "sha256:44d2bf329d49a244110a090b55f5f91aa82d9a9834ebfd30bcc73651e4a8cc40", size = 36745, upload-time = "2026-02-19T16:09:01.6Z" }, ] -[[package]] -name = "triton" -version = "3.6.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/8c/f7/f1c9d3424ab199ac53c2da567b859bcddbb9c9e7154805119f8bd95ec36f/triton-3.6.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a6550fae429e0667e397e5de64b332d1e5695b73650ee75a6146e2e902770bea", size = 188105201, upload-time = "2026-01-20T16:00:29.272Z" }, - { url = "https://files.pythonhosted.org/packages/e0/12/b05ba554d2c623bffa59922b94b0775673de251f468a9609bc9e45de95e9/triton-3.6.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e8e323d608e3a9bfcc2d9efcc90ceefb764a82b99dea12a86d643c72539ad5d3", size = 188214640, upload-time = "2026-01-20T16:00:35.869Z" }, - { url = "https://files.pythonhosted.org/packages/ab/a8/cdf8b3e4c98132f965f88c2313a4b493266832ad47fb52f23d14d4f86bb5/triton-3.6.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:74caf5e34b66d9f3a429af689c1c7128daba1d8208df60e81106b115c00d6fca", size = 188266850, upload-time = "2026-01-20T16:00:43.041Z" }, - { url = "https://files.pythonhosted.org/packages/f9/0b/37d991d8c130ce81a8728ae3c25b6e60935838e9be1b58791f5997b24a54/triton-3.6.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10c7f76c6e72d2ef08df639e3d0d30729112f47a56b0c81672edc05ee5116ac9", size = 188289450, upload-time = "2026-01-20T16:00:49.136Z" }, - { url = "https://files.pythonhosted.org/packages/35/f8/9c66bfc55361ec6d0e4040a0337fb5924ceb23de4648b8a81ae9d33b2b38/triton-3.6.0-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d002e07d7180fd65e622134fbd980c9a3d4211fb85224b56a0a0efbd422ab72f", size = 188400296, upload-time = "2026-01-20T16:00:56.042Z" }, -] - [[package]] name = "typer" version = "0.24.1" @@ -4305,6 +4079,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/39/08/aaaad47bc4e9dc8c725e68f9d04865dbcb2052843ff09c97b08904852d84/urllib3-2.6.3-py3-none-any.whl", hash = "sha256:bf272323e553dfb2e87d9bfd225ca7b0f467b919d7bbd355436d3fd37cb0acd4", size = 131584, upload-time = "2026-01-07T16:24:42.685Z" }, ] +[[package]] +name = "uv" +version = "0.11.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9b/7d/17750123a8c8e324627534fe1ae2e7a46689db8492f1a834ab4fd229a7d8/uv-0.11.7.tar.gz", hash = "sha256:46d971489b00bdb27e0aa715e4a5cd4ef2c28ea5b6ef78f2b67bf861eb44b405", size = 4083385, upload-time = "2026-04-15T21:42:55.474Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/5b/2bb2ab6fe6c78c2be10852482ef0cae5f3171460a6e5e24c32c9a0843163/uv-0.11.7-py3-none-linux_armv6l.whl", hash = "sha256:f422d39530516b1dfb28bb6e90c32bb7dacd50f6a383cd6e40c1a859419fbc8c", size = 23757265, upload-time = "2026-04-15T21:43:14.494Z" }, + { url = "https://files.pythonhosted.org/packages/b2/f5/36ff27b01e60a88712628c8a5a6003b8e418883c24e084e506095844a797/uv-0.11.7-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:8b2fe1ec6775dad10183e3fdce430a5b37b7857d49763c884f3a67eaa8ca6f8a", size = 23184529, upload-time = "2026-04-15T21:42:30.225Z" }, + { url = "https://files.pythonhosted.org/packages/8a/fa/f379be661316698f877e78f4c51e5044be0b6f390803387237ad92c4057f/uv-0.11.7-py3-none-macosx_11_0_arm64.whl", hash = "sha256:162fa961a9a081dcea6e889c79f738a5ae56507047e4672964972e33c301bea9", size = 21780167, upload-time = "2026-04-15T21:42:44.942Z" }, + { url = "https://files.pythonhosted.org/packages/f2/7f/fbed29775b0612f4f5679d3226268f1a347161abc1727b4080fb41d9f46f/uv-0.11.7-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:5985a15a92bd9a170fc1947abb1fbc3e9828c5a430ad85b5bed8356c20b67a71", size = 23609640, upload-time = "2026-04-15T21:42:22.57Z" }, + { url = "https://files.pythonhosted.org/packages/ad/de/989a69634a869a22322770120557c2d8cbba5b77ec7cfad326b4ec0f0547/uv-0.11.7-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:fab0bb43fbbc0ee5b5fee212078d2300c371b725faff7cf72eeaafa0bff0606b", size = 23322484, upload-time = "2026-04-15T21:43:26.52Z" }, + { url = "https://files.pythonhosted.org/packages/24/08/c1af05ea602eb4eb75d86badb6b0594cc104c3ca83ccf06d9ed4dd2186ad/uv-0.11.7-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:23d457d6731ebdb83f1bffebe4894edab2ef43c1ec5488433c74300db4958924", size = 23326385, upload-time = "2026-04-15T21:42:41.32Z" }, + { url = "https://files.pythonhosted.org/packages/68/99/e246962da06383e992ecab55000c62a50fb36efef855ea7264fad4816bf4/uv-0.11.7-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7d6a17507b8139b8803f445a03fd097f732ce8356b1b7b13cdb4dd8ef7f4b2e0", size = 24985751, upload-time = "2026-04-15T21:42:37.777Z" }, + { url = "https://files.pythonhosted.org/packages/45/2d/b0b68083859579ce811996c1480765ec6a2442b44c451eaef53e6218fbae/uv-0.11.7-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:dd48823ca4b505124389f49ae50626ba9f57212b9047738efc95126ed5f3844d", size = 25724160, upload-time = "2026-04-15T21:43:18.762Z" }, + { url = "https://files.pythonhosted.org/packages/4e/19/5970e89d9e458fd3c4966bbc586a685a1c0ab0a8bf334503f63fa20b925b/uv-0.11.7-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:eb91f52ee67e10d5290f2c2897e2171357f1a10966de38d83eefa93d96843b0c", size = 25028512, upload-time = "2026-04-15T21:43:02.721Z" }, + { url = "https://files.pythonhosted.org/packages/83/eb/4e1557daf6693cb446ed28185664ad6682fd98c6dbac9e433cbc35df450a/uv-0.11.7-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4e4d5e31bea86e1b6e0f5a0f95e14e80018e6f6c0129256d2915a4b3d793644d", size = 24933975, upload-time = "2026-04-15T21:42:18.828Z" }, + { url = "https://files.pythonhosted.org/packages/68/55/3b517ec8297f110d6981f525cccf26f86e30883fbb9c282769cffbcdcfca/uv-0.11.7-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:ceae53b202ea92bc954759bc7c7570cdcd5c3512fce15701198c19fd2dfb8605", size = 23706403, upload-time = "2026-04-15T21:43:10.664Z" }, + { url = "https://files.pythonhosted.org/packages/dc/30/7d93a0312d60e147722967036dc8ea37baab4802784bddc22464cb707deb/uv-0.11.7-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:f97e9f4e4d44fb5c4dfaa05e858ef3414a96416a2e4af270ecd88a3e5fb049a9", size = 24495797, upload-time = "2026-04-15T21:42:26.538Z" }, + { url = "https://files.pythonhosted.org/packages/8c/89/d49480bdab7725d36982793857e461d471bde8e1b7f438ffccee677a7bf8/uv-0.11.7-py3-none-manylinux_2_31_riscv64.whl", hash = "sha256:750ee5b96959b807cf442b73dd8b55111862d63f258f896787ea5f06b68aaca9", size = 24580471, upload-time = "2026-04-15T21:42:52.871Z" }, + { url = "https://files.pythonhosted.org/packages/b6/9f/c57dc03b48be17b564e304eb9ff982890c12dfb888b1ce370788733329ab/uv-0.11.7-py3-none-musllinux_1_1_i686.whl", hash = "sha256:f394331f0507e80ee732cb3df737589de53bed999dd02a6d24682f08c2f8ac4f", size = 24113637, upload-time = "2026-04-15T21:42:34.094Z" }, + { url = "https://files.pythonhosted.org/packages/13/ba/b87e358b629a68258527e3490e73b7b148770f4d2257842dea3b7981d4e8/uv-0.11.7-py3-none-musllinux_1_1_x86_64.whl", hash = "sha256:0df59ab0c6a4b14a763e8445e1c303af9abeb53cdfa4428daf9ff9642c0a3cce", size = 25119850, upload-time = "2026-04-15T21:43:22.529Z" }, + { url = "https://files.pythonhosted.org/packages/4b/74/16d229e1d8574bcbafa6dc643ac20b70c3e581f42ac31a6f4fd53035ffe3/uv-0.11.7-py3-none-win32.whl", hash = "sha256:553e67cc766d013ce24353fecd4ea5533d2aedcfd35f9fac430e07b1d1f23ed4", size = 22918454, upload-time = "2026-04-15T21:42:58.702Z" }, + { url = "https://files.pythonhosted.org/packages/a6/1d/b73e473da616ac758b8918fb218febcc46ddf64cba9e03894dfa226b28bd/uv-0.11.7-py3-none-win_amd64.whl", hash = "sha256:5674dfb5944513f4b3735b05c2deba6b1b01151f46729d533d413a9a905f8c5d", size = 25447744, upload-time = "2026-04-15T21:42:48.813Z" }, + { url = "https://files.pythonhosted.org/packages/1b/bb/e6bfdea92ed270f3445a5a3c17599d041b3f2dbc5026c09e02830a03bbaf/uv-0.11.7-py3-none-win_arm64.whl", hash = "sha256:6158b7e39464f1aa1e040daa0186cae4749a78b5cd80ac769f32ca711b8976b1", size = 23941816, upload-time = "2026-04-15T21:43:06.732Z" }, +] + [[package]] name = "uvicorn" version = "0.44.0" From e9a49890f173997b01818d104421327a72f8afb3 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Sat, 18 Apr 2026 17:22:23 +0000 Subject: [PATCH 20/30] [chore]: weekly bump of uv.lock on main (2026-04-18) (#1292) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Automated weekly update of uv.lock file for nSpect Scanning: - `uv.lock` — upgraded all transitive dependencies to latest compatible versions Signed-off-by: github-actions[bot] Co-authored-by: github-actions[bot] --- uv.lock | 363 +++++++++++++++++++++----------------------------------- 1 file changed, 133 insertions(+), 230 deletions(-) diff --git a/uv.lock b/uv.lock index e223ce17eae..e70f5d3b780 100644 --- a/uv.lock +++ b/uv.lock @@ -278,18 +278,6 @@ toml = [ { name = "tomli", marker = "python_full_version < '3.11'" }, ] -[[package]] -name = "cbcbox" -version = "2.929" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6a/9e/c41844f3a746500b88982817233f3f2f4fa6d3a554cedb9222bac9efcbea/cbcbox-2.929-py3-none-macosx_15_0_arm64.whl", hash = "sha256:610fd250f737b19f599d56cca299682a4519472b3ff6a13044502eb71e03931e", size = 59626733, upload-time = "2026-03-23T15:28:20.758Z" }, - { url = "https://files.pythonhosted.org/packages/3f/08/172af0d618ede8668862ae9c83ffcd929bc966610c1ce357b268edb33645/cbcbox-2.929-py3-none-macosx_15_0_x86_64.whl", hash = "sha256:b369071081176fd55ba9d104c2f2330133c36f0264ebe0215c50a366f290a050", size = 87883726, upload-time = "2026-03-23T15:28:24.874Z" }, - { url = "https://files.pythonhosted.org/packages/d4/c7/a57493dd959a57dcb8dc31aeef468c0e95d066d4bbf8fc59f9dda3a1fd97/cbcbox-2.929-py3-none-manylinux2014_aarch64.whl", hash = "sha256:62841ab4ed4e1c368bcba8d62c66cd7dd03f3dfaf1ecf4a01c6385a122d1946e", size = 145740150, upload-time = "2026-03-23T15:28:29.617Z" }, - { url = "https://files.pythonhosted.org/packages/4c/86/502b81d060603e1016a8cb471064e4e65d8009c2feed60ea878b8513a728/cbcbox-2.929-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a57a403cc44fb2af0aa53cb790a7642390f1836a5ff6a7a0b31b30eb6597e3e0", size = 181208771, upload-time = "2026-03-23T15:28:35.544Z" }, - { url = "https://files.pythonhosted.org/packages/0c/85/2fef9c7c3054ebfb9b0e84fabbeac776a18514b025efa5abd768c889f560/cbcbox-2.929-py3-none-win_amd64.whl", hash = "sha256:0a91420befb965ec1763aab62d7908d3a89a51af845e2a4e05c09d99a67f1151", size = 135305901, upload-time = "2026-03-23T15:28:42.021Z" }, -] - [[package]] name = "certifi" version = "2026.2.25" @@ -299,66 +287,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/9a/3c/c17fb3ca2d9c3acff52e30b309f538586f9f5b9c9cf454f3845fc9af4881/certifi-2026.2.25-py3-none-any.whl", hash = "sha256:027692e4402ad994f1c42e52a4997a9763c646b73e4096e4d5d6db8af1d6f0fa", size = 153684, upload-time = "2026-02-25T02:54:15.766Z" }, ] -[[package]] -name = "cffi" -version = "2.0.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/93/d7/516d984057745a6cd96575eea814fe1edd6646ee6efd552fb7b0921dec83/cffi-2.0.0-cp310-cp310-macosx_10_13_x86_64.whl", hash = "sha256:0cf2d91ecc3fcc0625c2c530fe004f82c110405f101548512cce44322fa8ac44", size = 184283, upload-time = "2025-09-08T23:22:08.01Z" }, - { url = "https://files.pythonhosted.org/packages/9e/84/ad6a0b408daa859246f57c03efd28e5dd1b33c21737c2db84cae8c237aa5/cffi-2.0.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:f73b96c41e3b2adedc34a7356e64c8eb96e03a3782b535e043a986276ce12a49", size = 180504, upload-time = "2025-09-08T23:22:10.637Z" }, - { url = "https://files.pythonhosted.org/packages/50/bd/b1a6362b80628111e6653c961f987faa55262b4002fcec42308cad1db680/cffi-2.0.0-cp310-cp310-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:53f77cbe57044e88bbd5ed26ac1d0514d2acf0591dd6bb02a3ae37f76811b80c", size = 208811, upload-time = "2025-09-08T23:22:12.267Z" }, - { url = "https://files.pythonhosted.org/packages/4f/27/6933a8b2562d7bd1fb595074cf99cc81fc3789f6a6c05cdabb46284a3188/cffi-2.0.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:3e837e369566884707ddaf85fc1744b47575005c0a229de3327f8f9a20f4efeb", size = 216402, upload-time = "2025-09-08T23:22:13.455Z" }, - { url = "https://files.pythonhosted.org/packages/05/eb/b86f2a2645b62adcfff53b0dd97e8dfafb5c8aa864bd0d9a2c2049a0d551/cffi-2.0.0-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:5eda85d6d1879e692d546a078b44251cdd08dd1cfb98dfb77b670c97cee49ea0", size = 203217, upload-time = "2025-09-08T23:22:14.596Z" }, - { url = "https://files.pythonhosted.org/packages/9f/e0/6cbe77a53acf5acc7c08cc186c9928864bd7c005f9efd0d126884858a5fe/cffi-2.0.0-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:9332088d75dc3241c702d852d4671613136d90fa6881da7d770a483fd05248b4", size = 203079, upload-time = "2025-09-08T23:22:15.769Z" }, - { url = "https://files.pythonhosted.org/packages/98/29/9b366e70e243eb3d14a5cb488dfd3a0b6b2f1fb001a203f653b93ccfac88/cffi-2.0.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fc7de24befaeae77ba923797c7c87834c73648a05a4bde34b3b7e5588973a453", size = 216475, upload-time = "2025-09-08T23:22:17.427Z" }, - { url = "https://files.pythonhosted.org/packages/21/7a/13b24e70d2f90a322f2900c5d8e1f14fa7e2a6b3332b7309ba7b2ba51a5a/cffi-2.0.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:cf364028c016c03078a23b503f02058f1814320a56ad535686f90565636a9495", size = 218829, upload-time = "2025-09-08T23:22:19.069Z" }, - { url = "https://files.pythonhosted.org/packages/60/99/c9dc110974c59cc981b1f5b66e1d8af8af764e00f0293266824d9c4254bc/cffi-2.0.0-cp310-cp310-musllinux_1_2_i686.whl", hash = "sha256:e11e82b744887154b182fd3e7e8512418446501191994dbf9c9fc1f32cc8efd5", size = 211211, upload-time = "2025-09-08T23:22:20.588Z" }, - { url = "https://files.pythonhosted.org/packages/49/72/ff2d12dbf21aca1b32a40ed792ee6b40f6dc3a9cf1644bd7ef6e95e0ac5e/cffi-2.0.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:8ea985900c5c95ce9db1745f7933eeef5d314f0565b27625d9a10ec9881e1bfb", size = 218036, upload-time = "2025-09-08T23:22:22.143Z" }, - { url = "https://files.pythonhosted.org/packages/e2/cc/027d7fb82e58c48ea717149b03bcadcbdc293553edb283af792bd4bcbb3f/cffi-2.0.0-cp310-cp310-win32.whl", hash = "sha256:1f72fb8906754ac8a2cc3f9f5aaa298070652a0ffae577e0ea9bd480dc3c931a", size = 172184, upload-time = "2025-09-08T23:22:23.328Z" }, - { url = "https://files.pythonhosted.org/packages/33/fa/072dd15ae27fbb4e06b437eb6e944e75b068deb09e2a2826039e49ee2045/cffi-2.0.0-cp310-cp310-win_amd64.whl", hash = "sha256:b18a3ed7d5b3bd8d9ef7a8cb226502c6bf8308df1525e1cc676c3680e7176739", size = 182790, upload-time = "2025-09-08T23:22:24.752Z" }, - { url = "https://files.pythonhosted.org/packages/12/4a/3dfd5f7850cbf0d06dc84ba9aa00db766b52ca38d8b86e3a38314d52498c/cffi-2.0.0-cp311-cp311-macosx_10_13_x86_64.whl", hash = "sha256:b4c854ef3adc177950a8dfc81a86f5115d2abd545751a304c5bcf2c2c7283cfe", size = 184344, upload-time = "2025-09-08T23:22:26.456Z" }, - { url = "https://files.pythonhosted.org/packages/4f/8b/f0e4c441227ba756aafbe78f117485b25bb26b1c059d01f137fa6d14896b/cffi-2.0.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:2de9a304e27f7596cd03d16f1b7c72219bd944e99cc52b84d0145aefb07cbd3c", size = 180560, upload-time = "2025-09-08T23:22:28.197Z" }, - { url = "https://files.pythonhosted.org/packages/b1/b7/1200d354378ef52ec227395d95c2576330fd22a869f7a70e88e1447eb234/cffi-2.0.0-cp311-cp311-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:baf5215e0ab74c16e2dd324e8ec067ef59e41125d3eade2b863d294fd5035c92", size = 209613, upload-time = "2025-09-08T23:22:29.475Z" }, - { url = "https://files.pythonhosted.org/packages/b8/56/6033f5e86e8cc9bb629f0077ba71679508bdf54a9a5e112a3c0b91870332/cffi-2.0.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:730cacb21e1bdff3ce90babf007d0a0917cc3e6492f336c2f0134101e0944f93", size = 216476, upload-time = "2025-09-08T23:22:31.063Z" }, - { url = "https://files.pythonhosted.org/packages/dc/7f/55fecd70f7ece178db2f26128ec41430d8720f2d12ca97bf8f0a628207d5/cffi-2.0.0-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:6824f87845e3396029f3820c206e459ccc91760e8fa24422f8b0c3d1731cbec5", size = 203374, upload-time = "2025-09-08T23:22:32.507Z" }, - { url = "https://files.pythonhosted.org/packages/84/ef/a7b77c8bdc0f77adc3b46888f1ad54be8f3b7821697a7b89126e829e676a/cffi-2.0.0-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:9de40a7b0323d889cf8d23d1ef214f565ab154443c42737dfe52ff82cf857664", size = 202597, upload-time = "2025-09-08T23:22:34.132Z" }, - { url = "https://files.pythonhosted.org/packages/d7/91/500d892b2bf36529a75b77958edfcd5ad8e2ce4064ce2ecfeab2125d72d1/cffi-2.0.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8941aaadaf67246224cee8c3803777eed332a19d909b47e29c9842ef1e79ac26", size = 215574, upload-time = "2025-09-08T23:22:35.443Z" }, - { url = "https://files.pythonhosted.org/packages/44/64/58f6255b62b101093d5df22dcb752596066c7e89dd725e0afaed242a61be/cffi-2.0.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:a05d0c237b3349096d3981b727493e22147f934b20f6f125a3eba8f994bec4a9", size = 218971, upload-time = "2025-09-08T23:22:36.805Z" }, - { url = "https://files.pythonhosted.org/packages/ab/49/fa72cebe2fd8a55fbe14956f9970fe8eb1ac59e5df042f603ef7c8ba0adc/cffi-2.0.0-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:94698a9c5f91f9d138526b48fe26a199609544591f859c870d477351dc7b2414", size = 211972, upload-time = "2025-09-08T23:22:38.436Z" }, - { url = "https://files.pythonhosted.org/packages/0b/28/dd0967a76aab36731b6ebfe64dec4e981aff7e0608f60c2d46b46982607d/cffi-2.0.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:5fed36fccc0612a53f1d4d9a816b50a36702c28a2aa880cb8a122b3466638743", size = 217078, upload-time = "2025-09-08T23:22:39.776Z" }, - { url = "https://files.pythonhosted.org/packages/2b/c0/015b25184413d7ab0a410775fdb4a50fca20f5589b5dab1dbbfa3baad8ce/cffi-2.0.0-cp311-cp311-win32.whl", hash = "sha256:c649e3a33450ec82378822b3dad03cc228b8f5963c0c12fc3b1e0ab940f768a5", size = 172076, upload-time = "2025-09-08T23:22:40.95Z" }, - { url = "https://files.pythonhosted.org/packages/ae/8f/dc5531155e7070361eb1b7e4c1a9d896d0cb21c49f807a6c03fd63fc877e/cffi-2.0.0-cp311-cp311-win_amd64.whl", hash = "sha256:66f011380d0e49ed280c789fbd08ff0d40968ee7b665575489afa95c98196ab5", size = 182820, upload-time = "2025-09-08T23:22:42.463Z" }, - { url = "https://files.pythonhosted.org/packages/95/5c/1b493356429f9aecfd56bc171285a4c4ac8697f76e9bbbbb105e537853a1/cffi-2.0.0-cp311-cp311-win_arm64.whl", hash = "sha256:c6638687455baf640e37344fe26d37c404db8b80d037c3d29f58fe8d1c3b194d", size = 177635, upload-time = "2025-09-08T23:22:43.623Z" }, - { url = "https://files.pythonhosted.org/packages/ea/47/4f61023ea636104d4f16ab488e268b93008c3d0bb76893b1b31db1f96802/cffi-2.0.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:6d02d6655b0e54f54c4ef0b94eb6be0607b70853c45ce98bd278dc7de718be5d", size = 185271, upload-time = "2025-09-08T23:22:44.795Z" }, - { url = "https://files.pythonhosted.org/packages/df/a2/781b623f57358e360d62cdd7a8c681f074a71d445418a776eef0aadb4ab4/cffi-2.0.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8eca2a813c1cb7ad4fb74d368c2ffbbb4789d377ee5bb8df98373c2cc0dee76c", size = 181048, upload-time = "2025-09-08T23:22:45.938Z" }, - { url = "https://files.pythonhosted.org/packages/ff/df/a4f0fbd47331ceeba3d37c2e51e9dfc9722498becbeec2bd8bc856c9538a/cffi-2.0.0-cp312-cp312-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:21d1152871b019407d8ac3985f6775c079416c282e431a4da6afe7aefd2bccbe", size = 212529, upload-time = "2025-09-08T23:22:47.349Z" }, - { url = "https://files.pythonhosted.org/packages/d5/72/12b5f8d3865bf0f87cf1404d8c374e7487dcf097a1c91c436e72e6badd83/cffi-2.0.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b21e08af67b8a103c71a250401c78d5e0893beff75e28c53c98f4de42f774062", size = 220097, upload-time = "2025-09-08T23:22:48.677Z" }, - { url = "https://files.pythonhosted.org/packages/c2/95/7a135d52a50dfa7c882ab0ac17e8dc11cec9d55d2c18dda414c051c5e69e/cffi-2.0.0-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:1e3a615586f05fc4065a8b22b8152f0c1b00cdbc60596d187c2a74f9e3036e4e", size = 207983, upload-time = "2025-09-08T23:22:50.06Z" }, - { url = "https://files.pythonhosted.org/packages/3a/c8/15cb9ada8895957ea171c62dc78ff3e99159ee7adb13c0123c001a2546c1/cffi-2.0.0-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:81afed14892743bbe14dacb9e36d9e0e504cd204e0b165062c488942b9718037", size = 206519, upload-time = "2025-09-08T23:22:51.364Z" }, - { url = "https://files.pythonhosted.org/packages/78/2d/7fa73dfa841b5ac06c7b8855cfc18622132e365f5b81d02230333ff26e9e/cffi-2.0.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:3e17ed538242334bf70832644a32a7aae3d83b57567f9fd60a26257e992b79ba", size = 219572, upload-time = "2025-09-08T23:22:52.902Z" }, - { url = "https://files.pythonhosted.org/packages/07/e0/267e57e387b4ca276b90f0434ff88b2c2241ad72b16d31836adddfd6031b/cffi-2.0.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3925dd22fa2b7699ed2617149842d2e6adde22b262fcbfada50e3d195e4b3a94", size = 222963, upload-time = "2025-09-08T23:22:54.518Z" }, - { url = "https://files.pythonhosted.org/packages/b6/75/1f2747525e06f53efbd878f4d03bac5b859cbc11c633d0fb81432d98a795/cffi-2.0.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:2c8f814d84194c9ea681642fd164267891702542f028a15fc97d4674b6206187", size = 221361, upload-time = "2025-09-08T23:22:55.867Z" }, - { url = "https://files.pythonhosted.org/packages/7b/2b/2b6435f76bfeb6bbf055596976da087377ede68df465419d192acf00c437/cffi-2.0.0-cp312-cp312-win32.whl", hash = "sha256:da902562c3e9c550df360bfa53c035b2f241fed6d9aef119048073680ace4a18", size = 172932, upload-time = "2025-09-08T23:22:57.188Z" }, - { url = "https://files.pythonhosted.org/packages/f8/ed/13bd4418627013bec4ed6e54283b1959cf6db888048c7cf4b4c3b5b36002/cffi-2.0.0-cp312-cp312-win_amd64.whl", hash = "sha256:da68248800ad6320861f129cd9c1bf96ca849a2771a59e0344e88681905916f5", size = 183557, upload-time = "2025-09-08T23:22:58.351Z" }, - { url = "https://files.pythonhosted.org/packages/95/31/9f7f93ad2f8eff1dbc1c3656d7ca5bfd8fb52c9d786b4dcf19b2d02217fa/cffi-2.0.0-cp312-cp312-win_arm64.whl", hash = "sha256:4671d9dd5ec934cb9a73e7ee9676f9362aba54f7f34910956b84d727b0d73fb6", size = 177762, upload-time = "2025-09-08T23:22:59.668Z" }, - { url = "https://files.pythonhosted.org/packages/4b/8d/a0a47a0c9e413a658623d014e91e74a50cdd2c423f7ccfd44086ef767f90/cffi-2.0.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:00bdf7acc5f795150faa6957054fbbca2439db2f775ce831222b66f192f03beb", size = 185230, upload-time = "2025-09-08T23:23:00.879Z" }, - { url = "https://files.pythonhosted.org/packages/4a/d2/a6c0296814556c68ee32009d9c2ad4f85f2707cdecfd7727951ec228005d/cffi-2.0.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:45d5e886156860dc35862657e1494b9bae8dfa63bf56796f2fb56e1679fc0bca", size = 181043, upload-time = "2025-09-08T23:23:02.231Z" }, - { url = "https://files.pythonhosted.org/packages/b0/1e/d22cc63332bd59b06481ceaac49d6c507598642e2230f201649058a7e704/cffi-2.0.0-cp313-cp313-manylinux1_i686.manylinux2014_i686.manylinux_2_17_i686.manylinux_2_5_i686.whl", hash = "sha256:07b271772c100085dd28b74fa0cd81c8fb1a3ba18b21e03d7c27f3436a10606b", size = 212446, upload-time = "2025-09-08T23:23:03.472Z" }, - { url = "https://files.pythonhosted.org/packages/a9/f5/a2c23eb03b61a0b8747f211eb716446c826ad66818ddc7810cc2cc19b3f2/cffi-2.0.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d48a880098c96020b02d5a1f7d9251308510ce8858940e6fa99ece33f610838b", size = 220101, upload-time = "2025-09-08T23:23:04.792Z" }, - { url = "https://files.pythonhosted.org/packages/f2/7f/e6647792fc5850d634695bc0e6ab4111ae88e89981d35ac269956605feba/cffi-2.0.0-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:f93fd8e5c8c0a4aa1f424d6173f14a892044054871c771f8566e4008eaa359d2", size = 207948, upload-time = "2025-09-08T23:23:06.127Z" }, - { url = "https://files.pythonhosted.org/packages/cb/1e/a5a1bd6f1fb30f22573f76533de12a00bf274abcdc55c8edab639078abb6/cffi-2.0.0-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.whl", hash = "sha256:dd4f05f54a52fb558f1ba9f528228066954fee3ebe629fc1660d874d040ae5a3", size = 206422, upload-time = "2025-09-08T23:23:07.753Z" }, - { url = "https://files.pythonhosted.org/packages/98/df/0a1755e750013a2081e863e7cd37e0cdd02664372c754e5560099eb7aa44/cffi-2.0.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c8d3b5532fc71b7a77c09192b4a5a200ea992702734a2e9279a37f2478236f26", size = 219499, upload-time = "2025-09-08T23:23:09.648Z" }, - { url = "https://files.pythonhosted.org/packages/50/e1/a969e687fcf9ea58e6e2a928ad5e2dd88cc12f6f0ab477e9971f2309b57c/cffi-2.0.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:d9b29c1f0ae438d5ee9acb31cadee00a58c46cc9c0b2f9038c6b0b3470877a8c", size = 222928, upload-time = "2025-09-08T23:23:10.928Z" }, - { url = "https://files.pythonhosted.org/packages/36/54/0362578dd2c9e557a28ac77698ed67323ed5b9775ca9d3fe73fe191bb5d8/cffi-2.0.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6d50360be4546678fc1b79ffe7a66265e28667840010348dd69a314145807a1b", size = 221302, upload-time = "2025-09-08T23:23:12.42Z" }, - { url = "https://files.pythonhosted.org/packages/eb/6d/bf9bda840d5f1dfdbf0feca87fbdb64a918a69bca42cfa0ba7b137c48cb8/cffi-2.0.0-cp313-cp313-win32.whl", hash = "sha256:74a03b9698e198d47562765773b4a8309919089150a0bb17d829ad7b44b60d27", size = 172909, upload-time = "2025-09-08T23:23:14.32Z" }, - { url = "https://files.pythonhosted.org/packages/37/18/6519e1ee6f5a1e579e04b9ddb6f1676c17368a7aba48299c3759bbc3c8b3/cffi-2.0.0-cp313-cp313-win_amd64.whl", hash = "sha256:19f705ada2530c1167abacb171925dd886168931e0a7b78f5bffcae5c6b5be75", size = 183402, upload-time = "2025-09-08T23:23:15.535Z" }, - { url = "https://files.pythonhosted.org/packages/cb/0e/02ceeec9a7d6ee63bb596121c2c8e9b3a9e150936f4fbef6ca1943e6137c/cffi-2.0.0-cp313-cp313-win_arm64.whl", hash = "sha256:256f80b80ca3853f90c21b23ee78cd008713787b1b1e93eae9f3d6a7134abd91", size = 177780, upload-time = "2025-09-08T23:23:16.761Z" }, -] - [[package]] name = "cfgv" version = "3.5.0" @@ -576,14 +504,15 @@ toml = [ [[package]] name = "cppimport" -version = "22.8.2" +version = "26.4.17" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, { name = "mako" }, { name = "pybind11" }, + { name = "setuptools" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/54/27/01d9078a77b9e31b79b9716e66ca4db74f4744c5232bcb3e8769395c4280/cppimport-22.8.2.tar.gz", hash = "sha256:bbb4957102db41bc99ad72c233bce92f9d1fd91be352fc07878c4361033a401f", size = 26635, upload-time = "2022-08-02T16:50:36.872Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3b/71/07103183044d791dc8971387d34453c5237270807a96a3548878d624cd78/cppimport-26.4.17.tar.gz", hash = "sha256:1fd9c1f16b2cf346507ef81b17f9b4947ddf0d5818b26f6c03078e8e7b366444", size = 28620, upload-time = "2026-04-17T08:38:06.466Z" } [[package]] name = "cuda-pathfinder" @@ -649,18 +578,18 @@ name = "deepspeed" version = "0.18.9" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "einops", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "hjson", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "msgpack", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "ninja", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "packaging", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "psutil", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "py-cpuinfo", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "pydantic", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "einops", marker = "sys_platform != 'win32'" }, + { name = "hjson", marker = "sys_platform != 'win32'" }, + { name = "msgpack", marker = "sys_platform != 'win32'" }, + { name = "ninja", marker = "sys_platform != 'win32'" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' and sys_platform != 'win32'" }, + { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11' and sys_platform != 'win32'" }, + { name = "packaging", marker = "sys_platform != 'win32'" }, + { name = "psutil", marker = "sys_platform != 'win32'" }, + { name = "py-cpuinfo", marker = "sys_platform != 'win32'" }, + { name = "pydantic", marker = "sys_platform != 'win32'" }, { name = "torch", marker = "sys_platform == 'never'" }, - { name = "tqdm", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "tqdm", marker = "sys_platform != 'win32'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/24/61/5ea1c63b139fe7530b196b68ce0bdffa9cde79882e527dcecae58bd6c770/deepspeed-0.18.9.tar.gz", hash = "sha256:ee4818dcf342794f74f429a0aeebef90291ec808fa82609c5140c23e665c4011", size = 1663466, upload-time = "2026-03-30T16:43:16.566Z" } @@ -950,7 +879,7 @@ wheels = [ [[package]] name = "huggingface-hub" -version = "1.10.2" +version = "1.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "filelock" }, @@ -963,9 +892,9 @@ dependencies = [ { name = "typer" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0c/4d/00734890c7fcfe2c7ff04f1c1a167186c42b19e370a2dd8cfd8c34fc92c4/huggingface_hub-1.10.2.tar.gz", hash = "sha256:4b276f820483b709dc86a53bcb8183ea496b8d8447c9f7f88a115a12b498a95f", size = 758428, upload-time = "2026-04-14T10:42:28.498Z" } +sdist = { url = "https://files.pythonhosted.org/packages/dc/89/e7aa12d8a6b9259bed10671abb25ae6fa437c0f88a86ecbf59617bae7759/huggingface_hub-1.11.0.tar.gz", hash = "sha256:15fb3713c7f9cdff7b808a94fd91664f661ab142796bb48c9cd9493e8d166278", size = 761749, upload-time = "2026-04-16T13:07:39.73Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5e/c9/4c1e1216b24bcab140c83acdf8bc89a846ea17cd8a06cd18e3fd308a297f/huggingface_hub-1.10.2-py3-none-any.whl", hash = "sha256:c26c908767cc711493978dc0b4f5747ba7841602997cc98bfd628450a28cf9bc", size = 642581, upload-time = "2026-04-14T10:42:26.563Z" }, + { url = "https://files.pythonhosted.org/packages/37/02/4f3f8997d1ea7fe0146b343e5e14bd065fa87af790d07e5576d31b31cc18/huggingface_hub-1.11.0-py3-none-any.whl", hash = "sha256:42a6de0afbfeb5e022222d36398f029679db4eb4778801aafda32257ae9131ab", size = 645499, upload-time = "2026-04-16T13:07:37.716Z" }, ] [[package]] @@ -1005,11 +934,11 @@ wheels = [ [[package]] name = "identify" -version = "2.6.18" +version = "2.6.19" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/46/c4/7fb4db12296cdb11893d61c92048fe617ee853f8523b9b296ac03b43757e/identify-2.6.18.tar.gz", hash = "sha256:873ac56a5e3fd63e7438a7ecbc4d91aca692eb3fefa4534db2b7913f3fc352fd", size = 99580, upload-time = "2026-03-15T18:39:50.319Z" } +sdist = { url = "https://files.pythonhosted.org/packages/52/63/51723b5f116cc04b061cb6f5a561790abf249d25931d515cd375e063e0f4/identify-2.6.19.tar.gz", hash = "sha256:6be5020c38fcb07da56c53733538a3081ea5aa70d36a156f83044bfbf9173842", size = 99567, upload-time = "2026-04-17T18:39:50.265Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/46/33/92ef41c6fad0233e41d3d84ba8e8ad18d1780f1e5d99b3c683e6d7f98b63/identify-2.6.18-py2.py3-none-any.whl", hash = "sha256:8db9d3c8ea9079db92cafb0ebf97abdc09d52e97f4dcf773a2e694048b7cd737", size = 99394, upload-time = "2026-03-15T18:39:48.915Z" }, + { url = "https://files.pythonhosted.org/packages/94/84/d9273cd09688070a6523c4aee4663a8538721b2b755c4962aafae0011e72/identify-2.6.19-py2.py3-none-any.whl", hash = "sha256:20e6a87f786f768c092a721ad107fc9df0eb89347be9396cadf3f4abbd1fb78a", size = 99397, upload-time = "2026-04-17T18:39:49.221Z" }, ] [[package]] @@ -1295,19 +1224,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, ] -[[package]] -name = "mip" -version = "1.17.6" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cbcbox" }, - { name = "cffi" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/96/15/7496ce00eafc5b99e98bb696588527754d7629f11ac6208fdc9355f6eec6/mip-1.17.6.tar.gz", hash = "sha256:9da8074b80bd3ef788513d5a214ef832916d82aa66487da11a49c7da9f89d270", size = 9443716, upload-time = "2026-03-23T16:20:09.009Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/65/1d/0eb2531e779be687a3249e7f2ba5e465dd34611072390252e5fcf89e7fef/mip-1.17.6-py3-none-any.whl", hash = "sha256:4fb7ff5d7beacbe7d007de172306a76aff8fb36760fa63a5d88a8cb3aa6f4a57", size = 88214, upload-time = "2026-03-23T16:20:07.289Z" }, -] - [[package]] name = "ml-dtypes" version = "0.5.4" @@ -1850,7 +1766,6 @@ all = [ { name = "immutabledict" }, { name = "lief" }, { name = "lru-dict" }, - { name = "mip" }, { name = "ml-dtypes" }, { name = "nltk" }, { name = "onnx" }, @@ -1888,7 +1803,6 @@ dev = [ { name = "immutabledict" }, { name = "lief" }, { name = "lru-dict" }, - { name = "mip" }, { name = "ml-dtypes" }, { name = "mypy" }, { name = "nltk" }, @@ -1995,7 +1909,6 @@ puzzletron = [ { name = "hydra-core" }, { name = "immutabledict" }, { name = "lru-dict" }, - { name = "mip" }, { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "pandas", version = "3.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "typeguard" }, @@ -2018,7 +1931,6 @@ requires-dist = [ { name = "immutabledict", marker = "extra == 'puzzletron'" }, { name = "lief", marker = "extra == 'onnx'" }, { name = "lru-dict", marker = "extra == 'puzzletron'" }, - { name = "mip", marker = "extra == 'puzzletron'" }, { name = "ml-dtypes", marker = "extra == 'onnx'" }, { name = "mypy", marker = "extra == 'dev-lint'", specifier = "==1.17.1" }, { name = "ninja" }, @@ -2484,7 +2396,7 @@ wheels = [ [[package]] name = "peft" -version = "0.19.0" +version = "0.19.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "accelerate" }, @@ -2499,9 +2411,9 @@ dependencies = [ { name = "tqdm" }, { name = "transformers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2a/58/2e758e0794daa49dd9a47c56da7e31ee16d66d09ec787bbe44b1486f84fd/peft-0.19.0.tar.gz", hash = "sha256:a2917070092184a462093443029bc4f9292a91b9b99880488e319309ff0a172d", size = 762553, upload-time = "2026-04-14T14:01:53.189Z" } +sdist = { url = "https://files.pythonhosted.org/packages/86/cf/037f1e3d5186496c05513a6754639e2dab3038a05f384284d49a9bd06a2d/peft-0.19.1.tar.gz", hash = "sha256:0d97542fe96dcdaa20d3b81c06f26f988618f416a73544ab23c3618ccb674a40", size = 763738, upload-time = "2026-04-16T15:46:45.105Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1f/fd/99e9beed55de9d54f6cd038880b4db4bacb45371dd54d40ea3fda7eaff5e/peft-0.19.0-py3-none-any.whl", hash = "sha256:7feca0f07bee9101807c7fd4601353d91161ea9e1f450150ee7859b2354c7690", size = 680671, upload-time = "2026-04-14T14:01:51.279Z" }, + { url = "https://files.pythonhosted.org/packages/e8/b6/f54d676ed93cc2dd2234c3b172ea9c8c3d7d29361e66b1b23dec57a67465/peft-0.19.1-py3-none-any.whl", hash = "sha256:2113f72a81621b5913ef28f9022204c742df111890c5f49d812716a4a301e356", size = 680692, upload-time = "2026-04-16T15:46:42.886Z" }, ] [[package]] @@ -2810,18 +2722,9 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ab/87/99f21e9b20899d6dc1bf7544cfe53e5fa17acc21bb267971a540425357d3/pybind11-3.0.3-py3-none-any.whl", hash = "sha256:fb5f8e4a64946b4dcc0451c83a8c384f803bc0a62dd1ba02f199e97dbc9aad4c", size = 313717, upload-time = "2026-03-31T23:42:04.814Z" }, ] -[[package]] -name = "pycparser" -version = "3.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/1b/7d/92392ff7815c21062bea51aa7b87d45576f649f16458d78b7cf94b9ab2e6/pycparser-3.0.tar.gz", hash = "sha256:600f49d217304a5902ac3c37e1281c9fe94e4d0489de643a9504c5cdfdfc6b29", size = 103492, upload-time = "2026-01-21T14:26:51.89Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/0c/c3/44f3fbbfa403ea2a7c779186dc20772604442dde72947e7d01069cbe98e3/pycparser-3.0-py3-none-any.whl", hash = "sha256:b727414169a36b7d524c1c3e31839a521725078d7b2ff038656844266160a992", size = 48172, upload-time = "2026-01-21T14:26:50.693Z" }, -] - [[package]] name = "pydantic" -version = "2.13.1" +version = "2.13.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "annotated-types" }, @@ -2829,95 +2732,95 @@ dependencies = [ { name = "typing-extensions" }, { name = "typing-inspection" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f3/6b/1353beb3d1cd5cf61cdec5b6f87a9872399de3bc5cae0b7ce07ff4de2ab0/pydantic-2.13.1.tar.gz", hash = "sha256:a0f829b279ddd1e39291133fe2539d2aa46cc6b150c1706a270ff0879e3774d2", size = 843746, upload-time = "2026-04-15T14:57:19.398Z" } +sdist = { url = "https://files.pythonhosted.org/packages/09/e5/06d23afac9973109d1e3c8ad38e1547a12e860610e327c05ee686827dc37/pydantic-2.13.2.tar.gz", hash = "sha256:b418196607e61081c3226dcd4f0672f2a194828abb9109e9cfb84026564df2d1", size = 843836, upload-time = "2026-04-17T09:31:59.636Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/81/5a/2225f4c176dbfed0d809e848b50ef08f70e61daa667b7fa14b0d311ae44d/pydantic-2.13.1-py3-none-any.whl", hash = "sha256:9557ecc2806faaf6037f85b1fbd963d01e30511c48085f0d573650fdeaad378a", size = 471917, upload-time = "2026-04-15T14:57:17.277Z" }, + { url = "https://files.pythonhosted.org/packages/77/ca/b45c378e6e8d0b90577288b533e04e95b7afd61bb1d51b6c263176435489/pydantic-2.13.2-py3-none-any.whl", hash = "sha256:a525087f4c03d7e7456a3de89b64cd693d2229933bb1068b9af6befd5563694e", size = 471947, upload-time = "2026-04-17T09:31:57.541Z" }, ] [[package]] name = "pydantic-core" -version = "2.46.1" +version = "2.46.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a1/93/f97a86a7eb28faa1d038af2fd5d6166418b4433659108a4c311b57128b2d/pydantic_core-2.46.1.tar.gz", hash = "sha256:d408153772d9f298098fb5d620f045bdf0f017af0d5cb6e309ef8c205540caa4", size = 471230, upload-time = "2026-04-15T14:49:34.52Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a2/a0/07f275411355b567b994e565bc5ea9dbf522978060c18e3b7edf646c0fc2/pydantic_core-2.46.1-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:84eb5414871fd0293c38d2075802f95030ff11a92cf2189942bf76fd181af77b", size = 2123782, upload-time = "2026-04-15T14:52:57.172Z" }, - { url = "https://files.pythonhosted.org/packages/ab/71/d027c7de46df5b9287ed6f0ef02346c84d61348326253a4f13695d54d66f/pydantic_core-2.46.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:5c75fb25db086bf504c55730442e471c12bc9bfae817dd359b1a36bc93049d34", size = 1948561, upload-time = "2026-04-15T14:53:12.07Z" }, - { url = "https://files.pythonhosted.org/packages/77/74/cba894bea0d51a3b2dcada9eb3af9c4cfaa271bf21123372dc82ccef029f/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:29dc09f0221425453fd9f73fd70bba15817d25b95858282702d7305a08d37306", size = 1974387, upload-time = "2026-04-15T14:50:14.048Z" }, - { url = "https://files.pythonhosted.org/packages/3b/ad/cc122887d6f20ac5d997928b0bf3016ac9c7bae07dce089333aa0c2e868b/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:139fd6722abc5e6513aa0a27b06ebeb997838c5b179cf5e83862ace45f281c56", size = 2054868, upload-time = "2026-04-15T14:49:51.912Z" }, - { url = "https://files.pythonhosted.org/packages/9f/09/22049b22d65a67253cbdced88dbce0e97162f35cc433917df37df794ede8/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:ba723fd8ef6011af71f92ed54adb604e7699d172f4273e4b46f1cfb8ee8d72fd", size = 2228717, upload-time = "2026-04-15T14:49:27.384Z" }, - { url = "https://files.pythonhosted.org/packages/e6/98/b35a8a187cf977462668b5064c606e290c88c2561e053883d86193ab9c51/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:828410e082555e55da9bbb5e6c17617386fe1415c4d42765a90d372ed9cce813", size = 2298261, upload-time = "2026-04-15T14:52:20.463Z" }, - { url = "https://files.pythonhosted.org/packages/98/ae/46f8d693caefc09d8e2d3f19a6b4f2252cf6542f0b555759f2b5ec2b4ca5/pydantic_core-2.46.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fb5cd53264c9906c163a71b489e9ac71b0ae13a2dd0241e6129f4df38ba1c814", size = 2094496, upload-time = "2026-04-15T14:49:59.711Z" }, - { url = "https://files.pythonhosted.org/packages/ee/40/7e4013639d316d2cb67dae288c768d49cc4a7a4b16ef869e486880db1a1f/pydantic_core-2.46.1-cp310-cp310-manylinux_2_31_riscv64.whl", hash = "sha256:4530a6594883d9d4a9c7ef68464ef6b4a88d839e3531c089a3942c78bffe0a66", size = 2144795, upload-time = "2026-04-15T14:52:44.731Z" }, - { url = "https://files.pythonhosted.org/packages/0d/87/c00f6450059804faf30f568009c8c98e72e6802c1ccd8b562da57953ad81/pydantic_core-2.46.1-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:ed1c71f60abbf9c9a440dc8fc6b1180c45dcab3a5e311250de99744a0166bc95", size = 2173108, upload-time = "2026-04-15T14:51:37.806Z" }, - { url = "https://files.pythonhosted.org/packages/46/15/7a8fb06c109a07dbc1f5f272b2da1290c8a25f5900a579086e433049fc1a/pydantic_core-2.46.1-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:254253491f1b8e3ba18c15fe924bb9b175f1a48413b74e8f0c67b8f51b6f726b", size = 2185687, upload-time = "2026-04-15T14:51:33.125Z" }, - { url = "https://files.pythonhosted.org/packages/d9/38/c52ead78febf23d32db898c7022173c674226cf3c8ee1645220ab9516931/pydantic_core-2.46.1-cp310-cp310-musllinux_1_1_armv7l.whl", hash = "sha256:dfcf6485ac38698a5b45f37467b8eb2f4f8e3edd5790e2579c5d52fdfffb2e3d", size = 2326273, upload-time = "2026-04-15T14:51:10.614Z" }, - { url = "https://files.pythonhosted.org/packages/1e/af/cb5ea2336e9938b3a0536ce4bfed4a342285caa8a6b8ff449a7bc2f179ec/pydantic_core-2.46.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:592b39150ab5b5a2cb2eb885097ee4c2e4d54e3b902f6ae32528f7e6e42c00fc", size = 2368428, upload-time = "2026-04-15T14:49:25.804Z" }, - { url = "https://files.pythonhosted.org/packages/a2/99/adcfbcbd96556120e7d795aab4fd77f5104a49051929c3805a9d736ec48f/pydantic_core-2.46.1-cp310-cp310-win32.whl", hash = "sha256:eb37b1369ad39ec046a36dc81ffd76870766bda2073f57448bbcb1fd3e4c5ad0", size = 1993405, upload-time = "2026-04-15T14:50:51.082Z" }, - { url = "https://files.pythonhosted.org/packages/c4/ff/2767be513a250293f80748740ce73b0f0677711fc791b1afab3499734dd2/pydantic_core-2.46.1-cp310-cp310-win_amd64.whl", hash = "sha256:c330dab8254d422880177436a5892ac6d9337afff9fe383fb1f8c6caedb685e1", size = 2068177, upload-time = "2026-04-15T14:52:29.899Z" }, - { url = "https://files.pythonhosted.org/packages/37/96/d83d23fc3c822326d808b8c0457d4f7afb1552e741a7c2378a974c522c63/pydantic_core-2.46.1-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:f0f84431981c6ae217ebb96c3eca8212f6f5edf116f62f62cc6c7d72971f826c", size = 2121938, upload-time = "2026-04-15T14:49:21.568Z" }, - { url = "https://files.pythonhosted.org/packages/11/44/94b1251825560f5d90e25ebcd457c4772e1f3e1a378f438c040fe2148f3e/pydantic_core-2.46.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:a05f60b36549f59ab585924410187276ec17a94bae939273a213cea252c8471e", size = 1946541, upload-time = "2026-04-15T14:49:57.925Z" }, - { url = "https://files.pythonhosted.org/packages/d6/8f/79aff4c8bd6fb49001ffe4747c775c0f066add9da13dec180eb0023ada34/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b2c93fd1693afdfae7b2897f7530ed3f180d9fc92ee105df3ebdff24d5061cc8", size = 1973067, upload-time = "2026-04-15T14:51:14.765Z" }, - { url = "https://files.pythonhosted.org/packages/56/01/826ab3afb1d43cbfdc2aa592bff0f1f6f4b90f5a801478ba07bde74e706f/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:0c19983759394c702a776f42f33df8d7bb7883aefaa44a69ba86356a9fd67367", size = 2053146, upload-time = "2026-04-15T14:51:48.847Z" }, - { url = "https://files.pythonhosted.org/packages/6c/32/be20ec48ccbd85cac3f8d96ca0a0f87d5c14fbf1eb438da0ac733f2546f2/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:6e8debf586d7d800a718194417497db5126d4f4302885a2dff721e9df3f4851c", size = 2227393, upload-time = "2026-04-15T14:51:53.218Z" }, - { url = "https://files.pythonhosted.org/packages/b5/8e/1fae21c887f363ed1a5cf9f267027700c796b7435313c21723cd3e8aeeb3/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:54160da754d63da7780b76e5743d44f026b9daffc6b8c9696a756368c0a298c9", size = 2296193, upload-time = "2026-04-15T14:50:31.065Z" }, - { url = "https://files.pythonhosted.org/packages/0a/29/e5637b539458ffb60ba9c204fc16c52ea36828427fa667e4f9c7d83cfea9/pydantic_core-2.46.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:74cee962c8b4df9a9b0bb63582e51986127ee2316f0c49143b2996f4b201bd9c", size = 2092156, upload-time = "2026-04-15T14:52:37.227Z" }, - { url = "https://files.pythonhosted.org/packages/bc/fa/3a453934af019c72652fb75489c504ae689de632fa2e037fec3195cd6948/pydantic_core-2.46.1-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:0ba3462872a678ebe21b15bd78eff40298b43ea50c26f230ec535c00cf93ec7e", size = 2142845, upload-time = "2026-04-15T14:51:04.847Z" }, - { url = "https://files.pythonhosted.org/packages/36/c2/71b56fa10a80b98036f4bf0fbb912833f8e9c61b15e66c236fadaf54c27c/pydantic_core-2.46.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:b718873a966d91514c5252775f568985401b54a220919ab22b19a6c4edd8c053", size = 2170756, upload-time = "2026-04-15T14:50:17.16Z" }, - { url = "https://files.pythonhosted.org/packages/e1/da/a4c761dc8d982e2c53f991c0c36d37f6fe308e149bf0a101c25b0750a893/pydantic_core-2.46.1-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:cb1310a9fd722da8cceec1fb59875e1c86bee37f0d8a9c667220f00ee722cc8f", size = 2183579, upload-time = "2026-04-15T14:51:20.888Z" }, - { url = "https://files.pythonhosted.org/packages/e5/d4/b0a6c00622e4afd9a807b8bb05ba8f1a0b69ca068ac138d9d36700fe767b/pydantic_core-2.46.1-cp311-cp311-musllinux_1_1_armv7l.whl", hash = "sha256:98e3ede76eb4b9db8e7b5efea07a3f3315135485794a5df91e3adf56c4d573b6", size = 2324516, upload-time = "2026-04-15T14:52:32.521Z" }, - { url = "https://files.pythonhosted.org/packages/45/f1/a4bace0c98b0774b02de99233882c48d94b399ba4394dd5e209665d05062/pydantic_core-2.46.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:780b8f24ff286e21fd010247011a68ea902c34b1eee7d775b598bc28f5f28ab6", size = 2367084, upload-time = "2026-04-15T14:50:37.832Z" }, - { url = "https://files.pythonhosted.org/packages/3a/54/ae827a3976b136d1c9a9a56c2299a8053605a69facaa0c7354ba167305eb/pydantic_core-2.46.1-cp311-cp311-win32.whl", hash = "sha256:1d452f4cad0f39a94414ca68cda7cc55ff4c3801b5ab0bc99818284a3d39f889", size = 1992061, upload-time = "2026-04-15T14:51:44.704Z" }, - { url = "https://files.pythonhosted.org/packages/55/ae/d85de69e0fdfafc0e87d88bd5d0c157a5443efaaef24eed152a8a8f8dfb6/pydantic_core-2.46.1-cp311-cp311-win_amd64.whl", hash = "sha256:f463fd6a67138d70200d2627676e9efbb0cee26d98a5d3042a35aa20f95ec129", size = 2065497, upload-time = "2026-04-15T14:51:17.077Z" }, - { url = "https://files.pythonhosted.org/packages/46/a7/9eb3b1038db630e1550924e81d1211b0dd70ac3740901fd95f30f5497990/pydantic_core-2.46.1-cp311-cp311-win_arm64.whl", hash = "sha256:155aec0a117140e86775eec113b574c1c299358bfd99467b2ea7b2ea26db2614", size = 2045914, upload-time = "2026-04-15T14:51:24.782Z" }, - { url = "https://files.pythonhosted.org/packages/ce/fb/caaa8ee23861c170f07dbd58fc2be3a2c02a32637693cbb23eef02e84808/pydantic_core-2.46.1-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:ae8c8c5eb4c796944f3166f2f0dab6c761c2c2cc5bd20e5f692128be8600b9a4", size = 2119472, upload-time = "2026-04-15T14:49:45.946Z" }, - { url = "https://files.pythonhosted.org/packages/fa/61/bcffaa52894489ff89e5e1cdde67429914bf083c0db7296bef153020f786/pydantic_core-2.46.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:daba6f5f5b986aa0682623a1a4f8d1ecb0ec00ce09cfa9ca71a3b742bc383e3a", size = 1951230, upload-time = "2026-04-15T14:52:27.646Z" }, - { url = "https://files.pythonhosted.org/packages/f8/95/80d2f43a2a1a1e3220fd329d614aa5a39e0a75d24353a3aaf226e605f1c2/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:0265f3a2460539ecc97817a80c7a23c458dd84191229b655522a2674f701f14e", size = 1976394, upload-time = "2026-04-15T14:50:32.742Z" }, - { url = "https://files.pythonhosted.org/packages/8d/31/2c5b1a207926b5fc1961a2d11da940129bc3841c36cc4df03014195b2966/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:bb16c0156c4b4e94aa3719138cc43c53d30ff21126b6a3af63786dcc0757b56e", size = 2068455, upload-time = "2026-04-15T14:50:01.286Z" }, - { url = "https://files.pythonhosted.org/packages/7d/36/c6aa07274359a51ac62895895325ce90107e811c6cea39d2617a99ef10d7/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:1b42d80fad8e4b283e1e4138f1142f0d038c46d137aad2f9824ad9086080dd41", size = 2239049, upload-time = "2026-04-15T14:53:02.216Z" }, - { url = "https://files.pythonhosted.org/packages/0a/3f/77cdd0db8bddc714842dfd93f737c863751cf02001c993341504f6b0cd53/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9cced85896d5b795293bc36b7e2fb0347a36c828551b50cbba510510d928548c", size = 2318681, upload-time = "2026-04-15T14:50:04.539Z" }, - { url = "https://files.pythonhosted.org/packages/a1/a3/09d929a40e6727274b0b500ad06e1b3f35d4f4665ae1c8ba65acbb17e9b5/pydantic_core-2.46.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a641cb1e74b44c418adaf9f5f450670dbec53511f030d8cde8d8accb66edc363", size = 2096527, upload-time = "2026-04-15T14:53:14.766Z" }, - { url = "https://files.pythonhosted.org/packages/89/ae/544c3a82456ebc254a9fcbe2715bab76c70acf9d291aaea24391147943e4/pydantic_core-2.46.1-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:191e7a122ab14eb12415fe3f92610fc06c7f1d2b4b9101d24d490d447ac92506", size = 2170407, upload-time = "2026-04-15T14:51:27.138Z" }, - { url = "https://files.pythonhosted.org/packages/9d/ce/0dfd881c7af4c522f47b325707bd9a2cdcf4f40e4f2fd30df0e9a3e8d393/pydantic_core-2.46.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4fe4ff660f7938b5d92f21529ce331b011aa35e481ab64b7cd03f52384e544bb", size = 2188578, upload-time = "2026-04-15T14:50:39.655Z" }, - { url = "https://files.pythonhosted.org/packages/a1/e9/980ea2a6d5114dd1a62ecc5f56feb3d34555f33bd11043f042e5f7f0724a/pydantic_core-2.46.1-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:18fcea085b3adc3868d8d19606da52d7a52d8bccd8e28652b0778dbe5e6a6660", size = 2188959, upload-time = "2026-04-15T14:52:42.243Z" }, - { url = "https://files.pythonhosted.org/packages/e7/f1/595e0f50f4bfc56cde2fe558f2b0978f29f2865da894c6226231e17464a5/pydantic_core-2.46.1-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:e8e589e7c9466e022d79e13c5764c2239b2e5a7993ba727822b021234f89b56b", size = 2339973, upload-time = "2026-04-15T14:52:10.642Z" }, - { url = "https://files.pythonhosted.org/packages/49/44/be9f979a6ab6b8c36865ccd92c3a38a760c66055e1f384665f35525134c4/pydantic_core-2.46.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:f78eb3d4027963bdc9baccd177f02a98bf8714bc51fe17153d8b51218918b5bc", size = 2385228, upload-time = "2026-04-15T14:51:00.77Z" }, - { url = "https://files.pythonhosted.org/packages/5b/d4/c826cd711787d240219f01d0d3ca116cb55516b8b95277820aa9c85e1882/pydantic_core-2.46.1-cp312-cp312-win32.whl", hash = "sha256:54fe30c20cab03844dc63bdc6ddca67f74a2eb8482df69c1e5f68396856241be", size = 1978828, upload-time = "2026-04-15T14:50:29.362Z" }, - { url = "https://files.pythonhosted.org/packages/22/05/8a1fcf8181be4c7a9cfc34e5fbf2d9c3866edc9dfd3c48d5401806e0a523/pydantic_core-2.46.1-cp312-cp312-win_amd64.whl", hash = "sha256:aea4e22ed4c53f2774221435e39969a54d2e783f4aee902cdd6c8011415de893", size = 2070015, upload-time = "2026-04-15T14:49:47.301Z" }, - { url = "https://files.pythonhosted.org/packages/61/d5/fea36ad2882b99c174ef4ffbc7ea6523f6abe26060fbc1f77d6441670232/pydantic_core-2.46.1-cp312-cp312-win_arm64.whl", hash = "sha256:f76fb49c34b4d66aa6e552ce9e852ea97a3a06301a9f01ae82f23e449e3a55f8", size = 2030176, upload-time = "2026-04-15T14:50:47.307Z" }, - { url = "https://files.pythonhosted.org/packages/ff/d2/bda39bad2f426cb5078e6ad28076614d3926704196efe0d7a2a19a99025d/pydantic_core-2.46.1-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:cdc8a5762a9c4b9d86e204d555444e3227507c92daba06259ee66595834de47a", size = 2119092, upload-time = "2026-04-15T14:49:50.392Z" }, - { url = "https://files.pythonhosted.org/packages/ee/f3/69631e64d69cb3481494b2bddefe0ddd07771209f74e9106d066f9138c2a/pydantic_core-2.46.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:ba381dfe9c85692c566ecb60fa5a77a697a2a8eebe274ec5e4d6ec15fafad799", size = 1951400, upload-time = "2026-04-15T14:51:06.588Z" }, - { url = "https://files.pythonhosted.org/packages/53/1c/21cb3db6ae997df31be8e91f213081f72ffa641cb45c89b8a1986832b1f9/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1593d8de98207466dc070118322fef68307a0cc6a5625e7b386f6fdae57f9ab6", size = 1976864, upload-time = "2026-04-15T14:50:54.804Z" }, - { url = "https://files.pythonhosted.org/packages/91/9c/05c819f734318ce5a6ca24da300d93696c105af4adb90494ee571303afd8/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8262c74a1af5b0fdf795f5537f7145785a63f9fbf9e15405f547440c30017ed8", size = 2066669, upload-time = "2026-04-15T14:51:42.346Z" }, - { url = "https://files.pythonhosted.org/packages/cb/23/fadddf1c7f2f517f58731aea9b35c914e6005250f08dac9b8e53904cdbaa/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:4b88949a24182e83fbbb3f7ca9b7858d0d37b735700ea91081434b7d37b3b444", size = 2238737, upload-time = "2026-04-15T14:50:45.558Z" }, - { url = "https://files.pythonhosted.org/packages/23/07/0cd4f95cb0359c8b1ec71e89c3777e7932c8dfeb9cd54740289f310aaead/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b8f3708cd55537aeaf3fd0ea55df0d68d0da51dcb07cbc8508745b34acc4c6e0", size = 2316258, upload-time = "2026-04-15T14:51:08.471Z" }, - { url = "https://files.pythonhosted.org/packages/0c/40/6fc24c3766a19c222a0d60d652b78f0283339d4cd4c173fab06b7ee76571/pydantic_core-2.46.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f79292435fff1d4f0c18d9cfaf214025cc88e4f5104bfaed53f173621da1c743", size = 2097474, upload-time = "2026-04-15T14:49:56.543Z" }, - { url = "https://files.pythonhosted.org/packages/4b/af/f39795d1ce549e35d0841382b9c616ae211caffb88863147369a8d74fba9/pydantic_core-2.46.1-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:a2e607aeb59cf4575bb364470288db3b9a1f0e7415d053a322e3e154c1a0802e", size = 2168383, upload-time = "2026-04-15T14:51:29.269Z" }, - { url = "https://files.pythonhosted.org/packages/e6/32/0d563f74582795779df6cc270c3fc220f49f4daf7860d74a5a6cda8491ff/pydantic_core-2.46.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:ec5ca190b75878a9f6ae1fc8f5eb678497934475aef3d93204c9fa01e97370b6", size = 2186182, upload-time = "2026-04-15T14:50:19.097Z" }, - { url = "https://files.pythonhosted.org/packages/5c/07/1c10d5ce312fc4cf86d1e50bdcdbb8ef248409597b099cab1b4bb3a093f7/pydantic_core-2.46.1-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:1f80535259dcdd517d7b8ca588d5ca24b4f337228e583bebedf7a3adcdf5f721", size = 2187859, upload-time = "2026-04-15T14:49:22.974Z" }, - { url = "https://files.pythonhosted.org/packages/92/01/e1f62d4cb39f0913dbf5c95b9b119ef30ddba9493dff8c2b012f0cdd67dc/pydantic_core-2.46.1-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:24820b3c82c43df61eca30147e42853e6c127d8b868afdc0c162df829e011eb4", size = 2338372, upload-time = "2026-04-15T14:49:53.316Z" }, - { url = "https://files.pythonhosted.org/packages/44/ed/218dfeea6127fb1781a6ceca241ec6edf00e8a8933ff331af2215975a534/pydantic_core-2.46.1-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:f12794b1dd8ac9fb66619e0b3a0427189f5d5638e55a3de1385121a9b7bf9b39", size = 2384039, upload-time = "2026-04-15T14:53:04.929Z" }, - { url = "https://files.pythonhosted.org/packages/6c/1e/011e763cd059238249fbd5780e0f8d0b04b47f86c8925e22784f3e5fc977/pydantic_core-2.46.1-cp313-cp313-win32.whl", hash = "sha256:9bc09aed935cdf50f09e908923f9efbcca54e9244bd14a5a0e2a6c8d2c21b4e9", size = 1977943, upload-time = "2026-04-15T14:52:17.969Z" }, - { url = "https://files.pythonhosted.org/packages/8c/06/b559a490d3ed106e9b1777b8d5c8112dd8d31716243cd662616f66c1f8ea/pydantic_core-2.46.1-cp313-cp313-win_amd64.whl", hash = "sha256:fac2d6c8615b8b42bee14677861ba09d56ee076ba4a65cfb9c3c3d0cc89042f2", size = 2068729, upload-time = "2026-04-15T14:53:07.288Z" }, - { url = "https://files.pythonhosted.org/packages/9f/52/32a198946e2e19508532aa9da02a61419eb15bd2d96bab57f810f2713e31/pydantic_core-2.46.1-cp313-cp313-win_arm64.whl", hash = "sha256:f978329f12ace9f3cb814a5e44d98bbeced2e36f633132bafa06d2d71332e33e", size = 2029550, upload-time = "2026-04-15T14:52:22.707Z" }, - { url = "https://files.pythonhosted.org/packages/44/4b/1952d38a091aa7572c13460db4439d5610a524a1a533fb131e17d8eff9c2/pydantic_core-2.46.1-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:c56887c0ffa05318128a80303c95066a9d819e5e66d75ff24311d9e0a58d6930", size = 2123089, upload-time = "2026-04-15T14:50:20.658Z" }, - { url = "https://files.pythonhosted.org/packages/90/06/f3623aa98e2d7cb4ed0ae0b164c5d8a1b86e5aca01744eba980eefcd5da4/pydantic_core-2.46.1-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:614b24b875c1072631065fa85e195b40700586afecb0b27767602007920dacf8", size = 1945481, upload-time = "2026-04-15T14:50:56.945Z" }, - { url = "https://files.pythonhosted.org/packages/69/f9/a9224203b8426893e22db2cf0da27cd930ad7d76e0a611ebd707e5e6c916/pydantic_core-2.46.1-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a6382f6967c48519b6194e9e1e579e5898598b682556260eeaf05910400d827e", size = 1986294, upload-time = "2026-04-15T14:49:31.839Z" }, - { url = "https://files.pythonhosted.org/packages/96/29/954d2174db68b9f14292cef3ae8a05a25255735909adfcf45ca768023713/pydantic_core-2.46.1-graalpy311-graalpy242_311_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:93cb8aa6c93fb833bb53f3a2841fbea6b4dc077453cd5b30c0634af3dee69369", size = 2144185, upload-time = "2026-04-15T14:52:39.449Z" }, - { url = "https://files.pythonhosted.org/packages/f4/97/95de673a1356a88b2efdaa120eb6af357a81555c35f6809a7a1423ff7aef/pydantic_core-2.46.1-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:5f9107a24a4bc00293434dfa95cf8968751ad0dd703b26ea83a75a56f7326041", size = 2107564, upload-time = "2026-04-15T14:50:49.14Z" }, - { url = "https://files.pythonhosted.org/packages/00/fc/a7c16d85211ea9accddc693b7d049f20b0c06440d9264d1e1c074394ee6c/pydantic_core-2.46.1-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:2b1801ba99876984d0a03362782819238141c4d0f3f67f69093663691332fc35", size = 1939925, upload-time = "2026-04-15T14:50:36.188Z" }, - { url = "https://files.pythonhosted.org/packages/2e/23/87841169d77820ddabeb81d82002c95dcb82163846666d74f5bdeeaec750/pydantic_core-2.46.1-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b7fd82a91a20ed6d54fa8c91e7a98255b1ff45bf09b051bfe7fe04eb411e232e", size = 1995313, upload-time = "2026-04-15T14:50:22.538Z" }, - { url = "https://files.pythonhosted.org/packages/ea/96/b46609359a354fa9cd336fc5d93334f1c358b756cc81e4b397347a88fa6f/pydantic_core-2.46.1-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f135bf07c92c93def97008bc4496d16934da9efefd7204e5f22a2c92523cb1f", size = 2151197, upload-time = "2026-04-15T14:51:22.925Z" }, - { url = "https://files.pythonhosted.org/packages/f5/e7/3d1d2999ad8e78b124c752e4fc583ecd98f3bea7cc42045add2fb6e31b62/pydantic_core-2.46.1-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:b44b44537efbff2df9567cd6ba51b554d6c009260a021ab25629c81e066f1683", size = 2121103, upload-time = "2026-04-15T14:52:59.537Z" }, - { url = "https://files.pythonhosted.org/packages/de/08/50a56632994007c7a58c86f782accccbe2f3bb7ca80f462533e26424cd18/pydantic_core-2.46.1-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:8f9ca3af687cc6a5c89aeaa00323222fcbceb4c3cdc78efdac86f46028160c04", size = 1952464, upload-time = "2026-04-15T14:52:04.001Z" }, - { url = "https://files.pythonhosted.org/packages/75/0b/3cf631e33a55b1788add3e42ac921744bd1f39279082a027b4ef6f48bd32/pydantic_core-2.46.1-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e2678a4cbc205f00a44542dca19d15c11ccddd7440fd9df0e322e2cae55bb67a", size = 2138504, upload-time = "2026-04-15T14:52:01.812Z" }, - { url = "https://files.pythonhosted.org/packages/fa/69/f96f3dfc939450b9aeb80d3fe1943e7bc0614b14e9447d84f48d65153e0c/pydantic_core-2.46.1-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:e5a98cbb03a8a7983b0fb954e0af5e7016587f612e6332c6a4453f413f1d1851", size = 2165467, upload-time = "2026-04-15T14:52:15.455Z" }, - { url = "https://files.pythonhosted.org/packages/a8/22/bb61cccddc2ce85b179cd81a580a1746e880870060fbf4bf6024dab7e8aa/pydantic_core-2.46.1-pp311-pypy311_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:b2f098b08860bd149e090ad232f27fffb5ecf1bfd9377015445c8e17355ec2d1", size = 2183882, upload-time = "2026-04-15T14:51:50.868Z" }, - { url = "https://files.pythonhosted.org/packages/0e/01/b9039da255c5fd3a7fd85344fda8861c847ad6d8fdd115580fa4505b2022/pydantic_core-2.46.1-pp311-pypy311_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:d2623606145b55a96efdd181b015c0356804116b2f14d3c2af4832fe4f45ed5f", size = 2323011, upload-time = "2026-04-15T14:49:40.32Z" }, - { url = "https://files.pythonhosted.org/packages/24/b1/f426b20cb72d0235718ccc4de3bc6d6c0d0c2a91a3fd2f32ae11b624bcc9/pydantic_core-2.46.1-pp311-pypy311_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:420f515c42aaec607ff720867b300235bd393abd709b26b190ceacb57a9bfc17", size = 2365696, upload-time = "2026-04-15T14:49:41.936Z" }, - { url = "https://files.pythonhosted.org/packages/ef/d2/d2b0025246481aa2ce6db8ba196e29b92063343ac76e675b3a1fa478ed4d/pydantic_core-2.46.1-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:375cfdd2a1049910c82ba2ff24f948e93599a529e0fdb066d747975ca31fc663", size = 2190970, upload-time = "2026-04-15T14:49:33.111Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/43/bb/4742f05b739b2478459bb16fa8470549518c802e06ddcf3f106c5081315e/pydantic_core-2.46.2.tar.gz", hash = "sha256:37bb079f9ee3f1a519392b73fda2a96379b31f2013c6b467fe693e7f2987f596", size = 471269, upload-time = "2026-04-17T09:10:07.017Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/f2/98f37e836c5ba0335432768e0d8645e6f50a3c838b48a74d9256256784fc/pydantic_core-2.46.2-cp310-cp310-macosx_10_12_x86_64.whl", hash = "sha256:160ef93541f4f84e3e5068e6c1f64d8fd6f57586e5853d609b467d3333f8146a", size = 2108178, upload-time = "2026-04-17T09:10:24.689Z" }, + { url = "https://files.pythonhosted.org/packages/55/69/975458de8e5453322cfc57d6c7029c3e66d9e7a4389c53ddd5ad02d5e5da/pydantic_core-2.46.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1a9124b63f4f40a12a0666df57450b4c24b98407ff74349221b869ec085a5d8e", size = 1949232, upload-time = "2026-04-17T09:11:39.536Z" }, + { url = "https://files.pythonhosted.org/packages/94/8d/938175e6e82d051ac4644765680db06571d7e106a42f760da09bd90f6525/pydantic_core-2.46.2-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:de12004a7da7f1eb67ece37439a5a23a915636085dd042176fda362e006e6940", size = 1974741, upload-time = "2026-04-17T09:13:01.922Z" }, + { url = "https://files.pythonhosted.org/packages/f2/38/7329f8ac5c732bddf15f939c2add40b95170e0ecca5ef124c12def3f78ba/pydantic_core-2.46.2-cp310-cp310-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a070c7769fec277409ad0b3d55b2f0a3703a6f00cf5031fe93090f155bf56382", size = 2041905, upload-time = "2026-04-17T09:11:11.94Z" }, + { url = "https://files.pythonhosted.org/packages/99/2c/47cfd069937ee5cbc0d9e18fa9795c8f80c49a6b4fc777d4cd870f2ade7b/pydantic_core-2.46.2-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:41d701bb34f81f0b11c724cc544b9a10b26a28f4d0d1197f2037c91225708706", size = 2222703, upload-time = "2026-04-17T09:10:31.196Z" }, + { url = "https://files.pythonhosted.org/packages/83/b0/7ed83ca8cd92c99bcab90cf42ed953723fbc19d8a20c8c12bb68c51febc1/pydantic_core-2.46.2-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:19631e7350b7a574fb6b6db222f4b17e8bd31803074b3307d07df62379d2b2e4", size = 2276317, upload-time = "2026-04-17T09:09:53.263Z" }, + { url = "https://files.pythonhosted.org/packages/85/70/50b1b62990996e7916aae2852b29cbf3ecc3fdae78209eb284cd61e2c918/pydantic_core-2.46.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:48b1059e4f2a6ec3e41983148eb1eec5ef9fa3a80bbc4ac0893ac76b115fe039", size = 2092152, upload-time = "2026-04-17T09:10:44.683Z" }, + { url = "https://files.pythonhosted.org/packages/c1/51/a062864e6b34ada7e343ad9ed29368e495620a8ef1c009b47a68b46e1634/pydantic_core-2.46.2-cp310-cp310-manylinux_2_31_riscv64.whl", hash = "sha256:df73724fce8ad53c670358c905b37930bd7b9d92e57db640a65c53b2706eee00", size = 2118091, upload-time = "2026-04-17T09:10:05.083Z" }, + { url = "https://files.pythonhosted.org/packages/07/e0/fcc97c4d0319615dc0b5b132b420904639652f8514e9c76482acb70ea1d4/pydantic_core-2.46.2-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a0891a9be0def16fb320af21a198ece052eed72bf44d73d8ff43f702bd26fd6b", size = 2174304, upload-time = "2026-04-17T09:11:00.54Z" }, + { url = "https://files.pythonhosted.org/packages/00/52/28f53796ca74b7e3dd45938f300517f04970e985ad600d0d0f36a11378bd/pydantic_core-2.46.2-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:2ca790779aa1cba1329b8dc42ccebada441d9ac1d932de980183d544682c646d", size = 2181444, upload-time = "2026-04-17T09:11:45.442Z" }, + { url = "https://files.pythonhosted.org/packages/22/49/164d5d3a7356d2607a72e77264a3b252a7c7d9362a81fc9df47bef7ae3aa/pydantic_core-2.46.2-cp310-cp310-musllinux_1_1_armv7l.whl", hash = "sha256:6b865eb702c3af71cf7331919a787563ce2413f7a54ef49ec6709a01b4f22ce6", size = 2328611, upload-time = "2026-04-17T09:10:08.574Z" }, + { url = "https://files.pythonhosted.org/packages/6b/77/6266bb3b79c27b533e5ee02c1e3da5848872112178880cc5006a84e857ac/pydantic_core-2.46.2-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:631bec5f951a30a4b332b4a57d0cdd5a2c8187eb71301f966425f2e54a697855", size = 2351070, upload-time = "2026-04-17T09:13:34.92Z" }, + { url = "https://files.pythonhosted.org/packages/10/7f/d4233852d16d8e85b034a524d8017e051a0aa4acd04c64c3a69a1a2a0ba6/pydantic_core-2.46.2-cp310-cp310-win32.whl", hash = "sha256:8cbd9d67357f3a925f2af1d44db3e8ef1ce1a293ea0add98081b072d4a12e3b4", size = 1976750, upload-time = "2026-04-17T09:13:15.537Z" }, + { url = "https://files.pythonhosted.org/packages/70/31/d65117cf5f89d81705da5b1dcdad8efa0a0b65dbbc7f13cafbabb7d01615/pydantic_core-2.46.2-cp310-cp310-win_amd64.whl", hash = "sha256:dd51dd16182b4bfdcefd27b39b856aa4a57b77f15b231a2d10c45391b0a02028", size = 2073989, upload-time = "2026-04-17T09:12:17.315Z" }, + { url = "https://files.pythonhosted.org/packages/89/91/089f517a725f29084364169437833ab0ae4da4d7a6ed9d4474db7f1412e6/pydantic_core-2.46.2-cp311-cp311-macosx_10_12_x86_64.whl", hash = "sha256:d8060f42db3cd204871db0afd51fef54a13fa544c4dd48cdcae2e174ef40c8ba", size = 2106218, upload-time = "2026-04-17T09:10:48.023Z" }, + { url = "https://files.pythonhosted.org/packages/a0/92/23858ed1b58f2a134e50c2fdd0e34ea72721ccb257e1e9346514e1ccb5b9/pydantic_core-2.46.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:73a9d2809bd8d4a7cda4d336dc996a565eb4feaaa39932f9d85a65fa18382f28", size = 1948087, upload-time = "2026-04-17T09:11:58.639Z" }, + { url = "https://files.pythonhosted.org/packages/5d/ac/e2240fccb4794e965817593d5a46cf5ea22f2001b73fe360b7578925b7d8/pydantic_core-2.46.2-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3b0a2dee92dfaabcfb93629188c3e9cf74fdfc0f22e7c369cb444a98814a1e50", size = 1972931, upload-time = "2026-04-17T09:13:13.304Z" }, + { url = "https://files.pythonhosted.org/packages/1a/da/3b11dab2aa15c5c8ed20a01eb7aa432a78b8e3a4713659f7e58490a020a5/pydantic_core-2.46.2-cp311-cp311-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3098446ba8cf774f61cb8d4008c1dba14a30426a15169cd95ac3392a461193b1", size = 2040454, upload-time = "2026-04-17T09:13:47.895Z" }, + { url = "https://files.pythonhosted.org/packages/d7/39/c4cf5e1f1c6c34c53c0902039c95d81dc15cdd1f03634bd1a93f33e70a72/pydantic_core-2.46.2-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:57c584af6c375ea3f826d8131a94cb212b3d9926eaff67117e3711bbff3a83a5", size = 2221320, upload-time = "2026-04-17T09:13:08.568Z" }, + { url = "https://files.pythonhosted.org/packages/c7/46/891035bc9e93538e754c3188424d24b5a69ec3ae5210fa01d483e99b3302/pydantic_core-2.46.2-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:547381cca999be88b4715a0ed7afa11f07fc7e53cb1883687b190d25a92c56cf", size = 2274559, upload-time = "2026-04-17T09:11:10.257Z" }, + { url = "https://files.pythonhosted.org/packages/ab/d0/7af0b905b3148152c159c9caf203e7ecd9b90b76389f0862e6ab0cf1b2a3/pydantic_core-2.46.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:caeed15dcb1233a5a94bc6ff37ef5393cf5b33a45e4bdfb2d6042f3d24e1cb27", size = 2089239, upload-time = "2026-04-17T09:13:06.326Z" }, + { url = "https://files.pythonhosted.org/packages/c5/bc/566afe02ba2de37712eece74ac7bfba322abd7916410bf90504f1b17ddad/pydantic_core-2.46.2-cp311-cp311-manylinux_2_31_riscv64.whl", hash = "sha256:c05f53362568c75476b5c96659377a5dfd982cfbe5a5c07de5106d08a04efc4f", size = 2116182, upload-time = "2026-04-17T09:11:33.738Z" }, + { url = "https://files.pythonhosted.org/packages/4e/5b/3fcb3a229bbfa23b0e3c65014057af0f9d51ec7a2d9f7adb282f41ff5ac8/pydantic_core-2.46.2-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:2643ac7eae296200dbd48762a1c852cf2cad5f5e3eba34e652053cebf03becf8", size = 2172346, upload-time = "2026-04-17T09:10:46.472Z" }, + { url = "https://files.pythonhosted.org/packages/43/9a/baa9e3aa70ea7bbcb9db0f87162a371649ac80c03e43eb54af193390cf17/pydantic_core-2.46.2-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:dc4620a47c6fe6a39f89392c00833a82fc050ce90169798f78a25a8d4df03b6e", size = 2179540, upload-time = "2026-04-17T09:11:21.881Z" }, + { url = "https://files.pythonhosted.org/packages/bd/46/912047a5427f949c909495704b3c8b9ead9d1c66f87e96606011beab1fcb/pydantic_core-2.46.2-cp311-cp311-musllinux_1_1_armv7l.whl", hash = "sha256:78cb0d2453b50bf2035f85fd0d9cfabdb98c47f9c53ddb7c23873cd83da9560b", size = 2327423, upload-time = "2026-04-17T09:13:40.291Z" }, + { url = "https://files.pythonhosted.org/packages/e9/bf/c5e661451dc9411c2ab88a244c1ba57644950c971486040dc200f77b69f4/pydantic_core-2.46.2-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:f0c1cbb7d6112932cc188c6be007a5e2867005a069e47f42fe67bf5f122b0908", size = 2348652, upload-time = "2026-04-17T09:10:37.76Z" }, + { url = "https://files.pythonhosted.org/packages/77/b3/3219e7c522af54b010cf7422dcb11cc6616a4414d1ccd628b0d3f61c6af6/pydantic_core-2.46.2-cp311-cp311-win32.whl", hash = "sha256:c1ce5b2366f85cfdbf7f0907755043707f86d09a5b1b1acebbb7bf1600d75c64", size = 1974410, upload-time = "2026-04-17T09:13:27.392Z" }, + { url = "https://files.pythonhosted.org/packages/e5/29/e5cfac8a74c59873dfd47d3a1477c39ad9247639a7120d3e251a9ff12417/pydantic_core-2.46.2-cp311-cp311-win_amd64.whl", hash = "sha256:f1a6197eadff5bd0bb932f12bb038d403cb75db5b0b391e70e816a647745ddaf", size = 2071158, upload-time = "2026-04-17T09:09:57.69Z" }, + { url = "https://files.pythonhosted.org/packages/6f/8b/b7b19b717cdb3675cb109de143f62d4dc62f5d4a0b9879b6f1ace62c6654/pydantic_core-2.46.2-cp311-cp311-win_arm64.whl", hash = "sha256:15e42885b283f87846ee79e161002c5c496ef747a73f6e47054f45a13d9035bc", size = 2043507, upload-time = "2026-04-17T09:09:51.828Z" }, + { url = "https://files.pythonhosted.org/packages/97/ec/2fafa4c86f5d2a69372c7cddef30925fd0e370b1efaf556609c1a0196d8a/pydantic_core-2.46.2-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:ea1ad8c89da31512fe2d249cf0638fb666925bda341901541bc5f3311c6fcc9e", size = 2101729, upload-time = "2026-04-17T09:12:30.042Z" }, + { url = "https://files.pythonhosted.org/packages/cf/55/be5386c2c4b49af346e8a26b748194ff25757bbb6cf544130854e997af7a/pydantic_core-2.46.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b308da17b92481e0587244631c5529e5d91d04cb2b08194825627b1eca28e21e", size = 1951546, upload-time = "2026-04-17T09:10:10.585Z" }, + { url = "https://files.pythonhosted.org/packages/29/92/89e273a055ce440e6636c756379af35ad86da9d336a560049c3ba5e41c80/pydantic_core-2.46.2-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d333a50bdd814a917d8d6a7ee35ba2395d53ddaa882613bc24e54a9d8b129095", size = 1976178, upload-time = "2026-04-17T09:11:49.619Z" }, + { url = "https://files.pythonhosted.org/packages/91/b3/e4664469cf70c0cb0f7b2f5719d64e5968bb6f38217042c2afa3d3c4ba17/pydantic_core-2.46.2-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:1d00b99590c5bd1fabbc5d28b170923e32c1b1071b1f1de1851a4d14d89eb192", size = 2051697, upload-time = "2026-04-17T09:12:04.917Z" }, + { url = "https://files.pythonhosted.org/packages/98/58/dbf68213ee06ce51cdd6d8c95f97980e646858c45bd96bd2dfb40433be73/pydantic_core-2.46.2-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9f0e686960ffe9e65066395af856ac2d52c159043144433602c50c221d81c1ba", size = 2233160, upload-time = "2026-04-17T09:12:00.956Z" }, + { url = "https://files.pythonhosted.org/packages/f5/d3/68092aa0ee6c60ff4de4740eb82db3d4ce338ec89b3cecb978c532472f12/pydantic_core-2.46.2-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2d1128da41c9cb474e0a4701f9c363ec645c9d1a02229904c76bf4e0a194fde2", size = 2298398, upload-time = "2026-04-17T09:10:29.694Z" }, + { url = "https://files.pythonhosted.org/packages/e4/51/5d6155eb737db55b0ad354ca5f333ef009f75feb67df2d79a84bace45af6/pydantic_core-2.46.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:48649cf2d8c358d79586e9fb2f8235902fcaa2d969ec1c5301f2d1873b2f8321", size = 2094058, upload-time = "2026-04-17T09:12:10.995Z" }, + { url = "https://files.pythonhosted.org/packages/6b/f3/eb4a986197d71319430464ff181226c95adc8f06d932189b158bae5a82f5/pydantic_core-2.46.2-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:b902f0fc7c2cf503865a05718b68147c6cd5d0a3867af38c527be574a9fa6e9d", size = 2130388, upload-time = "2026-04-17T09:12:41.159Z" }, + { url = "https://files.pythonhosted.org/packages/56/00/44a9c4fe6d0f64b5786d6a8c649d6f0e34ba6c89b3663add1066e54451a2/pydantic_core-2.46.2-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:e80011f808b03d1d87a8f1e76ae3da19a18eb706c823e17981dcf1fae43744fc", size = 2184245, upload-time = "2026-04-17T09:12:36.532Z" }, + { url = "https://files.pythonhosted.org/packages/78/6b/685b98a834d5e3d1c34a1bde1627525559dd223b75075bc7490cdb24eb33/pydantic_core-2.46.2-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b839d5c802e31348b949b6473f8190cddbf7d47475856d8ac995a373ee16ec59", size = 2186842, upload-time = "2026-04-17T09:13:04.054Z" }, + { url = "https://files.pythonhosted.org/packages/22/64/caa2f5a2ac8b6113adaa410ccdf31ba7f54897a6e54cd0d726fc7e780c88/pydantic_core-2.46.2-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:c6b1064f3f9cf9072e1d59dd2936f9f3b668bec1c37039708c9222db703c0d5b", size = 2336066, upload-time = "2026-04-17T09:12:13.006Z" }, + { url = "https://files.pythonhosted.org/packages/ee/f9/7d2701bf82945b5b9e7df8347be97ef6a36da2846bfe5b4afec299ffe27b/pydantic_core-2.46.2-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:37a68e6f2ac95578ce3c0564802404b27b24988649616e556c07e77111ed3f1d", size = 2363691, upload-time = "2026-04-17T09:13:42.972Z" }, + { url = "https://files.pythonhosted.org/packages/3b/65/0dab11574101522941055109419db3cc09db871643dc3fc74e2413215e5b/pydantic_core-2.46.2-cp312-cp312-win32.whl", hash = "sha256:d9ffa75a7ef4b97d6e5e205fabd4304ef01fec09e6f1bdde04b9ad1b07d20289", size = 1958801, upload-time = "2026-04-17T09:11:31.981Z" }, + { url = "https://files.pythonhosted.org/packages/13/2b/df84baa609c676f6450b8ecad44ea59146c805e3371b7b52443c0899f989/pydantic_core-2.46.2-cp312-cp312-win_amd64.whl", hash = "sha256:0551f2d2ddb68af5a00e26497f8025c538f73ef3cb698f8e5a487042cd2792a8", size = 2072634, upload-time = "2026-04-17T09:11:02.407Z" }, + { url = "https://files.pythonhosted.org/packages/d1/4e/e1ce8029fc438086a946739bf9d596f70ff470aad4a8345555920618cabe/pydantic_core-2.46.2-cp312-cp312-win_arm64.whl", hash = "sha256:83aef30f106edcc21a6a4cc44b82d3169a1dbe255508db788e778f3c804d3583", size = 2026188, upload-time = "2026-04-17T09:13:11.083Z" }, + { url = "https://files.pythonhosted.org/packages/07/2b/662e48254479a2d3450ba24b1e25061108b64339794232f503990c519144/pydantic_core-2.46.2-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:d26e9eea3715008a09a74585fe9becd0c67fbb145dc4df9756d597d7230a652c", size = 2101762, upload-time = "2026-04-17T09:10:13.87Z" }, + { url = "https://files.pythonhosted.org/packages/73/ab/bafd7c7503757ccc8ec4d1911e106fe474c629443648c51a88f08b0fe91a/pydantic_core-2.46.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:48b36e3235140510dc7861f0cd58b714b1cdd3d48f75e10ce52e69866b746f10", size = 1951814, upload-time = "2026-04-17T09:12:25.934Z" }, + { url = "https://files.pythonhosted.org/packages/92/cc/7549c2d57ba2e9a42caa5861a2d398dbe31c02c6aca783253ace59ce84f8/pydantic_core-2.46.2-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:36b1f99dc451f1a3981f236151465bcf995bbe712d0727c9f7b236fe228a8133", size = 1977329, upload-time = "2026-04-17T09:13:37.605Z" }, + { url = "https://files.pythonhosted.org/packages/18/50/7ed4a8a0d478a4dca8f0134a5efa7193f03cc8520dd4c9509339fb2e5002/pydantic_core-2.46.2-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8641c8d535c2d95b45c2e19b646ecd23ebba35d461e0ae48a3498277006250ab", size = 2051832, upload-time = "2026-04-17T09:12:49.771Z" }, + { url = "https://files.pythonhosted.org/packages/dc/16/bb35b193741c0298ddc5f5e4234269efdc0c65e2bcd198aa0de9b68845e4/pydantic_core-2.46.2-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:20fb194788a0a50993e87013e693494ba183a2af5b44e99cf060bbae10912b11", size = 2233127, upload-time = "2026-04-17T09:11:04.449Z" }, + { url = "https://files.pythonhosted.org/packages/91/a5/98f4b637149185addea19e1785ea20c373cca31b202f589111d8209d9873/pydantic_core-2.46.2-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9262d11d0cd11ee3303a95156939402bed6cedfe5ed0e331b95a283a4da6eb8b", size = 2297418, upload-time = "2026-04-17T09:11:25.929Z" }, + { url = "https://files.pythonhosted.org/packages/36/90/93a5d21990b152da7b7507b7fddb0b935f6a0984d57ac3ec45a6e17777a2/pydantic_core-2.46.2-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ac204542736aa295fa25f713b7fad6fc50b46ab7764d16087575c85f085174f3", size = 2093735, upload-time = "2026-04-17T09:12:06.908Z" }, + { url = "https://files.pythonhosted.org/packages/14/22/b8b1ffdddf08b4e84380bcb67f41dbbf4c171377c1d36fc6290794bb2094/pydantic_core-2.46.2-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:9a7c43a0584742dface3ca0daf6f719d46c1ac2f87cf080050f9ae052c75e1b2", size = 2127570, upload-time = "2026-04-17T09:11:53.906Z" }, + { url = "https://files.pythonhosted.org/packages/c6/26/e60d72b4e2d0ce1fa811044a974412ac1c567fe067d97b3e6b290530786e/pydantic_core-2.46.2-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fd05e1edb6a90ad446fa268ab09e59202766b837597b714b2492db11ee87fab9", size = 2183524, upload-time = "2026-04-17T09:11:30.092Z" }, + { url = "https://files.pythonhosted.org/packages/35/32/36bec7584a1eefb17dec4dfa1c946d3fe4440f466c5705b8adfda69c9a9f/pydantic_core-2.46.2-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:91155b110788b5501abc7ea954f1d08606219e4e28e3c73a94124307c06efb80", size = 2185408, upload-time = "2026-04-17T09:10:57.228Z" }, + { url = "https://files.pythonhosted.org/packages/fc/d6/1a5689d873620efd67d6b163db0c444c056adb0849b5bc33e2b9f09665a6/pydantic_core-2.46.2-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:e4e2c72a529fa03ff228be1d2b76944013f428220b764e03cc50ada67e17a42c", size = 2335171, upload-time = "2026-04-17T09:11:43.369Z" }, + { url = "https://files.pythonhosted.org/packages/3e/8e/675104802abe8ef502b072050ee5f2e915251aa1a3af87e1015ce31ec42d/pydantic_core-2.46.2-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:56291ec1a11c3499890c99a8fd9053b47e60fe837a77ec72c0671b1b8b3dce24", size = 2362743, upload-time = "2026-04-17T09:10:18.333Z" }, + { url = "https://files.pythonhosted.org/packages/8d/bc/86c5dde4fa6e24467680eef5047da3c1a19be0a527d0d8e14aa76b39307c/pydantic_core-2.46.2-cp313-cp313-win32.whl", hash = "sha256:b50f9c5f826ddca1246f055148df939f5f3f2d0d96db73de28e2233f22210d4c", size = 1958074, upload-time = "2026-04-17T09:12:38.622Z" }, + { url = "https://files.pythonhosted.org/packages/2a/97/2537e8c1282b2c4eb062580c0d7a4339e10b072b803d1ee0b7f1f0a5c22c/pydantic_core-2.46.2-cp313-cp313-win_amd64.whl", hash = "sha256:251a57788823230ca8cbc99e6245d1a2ed6e180ec4864f251c94182c580c7f2e", size = 2071741, upload-time = "2026-04-17T09:13:32.405Z" }, + { url = "https://files.pythonhosted.org/packages/da/aa/2ee75798706f9dbc4e76dbe59e41a396c5c311e3d6223b9cf6a5fa7780be/pydantic_core-2.46.2-cp313-cp313-win_arm64.whl", hash = "sha256:315d32d1a71494d6b4e1e14a9fa7a4329597b4c4340088ad7e1a9dafbeed92a9", size = 2025955, upload-time = "2026-04-17T09:10:15.567Z" }, + { url = "https://files.pythonhosted.org/packages/25/ec/e91aa08df1c33d5e3c2b60c07a1eca9f21809728a824c7b467bb3bda68b5/pydantic_core-2.46.2-graalpy311-graalpy242_311_native-macosx_10_12_x86_64.whl", hash = "sha256:7c5a5b3dbb9e8918e223be6580da5ffcf861c0505bbc196ebed7176ce05b7b4e", size = 2105046, upload-time = "2026-04-17T09:10:55.614Z" }, + { url = "https://files.pythonhosted.org/packages/f0/73/27112400a0452e375290e7c40aef5cc9844ac0920fb1029238cfc68121fa/pydantic_core-2.46.2-graalpy311-graalpy242_311_native-macosx_11_0_arm64.whl", hash = "sha256:bc1e8ce33d5a337f2ba862e0719b8201cd54aaed967406c748e009191d47efdd", size = 1940029, upload-time = "2026-04-17T09:12:21.5Z" }, + { url = "https://files.pythonhosted.org/packages/b1/44/3d39f782bc82ddd0b2d82bde83b408aa40a332cdf6f3018acb34e3d4dcfc/pydantic_core-2.46.2-graalpy311-graalpy242_311_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b737c0b280f41143266445de2689c0e49c79307e51c44ce3a77fef2bedad4994", size = 1987772, upload-time = "2026-04-17T09:10:02.357Z" }, + { url = "https://files.pythonhosted.org/packages/c4/1a/0242e5b7b6cf51dbccc065029f0420107b6bf7e191fcb918f5cb71218acf/pydantic_core-2.46.2-graalpy311-graalpy242_311_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1b877d597afb82b4898e35354bba55de6f7f048421ae0edadbb9886ec137b532", size = 2138468, upload-time = "2026-04-17T09:11:51.546Z" }, + { url = "https://files.pythonhosted.org/packages/f3/d2/66c146f421178641bda880b0267c0d57dd84f5fec9ecc8e46be17b480742/pydantic_core-2.46.2-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:e9fcabd1857492b5bf16f90258babde50f618f55d046b1309972da2396321ff9", size = 2091621, upload-time = "2026-04-17T09:12:47.501Z" }, + { url = "https://files.pythonhosted.org/packages/ee/b2/c28419aa9fc8055f4ac8e801d1d11c6357351bfa4321ed9bafab3eb98087/pydantic_core-2.46.2-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:fb3ec2c7f54c07b30d89983ce78dc32c37dd06a972448b8716d609493802d628", size = 1937059, upload-time = "2026-04-17T09:10:53.554Z" }, + { url = "https://files.pythonhosted.org/packages/30/ce/cd0824a2db213dc17113291b7a09b9b0ccd9fbf97daa4b81548703341baf/pydantic_core-2.46.2-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:130a6c837d819ef33e8c2bf702ed2c3429237ea69807f1140943d6f4bdaf52fa", size = 1997278, upload-time = "2026-04-17T09:12:23.784Z" }, + { url = "https://files.pythonhosted.org/packages/c9/69/47283fe3c0c967d3e9e9cd6c42b70907610c8a6f8d6e8381f1bb55f8006c/pydantic_core-2.46.2-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:c2e25417cec5cd9bddb151e33cb08c50160f317479ecc02b22a95ec18f8fe004", size = 2147096, upload-time = "2026-04-17T09:12:43.124Z" }, + { url = "https://files.pythonhosted.org/packages/16/d5/dec7c127fa722ff56e1ccf1e960ae1318a9f66742135e97bf9771447216f/pydantic_core-2.46.2-pp311-pypy311_pp73-macosx_10_12_x86_64.whl", hash = "sha256:c3ad79ed32004d9de91cacd4b5faaff44d56051392fe1d5526feda596f01af25", size = 2107613, upload-time = "2026-04-17T09:10:36.269Z" }, + { url = "https://files.pythonhosted.org/packages/bc/35/975c109b337260a71c93198baf663982b6b39fe3e584e279548a0969e5d4/pydantic_core-2.46.2-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:d157c48d28eebe5d46906de06a6a2f2c9e00b67d3e42de1f1b9c2d42b810f77c", size = 1947099, upload-time = "2026-04-17T09:12:15.304Z" }, + { url = "https://files.pythonhosted.org/packages/4e/11/52a971a0f9218631690274be533f05e5ddde5547f0823bb3e9dfd1be49f6/pydantic_core-2.46.2-pp311-pypy311_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7b42c6471288dedc979ac8400d9c9770f03967dd187db1f8d3405d4d182cc714", size = 2133866, upload-time = "2026-04-17T09:12:27.994Z" }, + { url = "https://files.pythonhosted.org/packages/fe/7a/33d94d0698602b2d1712e78c703a33952eb2ca69e02e8e4b208e7f6602b5/pydantic_core-2.46.2-pp311-pypy311_pp73-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:4f27bc4801358dc070d6697b41237fce9923d8e69a1ce1e95606ac36c1552dc1", size = 2161721, upload-time = "2026-04-17T09:11:16.111Z" }, + { url = "https://files.pythonhosted.org/packages/b0/cb/0df7ee0a148e9ce0968a80787967ddca9f6b3f8a49152a881b88da262701/pydantic_core-2.46.2-pp311-pypy311_pp73-musllinux_1_1_aarch64.whl", hash = "sha256:e094a8f85db41aa7f6a45c5dac2950afc9862e66832934231962252b5d284eed", size = 2180175, upload-time = "2026-04-17T09:11:41.577Z" }, + { url = "https://files.pythonhosted.org/packages/8e/a8/258a32878140347532be4e44c6f3b1ace3b52b9c9ca7548a65ce18adf4b4/pydantic_core-2.46.2-pp311-pypy311_pp73-musllinux_1_1_armv7l.whl", hash = "sha256:807eeda5551f6884d3b4421578be37be50ddb7a58832348e99617a6714a73748", size = 2319882, upload-time = "2026-04-17T09:10:21.872Z" }, + { url = "https://files.pythonhosted.org/packages/13/b9/5071c298a0f91314a5402b8c56e0efbcebe77085327d0b4df7dc9cb0b674/pydantic_core-2.46.2-pp311-pypy311_pp73-musllinux_1_1_x86_64.whl", hash = "sha256:fcaa1c3c846a7f6686b38fe493d1b2e8007380e293bfef6a9354563c026cbf36", size = 2348065, upload-time = "2026-04-17T09:11:08.263Z" }, + { url = "https://files.pythonhosted.org/packages/75/f3/0a7087e5f861d66ca64ce927230b397cc264c87b712156e6a93b26a459c8/pydantic_core-2.46.2-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:154dbfdfb11b8cbd8ff4d00d0b81e3d19f4cb4bedd5aa9f091060ba071474c6a", size = 2192159, upload-time = "2026-04-17T09:11:20.123Z" }, ] [[package]] @@ -3901,14 +3804,14 @@ name = "torch" version = "2.10.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "filelock", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "fsspec", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "jinja2", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.11' and sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "setuptools", marker = "(python_full_version >= '3.12' and platform_machine != 's390x' and sys_platform == 'darwin') or (python_full_version >= '3.12' and sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "sympy", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, - { name = "typing-extensions", marker = "(platform_machine != 's390x' and sys_platform == 'darwin') or (sys_platform != 'darwin' and sys_platform != 'win32')" }, + { name = "filelock" }, + { name = "fsspec" }, + { name = "jinja2" }, + { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "setuptools", marker = "python_full_version >= '3.12'" }, + { name = "sympy" }, + { name = "typing-extensions" }, ] [[package]] @@ -3949,7 +3852,7 @@ wheels = [ [[package]] name = "torchvision" -version = "0.25.0" +version = "0.26.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -3958,26 +3861,26 @@ dependencies = [ { name = "torch", marker = "sys_platform == 'never'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/50/ae/cbf727421eb73f1cf907fbe5788326a08f111b3f6b6ddca15426b53fec9a/torchvision-0.25.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a95c47abb817d4e90ea1a8e57bd0d728e3e6b533b3495ae77d84d883c4d11f56", size = 1874919, upload-time = "2026-01-21T16:27:47.617Z" }, - { url = "https://files.pythonhosted.org/packages/64/68/dc7a224f606d53ea09f9a85196a3921ec3a801b0b1d17e84c73392f0c029/torchvision-0.25.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:acc339aba4a858192998c2b91f635827e40d9c469d9cf1455bafdda6e4c28ea4", size = 2343220, upload-time = "2026-01-21T16:27:44.26Z" }, - { url = "https://files.pythonhosted.org/packages/f9/fa/8cce5ca7ffd4da95193232493703d20aa06303f37b119fd23a65df4f239a/torchvision-0.25.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:0d9a3f925a081dd2ebb0b791249b687c2ef2c2717d027946654607494b9b64b6", size = 8068106, upload-time = "2026-01-21T16:27:37.805Z" }, - { url = "https://files.pythonhosted.org/packages/8b/b9/a53bcf8f78f2cd89215e9ded70041765d50ef13bf301f9884ec6041a9421/torchvision-0.25.0-cp310-cp310-win_amd64.whl", hash = "sha256:b57430fbe9e9b697418a395041bb615124d9c007710a2712fda6e35fb310f264", size = 3697295, upload-time = "2026-01-21T16:27:36.574Z" }, - { url = "https://files.pythonhosted.org/packages/3e/be/c704bceaf11c4f6b19d64337a34a877fcdfe3bd68160a8c9ae9bea4a35a3/torchvision-0.25.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:db74a551946b75d19f9996c419a799ffdf6a223ecf17c656f90da011f1d75b20", size = 1874923, upload-time = "2026-01-21T16:27:46.574Z" }, - { url = "https://files.pythonhosted.org/packages/ae/e9/f143cd71232430de1f547ceab840f68c55e127d72558b1061a71d0b193cd/torchvision-0.25.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:f49964f96644dbac2506dffe1a0a7ec0f2bf8cf7a588c3319fed26e6329ffdf3", size = 2344808, upload-time = "2026-01-21T16:27:43.191Z" }, - { url = "https://files.pythonhosted.org/packages/43/ae/ad5d6165797de234c9658752acb4fce65b78a6a18d82efdf8367c940d8da/torchvision-0.25.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:153c0d2cbc34b7cf2da19d73450f24ba36d2b75ec9211b9962b5022fb9e4ecee", size = 8070752, upload-time = "2026-01-21T16:27:33.748Z" }, - { url = "https://files.pythonhosted.org/packages/23/19/55b28aecdc7f38df57b8eb55eb0b14a62b470ed8efeb22cdc74224df1d6a/torchvision-0.25.0-cp311-cp311-win_amd64.whl", hash = "sha256:ea580ffd6094cc01914ad32f8c8118174f18974629af905cea08cb6d5d48c7b7", size = 4038722, upload-time = "2026-01-21T16:27:41.355Z" }, - { url = "https://files.pythonhosted.org/packages/56/3a/6ea0d73f49a9bef38a1b3a92e8dd455cea58470985d25635beab93841748/torchvision-0.25.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c2abe430c90b1d5e552680037d68da4eb80a5852ebb1c811b2b89d299b10573b", size = 1874920, upload-time = "2026-01-21T16:27:45.348Z" }, - { url = "https://files.pythonhosted.org/packages/51/f8/c0e1ef27c66e15406fece94930e7d6feee4cb6374bbc02d945a630d6426e/torchvision-0.25.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:b75deafa2dfea3e2c2a525559b04783515e3463f6e830cb71de0fb7ea36fe233", size = 2344556, upload-time = "2026-01-21T16:27:40.125Z" }, - { url = "https://files.pythonhosted.org/packages/68/2f/f24b039169db474e8688f649377de082a965fbf85daf4e46c44412f1d15a/torchvision-0.25.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:f25aa9e380865b11ea6e9d99d84df86b9cc959f1a007cd966fc6f1ab2ed0e248", size = 8072351, upload-time = "2026-01-21T16:27:21.074Z" }, - { url = "https://files.pythonhosted.org/packages/ad/16/8f650c2e288977cf0f8f85184b90ee56ed170a4919347fc74ee99286ed6f/torchvision-0.25.0-cp312-cp312-win_amd64.whl", hash = "sha256:f9c55ae8d673ab493325d1267cbd285bb94d56f99626c00ac4644de32a59ede3", size = 4303059, upload-time = "2026-01-21T16:27:11.08Z" }, - { url = "https://files.pythonhosted.org/packages/f5/5b/1562a04a6a5a4cf8cf40016a0cdeda91ede75d6962cff7f809a85ae966a5/torchvision-0.25.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:24e11199e4d84ba9c5ee7825ebdf1cd37ce8deec225117f10243cae984ced3ec", size = 1874918, upload-time = "2026-01-21T16:27:39.02Z" }, - { url = "https://files.pythonhosted.org/packages/36/b1/3d6c42f62c272ce34fcce609bb8939bdf873dab5f1b798fd4e880255f129/torchvision-0.25.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5f271136d2d2c0b7a24c5671795c6e4fd8da4e0ea98aeb1041f62bc04c4370ef", size = 2309106, upload-time = "2026-01-21T16:27:30.624Z" }, - { url = "https://files.pythonhosted.org/packages/c7/60/59bb9c8b67cce356daeed4cb96a717caa4f69c9822f72e223a0eae7a9bd9/torchvision-0.25.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:855c0dc6d37f462482da7531c6788518baedca1e0847f3df42a911713acdfe52", size = 8071522, upload-time = "2026-01-21T16:27:29.392Z" }, - { url = "https://files.pythonhosted.org/packages/32/a5/9a9b1de0720f884ea50dbf9acb22cbe5312e51d7b8c4ac6ba9b51efd9bba/torchvision-0.25.0-cp313-cp313-win_amd64.whl", hash = "sha256:cef0196be31be421f6f462d1e9da1101be7332d91984caa6f8022e6c78a5877f", size = 4321911, upload-time = "2026-01-21T16:27:35.195Z" }, - { url = "https://files.pythonhosted.org/packages/52/99/dca81ed21ebaeff2b67cc9f815a20fdaa418b69f5f9ea4c6ed71721470db/torchvision-0.25.0-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:a8f8061284395ce31bcd460f2169013382ccf411148ceb2ee38e718e9860f5a7", size = 1896209, upload-time = "2026-01-21T16:27:32.159Z" }, - { url = "https://files.pythonhosted.org/packages/28/cc/2103149761fdb4eaed58a53e8437b2d716d48f05174fab1d9fcf1e2a2244/torchvision-0.25.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:146d02c9876858420adf41f3189fe90e3d6a409cbfa65454c09f25fb33bf7266", size = 2310735, upload-time = "2026-01-21T16:27:22.327Z" }, - { url = "https://files.pythonhosted.org/packages/76/ad/f4c985ad52ddd3b22711c588501be1b330adaeaf6850317f66751711b78c/torchvision-0.25.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:c4d395cb2c4a2712f6eb93a34476cdf7aae74bb6ea2ea1917f858e96344b00aa", size = 8089557, upload-time = "2026-01-21T16:27:27.666Z" }, - { url = "https://files.pythonhosted.org/packages/63/cc/0ea68b5802e5e3c31f44b307e74947bad5a38cc655231d845534ed50ddb8/torchvision-0.25.0-cp313-cp313t-win_amd64.whl", hash = "sha256:5e6b449e9fa7d642142c0e27c41e5a43b508d57ed8e79b7c0a0c28652da8678c", size = 4344260, upload-time = "2026-01-21T16:27:17.018Z" }, + { url = "https://files.pythonhosted.org/packages/74/b4/cdfee31e0402ea035135462cb0ab496e974d56fab6b4e7a1f0cbccb8cd28/torchvision-0.26.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:a06d4772a8e13e772906ed736cc53ec6639e5e60554f8e5fa6ca165aabebc464", size = 1863503, upload-time = "2026-03-23T18:13:01.384Z" }, + { url = "https://files.pythonhosted.org/packages/e4/74/11fee109841e80ad14e5ca2d80bff6b10eb11b7838ff06f35bfeaa9f7251/torchvision-0.26.0-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:2adfbe438473236191ff077a4a9a0c767436879c89628aa97137e959b0c11a94", size = 7766423, upload-time = "2026-03-23T18:12:56.049Z" }, + { url = "https://files.pythonhosted.org/packages/5e/00/24d8c7845c3f270153fb81395a5135b2778e2538e81d14c6aea5106c689c/torchvision-0.26.0-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:b6f9ad1ecc0eab52647298b379ee9426845f8903703e6127973f8f3d049a798b", size = 7518249, upload-time = "2026-03-23T18:12:51.743Z" }, + { url = "https://files.pythonhosted.org/packages/d7/ed/e53cd7c0da7ae002e5e929c1796ebbe7ec0c700c29f7a0a6696497fb3d8b/torchvision-0.26.0-cp310-cp310-win_amd64.whl", hash = "sha256:f13f12b3791a266de2d599cb8162925261622a037d87fc03132848343cf68f75", size = 3669784, upload-time = "2026-03-23T18:12:49.949Z" }, + { url = "https://files.pythonhosted.org/packages/b4/bd/d552a2521bade3295b2c6e7a4a0d1022261cab7ca7011f4e2a330dbb3caa/torchvision-0.26.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:55bd6ad4ae77be01ba67a410b05b51f53b0d0ee45f146eb6a0dfb9007e70ab3c", size = 1863499, upload-time = "2026-03-23T18:12:58.696Z" }, + { url = "https://files.pythonhosted.org/packages/33/bf/21b899792b08cae7a298551c68398a79e333697479ed311b3b067aab4bdc/torchvision-0.26.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:1c55dc8affbcc0eb2060fbabbe996ae9e5839b24bb6419777f17848945a411b1", size = 7767527, upload-time = "2026-03-23T18:12:44.348Z" }, + { url = "https://files.pythonhosted.org/packages/9a/45/57bbf9e216850d065e66dd31a50f57424b607f1d878ab8956e56a1f4e36b/torchvision-0.26.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:fd10b5f994c210f4f6d6761cf686f82d748554adf486cb0979770c3252868c8f", size = 7519925, upload-time = "2026-03-23T18:12:53.283Z" }, + { url = "https://files.pythonhosted.org/packages/10/58/ed8f7754299f3e91d6414b6dc09f62b3fa7c6e5d63dfe48d69ab81498a37/torchvision-0.26.0-cp311-cp311-win_amd64.whl", hash = "sha256:de6424b12887ad884f39a0ee446994ae3cd3b6a00a9cafe1bead85a031132af0", size = 3983834, upload-time = "2026-03-23T18:13:00.224Z" }, + { url = "https://files.pythonhosted.org/packages/ae/e7/56b47cc3b132aea90ccce22bcb8975dec688b002150012acc842846039d0/torchvision-0.26.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:c409e1c3fdebec7a3834465086dbda8bf7680eff79abf7fd2f10c6b59520a7a4", size = 1863502, upload-time = "2026-03-23T18:12:57.326Z" }, + { url = "https://files.pythonhosted.org/packages/f4/ec/5c31c92c08b65662fe9604a4067ae8232582805949f11ddc042cebe818ed/torchvision-0.26.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:406557718e62fdf10f5706e88d8a5ec000f872da913bf629aab9297622585547", size = 7767944, upload-time = "2026-03-23T18:12:42.805Z" }, + { url = "https://files.pythonhosted.org/packages/f5/d8/cb6ccda1a1f35a6597645818641701207b3e8e13553e75fce5d86bac74b2/torchvision-0.26.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:d61a5abb6b42a0c0c311996c2ac4b83a94418a97182c83b055a2a4ae985e05aa", size = 7522205, upload-time = "2026-03-23T18:12:54.654Z" }, + { url = "https://files.pythonhosted.org/packages/1c/a9/c272623a0f735c35f0f6cd6dc74784d4f970e800cf063bb76687895a2ab9/torchvision-0.26.0-cp312-cp312-win_amd64.whl", hash = "sha256:7993c01648e7c61d191b018e84d38fe0825c8fcb2720cd0f37caf7ba14404aa1", size = 4255155, upload-time = "2026-03-23T18:12:32.652Z" }, + { url = "https://files.pythonhosted.org/packages/da/80/0762f77f53605d10c9477be39bb47722cc8e383bbbc2531471ce0e396c07/torchvision-0.26.0-cp313-cp313-macosx_12_0_arm64.whl", hash = "sha256:5d63dd43162691258b1b3529b9041bac7d54caa37eae0925f997108268cbf7c4", size = 1860809, upload-time = "2026-03-23T18:12:47.629Z" }, + { url = "https://files.pythonhosted.org/packages/e6/81/0b3e58d1478c660a5af4268713486b2df7203f35abd9195fea87348a5178/torchvision-0.26.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:a39c7a26538c41fda453f9a9692b5ff9b35a5437db1d94f3027f6f509c160eac", size = 7727494, upload-time = "2026-03-23T18:12:46.062Z" }, + { url = "https://files.pythonhosted.org/packages/b6/dc/d9ab5d29115aa05e12e30f1397a3eeae1d88a511241dc3bce48dc4342675/torchvision-0.26.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:b7e6213620bbf97742e5f79832f9e9d769e6cf0f744c5b53dad80b76db633691", size = 7521747, upload-time = "2026-03-23T18:12:36.815Z" }, + { url = "https://files.pythonhosted.org/packages/a9/1b/f1bc86a918c5f6feab1eeff11982e2060f4704332e96185463d27855bdf5/torchvision-0.26.0-cp313-cp313-win_amd64.whl", hash = "sha256:4280c35ec8cba1fcc8294fb87e136924708726864c379e4c54494797d86bc474", size = 4319880, upload-time = "2026-03-23T18:12:38.168Z" }, + { url = "https://files.pythonhosted.org/packages/66/28/b4ad0a723ed95b003454caffcc41894b34bd8379df340848cae2c33871de/torchvision-0.26.0-cp313-cp313t-macosx_12_0_arm64.whl", hash = "sha256:358fc4726d0c08615b6d83b3149854f11efb2a564ed1acb6fce882e151412d23", size = 1951973, upload-time = "2026-03-23T18:12:48.781Z" }, + { url = "https://files.pythonhosted.org/packages/71/e2/7a89096e6cf2f3336353b5338ba925e0addf9d8601920340e6bdf47e8eb3/torchvision-0.26.0-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:3daf9cc149cf3cdcbd4df9c59dae69ffca86c6823250442c3bbfd63fc2e26c61", size = 7728679, upload-time = "2026-03-23T18:12:26.196Z" }, + { url = "https://files.pythonhosted.org/packages/69/1d/4e1eebc17d18ce080a11dcf3df3f8f717f0efdfa00983f06e8ba79259f61/torchvision-0.26.0-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:82c3965eca27e86a316e31e4c3e5a16d353e0bcbe0ef8efa2e66502c54493c4b", size = 7609138, upload-time = "2026-03-23T18:12:35.327Z" }, + { url = "https://files.pythonhosted.org/packages/f3/a4/f1155e943ae5b32400d7000adc81c79bb0392b16ceb33bcf13e02e48cced/torchvision-0.26.0-cp313-cp313t-win_amd64.whl", hash = "sha256:ebc043cc5a4f0bf22e7680806dbba37ffb19e70f6953bbb44ed1a90aeb5c9bea", size = 4248202, upload-time = "2026-03-23T18:12:41.423Z" }, ] [[package]] From c20f9c411d236fcba99845efe5baa9efd4f6ac37 Mon Sep 17 00:00:00 2001 From: kaix-nv Date: Sun, 19 Apr 2026 05:44:17 +0100 Subject: [PATCH 21/30] Add a standalone monitor skill for persistent job tracking (#1252) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: ? Add a standalone monitor skill for persistent job tracking across sessions, and integrate it with PTQ, evaluation, and deployment skills. Problem: Each skill had ad-hoc inline monitoring (squeue polling, nel status checks) that didn't survive session restarts and couldn't track multiple jobs. Users had to manually ask "check status" every time. Solution: A centralized monitor skill with: - Job registry (.claude/active_jobs.json): single source of truth for all active jobs - Durable recurring cron: polls every 15 min, survives session restarts, self-cleans when all jobs complete - User-initiated mode: works in new conversations by reading the registry - Aggregated reporting: "2 of 4 completed" instead of per-job noise ### Usage After any skill submits a job, the monitor skill automatically: 1. Registers the job in .claude/active_jobs.json 2. Sets up a durable cron to poll status every 15 minutes User can also trigger manually: User: "check my eval status" → reads registry, reports current state User: "is the PTQ done?" → finds job, checks status User: "what jobs are running?" → lists all registered jobs ### Testing ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A - Did you write any new necessary tests?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit * **New Features** * Added monitor skill for tracking SLURM jobs, NEL evaluations, and launcher experiments with persistent job registry. * **Documentation** * Updated deployment, evaluation, and PTQ documentation to use the new monitor skill. * Simplified diagnostic and troubleshooting instructions. --------- Signed-off-by: Kai Xu --- .claude/skills/deployment/SKILL.md | 2 +- .claude/skills/evaluation/SKILL.md | 72 +++++--------------- .claude/skills/monitor/SKILL.md | 102 +++++++++++++++++++++++++++++ .claude/skills/ptq/SKILL.md | 4 +- 4 files changed, 120 insertions(+), 60 deletions(-) create mode 100644 .claude/skills/monitor/SKILL.md diff --git a/.claude/skills/deployment/SKILL.md b/.claude/skills/deployment/SKILL.md index 6f3f9b56bde..5210eae6c3c 100644 --- a/.claude/skills/deployment/SKILL.md +++ b/.claude/skills/deployment/SKILL.md @@ -195,7 +195,7 @@ If a cluster config exists (`~/.config/modelopt/clusters.yaml` or `.claude/clust 3. **Deploy based on remote environment:** - - **SLURM** — see `skills/common/slurm-setup.md` for job script templates (container setup, account/partition discovery). The server command inside the container is the same as Step 4 (e.g., `python -m vllm.entrypoints.openai.api_server --model --quantization modelopt`). Use `remote_submit_job` and `remote_poll_job` to manage the job. Get the node hostname from `squeue -j $JOBID -o %N`. + - **SLURM** — see `skills/common/slurm-setup.md` for job script templates (container setup, account/partition discovery). The server command inside the container is the same as Step 4 (e.g., `python -m vllm.entrypoints.openai.api_server --model --quantization modelopt`). After submitting, register the job and set up monitoring per the **monitor skill**. Get the node hostname from `squeue -j $JOBID -o %N`. - **Bare metal / Docker** — use `remote_run` to start the server directly: diff --git a/.claude/skills/evaluation/SKILL.md b/.claude/skills/evaluation/SKILL.md index 41a59c59422..29b4790c85a 100644 --- a/.claude/skills/evaluation/SKILL.md +++ b/.claude/skills/evaluation/SKILL.md @@ -256,64 +256,24 @@ After the dry-run, check the output from `nel` for any problems with the config. **Monitoring Progress** -After job submission, you can monitor progress using: +After job submission, register the job per the **monitor skill** for durable cross-session tracking. For one-off queries (live status, debugging a failed run, analyzing results) use the **launching-evals skill**; for querying past runs in MLflow use **accessing-mlflow**. -1. **Check job status:** +**NEL-specific diagnostics** (for debugging failures): - ```bash - nel status - nel info - ``` - -2. **Stream logs** (Local execution only): - - ```bash - nel logs - ``` - - Note: `nel logs` is not supported for SLURM execution. - -3. **Inspect logs via SSH** (SLURM workaround): - - When `nel logs` is unavailable (SLURM), use SSH to inspect logs directly: - - First, get log locations: - - ```bash - nel info --logs - ``` - - Then, use SSH to view logs: - - **Check server deployment logs:** - - ```bash - ssh @ "tail -100 --logs`>/server--*.log" - ``` - - Shows vLLM server startup, model loading, and deployment errors (e.g., missing wget/curl). - - **Check evaluation client logs:** - - ```bash - ssh @ "tail -100 --logs`>/client-.log" - ``` - - Shows evaluation progress, task execution, and results. - - **Check SLURM scheduler logs:** - - ```bash - ssh @ "tail -100 --logs`>/slurm-.log" - ``` - - Shows job scheduling, health checks, and overall execution flow. - - **Search for errors:** - - ```bash - ssh @ "grep -i 'error\|warning\|failed' --logs`>/*.log" - ``` +```bash +# Quick status check +nel status +nel info + +# Get log paths +nel info --logs + +# Inspect logs via SSH +ssh @ "tail -100 /server--*.log" # deployment errors +ssh @ "tail -100 /client-.log" # evaluation errors +ssh @ "tail -100 /slurm-.log" # scheduling/walltime +ssh @ "grep -i 'error\|failed' /*.log" # search all logs +``` --- diff --git a/.claude/skills/monitor/SKILL.md b/.claude/skills/monitor/SKILL.md new file mode 100644 index 00000000000..cd896c347e9 --- /dev/null +++ b/.claude/skills/monitor/SKILL.md @@ -0,0 +1,102 @@ +--- +name: monitor +description: Monitor submitted jobs (PTQ, evaluation, deployment) on SLURM clusters. Use when the user asks "check job status", "is my job done", "monitor my evaluation", "what's the status of the PTQ", "check on job ", or after any skill submits a long-running job. Also triggers on "nel status", "squeue", or any request to check progress of a previously submitted job. +--- + +# Job Monitor + +Monitor jobs submitted to SLURM clusters — PTQ quantization, NEL evaluation, model deployment, or raw SLURM jobs. + +## When to use + +1. **Auto-monitor** — another skill (PTQ, evaluation, deployment) just submitted a job. Register the job and set up monitoring immediately. +2. **User-initiated** — user asks about a job status, possibly in a new conversation. Check the registry, identify the job, and report. + +--- + +## Job Registry + +All active jobs are tracked in `.claude/active_jobs.json`. This file is the single source of truth for what's being monitored. + +```json +[ + { + "type": "nel", + "id": "", + "host": "", + "user": "", + "submitted": "YYYY-MM-DD HH:MM", + "description": "", + "last_status": "" + } +] +``` + +`type` is one of: `nel`, `slurm`, `launcher`. + +--- + +## On Job Submission + +Every time a job is submitted (by any skill or manually): + +1. **Add an entry** to `.claude/active_jobs.json`. Create the file if it doesn't exist. +2. **Set up a durable recurring cron** (if one isn't already running) that polls all registered jobs every 15 minutes. The cron prompt should: read the registry, check each job, report state changes to the user, remove completed jobs, and delete itself when the registry is empty. + +Always do both steps. Don't try to predict job duration. + +--- + +## On Cron Fire / Status Check + +Whether triggered by the cron or by the user asking "check status": + +1. **Read the registry** from `.claude/active_jobs.json` +2. **Check each job** using the appropriate method (see below) +3. **Report only state changes** — compare against `last_status` in registry +4. **Update `last_status`** in the registry +5. **Remove completed jobs** — any job in a terminal state (COMPLETED, FAILED, CANCELLED, KILLED) +6. **If registry is empty** — delete the recurring cron + +--- + +## How to Check Each Job Type + +### NEL jobs (`type: nel`) + +- **Check:** `nel status ` +- **On completion:** `nel info ` to fetch results +- **On failure:** `nel info --logs` then inspect server/client/SLURM logs via SSH + +### Launcher jobs (`type: launcher`) + +- **Check:** Tail the launcher's background output file for key events +- **Key events:** experiment ID, SLURM job ID, container import, calibration progress, export path, final status +- **On failure:** Look for `Traceback`, `Error`, or `FAILED` in the output + +### Raw SLURM jobs (`type: slurm`) + +- **Check:** `ssh "squeue -j -h -o '%T %M %R'"` — if empty, job left the queue +- **On completion:** `ssh "sacct -j --format=State,ExitCode,Elapsed -n"` +- **On failure:** Check the job's output log file + +--- + +## Identifying Jobs (user-initiated, no ID given) + +When the user asks about a job without specifying an ID, check in order: + +1. `.claude/active_jobs.json` — most reliable, has context +2. `nel ls runs --since 1d` — recent NEL runs +3. `ssh "squeue -u "` — active SLURM jobs +4. `ls -lt tools/launcher/experiments/cicd/ | head -10` — recent launcher experiments + +--- + +## Reporting Guidelines + +- **Report state changes proactively** — PENDING → RUNNING, or job completes +- **Aggregate multiple jobs** — "2 of 4 completed (MMLU-Pro: 42.3%, GSM8K: 67.1%), 1 running, 1 pending" +- **Summarize, don't echo** — interpret events ("Calibration complete, exporting checkpoint") not raw logs +- **On failure, diagnose immediately** — check logs and report root cause without waiting for user to ask +- **Minimize noise** — don't report "still running" unless the user is actively asking diff --git a/.claude/skills/ptq/SKILL.md b/.claude/skills/ptq/SKILL.md index c4c70651a8f..b2b3be1d3fa 100644 --- a/.claude/skills/ptq/SKILL.md +++ b/.claude/skills/ptq/SKILL.md @@ -118,9 +118,7 @@ For SLURM, see `skills/common/slurm-setup.md` and `references/slurm-setup-ptq.md ### Monitoring -- **Launcher**: blocks and tails logs automatically -- **SLURM (manual)**: poll with `squeue -u $USER` + `sleep` (not cron or background tasks) -- **Local**: watch stdout +After job submission, register the job and set up monitoring per the **monitor skill**. ## Step 5 — Verify output From 26ae8da51756eb081da09ea1b5253855815e022f Mon Sep 17 00:00:00 2001 From: jingyu-ml <108295447+jingyu-ml@users.noreply.github.com> Date: Sat, 18 Apr 2026 23:50:14 -0700 Subject: [PATCH 22/30] [2/3] Implicit Gemm NVFP4 (#1227) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: new feature - Add Conv3D implicit GEMM kernel with BF16 WMMA tensor cores and fused NVFP4 activation quantization for video diffusion VAE layers - Integrate into _QuantConv3d via QuantModuleRegistry — automatically dispatched when NVFP4 quantization is applied to nn.Conv3d - Move kernel from `experimental/conv/ to modelopt/torch/kernels/conv/`; move tests to `tests/gpu/torch/quantization/kernels/` ### Testing - Added test cases to measure the difference between cuDNN and our CUDA implicit GEMM kernel - Added an NVFP4 fake quantization test using CUDA code ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ ### Additional Information ## Summary by CodeRabbit * **New Features** * Per-backbone quantization/export in a single run with per-backbone checkpoints and backbone-aware quant filters * Configurable NVFP4 block-size via CLI/config; improved NVFP4 Conv3D inference path and Wan 2.2 quantization support * **Bug Fixes** * Video-model calibration now respects extra params and forces video decoding during calibration * **Documentation** * Added comprehensive Conv3D implicit‑GEMM kernel documentation; removed experimental Conv3D prototype docs/benchmark * **Tests** * New Wan 2.2 quantization/export tests and expanded Conv3D/FP4 kernel test coverage --------- Signed-off-by: Jingyu Xin --- .github/codecov.yml | 12 - CHANGELOG.rst | 1 + examples/diffusers/README.md | 14 + .../diffusers/quantization/calibration.py | 13 +- .../diffusers/quantization/models_utils.py | 45 +- .../quantization/pipeline_manager.py | 85 ++- examples/diffusers/quantization/quantize.py | 136 ++-- .../diffusers/quantization/quantize_config.py | 3 +- examples/diffusers/quantization/utils.py | 18 +- experimental/conv/README.md | 105 ---- experimental/conv/bench_implicit_gemm.py | 208 ------- .../quantization/nn/modules/quant_conv.py | 79 ++- .../plugins/diffusion/diffusers.py | 35 ++ .../torch/quantization/src/conv/README.md | 141 +++++ .../src}/conv/implicit_gemm_binding.cpp | 0 .../src}/conv/implicit_gemm_cuda.py | 0 .../src}/conv/implicit_gemm_kernel.cu | 5 + tests/examples/diffusers/conftest.py | 31 + tests/examples/diffusers/test_diffusers.py | 113 ++++ .../test_export_diffusers_hf_ckpt.py | 81 +++ .../kernels}/test_implicit_gemm.py | 585 +++++++++++++++--- .../plugins/test_diffusers_wan_conv3d.py | 129 ++++ .../torch/quantization/test_quant_conv.py | 141 +++++ 23 files changed, 1458 insertions(+), 522 deletions(-) delete mode 100644 experimental/conv/README.md delete mode 100644 experimental/conv/bench_implicit_gemm.py create mode 100644 modelopt/torch/quantization/src/conv/README.md rename {experimental => modelopt/torch/quantization/src}/conv/implicit_gemm_binding.cpp (100%) rename {experimental => modelopt/torch/quantization/src}/conv/implicit_gemm_cuda.py (100%) rename {experimental => modelopt/torch/quantization/src}/conv/implicit_gemm_kernel.cu (98%) create mode 100644 tests/examples/diffusers/conftest.py rename {experimental/conv => tests/gpu/torch/quantization/kernels}/test_implicit_gemm.py (64%) create mode 100644 tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py diff --git a/.github/codecov.yml b/.github/codecov.yml index b4ac8367690..24756fdcbb2 100644 --- a/.github/codecov.yml +++ b/.github/codecov.yml @@ -11,15 +11,3 @@ coverage: target: auto threshold: 1% # Allow atmost 1% coverage drop from main branch. patch: false - -# Exclude GPU-only Triton kernel files from ALL codecov calculations (project -# and patch checks, all flags). Rationale: these files are dominated by -# @triton.jit kernel bodies that CPU unit tests cannot exercise. GPU tests -# cover them end-to-end (see tests/gpu/torch/sparsity/attention_sparsity/) but -# the `gpu`-flag upload may race with the PR status check, so relying on flag -# combination alone leaves the project check flaky. Dropping these files here -# makes the check deterministic — local `pytest --cov` and GPU runs still -# measure them; only the codecov PR status ignores them. -ignore: - - "modelopt/torch/kernels/triton_fa.py" - - "modelopt/torch/kernels/hf_triton_attention.py" diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 80dea0e43e4..20a677d0a0b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -16,6 +16,7 @@ Changelog - Add support for vLLM fakequant reload using ModelOpt state for HF models. See `examples/vllm_serve/README.md `_ for more details. - [Early Testing] Add Claude Code PTQ skill (``.claude/skills/ptq/``) for agent-assisted post-training quantization. The skill guides the agent through environment detection, model support checking, format selection, and execution via the launcher or manual SLURM/Docker/bare GPU paths. Includes handling for unlisted models with custom module patching. This feature is in early testing — use with caution. - Add performant layerwise calibration for large models that don't fit on GPU (e.g. DeepSeek-R1, Kimi-K2). See `modelopt_recipes/general/ptq/nvfp4_experts_only-fp8_kv.yaml `_ for usage. Layerwise calibration also supports PTQ with intermediate progress saving — useful when long PTQ runs get hit with Slurm timeouts. See `modelopt_recipes/general/ptq/nvfp4_default-none_kv_gptq.yaml `_ for usage. +- Add implicit GEMM CUDA kernel for Conv3D with fused NVFP4 fake quantization (``modelopt.torch.quantization.src.conv``). When NVFP4 quantization is applied to an ``nn.Conv3d`` layer via ModelOpt PTQ, the implicit GEMM path is used automatically instead of cuDNN. Uses BF16 WMMA tensor cores (SM80+) with FP32 accumulation and in-kernel FP4 (E2M1) activation quantization. Grouped convolution (``groups > 1``) falls back to the default cuDNN path. Inference only — training mode falls back to cuDNN with a warning. **Backward Breaking Changes** diff --git a/examples/diffusers/README.md b/examples/diffusers/README.md index 84f248bfb15..ac14d982279 100644 --- a/examples/diffusers/README.md +++ b/examples/diffusers/README.md @@ -117,6 +117,20 @@ python quantize.py \ --hf-ckpt-dir ./hf_ckpt ``` +#### Wan 2.2 VAE NVFP4 (Conv3D Implicit GEMM) + +The Wan 2.2 VAE (`AutoencoderKLWan`, shared between the 5B and 14B pipelines) is built from 3D convolutions. When quantizing the VAE with NVFP4, the `Conv3d` layers are automatically dispatched through a custom BF16 WMMA implicit-GEMM kernel with fused FP4 activation quantization. Requires SM80+ (Ampere or newer). See [`modelopt/torch/quantization/src/conv/README.md`](../../modelopt/torch/quantization/src/conv/README.md) for kernel details. + +```sh +python quantize.py \ + --model {wan2.2-t2v-14b|wan2.2-t2v-5b} \ + --backbone vae \ + --format fp4 --quant-algo max --collect-method default \ + --model-dtype BFloat16 --trt-high-precision-dtype BFloat16 \ + --batch-size 1 --calib-size 32 --n-steps 30 \ + --quantized-torch-ckpt-save-path ./wan22_vae_fp4.pt +``` + #### [LTX-2](https://github.com/Lightricks/LTX-2) FP4 > [!WARNING] diff --git a/examples/diffusers/quantization/calibration.py b/examples/diffusers/quantization/calibration.py index 7464216cc83..27b1ec22436 100644 --- a/examples/diffusers/quantization/calibration.py +++ b/examples/diffusers/quantization/calibration.py @@ -108,11 +108,12 @@ def run_calibration(self, batched_prompts: list[list[str]]) -> None: def _run_wan_video_calibration( self, prompt_batch: list[str], extra_args: dict[str, Any] ) -> None: + extra_params = self.pipeline_manager.config.extra_params kwargs = {} kwargs["negative_prompt"] = extra_args["negative_prompt"] - kwargs["height"] = extra_args["height"] - kwargs["width"] = extra_args["width"] - kwargs["num_frames"] = extra_args["num_frames"] + kwargs["height"] = extra_params.get("height", extra_args["height"]) + kwargs["width"] = extra_params.get("width", extra_args["width"]) + kwargs["num_frames"] = extra_params.get("num_frames", extra_args["num_frames"]) kwargs["guidance_scale"] = extra_args["guidance_scale"] if "guidance_scale_2" in extra_args: kwargs["guidance_scale_2"] = extra_args["guidance_scale_2"] @@ -154,7 +155,11 @@ def _run_ltx2_calibration(self, prompt_batch: list[str], extra_args: dict[str, A "images": extra_params.get("images", []), "tiling_config": extra_params.get("tiling_config", TilingConfig.default()), } - self.pipe(prompt=prompt, **kwargs) + decoded_video, decoded_audio = self.pipe(prompt=prompt, **kwargs) + # vae_decode_video returns a lazy generator — consume it so the + # video decoder's forward() actually runs during calibration. + for _ in decoded_video: + pass def _run_ltx_video_calibration( self, prompt_batch: list[str], extra_args: dict[str, Any] diff --git a/examples/diffusers/quantization/models_utils.py b/examples/diffusers/quantization/models_utils.py index 1e90c547148..b59744282f6 100644 --- a/examples/diffusers/quantization/models_utils.py +++ b/examples/diffusers/quantization/models_utils.py @@ -33,7 +33,9 @@ from utils import ( filter_func_default, filter_func_flux_dev, + filter_func_ltx2_vae, filter_func_ltx_video, + filter_func_wan_vae, filter_func_wan_video, ) @@ -54,31 +56,30 @@ class ModelType(str, Enum): WAN22_T2V_5b = "wan2.2-t2v-5b" -def get_model_filter_func(model_type: ModelType) -> Callable[[str], bool]: - """ - Get the appropriate filter function for a given model type. +_FILTER_FUNC_MAP: dict[ModelType, Callable[[str], bool]] = { + ModelType.FLUX_DEV: filter_func_flux_dev, + ModelType.FLUX2_DEV: filter_func_flux_dev, + ModelType.LTX_VIDEO_DEV: filter_func_ltx_video, + ModelType.LTX2: filter_func_ltx_video, + ModelType.WAN22_T2V_14b: filter_func_wan_video, + ModelType.WAN22_T2V_5b: filter_func_wan_video, +} - Args: - model_type: The model type enum +_VAE_FILTER_FUNC_MAP: dict[tuple[ModelType, str], Callable[[str], bool]] = { + (ModelType.LTX2, "video_decoder"): filter_func_ltx2_vae, + (ModelType.WAN22_T2V_14b, "vae"): filter_func_wan_vae, + (ModelType.WAN22_T2V_5b, "vae"): filter_func_wan_vae, +} - Returns: - A filter function appropriate for the model type - """ - filter_func_map = { - ModelType.FLUX_DEV: filter_func_flux_dev, - ModelType.FLUX_SCHNELL: filter_func_default, - ModelType.FLUX2_DEV: filter_func_flux_dev, - ModelType.SDXL_BASE: filter_func_default, - ModelType.SDXL_TURBO: filter_func_default, - ModelType.SD3_MEDIUM: filter_func_default, - ModelType.SD35_MEDIUM: filter_func_default, - ModelType.LTX_VIDEO_DEV: filter_func_ltx_video, - ModelType.LTX2: filter_func_ltx_video, - ModelType.WAN22_T2V_14b: filter_func_wan_video, - ModelType.WAN22_T2V_5b: filter_func_wan_video, - } - return filter_func_map.get(model_type, filter_func_default) +def get_model_filter_func( + model_type: ModelType, backbone_name: str = "transformer" +) -> Callable[[str], bool]: + """Get the appropriate filter function for a given model type and backbone.""" + vae_func = _VAE_FILTER_FUNC_MAP.get((model_type, backbone_name)) + if vae_func is not None: + return vae_func + return _FILTER_FUNC_MAP.get(model_type, filter_func_default) # Model registry with HuggingFace model IDs diff --git a/examples/diffusers/quantization/pipeline_manager.py b/examples/diffusers/quantization/pipeline_manager.py index 45f2fb63987..f62aeffca98 100644 --- a/examples/diffusers/quantization/pipeline_manager.py +++ b/examples/diffusers/quantization/pipeline_manager.py @@ -42,6 +42,7 @@ def __init__(self, config: ModelConfig, logger: logging.Logger): self.pipe: Any | None = None self.pipe_upsample: LTXLatentUpsamplePipeline | None = None # For LTX-Video upsampling self._transformer: torch.nn.Module | None = None + self._video_decoder: torch.nn.Module | None = None @staticmethod def create_pipeline_from( @@ -58,23 +59,20 @@ def create_pipeline_from( Raises: ValueError: If model type is unsupported """ - try: - pipeline_cls = MODEL_PIPELINE[model_type] - if pipeline_cls is None: - raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.") - model_id = ( - MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path - ) - pipe = pipeline_cls.from_pretrained( - model_id, - torch_dtype=torch_dtype, - use_safetensors=True, - **MODEL_DEFAULTS[model_type].get("from_pretrained_extra_args", {}), - ) - pipe.set_progress_bar_config(disable=True) - return pipe - except Exception as e: - raise e + pipeline_cls = MODEL_PIPELINE[model_type] + if pipeline_cls is None: + raise ValueError(f"Model type {model_type.value} does not use diffusers pipelines.") + model_id = ( + MODEL_REGISTRY[model_type] if override_model_path is None else override_model_path + ) + pipe = pipeline_cls.from_pretrained( + model_id, + torch_dtype=torch_dtype, + use_safetensors=True, + **MODEL_DEFAULTS[model_type].get("from_pretrained_extra_args", {}), + ) + pipe.set_progress_bar_config(disable=True) + return pipe def create_pipeline(self) -> Any: """ @@ -157,42 +155,32 @@ def setup_device(self) -> None: self.logger.info("Enabling VAE tiling for LTX-Video") self.pipe.vae.enable_tiling() - def get_backbone(self) -> torch.nn.Module: - """ - Get the backbone model (transformer or UNet). - - Returns: - Backbone model module - """ - if not self.pipe: - raise RuntimeError("Pipeline not created. Call create_pipeline() first.") - - backbone_pairs = list(self.iter_backbones()) - if len(backbone_pairs) == 1: - return backbone_pairs[0][1] - return torch.nn.ModuleList([module for _, module in backbone_pairs]) - def iter_backbones(self) -> Iterator[tuple[str, torch.nn.Module]]: """ - Yield backbone modules by name, based on a backbone spec. - - Yields: - (backbone_name, module) pairs + Yield (backbone_name, module) pairs. """ if not self.pipe: raise RuntimeError("Pipeline not created. Call create_pipeline() first.") names = list(self.config.backbone) + if not names: + raise RuntimeError("No backbone names provided.") if self.config.model_type == ModelType.LTX2: - self._ensure_ltx2_transformer_cached() - name = names[0] if names else "transformer" - yield name, self._transformer + for name in names: + if name == "video_decoder": + self._ensure_ltx2_video_decoder_cached() + yield name, self._video_decoder + elif name == "transformer": + self._ensure_ltx2_transformer_cached() + yield name, self._transformer + else: + raise ValueError( + f"Unsupported LTX-2 backbone name '{name}'. " + "Expected 'transformer' or 'video_decoder'." + ) return - if not names: - raise RuntimeError("No backbone names provided.") - for name in names: module = getattr(self.pipe, name, None) if module is None: @@ -207,6 +195,16 @@ def _ensure_ltx2_transformer_cached(self) -> None: self.pipe.stage_1_model_ledger.transformer = lambda: transformer self._transformer = transformer + def _ensure_ltx2_video_decoder_cached(self) -> None: + if not self.pipe: + raise RuntimeError("Pipeline not created. Call create_pipeline() first.") + if self._video_decoder is None: + video_decoder = self.pipe.stage_1_model_ledger.video_decoder() + # Cache it so subsequent calls return the same (quantized) instance + self.pipe.stage_1_model_ledger.video_decoder = lambda: video_decoder + self.pipe.stage_2_model_ledger.video_decoder = lambda: video_decoder + self._video_decoder = video_decoder + def _create_ltx2_pipeline(self) -> Any: params = dict(self.config.extra_params) checkpoint_path = params.pop("checkpoint_path", None) @@ -261,7 +259,6 @@ def _create_ltx2_pipeline(self) -> Any: return TI2VidTwoStagesPipeline(**pipeline_kwargs) def print_quant_summary(self): - backbone_pairs = list(self.iter_backbones()) - for name, backbone in backbone_pairs: + for name, backbone in self.iter_backbones(): self.logger.info(f"{name} quantization info:") mtq.print_quant_summary(backbone) diff --git a/examples/diffusers/quantization/quantize.py b/examples/diffusers/quantization/quantize.py index 953888f62b1..2a3c947a2d6 100644 --- a/examples/diffusers/quantization/quantize.py +++ b/examples/diffusers/quantization/quantize.py @@ -116,33 +116,48 @@ def get_quant_config(self, n_steps: int, backbone: torch.nn.Module) -> Any: if self.config.format == QuantFormat.INT8: if self.config.algo == QuantAlgo.SMOOTHQUANT: - quant_config = mtq.INT8_SMOOTHQUANT_CFG + base_cfg = mtq.INT8_SMOOTHQUANT_CFG else: - quant_config = INT8_DEFAULT_CONFIG + base_cfg = INT8_DEFAULT_CONFIG if self.config.collect_method != CollectMethod.DEFAULT: reset_set_int8_config( - quant_config, + base_cfg, self.config.percentile, n_steps, collect_method=self.config.collect_method.value, backbone=backbone, ) elif self.config.format == QuantFormat.FP8: - quant_config = FP8_DEFAULT_CONFIG + base_cfg = FP8_DEFAULT_CONFIG elif self.config.format == QuantFormat.FP4: if self.model_config.model_type.value.startswith("flux"): - quant_config = NVFP4_FP8_MHA_CONFIG + base_cfg = NVFP4_FP8_MHA_CONFIG else: - quant_config = NVFP4_DEFAULT_CONFIG + base_cfg = NVFP4_DEFAULT_CONFIG else: raise NotImplementedError(f"Unknown format {self.config.format}") + + # Build a fresh config dict so we never mutate the global constants. + quant_cfg_list = list(base_cfg["quant_cfg"]) + + if self.config.format == QuantFormat.FP4: + for i, entry in enumerate(quant_cfg_list): + if isinstance(entry, dict) and "block_sizes" in entry.get("cfg", {}): + new_block_sizes = {**entry["cfg"]["block_sizes"], -1: self.config.block_size} + quant_cfg_list[i] = { + **entry, + "cfg": {**entry["cfg"], "block_sizes": new_block_sizes}, + } + if self.config.quantize_mha: - quant_config["quant_cfg"].append( + quant_cfg_list.append( { "quantizer_name": "*[qkv]_bmm_quantizer", "cfg": {"num_bits": (4, 3), "axis": None}, } ) + + quant_config = {**base_cfg, "quant_cfg": quant_cfg_list} set_quant_config_attr( quant_config, self.model_config.trt_high_precision_dtype.value, @@ -158,6 +173,7 @@ def quantize_model( backbone: torch.nn.Module, quant_config: Any, forward_loop: callable, # type: ignore[valid-type] + backbone_name: str = "transformer", ) -> torch.nn.Module: """ Apply quantization to the model. @@ -166,15 +182,18 @@ def quantize_model( backbone: Model backbone to quantize quant_config: Quantization configuration forward_loop: Forward pass function for calibration + backbone_name: Name of the backbone being quantized """ self.logger.info("Checking for LoRA layers...") check_lora(backbone) - self.logger.info("Starting model quantization...") + self.logger.info(f"Starting model quantization for {backbone_name}...") mtq.quantize(backbone, quant_config, forward_loop) # Get model-specific filter function - model_filter_func = get_model_filter_func(self.model_config.model_type) - self.logger.info(f"Using filter function for {self.model_config.model_type.value}") + model_filter_func = get_model_filter_func(self.model_config.model_type, backbone_name) + self.logger.info( + f"Using filter function for {self.model_config.model_type.value}/{backbone_name}" + ) self.logger.info("Disabling specific quantizers...") mtq.disable_quantizer(backbone, model_filter_func) @@ -221,20 +240,27 @@ def _has_conv_layers(self, model: torch.nn.Module) -> bool: return True return False - def save_checkpoint(self, backbone: torch.nn.Module) -> None: + def save_checkpoint( + self, + backbone: torch.nn.Module, + backbone_name: str | None = None, + ) -> None: """ Save quantized model checkpoint. Args: backbone: The quantized backbone module to save (must be the same instance that was passed to mtq.quantize, as it carries the _modelopt_state). + backbone_name: Optional name for the backbone file (defaults to "backbone"). """ if not self.config.quantized_torch_ckpt_path: return ckpt_path = self.config.quantized_torch_ckpt_path ckpt_path.mkdir(parents=True, exist_ok=True) - target_path = ckpt_path / "backbone.pt" + filename = f"{backbone_name}.pt" if backbone_name else "backbone.pt" + target_path = ckpt_path / filename + self.logger.info(f"Saving backbone to {target_path}") mto.save(backbone, str(target_path)) @@ -292,14 +318,19 @@ def restore_checkpoint(self) -> None: if self.pipeline_manager is None: raise RuntimeError("Pipeline manager is required for per-backbone checkpoints.") - backbone = self.pipeline_manager.get_backbone() - if restore_path.exists() and restore_path.is_dir(): - source_path = restore_path / "backbone.pt" + if not restore_path.exists() or not restore_path.is_dir(): + raise FileNotFoundError(f"Checkpoint directory not found: {restore_path}") + + for backbone_name, backbone in self.pipeline_manager.iter_backbones(): + source_path = restore_path / f"{backbone_name}.pt" if not source_path.exists(): - raise FileNotFoundError(f"Backbone checkpoint not found: {source_path}") - self.logger.info(f"Restoring backbone from {source_path}") + raise FileNotFoundError( + f"Checkpoint not found for '{backbone_name}' in {restore_path}" + ) + self.logger.info(f"Restoring {backbone_name} from {source_path}") mto.restore(backbone, str(source_path)) - self.logger.info("Backbone checkpoints restored successfully") + + self.logger.info("Checkpoints restored successfully") # TODO: should not do the any data type def export_hf_ckpt(self, pipe: Any, model_config: ModelConfig | None = None) -> None: @@ -368,9 +399,9 @@ def create_argument_parser() -> argparse.ArgumentParser: nargs="+", default=None, help=( - "Model backbone(s) in the DiffusionPipeline to work on. " - "Provide one name or multiple names separated by space or comma. " - "If not provided use default based on model type." + "Model backbone(s) in the DiffusionPipeline to quantize. " + "Provide one or more names (e.g., 'transformer', 'video_decoder'). " + "If not provided, uses default based on model type." ), ) model_group.add_argument( @@ -448,6 +479,12 @@ def create_argument_parser() -> argparse.ArgumentParser: action="store_true", help="Compress quantized weights to reduce memory footprint (FP8/FP4 only)", ) + quant_group.add_argument( + "--block-size", + type=int, + default=16, + help="Block size for NVFP4 quantization (default: 16)", + ) calib_group = parser.add_argument_group("Calibration Configuration") calib_group.add_argument("--batch-size", type=int, default=2, help="Batch size for calibration") @@ -535,6 +572,7 @@ def main() -> None: lowrank=args.lowrank, quantize_mha=args.quantize_mha, compress=args.compress, + block_size=args.block_size, ) if args.prompts_file is not None: @@ -571,7 +609,6 @@ def main() -> None: pipe = pipeline_manager.create_pipeline() pipeline_manager.setup_device() - backbone = pipeline_manager.get_backbone() export_manager = ExportManager(export_config, logger, pipeline_manager) if export_config.restore_from and export_config.restore_from.exists(): @@ -581,37 +618,48 @@ def main() -> None: logger.info("Initializing calibration...") calibrator = Calibrator(pipeline_manager, calib_config, model_config.model_type, logger) batched_prompts = calibrator.load_and_batch_prompts() - quantizer = Quantizer(quant_config, model_config, logger) - backbone_quant_config = quantizer.get_quant_config(calib_config.n_steps, backbone) - # Pipe loads the ckpt just before the inference. - def forward_loop(mod): - calibrator.run_calibration(batched_prompts) + for backbone_name, backbone in pipeline_manager.iter_backbones(): + logger.info(f"Quantizing backbone: {backbone_name}") + backbone_quant_config = quantizer.get_quant_config(calib_config.n_steps, backbone) - quantizer.quantize_model(backbone, backbone_quant_config, forward_loop) + # Calibration runs the full pipeline (not just `mod`), so the + # closure intentionally ignores the backbone argument. + def forward_loop(mod): + calibrator.run_calibration(batched_prompts) - # Compress model weights if requested (only for FP8/FP4) - if quant_config.compress: - logger.info("Compressing model weights to reduce memory footprint...") - mtq.compress(backbone) - logger.info("Model compression completed") + quantizer.quantize_model( + backbone, + backbone_quant_config, + forward_loop, + backbone_name=backbone_name, + ) - export_manager.save_checkpoint(backbone) + # Compress model weights if requested (only for FP8/FP4) + if quant_config.compress: + logger.info(f"Compressing {backbone_name} weights...") + mtq.compress(backbone) + logger.info(f"{backbone_name} compression completed") - # TODO (Jingyu): To update this function, as we are focusing more on the torch deployment side. - check_conv_and_mha( - backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha - ) + # For VAE backbones, skip check_conv_and_mha — the whole point + # of VAE quantization is to quantize Conv layers. + if backbone_name not in ("video_decoder", "vae"): + check_conv_and_mha( + backbone, quant_config.format == QuantFormat.FP4, quant_config.quantize_mha + ) + + export_manager.save_checkpoint(backbone, backbone_name) pipeline_manager.print_quant_summary() - export_manager.export_onnx( - pipe, - backbone, - model_config.model_type, - quant_config.format, - ) + for backbone_name, backbone in pipeline_manager.iter_backbones(): + export_manager.export_onnx( + pipe, + backbone, + model_config.model_type, + quant_config.format, + ) export_manager.export_hf_ckpt(pipe, model_config) diff --git a/examples/diffusers/quantization/quantize_config.py b/examples/diffusers/quantization/quantize_config.py index 606debf9e32..a92dd4e8147 100644 --- a/examples/diffusers/quantization/quantize_config.py +++ b/examples/diffusers/quantization/quantize_config.py @@ -79,6 +79,7 @@ class QuantizationConfig: lowrank: int = 32 # SVDQuant lowrank quantize_mha: bool = False compress: bool = False + block_size: int = 16 # NVFP4 block size def validate(self) -> None: """Validate configuration consistency.""" @@ -120,7 +121,7 @@ class ModelConfig: model_type: ModelType = ModelType.FLUX_DEV model_dtype: dict[str, torch.dtype] = field(default_factory=lambda: {"default": torch.float16}) - backbone: str = "" + backbone: list[str] = field(default_factory=list) trt_high_precision_dtype: DataType = DataType.HALF override_model_path: Path | None = None cpu_offloading: bool = False diff --git a/examples/diffusers/quantization/utils.py b/examples/diffusers/quantization/utils.py index 0c38fc28606..d102e83e068 100644 --- a/examples/diffusers/quantization/utils.py +++ b/examples/diffusers/quantization/utils.py @@ -46,7 +46,7 @@ def filter_func_default(name: str) -> bool: def check_conv_and_mha(backbone, if_fp4, quantize_mha): for name, module in backbone.named_modules(): - if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d)) and if_fp4: + if isinstance(module, (torch.nn.Conv1d, torch.nn.Conv2d)) and if_fp4: module.weight_quantizer.disable() module.input_quantizer.disable() @@ -87,6 +87,22 @@ def filter_func_flux_dev(name: str) -> bool: return pattern.match(name) is not None +def filter_func_ltx2_vae(name: str) -> bool: + """Filter for LTX-2 VAE: keeps only conv1/conv2 in up_blocks resnets.""" + keep = re.compile(r".*up_blocks\.\d+\.resnets\.\d+\.conv[12](?:\.|$)") + return not keep.match(name) + + +def filter_func_wan_vae(name: str) -> bool: + """Filter for Wan 2.2 VAE: keeps only conv1/conv2 in resnet blocks.""" + keep = re.compile( + r".*(down_blocks\.\d+\.(?:resnets\.\d+\.)?conv[12]" + r"|mid_block\.resnets\.\d+\.conv[12]" + r"|up_blocks\.\d+\.resnets\.\d+\.conv[12])(?:\.|$)" + ) + return not keep.match(name) + + def filter_func_wan_video(name: str) -> bool: """Filter function specifically for WAN-Video models.""" pattern = re.compile( diff --git a/experimental/conv/README.md b/experimental/conv/README.md deleted file mode 100644 index 65b7cc5563e..00000000000 --- a/experimental/conv/README.md +++ /dev/null @@ -1,105 +0,0 @@ -# Conv3D Implicit GEMM (Experimental) - -Experimental Conv3D kernel prototype using implicit GEMM, with optional fused FP4 fake quantization for activations. - -This code is kept under `experimental/` by design and is **not** part of the stable `modelopt.torch.quantization` API. - -## Model Support - -| Model/Framework | Supported | Notes | -|-----------------|-----------|-------| -| Video diffusion VAE Conv3D layers | Tested | Validated on VAE encoder/decoder Conv3D layers in video diffusion models | -| Generic LLM backbones | No | Conv3D path is not relevant | -| End-to-end ModelOpt PTQ/QAT pipeline | No | Not wired into formal quantization/export/compress flows | - -## Deployment - -| Framework | Supported | Notes | -|-----------|-----------|-------| -| TensorRT-LLM | No | No formal export integration for this kernel path | -| vLLM | No | No integration | -| SGLang | No | No integration | -| PyTorch runtime (CUDA) | Yes (experimental) | JIT-compiles CUDA extension on first use | - -## Usage - -```python -import torch - -from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda -from modelopt.torch.quantization.tensor_quant import dynamic_block_quantize_op - -x = torch.randn(1, 128, 21, 60, 106, device="cuda") -w = torch.randn(512, 128, 3, 3, 3, device="cuda") -block_size = 128 - -# Without FP4 activation quantization (drop-in-style Conv3D call) -out = conv3d_implicit_gemm_cuda(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) - -# Optional FP4 block quantization of weights along the GEMM K dimension. -# The kernel's A-tile (activations) is quantized along K = Cin*kD*kH*kW, -# so weights must be flattened to [Cout, K] before quantizing to match. -Cout, Cin = w.shape[:2] -K = Cin * w.shape[2] * w.shape[3] * w.shape[4] -w_flat = w.reshape(Cout, K) -w_q_flat = dynamic_block_quantize_op( - w_flat, - block_size, - w_flat.abs().max().unsqueeze(0), - 4, # num_bits - 2, # exponent_bits - 8, # scale_num_bits - 4, # scale_exponent_bits -) -w_q = w_q_flat.reshape_as(w) - -# With FP4 activation fake quantization -out_q = conv3d_implicit_gemm_cuda( - x, - w_q, - stride=(1, 1, 1), - padding=(1, 1, 1), - act_amax=x.abs().max().unsqueeze(0), - quant_act=True, - fp4_block_size=block_size, # 16, 32, 64, 128, or 256 -) -``` - -## API - -Function: `conv3d_implicit_gemm_cuda(...)` from `experimental/conv/implicit_gemm_cuda.py` - -| Parameter | Description | -|-----------|-------------| -| `x` | Input tensor `[N, Cin, D, H, W]` | -| `w` | Weight tensor `[Cout, Cin, kD, kH, kW]` | -| `bias` | Optional bias `[Cout]` | -| `stride` | Convolution stride `(D, H, W)` | -| `padding` | Convolution padding `(D, H, W)` | -| `dilation` | Convolution dilation `(D, H, W)` | -| `act_amax` | Activation abs-max scalar tensor (required when `quant_act=True`) | -| `quant_act` | Enable FP4 fake quantization on activations | -| `fp4_block_size` | FP4 quantization block size (`16`, `32`, `64`, `128`, or `256`) | - -## Status - -Current state: **Prototype** - -Known limitations: - -- API is unstable and may change without notice. -- Not registered in core quantization module registries. -- Not covered by formal export/compress integration. -- CUDA extension compile latency on first invocation. -- Validation and performance coverage are limited to local experiments. - -## Notes - -- The CUDA kernel is JIT-compiled on first call (can take several seconds). -- Output shape matches `torch.nn.functional.conv3d`. -- FP4 path applies quantize-dequantize in-kernel for activation tiles. - -## References - -- Implicit GEMM-based convolution design patterns in GPU kernels. -- ModelOpt FP4-related quantization utilities in `modelopt.torch.quantization.tensor_quant`. diff --git a/experimental/conv/bench_implicit_gemm.py b/experimental/conv/bench_implicit_gemm.py deleted file mode 100644 index 164c0744674..00000000000 --- a/experimental/conv/bench_implicit_gemm.py +++ /dev/null @@ -1,208 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Latency benchmark: implicit GEMM (quant / non-quant) vs cuDNN conv3d. - -Usage: - python -m experimental.conv.bench_implicit_gemm - python -m experimental.conv.bench_implicit_gemm --shapes wan22 - python -m experimental.conv.bench_implicit_gemm --shapes all --warmup 20 --iters 100 -""" - -import argparse - -import torch -import torch.nn.functional as F - -# --------------------------------------------------------------------------- -# Benchmark shapes -# --------------------------------------------------------------------------- - -# (name, N, Cin, D, H, W, Cout, kD, kH, kW, stride, padding, dilation) -SHAPES = { - "small": [ - ("small_16x32_3x3x3", 1, 16, 8, 8, 8, 32, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ], - "medium": [ - ("med_64x128_3x3x3", 1, 64, 16, 32, 32, 128, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ("med_128x256_3x3x3", 1, 128, 8, 16, 16, 256, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ("med_128x128_1x3x3", 1, 128, 16, 32, 32, 128, 1, 3, 3, (1, 1, 1), (0, 1, 1), (1, 1, 1)), - ], - "wan22": [ - ("wan22_128x512", 1, 128, 21, 60, 106, 512, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ("wan22_512x512", 1, 512, 21, 60, 106, 512, 1, 1, 1, (1, 1, 1), (0, 0, 0), (1, 1, 1)), - ("wan22_512x128", 1, 512, 21, 60, 106, 128, 3, 3, 3, (1, 1, 1), (1, 1, 1), (1, 1, 1)), - ], - "stride": [ - ("stride2_64x128", 1, 64, 16, 32, 32, 128, 3, 3, 3, (2, 2, 2), (1, 1, 1), (1, 1, 1)), - ("stride2_128x256", 1, 128, 16, 32, 32, 256, 3, 3, 3, (2, 2, 2), (1, 1, 1), (1, 1, 1)), - ], -} - - -def get_shapes(name: str): - """Return list of benchmark shapes by name or all shapes.""" - if name == "all": - result = [] - for v in SHAPES.values(): - result.extend(v) - return result - return SHAPES[name] - - -# --------------------------------------------------------------------------- -# Timing utility -# --------------------------------------------------------------------------- - - -def bench_fn(fn, warmup: int, iters: int) -> float: - """Benchmark a callable, return median time in ms.""" - for _ in range(warmup): - fn() - torch.cuda.synchronize() - - times = [] - for _ in range(iters): - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - fn() - end.record() - torch.cuda.synchronize() - times.append(start.elapsed_time(end)) - - times.sort() - return times[len(times) // 2] # median - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - - -def run_benchmark(shapes_name: str, warmup: int, iters: int, fp4_block_size: int): - """Run latency benchmark for the given shapes.""" - from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda - - shapes = get_shapes(shapes_name) - - # Header - print(f"\n{'=' * 100}") - print( - f"Conv3D Latency Benchmark | warmup={warmup} iters={iters} fp4_block_size={fp4_block_size}" - ) - print(f"GPU: {torch.cuda.get_device_name()}") - print(f"{'=' * 100}") - print( - f"{'Shape':<25} {'M':>10} {'K':>8} {'N':>6} " - f"{'cuDNN':>9} {'GEMM':>9} {'GEMM+FP4':>9} " - f"{'GEMM/cuDNN':>11} {'FP4/cuDNN':>10}" - ) - print("-" * 100) - - for name, n, cin, d, h, w, cout, kd, kh, kw, stride, padding, dilation in shapes: - torch.manual_seed(42) - x = torch.randn(n, cin, d, h, w, device="cuda", dtype=torch.float32) - weight = torch.randn(cout, cin, kd, kh, kw, device="cuda", dtype=torch.float32) - act_amax = x.abs().max().unsqueeze(0) - - # Compute GEMM dimensions for display - sd, sh, sw = stride - dd, dh, dw = dilation - pd, ph, pw = padding - od = (d + 2 * pd - dd * (kd - 1) - 1) // sd + 1 - oh = (h + 2 * ph - dh * (kh - 1) - 1) // sh + 1 - ow = (w + 2 * pw - dw * (kw - 1) - 1) // sw + 1 - gemm_m = n * od * oh * ow - gemm_k = cin * kd * kh * kw - gemm_n = cout - - # cuDNN (torch.nn.functional.conv3d) - t_cudnn = bench_fn( - lambda: F.conv3d(x, weight, stride=stride, padding=padding, dilation=dilation), - warmup, - iters, - ) - - # Implicit GEMM (non-quantized) - t_gemm = bench_fn( - lambda: conv3d_implicit_gemm_cuda( - x, - weight, - stride=stride, - padding=padding, - dilation=dilation, - quant_act=False, - fp4_block_size=fp4_block_size, - ), - warmup, - iters, - ) - - # Implicit GEMM (FP4 quantized) - t_fp4 = bench_fn( - lambda: conv3d_implicit_gemm_cuda( - x, - weight, - stride=stride, - padding=padding, - dilation=dilation, - act_amax=act_amax, - quant_act=True, - fp4_block_size=fp4_block_size, - ), - warmup, - iters, - ) - - ratio_gemm = t_gemm / t_cudnn - ratio_fp4 = t_fp4 / t_cudnn - - print( - f"{name:<25} {gemm_m:>10,} {gemm_k:>8,} {gemm_n:>6,} " - f"{t_cudnn:>8.3f}ms {t_gemm:>8.3f}ms {t_fp4:>8.3f}ms " - f"{ratio_gemm:>10.2f}x {ratio_fp4:>9.2f}x" - ) - - print(f"{'=' * 100}") - print("Ratios > 1.0x mean slower than cuDNN; < 1.0x mean faster.") - print() - - -def main(): - """Entry point for the benchmark CLI.""" - parser = argparse.ArgumentParser(description="Conv3D latency benchmark") - parser.add_argument( - "--shapes", - default="all", - choices=[*list(SHAPES.keys()), "all"], - help="Which shape set to benchmark (default: all)", - ) - parser.add_argument("--warmup", type=int, default=20, help="Warmup iterations") - parser.add_argument("--iters", type=int, default=100, help="Benchmark iterations") - parser.add_argument( - "--fp4-block-size", - type=int, - default=128, - choices=[128, 256], - help="FP4 block size (default: 128)", - ) - args = parser.parse_args() - - run_benchmark(args.shapes, args.warmup, args.iters, args.fp4_block_size) - - -if __name__ == "__main__": - main() diff --git a/modelopt/torch/quantization/nn/modules/quant_conv.py b/modelopt/torch/quantization/nn/modules/quant_conv.py index 44f0ae663c5..ed165556249 100644 --- a/modelopt/torch/quantization/nn/modules/quant_conv.py +++ b/modelopt/torch/quantization/nn/modules/quant_conv.py @@ -15,8 +15,12 @@ """Quantized convolution.""" +import warnings + import torch.nn as nn +from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda + from ... import tensor_quant from .quant_module import QuantLinearConvBase, QuantModuleRegistry, _LegacyQuantLinearConvBaseMixin @@ -62,12 +66,85 @@ class QuantConv2d(_LegacyQuantLinearConvBaseMixin, nn.Conv2d): default_quant_desc_weight = _QuantConv2d.default_quant_desc_weight +def _is_nvfp4_quantizer(quantizer) -> bool: + """Check if a TensorQuantizer is configured for NVFP4 dynamic block quantization.""" + return ( + quantizer.num_bits == (2, 1) + and quantizer.block_sizes is not None + and quantizer.block_sizes.get("scale_bits") == (4, 3) + and quantizer.block_sizes.get("type") == "dynamic" + ) + + +def _nvfp4_quantize_weight_along_k(weight, weight_quantizer): + """Apply NVFP4 fake quantization to Conv3D weight along the GEMM K dimension.""" + w_flat = weight.reshape(weight.shape[0], -1) + return weight_quantizer(w_flat).reshape_as(weight) + + @QuantModuleRegistry.register({nn.Conv3d: "nn.Conv3d"}) class _QuantConv3d(QuantLinearConvBase): - """Quantized 3D convolution.""" + """Quantized 3D convolution. + + For NVFP4, uses a fused implicit GEMM kernel with activation FP4 quantization + inside the kernel. For all other configs, the default cuDNN path is used. + """ default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL + def _should_use_implicit_gemm(self): + """Check if both quantizers are NVFP4 and the implicit GEMM kernel is available.""" + return ( + hasattr(self, "input_quantizer") + and hasattr(self, "weight_quantizer") + and _is_nvfp4_quantizer(self.input_quantizer) + and _is_nvfp4_quantizer(self.weight_quantizer) + and self.groups == 1 + ) + + def _implicit_gemm_forward(self, input): + """Run NVFP4 implicit GEMM kernel. Input may already be padded.""" + # _get_amax is an internal TensorQuantizer method with no public equivalent; + # block_sizes is a public property. + act_amax = self.input_quantizer._get_amax(input) + weight = _nvfp4_quantize_weight_along_k(self.weight, self.weight_quantizer) + fp4_block_size = self.input_quantizer.block_sizes.get(-1, 16) + + output = conv3d_implicit_gemm_cuda( + input, + weight, + bias=self.bias, + stride=self.stride, + padding=self.padding, + dilation=self.dilation, + act_amax=act_amax, + quant_act=self.input_quantizer.is_enabled, + fp4_block_size=fp4_block_size, + ) + return self.output_quantizer(output) + + def forward(self, input, *args, **kwargs): + """Forward with implicit GEMM for NVFP4, default path otherwise.""" + if not self._should_use_implicit_gemm(): + return super().forward(input, *args, **kwargs) + + if self.training: + warnings.warn( + "Implicit GEMM Conv3D kernel is inference-only and does not support training. " + "Falling back to the default cuDNN quantization path, which could produce " + "different numerics.", + stacklevel=2, + ) + return super().forward(input, *args, **kwargs) + + # During calibration, only collect amax — use the faster cuDNN path. + # _if_calib/_if_quant are internal TensorQuantizer state with no public property; + # toggled via enable_calib()/disable_calib()/enable_quant()/disable_quant(). + if self.input_quantizer._if_calib and not self.input_quantizer._if_quant: + return super().forward(input, *args, **kwargs) + + return self._implicit_gemm_forward(input) + class QuantConv3d(_LegacyQuantLinearConvBaseMixin, nn.Conv3d): """Quantized 3D convolution.""" diff --git a/modelopt/torch/quantization/plugins/diffusion/diffusers.py b/modelopt/torch/quantization/plugins/diffusion/diffusers.py index f9ae55b3e2d..f2f6a702479 100644 --- a/modelopt/torch/quantization/plugins/diffusion/diffusers.py +++ b/modelopt/torch/quantization/plugins/diffusion/diffusers.py @@ -63,6 +63,7 @@ QuantModuleRegistry, TensorQuantizer, ) +from ...nn.modules.quant_conv import _QuantConv3d from ..custom import _QuantFunctionalMixin onnx_dtype_map = { @@ -278,3 +279,37 @@ def symbolic( high_precision_flag, disable_fp8_mha, ) + + +# WanCausalConv3d quantization support (diffusers VAE) +try: + from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d + + @QuantModuleRegistry.register({WanCausalConv3d: "WanCausalConv3d"}) + class _QuantDiffusersWanCausalConv3d(_QuantConv3d): + """Quantized WanCausalConv3d — applies causal padding before quantized conv.""" + + def forward(self, x, cache_x=None): + # Apply WanCausalConv3d-specific causal padding + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + cache_x = cache_x.to(x.device) + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + # NVFP4 implicit GEMM path (self.padding is (0,0,0) since padding already applied) + if self._should_use_implicit_gemm(): + if not (self.input_quantizer._if_calib and not self.input_quantizer._if_quant): + return self._implicit_gemm_forward(x) + + # Default quantized conv path (skip WanCausalConv3d.forward to avoid double-padding) + with self.quantize_weight(): + input = self.input_quantizer(x) + output = torch.nn.Conv3d.forward(self, input) + if isinstance(output, tuple): + return (self.output_quantizer(output[0]), *output[1:]) + return self.output_quantizer(output) + +except ImportError: + pass diff --git a/modelopt/torch/quantization/src/conv/README.md b/modelopt/torch/quantization/src/conv/README.md new file mode 100644 index 00000000000..6b14fd5953b --- /dev/null +++ b/modelopt/torch/quantization/src/conv/README.md @@ -0,0 +1,141 @@ +# Conv3D Implicit GEMM + +Conv3D kernel using implicit GEMM with BF16 WMMA tensor cores and optional fused FP4 (E2M1) fake quantization. + +This kernel is integrated into `modelopt.torch.quantization` via `_QuantConv3d` — when NVFP4 quantization is applied to an `nn.Conv3d` layer through ModelOpt PTQ, the implicit GEMM path is used automatically. We have only tested it on VAE Conv3D layers from video generation models (e.g. Wan2.2). + +## Requirements + +- **GPU:** SM80+ (Ampere or newer) for BF16 WMMA tensor cores +- **PyTorch:** CUDA toolkit with JIT C++ extension support (`torch.utils.cpp_extension`) +- **Grouped convolution is not supported** (groups must be 1) + +## Data Types + +| Stage | Precision | +|-------|-----------| +| Input / output tensors | FP32, FP16, or BF16 (dtype is preserved) | +| Internal compute | BF16 via WMMA m16n16k16 tensor cores | +| Accumulation | FP32 | +| FP4 activation quantization | E2M1 values, FP8 E4M3 scales | + +## Integration with ModelOpt Quantization + +When NVFP4 quantization is configured on a `Conv3d` layer via ModelOpt PTQ, the implicit GEMM kernel is used automatically during quantized inference. The integration is in `_QuantConv3d` (`modelopt/torch/quantization/nn/modules/quant_conv.py`): + +- During **calibration**, the standard cuDNN path is used (faster). +- During **quantized inference** with NVFP4 input and weight quantizers, the kernel fuses activation FP4 quantization inside the GEMM. +- For all other quantization configs, the default cuDNN path is used as fallback. + +## Usage + +```python +import torch + +from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda +from modelopt.torch.quantization.tensor_quant import dynamic_block_quantize_op + +x = torch.randn(1, 128, 21, 60, 106, device="cuda") +w = torch.randn(512, 128, 3, 3, 3, device="cuda") +block_size = 128 + +# Without FP4 activation quantization (drop-in-style Conv3D call) +out = conv3d_implicit_gemm_cuda(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + +# Optional FP4 block quantization of weights along the GEMM K dimension. +# The kernel's A-tile (activations) is quantized along K = Cin*kD*kH*kW, +# so weights must be flattened to [Cout, K] before quantizing to match. +Cout, Cin = w.shape[:2] +K = Cin * w.shape[2] * w.shape[3] * w.shape[4] +w_flat = w.reshape(Cout, K) +w_q_flat = dynamic_block_quantize_op( + w_flat, + block_size, + w_flat.abs().max().unsqueeze(0), + 4, # num_bits + 2, # exponent_bits + 8, # scale_num_bits + 4, # scale_exponent_bits +) +w_q = w_q_flat.reshape_as(w) + +# With FP4 activation fake quantization +out_q = conv3d_implicit_gemm_cuda( + x, + w_q, + stride=(1, 1, 1), + padding=(1, 1, 1), + act_amax=x.abs().max().unsqueeze(0), + quant_act=True, + fp4_block_size=block_size, # 16, 32, 64, 128, or 256 +) +``` + +## API + +### `conv3d_implicit_gemm_cuda` + +`from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda` + +| Parameter | Description | +|-----------|-------------| +| `x` | Input tensor `[N, Cin, D, H, W]` | +| `w` | Weight tensor `[Cout, Cin, kD, kH, kW]` | +| `bias` | Optional bias `[Cout]` | +| `stride` | Convolution stride `(D, H, W)` | +| `padding` | Convolution padding `(D, H, W)` | +| `dilation` | Convolution dilation `(D, H, W)` | +| `act_amax` | Activation abs-max scalar tensor (required when `quant_act=True`) | +| `quant_act` | Enable FP4 fake quantization on activations | +| `fp4_block_size` | FP4 quantization block size (`16`, `32`, `64`, `128`, or `256`) | + +### `fp4_fake_quant` + +`from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import fp4_fake_quant` + +Standalone FP4 (E2M1) blockwise fake quantization with FP8 E4M3 scale quantization. Uses the same CUDA device functions as the fused path inside the GEMM kernel. + +| Parameter | Description | +|-----------|-------------| +| `x` | Input tensor (any shape; `numel` must be divisible by `block_size`) | +| `global_amax` | Scalar tensor — global abs max for scale computation | +| `block_size` | Number of elements per FP4 quantization block (default `16`) | + +## Testing + +```bash +# Run tests (requires GPU) +python -m pytest tests/gpu/torch/quantization/kernels/test_implicit_gemm.py -v +``` + +## Status + +Current state: **Integrated** (registered in `QuantModuleRegistry`, auto-dispatched for NVFP4 Conv3D) + +Known limitations: + +- CUDA extension compile latency on first invocation (~seconds). +- Grouped convolution (`groups > 1`) is not supported. In the ModelOpt E2E flow, `_QuantConv3d` automatically falls back to the default cuDNN path for grouped convolutions. +- BF16 rounding error accumulates with the K dimension — expect max abs diff scaling roughly as `sqrt(K)` compared to cuDNN FP32. +- Inference only (`@torch.no_grad`) — not suitable for QAT backward pass. + +## Notes + +- The CUDA kernel is JIT-compiled on first call via `torch.utils.cpp_extension.load()`. +- Output shape matches `torch.nn.functional.conv3d`. +- FP4 path applies quantize-dequantize in-kernel for activation tiles (no extra global memory pass). +- Tile config: BLOCK_M=64, BLOCK_N=64, BLOCK_K=256, 8 warps (256 threads), ~70 KB shared memory per block. +- The kernel body is guarded by `#if __CUDA_ARCH__ >= 800` so it compiles as an empty stub when nvcc targets pre-Ampere archs (PyTorch's default `-gencode` list can include sm_75, which lacks BF16 WMMA fragments). Dispatch is enforced at runtime by `_get_cuda_module()` via `_MIN_SM_MAJOR = 8`. + +## Files + +| File | Role | +|------|------| +| `implicit_gemm_cuda.py` | Python API and JIT compilation | +| `implicit_gemm_kernel.cu` | CUDA kernel (BF16 WMMA + FP4 quantization) | +| `implicit_gemm_binding.cpp` | PyTorch C++ extension binding | + +## References + +- Implicit GEMM-based convolution design patterns in GPU kernels. +- ModelOpt FP4-related quantization utilities in `modelopt.torch.quantization.tensor_quant`. diff --git a/experimental/conv/implicit_gemm_binding.cpp b/modelopt/torch/quantization/src/conv/implicit_gemm_binding.cpp similarity index 100% rename from experimental/conv/implicit_gemm_binding.cpp rename to modelopt/torch/quantization/src/conv/implicit_gemm_binding.cpp diff --git a/experimental/conv/implicit_gemm_cuda.py b/modelopt/torch/quantization/src/conv/implicit_gemm_cuda.py similarity index 100% rename from experimental/conv/implicit_gemm_cuda.py rename to modelopt/torch/quantization/src/conv/implicit_gemm_cuda.py diff --git a/experimental/conv/implicit_gemm_kernel.cu b/modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu similarity index 98% rename from experimental/conv/implicit_gemm_kernel.cu rename to modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu index a3b40f48481..10d20c2e379 100644 --- a/experimental/conv/implicit_gemm_kernel.cu +++ b/modelopt/torch/quantization/src/conv/implicit_gemm_kernel.cu @@ -145,6 +145,10 @@ __global__ void __launch_bounds__(WARPS_M * WARPS_N * 32, 2) const float *__restrict__ act_amax, int Cin, int Dp, int Hp, int Wp, int Cout, int OD, int OH, int OW, int kD, int kH, int kW, int sd, int sh, int sw, int dd, int dh, int dw, int M, int K) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + // BF16 WMMA fragments are only available on sm_80+. On earlier archs this + // kernel compiles as an empty stub; the runtime gate in _get_cuda_module() + // prevents dispatching it on unsupported hardware. // Derived constants constexpr int NUM_WARPS = WARPS_M * WARPS_N; constexpr int NUM_THREADS = NUM_WARPS * 32; @@ -469,6 +473,7 @@ __global__ void __launch_bounds__(WARPS_M * WARPS_N * 32, 2) y[m_idx * Cout + n_idx] = result; } } +#endif // __CUDA_ARCH__ >= 800 } // ============================================================================= diff --git a/tests/examples/diffusers/conftest.py b/tests/examples/diffusers/conftest.py new file mode 100644 index 00000000000..8893d188d9e --- /dev/null +++ b/tests/examples/diffusers/conftest.py @@ -0,0 +1,31 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + + +@pytest.fixture(scope="session") +def tiny_wan22_path(tmp_path_factory): + """Create a tiny Wan 2.2 (14B-style) pipeline and return its path. + + Built once per session and shared across all tests that need it. + """ + try: + from _test_utils.torch.diffusers_models import create_tiny_wan22_pipeline_dir + except ImportError: + pytest.skip("Wan 2.2 diffusers models not available (requires diffusers with WanPipeline)") + + tmp_path = tmp_path_factory.mktemp("wan22") + return str(create_tiny_wan22_pipeline_dir(tmp_path)) diff --git a/tests/examples/diffusers/test_diffusers.py b/tests/examples/diffusers/test_diffusers.py index 5bc8f981ecd..5b117b41b3f 100644 --- a/tests/examples/diffusers/test_diffusers.py +++ b/tests/examples/diffusers/test_diffusers.py @@ -151,6 +151,119 @@ def test_diffusers_quantization( model.inference(tmp_path) +class Wan22Model(NamedTuple): + model: str + backbone: str | None + format_type: str + quant_algo: str + collect_method: str + + def _ckpt_path(self, tmp_path: Path) -> str: + stem = self.model.replace("wan2.2-t2v-", "") + parts = [stem, *([self.backbone] if self.backbone else []), self.format_type] + return str(tmp_path / f"wan22_{'_'.join(parts)}.pt") + + def _common_args(self, tiny_wan22_path: str) -> list[str]: + cmd_args = [ + "python", + "quantize.py", + "--model", + self.model, + "--override-model-path", + tiny_wan22_path, + "--format", + self.format_type, + "--quant-algo", + self.quant_algo, + "--collect-method", + self.collect_method, + "--model-dtype", + "BFloat16", + "--trt-high-precision-dtype", + "BFloat16", + "--calib-size", + "2", + "--batch-size", + "1", + "--n-steps", + "2", + # Tiny video dims — override MODEL_DEFAULTS for fast CI. + "--extra-param", + "height=16", + "--extra-param", + "width=16", + "--extra-param", + "num_frames=5", + ] + if self.backbone is not None: + cmd_args.extend(["--backbone", self.backbone]) + return cmd_args + + def quantize(self, tiny_wan22_path: str, tmp_path: Path) -> None: + run_example_command( + [ + *self._common_args(tiny_wan22_path), + "--quantized-torch-ckpt-save-path", + self._ckpt_path(tmp_path), + ], + "diffusers/quantization", + ) + + def restore(self, tiny_wan22_path: str, tmp_path: Path) -> None: + run_example_command( + [*self._common_args(tiny_wan22_path), "--restore-from", self._ckpt_path(tmp_path)], + "diffusers/quantization", + ) + + +# The VAE (``AutoencoderKLWan``) is shared between Wan 2.2 14B and 5B, so the +# Conv3D NVFP4 implicit-GEMM dispatch exercises the same kernel either way; we +# parametrize both ``--model`` values to also cover the ``quantize.py`` dispatch +# for each. +@pytest.mark.parametrize( + "wan_model", + [ + Wan22Model("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean"), + pytest.param( + Wan22Model("wan2.2-t2v-14b", None, "fp8", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-14b", None, "fp4", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-14b", "vae", "fp8", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-14b", "vae", "fp4", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-5b", "vae", "fp8", "max", "default"), + marks=minimum_sm(89), + ), + pytest.param( + Wan22Model("wan2.2-t2v-5b", "vae", "fp4", "max", "default"), + marks=minimum_sm(89), + ), + ], + ids=[ + "wan22_14b_transformer_int8_smoothquant", + "wan22_14b_transformer_fp8_max", + "wan22_14b_transformer_fp4_max", + "wan22_14b_vae_fp8_max", + "wan22_14b_vae_fp4_max", + "wan22_5b_vae_fp8_max", + "wan22_5b_vae_fp4_max", + ], +) +def test_wan22_quantization(wan_model: Wan22Model, tiny_wan22_path: str, tmp_path: Path) -> None: + wan_model.quantize(tiny_wan22_path, tmp_path) + wan_model.restore(tiny_wan22_path, tmp_path) + + @pytest.mark.parametrize( ("model_name", "model_path", "torch_compile"), [ diff --git a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py index 6db1eaeb68d..a5c81d36937 100644 --- a/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py +++ b/tests/examples/diffusers/test_export_diffusers_hf_ckpt.py @@ -128,3 +128,84 @@ def test_diffusers_hf_ckpt_export(model: DiffuserHfExportModel, tmp_path: Path) weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin")) assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}" + + +class Wan22HfExportModel(NamedTuple): + model: str + backbone: str | None + format_type: str + quant_algo: str + collect_method: str + + def _suffix(self) -> str: + stem = self.model.replace("wan2.2-t2v-", "") + parts = [stem, *([self.backbone] if self.backbone else []), self.format_type] + return "_".join(parts) + + def quantize_and_export_hf(self, tiny_wan22_path: str, tmp_path: Path) -> Path: + hf_ckpt_dir = tmp_path / f"wan22_{self._suffix()}_hf_ckpt" + cmd_args = [ + "python", + "quantize.py", + "--model", + self.model, + "--override-model-path", + tiny_wan22_path, + "--format", + self.format_type, + "--quant-algo", + self.quant_algo, + "--collect-method", + self.collect_method, + "--model-dtype", + "BFloat16", + "--trt-high-precision-dtype", + "BFloat16", + "--calib-size", + "2", + "--batch-size", + "1", + "--n-steps", + "2", + # Tiny video dims — override MODEL_DEFAULTS for fast CI. + "--extra-param", + "height=16", + "--extra-param", + "width=16", + "--extra-param", + "num_frames=5", + "--hf-ckpt-dir", + str(hf_ckpt_dir), + ] + if self.backbone is not None: + cmd_args.extend(["--backbone", self.backbone]) + run_example_command(cmd_args, "diffusers/quantization") + return hf_ckpt_dir + + +@pytest.mark.parametrize( + "wan_model", + [ + Wan22HfExportModel("wan2.2-t2v-14b", None, "int8", "smoothquant", "min-mean"), + pytest.param( + Wan22HfExportModel("wan2.2-t2v-14b", None, "fp8", "max", "default"), + marks=minimum_sm(89), + ), + ], + ids=[ + "wan22_14b_transformer_int8_smoothquant", + "wan22_14b_transformer_fp8_max", + ], +) +def test_wan22_hf_ckpt_export( + wan_model: Wan22HfExportModel, tiny_wan22_path: str, tmp_path: Path +) -> None: + hf_ckpt_dir = wan_model.quantize_and_export_hf(tiny_wan22_path, tmp_path) + + assert hf_ckpt_dir.exists(), f"HF checkpoint directory was not created: {hf_ckpt_dir}" + + config_files = list(hf_ckpt_dir.rglob("config.json")) + assert len(config_files) > 0, f"No config.json found in {hf_ckpt_dir}" + + weight_files = list(hf_ckpt_dir.rglob("*.safetensors")) + list(hf_ckpt_dir.rglob("*.bin")) + assert len(weight_files) > 0, f"No weight files (.safetensors or .bin) found in {hf_ckpt_dir}" diff --git a/experimental/conv/test_implicit_gemm.py b/tests/gpu/torch/quantization/kernels/test_implicit_gemm.py similarity index 64% rename from experimental/conv/test_implicit_gemm.py rename to tests/gpu/torch/quantization/kernels/test_implicit_gemm.py index af52660e42e..56ceaacc01f 100644 --- a/experimental/conv/test_implicit_gemm.py +++ b/tests/gpu/torch/quantization/kernels/test_implicit_gemm.py @@ -28,7 +28,7 @@ @pytest.fixture(scope="module") def cuda_conv3d(): """Import and return the CUDA implicit GEMM conv3d function.""" - from experimental.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda + from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import conv3d_implicit_gemm_cuda return conv3d_implicit_gemm_cuda @@ -303,9 +303,9 @@ def test_deterministic(self, cuda_conv3d): @pytest.fixture(scope="module") -def cuda_fp4_quant(): - """Import FP4 fake quant for reference comparisons.""" - from experimental.conv.implicit_gemm_cuda import fp4_fake_quant +def cuda_fp4(): + """Import and return the CUDA FP4 fake quant function.""" + from modelopt.torch.quantization.src.conv.implicit_gemm_cuda import fp4_fake_quant return fp4_fake_quant @@ -361,7 +361,7 @@ def test_quant_deterministic(self, cuda_conv3d, fp4_block_size): assert torch.equal(out1, out2), f"Non-deterministic for fp4_block_size={fp4_block_size}" @pytest.mark.parametrize("fp4_block_size", [16, 32, 64, 128, 256]) - def test_quant_vs_unfused_reference(self, cuda_conv3d, cuda_fp4_quant, fp4_block_size): + def test_quant_vs_unfused_reference(self, cuda_conv3d, cuda_fp4, fp4_block_size): """Compare fused kernel vs unfused: fp4(im2col) @ fp4(weight). Uses a shape where K is a multiple of 256 so all K-tiles are full @@ -382,9 +382,9 @@ def test_quant_vs_unfused_reference(self, cuda_conv3d, cuda_fp4_quant, fp4_block im2col = x.permute(0, 2, 3, 4, 1).reshape(-1, cin) # [M, K] # 2. FP4 fake-quant both matrices along K with the same block_size - im2col_q = cuda_fp4_quant(im2col, act_amax, fp4_block_size) + im2col_q = cuda_fp4(im2col, act_amax, fp4_block_size) w_flat = w.reshape(cout, cin).transpose(0, 1).contiguous() # [K, Cout] - w_flat_q = cuda_fp4_quant(w_flat, w_amax, fp4_block_size) + w_flat_q = cuda_fp4(w_flat, w_amax, fp4_block_size) # 3. Matmul (in BF16 to match kernel's WMMA path) ref_out = (im2col_q.bfloat16() @ w_flat_q.bfloat16()).float() @@ -421,81 +421,33 @@ def test_smaller_block_less_error(self, cuda_conv3d): """Smaller FP4 block sizes should generally produce lower quantization error. Finer-grained blocks capture local ranges better, reducing quant error vs cuDNN. - Test monotonicity: error(16) <= error(32) <= ... <= error(256) (with some tolerance). - Reports detailed accuracy metrics for each block size vs cuDNN baseline. + Test monotonicity on a medium config: error(16) <= error(64) <= error(256) (with 1.2x slack). """ torch.manual_seed(42) - # Test multiple shapes to get a comprehensive picture - configs = [ - ("Small K=432", 1, 16, 8, 8, 8, 32, 3, 3, 3), - ("Medium K=1728", 1, 64, 8, 8, 8, 64, 3, 3, 3), - ("Large K=3456", 1, 128, 5, 8, 8, 256, 3, 3, 3), - ] + # Medium K=1728: Cin=64, 3x3x3 kernel + cin, cout = 64, 64 + x = torch.randn(1, cin, 8, 8, 8, device="cuda", dtype=torch.float32) + w = torch.randn(cout, cin, 3, 3, 3, device="cuda", dtype=torch.float32) + act_amax = x.abs().max().unsqueeze(0) + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1)) block_sizes = [16, 32, 64, 128, 256] - all_errors = {} - - for desc, n, cin, d, h, w_s, cout, kd, kh, kw in configs: - x = torch.randn(n, cin, d, h, w_s, device="cuda", dtype=torch.float32) - w = torch.randn(cout, cin, kd, kh, kw, device="cuda", dtype=torch.float32) - act_amax = x.abs().max().unsqueeze(0) - k_size = cin * kd * kh * kw - - ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1)) - ref_abs_mean = ref.abs().mean().item() - - # Also compute no-quant baseline (BF16 rounding only) - out_nq = cuda_conv3d( + errors = {} + for bs in block_sizes: + out = cuda_conv3d( x, w, stride=(1, 1, 1), padding=(1, 1, 1), dilation=(1, 1, 1), - quant_act=False, - ) - nq_diff = (out_nq - ref).abs() - - print( - f"\n {desc} (K={k_size}), output range [{ref.min().item():.1f}, {ref.max().item():.1f}]" - ) - print( - f" {'Block Size':>10} | {'Max Diff':>10} | {'Mean Diff':>10} | {'RMSE':>10} | {'Rel Err%':>8}" - ) - print(f" {'-' * 10}-+-{'-' * 10}-+-{'-' * 10}-+-{'-' * 10}-+-{'-' * 8}") - print( - f" {'no-quant':>10} | {nq_diff.max().item():>10.4f} | " - f"{nq_diff.mean().item():>10.6f} | " - f"{((out_nq - ref) ** 2).mean().sqrt().item():>10.4f} | " - f"{nq_diff.mean().item() / ref_abs_mean * 100:>7.3f}%" + act_amax=act_amax, + quant_act=True, + fp4_block_size=bs, ) + errors[bs] = (out - ref).abs().mean().item() - errors = {} - for bs in block_sizes: - out = cuda_conv3d( - x, - w, - stride=(1, 1, 1), - padding=(1, 1, 1), - dilation=(1, 1, 1), - act_amax=act_amax, - quant_act=True, - fp4_block_size=bs, - ) - diff = (out - ref).abs() - max_d = diff.max().item() - mean_d = diff.mean().item() - rmse = ((out - ref) ** 2).mean().sqrt().item() - rel_err = mean_d / ref_abs_mean * 100 - errors[bs] = mean_d - print( - f" {bs:>10} | {max_d:>10.4f} | {mean_d:>10.6f} | " - f"{rmse:>10.4f} | {rel_err:>7.3f}%" - ) - all_errors[desc] = errors - - # Monotonicity check on the medium config - errors = all_errors["Medium K=1728"] + # Monotonicity: smaller blocks should have equal or lower error for smaller, larger in [(16, 64), (16, 256), (32, 256), (64, 256)]: assert errors[smaller] <= errors[larger] * 1.2, ( f"Expected error({smaller})={errors[smaller]:.6f} <= " @@ -590,14 +542,6 @@ def test_quant_realistic_shape(self, cuda_conv3d, fp4_block_size): # ============================================================================= -@pytest.fixture(scope="module") -def cuda_fp4(): - """Import and return the CUDA FP4 fake quant function.""" - from experimental.conv.implicit_gemm_cuda import fp4_fake_quant - - return fp4_fake_quant - - def _py_fp4_fake_quant_ref(x_flat, global_amax, block_size): """Pure Python reference for FP4 fake quant (no BF16 rounding). @@ -909,8 +853,8 @@ def _modelopt_dynamic_block_quantize_available(): class TestFP4FakeQuantVsModelopt: """Compare experimental CUDA FP4 fake quant against all modelopt FP4 implementations. - This ensures the standalone FP4 kernel in experimental/conv produces the same - results as the official modelopt quantization paths: + This ensures the standalone FP4 kernel produces the same results as the + other modelopt quantization paths: 1. Triton fp4_fake_quant_block (Hopper+ dynamic blockwise) 2. cuda_ext_mx.fused_amax_convert (CUDA extension fallback) 3. dynamic_block_quantize_op (high-level API that dispatches to either) @@ -1065,3 +1009,484 @@ def test_vs_triton_input_dtypes(self, cuda_fp4, dtype): # BF16/FP16 input rounding may cause small diffs tol = 1e-2 if dtype != torch.float32 else 1e-5 assert max_diff < tol, f"dtype={dtype}: experimental vs Triton max diff: {max_diff:.6e}" + + +# ============================================================================= +# Input Validation / Error Path Tests +# ============================================================================= + + +class TestConv3dInputValidation: + """Verify error paths raise appropriate exceptions.""" + + def test_invalid_fp4_block_size(self, cuda_conv3d): + """fp4_block_size not in {16, 32, 64, 128, 256} should raise ValueError.""" + x = torch.randn(1, 4, 4, 4, 4, device="cuda") + w = torch.randn(8, 4, 3, 3, 3, device="cuda") + with pytest.raises(ValueError, match="fp4_block_size"): + cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), fp4_block_size=7) + + def test_non_5d_input(self, cuda_conv3d): + """Non-5D tensors should raise ValueError.""" + x = torch.randn(1, 4, 4, 4, device="cuda") # 4D + w = torch.randn(8, 4, 3, 3, 3, device="cuda") + with pytest.raises(ValueError, match="5D"): + cuda_conv3d(x, w) + + def test_non_5d_weight(self, cuda_conv3d): + """Non-5D weight should raise ValueError.""" + x = torch.randn(1, 4, 4, 4, 4, device="cuda") + w = torch.randn(8, 4, 3, 3, device="cuda") # 4D + with pytest.raises(ValueError, match="5D"): + cuda_conv3d(x, w) + + def test_grouped_conv_error(self, cuda_conv3d): + """Mismatched Cin (groups > 1) should raise ValueError.""" + x = torch.randn(1, 8, 4, 4, 4, device="cuda") + w = torch.randn(8, 4, 3, 3, 3, device="cuda") # Cin=4 != x.Cin=8 + with pytest.raises(ValueError, match="Grouped convolution"): + cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + + def test_quant_act_without_amax(self, cuda_conv3d): + """quant_act=True without act_amax should raise ValueError.""" + x = torch.randn(1, 4, 4, 4, 4, device="cuda") + w = torch.randn(8, 4, 3, 3, 3, device="cuda") + with pytest.raises(ValueError, match="act_amax"): + cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=True, act_amax=None) + + def test_fp4_numel_not_divisible(self, cuda_fp4): + """fp4_fake_quant should error when numel is not divisible by block_size.""" + inp = torch.randn(17, device="cuda") + amax = torch.tensor([1.0], device="cuda") + with pytest.raises(AssertionError, match="divisible"): + cuda_fp4(inp, amax, block_size=16) + + +# ============================================================================= +# Input Dtype Tests +# ============================================================================= + + +class TestConv3dInputDtypes: + """Verify conv3d works with non-float32 inputs and preserves output dtype.""" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_dtype_preservation(self, cuda_conv3d, dtype): + """Output dtype should match input dtype.""" + x = torch.randn(1, 16, 8, 8, 8, device="cuda", dtype=dtype) + w = torch.randn(32, 16, 3, 3, 3, device="cuda", dtype=dtype) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_dtype_correctness(self, cuda_conv3d, dtype): + """Non-float32 inputs should produce correct results (vs F.conv3d in float32).""" + torch.manual_seed(42) + x_fp32 = torch.randn(1, 16, 8, 8, 8, device="cuda") + w_fp32 = torch.randn(32, 16, 3, 3, 3, device="cuda") + x = x_fp32.to(dtype) + w = w_fp32.to(dtype) + + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + ref = F.conv3d(x_fp32, w_fp32, stride=(1, 1, 1), padding=(1, 1, 1)) + + # Both BF16 input rounding and internal BF16 WMMA contribute to error + max_diff = (out.float() - ref).abs().max().item() + k_size = 16 * 27 + scaled_atol = ATOL * (k_size / 1000.0) ** 0.5 * 2 # extra slack for input rounding + assert max_diff < scaled_atol, f"dtype={dtype}: max diff {max_diff:.4f} > {scaled_atol:.4f}" + + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) + def test_dtype_quant_path(self, cuda_conv3d, dtype): + """FP4 quantized path should also work with non-float32 inputs.""" + x = torch.randn(1, 16, 8, 8, 8, device="cuda", dtype=dtype) + w = torch.randn(32, 16, 3, 3, 3, device="cuda", dtype=dtype) + act_amax = x.float().abs().max().unsqueeze(0) + + out = cuda_conv3d( + x, + w, + stride=(1, 1, 1), + padding=(1, 1, 1), + act_amax=act_amax, + quant_act=True, + fp4_block_size=16, + ) + assert out.dtype == dtype + assert not torch.isnan(out).any() + assert out.abs().max() > 0 + + +# ============================================================================= +# Non-Contiguous Input Tests +# ============================================================================= + + +class TestConv3dNonContiguous: + """Verify kernel handles non-contiguous tensors (via internal .contiguous() calls).""" + + def test_non_contiguous_input(self, cuda_conv3d): + """Permuted (non-contiguous) input should produce correct results.""" + torch.manual_seed(42) + # Create non-contiguous tensor via permute + permute back + x_base = torch.randn(1, 8, 8, 8, 16, device="cuda") + x = x_base.permute(0, 4, 1, 2, 3) # [1, 16, 8, 8, 8] but non-contiguous + assert not x.is_contiguous() + + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + x_contig = x.contiguous() + + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + ref = cuda_conv3d(x_contig, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert torch.equal(out, ref), "Non-contiguous input produced different results" + + def test_non_contiguous_weight(self, cuda_conv3d): + """Transposed (non-contiguous) weight should produce correct results.""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 8, 8, device="cuda") + # Create non-contiguous weight + w_base = torch.randn(16, 32, 3, 3, 3, device="cuda") + w = w_base.transpose(0, 1) # [32, 16, ...] but non-contiguous + assert not w.is_contiguous() + + w_contig = w.contiguous() + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + ref = cuda_conv3d(x, w_contig, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert torch.equal(out, ref), "Non-contiguous weight produced different results" + + +# ============================================================================= +# Combined Conv Parameter Tests +# ============================================================================= + + +class TestConv3dCombinedParams: + """Test combinations of stride + dilation + padding that were never combined.""" + + def test_stride_and_dilation(self, cuda_conv3d): + """Stride > 1 and dilation > 1 simultaneously.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (2, 2, 2)) + + def test_asymmetric_stride_and_padding(self, cuda_conv3d): + """Asymmetric stride with asymmetric padding.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 2, 2), (0, 1, 2), (1, 1, 1)) + + def test_all_non_default(self, cuda_conv3d): + """Non-default stride + padding + dilation all at once.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 2, 1), (1, 0, 1), (1, 2, 1)) + + def test_bias_with_stride(self, cuda_conv3d): + """Bias with non-default stride.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + b = torch.randn(32, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, b, (2, 2, 2), (1, 1, 1), (1, 1, 1)) + + def test_bias_with_dilation(self, cuda_conv3d): + """Bias with non-default dilation.""" + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + b = torch.randn(32, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, b, (1, 1, 1), (2, 2, 2), (2, 2, 2)) + + +# ============================================================================= +# FP4 Quantized Path: Advanced Conv Params +# ============================================================================= + + +def _run_quant_smoke_test(cuda_conv3d, x, w, bias, stride, padding, dilation, fp4_block_size=16): + """Helper: run FP4-quantized conv and verify basic sanity.""" + act_amax = x.abs().max().unsqueeze(0) + out = cuda_conv3d( + x, + w, + bias=bias, + stride=stride, + padding=padding, + dilation=dilation, + act_amax=act_amax, + quant_act=True, + fp4_block_size=fp4_block_size, + ) + ref = F.conv3d(x, w, bias=bias, stride=stride, padding=padding, dilation=dilation) + assert out.shape == ref.shape, f"Shape mismatch: {out.shape} vs {ref.shape}" + assert not torch.isnan(out).any(), "Output contains NaN" + assert not torch.isinf(out).any(), "Output contains Inf" + # Quantized output should be in a reasonable range relative to reference + if ref.abs().max() > 0: + ratio = out.abs().max().item() / ref.abs().max().item() + assert 0.01 < ratio < 100, f"Output magnitude ratio {ratio:.2f} is unreasonable" + return out + + +class TestConv3dFP4QuantAdvanced: + """FP4 quantized path with non-trivial stride, dilation, and kernel shapes.""" + + def test_quant_with_stride(self, cuda_conv3d): + """FP4 quant with stride=(2,2,2).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (1, 1, 1)) + + def test_quant_with_dilation(self, cuda_conv3d): + """FP4 quant with dilation=(2,2,2).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (2, 2, 2), (2, 2, 2)) + + def test_quant_with_asymmetric_kernel(self, cuda_conv3d): + """FP4 quant with 1x3x3 kernel.""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 16, 16, device="cuda") + w = torch.randn(32, 16, 1, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 1, 1), (1, 1, 1)) + + def test_quant_with_stride_and_dilation(self, cuda_conv3d): + """FP4 quant with both stride>1 and dilation>1.""" + torch.manual_seed(42) + x = torch.randn(1, 16, 16, 16, 16, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (2, 2, 2), (1, 1, 1), (2, 2, 2)) + + def test_quant_with_no_padding(self, cuda_conv3d): + """FP4 quant with padding=(0,0,0).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 8, 8, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + _run_quant_smoke_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1)) + + def test_quant_bias_reference(self, cuda_conv3d, cuda_fp4): + """FP4 quant + bias: verify bias is added correctly by comparing with/without. + + The difference between bias and no-bias output should equal the bias broadcast. + """ + torch.manual_seed(42) + cin, cout = 256, 64 + x = torch.randn(1, cin, 4, 4, 4, device="cuda") + w = torch.randn(cout, cin, 1, 1, 1, device="cuda") + b = torch.randn(cout, device="cuda") + act_amax = x.abs().max().unsqueeze(0) + + kwargs = { + "stride": (1, 1, 1), + "padding": (0, 0, 0), + "dilation": (1, 1, 1), + "act_amax": act_amax, + "quant_act": True, + "fp4_block_size": 16, + } + out_bias = cuda_conv3d(x, w, bias=b, **kwargs) + out_no_bias = cuda_conv3d(x, w, bias=None, **kwargs) + + # Difference should be the bias broadcast over spatial dims + diff = out_bias - out_no_bias # [1, Cout, D, H, W] + expected_bias = b.view(1, -1, 1, 1, 1).expand_as(diff) + assert torch.allclose(diff, expected_bias, atol=1e-5), ( + f"Bias diff mismatch: max {(diff - expected_bias).abs().max().item():.6e}" + ) + + +# ============================================================================= +# Zero / Degenerate Input Tests +# ============================================================================= + + +class TestConv3dZeroInputs: + """Tests with zero and degenerate inputs.""" + + def test_zero_input(self, cuda_conv3d): + """Zero activation tensor should produce zero (or bias-only) output.""" + x = torch.zeros(1, 16, 8, 8, 8, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert torch.allclose(out, ref, atol=1e-5), f"Max diff: {(out - ref).abs().max().item()}" + + def test_zero_weight(self, cuda_conv3d): + """Zero weight tensor should produce zero output.""" + x = torch.randn(1, 16, 8, 8, 8, device="cuda") + w = torch.zeros(32, 16, 3, 3, 3, device="cuda") + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert torch.allclose(out, torch.zeros_like(out), atol=1e-5) + + def test_zero_input_quant(self, cuda_conv3d): + """Zero input with FP4 quant should not produce NaN.""" + x = torch.zeros(1, 16, 8, 8, 8, device="cuda") + w = torch.randn(32, 16, 3, 3, 3, device="cuda") + # act_amax=0 is a tricky edge case — the kernel's scale guard should handle it + act_amax = torch.tensor([1e-10], device="cuda") # near-zero but not exactly 0 + out = cuda_conv3d( + x, + w, + stride=(1, 1, 1), + padding=(1, 1, 1), + act_amax=act_amax, + quant_act=True, + fp4_block_size=16, + ) + assert not torch.isnan(out).any(), "Zero input with quant produced NaN" + + +# ============================================================================= +# Numerical Stability Tests +# ============================================================================= + + +class TestConv3dNumericalStability: + """Test with extreme value ranges.""" + + def test_large_values(self, cuda_conv3d): + """Large input values (randn * 100).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 8, 8, device="cuda") * 100 + w = torch.randn(32, 16, 3, 3, 3, device="cuda") * 100 + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + # With large values, BF16 rounding error scales proportionally + rel_err = (out - ref).abs().max().item() / ref.abs().max().item() + assert rel_err < 0.05, f"Relative error {rel_err:.4f} too high for large values" + + def test_small_values(self, cuda_conv3d): + """Small input values (randn * 1e-3).""" + torch.manual_seed(42) + x = torch.randn(1, 16, 8, 8, 8, device="cuda") * 1e-3 + w = torch.randn(32, 16, 3, 3, 3, device="cuda") * 1e-3 + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + assert out.shape == ref.shape + # Small values: absolute error is small, relative error may be larger due to BF16 + max_diff = (out - ref).abs().max().item() + assert max_diff < 1e-5, f"Max diff {max_diff:.6e} for small values" + + def test_uniform_input(self, cuda_conv3d): + """Uniform input (all ones) — exposes accumulation patterns.""" + x = torch.ones(1, 16, 8, 8, 8, device="cuda") + w = torch.ones(32, 16, 3, 3, 3, device="cuda") + ref = F.conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1)) + out = cuda_conv3d(x, w, stride=(1, 1, 1), padding=(1, 1, 1), quant_act=False) + k_size = 16 * 27 + scaled_atol = ATOL * (k_size / 1000.0) ** 0.5 + max_diff = (out - ref).abs().max().item() + assert max_diff < scaled_atol, f"Uniform input: max diff {max_diff:.4f}" + + +# ============================================================================= +# Exact Block Boundary Tests +# ============================================================================= + + +class TestConv3dExactBoundaries: + """Shapes that land exactly on BLOCK_M=64, BLOCK_N=64, BLOCK_K=256 boundaries.""" + + def test_m_exact_128(self, cuda_conv3d): + """M = 128 = 2 * BLOCK_M (exactly 2 M-tiles, no remainder).""" + # batch=1, output 4x4x8 = 128 with kernel 1x1x1 + x = torch.randn(1, 32, 4, 4, 8, device="cuda") + w = torch.randn(64, 32, 1, 1, 1, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1)) + + def test_k_exact_512(self, cuda_conv3d): + """K = 512 = 2 * BLOCK_K (exactly 2 K-tiles, no remainder).""" + # Cin=512, kernel 1x1x1 -> K=512 + x = torch.randn(1, 512, 4, 4, 4, device="cuda") + w = torch.randn(64, 512, 1, 1, 1, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1)) + + def test_cout_exact_64(self, cuda_conv3d): + """Cout = 64 = 1 * BLOCK_N (exactly 1 N-tile).""" + x = torch.randn(1, 16, 8, 8, 8, device="cuda") + w = torch.randn(64, 16, 3, 3, 3, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (1, 1, 1), (1, 1, 1)) + + def test_all_exact_multiples(self, cuda_conv3d): + """M=64, N=64, K=256 — single tile in each dimension.""" + # batch=1, Cin=256, kernel 1x1x1 -> K=256; output 4x4x4=64; Cout=64 + x = torch.randn(1, 256, 4, 4, 4, device="cuda") + w = torch.randn(64, 256, 1, 1, 1, device="cuda") + _run_conv3d_test(cuda_conv3d, x, w, None, (1, 1, 1), (0, 0, 0), (1, 1, 1)) + + +# ============================================================================= +# FP4 Fake Quant: Shape and Edge Case Tests +# ============================================================================= + + +class TestFP4FakeQuantShapes: + """Test fp4_fake_quant with multi-dimensional inputs.""" + + def test_3d_shape_preservation(self, cuda_fp4): + """3D input should preserve shape after quantization.""" + inp = torch.randn(4, 8, 32, device="cuda") # numel=1024 + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=16) + assert out.shape == (4, 8, 32) + + def test_4d_shape_preservation(self, cuda_fp4): + """4D input should preserve shape.""" + inp = torch.randn(2, 4, 8, 16, device="cuda") # numel=1024 + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=16) + assert out.shape == (2, 4, 8, 16) + + def test_5d_shape_preservation(self, cuda_fp4): + """5D input (like a Conv3D activation) should preserve shape.""" + inp = torch.randn(1, 4, 4, 4, 16, device="cuda") # numel=1024 + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=16) + assert out.shape == (1, 4, 4, 4, 16) + + def test_multidim_correctness(self, cuda_fp4): + """Multi-dim quantization should equal flatten -> quant -> reshape.""" + torch.manual_seed(42) + inp = torch.randn(4, 8, 32, device="cuda") + amax = inp.abs().max().unsqueeze(0) + + out_3d = cuda_fp4(inp, amax, block_size=16) + out_flat = cuda_fp4(inp.reshape(-1), amax, block_size=16).reshape(4, 8, 32) + assert torch.equal(out_3d, out_flat) + + +class TestFP4FakeQuantEdgeCases: + """Edge cases for fp4_fake_quant.""" + + def test_very_large_values(self, cuda_fp4): + """Very large input values should saturate to max E2M1 level, not produce NaN.""" + inp = torch.tensor([1e6, -1e6, 5e5, -5e5, 1e4, -1e4, 100, -100], device="cuda") + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=8) + assert not torch.isnan(out).any() + assert not torch.isinf(out).any() + + def test_very_small_values(self, cuda_fp4): + """Very small input values should quantize to zero or near-zero.""" + inp = torch.tensor([1e-8, -1e-8, 1e-10, -1e-10, 1e-6, -1e-6, 0, 0], device="cuda") + amax = torch.tensor([1.0], device="cuda") + out = cuda_fp4(inp, amax, block_size=8) + assert not torch.isnan(out).any() + # Very small values relative to amax should quantize to ~0 + assert out.abs().max() < 1e-3 + + def test_uniform_block(self, cuda_fp4): + """All-same-value block.""" + inp = torch.full((16,), 3.0, device="cuda") + amax = inp.abs().max().unsqueeze(0) + out = cuda_fp4(inp, amax, block_size=16) + # All elements are the same, so they should all quantize to the same E2M1 level + assert (out == out[0]).all(), f"Uniform block produced non-uniform output: {out}" + + def test_near_zero_amax(self, cuda_fp4): + """Very small global_amax should not produce NaN/Inf.""" + inp = torch.randn(16, device="cuda") * 1e-8 + amax = torch.tensor([1e-10], device="cuda") + out = cuda_fp4(inp, amax, block_size=16) + assert not torch.isnan(out).any(), "Near-zero amax produced NaN" + assert not torch.isinf(out).any(), "Near-zero amax produced Inf" diff --git a/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py b/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py new file mode 100644 index 00000000000..dcd8f24ab13 --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_diffusers_wan_conv3d.py @@ -0,0 +1,129 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for the quantized WanCausalConv3d wrapper in the diffusers plugin. + +WanCausalConv3d applies asymmetric causal padding before the underlying Conv3D, +so the quantized subclass has to replicate that padding logic around the +quantized forward. These tests run on CPU with the NVFP4 dispatch disabled (by +only exercising the default quantized path); the NVFP4 implicit-GEMM path is +covered by the GPU tests under ``tests/gpu/torch/quantization/kernels``. +""" + +import pytest +import torch + +pytest.importorskip("onnx") + +from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d + +# Triggers registration of _QuantDiffusersWanCausalConv3d. +import modelopt.torch.quantization.plugins.diffusion.diffusers # noqa: F401 +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins.diffusion.diffusers import _QuantDiffusersWanCausalConv3d + + +def _make_quantized(in_ch: int = 4, out_ch: int = 6, padding=(1, 1, 1)) -> torch.nn.Module: + m = WanCausalConv3d(in_ch, out_ch, kernel_size=3, padding=padding) + m.eval() + # Convert via the registry so the generated class picks up WanCausalConv3d + # in its MRO (needed for ``nn.Conv3d._conv_forward`` to be reachable). + return QuantModuleRegistry.convert(m) + + +def _strip_quant_state(state_dict: dict) -> dict: + return { + k: v + for k, v in state_dict.items() + if not any( + k.startswith(p) for p in ("input_quantizer", "weight_quantizer", "output_quantizer") + ) + } + + +class TestQuantWanCausalConv3dRegistration: + def test_registered(self): + assert WanCausalConv3d in QuantModuleRegistry + # Registry-returned class is a generated subclass of our dm class. + assert issubclass(QuantModuleRegistry[WanCausalConv3d], _QuantDiffusersWanCausalConv3d) + + def test_convert_preserves_type_identity(self): + m = _make_quantized() + mro = [c.__name__ for c in type(m).__mro__] + assert "WanCausalConv3d" in mro + assert "_QuantDiffusersWanCausalConv3d" in mro + + def test_no_implicit_gemm_without_nvfp4(self): + # Default quantizer config is INT8 — must NOT route through implicit GEMM. + m = _make_quantized() + assert not m._should_use_implicit_gemm() + + +class TestQuantWanCausalConv3dForward: + """Exercise the default quantized path (NVFP4 kernel is GPU-only). + + We disable the quantizers and assert output matches an unquantized + ``WanCausalConv3d`` with the same weights — this verifies the causal-padding + logic in the overridden ``forward`` is preserved after conversion. + """ + + @pytest.mark.parametrize("padding", [(1, 1, 1), (2, 0, 0), (0, 1, 1)]) + def test_matches_unquantized_no_cache(self, padding): + torch.manual_seed(0) + m = _make_quantized(in_ch=4, out_ch=6, padding=padding) + m.input_quantizer.disable() + m.weight_quantizer.disable() + + m_ref = WanCausalConv3d(4, 6, kernel_size=3, padding=padding) + m_ref.eval() + m_ref.load_state_dict(_strip_quant_state(m.state_dict()), strict=False) + + x = torch.randn(1, 4, 5, 6, 6) + out = m(x) + out_ref = m_ref(x) + assert torch.allclose(out, out_ref, atol=1e-5), ( + f"Max diff: {(out - out_ref).abs().max().item()}" + ) + + def test_matches_unquantized_with_cache_x(self): + """cache_x is the temporal-cache branch used during causal decoding.""" + torch.manual_seed(0) + m = _make_quantized(padding=(1, 1, 1)) # _padding[4] == 2 > 0 → cache path active + m.input_quantizer.disable() + m.weight_quantizer.disable() + + m_ref = WanCausalConv3d(4, 6, kernel_size=3, padding=(1, 1, 1)) + m_ref.eval() + m_ref.load_state_dict(_strip_quant_state(m.state_dict()), strict=False) + + x = torch.randn(1, 4, 5, 6, 6) + cache_x = torch.randn(1, 4, 1, 6, 6) + out = m(x, cache_x=cache_x) + out_ref = m_ref(x, cache_x=cache_x) + assert torch.allclose(out, out_ref, atol=1e-5) + + def test_output_quantizer_applied(self): + """Enabling the output quantizer must change the forward output.""" + torch.manual_seed(0) + m = _make_quantized() + m.input_quantizer.disable() + m.weight_quantizer.disable() + # Output quantizer is disabled by default; enable it and check it takes effect. + x = torch.randn(1, 4, 3, 5, 5) + out_disabled = m(x) + m.output_quantizer.enable() + out_enabled = m(x) + # INT8 default config clamps and rounds; at least some elements differ. + assert not torch.allclose(out_disabled, out_enabled) diff --git a/tests/unit/torch/quantization/test_quant_conv.py b/tests/unit/torch/quantization/test_quant_conv.py index be872de206d..ee1d449d6e6 100644 --- a/tests/unit/torch/quantization/test_quant_conv.py +++ b/tests/unit/torch/quantization/test_quant_conv.py @@ -15,6 +15,8 @@ """Tests of QuantConv module.""" +import warnings + import pytest import torch import torch.nn.functional as F @@ -23,11 +25,19 @@ from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.nn.modules import quant_conv +from modelopt.torch.quantization.nn.modules.quant_conv import ( + _is_nvfp4_quantizer, + _nvfp4_quantize_weight_along_k, +) from modelopt.torch.quantization.nn.modules.tensor_quantizer import TensorQuantizer NUM_IN_CHANNELS = 3 NUM_OUT_CHANNELS = 5 +_NVFP4_CFG = QuantizerAttributeConfig( + num_bits=(2, 1), block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)} +) + class TestQuantConvND: @pytest.mark.parametrize( @@ -316,3 +326,134 @@ def test_against_unquantized(self, conv_cls, nn_conv_cls, input_shape): output = conv(test_input) assert torch.allclose(quant_output, output) + + +class TestQuantConv3dNVFP4: + """Tests for the NVFP4 implicit GEMM dispatch path in ``_QuantConv3d``. + + The CUDA kernel itself is GPU-only (covered in ``tests/gpu``); here we only + exercise the CPU-side predicate and fallback branches that currently lack + coverage. + """ + + @staticmethod + def _make_nvfp4_conv3d(groups: int = 1, bias: bool = False) -> quant_conv.QuantConv3d: + return quant_conv.QuantConv3d( + NUM_IN_CHANNELS * groups, + NUM_OUT_CHANNELS * groups, + kernel_size=3, + groups=groups, + bias=bias, + quant_desc_input=_NVFP4_CFG, + quant_desc_weight=_NVFP4_CFG, + ) + + def test_is_nvfp4_quantizer_true(self): + q = TensorQuantizer(_NVFP4_CFG) + assert _is_nvfp4_quantizer(q) + + def test_is_nvfp4_quantizer_false_for_int8(self): + q = TensorQuantizer(QuantizerAttributeConfig(num_bits=8)) + assert not _is_nvfp4_quantizer(q) + + def test_is_nvfp4_quantizer_false_for_fp8(self): + # FP8 E4M3: num_bits == (4, 3), block_sizes is None + q = TensorQuantizer(QuantizerAttributeConfig(num_bits=(4, 3))) + assert not _is_nvfp4_quantizer(q) + + def test_is_nvfp4_quantizer_false_for_static_block(self): + static_cfg = QuantizerAttributeConfig( + num_bits=(2, 1), + block_sizes={-1: 16, "type": "static", "scale_bits": (4, 3)}, + ) + q = TensorQuantizer(static_cfg) + assert not _is_nvfp4_quantizer(q) + + def test_should_use_implicit_gemm_true(self): + m = self._make_nvfp4_conv3d() + assert m._should_use_implicit_gemm() + + def test_should_use_implicit_gemm_false_groups_gt_1(self): + m = self._make_nvfp4_conv3d(groups=NUM_IN_CHANNELS) + # Groups > 1 disqualifies the fused kernel even with NVFP4 quantizers. + assert not m._should_use_implicit_gemm() + + def test_should_use_implicit_gemm_false_non_nvfp4(self): + # INT8 per-tensor config — default Conv3D quantizers, no NVFP4. + m = quant_conv.QuantConv3d(NUM_IN_CHANNELS, NUM_OUT_CHANNELS, kernel_size=3, bias=False) + assert not m._should_use_implicit_gemm() + + def test_nvfp4_quantize_weight_along_k_reshape(self): + """Verify weight is flattened/restored along the K (input) dimension. + + Uses a stub quantizer so this test stays CPU-only (the real NVFP4 dynamic + quantizer requires a CUDA tensor). + """ + + def identity_quantizer(x): + # K-dim must be the last axis when passed to the quantizer so NVFP4's + # block-wise scaling aligns with the GEMM reduction axis. + assert x.dim() == 2 and x.shape[0] == NUM_OUT_CHANNELS + return x + + w = torch.randn(NUM_OUT_CHANNELS, NUM_IN_CHANNELS, 3, 3, 3) + qw = _nvfp4_quantize_weight_along_k(w, identity_quantizer) + assert qw.shape == w.shape + assert torch.equal(qw, w) + + def test_forward_non_nvfp4_matches_unquantized(self): + # Disabled quantizers: forward must match plain conv3d exactly. + m = quant_conv.QuantConv3d(NUM_IN_CHANNELS, NUM_OUT_CHANNELS, kernel_size=3, bias=False) + m.input_quantizer.disable() + m.weight_quantizer.disable() + x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4) + out = m(x) + ref = F.conv3d(x, m.weight) + assert torch.allclose(out, ref) + + def test_forward_nvfp4_training_warns_and_falls_back(self): + """Training mode must fall back to the default (cuDNN) path with a warning. + + The implicit-GEMM kernel is inference-only; this exercises the CPU-visible + training-fallback branch. We disable the quantizers so the default-path + ``super().forward()`` does not try to run NVFP4 dynamic quantization + (which requires CUDA). + """ + m = self._make_nvfp4_conv3d() + # NVFP4 predicate reads configuration (num_bits/block_sizes), not enable + # state, so disabling the quantizers still routes through the NVFP4 branch. + m.input_quantizer.disable() + m.weight_quantizer.disable() + assert m._should_use_implicit_gemm() + m.train() + x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + out = m(x) + assert out.shape == (1, NUM_OUT_CHANNELS, 2, 2, 2) + assert any("inference-only" in str(w.message) for w in caught), ( + f"Expected an 'inference-only' warning, got: {[str(w.message) for w in caught]}" + ) + + def test_forward_nvfp4_calib_only_uses_default_path(self): + """Calibration-only mode must NOT try to invoke the CUDA kernel. + + When the input quantizer is in calibration mode without quant enabled, + the forward must fall back to the default path. This exercises the + calib-only early-return branch; we assert that the output matches the + default-path output exactly (i.e. the implicit-GEMM path wasn't used). + """ + m = self._make_nvfp4_conv3d() + m.eval() + # Match the state toggled by TensorQuantizer.disable_quant()/enable_calib(). + m.input_quantizer.disable_quant() + m.input_quantizer.enable_calib() + m.weight_quantizer.disable_quant() + m.weight_quantizer.enable_calib() + assert m.input_quantizer._if_calib and not m.input_quantizer._if_quant + + x = torch.randn(1, NUM_IN_CHANNELS, 4, 4, 4) + out = m(x) + # Default path with quant disabled should equal plain conv3d. + ref = F.conv3d(x, m.weight.detach()) + assert torch.allclose(out, ref, atol=1e-5) From 010b220dc09890bf7646f8fad5b69fd0bc41ac5b Mon Sep 17 00:00:00 2001 From: kinjalpatel27 <31936134+kinjalpatel27@users.noreply.github.com> Date: Sun, 19 Apr 2026 22:12:45 -0700 Subject: [PATCH 23/30] vLLM fakequant export update for AWQ checkpoint (#1242) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug Enables end-to-end AWQ checkpoint export and reload in the vLLM fake-quant serving path (`MODELOPT_STATE_PATH`). Previously, the `input_quantizer` was using incorrect `pre_quant_scale` especially with grouped quantizers like `qkv_proj`, using simply the first `input_quantizer.pre_quant_scale`. This MR adds `_resmooth_experts_for_export` that non-mutatively averages `pre_quant_scale` across MoE experts and unifies input `_amax`, required because vLLM uses a single input quantizer per expert group. Adds `merge_amax_tensors_for_group` (element-wise max for same-shape, `cat` for GQA, scalar-max fallback) replacing the scalar-collapsing `torch.stack().max()` that dropped per-channel `_amax` structure. ### Usage ```python # Export AWQ checkpoint from HF model from modelopt.torch.export.plugins.vllm_fakequant_hf import export_hf_vllm_fq_checkpoint export_hf_vllm_fq_checkpoint(model, export_dir="./awq_vllm_checkpoint") ``` ### Testing **Step 1 — Export the quantized checkpoint:** ```bash python examples/llm_ptq/hf_ptq.py \ --pyt_ckpt_path \ --recipe \ --calib_size 512 \ --export_path \ --vllm_fakequant_export ``` This produces `/vllm_fq_modelopt_state.pth` with the averaged per-expert pre_quant_scale and unified _amax now included. Step 2 — Serve via vLLM fakequant worker: ```bash MODELOPT_STATE_PATH=/vllm_fq_modelopt_state.pth \ python examples/vllm_serve/vllm_serve_fakequant.py \ --tensor-parallel-size ``` Tested for quantization configurations: ``` FP8_DEFAULT_CFG FP8_DEFAULT_CFG (input_q disabled) INT8_SMOOTHQUANT_CFG INT8_WEIGHT_ONLY_CFG NVFP4_DEFAULT_CFG NVFP4_AWQ_LITE_CFG INT4_AWQ_CFG NVFP4_AWQ_CFG NVFP4_DEFAULT_CFG (input_q disabled) ``` ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A ### Additional Information ## Summary by CodeRabbit * **New Features** * Added Nemotron-style MoE export support and group-aware AWQ resmoothing with optional requantization during export. * Improved handling for shared-input / expert groups and tensor-parallel sharding of pre-quantization scales. * **Bug Fixes** * Removed AWQ reload limitation from known issues; improved checkpoint validation and safer save/load behavior. * Better detection and handling of enabled weight-quantizers and clearer warnings for mismatched checkpoint keys. --------- Signed-off-by: Kinjal Patel --- examples/vllm_serve/README.md | 5 +- examples/vllm_serve/fakequant_worker.py | 53 +- examples/vllm_serve/vllm_reload_utils.py | 229 +++++-- modelopt/torch/export/layer_utils.py | 19 +- .../torch/export/plugins/vllm_fakequant_hf.py | 562 +++++++++++++++--- modelopt/torch/export/unified_export_hf.py | 6 +- .../export/test_vllm_fakequant_hf_export.py | 42 +- .../export/test_vllm_quantizer_reload.py | 156 +++++ 8 files changed, 911 insertions(+), 161 deletions(-) create mode 100644 tests/unit/torch/export/test_vllm_quantizer_reload.py diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 3d96ebb46a2..6513b5b04dc 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -98,6 +98,5 @@ QUANT_CFG= QUANT_FILE_PATH= python vllm_serve_fa ## Known Problems 1. **MCore reload does not use `MODELOPT_STATE_PATH`**; use `QUANT_FILE_PATH` and make sure `QUANT_CFG` matches the quantization recipe used for the original MCore model (otherwise quantizer keys/config won’t align). -2. AWQ reload is not supported yet -3. KV cache quantization export and reload is not supported in MCore yet. -4. **`NVFP4_KV_CFG` and `NVFP4_AFFINE_KV_CFG` require `--enforce-eager`**; these configs use a dynamic-block Triton kernel for KV-cache quantization that is incompatible with CUDA graph capture (the kernel grid is computed from Python-level tensor shapes, which get baked in at capture time). Without `--enforce-eager`, the captured grid will be wrong for different batch sizes, producing incorrect outputs. +2. KV cache quantization export and reload is not supported in MCore yet. +3. **`NVFP4_KV_CFG` and `NVFP4_AFFINE_KV_CFG` require `--enforce-eager`**; these configs use a dynamic-block Triton kernel for KV-cache quantization that is incompatible with CUDA graph capture (the kernel grid is computed from Python-level tensor shapes, which get baked in at capture time). Without `--enforce-eager`, the captured grid will be wrong for different batch sizes, producing incorrect outputs. diff --git a/examples/vllm_serve/fakequant_worker.py b/examples/vllm_serve/fakequant_worker.py index b88af9c72ee..4f84df0581d 100644 --- a/examples/vllm_serve/fakequant_worker.py +++ b/examples/vllm_serve/fakequant_worker.py @@ -15,6 +15,7 @@ import os +import warnings from typing import Any import torch @@ -26,13 +27,16 @@ convert_modelopt_state_to_vllm, load_state_dict_from_path, restore_from_modelopt_state_vllm, + shard_pre_quant_scale_for_tp, ) import modelopt.torch.quantization as mtq +from modelopt.torch.export.plugins.vllm_fakequant_hf import is_weight_quantizer_state_key from modelopt.torch.quantization.plugins.vllm import ( disable_compilation, post_restore_vllm_parallel_linears, ) +from modelopt.torch.utils import safe_load from modelopt.torch.utils.dataset_utils import get_dataset_dataloader quant_config: dict[str, Any] = { @@ -61,28 +65,48 @@ def _fakequant_run_prolog_worker(self) -> None: model = model.unwrap() if quant_config["modelopt_state_path"]: print(f"Loading modelopt state from {quant_config['modelopt_state_path']}") - # Load on CPU to avoid failures when the checkpoint was saved from a different - # GPU mapping - modelopt_state = torch.load( - quant_config["modelopt_state_path"], weights_only=True, map_location="cpu" - ) + # Load on CPU to avoid failures when the checkpoint was saved from a different GPU mapping. + modelopt_state = safe_load(quant_config["modelopt_state_path"], map_location="cpu") modelopt_weights = modelopt_state.pop("modelopt_state_weights", None) map_fun = ( self.model_runner.model.hf_to_vllm_mapper.apply_dict if hasattr(self.model_runner.model, "hf_to_vllm_mapper") else None ) - # convert modelopt state to vllm format modelopt_state = convert_modelopt_state_to_vllm(modelopt_state, map_fun=map_fun) - # restore model from modelopt state restore_from_modelopt_state_vllm(model, modelopt_state) if modelopt_weights is not None: - # convert quantizer state values to vllm format modelopt_weights = convert_dict_to_vllm(modelopt_weights, map_fun=map_fun) mtq.utils.set_quantizer_state_dict(model, modelopt_weights) - # set_quantizer_state_dict does not invoke modelopt_post_restore (unlike restore_quantizer_state). + if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + from modelopt.torch.quantization.nn import TensorQuantizer + from modelopt.torch.utils import get_unwrapped_name + + loaded_keys = { + get_unwrapped_name(n, model) + for n, m in model.named_modules() + if isinstance(m, TensorQuantizer) + } + # Same namespace as ``loaded_keys``: checkpoint keys may include DDP/FSDP + # prefixes that ``convert_dict_to_vllm`` does not strip. + pqs_in_weights = { + get_unwrapped_name(k, model) + for k, v in modelopt_weights.items() + if isinstance(v, dict) and "_pre_quant_scale" in v + } + unmatched_pqs = pqs_in_weights - loaded_keys + if unmatched_pqs: + sample = sorted(unmatched_pqs)[:20] + warnings.warn( + f"{len(unmatched_pqs)} checkpoint pre_quant_scale key(s) have no " + f"matching TensorQuantizer in the model (showing up to 20): {sample}", + stacklevel=2, + ) + # set_quantizer_state_dict does not run modelopt_post_restore (unlike restore_quantizer_state). post_restore_vllm_parallel_linears(model) + # Must follow post_restore: shard_pre_quant_scale_for_tp uses weight H_in vs pqs length. + shard_pre_quant_scale_for_tp(model) else: if quant_config["quant_file_path"]: @@ -101,15 +125,13 @@ def _fakequant_run_prolog_worker(self) -> None: quant_cfg = get_quant_config(quant_config, model) - # quantize model with disable_compilation(model): print("Quantizing model...") mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop) quantizer_file_path = quant_config["quant_file_path"] if quantizer_file_path: - # Get amax and other quantizer state from the quantizer file - # this can be used with Megatron-LM exported model using export_mcore_gpt_to_hf_vllm_fq + self.model_runner._dummy_run(1) current_state_dict = load_state_dict_from_path(self, quantizer_file_path, model) model.load_state_dict(current_state_dict) @@ -122,8 +144,11 @@ def _fakequant_run_prolog_worker(self) -> None: mtq.fold_weight(model) for name, module in model.named_modules(): - if name.endswith("weight_quantizer"): - assert not module.is_enabled, f"quantizer {name} is still enabled" + if is_weight_quantizer_state_key(name) and module.is_enabled: + raise RuntimeError( + f"Weight quantizer {name!r} is still enabled after fold_weight — " + "double-quantization would corrupt activations." + ) class FakeQuantWorker(BaseWorker): diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 2b59d1be2bd..aa8d3a5388b 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -22,6 +22,11 @@ import torch from vllm.distributed.parallel_state import get_tp_group +from modelopt.torch.export.plugins.vllm_fakequant_hf import ( + infer_quantizer_prefix_remap, + is_weight_quantizer_state_key, + merge_amax_tensors_for_group, +) from modelopt.torch.opt.conversion import ( ModelLikeModule, ModeloptStateManager, @@ -84,7 +89,7 @@ def _convert_key_for_vllm(key: str, value: Any) -> tuple[str, str | None, Any]: if "quantizer" not in key: return ("copy", key, value) - # Skip softmax_quantizer and lm_head quantizers(not needed in vLLM) + # Skip softmax_quantizer and lm_head quantizers (not needed in vLLM). if "softmax_quantizer" in key or (key.startswith("lm_head.") and "quantizer" in key): return ("skip", None, None) @@ -95,8 +100,7 @@ def _convert_key_for_vllm(key: str, value: Any) -> tuple[str, str | None, Any]: group_key = qkv_match.group(1) + "qkv_proj." + qkv_match.group(3) + suffix return ("group", group_key, value) - # Check if this is an expert gate/up projection - # if "mixer" not in key: + # Expert gate/up (per-expert) → w13 merge expert_gate_up_match = re.search( r"(.*\.experts)\.\d+\.(gate|up)_proj\.([^.]+_quantizer)(\..+)?$", key ) @@ -113,8 +117,6 @@ def _convert_key_for_vllm(key: str, value: Any) -> tuple[str, str | None, Any]: group_key = gate_up_match.group(1) + "gate_up_proj." + gate_up_match.group(3) + suffix return ("group", group_key, value) - # Check if this is an expert down_proj - # if "mixer" not in key: expert_down_match = re.search(r"(.*\.experts)\.\d+\.down_proj\.([^.]+_quantizer)(\..+)?$", key) if expert_down_match: suffix = expert_down_match.group(3) or "" @@ -148,9 +150,10 @@ def _group_keys_for_vllm( for key, value in state_dict.items(): action, new_key, new_value = _convert_key_for_vllm(key, value) if new_key is None or new_value is None: - assert action == "skip", ( - f"Expected action to be 'skip' for key {key}, value {value}, got {action}" - ) + if action != "skip": + raise RuntimeError( + f"Expected action to be 'skip' for key {key}, value {value}, got {action}" + ) continue if action == "copy": vllm_state_dict[new_key] = new_value @@ -176,7 +179,7 @@ def _merge_values_by_max_or_concat(merged_key: str, key_value_pairs: list[tuple[ for dict_key in values[0]: tensors = [v[dict_key] for v in values] if "_amax" in dict_key: - merged_value[dict_key] = torch.stack(tensors).max(dim=0)[0] + merged_value[dict_key] = merge_amax_tensors_for_group(tensors) elif "_pre_quant_scale" in dict_key: # _pre_quant_scale is per-input-channel: identical across q/k/v projections # since they share the same input. Do not concatenate; take the first value. @@ -187,7 +190,7 @@ def _merge_values_by_max_or_concat(merged_key: str, key_value_pairs: list[tuple[ else: # Values are tensors directly if "_amax" in merged_key: - merged_value = torch.stack(values).max(dim=0)[0] + merged_value = merge_amax_tensors_for_group(values) else: merged_value = torch.cat(values, dim=0) return merged_value @@ -231,6 +234,25 @@ def convert_dict_to_vllm( max_or_concat: Whether to merge grouped values by taking max/concatenate or require identical map_fun: Function to map the state dict to vLLM format """ + # If map_fun is provided, pre-transform quantizer key module-path prefixes so that + # HF→vLLM model renames (e.g. backbone.layers → model.layers) are applied before + # key grouping (q/k/v → qkv, experts.N.up_proj → experts.w13, etc.). + # This is necessary for models where the HF root module differs from vLLM's (e.g. + # NemotronH uses backbone.layers in HF but model.layers in vLLM), and for + # modelopt_state_weights where ALL keys are quantizer keys so map_fun is never + # invoked on non-quantizer keys. + if map_fun is not None: + q_only = {k: v for k, v in state_dict.items() if "_quantizer" in k} + prefix_remap = infer_quantizer_prefix_remap(q_only, map_fun) + if prefix_remap: + renamed = {} + for k, v in state_dict.items(): + if "_quantizer" in k: + first = k.split(".")[0] + k = prefix_remap.get(first, first) + k[len(first) :] + renamed[k] = v + state_dict = renamed + vllm_state_dict, merge_groups = _group_keys_for_vllm(state_dict) merge_fn = _merge_values_by_max_or_concat if max_or_concat else _merge_values_require_identical @@ -340,7 +362,26 @@ def _has_buffers(state: dict) -> bool: } # Add state for quantizers in model but not in metadata (e.g. disabled/excluded) for k in model_keys - filtered.keys(): - filtered[k] = model_qstate[k] + state = model_qstate[k] + # Weight quantizers absent from exported metadata were disabled during export + # (weights are already fake-quantized and pre_quant_scale is folded in). + # Keep them disabled on reload so fold_weight does not re-quantize the + # already-folded weights (re-quantizing distorts the pqs-scaled values). + if is_weight_quantizer_state_key(k) and not state.get("_disabled"): + state = {**state, "_disabled": True} + filtered[k] = state + + # Invariant: weight quantizers absent from export must be _disabled. + for wq_k in model_keys: + if not is_weight_quantizer_state_key(wq_k): + continue + wq_state = filtered[wq_k] + if wq_k not in saved and not wq_state.get("_disabled"): + raise RuntimeError( + f"Weight quantizer {wq_k!r} is missing from saved quantizer_state but " + f"is not marked _disabled (got _disabled={wq_state.get('_disabled')!r}). " + f"vLLM fakequant export omits weight quantizer keys when weights are folded." + ) metadata["quantizer_state"] = filtered @@ -379,10 +420,123 @@ def restore_from_modelopt_state_vllm( if not manager.has_state and isinstance(model, ModelLikeModule): model = model.init_modellike() - assert not isinstance(model, ModelLikeModule), "Model must be a regular Module now!" + if isinstance(model, ModelLikeModule): + raise RuntimeError("Model must be a regular Module after restore, got ModelLikeModule") return model +def _tp_concat_shard_dims( + value_shape: tuple[int, ...], + expected_shape: tuple[int, ...], + tp_world_size: int, +) -> list[int]: + """Dims ``d`` where checkpoint looks like TP concat: ``value[d] == expected[d] * tp_world_size``.""" + return [ + d for d in range(len(expected_shape)) if value_shape[d] == expected_shape[d] * tp_world_size + ] + + +def _narrow_tensor_to_tp_local_shard( + value: torch.Tensor, + expected_shape: tuple[int, ...] | torch.Size, + tp_rank: int, + tp_world_size: int, + *, + context: str, +) -> torch.Tensor: + """Slice ``value`` to this TP rank when it is the concat of per-rank shards along one dim.""" + value_shape = value.shape + expected_shape = tuple(expected_shape) + if value_shape == expected_shape: + return value + if len(value_shape) != len(expected_shape): + raise ValueError( + f"{context}: rank mismatch (checkpoint={tuple(value_shape)}, expected={tuple(expected_shape)})" + ) + shard_dims = _tp_concat_shard_dims(value_shape, expected_shape, tp_world_size) + if len(shard_dims) != 1: + raise ValueError( + f"{context}: cannot infer TP shard dim " + f"(expected={tuple(expected_shape)}, checkpoint={tuple(value_shape)}, tp={tp_world_size})" + ) + d = shard_dims[0] + shard_size = expected_shape[d] + start = tp_rank * shard_size + if start + shard_size > value_shape[d]: + raise ValueError( + f"{context}: TP shard out of bounds " + f"(expected={tuple(expected_shape)}, checkpoint={tuple(value_shape)})" + ) + return value.narrow(d, start, shard_size).contiguous() + + +def _pqs_local_expected_shape(pqs: torch.Tensor, expected_in: int) -> tuple[int, ...] | None: + """Local per-rank shape for ``_pre_quant_scale`` (1-D ``[H]`` or broadcast 2-D ``[1, H]``).""" + if pqs.ndim == 1: + return (expected_in,) + if pqs.ndim == 2 and pqs.shape[0] == 1: + return (1, expected_in) + return None + + +def _expected_in_features_for_input_quantizer(parent: Any, input_quantizer_attr: str) -> int | None: + """Input feature count for the weight paired with ``*_input_quantizer`` (Linear or FusedMoE).""" + stem = input_quantizer_attr[: -len("_input_quantizer")] + w = getattr(parent, (stem + "_weight") if stem else "weight", None) + if w is None or not isinstance(w, torch.Tensor) or w.is_meta: + return None + return int(w.shape[-1] if w.ndim == 3 else w.shape[1]) + + +def shard_pre_quant_scale_for_tp(model: Any) -> None: + """Shard ``_pre_quant_scale`` in-place for the local TP rank (row-parallel inputs). + + HF exports often store full (unsharded) scales; after load, row-parallel layers need + ``pqs`` narrowed to ``H_in / tp`` when ``len(pqs) == H_in * tp_world_size``. + + Call after parallel linear modules expose TP-sharded weight shapes (e.g. + ``post_restore_vllm_parallel_linears``). If run earlier, ``expected_in`` inferred from + weights can match an unsharded checkpoint and a second call becomes a no-op even when + pqs should still be narrowed. + + Args: + model: vLLM model with ``TensorQuantizer`` submodules. + """ + from modelopt.torch.quantization.nn import TensorQuantizer + + tp_group = get_tp_group() + tp_rank, tp_world_size = tp_group.rank_in_group, tp_group.world_size + if tp_world_size == 1: + return + + for qname, quantizer in model.named_modules(): + if not isinstance(quantizer, TensorQuantizer): + continue + pqs = getattr(quantizer, "_pre_quant_scale", None) + if pqs is None: + continue + last = qname.rfind(".") + if last == -1 or not qname[last + 1 :].endswith("input_quantizer"): + continue + try: + parent = model.get_submodule(qname[:last]) + except (AttributeError, LookupError): + continue + expected_in = _expected_in_features_for_input_quantizer(parent, qname[last + 1 :]) + if expected_in is None: + continue + expected_shape = _pqs_local_expected_shape(pqs, expected_in) + if expected_shape is None: + continue + quantizer._pre_quant_scale = _narrow_tensor_to_tp_local_shard( + pqs, + expected_shape, + tp_rank, + tp_world_size, + context=f"{qname}._pre_quant_scale", + ) + + def process_state_dict_for_tp(saved_qstate_dict, current_state_dict): """Shard quantizer tensors for tensor parallelism by matching expected shapes.""" tp_group = get_tp_group() @@ -393,42 +547,14 @@ def process_state_dict_for_tp(saved_qstate_dict, current_state_dict): for key, value in saved_qstate_dict.items(): if key in current_state_dict: expected = current_state_dict[key] - if not hasattr(value, "shape") or not hasattr(expected, "shape"): - result[key] = value - continue - expected_shape = expected.shape - value_shape = value.shape - if value_shape != expected_shape: - # Verify compatible rank before indexing - if len(value_shape) != len(expected_shape): - raise ValueError( - f"Cannot infer TP shard dim for {key}: rank mismatch " - f"(checkpoint rank={len(value_shape)}, expected rank={len(expected_shape)})" - ) - # Find the dimension that was tensor-parallel sharded. - # We expect exactly one dimension to satisfy: - # checkpoint_dim == expected_dim * tp_world_size - shard_dims = [ - d - for d in range(len(expected_shape)) - if value_shape[d] == expected_shape[d] * tp_world_size - ] - if len(shard_dims) != 1: - raise ValueError( - f"Cannot infer TP shard dim for {key}: " - f"expected_shape={tuple(expected_shape)}, checkpoint_shape={tuple(value_shape)}" - ) - - shard_dim = shard_dims[0] - shard_size = expected_shape[shard_dim] - start = tp_rank * shard_size - end = start + shard_size - if end > value_shape[shard_dim]: - raise ValueError( - f"TP shard out of bounds for {key}: " - f"expected_shape={tuple(expected_shape)}, checkpoint_shape={tuple(value_shape)}" - ) - value = value.narrow(shard_dim, start, shard_size).contiguous() + if hasattr(value, "shape") and hasattr(expected, "shape"): + value = _narrow_tensor_to_tp_local_shard( + value, + expected.shape, + tp_rank, + tp_world_size, + context=f"Key {key!r}", + ) result[key] = value return result @@ -437,12 +563,8 @@ def process_state_dict_for_tp(saved_qstate_dict, current_state_dict): def load_state_dict_from_path( fakequant_runner: Any, quantizer_file_path: str, model: Any ) -> dict[str, Any]: - fakequant_runner.model_runner._dummy_run(1) - print(f"Loading quantizer values from {quantizer_file_path}") - # Load on CPU to avoid failures when the checkpoint was saved from a different - # GPU mapping + # Load on CPU to avoid failures when the checkpoint was saved from a different GPU mapping. saved_quant_dict = torch.load(quantizer_file_path, weights_only=True, map_location="cpu") - # convert quant keys to vLLM format if hasattr(fakequant_runner.model_runner.model, "hf_to_vllm_mapper"): saved_quant_dict = fakequant_runner.model_runner.model.hf_to_vllm_mapper.apply_dict( saved_quant_dict @@ -455,7 +577,6 @@ def load_state_dict_from_path( saved_quant_dict = convert_dict_to_vllm(saved_quant_dict) current_state_dict = model.state_dict() - # Count quant keys in checkpoint and model checkpoint_quant_keys = [key for key in saved_quant_dict if "quantizer" in key] model_quant_keys = [key for key in current_state_dict if "quantizer" in key] ckpt_key_set = set(checkpoint_quant_keys) diff --git a/modelopt/torch/export/layer_utils.py b/modelopt/torch/export/layer_utils.py index 7726bf61af7..e8ee5afd451 100755 --- a/modelopt/torch/export/layer_utils.py +++ b/modelopt/torch/export/layer_utils.py @@ -81,8 +81,16 @@ has_mcore = True -def get_experts_list(module: torch.nn.Module, model_type: str): - """Returns list of grouped experts by linear name for given module.""" +def get_experts_list( + module: torch.nn.Module, + model_type: str, +): + """Returns list of grouped experts by linear name for given module. + + Args: + module: MoE block (e.g. MixtralSparseMoeBlock, NemotronHMOE). + model_type: `type(root_model).__name__.lower()` (may change after ModelOpt quantize). + """ experts_list = [] # Define linear layer names for different model types @@ -98,6 +106,8 @@ def get_experts_list(module: torch.nn.Module, model_type: str): ] ): linear_names = ["gate_proj", "down_proj", "up_proj"] + elif "nemotronhforcausallm" in model_type: + linear_names = ["up_proj", "down_proj"] else: raise NotImplementedError(f" {model_type} not supported") @@ -305,7 +315,7 @@ def is_moe(module: nn.Module) -> bool: if name.endswith("sparsemoeblock") or "moelayer" in name: return True # Explicit matches for non-standard naming - return any(key in name for key in ["arcticmoe", "deepseekmoe", "dbrxffn"]) + return any(key in name for key in ["arcticmoe", "deepseekmoe", "dbrxffn", "nemotronhmoe"]) def is_quantlinear(module: nn.Module) -> bool: @@ -994,6 +1004,9 @@ def module_match_name_list(module, name_list): return ["w1_linear", "w2_linear", "v1_linear"] elif module_match_name_list(module, ["GptOssMoE"]): return ["gate_up_proj", "down_proj"] + elif module_match_name_list(module, ["NemotronHMOE"]): + # NemotronHMOE experts (NemotronHMLP) use up_proj and down_proj only (no gate). + return ["up_proj", "down_proj"] else: # assuming w1, w2, w3 by default return ["w1", "w2", "w3"] diff --git a/modelopt/torch/export/plugins/vllm_fakequant_hf.py b/modelopt/torch/export/plugins/vllm_fakequant_hf.py index 786f9cdf593..42baad912b8 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_hf.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_hf.py @@ -14,7 +14,14 @@ # limitations under the License. """Export HuggingFace model to vLLM fakequant checkpoint.""" +import copy +import logging +import re +import warnings +from collections.abc import Callable +from contextlib import ExitStack, contextmanager from pathlib import Path +from typing import Any import torch import torch.nn as nn @@ -22,13 +29,108 @@ import modelopt.torch.opt as mto from modelopt.torch.quantization.config import RotateConfig from modelopt.torch.quantization.conversion import quantizer_state -from modelopt.torch.quantization.nn import QuantModule, TensorQuantizer +from modelopt.torch.quantization.model_calib import enable_stats_collection, finish_stats_collection +from modelopt.torch.quantization.nn import QuantModule, SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import get_quantizer_state_dict from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector -from modelopt.torch.utils import get_unwrapped_name +from modelopt.torch.utils import get_unwrapped_name, safe_save -__all__ = ["export_hf_vllm_fq_checkpoint"] +from ..layer_utils import get_experts_list, is_moe +from ..quant_utils import get_quantization_format +from ..unified_export_hf import collect_shared_input_modules + +__all__ = [ + "export_hf_vllm_fq_checkpoint", + "infer_quantizer_prefix_remap", + "is_weight_quantizer_state_key", + "merge_amax_tensors_for_group", +] + +# Matches ``…weight_quantizer``, ``…weight_quantizer.0``, ``…w13_weight_quantizer.0``, etc. +_WEIGHT_QUANTIZER_STATE_KEY = re.compile(r"(?:^|\.)(?:\w+_)?weight_quantizer(?:\.\d+)*$") + + +def is_weight_quantizer_state_key(key: str) -> bool: + """Return True for weight-quantizer state keys, including SequentialQuantizer entries. + + Matches ``weight_quantizer``, ``w13_weight_quantizer``, ``weight_quantizer.0``, etc. + """ + return bool(_WEIGHT_QUANTIZER_STATE_KEY.search(key)) + + +def infer_quantizer_prefix_remap( + quantizer_keys: dict[str, Any], + map_fun: Callable[[dict[str, Any]], dict[str, Any]], +) -> dict[str, str]: + """Infer HF root name → vLLM root (e.g. ``backbone`` → ``model``) for reload/export. + + Map HF root → vLLM root (e.g. ``backbone`` → ``model``) by probing ``map_fun`` with + synthetic ``.weight`` keys and a 2-D placeholder (quantizer paths are not weight + keys). Keys under the same HF root must agree on the target root or :exc:`ValueError` is + raised; failed probes are skipped. Returns ``{hf_root: vllm_root}`` only where the root + renames; not for arbitrary layer rewrites. + + Args: + quantizer_keys: HF quantizer state paths as keys (values unused). + map_fun: HF→vLLM weight ``state_dict`` mapper, same as for ``convert_dict_to_vllm``. + + Returns: + ``{hf_root: vllm_root}`` for roots that rename; omits identity pairs. + """ + logger = logging.getLogger(__name__) + probe_weight = torch.empty((1, 1)) + observed_vllm_root: dict[str, str] = {} + + for key in quantizer_keys: + first_component = key.split(".")[0] + last_dot = key.rfind(".") + if last_dot == -1: + continue + probe_key = key[:last_dot] + ".weight" + try: + result = map_fun({probe_key: probe_weight}) + if not result: + continue + new_key = next(iter(result)) + new_first = new_key.split(".")[0] + except Exception as e: + logger.debug("prefix-remap probe failed for %r: %s", probe_key, e) + continue + + if first_component not in observed_vllm_root: + observed_vllm_root[first_component] = new_first + elif observed_vllm_root[first_component] != new_first: + raise ValueError( + "Inconsistent HF→vLLM prefix remap for " + f"{first_component!r}: probes implied " + f"{observed_vllm_root[first_component]!r} and {new_first!r}. " + "map_fun must apply one target root per HF root, or use explicit quantizer " + "key remapping." + ) + + return { + hf_root: vllm_root + for hf_root, vllm_root in observed_vllm_root.items() + if hf_root != vllm_root + } + + +def _check_all_weight_quantizers_disabled(model: nn.Module) -> None: + """Export invariant before writing metadata: every weight quantizer must be off.""" + for _, module in model.named_modules(): + if not isinstance(module, QuantModule): + continue + for attr_name, quantizer in module.named_children(): + if attr_name.endswith("weight_quantizer") and isinstance( + quantizer, (TensorQuantizer, SequentialQuantizer) + ): + if quantizer.is_enabled: + raise RuntimeError( + f"vLLM fakequant export: {attr_name!r} must be disabled before saving " + f"quantizer_state (weights already folded). " + f"See filter_modelopt_state_quantizer_state_for_model in vllm_reload_utils." + ) def disable_rotate(quantizer: TensorQuantizer): @@ -47,6 +149,7 @@ def _fakequant_module_weights( state_dict: dict | None, input_quantizers_folded_pqs: set, fakequant_weights: set, + requant_weights: set[str], inplace: bool, ): """Apply fake-quant to a single QuantModule's weights. @@ -67,17 +170,25 @@ def _fakequant_module_weights( weight_name = attr_name.removesuffix("_quantizer") prefix = f"{module_name}." if module_name else "" sd_key = f"{prefix}{weight_name}" - assert sd_key not in fakequant_weights, f"Weight {sd_key} has already been fakequantized" + if sd_key in fakequant_weights: + raise RuntimeError(f"Weight {sd_key} has already been fakequantized") if inplace: w = getattr(module, weight_name) - w_quant = quantizer(w.float()).to(w.dtype) + if sd_key in requant_weights: + w_quant = requant_weights_for_export(quantizer, w, copy_quantizer=False) + else: + w_quant = quantizer(w.float()).to(w.dtype) else: - assert state_dict is not None + if state_dict is None: + raise RuntimeError("state_dict is required when inplace=False for fakequant export") if sd_key not in state_dict: continue w = state_dict[sd_key] - w_quant = quantizer(w.float()).to(w.dtype) + if sd_key in requant_weights: + w_quant = requant_weights_for_export(quantizer, w) + else: + w_quant = quantizer(w.float()).to(w.dtype) # Fold pre_quant_scale: (x*s)@fake_quant(W) = x@(fake_quant(W)*s) # Only valid when input_quantizer does NOT fake-quant activations. If it does @@ -88,7 +199,7 @@ def _fakequant_module_weights( if ( hasattr(inp_q, "_pre_quant_scale") and inp_q._pre_quant_scale is not None - and inp_q._disabled + and not inp_q.is_enabled ): scale = inp_q._pre_quant_scale.squeeze().to(device=w_quant.device) w_quant = (w_quant * scale[None, :]).to(w_quant.dtype) @@ -100,11 +211,266 @@ def _fakequant_module_weights( if inplace: w.data.copy_(w_quant) else: - assert state_dict is not None + if state_dict is None: + raise RuntimeError("state_dict is required when inplace=False for fakequant export") state_dict[sd_key] = w_quant.cpu() fakequant_weights.add(sd_key) +def _collect_group_pre_quant_scales( + experts: list[nn.Module], +) -> list[torch.Tensor] | None: + """Return per-expert ``pre_quant_scale`` tensors if every expert can be averaged; else None. + + Skips groups where any expert has no input quantizer, no pqs (e.g. weight-only AWQ INT4), + or a disabled input quantizer (pqs already folded / not used). + """ + pre_quant_scales: list[torch.Tensor] = [] + for expert_module in experts: + input_quantizer = getattr(expert_module, "input_quantizer", None) + if ( + input_quantizer is None + or not input_quantizer.is_enabled + or input_quantizer.pre_quant_scale is None + ): + return None + pre_quant_scales.append(input_quantizer.pre_quant_scale) + return pre_quant_scales + + +def requant_weights_for_export( + quantizer: TensorQuantizer | SequentialQuantizer, + weight: torch.Tensor, + copy_quantizer: bool = True, +) -> torch.Tensor: + """Requantize folded weights after resmooth (``TensorQuantizer`` or ``SequentialQuantizer``). + + A single ``TensorQuantizer`` is treated as a one-stage chain so the same + calibrate-then-apply steps cover W4A8-style sequential weights (e.g. INT4→FP8). + + Deepcopy may leave buffers on the original device; ``.to(device=w.device)`` aligns with + ``w`` (e.g. CPU offload). + """ + if copy_quantizer: + copied = copy.deepcopy(quantizer).to(device=weight.device) + else: + copied = quantizer + quantizers: list[TensorQuantizer] = ( + list(copied) if isinstance(copied, SequentialQuantizer) else [copied] + ) + + for quantizer_copy in quantizers: + quantizer_copy.eval() + quantizer_copy.reset_amax() + enable_stats_collection(quantizer_copy) + weight_quantized = weight + for quantizer_copy in quantizers: + weight_quantized = quantizer_copy(weight_quantized) + for quantizer_copy in quantizers: + finish_stats_collection(quantizer_copy) + # Re-run application pass to get the quantized output with the freshly collected amax. + # The calibration forward above only collected stats; its output is intentionally discarded. + weight_quantized = weight + for quantizer_copy in quantizers: + weight_quantized = quantizer_copy(weight_quantized) + return weight_quantized.to(weight.dtype) + + +def merge_amax_tensors_for_group(tensors: list[torch.Tensor]) -> torch.Tensor: + """Combine `_amax` buffers from a merge group into a single tensor. + + Used when HuggingFace module names are folded to vLLM names (e.g. q/k/v → qkv_proj). + + - If every tensor has the same shape, take the element-wise maximum over the group + (conservative when each branch carried the same axis layout). + - If shapes differ: ``torch.cat(..., dim=0)`` assumes **1D per-channel** amaxes in + fused order (e.g. GQA q/k/v → ``[N_q]`` + ``[N_kv]`` + ``[N_kv]``), matching vLLM’s + grouped quantizer. Not valid for 2D blockwise amax; on failure, **scalar** + max (drops channel structure). + """ + if not tensors: + raise ValueError("merge_amax_tensors_for_group: expected at least one tensor") + if len(tensors) == 1: + return tensors[0] + + first = tensors[0] + if all(t.shape == first.shape for t in tensors): + stacked = torch.stack([t.float() for t in tensors], dim=0) + return torch.amax(stacked, dim=0).to(dtype=first.dtype, device=first.device) + + try: + return torch.cat(tensors, dim=0).to(dtype=first.dtype, device=first.device) + except RuntimeError: + shapes = [tuple(t.shape) for t in tensors] + warnings.warn( + f"merge_amax_tensors_for_group: torch.cat failed for shapes {shapes}; " + "falling back to scalar max which loses per-channel amax structure.", + stacklevel=2, + ) + flat = torch.cat([t.reshape(-1).float() for t in tensors]) + return torch.max(flat).to(dtype=first.dtype, device=first.device) + + +@contextmanager +def _enable_writeback_for_group( + group: list[nn.Module], + root_model: nn.Module, + name_to_module: dict[str, nn.Module], +): + """Nest ``enable_weight_access_and_writeback`` for every module in ``group`` (one ``with``). + + The stdlib pattern for a *variable* number of context managers is :class:`ExitStack`; + wrapping it here keeps call sites readable. + """ + with ExitStack() as stack: + for m in group: + stack.enter_context(enable_weight_access_and_writeback(m, root_model, name_to_module)) + yield + + +def _resmooth_experts_for_export( + model: nn.Module, + state_dict: dict[str, Any] | None, + *, + inplace: bool = False, +) -> tuple[dict[str, tuple[torch.Tensor, torch.Tensor | None]], set[str]]: + """Prepare AWQ weights for vLLM fakequant export when several linears share one input quantizer. + + PTQ can assign a different ``pre_quant_scale`` per branch (per expert, or per + q/k/v projection) even though they see the same activation. vLLM’s fused kernels expose a + **single** input quantizer for that fused group, so reload must use one scale — otherwise + activations are scaled wrong for k/v or non-primary experts. + + For each group (MoE experts via ``get_experts_list``; dense shared-input linears + via ``collect_shared_input_modules`` / hooks), average ``pre_quant_scale``, set weights to + ``W' = W * old_pqs / avg_pqs`` so the net is unchanged, merge input ``amax`` where needed, + and return per-``input_quantizer`` tensor overrides for ``modelopt_state_weights``. + + Runs only for AWQ with **enabled** input quantizers (e.g. activation-aware); if inputs are + disabled and PQS was folded into weights only, there is nothing to unify. + + ``inplace=False`` — adjust a detached ``state_dict`` copy (``state_dict`` required). + ``inplace=True`` — pass ``state_dict=None``; update live ``nn.Parameter`` data under + ``_enable_writeback_for_group`` (nested writeback per module so offloaded/meta weights + materialize before ``copy_``). + """ + if not inplace and state_dict is None: + raise ValueError("state_dict is required when inplace=False") + qfmt = get_quantization_format(model) + if qfmt is None or "awq" not in qfmt.lower(): + return {}, set() + + name_to_module = dict(model.named_modules()) if inplace else None + + model_type = type(model).__name__.lower() + id_to_name: dict[int, str] = {id(m): n for n, m in model.named_modules()} + out: dict[str, tuple[torch.Tensor, torch.Tensor | None]] = {} + requant_weights: set[str] = set() + + def _process_group(modules: list[nn.Module]) -> None: + pqs_list = _collect_group_pre_quant_scales(modules) + if pqs_list is None: + return + + # Mean and clamp in float32: fp16/bf16 would underflow float32.tiny to 0 and divide by zero. + pqs_dtype = pqs_list[0].dtype + avg_pqs = torch.stack([p.float() for p in pqs_list]).mean(0) + avg_pqs = avg_pqs.clamp(min=torch.finfo(torch.float32).tiny) + + for m in modules: + nm = id_to_name.get(id(m)) + if nm is None or not hasattr(m, "weight"): + continue + w_key = f"{nm}.weight" + old_pqs = m.input_quantizer._pre_quant_scale + avg_pqs_dev = avg_pqs.to(device=old_pqs.device, dtype=old_pqs.dtype) + if torch.equal(old_pqs, avg_pqs_dev): + continue + if inplace: + w_param = m.weight + ratio = old_pqs.to(dtype=torch.float32, device=w_param.device) / avg_pqs.to( + device=w_param.device + ) + w_param.data.copy_((w_param.to(torch.float32) * ratio).to(w_param.dtype)) + else: + if state_dict is None: + raise RuntimeError( + "state_dict is required when inplace=False in _resmooth_experts_for_export" + ) + weight = state_dict[w_key] + ratio = old_pqs.to(dtype=torch.float32, device=weight.device) / avg_pqs.to( + device=weight.device + ) + state_dict[w_key] = (weight.to(torch.float32) * ratio).to(weight.dtype) + requant_weights.add(w_key) + + synced_amax: torch.Tensor | None = None + amaxes = [m.input_quantizer.amax for m in modules] + if all(a is not None for a in amaxes): + synced_amax = merge_amax_tensors_for_group(amaxes) + + avg_pqs_out = avg_pqs.detach().to(pqs_dtype).clone() + for m in modules: + nm = id_to_name.get(id(m)) + if nm is None: + continue + out[get_unwrapped_name(f"{nm}.input_quantizer", model)] = (avg_pqs_out, synced_amax) + + # MoE expert groups — must be enumerated by name because MoE routing sends + # different tokens to each expert, so forward hooks cannot detect them as + # sharing the same input tensor. + for _, module in model.named_modules(): + if not is_moe(module): + continue + try: + expert_groups = get_experts_list(module, model_type) + except NotImplementedError: + continue + for experts in expert_groups: + if not experts: + continue + if inplace: + if name_to_module is None: + raise RuntimeError( + "name_to_module is required when inplace=True in _resmooth_experts_for_export" + ) + with _enable_writeback_for_group(experts, model, name_to_module): + _process_group(experts) + else: + _process_group(experts) + + # Dense shared-input groups (e.g. q/k/v in GQA attention) — detected via forward + # hooks so any architecture is covered regardless of projection attribute names. + + dev = next(model.parameters()).device + + def _dummy_forward() -> None: + # Partial forward is OK: hooks record layers reached before failure. + with torch.inference_mode(): + try: + model(torch.ones([1, 2], dtype=torch.long, device=dev)) + except Exception as e: + logging.getLogger(__name__).debug( + "Dummy forward for shared-input detection failed (expected for VLMs): %s", e + ) + + input_to_linear, _ = collect_shared_input_modules(model, _dummy_forward) + for modules in input_to_linear.values(): + if len(modules) <= 1: + continue + if inplace: + if name_to_module is None: + raise RuntimeError( + "name_to_module is required when inplace=True in _resmooth_experts_for_export" + ) + with _enable_writeback_for_group(modules, model, name_to_module): + _process_group(modules) + else: + _process_group(modules) + + return out, requant_weights + + def export_hf_vllm_fq_checkpoint( model: nn.Module, export_dir: Path | str, @@ -115,8 +481,15 @@ def export_hf_vllm_fq_checkpoint( Folds fake-quant weights into a ``state_dict()`` copy (optional ``pre_quant_scale`` into weight when input fake-quant is off), drops quantizer keys from the HF save, briefly disables weight quantizers to snapshot - ModelOpt/quantizer state, then re-enables them. Writes ``export_dir`` via - ``save_pretrained(..., save_modelopt_state=False)``. + ModelOpt/quantizer state, then re-enables them. Weight files are written with an + explicit ``state_dict`` (and ``hf_quantizer`` cleared during save) so safetensors + do not pick up live quantizer buffers. + + For MoE models with AWQ quantization, pre_quant_scale is averaged across experts + and input amax is unified — required because vLLM uses a single input quantizer + per expert group. By default this updates only a detached ``state_dict`` copy. + With ``inplace_mem_efficient=True``, resmooth runs **in place** on materialized + weight parameters only (no ``state_dict``), before the inplace fakequant loop. Args: model: In-memory quantized model. @@ -129,16 +502,19 @@ def export_hf_vllm_fq_checkpoint( export_dir = Path(export_dir) export_dir.mkdir(parents=True, exist_ok=True) - # Step 1: Build the folded HF state dict. - fakequant_weights = set() - input_quantizers_folded_pqs = set() + fakequant_weights: set[str] = set() + # Input quantizer keys whose _pre_quant_scale was folded into the weight above. + input_quantizers_folded_pqs: set[str] = set() with torch.inference_mode(): if inplace_mem_efficient: + # Resmooth shared-input groups, then fakequant (state dict and/or params). + pqs_overrides, requant_weights = _resmooth_experts_for_export(model, None, inplace=True) # Inplace path: iterate decoder layers, one offload<->onload per layer. decoder_layers = LayerActivationCollector.get_decoder_layers(model) - assert decoder_layers is not None, ( - "inplace_mem_efficient=True requires a model with discoverable decoder layers" - ) + if decoder_layers is None: + raise RuntimeError( + "inplace_mem_efficient=True requires a model with discoverable decoder layers" + ) for name, module in model.named_modules(): if module not in decoder_layers: continue @@ -152,14 +528,21 @@ def export_hf_vllm_fq_checkpoint( None, input_quantizers_folded_pqs, fakequant_weights, + requant_weights, inplace=True, ) # Meta tensors for offloaded weights (free); offload maps now have # fakequanted values via writeback. state_dict = model.state_dict() else: - # Default path: full state_dict copy, fakequant into the copy. state_dict = model.state_dict() + # Resmooth shared-input groups, then fakequant (state dict and/or params). + pqs_overrides, requant_weights = _resmooth_experts_for_export( + model, state_dict, inplace=False + ) + + # Default path: fakequant into the resmoothed state_dict copy (do not refresh + # from model.state_dict() or resmooth is lost). for module_name, module in model.named_modules(): with enable_weight_access_and_writeback(module, model): _fakequant_module_weights( @@ -169,6 +552,7 @@ def export_hf_vllm_fq_checkpoint( state_dict, input_quantizers_folded_pqs, fakequant_weights, + requant_weights, inplace=False, ) @@ -188,63 +572,87 @@ def export_hf_vllm_fq_checkpoint( # attention quantizers remain active. # Rotation is also cleared: the weight was already folded with rotation applied, # so if fold_weight is called on reload it must not re-rotate the exported weight. - wqs_to_restore = [] - for _, module in model.named_modules(): - if isinstance(module, QuantModule): - for attr_name, quantizer in module.named_children(): - if ( - attr_name.endswith("weight_quantizer") - and isinstance(quantizer, TensorQuantizer) - and quantizer.is_enabled - ): - quantizer.disable() - orig_rotate = quantizer._rotate - if quantizer.rotate_is_enabled: - quantizer._rotate = disable_rotate(quantizer) - wqs_to_restore.append((quantizer, orig_rotate)) - - quantizer_state_dict = get_quantizer_state_dict(model) - for key in list(quantizer_state_dict): - if key.endswith("weight_quantizer"): - # Fakequant amax is folded into HF weights; do not reload weight quantizer tensors. - quantizer_state_dict.pop(key) - elif key in input_quantizers_folded_pqs: - # pre_quant_scale was folded into the weight; keep the buffer for strict load but - # save identity so activations are not scaled twice. - qstate_val = quantizer_state_dict[key] - if isinstance(qstate_val, dict) and "_pre_quant_scale" in qstate_val: - quantizer_state_dict[key]["_pre_quant_scale"] = torch.ones_like( - qstate_val["_pre_quant_scale"] - ) - modelopt_state = mto.modelopt_state(model) - # ``modelopt_state`` may be stale if another mode (e.g. calibrate) ran last. Rebuild - # ``quantizer_state`` and drop disabled weight quantizer entries (weights already folded). - qstate = quantizer_state(model) - for key in list(qstate): - if key.endswith("weight_quantizer") and qstate[key].get("_disabled"): - qstate.pop(key) - - for mode_str, m_state in modelopt_state.get("modelopt_state_dict", []): - if mode_str == "quantize" and "metadata" in m_state: - m_state["metadata"]["quantizer_state"] = qstate - break - - # Per-quantizer tensor dict loaded alongside metadata on reload. - modelopt_state["modelopt_state_weights"] = quantizer_state_dict - torch.save(modelopt_state, export_dir / "vllm_fq_modelopt_state.pth") - - # Step 3: Save HF weights. - if inplace_mem_efficient: - prev_ignore = getattr(model, "_keys_to_ignore_on_save", None) - model._keys_to_ignore_on_save = quantizer_keys - try: - model.save_pretrained(export_dir, save_modelopt_state=False) - finally: - model._keys_to_ignore_on_save = prev_ignore - else: - model.save_pretrained(export_dir, state_dict=clean_sd, save_modelopt_state=False) + wqs_to_restore: list[tuple[TensorQuantizer, Any]] = [] + try: + for _, module in model.named_modules(): + if isinstance(module, QuantModule): + for attr_name, quantizer in module.named_children(): + if not (attr_name.endswith("weight_quantizer") and quantizer.is_enabled): + continue + if isinstance(quantizer, SequentialQuantizer): + quantizer.disable() + for sub in quantizer: + orig_rotate = sub._rotate + if sub.rotate_is_enabled: + sub._rotate = disable_rotate(sub) + wqs_to_restore.append((sub, orig_rotate)) + elif isinstance(quantizer, TensorQuantizer): + quantizer.disable() + orig_rotate = quantizer._rotate + if quantizer.rotate_is_enabled: + quantizer._rotate = disable_rotate(quantizer) + wqs_to_restore.append((quantizer, orig_rotate)) + + quantizer_state_dict = get_quantizer_state_dict(model) + for key in list(quantizer_state_dict): + if is_weight_quantizer_state_key(key): + # Fakequant amax is folded into HF weights; do not reload weight quantizer tensors. + # Reload must force-disable WQs missing from saved state (see + # ``filter_modelopt_state_quantizer_state_for_model`` assertion in vllm_reload_utils). + quantizer_state_dict.pop(key) + elif key in input_quantizers_folded_pqs: + # pre_quant_scale was folded into the weight; keep the buffer for strict load but + # save identity so activations are not scaled twice. + qstate_val = quantizer_state_dict[key] + if isinstance(qstate_val, dict) and "_pre_quant_scale" in qstate_val: + quantizer_state_dict[key]["_pre_quant_scale"] = torch.ones_like( + qstate_val["_pre_quant_scale"] + ) + + # Patch input quantizers with averaged pqs and unified amax so that vLLM's single + # per-group input quantizer sees consistent values (covers both dense qkv and MoE experts). + for iq_key, (avg_pqs, max_input_amax) in pqs_overrides.items(): + if iq_key in quantizer_state_dict: + qstate_val = quantizer_state_dict[iq_key] + if isinstance(qstate_val, dict): + if "_pre_quant_scale" in qstate_val: + qstate_val["_pre_quant_scale"] = avg_pqs + if max_input_amax is not None and "_amax" in qstate_val: + qstate_val["_amax"] = max_input_amax + + modelopt_state = mto.modelopt_state(model) + _check_all_weight_quantizers_disabled(model) + # Rebuild quantizer_state from the live model (post-disable) and strip weight-quantizer + # entries. Apply to every mode that carries quantizer_state so that stale entries from + # a calibrate pass (which also stores quantizer_state in its metadata) are cleaned up. + # Reload synthesizes missing WQ rows with ``_disabled`` via + # ``filter_modelopt_state_quantizer_state_for_model``. + qstate = quantizer_state(model) + for key in list(qstate): + if is_weight_quantizer_state_key(key): + qstate.pop(key) + for _mode_str, m_state in modelopt_state.get("modelopt_state_dict", []): + md = m_state.get("metadata", {}) + if "quantizer_state" in md: + md["quantizer_state"] = qstate + + # Per-quantizer tensor dict loaded alongside metadata on reload. + modelopt_state["modelopt_state_weights"] = quantizer_state_dict + safe_save(modelopt_state, export_dir / "vllm_fq_modelopt_state.pth") + + # Step 3: Save HF weights. + if inplace_mem_efficient: + prev_ignore = getattr(model, "_keys_to_ignore_on_save", None) + model._keys_to_ignore_on_save = quantizer_keys + try: + model.save_pretrained(export_dir, save_modelopt_state=False) + finally: + model._keys_to_ignore_on_save = prev_ignore + else: + model.save_pretrained(export_dir, state_dict=clean_sd, save_modelopt_state=False) - if not inplace_mem_efficient: - for wq, orig_rotate in wqs_to_restore: - wq.enable() - wq._rotate = orig_rotate + finally: + if not inplace_mem_efficient: + for wq, orig_rotate in wqs_to_restore: + wq.enable() + wq._rotate = orig_rotate diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 22d87e303ff..af936a3002a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -163,7 +163,7 @@ def _save_component_state_dict_safetensors( json.dump(metadata, f, indent=4) -def _collect_shared_input_modules( +def collect_shared_input_modules( model: nn.Module, dummy_forward_fn: Callable[[], None], collect_layernorms: bool = False, @@ -387,7 +387,7 @@ def llm_dummy_forward(): else: model(fake_input) - input_to_linear, output_to_layernorm = _collect_shared_input_modules( + input_to_linear, output_to_layernorm = collect_shared_input_modules( model, llm_dummy_forward, collect_layernorms=True ) @@ -862,7 +862,7 @@ def _fuse_qkv_linears_diffusion( # Collect modules sharing the same input try: - input_to_linear, _ = _collect_shared_input_modules( + input_to_linear, _ = collect_shared_input_modules( model, dummy_forward_fn, collect_layernorms=False ) except Exception as e: diff --git a/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py b/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py index df5610ca3df..86a02b0ed89 100644 --- a/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py +++ b/tests/gpu/torch/export/test_vllm_fakequant_hf_export.py @@ -17,7 +17,8 @@ import pytest import torch -from _test_utils.torch.transformers_models import create_tiny_llama_dir +import transformers +from _test_utils.torch.transformers_models import create_tiny_llama_dir, create_tiny_qwen3_moe_dir from accelerate import init_empty_weights, load_checkpoint_and_dispatch from transformers import AutoConfig, AutoModelForCausalLM @@ -28,8 +29,7 @@ from modelopt.torch.utils import safe_load -@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG]) -def test_hf_vllm_export(tmp_path, quant_cfg): +def _test_hf_vllm_export(tmp_path, quant_cfg, model_dir): """Test HuggingFace model export for vLLM with fake quantization. This test verifies: @@ -39,11 +39,8 @@ def test_hf_vllm_export(tmp_path, quant_cfg): 4. Weight quantizer states are empty in saved state dict; input quantizer amaxes preserved """ - # Create a tiny LLaMA model for testing - tiny_model_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=2) - # Load the model - model = AutoModelForCausalLM.from_pretrained(tiny_model_dir) + model = AutoModelForCausalLM.from_pretrained(model_dir) model = model.cuda() model.eval() @@ -60,6 +57,23 @@ def forward_loop(model): folded_model = deepcopy(model) fold_weight(folded_model) expected_weights = {k: v for k, v in folded_model.state_dict().items() if "quantizer" not in k} + # fold_weight only applies the weight quantizer's fake-quant; it does NOT fold + # input_quantizer.pre_quant_scale into the weight. The export path does: + # w_exported = fake_quant(W) * pqs[None, :] + # for modules where input_quantizer is disabled but has pqs (AWQ weight-only). + # Apply the same pqs fold here so expected_weights matches the export output. + for module_name, module in folded_model.named_modules(): + inp_q = getattr(module, "input_quantizer", None) + if ( + inp_q is not None + and not inp_q.is_enabled + and getattr(inp_q, "_pre_quant_scale", None) is not None + ): + w_key = f"{module_name}.weight" if module_name else "weight" + if w_key in expected_weights: + w = expected_weights[w_key] + scale = inp_q._pre_quant_scale.squeeze().to(device=w.device) + expected_weights[w_key] = (w * scale[None, :]).to(w.dtype) del folded_model # Snapshot model state before export to verify it is not mutated @@ -231,3 +245,17 @@ def forward_loop(model): "_amax" in k for k in quantizer_state_dict_before[name] ): assert any("_amax" in k for k in state), f"input quantizer {name} should preserve _amax" + + +@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.INT4_AWQ_CFG]) +def test_hf_vllm_export_tiny_llama(tmp_path, quant_cfg): + tiny_model_dir = create_tiny_llama_dir(tmp_path, num_hidden_layers=2) + _test_hf_vllm_export(tmp_path, quant_cfg, tiny_model_dir) + + +@pytest.mark.parametrize("quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.INT4_AWQ_CFG]) +def test_hf_vllm_export_tiny_qwen3_moe(tmp_path, quant_cfg): + if quant_cfg == mtq.INT4_AWQ_CFG and transformers.__version__.startswith("5."): + pytest.skip("INT4_AWQ_CFG is not supported for Qwen3 MoE in transformers > 5.x") + tiny_model_dir = create_tiny_qwen3_moe_dir(tmp_path, num_hidden_layers=2) + _test_hf_vllm_export(tmp_path, quant_cfg, tiny_model_dir) diff --git a/tests/unit/torch/export/test_vllm_quantizer_reload.py b/tests/unit/torch/export/test_vllm_quantizer_reload.py new file mode 100644 index 00000000000..49fb34c25ad --- /dev/null +++ b/tests/unit/torch/export/test_vllm_quantizer_reload.py @@ -0,0 +1,156 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from modelopt.torch.export.plugins.vllm_fakequant_hf import ( + infer_quantizer_prefix_remap, + merge_amax_tensors_for_group, +) + + +def _map_backbone_to_model(sd: dict) -> dict: + """Test mapper: rename top-level ``backbone.`` to ``model.`` (typical HF vs vLLM).""" + out = {} + for k, v in sd.items(): + if k.startswith("backbone."): + out["model." + k[len("backbone.") :]] = v + else: + out[k] = v + return out + + +def test_infer_prefix_remap_simple_root_rename(): + """``infer_quantizer_prefix_remap`` infers one HF root → vLLM root from ``*.weight`` probes.""" + q = { + "backbone.layers.0.mlp.gate_proj.input_quantizer": {}, + "backbone.layers.1.self_attn.q_proj.weight_quantizer": {}, + } + rem = infer_quantizer_prefix_remap(q, _map_backbone_to_model) + assert rem == {"backbone": "model"} + + +def test_infer_prefix_remap_multiple_probes_same_root_agree(): + """Regression: every quantizer key under the same HF root must agree on the mapped vLLM root.""" + q = { + "backbone.a.w.input_quantizer": {}, + "backbone.b.w.weight_quantizer": {}, + } + rem = infer_quantizer_prefix_remap(q, _map_backbone_to_model) + assert rem == {"backbone": "model"} + + +def test_infer_prefix_remap_raises_on_inconsistent_root(): + """If ``map_fun`` maps the same HF root to different vLLM roots, raise with a clear error.""" + + def bad_map(sd: dict) -> dict: + out = {} + for k, v in sd.items(): + if "layers.0" in k: + out[k.replace("backbone.", "model.")] = v + elif "head" in k: + out[k.replace("backbone.", "encoder.")] = v + else: + out[k] = v + return out + + q = { + "backbone.layers.0.mlp.gate_proj.input_quantizer": {}, + "backbone.head.proj.input_quantizer": {}, + } + with pytest.raises(ValueError, match="Inconsistent HF→vLLM prefix remap"): + infer_quantizer_prefix_remap(q, bad_map) + + +def test_infer_prefix_remap_identity_empty(): + """When keys already match the mapper output, the inferred remap is empty (no rename).""" + q = {"model.layers.0.foo.input_quantizer": {}} + rem = infer_quantizer_prefix_remap(q, lambda d: dict(d)) + assert rem == {} + + +def test_infer_prefix_remap_probe_failure_skipped(): + """A probe that raises does not block remap if another key under the same root succeeds.""" + + def map_drop_layers0(sd: dict) -> dict: + out = {} + for k, v in sd.items(): + if "layers.0" in k: + raise RuntimeError("simulate missing layer") + if k.startswith("backbone."): + out["model." + k[len("backbone.") :]] = v + else: + out[k] = v + return out + + q = { + "backbone.layers.0.mlp.gate_proj.input_quantizer": {}, + "backbone.layers.1.mlp.gate_proj.input_quantizer": {}, + } + rem = infer_quantizer_prefix_remap(q, map_drop_layers0) + assert rem == {"backbone": "model"} + + +def test_infer_prefix_remap_no_quantizer_segment_still_probes_weight_path(): + """Short paths (e.g. ``embed.weight_quantizer``) still build a ``.weight`` probe path.""" + q = {"backbone.embed.weight_quantizer": {}} + rem = infer_quantizer_prefix_remap(q, _map_backbone_to_model) + assert rem == {"backbone": "model"} + + +def test_infer_prefix_remap_complex_mapper_not_one_root_raises_or_wrong(): + """Same HF root ``x`` mapping to different first components (``va.*`` vs ``vb.*``) must error.""" + + def split_map(sd: dict) -> dict: + k = next(iter(sd)) + v = sd[k] + if "branch_a" in k: + return {"va." + k[2:]: v} # x.branch_a... -> va.branch_a... + return {"vb." + k[2:]: v} + + q = { + "x.branch_a.mlp.w.input_quantizer": {}, + "x.branch_b.mlp.w.input_quantizer": {}, + } + with pytest.raises(ValueError, match="Inconsistent HF→vLLM prefix remap"): + infer_quantizer_prefix_remap(q, split_map) + + +def test_merge_amax_same_shape_elementwise_max(): + """``merge_amax_tensors_for_group``: identical shapes → element-wise max (stack then amax).""" + a = torch.tensor([1.0, 4.0, 2.0]) + b = torch.tensor([2.0, 3.0, 5.0]) + out = merge_amax_tensors_for_group([a, b]) + assert torch.allclose(out, torch.tensor([2.0, 4.0, 5.0])) + + +def test_merge_amax_different_1d_lengths_uses_cat(): + """``merge_amax_tensors_for_group``: mismatched 1-D lengths (e.g. GQA q/k/v) → ``cat`` on dim 0.""" + q = torch.tensor([1.0, 2.0, 3.0]) # e.g. 3 heads + k = torch.tensor([0.5, 0.5]) # 2 KV heads + v = torch.tensor([0.5, 0.5]) + out = merge_amax_tensors_for_group([q, k, v]) + assert out.shape == (7,) + assert torch.allclose(out, torch.cat([q, k, v])) + + +def test_merge_amax_incompatible_shapes_scalar_fallback(): + """``merge_amax_tensors_for_group``: when ``cat`` fails, fall back to a scalar global max.""" + a = torch.ones(2, 3) + b = torch.ones(2, 2) # cannot cat along dim=0 with matching trailing dims + out = merge_amax_tensors_for_group([a, b]) + assert out.shape == () + assert out.item() == 1.0 From 97d153118e9f43b9abc3af508d4702a6e01308a7 Mon Sep 17 00:00:00 2001 From: Frida Hou <201670829+Fridah-nv@users.noreply.github.com> Date: Mon, 20 Apr 2026 09:24:17 -0700 Subject: [PATCH 24/30] [minor] Add custom calibration backend registry (#1281) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: ? ### Usage ```python # Add a code snippet demonstrating how to use this ``` ### Testing ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A - Did you write any new necessary tests?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit * **New Features** * Added a public backend-specific calibrator registration API to support FP8 scale-sweep calibration, allowing backends to supply custom calibrators used during FP8 tuning. * **Tests** * Added unit tests confirming registry insertion/overwrite, that registered calibrators are invoked when FP8 scale-sweep is enabled, are not invoked when disabled, and that calibration falls back to defaults when no backend is registered. --------- Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Co-authored-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- modelopt/torch/quantization/model_calib.py | 42 ++++++- .../torch/quantization/test_mse_calibrator.py | 114 ++++++++++++++++++ 2 files changed, 155 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index b653369693d..9b1cc5bc0c6 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -20,6 +20,7 @@ import warnings from collections.abc import Callable from functools import partial +from typing import TypeAlias import torch import torch.distributed as dist @@ -36,7 +37,7 @@ from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method -from .calib import MseCalibrator, NVFP4MSECalibrator +from .calib import MseCalibrator, NVFP4MSECalibrator, _Calibrator from .conversion import create_and_replace_svdquant_linear_on_the_fly, set_quantizer_by_cfg_context from .nn import NVFP4StaticQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer from .utils import ( @@ -56,6 +57,7 @@ from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, GPTQHelper __all__ = [ + "CalibratorFactory", "awq", "layerwise_calibrate", "local_hessian_calibrate", @@ -64,6 +66,28 @@ "svdquant", ] +CalibratorFactory: TypeAlias = Callable[ + [torch.Tensor, int | tuple | list | None, Callable[..., torch.Tensor]], _Calibrator +] + +_FP8_SWEEP_CALIBRATOR_REGISTRY: dict[str, CalibratorFactory] = {} + + +def _register_fp8_sweep_calibrator(backend: str, calibrator_factory: CalibratorFactory) -> None: + """Register a custom calibrator factory for a quantization backend. + + When ``fp8_scale_sweep=True`` is passed to :func:`mse_calibrate`, any weight + quantizer whose ``backend`` attribute matches a registered key will use the + corresponding factory instead of the default :class:`MseCalibrator`. + + Args: + backend: Backend name string (must match ``TensorQuantizer.backend``). + calibrator_factory: Callable with signature + ``(amax: Tensor, axis: int | tuple | list | None, quant_func: Callable)`` + that returns a :class:`_Calibrator` instance. + """ + _FP8_SWEEP_CALIBRATOR_REGISTRY[backend] = calibrator_factory + def weight_only_quantize(model: nn.Module): """Just quantize the weights of the model.""" @@ -341,6 +365,22 @@ def mse_calibrate( # Convert to NVFP4StaticQuantizer in-place NVFP4StaticQuantizer.from_tensor_quantizer(module, global_amax=global_amax) + if fp8_scale_sweep: + # Check if backend has a registered custom calibrator factory. + _backend: str | None = getattr(module, "backend", None) + backend_factory = ( + _FP8_SWEEP_CALIBRATOR_REGISTRY.get(_backend) + if _backend is not None + else None + ) + if backend_factory is not None: + module._calibrator = backend_factory( + initial_amax, + module._calibrator._axis, + partial(_mse_quant_func, quantizer=module), + ) + continue + if fp8_scale_sweep and is_nvfp4_static: # Replace calibrator with NVFP4MSECalibrator module._calibrator = NVFP4MSECalibrator( diff --git a/tests/unit/torch/quantization/test_mse_calibrator.py b/tests/unit/torch/quantization/test_mse_calibrator.py index 5e55465120c..4332b093861 100644 --- a/tests/unit/torch/quantization/test_mse_calibrator.py +++ b/tests/unit/torch/quantization/test_mse_calibrator.py @@ -526,3 +526,117 @@ def quant_func(x, amax): assert a_best.numel() == 2 assert torch.all(torch.isfinite(a_best)) assert torch.all(a_best > 0) + + +class TestRegisterFP8SweepCalibrator: + """Tests for _register_fp8_sweep_calibrator and its dispatch in mse_calibrate.""" + + def setup_method(self): + from modelopt.torch.quantization.model_calib import _FP8_SWEEP_CALIBRATOR_REGISTRY + from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + _QUANT_FUNCTIONAL_BACKENDS, + ) + + self._orig_fp8_registry = dict(_FP8_SWEEP_CALIBRATOR_REGISTRY) + self._orig_quant_backends = dict(_QUANT_FUNCTIONAL_BACKENDS) + + def teardown_method(self): + from modelopt.torch.quantization.model_calib import _FP8_SWEEP_CALIBRATOR_REGISTRY + from modelopt.torch.quantization.nn.modules.tensor_quantizer import ( + _QUANT_FUNCTIONAL_BACKENDS, + ) + + _FP8_SWEEP_CALIBRATOR_REGISTRY.clear() + _FP8_SWEEP_CALIBRATOR_REGISTRY.update(self._orig_fp8_registry) + _QUANT_FUNCTIONAL_BACKENDS.clear() + _QUANT_FUNCTIONAL_BACKENDS.update(self._orig_quant_backends) + + def _quantize_and_calibrate(self, backend_name, fp8_scale_sweep=True): + """Quantize a small Linear with the given backend and run mse_calibrate.""" + import modelopt.torch.quantization as mtq + from modelopt.torch.quantization.model_calib import mse_calibrate + from modelopt.torch.quantization.nn.modules.tensor_quantizer import register_quant_backend + + register_quant_backend(backend_name, lambda x, tq: x) + model = torch.nn.Linear(8, 8, bias=False) + inputs = torch.randn(1, 8) + config = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + { + "quantizer_name": "*weight_quantizer", + "cfg": {"num_bits": 8, "axis": None, "backend": backend_name}, + }, + ], + "algorithm": "max", + } + mtq.quantize(model, config, forward_loop=lambda m: m(inputs)) + mse_calibrate(model, lambda m: m(inputs), fp8_scale_sweep=fp8_scale_sweep) + return model + + def test_register(self): + """_register_fp8_sweep_calibrator stores factories by backend key and allows overwrite.""" + from modelopt.torch.quantization.model_calib import ( + _FP8_SWEEP_CALIBRATOR_REGISTRY, + _register_fp8_sweep_calibrator, + ) + + def factory_a(amax, axis, qf): + return None + + def factory_b(amax, axis, qf): + return None + + _register_fp8_sweep_calibrator("backend_x", factory_a) + assert _FP8_SWEEP_CALIBRATOR_REGISTRY["backend_x"] is factory_a + + _register_fp8_sweep_calibrator("backend_x", factory_b) + assert _FP8_SWEEP_CALIBRATOR_REGISTRY["backend_x"] is factory_b + + def test_mse_calibrate_dispatches_to_registered_factory(self): + """mse_calibrate with fp8_scale_sweep=True calls the registered factory once per quantizer.""" + from modelopt.torch.quantization.calib.mse import MseCalibrator + from modelopt.torch.quantization.model_calib import _register_fp8_sweep_calibrator + + factory_calls: list = [] + + class _RecordingCalibrator(MseCalibrator): + def collect(self, x): + pass + + def compute_amax(self, verbose=False): + return self._initial_amax + + def my_factory(amax, axis, quant_func): + factory_calls.append(amax) + return _RecordingCalibrator(amax=amax, axis=axis, quant_func=quant_func) + + _register_fp8_sweep_calibrator("_test_dispatch", my_factory) + self._quantize_and_calibrate("_test_dispatch", fp8_scale_sweep=True) + + assert len(factory_calls) == 1 + + def test_mse_calibrate_skips_registry_when_fp8_sweep_false(self): + """Registry factory is not invoked when fp8_scale_sweep=False.""" + from modelopt.torch.quantization.model_calib import _register_fp8_sweep_calibrator + + factory_calls: list = [] + + def my_factory(amax, axis, quant_func): + factory_calls.append(amax) + return calib.MseCalibrator(amax=amax, axis=axis, quant_func=quant_func) + + _register_fp8_sweep_calibrator("_test_no_sweep", my_factory) + self._quantize_and_calibrate("_test_no_sweep", fp8_scale_sweep=False) + + assert len(factory_calls) == 0 + + def test_unregistered_backend_uses_default_mse_calibrator(self): + """A quantizer with an unregistered backend falls through to MseCalibrator.""" + from modelopt.torch.quantization.calib.mse import MseCalibrator + + model = self._quantize_and_calibrate("_test_unregistered", fp8_scale_sweep=True) + for module in model.modules(): + if isinstance(module, TensorQuantizer) and module.is_enabled: + if getattr(module, "_calibrator", None) is not None: + assert isinstance(module._calibrator, MseCalibrator) From 289a239ca56276819074f8132410c6002729ffab Mon Sep 17 00:00:00 2001 From: yeyu-nvidia Date: Mon, 20 Apr 2026 12:10:02 -0700 Subject: [PATCH 25/30] fix: use data_dir for directory paths in ShardedDataset (#1301) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - `datasets`' `resolve_pattern` only matches entries with `type=="file"`, so passing a bare directory path as `data_files` to `load_dataset` results in `FileNotFoundError` even when the directory exists on disk - Detect directory paths in `ShardedDataset._load_dataset()` and pass them via `data_dir` instead of `data_files` ## Reproduction ```python from datasets import load_dataset # This fails with FileNotFoundError: load_dataset("json", data_files="/path/to/data_directory") # This works: load_dataset("json", data_dir="/path/to/data_directory") ``` ## Test plan - [ ] Verify existing EAGLE3/DFlash training pipelines that pass directory paths work - [ ] Verify file path and glob patterns still work (falls through to `data_files`) - [ ] Verify `data_files=None` (no data_files arg) still works 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit ## Bug Fixes * Fixed an issue with dataset loading that prevented proper handling of directory-based data sources. Directories are now correctly detected and processed during dataset initialization. Signed-off-by: Ye Yu Co-authored-by: Claude Opus 4.6 --- modelopt/torch/utils/plugins/transformers_dataset.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/utils/plugins/transformers_dataset.py b/modelopt/torch/utils/plugins/transformers_dataset.py index 547c0ae06d5..56b1e4f07b1 100644 --- a/modelopt/torch/utils/plugins/transformers_dataset.py +++ b/modelopt/torch/utils/plugins/transformers_dataset.py @@ -73,10 +73,20 @@ def __getitem__(self, index): return self._raw_samples[index] def _load_dataset(self): + # datasets' resolve_pattern only matches entries with type=="file", so passing + # a bare directory path as data_files results in FileNotFoundError. + # Use data_dir for directory paths instead. + data_dir = None + data_files = self.data_files + if data_files and os.path.isdir(data_files): + data_dir = data_files + data_files = None + dataset = load_dataset( self.name, self.subset, - data_files=self.data_files, + data_files=data_files, + data_dir=data_dir, split=self.split, # num_proc=4, # TODO: Make this configurable streaming=self.num_streaming_samples is not None, From 355c6b7883789d738bc3b5a91bdc73d7e0313d8d Mon Sep 17 00:00:00 2001 From: "Chenhan D. Yu" <5185878+ChenhanYu@users.noreply.github.com> Date: Mon, 20 Apr 2026 12:43:58 -0700 Subject: [PATCH 26/30] fix: PTQ 1GPU, export PP divisibility, hidden states conversations key (#1293) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - **megatron_lm_ptq.yaml**: Qwen3-8B PTQ to single GPU for L40 clusters (TP=1, all tasks) - **quantize.sh**: Auto-find largest PP dividing model's `num_hidden_layers` for export step. Qwen3-8B has 36 layers which isn't divisible by 8, causing `AssertionError` on 8-GPU nodes - **compute_hidden_states_trtllm.py**: Use `messages` with `conversations` fallback, matching the HF version. Fixes `KeyError: 'conversations'` when data uses OpenAI `messages` format ## Test plan - [x] Qwen3-8B PTQ runs on single L40 GPU - [x] Export PP auto-selects valid divisor (36 layers → PP=6 on 8 GPUs, PP=4 on 4 GPUs, PP=1 on 1 GPU) - [x] EAGLE3 offline pipeline reads data with `messages` field 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **New Features** * Dataset input handling now supports multiple field formats for enhanced compatibility. * **Bug Fixes** * Optimized GPU resource allocation during model quantization with improved pipeline parallelism computation. * Updated quantization configuration for more efficient resource utilization. Signed-off-by: Chenhan Yu Co-authored-by: Claude Opus 4.6 (1M context) --- .../compute_hidden_states_trtllm.py | 2 +- .../common/megatron_lm/quantize/quantize.sh | 17 ++++++++++++++--- .../examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml | 16 ++++++++-------- 3 files changed, 23 insertions(+), 12 deletions(-) diff --git a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py index 0bf68e430fb..06531a16771 100644 --- a/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py +++ b/examples/speculative_decoding/collect_hidden_states/compute_hidden_states_trtllm.py @@ -256,7 +256,7 @@ async def submit_generates(): for entry in dataset: conversation_id = entry.get("conversation_id", entry.get("uuid")) - conversations = entry["conversations"] + conversations = entry.get("messages") or entry.get("conversations") if not conversations or not isinstance(conversations, list): num_invalid += 1 continue diff --git a/tools/launcher/common/megatron_lm/quantize/quantize.sh b/tools/launcher/common/megatron_lm/quantize/quantize.sh index 1bb0d60e80d..407a6743780 100755 --- a/tools/launcher/common/megatron_lm/quantize/quantize.sh +++ b/tools/launcher/common/megatron_lm/quantize/quantize.sh @@ -41,11 +41,22 @@ TP=${TP:-1} PP=${PP:-1} EP=${EP:-1} ETP=${ETP:-1} ${QUANTIZE_EXE} ${MLM_MODEL_CF export MLM_EXTRA_ARGS="--mmlu-dataset ${MMLU_DATASET:-/hf-local/cais/mmlu} --fraction 0.01 --lower-bound ${MMLU_LOWER_BOUND:-0.38} --disable-tqdm" TP=${TP:-1} PP=${PP:-1} EP=${EP:-1} ETP=${ETP:-1} MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${MMLU_EXE} ${MLM_MODEL_CFG} -# Export quantized checkpoint to HF format (PP=all GPUs) +# Export quantized checkpoint to HF format +# Use largest PP <= total GPUs that divides the model's num_hidden_layers TOTAL_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())" 2>/dev/null || echo ${NUM_GPUS:-1}) -echo "=== Exporting ${MLM_MODEL_CFG} ${QUANT_CFG} (PP=${TOTAL_GPUS}) ===" +EXPORT_PP=$(python3 -c " +import json, os +cfg = os.path.join('${HF_MODEL_CKPT}', 'config.json') +n_layers = json.load(open(cfg)).get('num_hidden_layers', 1) if os.path.exists(cfg) else 1 +gpus = ${TOTAL_GPUS} +pp = gpus +while pp > 1 and n_layers % pp != 0: + pp -= 1 +print(pp) +" 2>/dev/null || echo ${TOTAL_GPUS}) +echo "=== Exporting ${MLM_MODEL_CFG} ${QUANT_CFG} (PP=${EXPORT_PP}, ${TOTAL_GPUS} GPUs) ===" export MLM_EXTRA_ARGS= -TP=1 PP=${TOTAL_GPUS} EP=1 ETP=1 MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${EXPORT_EXE} ${MLM_MODEL_CFG} +TP=1 PP=${EXPORT_PP} EP=1 ETP=1 MLM_MODEL_CKPT=${MLM_MODEL_SAVE} ${EXPORT_EXE} ${MLM_MODEL_CFG} ls ${EXPORT_DIR} cat ${EXPORT_DIR}/hf_quant_config.json diff --git a/tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml b/tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml index 33b9da18e66..ff55a92e39f 100644 --- a/tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml +++ b/tools/launcher/examples/Qwen/Qwen3-8B/megatron_lm_ptq.yaml @@ -24,7 +24,7 @@ pipeline: config: model: Qwen/Qwen3-8B quant_cfg: NVFP4_DEFAULT_CFG - tp: 8 + tp: 1 calib_dataset: abisee/cnn_dailymail calib_size: 32 mmlu_dataset: cais/mmlu @@ -33,15 +33,15 @@ pipeline: slurm_config: _factory_: "slurm_factory" nodes: 1 - ntasks_per_node: 8 - gpus_per_node: 8 + ntasks_per_node: 1 + gpus_per_node: 1 task_1: _target_: common.megatron_lm.quantize.task.MegatronLMQuantizeTask config: model: Qwen/Qwen3-8B quant_cfg: FP8_DEFAULT_CFG - tp: 8 + tp: 1 calib_dataset: abisee/cnn_dailymail calib_size: 32 mmlu_dataset: cais/mmlu @@ -50,18 +50,18 @@ pipeline: slurm_config: _factory_: "slurm_factory" nodes: 1 - ntasks_per_node: 8 - gpus_per_node: 8 + ntasks_per_node: 1 + gpus_per_node: 1 # Step 3: TRT-LLM eval MMLU on all exported checkpoints task_2: script: common/tensorrt_llm/eval.sh environment: - HF_MODEL_CKPT: /scratchspace/export - - TP: "8" + - TP: "1" - EP: "1" slurm_config: _factory_: "slurm_factory" nodes: 1 ntasks_per_node: 1 - gpus_per_node: 8 + gpus_per_node: 1 From 2fef374deda8ea0374ff89e54e566af79b683a5d Mon Sep 17 00:00:00 2001 From: yeyu-nvidia Date: Mon, 20 Apr 2026 13:39:36 -0700 Subject: [PATCH 27/30] fix: auto-compute dp_replicate_size from world_size (#1302) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary - When `dp_shard_size < world_size` (e.g., `dp_shard_size=4` on 8 GPUs across 2 nodes), `ParallelismConfig` raises `total_size (4) does not match num_processes (8)` because `dp_replicate_size` defaults to 1 - Auto-compute `dp_replicate_size = world_size // (dp_shard_size * cp_size)` so intra-node FSDP2 sharding + inter-node data-parallel replication works without manual config - This enables `dp_shard_size` to be set to per-node GPU count (better NVLink utilization) while automatically creating replicas across nodes ## Test plan - [ ] Verify single-node training (dp_shard_size == world_size, dp_replicate_size == 1) unchanged - [ ] Verify multi-node with dp_shard_size < world_size creates correct replica groups - [ ] Verify existing EAGLE3/DFlash configs still work 🤖 Generated with [Claude Code](https://claude.com/claude-code) ## Summary by CodeRabbit * **Refactor** * Enhanced parallelism configuration initialization in the speculative decoding example to better handle distributed training scenarios. Signed-off-by: Ye Yu Co-authored-by: Claude Opus 4.6 --- examples/speculative_decoding/main.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/examples/speculative_decoding/main.py b/examples/speculative_decoding/main.py index efc4ba82bd3..31c73d04275 100644 --- a/examples/speculative_decoding/main.py +++ b/examples/speculative_decoding/main.py @@ -212,8 +212,23 @@ def train(): "Either data.data_path or data.offline_data_path must be set in the config." ) if training_args.cp_size > 1 or training_args.dp_shard_size > 1: + # Auto-compute dp_replicate_size so that + # dp_replicate_size * dp_shard_size * cp_size == world_size. + # Note: torch.cuda.device_count() returns per-node GPU count, not world_size. + # WORLD_SIZE (set by torchrun/accelerate) gives the correct multi-node total. + world_size = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count())) + parallel_size = training_args.dp_shard_size * training_args.cp_size + if world_size % parallel_size != 0: + raise ValueError( + f"world_size ({world_size}) must be divisible by " + f"dp_shard_size ({training_args.dp_shard_size}) * cp_size ({training_args.cp_size}) " + f"= {parallel_size}" + ) + dp_replicate_size = world_size // parallel_size training_args.parallelism_config = ParallelismConfig( - cp_size=training_args.cp_size, dp_shard_size=training_args.dp_shard_size + cp_size=training_args.cp_size, + dp_shard_size=training_args.dp_shard_size, + dp_replicate_size=dp_replicate_size, ) if training_args.cp_size > 1: patch_ring_attention_for_ttt() From 5ffb8487d921bc7eb323edf9d060792460e69ead Mon Sep 17 00:00:00 2001 From: sychen52 <41452870+sychen52@users.noreply.github.com> Date: Tue, 21 Apr 2026 10:50:04 -0700 Subject: [PATCH 28/30] add gptq fused kernel (#1291) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Add gptq fused kernel to improve speed. ### Usage check unittest ### Testing added a unittest ### Before your PR is "*Ready for review*" Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md) and your commits are signed (`git commit -s -S`). Make sure you read and follow the [Security Best Practices](https://github.com/NVIDIA/Model-Optimizer/blob/main/SECURITY.md#security-coding-practices-for-contributors) (e.g. avoiding hardcoded `trust_remote_code=True`, `torch.load(..., weights_only=False)`, `pickle`, etc.). - Is this change backward compatible?: ✅ / ❌ / N/A - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: ✅ / ❌ / N/A - Did you write any new necessary tests?: ✅ / ❌ / N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ✅ / ❌ / N/A ### Additional Information ## Summary by CodeRabbit * **New Features** * Fused GPTQ backend for faster blockwise weight updates, toggleable via a new "fused" option. * Shared NVFP4 quantization primitives exposed for reuse. * **Refactor** * Consolidated FP4 scale/quantization logic into reusable utilities and centralized Hessian inversion handling. * **Tests** * Expanded GPU tests comparing fused vs unfused GPTQ, added Triton-availability gating and a local benchmark entrypoint. --------- Signed-off-by: Shiyang Chen --- modelopt/torch/quantization/config.py | 6 + modelopt/torch/quantization/model_calib.py | 4 +- .../torch/quantization/triton/fp4_kernel.py | 93 ++++----- .../quantization/triton/fp4_kernel_hopper.py | 31 +-- .../quantization/triton/gptq_fused_kernel.py | 136 +++++++++++++ .../torch/quantization/triton/nvfp4_quant.py | 144 +++++++++++++ .../torch/quantization/utils/calib_utils.py | 189 +++++++++++++----- tests/gpu/torch/quantization/conftest.py | 9 + tests/gpu/torch/quantization/test_gptq.py | 179 ++++++++++++++++- 9 files changed, 652 insertions(+), 139 deletions(-) create mode 100644 modelopt/torch/quantization/triton/gptq_fused_kernel.py create mode 100644 modelopt/torch/quantization/triton/nvfp4_quant.py diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 3f24ac09a41..186ff1c7edd 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -1549,6 +1549,12 @@ class GPTQCalibConfig(QuantizeAlgorithmConfig): description="""The block size for GPTQ weight update, which must be a multiple of the group_size used in the quantization.""", ) + fused: bool = ModeloptField( + default=False, + title="Use fused Triton kernel for GPTQ.", + description="""When True, use a fused Triton kernel that combines quantization and + per-column error propagation into one launch per GPTQ block.""", + ) QuantizeQuantCfgType = list[QuantizerCfgEntry] diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 9b1cc5bc0c6..04aaa88a519 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -1698,6 +1698,7 @@ def gptq( forward_loop: ForwardLoop, perc_damp: float = 0.01, block_size: int = 128, + fused: bool = False, ): """GPTQ quantization. @@ -1723,6 +1724,7 @@ def gptq( forward_loop: Callable that replays calibration inputs through *model*. perc_damp: Percentage of avg Hessian diagonal for damping (default: 0.01). block_size: Block size for GPTQ weight update. + fused: If True, use fused Triton kernel for NVFP4 static quantization. """ total_start = time.time() @@ -1745,7 +1747,7 @@ def _make_gptq_handle(name, m): cls = GPTQHelper else: cls = _GPTQ_HELPER_REGISTRY.get(backend, GPTQHelper) - return cls(m, name, offload_to_cpu=True) + return cls(m, name, offload_to_cpu=True, fused=fused) gptq_handles = {name: _make_gptq_handle(name, m) for name, m in quantized_layers} for handle in gptq_handles.values(): diff --git a/modelopt/torch/quantization/triton/fp4_kernel.py b/modelopt/torch/quantization/triton/fp4_kernel.py index 63a8b3dcb72..9eb6b2d49f5 100644 --- a/modelopt/torch/quantization/triton/fp4_kernel.py +++ b/modelopt/torch/quantization/triton/fp4_kernel.py @@ -24,7 +24,9 @@ import triton import triton.language as tl -__all__ = ["fp4_dequantize", "static_blockwise_fp4_fake_quant"] +from .nvfp4_quant import nvfp4_scalar_quant + +__all__ = ["compute_fp4_scales", "fp4_dequantize", "static_blockwise_fp4_fake_quant"] _TORCH_TO_TL_DTYPE = { @@ -198,52 +200,47 @@ def static_blockwise_fp4_fake_quant_kernel( idx = block_offset + tl.arange(0, BLOCK_SIZE) scale = tl.load(scale_ptr + pid).to(tl.float32) - x = tl.load(x_ptr + idx).to(tl.float32) - x_abs = tl.abs(x) - # If scale is 0, inf, or nan, use 1.0 (matching CUDA kernel behavior) - # Note: (x != x) checks if x is NaN per IEEE 754 - scale_safe = tl.where( - (scale == 0) | (scale != scale) | (tl.abs(scale) == float("inf")), # noqa: PLR0124 - 1.0, - scale, - ) - abs_scaled = x_abs / scale_safe - - # FP4 values: 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0 - q_val = tl.where( - abs_scaled <= 0.25, - 0.0, - tl.where( - abs_scaled < 0.75, - 0.5, - tl.where( - abs_scaled <= 1.25, - 1.0, - tl.where( - abs_scaled < 1.75, - 1.5, - tl.where( - abs_scaled <= 2.5, - 2.0, - tl.where( - abs_scaled < 3.5, - 3.0, - tl.where(abs_scaled <= 5.0, 4.0, 6.0), - ), - ), - ), - ), - ), - ) - - x_rescaled = q_val * scale_safe - x_quant = tl.where(x >= 0, x_rescaled, -x_rescaled) + x_quant = nvfp4_scalar_quant(x, scale, BLOCK_SIZE) tl.store(y_ptr + idx, x_quant.to(OUT_DTYPE)) +def compute_fp4_scales( + amax: torch.Tensor, + global_amax: torch.Tensor | None = None, + quantize_block_scales: bool = True, +) -> torch.Tensor: + """Compute per-block FP4 scales from amax values. + + ``scale = amax / 6.0``, optionally quantized to FP8 E4M3. + + Args: + amax: Per-block amax values (any shape). + global_amax: Global amax for FP8 two-level scaling. Computed from *amax* if None. + quantize_block_scales: If True, quantize scales to FP8 E4M3. + + Returns: + Per-block scales (same shape as *amax*), float32. + """ + amax = amax.float() + scale = amax / 6.0 # FP4 max representable value is 6.0 + + if quantize_block_scales: + from modelopt.torch.quantization.tensor_quant import scaled_e4m3_impl + from modelopt.torch.quantization.utils import reduce_amax + + if global_amax is None: + global_amax = reduce_amax(amax, axis=None, keepdims=False, squeeze_scalar=True) + + global_amax = global_amax.float() + scale_fp8_quant_amax = global_amax / 6.0 + scale = scaled_e4m3_impl(scale, scale_fp8_quant_amax) + + return scale + + def static_blockwise_fp4_fake_quant( x: torch.Tensor, amax: torch.Tensor, @@ -266,19 +263,7 @@ def static_blockwise_fp4_fake_quant( if out_dtype is None: out_dtype = x.dtype - amax = amax.float() # Requires to be in float32 - scale = amax / 6.0 # FP4 max representable value is 6.0 - - if quantize_block_scales: - from modelopt.torch.quantization.tensor_quant import scaled_e4m3_impl - from modelopt.torch.quantization.utils import reduce_amax - - if global_amax is None: - global_amax = reduce_amax(amax, axis=None, keepdims=False, squeeze_scalar=True) - - global_amax = global_amax.float() - scale_fp8_quant_amax = global_amax / 6.0 - scale = scaled_e4m3_impl(scale, scale_fp8_quant_amax) + scale = compute_fp4_scales(amax, global_amax, quantize_block_scales) x_flat = x.contiguous().view(-1) y_flat = torch.empty_like(x_flat, dtype=out_dtype) diff --git a/modelopt/torch/quantization/triton/fp4_kernel_hopper.py b/modelopt/torch/quantization/triton/fp4_kernel_hopper.py index 2ec31863efc..624e723b957 100644 --- a/modelopt/torch/quantization/triton/fp4_kernel_hopper.py +++ b/modelopt/torch/quantization/triton/fp4_kernel_hopper.py @@ -24,6 +24,7 @@ import triton.language as tl from .fp4_kernel import _torch_dtype_to_tl +from .nvfp4_quant import fp4_round_magnitude, fp8_quantize_scale __all__ = ["fp4_fake_quant_block"] @@ -79,9 +80,7 @@ def fp4_fake_quant_kernel( block_max = tl.max(x_abs, axis=2, keep_dims=True) - block_max_scaled = block_max / (6.0 * global_scale_safe) - block_max_scaled = tl.minimum(block_max_scaled, 448.0) - block_max_quant = block_max_scaled.to(tl.float8e4nv).to(tl.float32) * global_scale + block_max_quant = fp8_quantize_scale(block_max, global_scale_safe) block_max_quant = tl.where(block_max_quant >= 1e-5, block_max_quant, 1.0) block_max_quant_broadcast = tl.broadcast_to( @@ -90,31 +89,7 @@ def fp4_fake_quant_kernel( abs_scaled = x_abs / block_max_quant_broadcast - q_val = tl.where( - abs_scaled <= 0.25, - 0.0, - tl.where( - abs_scaled < 0.75, - 0.5, - tl.where( - abs_scaled <= 1.25, - 1.0, - tl.where( - abs_scaled < 1.75, - 1.5, - tl.where( - abs_scaled <= 2.5, - 2.0, - tl.where( - abs_scaled < 3.5, - 3.0, - tl.where(abs_scaled <= 5.0, 4.0, 6.0), - ), - ), - ), - ), - ), - ) + q_val = fp4_round_magnitude(abs_scaled) x_rescaled = q_val * block_max_quant_broadcast x_rescaled = tl.where(tile_reshaped >= 0, x_rescaled, -x_rescaled) diff --git a/modelopt/torch/quantization/triton/gptq_fused_kernel.py b/modelopt/torch/quantization/triton/gptq_fused_kernel.py new file mode 100644 index 00000000000..c070eac8800 --- /dev/null +++ b/modelopt/torch/quantization/triton/gptq_fused_kernel.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fused Triton kernels for GPTQ blockwise weight-update. + +A kernel for scalar (NVFP4) quantization with inline two-level scale computation. +Fuses scale computation + quantization + per-column GPTQ error propagation into +one launch per GPTQ block, avoiding the Python-level per-column loop. + +Architecture: + - One Triton program per output row. + - ``w_full [BLOCK_SIZE]`` register tensor holds working weights. + - Per-column: calls ``nvfp4_scalar_qdq()`` for FP4 QDQ with inline scale + computation, then propagates error via ``w_full -= err * h_inv_row``. +""" + +import torch +import triton +import triton.language as tl + +from .nvfp4_quant import nvfp4_scalar_qdq + +__all__ = ["gptq_fused_block_scalar"] + + +# --------------------------------------------------------------------------- +# Scalar kernel — NVFP4 QDQ + error propagation +# --------------------------------------------------------------------------- + + +@triton.jit +def _gptq_scalar_kernel( + w_ptr, + qw_ptr, + err_ptr, + amax_ptr, + global_scale, + hinv_ptr, + num_rows, + n_amax_blocks, + quant_block_size, + block_start, + BLOCK_SIZE: tl.constexpr, +): + row = tl.program_id(0) + if row >= num_rows: + return + + w_base = w_ptr + row * BLOCK_SIZE + qw_base = qw_ptr + row * BLOCK_SIZE + err_base = err_ptr + row * BLOCK_SIZE + amax_base = amax_ptr + row * n_amax_blocks + + j_range = tl.arange(0, BLOCK_SIZE) + w_full = tl.load(w_base + j_range) + + for col in range(0, BLOCK_SIZE, 1): + block_amax = tl.load(amax_base + (block_start + col) // quant_block_size) + + w_scalar = tl.sum(tl.where(j_range == col, w_full, 0.0)) + q_scalar = tl.sum( + nvfp4_scalar_qdq( + tl.full([1], w_scalar, dtype=tl.float32), + block_amax, + global_scale, + 1, + ) + ) + + d_val = tl.load(hinv_ptr + col * BLOCK_SIZE + col) + err_val = (w_scalar - q_scalar) / d_val + tl.store(err_base + col, err_val) + tl.store(qw_base + col, q_scalar) + + remaining = j_range > col + hinv_row = tl.load(hinv_ptr + col * BLOCK_SIZE + j_range, mask=remaining, other=0.0) + w_full = w_full - err_val * hinv_row + + +def gptq_fused_block_scalar( + w_block: torch.Tensor, + block_amax: torch.Tensor, + global_scale: float, + h_inv_cho_blk: torch.Tensor, + quant_block_size: int, + block_start: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run scalar GPTQ (NVFP4) column loop for one block in a single Triton kernel launch. + + Computes FP8-quantized scales from per-block amax inline via + :func:`nvfp4_scalar_qdq`, then performs NVFP4 fake quantization and + GPTQ error propagation per column. + + Args: + w_block: Working weights ``[num_rows, block_size]`` (float32). + block_amax: Per-block amax values ``[num_rows, n_amax_blocks]`` (float32). + global_scale: Pre-computed ``global_amax / (6.0 * 448.0)`` (scalar). + h_inv_cho_blk: Block of upper-Cholesky inverse Hessian ``[block_size, block_size]``. + quant_block_size: Number of elements sharing one scale factor. + block_start: Column offset of this block in the full weight matrix. + + Returns: + ``(qw_block, err_block)`` each ``[num_rows, block_size]``. + """ + num_rows, block_size = w_block.shape + + qw_block = torch.empty_like(w_block) + err_block = torch.empty_like(w_block) + + _gptq_scalar_kernel[(num_rows,)]( + w_block.contiguous(), + qw_block, + err_block, + block_amax.contiguous(), + global_scale, + h_inv_cho_blk.contiguous(), + num_rows, + block_amax.shape[1], + quant_block_size, + block_start, + BLOCK_SIZE=block_size, + ) + + return qw_block, err_block diff --git a/modelopt/torch/quantization/triton/nvfp4_quant.py b/modelopt/torch/quantization/triton/nvfp4_quant.py new file mode 100644 index 00000000000..32ab776b2b4 --- /dev/null +++ b/modelopt/torch/quantization/triton/nvfp4_quant.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Composable Triton JIT functions for NVFP4 (E2M1) fake quantization. + +Single source of truth for FP4 decision-boundary rounding. Used by: + - ``fp4_kernel.py`` (standalone blockwise fake quant) + - ``fp4_kernel_hopper.py`` (Hopper block-pointer variant) + - ``gptq_fused_kernel.py`` (fused GPTQ scalar path) + +FP4 (E2M1) representable magnitudes: {0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0} +""" + +import triton +import triton.language as tl +from triton.language.extra.cuda import libdevice + + +@triton.jit +def fp4_round_magnitude(abs_scaled): + """Round ``|x| / scale`` to the nearest FP4 (E2M1) magnitude. + + Works with any tensor shape — the caller is responsible for computing + ``abs_scaled = |x| / scale`` beforehand. + + Returns: + Tensor of same shape as *abs_scaled* with values in + {0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0}. + """ + return tl.where( + abs_scaled <= 0.25, + 0.0, + tl.where( + abs_scaled < 0.75, + 0.5, + tl.where( + abs_scaled <= 1.25, + 1.0, + tl.where( + abs_scaled < 1.75, + 1.5, + tl.where( + abs_scaled <= 2.5, + 2.0, + tl.where(abs_scaled < 3.5, 3.0, tl.where(abs_scaled <= 5.0, 4.0, 6.0)), + ), + ), + ), + ), + ) + + +@triton.jit +def nvfp4_scalar_quant( + x, # [N] float32, already loaded + scale, # float32 scalar: pre-computed block scale (amax / 6.0) + N: tl.constexpr, +): + """NVFP4 scalar fake quantization for a group of elements sharing one scale. + + Quantizes each element independently: divide by scale, round to nearest + FP4 (E2M1) value via ``fp4_round_magnitude``, multiply by scale. + + Args: + x: [N] float32 tensor of values to quantize (already in registers). + scale: float32 scalar block scale. + N: Compile-time number of elements. + + Returns: + x_quant: [N] float32, fake-quantized values. + """ + x_abs = tl.abs(x) + # Guard against degenerate scale (matching CUDA kernel behavior) + scale_safe = tl.where( + (scale == 0.0) | libdevice.isnan(scale) | (tl.abs(scale) == float("inf")), + 1.0, + scale, + ) + abs_scaled = x_abs / scale_safe + q_val = fp4_round_magnitude(abs_scaled) + x_rescaled = q_val * scale_safe + x_quant = tl.where(x >= 0, x_rescaled, -x_rescaled) + return x_quant + + +@triton.jit +def fp8_quantize_scale(block_amax, global_scale): + """FP8 E4M3 fake-quantize the per-block NVFP4 scale. + + Computes ``scale = block_amax / 6.0``, then round-trips it through + FP8 E4M3 using ``global_scale`` for the second-level scaling. + + Works with any tensor shape (scalar, 1-D, or higher) since all ops + are element-wise. + + Args: + block_amax: Per-block amax value(s). + global_scale: Pre-computed ``global_amax / (6.0 * 448.0)``. + + Returns: + FP8-quantized per-block scale(s), same shape as ``block_amax``. + """ + FP8_E4M3_MAX: tl.constexpr = 448.0 + scale_in_fp8_range = block_amax / (6.0 * global_scale) + scale_clamped = tl.minimum(scale_in_fp8_range, FP8_E4M3_MAX) + return scale_clamped.to(tl.float8e4nv).to(tl.float32) * global_scale + + +@triton.jit +def nvfp4_scalar_qdq( + x, # [N] float32, already loaded + block_amax, # float32 scalar: per-block amax + global_scale, # float32 scalar: pre-computed global_amax / (6.0 * 448.0) + N: tl.constexpr, +): + """NVFP4 scalar fake quantization with inline two-level scale computation. + + Computes the per-block FP8-quantized scale from ``block_amax`` via + :func:`fp8_quantize_scale`, then quantizes each element to the nearest + FP4 (E2M1) value. + + Args: + x: [N] float32 tensor of values to quantize. + block_amax: Per-block amax (absolute maximum of the block). + global_scale: Pre-computed ``global_amax / (6.0 * 448.0)``. + N: Compile-time number of elements. + + Returns: + x_quant: [N] float32, fake-quantized values. + """ + scale = fp8_quantize_scale(block_amax, global_scale) + return nvfp4_scalar_quant(x, scale, N) diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 252f0af6fc8..ac2ec7a2553 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -74,6 +74,42 @@ def update_hessian(input, hessian, n_samples): return hessian, n_samples +def compute_hessian_inverse(hessian, weight, perc_damp): + """Compute damped upper-Cholesky inverse Hessian. + + Dead-neuron columns (all-zero in ``weight``) are zeroed in the + Hessian before inversion, matching the FP-Quant reference: + https://github.com/IST-DASLab/FP-Quant/blob/d2e3092f968262c4de5fb050e1aef568a280dadd/src/quantization/gptq.py#L200 + + Args: + hessian: Hessian matrix ``[in_features, in_features]``. + weight: Weight matrix ``[out_features, in_features]`` for dead-neuron detection. + perc_damp: Percentage of average Hessian diagonal for damping. + + Returns: + Upper-triangular Cholesky factor of the damped inverse Hessian + ``[in_features, in_features]``. Falls back to the identity matrix + when the Hessian is not positive definite. + """ + h = hessian.clone() + zero_cols = torch.nonzero(weight.eq(0).all(dim=0)).unsqueeze(-1) + + h[zero_cols, :] = 0 + h[:, zero_cols] = 0 + h[zero_cols, zero_cols] = 1 + + damp = perc_damp * torch.mean(torch.diag(h)) + diag_indices = torch.arange(h.shape[0], device=h.device) + h[diag_indices, diag_indices] += damp + + try: + h = torch.cholesky_inverse(torch.linalg.cholesky(h)) + return torch.linalg.cholesky(h, upper=True) + except (RuntimeError, torch.linalg.LinAlgError): + print_rank_0("Warning: Hessian is not positive definite, using identity matrix") + return torch.eye(h.shape[0], device=h.device, dtype=h.dtype) + + class GPTQHelper: """Encapsulates per-module GPTQ state and operations. @@ -90,10 +126,11 @@ class GPTQHelper: CACHE_NAME = "_forward_no_gptq_hessian" - def __init__(self, module, name, offload_to_cpu=False): + def __init__(self, module, name, offload_to_cpu=False, fused=False): """Initialize GPTQHelper with module state and Hessian storage.""" self.module = module self.name = name + self.fused = fused in_features = module.weight.shape[-1] device = module.weight.device if device.type == "meta" or (offload_to_cpu and get_used_gpu_mem_fraction(device) > 0.65): @@ -154,73 +191,40 @@ def update_weights(self, block_size, perc_damp): # ------------------------------------------------------------------ def _prepare_hessian_inverse(self, hessian, perc_damp): - """Compute damped inverse Hessian and store as ``self.h_inv``. - - Dead-neuron columns (all-zero in ``self.weight``) are zeroed in the - Hessian before inversion, matching the FP-Quant reference: - https://github.com/IST-DASLab/FP-Quant/blob/d2e3092f968262c4de5fb050e1aef568a280dadd/src/quantization/gptq.py#L200 - """ + """Compute damped inverse Hessian and store as ``self.h_inv``.""" assert self.weight is not None, "_prepare_hessian_inverse called before update_weights()" - h = hessian.clone() - zero_cols = torch.nonzero(self.weight.eq(0).all(dim=0)).unsqueeze(-1) - - h[zero_cols, :] = 0 - h[:, zero_cols] = 0 - h[zero_cols, zero_cols] = 1 - - damp = perc_damp * torch.mean(torch.diag(h)) - diag_indices = torch.arange(h.shape[0], device=h.device) - h[diag_indices, diag_indices] += damp - - try: - h = torch.cholesky_inverse(torch.linalg.cholesky(h)) - self.h_inv = torch.linalg.cholesky(h, upper=True) - except (RuntimeError, torch.linalg.LinAlgError): - print_rank_0("Warning: Hessian is not positive definite, using identity matrix") - self.h_inv = torch.eye(h.shape[0], device=h.device, dtype=h.dtype) + self.h_inv = compute_hessian_inverse(hessian, self.weight, perc_damp) def _blockwise_update(self, block_size): - """Column-wise GPTQ update using full-matrix QDQ. + """Column-wise GPTQ update. - For each column, quantizes the full weight matrix via the quantizer and - extracts the quantized column. This is the standard GPTQ approach. - - Reads/writes ``self.weight`` and ``self.h_inv`` in-place. + When ``self.fused`` is True and the weight quantizer is an + ``NVFP4StaticQuantizer``, uses :func:`gptq_blockwise_update_fused_scalar` + (a fused Triton kernel). Otherwise falls back to + :func:`gptq_blockwise_update` (unfused column-by-column loop). """ assert self.weight is not None and self.h_inv is not None, ( "_blockwise_update called before _prepare_hessian_inverse()" ) quantizer = self.module.weight_quantizer - block_sizes = getattr(quantizer, "block_sizes", None) - if block_sizes is not None: - group_size = block_sizes.get(-1) - if group_size is not None and block_size % group_size != 0: + + if self.fused and getattr(quantizer, "_is_nvfp4_static_quantizer", False): + block_sizes = quantizer.block_sizes + quant_block_size = block_sizes.get(-1) or block_sizes.get(1) + if quant_block_size is not None and block_size % quant_block_size != 0: raise ValueError( f"GPTQ block_size ({block_size}) must be divisible by the quantizer" - f" group_size ({group_size})" + f" group_size ({quant_block_size})" ) - num_cols = self.weight.shape[1] - - for block_start in range(0, num_cols, block_size): - block_end = min(block_start + block_size, num_cols) - n_cols_blk = block_end - block_start - h_inv_cho_blk = self.h_inv[block_start:block_end, block_start:block_end] - - wblk = self.weight.clone() - errs = torch.zeros_like(wblk[:, block_start:block_end]) - - for i in range(n_cols_blk): - w_ci = wblk[:, block_start + i] - d = h_inv_cho_blk[i, i] - qdq = quantizer(wblk) - self.weight[:, block_start + i] = qdq[:, block_start + i] - err = (w_ci - qdq[:, block_start + i]) / d - wblk[:, block_start + i : block_end].addr_(err, h_inv_cho_blk[i, i:], alpha=-1) - errs[:, i] = err - - self.weight[:, block_end:].addmm_( - errs, self.h_inv[block_start:block_end, block_end:], alpha=-1 + out_features, num_cols = self.weight.shape + n_blocks = num_cols // quant_block_size + block_amax = quantizer.amax.reshape(out_features, n_blocks).float() + global_scale = quantizer.global_amax.float().item() / (6.0 * 448.0) + gptq_blockwise_update_fused_scalar( + self.weight, block_amax, global_scale, self.h_inv, block_size, quant_block_size ) + else: + gptq_blockwise_update(self.weight, self.h_inv, block_size, quantizer) def _print_mse_error(self, hessian): """Log Hessian-weighted relative MSE between ``self.weight`` and original weights.""" @@ -231,6 +235,81 @@ def _print_mse_error(self, hessian): print_rank_0(f"[{self.name}] Relative MSE error: {mse.item():.2e}{suffix}") +def gptq_blockwise_update(weight, h_inv, block_size, quantize_fn): + """Column-wise GPTQ update using full-matrix fake quantization. + + For each column, quantizes the full weight matrix via ``quantize_fn`` and + extracts the quantized column. Error is propagated to remaining columns + within the block and then to all subsequent columns via the inverse Hessian. + + Args: + weight: Weight tensor ``[out_features, in_features]``, modified **in-place** + with fake-quantized values. + h_inv: Upper-triangular Cholesky factor of the damped inverse Hessian + ``[in_features, in_features]``. + block_size: Number of columns to process per GPTQ block. + quantize_fn: Callable ``(weight) -> qdq_weight`` that fake-quantizes + the full weight matrix. + """ + num_cols = weight.shape[1] + + for block_start in range(0, num_cols, block_size): + block_end = min(block_start + block_size, num_cols) + n_cols_blk = block_end - block_start + h_inv_cho_blk = h_inv[block_start:block_end, block_start:block_end] + + wblk = weight.clone() + errs = torch.zeros_like(weight[:, block_start:block_end]) + + for i in range(n_cols_blk): + w_ci = wblk[:, block_start + i] + d = h_inv_cho_blk[i, i] + qdq = quantize_fn(wblk) + weight[:, block_start + i] = qdq[:, block_start + i] + err = (w_ci - qdq[:, block_start + i]) / d + wblk[:, block_start + i : block_end].addr_(err, h_inv_cho_blk[i, i:], alpha=-1) + errs[:, i] = err + + weight[:, block_end:].addmm_(errs, h_inv[block_start:block_end, block_end:], alpha=-1) + + +def gptq_blockwise_update_fused_scalar( + weight, block_amax, global_scale, h_inv, block_size, quant_block_size +): + """Fused GPTQ blockwise update for NVFP4 scalar quantization. + + Uses a fused Triton kernel that combines scale computation, quantization, + and per-column error propagation into one launch per GPTQ block, avoiding + the Python-level per-column loop in :func:`gptq_blockwise_update`. + + Args: + weight: Weight tensor ``[out_features, in_features]``, modified **in-place** + with fake-quantized values. + block_amax: Per-block amax values ``[out_features, n_amax_blocks]``. + global_scale: Pre-computed ``global_amax / (6.0 * 448.0)`` (scalar). + h_inv: Upper-triangular Cholesky factor of the damped inverse Hessian + ``[in_features, in_features]``. + block_size: Number of columns to process per GPTQ block. + quant_block_size: Number of elements sharing one quantization scale factor. + """ + from modelopt.torch.quantization.triton.gptq_fused_kernel import gptq_fused_block_scalar + + num_cols = weight.shape[1] + for bs in range(0, num_cols, block_size): + be = min(bs + block_size, num_cols) + qw, err = gptq_fused_block_scalar( + weight[:, bs:be].clone().contiguous(), + block_amax, + global_scale, + h_inv[bs:be, bs:be].contiguous(), + quant_block_size, + bs, + ) + weight[:, bs:be] = qw + if be < num_cols: + weight[:, be:].addmm_(err, h_inv[bs:be, be:], alpha=-1) + + _GPTQ_HELPER_REGISTRY: dict[str, type[GPTQHelper]] = {} diff --git a/tests/gpu/torch/quantization/conftest.py b/tests/gpu/torch/quantization/conftest.py index 95f2667bee6..9e34e5ef680 100644 --- a/tests/gpu/torch/quantization/conftest.py +++ b/tests/gpu/torch/quantization/conftest.py @@ -13,9 +13,18 @@ # See the License for the specific language governing permissions and # limitations under the License. + import pytest +from modelopt.torch.quantization import triton as triton_kernel + @pytest.fixture(autouse=True) def env_setup(monkeypatch): monkeypatch.setenv("NVIDIA_TF32_OVERRIDE", "0") + + +requires_triton = pytest.mark.skipif( + not triton_kernel.IS_AVAILABLE, + reason="Triton not available", +) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index 2d5f9d6d707..f04be0e4729 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -14,16 +14,24 @@ # limitations under the License. import copy +import time import pytest import torch from _test_utils.torch.transformers_models import get_tiny_llama, get_tiny_tokenizer +from conftest import requires_triton import modelopt.torch.quantization as mtq from modelopt.torch.export.unified_export_hf import _export_quantized_weight from modelopt.torch.quantization.model_calib import gptq from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor -from modelopt.torch.quantization.utils.calib_utils import update_hessian +from modelopt.torch.quantization.utils import promote_nvfp4_static_quantizers +from modelopt.torch.quantization.utils.calib_utils import ( + compute_hessian_inverse, + gptq_blockwise_update, + gptq_blockwise_update_fused_scalar, + update_hessian, +) from modelopt.torch.utils.dataset_utils import create_forward_loop, get_dataset_dataloader RAND_SEED = 42 @@ -231,3 +239,172 @@ def test_gptq_e2e_flow(quant_cfg): calibrate_loop = create_forward_loop(dataloader=calib_dataloader) model = mtq.quantize(model, quant_cfg, forward_loop=calibrate_loop) + + +# --------------------------------------------------------------------------- +# Fused Triton GPTQ kernel tests for NVFP4 scalar quantization +# --------------------------------------------------------------------------- + + +def _make_nvfp4_test_data(quant_block_size, out_features, dim): + """Create weight, weight_quantizer, block_amax, global_scale, and h_inv for NVFP4 GPTQ tests.""" + # Build a quantized Linear with NVFP4 static config at the desired block size + model = torch.nn.Linear(dim, out_features, bias=False, device="cuda") + weight = model.weight.data.clone() + + nvfp4_static_cfg = { + "num_bits": (2, 1), + "block_sizes": {-1: quant_block_size, "type": "static", "scale_bits": (4, 3)}, + } + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*weight_quantizer", "cfg": nvfp4_static_cfg}, + ], + "algorithm": "max", + } + inp = torch.randn(4, 32, dim, device="cuda") + mtq.quantize(model, quant_cfg, forward_loop=lambda m: m(inp)) + promote_nvfp4_static_quantizers(model) + + # Restore original weight (GPTQ operates on original weights) + model.weight.data = weight.clone() + + weight_quantizer = model.weight_quantizer + block_amax = weight_quantizer.amax.reshape(out_features, -1).float() + global_scale = weight_quantizer.global_amax.float().item() / (6.0 * 448.0) + + # Compute Hessian + hessian = torch.zeros(dim, dim, dtype=torch.float32) + hessian, _ = update_hessian(inp, hessian, 0) + hessian = hessian.to("cuda") + h_inv = compute_hessian_inverse(hessian, weight, perc_damp=0.01) + + return weight, weight_quantizer, block_amax, global_scale, h_inv + + +def _run_unfused_gptq_nvfp4(weight, weight_quantizer, h_inv, gptq_block_size): + """Unfused NVFP4 GPTQ using the production blockwise update with weight_quantizer.""" + w = weight.float().clone() + gptq_blockwise_update(w, h_inv, gptq_block_size, weight_quantizer) + return w + + +def _run_fused_gptq_nvfp4( + weight, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size +): + """Fused Triton GPTQ for NVFP4 using the production fused update.""" + w = weight.float().clone() + gptq_blockwise_update_fused_scalar( + w, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size + ) + return w + + +_NVFP4_QUANT_BLOCK_SIZES = [16, 128] +_NVFP4_GPTQ_BLOCK_SIZES = [16, 128] + + +@requires_triton +@pytest.mark.parametrize("quant_block_size", _NVFP4_QUANT_BLOCK_SIZES) +@pytest.mark.parametrize("gptq_block_size", _NVFP4_GPTQ_BLOCK_SIZES) +def test_fused_vs_unfused_nvfp4(quant_block_size, gptq_block_size): + """Fused Triton NVFP4 GPTQ must match unfused production reference.""" + torch.manual_seed(42) + dim = max(256, quant_block_size * 4) + out_features = 64 + + weight, weight_quantizer, block_amax, global_scale, h_inv = _make_nvfp4_test_data( + quant_block_size, + out_features, + dim, + ) + + weight_fused = _run_fused_gptq_nvfp4( + weight, + block_amax, + global_scale, + h_inv, + gptq_block_size, + quant_block_size, + ) + weight_unfused = _run_unfused_gptq_nvfp4( + weight, + weight_quantizer, + h_inv, + gptq_block_size, + ) + + assert not torch.equal(weight_fused, weight.float()), "Fused did not update weights" + assert not torch.equal(weight_unfused, weight.float()), "Unfused did not update weights" + + diff = (weight_fused - weight_unfused).abs() + max_abs = diff.max().item() + mean_abs = diff.mean().item() + denom = weight_unfused.abs().max().item() + rel_max = max_abs / denom if denom > 0 else 0.0 + + print( + f"\n[nvfp4] gptq_bs={gptq_block_size} quant_bs={quant_block_size}: " + f"max_abs={max_abs:.2e} mean_abs={mean_abs:.2e} rel_max={rel_max:.2e}" + ) + + torch.testing.assert_close(weight_fused, weight_unfused, atol=1e-4, rtol=1e-4) + + +_NVFP4_BENCH_CONFIGS = [ + (16, 128, 256, 512), + (16, 128, 256, 2048), + (16, 128, 256, 4096), + (128, 128, 256, 512), + (128, 128, 256, 2048), + (128, 128, 256, 4096), +] + + +def bench_fused_nvfp4(): + """Benchmark fused Triton NVFP4 GPTQ vs unfused production loop (informational-only). + + Not collected by pytest. Run directly: ``python tests/gpu/torch/quantization/test_gptq.py`` + """ + + def _bench(fn, n_warmup=2, n_iters=5): + for _ in range(n_warmup): + fn() + torch.cuda.synchronize() + total = 0.0 + for _ in range(n_iters): + torch.cuda.synchronize() + t0 = time.perf_counter() + fn() + torch.cuda.synchronize() + total += time.perf_counter() - t0 + return total / n_iters + + for quant_block_size, gptq_block_size, out_features, dim in _NVFP4_BENCH_CONFIGS: + torch.manual_seed(42) + weight, weight_quantizer, block_amax, global_scale, h_inv = _make_nvfp4_test_data( + quant_block_size, out_features, dim + ) + + def run_fused(): + return _run_fused_gptq_nvfp4( + weight, block_amax, global_scale, h_inv, gptq_block_size, quant_block_size + ) + + def run_unfused(): + return _run_unfused_gptq_nvfp4(weight, weight_quantizer, h_inv, gptq_block_size) + + t_fused = _bench(run_fused) + t_unfused = _bench(run_unfused) + speedup = t_unfused / t_fused if t_fused > 0 else float("inf") + + tag = f"qbs{quant_block_size}_gbs{gptq_block_size}_{out_features}x{dim}" + print( + f"[{tag}] Fused: {t_fused * 1e3:8.2f} ms | " + f"Unfused: {t_unfused * 1e3:8.2f} ms | Speedup: {speedup:.1f}x" + ) + + +if __name__ == "__main__": + bench_fused_nvfp4() From c51c1762b3e124b9eecc5b7e09d4e7374d734a03 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 22 Apr 2026 01:49:08 +0530 Subject: [PATCH 29/30] fix: prevent gh-pages repo bloat from doc preview artifacts (#1309) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### What does this PR do? Type of change: Bug fix Fixes gh-pages branch bloat that grew from ~26 MB to ~441 MB in four weeks (nvbug 6099503). Three compounding causes were identified and addressed: 1. **Sphinx `.doctrees/` cache published to gh-pages** — `sphinx-build` was writing its build cache inside `build/html/` which was then uploaded verbatim. Accounts for ~3.3 GB uncompressed across history. 2. **`JamesIves/github-pages-deploy-action` appending a commit on every push** — main-site files accumulated forever with `single-commit: false` (default). 3. **PR preview deploying on every `synchronize` event for all PRs** — `rossjrw/pr-preview-action` re-deployed the full site for every push to any PR regardless of whether docs changed (e.g. PR #1128 triggered 64 preview deploys × ~11 MB each). Changes: - Pass `-d /tmp/doctrees` to `sphinx-build` so `.doctrees/` is never written into `build/html/` - Add `paths: [docs/**, modelopt/**]` filter to `pull_request` trigger so the docs workflow only runs on PRs that touch docs or source code - Set `single-commit: true` on the deploy action so main-site pushes squash into one commit - Deduplicate docs build: `deploy-preview` now downloads the artifact from `build-docs` instead of running a second `sphinx-build` - Set `retention-days: 1` on the artifact since it is only needed for the duration of the workflow run The one-time cleanup (force-push squashed orphan to gh-pages) was already applied separately — repo is now ~59 MB for a full clone vs ~441 MB before. ### Usage N/A — CI/workflow change only. ### Testing - Workflow logic reviewed manually. - The one-time cleanup was verified: `git rev-list --objects --disk-usage origin/gh-pages` now reports ~28 MB; full clone is ~59 MB. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: N/A - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: N/A ### Additional Information nvbug 6099503 ## Summary by CodeRabbit * **Chores** * Optimized documentation build and deployment workflow in CI/CD pipeline. * Improved pull request documentation preview handling with faster build timeouts and refined artifact management. * Enhanced GitHub Pages deployment configuration for better consistency. Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .github/workflows/pages.yml | 37 ++++++++++++++++++++++++++++++------- noxfile.py | 2 ++ 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/.github/workflows/pages.yml b/.github/workflows/pages.yml index 43e2b5cc78d..1e2ddc75ab9 100644 --- a/.github/workflows/pages.yml +++ b/.github/workflows/pages.yml @@ -23,33 +23,55 @@ permissions: jobs: build-docs: + if: github.event.action != 'closed' runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 10 steps: - uses: actions/checkout@v6 - uses: ./.github/actions/ubuntu-setup - name: Build docs run: pip install nox uv && nox -s docs - name: Upload docs artifact - if: github.event_name == 'push' || github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' uses: actions/upload-artifact@v4 with: name: docs-html path: docs/build/html + retention-days: 1 + + changes: + if: github.event_name == 'pull_request' && github.event.action != 'closed' + runs-on: ubuntu-latest + outputs: + docs: ${{ steps.filter.outputs.docs }} + steps: + - uses: dorny/paths-filter@v3 + id: filter + with: + filters: | + docs: + - 'docs/**' + - 'modelopt/**' + - '.github/workflows/pages.yml' deploy-preview: - if: github.event_name == 'pull_request' + if: | + always() && + github.event_name == 'pull_request' && + (github.event.action == 'closed' || needs.changes.outputs.docs == 'true') + needs: [build-docs, changes] runs-on: ubuntu-latest - timeout-minutes: 30 + timeout-minutes: 10 # Per-PR concurrency without cancel-in-progress so 'closed' cleanup always runs concurrency: group: pr-preview-${{ github.event.pull_request.number }} steps: - uses: actions/checkout@v6 - - uses: ./.github/actions/ubuntu-setup - - name: Build docs + - name: Download docs artifact if: github.event.action != 'closed' - run: pip install nox uv && nox -s docs + uses: actions/download-artifact@v4 + with: + name: docs-html + path: docs/build/html - name: Deploy / remove PR preview uses: rossjrw/pr-preview-action@v1 with: @@ -70,5 +92,6 @@ jobs: uses: JamesIves/github-pages-deploy-action@v4 with: folder: docs/build/html + single-commit: true # Preserve PR preview subdirectories deployed by the deploy-preview job clean-exclude: pr-preview diff --git a/noxfile.py b/noxfile.py index fcef3d30875..96db23e1eee 100644 --- a/noxfile.py +++ b/noxfile.py @@ -164,6 +164,8 @@ def docs(session): with session.chdir("docs"): session.run( "sphinx-build", + "-d", + "/tmp/doctrees", "source", "build/html", "--fail-on-warning", From 785d3a2df68e0becbd5ff03d12b48240cde85584 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 22 Apr 2026 09:14:12 +0530 Subject: [PATCH 30/30] [CI] Bump test containers to latest (#1299) - Use latest containers for testing in CICD ## Summary by CodeRabbit * **Chores** * Bumped TensorRT-LLM Docker images to 1.3.0rc12 in example and GPU test workflows. * Updated PyTorch container image from 26.01 to 26.03 for GPU tests. * Captured uv lock upgrade output to a temp file, inlined it into PR bodies, and adjusted workflow heredoc/templating and step behavior. * **Documentation** * Clarified an inline comment and simplified a warning message for an ONNX quantization extension. --------- Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- .github/workflows/bump_uv_lock.yml | 19 ++++++++++++++++--- .github/workflows/example_tests.yml | 4 ++-- .github/workflows/gpu_tests.yml | 5 ++--- modelopt/onnx/quantization/extensions.py | 4 +--- 4 files changed, 21 insertions(+), 11 deletions(-) diff --git a/.github/workflows/bump_uv_lock.yml b/.github/workflows/bump_uv_lock.yml index e0933418b6e..47360542000 100644 --- a/.github/workflows/bump_uv_lock.yml +++ b/.github/workflows/bump_uv_lock.yml @@ -30,7 +30,10 @@ jobs: t['tool']['uv']['override-dependencies'] = ['torch; sys_platform == \"never\"'] toml.dump(t, open('pyproject.toml', 'w')) " - - run: uv lock --upgrade + - name: Run uv lock upgrade + run: | + set -o pipefail + uv lock --upgrade 2>&1 | tee /tmp/uv_lock_output.txt - name: Check for changes id: changes run: | @@ -53,12 +56,22 @@ jobs: git config user.email "github-actions[bot]@users.noreply.github.com" git commit -s -m "[chore]: bump uv.lock" git push origin "$BRANCH" + UV_OUTPUT=$(cat /tmp/uv_lock_output.txt) gh pr create \ --title "[chore]: weekly bump of uv.lock on ${BASE} ($(date +%Y-%m-%d))" \ - --body "$(cat <<'EOF' + --body "$(cat < + uv lock --upgrade output + + \`\`\` + ${UV_OUTPUT} + \`\`\` + + EOF )" \ --base "$BASE" diff --git a/.github/workflows/example_tests.yml b/.github/workflows/example_tests.yml index e8d307ccef9..2e6bfa690eb 100644 --- a/.github/workflows/example_tests.yml +++ b/.github/workflows/example_tests.yml @@ -60,7 +60,7 @@ jobs: uses: ./.github/workflows/_example_tests_runner.yml secrets: inherit with: - docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10" + docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc12" example: ${{ matrix.example }} pip_install_extras: "[hf,dev-test]" runner: linux-amd64-gpu-rtxpro6000-latest-1 @@ -74,7 +74,7 @@ jobs: uses: ./.github/workflows/_example_tests_runner.yml secrets: inherit with: - docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10" + docker_image: "nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc12" example: ${{ matrix.example }} pip_install_extras: "[hf,dev-test]" runner: linux-amd64-gpu-rtxpro6000-latest-2 diff --git a/.github/workflows/gpu_tests.yml b/.github/workflows/gpu_tests.yml index 628aead7ee8..f786d7f33d8 100644 --- a/.github/workflows/gpu_tests.yml +++ b/.github/workflows/gpu_tests.yml @@ -40,14 +40,13 @@ jobs: include: - example: gpu timeout: 60 - container_image: pytorch:26.01-py3 - # tests/gpu/_extensions/test_onnx_extensions.py fails for newer containers until https://github.com/tbenthompson/cppimport/pull/98 + container_image: pytorch:26.03-py3 - example: gpu_megatron timeout: 45 container_image: nemo:26.04 - example: gpu_trtllm timeout: 30 - container_image: tensorrt-llm/release:1.3.0rc10 + container_image: tensorrt-llm/release:1.3.0rc12 runs-on: ${{ startsWith(github.ref, 'refs/heads/pull-request/') && 'linux-amd64-gpu-rtxpro6000-latest-1' || 'linux-amd64-gpu-rtxpro6000-latest-2' }} timeout-minutes: ${{ matrix.timeout }} container: diff --git a/modelopt/onnx/quantization/extensions.py b/modelopt/onnx/quantization/extensions.py index 68facdaac85..ba2c0e137dd 100644 --- a/modelopt/onnx/quantization/extensions.py +++ b/modelopt/onnx/quantization/extensions.py @@ -18,7 +18,7 @@ import os import sys -# TODO: cppimport is no longer maintained, switch to a different library +# TODO: cppimport is not actively maintained, consider a better alternative library import cppimport from modelopt.onnx.logging_config import logger @@ -32,7 +32,5 @@ except Exception as e: logger.warning( f"{e}\nUnable to load `modelopt_round_and_pack_ext', falling back to python based optimized version. " - "If you see `copy_file() got an unexpected keyword argument 'dry_run'`, you will need " - "https://github.com/tbenthompson/cppimport/pull/98 or downgrade setuptools until we have a workaround" ) round_and_pack_ext = None