From 71cde8864c2e296e39a1183e7cc1fea9013a4c05 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Fri, 10 Apr 2026 17:24:15 +0000 Subject: [PATCH 1/8] support for folding weights megatron-vllm export Signed-off-by: Kinjal Patel --- examples/vllm_serve/fakequant_worker.py | 1 + examples/vllm_serve/vllm_reload_utils.py | 27 +++++++++++--- .../export/plugins/vllm_fakequant_megatron.py | 35 +++++++++++++++++-- 3 files changed, 57 insertions(+), 6 deletions(-) diff --git a/examples/vllm_serve/fakequant_worker.py b/examples/vllm_serve/fakequant_worker.py index ec2b1f4033e..6ac6cacc78b 100644 --- a/examples/vllm_serve/fakequant_worker.py +++ b/examples/vllm_serve/fakequant_worker.py @@ -48,6 +48,7 @@ def _fakequant_run_prolog_worker(self) -> None: trust_remote_code = os.environ.get("TRUST_REMOTE_CODE", "false").lower() == "true" + print("tokenizer path: ", self.model_runner.model_config.tokenizer, "\n\n\n") 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..478f4f0a5ed 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -438,18 +438,37 @@ def load_state_dict_from_path( 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") + # 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. Any other missing quantizer key is still an error. + missing_wq_module_paths: set[str] = set() 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") + if "weight_quantizer" in key: + parts = key.split(".") + for i, part in enumerate(parts): + if part.endswith("weight_quantizer"): + missing_wq_module_paths.add(".".join(parts[: i + 1])) + break + else: + raise ValueError( + f"Key {key} not found in checkpoint state dict, but exists in model" + ) + + for name, module in model.named_modules(): + if name in missing_wq_module_paths and hasattr(module, "disable"): + module.disable() checkpoint_quant_count = len(checkpoint_quant_keys) model_quant_count = len(model_quant_keys) - # Ensure counts match - if checkpoint_quant_count != model_quant_count: + # Ensure counts match (excluding weight quantizer keys, which may be absent when weights + # were pre-folded at export) + model_non_wq_quant_count = sum(1 for k in model_quant_keys if "weight_quantizer" not in k) + if checkpoint_quant_count != model_non_wq_quant_count: 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"quant keys but model has {model_non_wq_quant_count} non-weight quantizer state keys. " f"This can happen if the model is using PP." ) diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 9a41ae2baff..db392412c51 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,41 @@ 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(): + # Some quantization kernels (e.g. NVFP4 dynamic block quant) require + # CUDA for both the weight and any stored amax buffers. During Megatron + # export the whole model (including quantizer buffers) may be on CPU + # after TP gather. Snapshot buffer devices, move everything to CUDA for + # the forward pass, then restore. + quant_device = ( + torch.device("cuda") if torch.cuda.is_available() else weight.device + ) + buf_devices = [(buf, buf.device) for buf in weight_quantizer.buffers()] + for buf, _ in buf_devices: + buf.data = buf.data.to(quant_device) + try: + weight = weight_quantizer(weight.to(quant_device)).to(dtype) + finally: + for buf, orig_device in buf_devices: + buf.data = buf.data.to(orig_device) + 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 b53afa63b61ed54292ce07f8a066aa4175a087d8 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Mon, 13 Apr 2026 03:42:38 +0000 Subject: [PATCH 2/8] minor Signed-off-by: Kinjal Patel --- examples/vllm_serve/fakequant_worker.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/vllm_serve/fakequant_worker.py b/examples/vllm_serve/fakequant_worker.py index 6ac6cacc78b..0cbb6273042 100644 --- a/examples/vllm_serve/fakequant_worker.py +++ b/examples/vllm_serve/fakequant_worker.py @@ -48,7 +48,7 @@ def _fakequant_run_prolog_worker(self) -> None: trust_remote_code = os.environ.get("TRUST_REMOTE_CODE", "false").lower() == "true" - print("tokenizer path: ", self.model_runner.model_config.tokenizer, "\n\n\n") + tokenizer = AutoTokenizer.from_pretrained( self.model_runner.model_config.tokenizer, trust_remote_code=trust_remote_code, From 8c0af4629bbd199088a73c3bed051f50023b9fd0 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Mon, 13 Apr 2026 15:24:18 +0000 Subject: [PATCH 3/8] minor Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_reload_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 478f4f0a5ed..0b42fdb86ee 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -460,8 +460,7 @@ def load_state_dict_from_path( module.disable() checkpoint_quant_count = len(checkpoint_quant_keys) - model_quant_count = len(model_quant_keys) - + # Ensure counts match (excluding weight quantizer keys, which may be absent when weights # were pre-folded at export) model_non_wq_quant_count = sum(1 for k in model_quant_keys if "weight_quantizer" not in k) From b147d49fe3637a5fcb478c49398528d417b0f59c Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 14 Apr 2026 22:17:34 +0000 Subject: [PATCH 4/8] minor Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_reload_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 0b42fdb86ee..8ae72a452c6 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -460,7 +460,7 @@ def load_state_dict_from_path( module.disable() checkpoint_quant_count = len(checkpoint_quant_keys) - + # Ensure counts match (excluding weight quantizer keys, which may be absent when weights # were pre-folded at export) model_non_wq_quant_count = sum(1 for k in model_quant_keys if "weight_quantizer" not in k) From 6016cfe67d0fdbd34c9d27d96827bfccd32bea64 Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 14 Apr 2026 22:53:47 +0000 Subject: [PATCH 5/8] minor Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_reload_utils.py | 25 ++++++------------- .../export/plugins/vllm_fakequant_megatron.py | 12 ++++----- 2 files changed, 14 insertions(+), 23 deletions(-) diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 8ae72a452c6..301c75712b7 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -14,7 +14,6 @@ # limitations under the License. import re -import warnings from collections import defaultdict from collections.abc import Callable from typing import Any @@ -445,11 +444,15 @@ def load_state_dict_from_path( for key in model_quant_keys: if key not in checkpoint_quant_keys: if "weight_quantizer" in key: + # State dict keys continue past the submodule (e.g. ...weight_quantizer._amax). + # named_modules() names stop at the weight_quantizer module; strip the suffix. parts = key.split(".") - for i, part in enumerate(parts): - if part.endswith("weight_quantizer"): - missing_wq_module_paths.add(".".join(parts[: i + 1])) - break + wq_i = next( + (i for i, p in enumerate(parts) if p.endswith("weight_quantizer")), + None, + ) + if wq_i is not None: + missing_wq_module_paths.add(".".join(parts[: wq_i + 1])) else: raise ValueError( f"Key {key} not found in checkpoint state dict, but exists in model" @@ -459,18 +462,6 @@ def load_state_dict_from_path( if name in missing_wq_module_paths and hasattr(module, "disable"): module.disable() - checkpoint_quant_count = len(checkpoint_quant_keys) - - # Ensure counts match (excluding weight quantizer keys, which may be absent when weights - # were pre-folded at export) - model_non_wq_quant_count = sum(1 for k in model_quant_keys if "weight_quantizer" not in k) - if checkpoint_quant_count != model_non_wq_quant_count: - warnings.warn( - f"Mismatch in quantizer state key counts: checkpoint has {checkpoint_quant_count} " - f"quant keys but model has {model_non_wq_quant_count} non-weight quantizer state keys. " - f"This can happen if the model is using PP." - ) - # 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 db392412c51..33f673a3f12 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -144,13 +144,13 @@ def _get_quantized_state( weight_quantizer = getattr(module, "weight_quantizer", None) if weight_quantizer is not None and weight_quantizer.is_enabled: with torch.no_grad(): - # Some quantization kernels (e.g. NVFP4 dynamic block quant) require - # CUDA for both the weight and any stored amax buffers. During Megatron - # export the whole model (including quantizer buffers) may be on CPU - # after TP gather. Snapshot buffer devices, move everything to CUDA for - # the forward pass, then restore. + # Some quantizers (e.g. NVFP4) require CUDA. If the model landed on + # CPU after TP gather, lift to the current CUDA device for the forward + # pass, then restore buffer devices. quant_device = ( - torch.device("cuda") if torch.cuda.is_available() else weight.device + torch.device("cuda", torch.cuda.current_device()) + if weight.device.type == "cpu" and torch.cuda.is_available() + else weight.device ) buf_devices = [(buf, buf.device) for buf in weight_quantizer.buffers()] for buf, _ in buf_devices: From 0ca8f87498eef8a3809f6701d4ed8cd691c0cf7d Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Tue, 14 Apr 2026 23:21:13 +0000 Subject: [PATCH 6/8] minor Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_reload_utils.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index 301c75712b7..b4dca372418 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -453,6 +453,11 @@ def load_state_dict_from_path( ) if wq_i is not None: missing_wq_module_paths.add(".".join(parts[: wq_i + 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." + ) else: raise ValueError( f"Key {key} not found in checkpoint state dict, but exists in model" From 3e2ec17fed0bf1bde6ad798762478a7296f7bebd Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Thu, 16 Apr 2026 00:33:47 +0000 Subject: [PATCH 7/8] addressed comments Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_reload_utils.py | 85 ++++++++++++++----- .../export/plugins/vllm_fakequant_megatron.py | 23 +++-- 2 files changed, 78 insertions(+), 30 deletions(-) diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index b4dca372418..e417461f49f 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging import re from collections import defaultdict from collections.abc import Callable @@ -30,8 +31,34 @@ convert_to_quantized_model, restore_quantizer_state, ) +from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import is_quantized +logger = logging.getLogger(__name__) + + +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: + logger.warning( + "Could not all_gather quantizer key lists across ranks (%s); using this rank's keys only.", + e, + ) + return local + def _values_equal(v1: Any, v2: Any) -> bool: """Compare values, handling dicts with tensors.""" @@ -434,37 +461,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. Any other missing quantizer key is still an error. + # 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: - if "weight_quantizer" in key: - # State dict keys continue past the submodule (e.g. ...weight_quantizer._amax). - # named_modules() names stop at the weight_quantizer module; strip the suffix. - parts = key.split(".") - wq_i = next( - (i for i, p in enumerate(parts) if p.endswith("weight_quantizer")), - None, - ) - if wq_i is not None: - missing_wq_module_paths.add(".".join(parts[: wq_i + 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." - ) + 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"Key {key} not found in checkpoint state dict, but exists in model" + 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 + logger.warning( + "%s quantizer key(s) missing from every rank's checkpoint (after all_gather): %s%s", + n, + sample, + f" ... (+{rest} more)" if rest > 0 else "", + ) for name, module in model.named_modules(): - if name in missing_wq_module_paths and hasattr(module, "disable"): + if ( + name in missing_wq_module_paths + and isinstance(module, TensorQuantizer) + and hasattr(module, "disable") + ): module.disable() # Update quant values diff --git a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py index 33f673a3f12..c8e45be3650 100644 --- a/modelopt/torch/export/plugins/vllm_fakequant_megatron.py +++ b/modelopt/torch/export/plugins/vllm_fakequant_megatron.py @@ -144,22 +144,29 @@ def _get_quantized_state( weight_quantizer = getattr(module, "weight_quantizer", None) if weight_quantizer is not None and weight_quantizer.is_enabled: with torch.no_grad(): - # Some quantizers (e.g. NVFP4) require CUDA. If the model landed on - # CPU after TP gather, lift to the current CUDA device for the forward - # pass, then restore buffer devices. + # 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 ) - buf_devices = [(buf, buf.device) for buf in weight_quantizer.buffers()] - for buf, _ in buf_devices: - buf.data = buf.data.to(quant_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: - for buf, orig_device in buf_devices: - buf.data = buf.data.to(orig_device) + if need_move: + weight_quantizer.to(wq_dev) name_to_value["weight"] = weight.cpu() else: return name_to_value, qformat, block_size From e454b79ba6c440c8a72ee64b3b323a8dd7445cff Mon Sep 17 00:00:00 2001 From: Kinjal Patel Date: Thu, 16 Apr 2026 01:02:35 +0000 Subject: [PATCH 8/8] minor Signed-off-by: Kinjal Patel --- examples/vllm_serve/vllm_reload_utils.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/examples/vllm_serve/vllm_reload_utils.py b/examples/vllm_serve/vllm_reload_utils.py index e417461f49f..2b59d1be2bd 100644 --- a/examples/vllm_serve/vllm_reload_utils.py +++ b/examples/vllm_serve/vllm_reload_utils.py @@ -13,8 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -import logging import re +import warnings from collections import defaultdict from collections.abc import Callable from typing import Any @@ -34,8 +34,6 @@ from modelopt.torch.quantization.nn import SequentialQuantizer, TensorQuantizer from modelopt.torch.quantization.utils import is_quantized -logger = logging.getLogger(__name__) - 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).""" @@ -53,9 +51,8 @@ def _union_quantizer_keys_across_ranks(local_quantizer_keys: list[str]) -> set[s out.update(g) return out except Exception as e: - logger.warning( - "Could not all_gather quantizer key lists across ranks (%s); using this rank's keys only.", - e, + warnings.warn( + f"Could not all_gather quantizer key lists across ranks ({e}); using this rank's keys only." ) return local @@ -311,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) @@ -493,11 +490,9 @@ def load_state_dict_from_path( keys = sorted(global_missing_non_wq) n = len(keys) sample, rest = keys[:8], n - 8 - logger.warning( - "%s quantizer key(s) missing from every rank's checkpoint (after all_gather): %s%s", - n, - sample, - f" ... (+{rest} more)" if rest > 0 else "", + warnings.warn( + 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():