From a228ceac2a82326146d986fb3d3eb32dbf453233 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 13 Aug 2026 01:17:19 -0700 Subject: [PATCH 1/5] fix(export): find decoder layers through arbitrarily nested wrappers get_homogeneous_hf_decoder_layers unwrapped .model then .language_model once each, so it only reached layers exactly two wrappers deep and in that order. Kimi-K3 keeps its decoder at language_model.model.layers, so the walk stopped on the intermediate wrapper, returned None, and the architecture was reported unsupported -- which takes layerwise calibration and per-layer export out of reach for the model that needs them most. Unwrap iteratively instead, bounded so a cycle cannot hang. Also re-apply --attn_implementation after model construction. Kimi-K3's remote code overwrites _attn_implementation to flash_attention_2 unconditionally in __init__, ignoring the flag; export runs a trace forward, so an unavailable backend fails there rather than at load. Sub-configs are walked because remote code typically rewrites the nested text_config, and layer modules hold a reference to the same object. Only applied when the caller passed the flag explicitly, so nothing changes by default. Both were found in the previous Kimi-K3 run (PR #2008 workflow) but were never upstreamed; the branch carrying them was deleted after that PR merged. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 9f86b2857cfdbf13aed2c1f75c2d97c46d0395fb) --- examples/hf_ptq/example_utils.py | 42 +++++++++++++++++++ .../torch/quantization/plugins/huggingface.py | 25 ++++++++--- 2 files changed, 61 insertions(+), 6 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index 862bc0f7fe5..ded65ee14e5 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -703,6 +703,42 @@ def _resolve_init_config(hf_config, auto_model_module, ckpt_path, config_kwargs) return hf_config +def _force_attn_implementation(model, attn_implementation: str) -> None: + """Set ``_attn_implementation`` on the model config and every nested sub-config. + + Some remote modeling code overrides the requested attention implementation inside + ``__init__`` -- Kimi-K3 rewrites it to ``flash_attention_2`` unconditionally, ignoring + ``--attn_implementation``. Export runs a trace forward, so a backend whose compiled + extension is unavailable in this environment fails there rather than at load. + + Sub-configs are walked because multimodal models keep separate ones per tower, and + remote code typically rewrites the *nested* config (Kimi-K3 rewrites ``text_config`` + from ``KimiLinearModel.__init__``). Layer modules hold a reference to the same config + object, so overriding it here takes effect on the next forward. Only applied when the + caller asked for an implementation explicitly. + """ + pending, seen, changed = [model.config], set(), [] + while pending: + cfg = pending.pop() + if cfg is None or id(cfg) in seen: + continue + seen.add(id(cfg)) + current = getattr(cfg, "_attn_implementation", None) + if current is not None and current != attn_implementation: + try: + cfg._attn_implementation = attn_implementation + changed.append(f"{type(cfg).__name__}: {current} -> {attn_implementation}") + except Exception as e: # pragma: no cover - depends on remote config class + warnings.warn(f"Could not apply attn_implementation on {type(cfg).__name__}: {e}") + for sub in ("text_config", "vision_config", "audio_config", "decoder", "encoder"): + pending.append(getattr(cfg, sub, None)) + + if changed: + print("Re-applied the requested attention implementation after model init:") + for line in changed: + print(f" {line}") + + def _get_config_dtype(config): config_dtype = ( getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.bfloat16 @@ -967,6 +1003,12 @@ def has_pack_quantized_config(config): **model_kwargs2, ) model.eval() + + # Honour the caller's explicit choice even when remote modeling code overwrote it + # during __init__ (see _force_attn_implementation). + if attn_implementation is not None: + _force_attn_implementation(model, attn_implementation) + if has_pack_quantized_config(hf_config): _unpack_compressed_linear_weights(model, ckpt_path) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 4acb4d30dfa..d02a514c213 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1796,17 +1796,30 @@ def is_homogeneous_hf_model(model: nn.Module) -> bool: return len(layer_classes) == 1 +#: How deep to unwrap before giving up. Each level is one of the wrappers below, and no +#: released architecture nests more than a handful; the bound only stops a cycle. +_MAX_DECODER_UNWRAP_DEPTH = 8 + + def get_homogeneous_hf_decoder_layers(model: nn.Module) -> nn.ModuleList | None: if not _is_supported_hf_model(model): return None + # Unwrap iteratively rather than testing each wrapper once: multimodal models nest them + # in either order and to arbitrary depth. Kimi-K3 keeps its layers at + # ``language_model.model.layers``, so a single ``.model`` then ``.language_model`` walk + # stops on the intermediate wrapper and reports the architecture unsupported. decoder = model - if hasattr(decoder, "model"): - decoder = decoder.model - if hasattr(decoder, "language_model"): - decoder = decoder.language_model - if hasattr(decoder, "layers"): - return decoder.layers + for _ in range(_MAX_DECODER_UNWRAP_DEPTH): + if hasattr(decoder, "layers"): + return decoder.layers + for attr in ("model", "language_model"): + inner = getattr(decoder, attr, None) + if isinstance(inner, nn.Module): + decoder = inner + break + else: + return None return None From 74b0e2a95b6dca054af7dea216c96468a994147c Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 13 Aug 2026 01:27:18 -0700 Subject: [PATCH 2/5] feat(recipes): experts-only NVFP4 layerwise fused export for offloaded models Pairs layerwise.export_dir with checkpoint_dir for a PTQ run too large to hold resident, so an interrupted run resumes without recalibrating or re-exporting finished layers. That combination is the point for a run that outlasts its GPU session -- a kill loses at most the in-flight layer. Scopes experts as '*.experts.*' rather than '*block_sparse_moe*'. On fine-grained MoE the broader glob also matches shared_experts.*, routed_expert_{up,down}_proj and routed_expert_norm; on Kimi-K3 that is 552 additional modules the vendor left unquantized, one of which is an RMSNorm. shared_experts is missed by the narrow glob because its path contains '_experts.' rather than '.experts.', per fnmatch semantics. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 02ed6b545ad9f96130e85a4a6296ead883022c2c) --- ..._only-kv_fp8_layerwise_export_offload.yaml | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) create mode 100644 modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml new file mode 100644 index 00000000000..2be18d6522c --- /dev/null +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml @@ -0,0 +1,62 @@ +# 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. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: > + NVFP4 W4A4 on routed experts only, FP8 KV cache, max layerwise calibration, with each + decoder layer exported to its own shard as soon as it is calibrated. Same intent as + nvfp4_experts_only-kv_fp8_layerwise_export, but scoped and paired for a model too large + to hold resident: use with --offload_folder and per-device memory budgets. + + Expert scoping is '*.experts.*' rather than '*block_sparse_moe*'. On fine-grained MoE + models the broader glob also matches shared_experts.*, routed_expert_{up,down}_proj and + routed_expert_norm -- on Kimi-K3 that is 552 extra modules the vendor deliberately left + unquantized, one of which is an RMSNorm. '*.experts.*' matches only the routed expert + projections; shared_experts is missed because its path contains '_experts.' rather than + '.experts.' (fnmatch semantics, conversion.py). + + Pairs checkpoint_dir with export_dir so an interrupted run resumes without recalibrating + or re-exporting finished layers. That is the point of the combination for a run that + outlasts its GPU session: a killed run loses at most the in-flight layer. + + A resumed run never recalibrates the layers it skipped, so the exported checkpoint is + complete but the in-memory model is not and must not be used for inference. +quantize: + algorithm: + method: max + layerwise: + enable: true + # max only updates _amax, so the exported shard stays valid for its layer. + calib_mutates_weights: false + checkpoint_dir: /tmp/modelopt_layerwise_ckpt + # Presence enables per-layer export; the value is replaced with --export_path. + export_dir: /tmp/modelopt_layerwise_export + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*.experts.*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*.experts.*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers From c075cab85de0bf12f6473f93c3a9abd9fddeaa9b Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Thu, 13 Aug 2026 14:13:02 -0700 Subject: [PATCH 3/5] fix(hf_ptq): pin offloaded weights that are read outside their own forward accelerate materializes an offloaded weight in that module's pre-forward hook and returns it to meta in the matching post-forward. A weight read from a *sibling's* forward is therefore on meta at the moment it is used. Kimi-K3 does exactly this: _apply_attn_res computes norm.weight.float() * proj.weight.squeeze(0).float() from the decoder layer's forward, reaching into six children (three call sites in modeling_kimi_linear.py). The result is RuntimeError: Tensor on device meta is not on the expected device cuda:0! which is why a disk-offloaded K3 could not run a forward at all. Setting the tensor resident is not enough on its own -- post_forward walks the module's tensors and pushes every one back to meta, so the hook has to go. Detaching alone is not enough either: AlignDevicesHook.detach_hook restores each tensor to original_devices[name] and skips meta, which is precisely what a disk-offloaded param has, so it would be left on meta. Retargeting original_devices at the execution device first makes detach materialize the values itself, through accelerate's own code path rather than a hand-rolled copy. Safe because these modules' forward is never called -- only their raw .weight is read -- so removing the hook removes nothing that was doing work. The tensors are one row each, (1, hidden) and (hidden,), so pinning all six on a 93-layer model costs single-digit MB. Verified on a disk-offloaded tiny Kimi-K3: 10 externally-read params on meta before, 0 after, and the forward completes with finite logits where it previously raised. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 4c50e132ab1995d09f59e45d0ebe768facaa4911) --- examples/hf_ptq/example_utils.py | 67 +++++++++++++++++++++++++++++++- 1 file changed, 65 insertions(+), 2 deletions(-) diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index ded65ee14e5..29205134846 100755 --- a/examples/hf_ptq/example_utils.py +++ b/examples/hf_ptq/example_utils.py @@ -730,8 +730,10 @@ def _force_attn_implementation(model, attn_implementation: str) -> None: changed.append(f"{type(cfg).__name__}: {current} -> {attn_implementation}") except Exception as e: # pragma: no cover - depends on remote config class warnings.warn(f"Could not apply attn_implementation on {type(cfg).__name__}: {e}") - for sub in ("text_config", "vision_config", "audio_config", "decoder", "encoder"): - pending.append(getattr(cfg, sub, None)) + pending.extend( + getattr(cfg, sub, None) + for sub in ("text_config", "vision_config", "audio_config", "decoder", "encoder") + ) if changed: print("Re-applied the requested attention implementation after model init:") @@ -739,6 +741,61 @@ def _force_attn_implementation(model, attn_implementation: str) -> None: print(f" {line}") +#: Modules whose weights are read from *another* module's forward, so accelerate never +#: materializes them. Kimi-K3's ``_apply_attn_res`` does +#: ``norm.weight.float() * proj.weight.squeeze(0).float()`` from the decoder layer's +#: forward, reaching into these six children (``modeling_kimi_linear.py``, three call +#: sites). Each is one row -- ``(1, hidden)`` and ``(hidden,)`` -- so pinning all of them +#: on a 93-layer model costs single-digit MB. +_EXTERNALLY_READ_PARAM_SUFFIXES = ( + "self_attention_res_proj", + "self_attention_res_norm", + "mlp_res_proj", + "mlp_res_norm", + "output_attn_res_proj", + "output_attn_res_norm", +) + + +def _pin_externally_read_params( + model, suffixes: tuple[str, ...] = _EXTERNALLY_READ_PARAM_SUFFIXES +) -> int: + """Make offloaded weights that are read outside their own forward permanently resident. + + accelerate materializes an offloaded weight in *that module's* pre-forward hook and + returns it to meta in the matching post-forward. A weight read from a sibling's forward + is therefore on meta at the moment it is used, which surfaces as + ``Tensor on device meta is not on the expected device cuda:0``. + + Setting the tensor is not enough on its own: ``post_forward`` walks the module's tensors + and pushes every one back to meta, so the hook has to go. Detaching alone is not enough + either -- ``AlignDevicesHook.detach_hook`` restores each tensor to + ``original_devices[name]`` and *skips* meta, which is precisely what a disk-offloaded + param has, so it would be left on meta. Retargeting ``original_devices`` at the + execution device first makes detach do the materialization itself, using accelerate's + own code path rather than a hand-rolled copy. + + Safe because these modules' ``forward`` is never called -- only their raw ``.weight`` is + read -- so removing the hook removes nothing that was doing work. + + Returns the number of modules pinned. + """ + from accelerate.hooks import remove_hook_from_module + + pinned = 0 + for name, module in model.named_modules(): + if not name.endswith(suffixes): + continue + hook = getattr(module, "_hf_hook", None) + if hook is None or not getattr(hook, "offload", False): + continue + device = hook.execution_device + hook.original_devices = dict.fromkeys(getattr(hook, "original_devices", {}), device) + remove_hook_from_module(module) + pinned += 1 + return pinned + + def _get_config_dtype(config): config_dtype = ( getattr(config, "dtype", None) or getattr(config, "torch_dtype", None) or torch.bfloat16 @@ -1009,6 +1066,12 @@ def has_pack_quantized_config(config): if attn_implementation is not None: _force_attn_implementation(model, attn_implementation) + # Offloaded weights that a sibling's forward reads would otherwise be on meta when used. + if _disk_offload: + n_pinned = _pin_externally_read_params(model) + if n_pinned: + print(f"Pinned {n_pinned} externally-read modules so offload cannot meta them.") + if has_pack_quantized_config(hf_config): _unpack_compressed_linear_weights(model, ckpt_path) From 71e6f94f0fc2d1e2d716c009accf1d761942bf18 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:09:19 +0000 Subject: [PATCH 4/5] docs(recipes): document the offloaded layerwise export recipe The recipe shipped undocumented; the parent branch's docs test now requires a row and the count kept in step. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- modelopt_recipes/ptq.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 1735289ac3f..3f85166158e 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -29,7 +29,7 @@ supported combinations. ### The shipped recipes
-All 25 general/ptq/ recipes (click to expand) +All 26 general/ptq/ recipes (click to expand) | Recipe | Model body | KV cache | Calibration | |--------|-----------|----------|-------------| @@ -49,6 +49,7 @@ supported combinations. | `nvfp4_experts_only-kv_fp8_layerwise` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise | | `nvfp4_experts_only-kv_fp8_layerwise_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (non-mutating, for disk offload) | | `nvfp4_experts_only-kv_fp8_layerwise_export` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise (exports each layer as it is calibrated) | +| `nvfp4_experts_only-kv_fp8_layerwise_export_offload` | NVFP4 W4A4, MoE experts only | FP8 (calibrated) | max, layerwise, tuned for accelerate-offloaded models | | `nvfp4_experts_only_mse-kv_fp8_cast` | NVFP4 W4A4, MoE experts only | FP8 (constant amax) | MSE + FP8 sweep | | `nvfp4_experts_only_input_scale1-kv_fp8_cast` | NVFP4 W4A4, MoE experts only, expert `input_scale` pinned to 1.0 | FP8 (constant amax) | max (weights); expert activations uncalibrated | | `nvfp4_omlp_only-kv_fp8` | NVFP4 W4A4, o_proj + MLP/MoE | FP8 (calibrated) | max | From 33e828b36c01d7ba47cdf955519c1b38674b6f58 Mon Sep 17 00:00:00 2001 From: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> Date: Fri, 21 Aug 2026 23:32:59 +0000 Subject: [PATCH 5/5] fix(recipes): stop the offload recipe pinning resume state to /tmp The recipe shipped checkpoint_dir: /tmp/modelopt_layerwise_ckpt, which is container-local -- a run that outlasts its GPU session came back to a wiped manifest and restarted at layer 0, the exact failure this recipe exists to avoid. Leaving it unset lets hf_ptq derive .layerwise_resume, which lives next to the shards it describes. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> --- ...fp4_experts_only-kv_fp8_layerwise_export_offload.yaml | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml index 2be18d6522c..b041123cfa7 100644 --- a/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml +++ b/modelopt_recipes/general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload.yaml @@ -34,9 +34,11 @@ metadata: projections; shared_experts is missed because its path contains '_experts.' rather than '.experts.' (fnmatch semantics, conversion.py). - Pairs checkpoint_dir with export_dir so an interrupted run resumes without recalibrating - or re-exporting finished layers. That is the point of the combination for a run that - outlasts its GPU session: a killed run loses at most the in-flight layer. + An interrupted run resumes without recalibrating or re-exporting finished layers, losing + at most the in-flight one -- the point of the combination for a run that outlasts its GPU + session. The resume state lives beside the checkpoint at .layerwise_resume + unless you set layerwise.checkpoint_dir yourself; do not point it at container-local + storage, or a run that survives its session comes back to a wiped manifest. A resumed run never recalibrates the layers it skipped, so the exported checkpoint is complete but the in-memory model is not and must not be used for inference. @@ -47,7 +49,6 @@ quantize: enable: true # max only updates _amax, so the exported shard stays valid for its layer. calib_mutates_weights: false - checkpoint_dir: /tmp/modelopt_layerwise_ckpt # Presence enables per-layer export; the value is replaced with --export_path. export_dir: /tmp/modelopt_layerwise_export quant_cfg: