Kimi-K3 on layerwise fused export: 3 defect fixes, multimodal support, single-B200 NVFP4 - #2218
Draft
Fridah-nv wants to merge 31 commits into
Draft
Kimi-K3 on layerwise fused export: 3 defect fixes, multimodal support, single-B200 NVFP4#2218Fridah-nv wants to merge 31 commits into
Fridah-nv wants to merge 31 commits into
Conversation
Contributor
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
_is_layerwise probed the recipe's algorithm with getattr(obj, "layerwise", None), but an algorithm loaded from YAML is a plain dict, where getattr always returns None. It therefore answered False for every layerwise recipe in the repo. The one thing it gates is whether --batch_size 0 skips auto batch-size probing, which its own comment says must be skipped because the probe "runs a full-model forward which defeats the point and can OOM on very large models". That protection has never engaged for the recipes it was written for. Replace it with an accessor that handles both shapes: dicts from YAML, and the config objects the deprecated --auto_quantize_* path still builds. Behaviour change: with --batch_size 0, layerwise recipes now use batch_size=1 instead of probing, which is what the existing comment intends. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…shes it A PTQ run that outlasts its GPU session loses every calibrated layer, and a completed one still pays for a second whole-model export pass over a full-precision intermediate checkpoint. Add layerwise.export_dir: each decoder layer's final quantized tensors are flushed to their own shard the moment calibration finishes with it, so the shards double as the resume artifact. Combined with layerwise.checkpoint_dir a restarted run skips layers already on disk instead of recalibrating them, and the per-layer weights.pt / quantizer_buffers.pt files are no longer written. When the last layer lands the directory is already a complete, loadable checkpoint, so export_hf_checkpoint() is skipped. Setting the field is the whole switch; hf_ptq.py rewrites it to --export_path, as it already does for checkpoint_dir. One layer per shard is what makes that work: a shard is only ever written whole, its name derives from the layer index, and the index is rebuilt at the end from the shards on disk. A crash can lose the layer in flight but never corrupt an earlier one, and a re-run overwrites in place rather than appending duplicates. Resident modules have no materialization window to discard export's damage, so transient_module_state snapshots the layer subtree's _parameters/_buffers/_modules and restores them, leaving calibration free to run every later layer through it. Tied-weight dedup is off here for the reason registry.py already turns it off for offload: data_ptr() cannot identify a tensor across an export that keeps rolling packed weights back. Weight-tied quantized modules are refused up front anyway. NVFP4 works because export_layer rediscovers the scale-fusion groups itself: the groups _fuse_shared_input_modules operates on -- q/k/v behind input_layernorm, gate/up behind post_attention_layernorm -- never cross a layer boundary, so a probe forward over one layer finds them. The probe uses that layer's real cached activations rather than the synthetic input the whole-model pass builds. Scope is resident, single-process models. AWQ and SVDQuant additionally need requantize_resmooth_fused_llm_layers' pre-quant-scale steps, which are still whole-model, so they stay refused along with accelerate offload, multi-process jobs, weight-tied quantized modules, multimodal models, MTP models and speculative decoding are refused with NotImplementedError before calibration starts; each would otherwise produce a silently different checkpoint rather than fail. A resumed run never recalibrates the layers it skipped, so the in-memory model is not valid for inference and layerwise.export_dir implies --skip_generate. Two pieces move out of this path to avoid duplicating what already exists: save_non_weight_artifacts() is extracted from the streaming exporter's tail, and FUSION_FREE_FORMATS moves to model_config.py, where _fuse_shared_input_modules had held the same set inline as a literal list. Verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only, FP8, 256 experts per layer): 93,273 tensors, 0 mismatched, and identical again after a simulated mid-model resume. NVFP4 equivalence is covered by a GPU test. The NVFP4 test leaves o_proj unquantized: layerwise calibration leaves its input amax at 0 on every layer but the last, which no export path can write. That is a pre-existing bug, unrelated to per-layer export. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Per-layer export exists so a run that outlasts its GPU session keeps its finished layers, and offloaded single-GPU runs on very large models are exactly the multi-hour runs that hit session limits -- but offload was the one case refused. Calibration already handled it; only export did not. finalize()'s tail walked model.state_dict() directly, and _collect drops meta tensors, so an offloaded model's embeddings, norms and lm_head were skipped with no error. Give them the same per-module materialization window the streaming exporter uses. Tail collection splits in two: modules needing a window, then everything already resident, which on a non-offloaded model is the whole tail. Tie detection cannot run under offload -- data_ptr() cannot group weights that are not resident -- so it now says so instead of reporting a clean bill of health. Resolving ties by name would fix it properly, the same way ExportContext.__post_init__ already has a TODO for. Also fix the fusion gate added with NVFP4 support: it asked get_quantization_format(layer), which returns the first format found, so a layer with FP8 attention and NVFP4 experts reported fp8 and silently skipped fusing its NVFP4 groups. Use the per-module scan the model-level gate already uses. The existing NVFP4 tests could not catch this -- both were single-format layers -- so this adds a mixed FP8/NVFP4 case, which fails without the fix on mlp.up_proj.weight_scale_2. Verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only, experts-only NVFP4, accelerate disk offload, max_memory 20GiB): 123,493 tensors, 0 mismatched, with embed_tokens, lm_head and norm correctly captured in the tail shard. Per-layer export is ~63s slower than the streaming export for this model (271s vs 208s). The cause is the per-layer fusion probe forward, not offload; the argument for this path is durability, not speed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Seven fixes from a review of the layerwise export path. Each would fail late or silently rather than at the point of the mistake. hf_ptq only retargets export_dir and runs the compatibility refusals on the mono-quantize path, so an AutoQuantize recipe carrying export_dir exported to the recipe's placeholder directory and then skipped export_hf_checkpoint(), leaving --export_path with no weights and printing success. Refuse that combination. set_layerwise_export_dir indexed algorithm as a dict, but detection accepts a list of algorithms too, so a list-shaped recipe died on a str index before calibration. Handle both shapes. The refusal loop listed --sparsity_fmt as the only route to the TRT-LLM exporter; int8_sq and encoder-decoder model_type reach it as well, and a Whisper model has discoverable decoder layers, so shards were written and then overwritten by a second checkpoint. resolve_checkpoint_dir hashed the config while it still held the recipe's placeholder export_dir, so two runs to different --export_path values shared one checkpoint dir and the second resumed against the wrong shards. Retarget first. _fusion_probe replayed a cached batch without the past_key_values reset _layer_forward_loop performs for the same tuples, so the probe could see kv_len at twice the attention mask width. save_file rejects two keys backed by one storage, and _collect's .cpu() is a no-op when the tensor is already there; copy aliases before writing, as _StreamingShardWriter already does. Finally, finalize()'s resident tail loop could still reach a module holding meta tensors, where packing raises deep inside the export handler. Raise there instead with the module name -- skipping it would drop weights silently, which is the failure this path exists to prevent. Re-verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only, experts-only NVFP4, accelerate disk offload): 123,493 tensors, 0 mismatched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…ise export LayerwiseExporter kept model, dtype and is_modelopt_qlora both on self and inside the ExportContext it builds from them, leaving two sources of truth for the same facts across 9 read sites. Build the context first in __init__ and read everything through it. No behaviour change. This is what docs/design/export-feature-scoping.md asks for in phase 1 -- construct the context once and thread it through, rather than rebuilding it or shadowing it per phase. Doing it now means that phase is already satisfied for this branch when the CheckpointExporter base class lands, instead of being work the refactor has to carry. Re-verified byte-identical to export_hf_checkpoint on Qwen3.6-35B-A3B (text-only, experts-only NVFP4, accelerate disk offload): 123,493 tensors, 0 mismatched. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Comments and docstrings only, no code change. State the reason and stop; drop restatements of what the code already says. Also corrects a stale claim in assert_layerwise_export_supported: it still said restricting to FUSION_FREE_FORMATS is what makes skipping requantize_resmooth_fused_llm_layers safe, which stopped being true when NVFP4 gained per-layer fusion. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Per-layer export presents its shards as the resume artifact, but the resume *point* does not come from them. _CheckpointState.from_folder returns `start = info[0] if info else 0`, read from checkpoint_dir's manifest, and nothing derives it from the shards on disk. So the manifest and the shards have to share a lifetime. They did not. Both shipped layerwise-export recipes default checkpoint_dir to /tmp/modelopt_layerwise_ckpt, and /tmp is container-local on a Slurm/Pyxis node. A run that outlasts its GPU session -- precisely the case this path exists for -- comes back to a wiped manifest, restarts calibration at layer 0, and overwrites every finished shard. assert_shards_present(0) passes trivially on the way through, so nothing warns: the run just silently redoes hours of work. Found on the Kimi-K3 run, where the manifest landed in /tmp/modelopt_layerwise_ckpt/Kimi-K3-bf16_4abe5702/ while the shards were on shared storage. At ~6 min per layer over 93 layers that would have cost a full session on every restart. Co-locate the checkpoint dir under --export_path when per-layer export is enabled, so the invariant is structural instead of something the user has to know. Only applied in that mode; otherwise checkpoint_dir is an ordinary resume directory and its placement is the caller's business. Verified on midi-K3: the manifest now resolves under <export_path>/.layerwise_checkpoint/<model>_<hash>/, and a run whose manifest is rewound to layer 3 reports "resuming layerwise calibration from layer 4/8" and skips the finished layers rather than recalculating them. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 9d0b846)
Name-based tied-weight dedup (#2194) removed tied_cache and moe_tied_cache from ExportContext, so constructing one with them now raises TypeError. Drop them. That change also supersedes this path's tie detection. Grouping quantized modules by weight.data_ptr() sees nothing once the weights are on meta, so under offload the check passed vacuously and only emitted a warning saying so. TiedWeightMap is keyed by name and survives offload, so group by that instead and delete the warning. data_ptr remains as a fallback for transformers <5.0, which publishes no map. The alias-key dedup in _collect deliberately stays as it is, matching the TODO(tied-map) note the streaming path carries: swapping that one needs offload-specific validation (meta tensors, per-tensor order, disk round-trip), and this path is per-tensor and offloaded too. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…is gone The shards are the resume artifact but not the resume point: restarting at layer K needs the cached activations feeding it and every skipped layer's output_meta, neither reconstructible from quantized weights. So the point comes from the checkpoint manifest, which makes one state dangerous -- shards on disk, manifest gone. start_layer is then 0, assert_shards_present checks an empty range and passes, and calibration recalculates and overwrites every finished layer without a word. A checkpoint_dir on ephemeral storage is the usual cause, and it is exactly the long run this feature exists for. Raise instead. A manifest that merely lags the shards stays allowed: that is a mid-window interrupt, and re-exporting those layers is idempotent. Adapted from 64ea3b8 on the stacked Kimi-K3 branch, without its multimodal machinery. Its guard also assumed checkpoint_dir is set, which raises TypeError in the export-without-resume mode; that case is now skipped, since with no checkpoint_dir there is no resume to lose and re-exporting is documented behaviour. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Two CI failures, both introduced here. build-docs: assert_layerwise_export_supported's docstring used a `.. todo::` directive. The convention was copied from _CheckpointState, which gets away with it by being private and never rendered; this function is in __all__, so autodoc renders it and sphinx.ext.todo is not enabled -- "Unknown directive type", fatal under --fail-on-warning. Plain prose instead. code-quality: colocate_layerwise_checkpoint_dir called _layerwise_checkpoint_dir_location, which does not exist on this base. The function came from the stacked Kimi-K3 branch, written when that helper returned a (shape, current) tuple; main has since replaced it with _layerwise_checkpoint_dir returning the directory. Rewritten against the current API, which also drops the dead legacy-flat branch. This was not only a type error: it would have raised NameError the first time a recipe set checkpoint_dir. Neither showed up locally because pre-commit was run over changed files while CI runs --all-files, and the docs were never built. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
transient_module_state restored which tensor each name points at, not tensor contents, and its docstring claimed export only ever rebinds. Scale fusion disproves that: preprocess_linear_fusion unifies a group's amax through _amax.data.copy_(), an in-place write the dict restore cannot undo. Measured on NVFP4, five of the model's amax buffers came back carrying the fused group maxima. Buffers are now clone-restored; parameters are not, since export replaces weights by rebinding and cloning them would copy the whole layer. The earlier "model unchanged" check used FP8, which never fuses, so it could not see this. detect_resume_point returns None once the manifest is complete, so start_layer fell back to 0 and a rerun recalibrated and overwrote every finished shard -- whether the previous run crashed between the last ckpt.save and finalize, or simply finished. In export mode a complete manifest now means finalize-only: the shards are already on disk and only the tail, index and configs can be missing. The layer-order guard becomes a raise. As an assert, python -O strips it and layer N's tensors land in layer M's shard under a well-formed index -- silent corruption, which is the failure mode the rest of this file refuses. set_layerwise_export_dir now retargets only entries that already declare export_dir, rather than switching per-layer export on for any entry that happens to carry a layerwise block. Document the new recipe in modelopt_recipes/ptq.md: its absence was failing test_every_general_ptq_recipe_is_documented and the recipe-count check. Its own metadata also still said resident-only, which offload support had made stale. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…he checkpoint Four more review findings. Nothing tied the exported shards to the run that produced them: the resume manifest records no model or quantization identity, and assert_shards_present only checks that files exist. Two runs pointed at one export directory with the same layer count could pair one run's manifest with the other's shards and finalize a structurally valid, wrong checkpoint. The exporter now writes .layerwise_export.json (model class, layer count, formats, kv-cache format) on first export and refuses when an existing one disagrees, naming the fields that differ. The CLI's config-hash checkpoint subdirectory made this hard to hit there, but direct library use had nothing. set_layerwise_export_dir returned the config untouched when it found nothing to retarget. Since the caller decides to skip the real export from a separately parsed recipe object, a mismatch would send shards to the recipe's placeholder path and leave --export_path empty on a run that reports success. It raises now. The co-located resume directory was a child of --export_path, so its cached activations shipped inside the deliverable; nothing deletes it on success, and deleting it would make re-running a finished command trip assert_no_orphan_shards, which cannot tell "already done" from "resume record lost". It is a sibling now, and named .layerwise_resume: with per-layer export on the shards are the weights, so that directory holds no checkpoint at all -- only the resume point, the next layer's cached activations, and skip-mode shape metadata. The changelog advertised a --layerwise_export flag removed earlier in this PR in favour of the recipe field, so it pointed users at an argparse error. Rewritten, and cut to two sentences per CONTRIBUTING.md. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
The shipped recipe was the only one in the repo to set layerwise.checkpoint_dir, and it set it to /tmp -- container-local, so a run that outlasted its GPU session came back to a wiped manifest, restarted at layer 0 and overwrote every finished shard. colocate_layerwise_checkpoint_dir existed to survive that by rewriting the path unconditionally, which also discarded any path the user had chosen deliberately. Removing the default removes the reason for the override. The recipe now ships no checkpoint_dir, matching every other layerwise recipe, and the helper becomes default_layerwise_resume_dir: it fills the field only when unset, so an explicit choice is honoured, and resume still works out of the box because leaving it unset would otherwise mean no resume at all -- the one thing this recipe is for. Three review findings alongside it. The int8 refusal tested "int8_sq" against --qformat, but that is the export format constant; the qformat preset is int8_smoothquant, and "int8_sq" is not a substring of it, so the check could never fire and an int8 run would have reached the TensorRT-LLM exporter and overwritten the shards. Matches both tokens now. The same dead test exists in export_quantized's own branch, left alone here. The export identity recorded only model class, layer count, format names and kv-cache format, so two FP8 runs differing in which modules were quantized shared one identity. It now carries a digest of the resolved quant config, which covers module selection and block settings, plus the source path. Source weights are not digested -- that would mean reading the whole model -- so differently-trained weights at the same path still compare equal; the code says so. The changelog entry was filed under Megatron Framework rather than Quantization. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
State why rather than what, and drop the restated-signature prose. No behavior change. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Export retargeting already walked a list-form algorithm, but the checkpoint-dir helpers only looked at a dict, so such a recipe calibrated with no resume manifest and lost its progress on interrupt. One list-aware accessor now backs all three helpers. Also correct the recipe's equivalence claim: layerwise export writes one shard per layer, so it matches whole-model export tensor for tensor, not byte for byte. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Two failures found in review, both on the per-layer export path. cache_outputs_for_next_layer_calib leaves the layer it just calibrated in 'run' mode with a drained deque. The qdq_from_prev=False branch already reset it; the True branch did not, because until this PR nothing ran the layer afterwards. The fusion probe does, so every fusing format asserted on layer 0. Reset it there too, and cover the combination -- it is the one layerwise axis the tests missed. The run identity is written before layer 0, so a run that died early left a directory holding only .layerwise_export.json, and re-running into it after a config change was refused as if shards were at stake. Bind it only once shards exist, and rewrite a stale one instead of keeping it forever. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
A sequential algorithm list made export ownership ambiguous. Detection read the first layerwise block, so an exporting block behind a non-exporting one missed the opt-in and exported to the recipe's placeholder; deriving the resume dir had the mirror bug, letting another pass's explicit checkpoint_dir stand in for the exporting one's, and giving every block one shared directory when none was set. Export finalizes shards as calibration walks the layers, so only the final pass can own them: layerwise_export_block() enforces exactly one export_dir on the last entry and every export-driven helper reads through it. resolve_checkpoint_dir now resolves each block against its own base, so two passes cannot collide on one manifest. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
AWQ and SVDQuant are only distinguishable after calibration -- their discriminators are registered by the calibrator -- so the constructor's gate saw plain NVFP4 and passed. An NVFP4-AWQ recipe then exported unfused pre_quant_scale plus the layernorm the whole-model path folds it into: 24 extra keys and 24 mismatched tensors against export_hf_checkpoint on a tiny Llama. Re-checking on the first exported layer catches any format that only becomes visible after calibration, not just AWQ by name. Also state what per-layer export leaves behind. The tail is converted in place, because its tensors are read from model.state_dict() after dispatch and restoring first would collect unpacked weights -- so the model is unusable for inference on every run, not only a resumed one. Ungate that warning and correct the docstrings that promised otherwise, including transient_module_state, which does not restore attributes outside _parameters/_buffers/_modules. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Five tests built a whole-model baseline and compared it to a per-layer export, differing only in the quant config; the two AWQ refusals and the two run-identity refusals were paired the same way. Parameterize each family so the axis under test reads as a table, and name the config variants once instead of inlining them. Same 17 cases collected. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…view gaps The auto-derived resume dir was never opted into, and kept one full activation set per layer -- on a 40-layer model that dwarfs the checkpoint it sits beside. Prune every boundary but the committed one, after the manifest commits it so a crash mid-write still resumes. Scoped to the per-layer export path, leaving the explicit checkpoint_dir semantics untouched. Fuse the sibling experts the probe never routed to, mirroring the replay requantize_resmooth_fused_llm_layers does for the same reason: sync_moe_gate_up_amax covers weight_quantizer.amax but not a static quantizer's global_amax, so unrouted gate/up pairs would keep unmerged scales. Drop the format re-check latch -- it inspected only one layer, so a recipe applying a pre-quant-scale format to a subset still shipped unfused pre_quant_scale. Refuse MTP checkpoints before calibration on the normal load path: only the FSDP2 loader flags the prefixes early, so the run used to write a complete-looking checkpoint and fail afterwards. Also: no MoE gate/up warning at finalize (the sync happens inside transient_module_state, so it fired on every MoE run and reported a miss that did not happen); delete out-of-range shards from a longer previous run; null-safe manifest read plus a num_layers check; drop __all__, which advertised names the package never exported. The resume test now interrupts for real rather than rewinding a finished run's manifest -- a state a crash cannot produce. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
… config The format gate lived in two places with a duplicated message, because AWQ and SVDQuant only become visible once the calibrator registers their discriminators. assert_formats_supported() is now that check, called before calibration to fail early and on each exported layer as the authority; assert_layerwise_export_supported keeps the structural cases. Record the parity with requantize_resmooth_fused_llm_layers as a table next to _fuse_shared_inputs: every step of that pass is intra-layer, so each maps to a per-layer equivalent or an explicit refusal, and a gap is visible rather than discovered. Resume is an axis of the equivalence oracle, not one example: every config now runs fresh and after a real interruption, since a resumed run re-enters export part-way through the model. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
The reset exists so the fusion probe hits a real forward. Without an exporter nothing runs the layer before _set_layer_states does, and layerwise calibration with get_qdq_activations_from_prev_layer=True has no test coverage, so leave that path exactly as it was. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…g it Tensor equality against a whole-model export never resolves the index, so a weight_map entry naming the wrong shard compares equal and still fails to load. That mapping is the one artifact per-layer export builds differently. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Added to mirror requantize_resmooth_fused_llm_layers, but sync_moe_gate_up_amax already walks every expert routed or not, and gate/up is the only input-sharing group among expert linears -- so the replay can only revisit pairs that sync has handled. Three runs on Qwen3-30B-A3B with static expert weights, replay on and off, were bitwise identical to the whole-model export. The hazard it was meant to cover is real but unreachable here: a static quantizer's weight_scale_2 reads global_amax, which sync does not write, but promote_static_block_weight_quantizers has no caller on this path, so global_amax stays unset and weight_scale_2 falls back to the amax sync does write. Fixing that belongs in the sync, not in a per-layer replay that no test can reach. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
…reason The flag flips silently, so a user who expected the post-quantization sample just sees it missing; every other consequence of export_dir announces itself. The comment also predated the tail-mutation fix -- finalize() converts the non-decoder modules in place on every run, not only a resumed one, which is how model_calib already words it. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
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) <noreply@anthropic.com> Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com> (cherry picked from commit 9f86b28)
…d 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) <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
(cherry picked from commit 02ed6b5)
…rward
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) <noreply@anthropic.com>
Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
(cherry picked from commit 4c50e13)
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>
Fridah-nv
force-pushed
the
fridah/k3-layerwise-fused-export
branch
from
August 21, 2026 23:13
521be11 to
d4f0d50
Compare
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
|
Calibration runs on the extracted language model, so the exporter has to be told which model the checkpoint describes. EXPORT_PARENT_ATTR carries that link and resolve_export_parent() turns it into a key prefix, found by identity so a module that merely looks like the language model cannot be mistaken for it. Three things then move into the parent's namespace: tensor keys gain the prefix, the vision tower and projector are collected in finalize() -- calibration never saw them and every other pass walks the submodel -- and the quant config's module references are rewritten, with those towers added to exclude_modules so a loader does not read plain BF16 as quantized. The config artifacts are written from the parent, and the VLM path no longer overwrites config.json afterwards, which would have stripped quantization_config off a finished checkpoint. The refusal narrows rather than disappears: a VLM whose language model is not reachable from the full model still cannot be exported, because the prefix would be undefined and the shards would silently describe the submodel alone. Reimplemented against the current exporter rather than cherry-picked -- the name mapper is now built from the export model, since Gemma3-VL stores the decoder at model.language_model.layers but publishes it as language_model.model.layers, and a submodel-scoped mapper cannot produce the published name. build_legacy_name_mapper covers transformers < 5, where save_pretrained rather than the exporter reverses _checkpoint_conversion_mapping. The test fails without the parent link: the vision tower is absent entirely. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
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 <export_path>.layerwise_resume, which lives next to the shards it describes. Signed-off-by: Fridah-nv <201670829+Fridah-nv@users.noreply.github.com>
Contributor
|
Curious to know if Kimi-K3 can be loaded with a single B200 node? It seems difficult given its size. |
Contributor
Author
We are able to do that with layerwise, the tradeoff is calibration speed |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Type of change: Bug fix + new feature
Stacked on #2136. Brings per-layer fused export up on a model that genuinely needs it —
moonshotai/Kimi-K3: 1.5 TB, 896 experts, 93 layers, a VLM — quantized to NVFP4 on asingle B200.
Rebased onto #2136 (2026-08-21). #2136's branch was rewritten, so six of this PR's
original commits were stale copies of work that has since landed there — the resume-manifest
and orphan-shard fixes among them. Those are dropped, not lost: they are #2136's now. The
pre-rebase branch is preserved at
fridah/k3-layerwise-fused-export-prerebase.What remains here
get_homogeneous_hf_decoder_layersunwrapped.modelthen.language_modelonce each, so it only found layers exactly two wrappersdeep in that order. K3 keeps its decoder at
language_model.model.layers, so the walkstopped on the intermediate wrapper and reported the architecture unsupported. Now
unwraps iteratively, bounded against cycles. Not K3-specific.
language model, so the exporter is told which model the checkpoint describes:
EXPORT_PARENT_ATTRcarries the link and the parent is located by identity. Tensor keysgain the prefix, the unquantized towers are exported (they were dropped entirely) and
added to
exclude_modules(or a loader reads plain BF16 as NVFP4), and the configartifacts come from the parent. The refusal narrows rather than disappears: a VLM whose
language model is not reachable from the full model is still refused, since the prefix
would be undefined. Includes a transformers-4 hub-name fallback — also not
K3-specific: any transformers-4 model with a
_checkpoint_conversion_mappinggotin-memory names from per-layer export and hub names from whole-model export.
its own module's pre-forward hook, so a weight read from a sibling's forward is on meta
when used. K3's
_apply_attn_resdoes exactly that; a disk-offloaded K3 could not run aforward at all.
*.experts.*ratherthan
*block_sparse_moe*(the broad glob also matchesshared_experts.*androuted_expert_*_proj, 552 modules the vendor left unquantized), plus itsptq.mdrow.Also fixed during the rebase: the recipe shipped
checkpoint_dir: /tmp/..., which iscontainer-local — a run that outlasts its GPU session came back to a wiped manifest and
restarted at layer 0, the exact failure the recipe exists to avoid. Left unset, #2136 derives
<export_path>.layerwise_resume, next to the shards it describes.Usage
python examples/hf_ptq/hf_ptq.py \ --pyt_ckpt_path <bf16_ckpt> \ --recipe general/ptq/nvfp4_experts_only-kv_fp8_layerwise_export_offload \ --export_path <out> \ --qformat nvfp4 --trust_remote_code --attn_implementation eager \ --offload_folder <scratch> --max_gpu_memory_gb 140 --max_cpu_memory_gb 1700 \ --calib_size 256 --batch_size 8 --skip_generate # Re-running the same command IS the resume path: it reads the manifest beside the # shards, skips finished layers, and continues.Testing
tests/gpu/torch/export/test_layerwise_export.py— 25 passed (24 inherited from #2136,plus multimodal equivalence: VLM namespace, towers present, config from the parent). That
test is not vacuous: with the parent link stubbed out it fails with "vision tower missing
from the checkpoint".
tests/unit/recipe283 ·tests/unit/torch/export172 ·tests/examples/hf_ptq36 ·pre-commit clean.
Validated on real models, not only fixtures:
projections (= 92 × 896 × 3) carrying a calibrated
input_scale.printed
Checkpoint: resuming layerwise calibration from layer 13/93and skipped thefinished work.
quant_algo=NVFP4,kv_cache_dtype=fp8_e4m3, andselects the
FLASHINFER_TRTLLMNvFp4 MoE kernel (not the emulation fallback).Re-run pending after the rebase. The Kimi-K3 numbers above predate it. The multimodal
path was reimplemented against #2136's current exporter rather than cherry-picked, so it
needs one full K3 run to confirm before this leaves draft.
Not validated: generation and accuracy. The checkpoint is 1.65 TB against 1.46 TB of HBM
on 8× B200, and vLLM's
cpu_offload_gbis a no-op for this model (80 and 200 givebyte-identical on-device memory). That is a hardware gap, not a checkpoint defect.
Before your PR is "Ready for review"
layerwise.export_diris opt-in andunset by default. One behaviour change relative to feat(export): export each decoder layer as layerwise calibration finishes it #2136: multimodal models are now
supported rather than refused, so feat(export): export each decoder layer as layerwise calibration finishes it #2136's VLM-refusal test needs reconciling with this PR
— see below.
guidance in
CONTRIBUTING.md: N/A — no new dependencies, no copied code.layerwise.export_direntry; this PRneeds it amended to drop "multimodal models are refused" once the merge order is settled.
Additional Information
Draft: depends on #2136 and must not merge before it.
is covered there by a unit test and verified on a real Qwen3.6-35B checkpoint. This PR
replaces it with support. Whichever lands second needs that test reconciled; the cleanest
order is #2136 first, then this PR updating the refusal test alongside the feature.
-s) but not GPG signed — consistent with the #2136 branch,flagged since the template asks.