Kimi-K3 on layerwise fused export: 3 defect fixes, multimodal support, single-B200 NVFP4 - #2218
Draft
Fridah-nv wants to merge 6 commits into
Draft
Kimi-K3 on layerwise fused export: 3 defect fixes, multimodal support, single-B200 NVFP4#2218Fridah-nv wants to merge 6 commits into
Fridah-nv wants to merge 6 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 |
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>
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
1 participant
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.