[WIP] [Kimi-K3] - #2978
Conversation
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
|
Important Review skippedIgnore keyword(s) in the title. ⛔ Ignored keywords (3)
Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds support for the Kimi-K3 and Inkling models, introducing their configurations, modeling implementations, processors, tokenizers, and quantization examples. It also updates the sequential pipeline and quantization modifiers to handle sequential epoch start and optimization events. The review feedback highlights several critical bugs, including an incorrect parameter name (allowed_modules instead of allow_modules) in the quantization modifier, a potential TypeError when subscripting past_key_values in the Kimi-K3 model, and a hardcoded if False: condition that bypasses the MoE inference path. Other recommended improvements include removing hardcoded absolute paths in example scripts, restoring commented-out test cases, replacing a bare except: clause, and using getattr defensively to avoid potential AttributeErrors.
| ): | ||
| def on_sequential_epoch_start(self, state, event, modules, **kwargs): | ||
| model = state.model | ||
| apply_quantization_config(model, self.resolved_config, allowed_modules=modules) |
There was a problem hiding this comment.
The argument name allowed_modules is incorrect and does not match the allow_modules parameter expected by apply_quantization_config. This mismatch will cause a runtime TypeError when on_sequential_epoch_start is executed.
| apply_quantization_config(model, self.resolved_config, allowed_modules=modules) | |
| apply_quantization_config(model, self.resolved_config, allow_modules=modules) |
| ): | ||
| # Retrieve the first layer to inspect the logits and mask out the hidden states | ||
| # that are set to 0 | ||
| first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] |
There was a problem hiding this comment.
Subscripting past_key_values directly will raise a TypeError because KimiDynamicCache does not implement __getitem__. You should access the key cache directly via past_key_values.key_cache[0] or implement __getitem__ in KimiDynamicCache.
| first_layer_past_key_value = past_key_values[0][0][:, :, :, 0] | |
| first_layer_past_key_value = past_key_values.key_cache[0][:, :, :, 0] |
| if self.use_latent_moe: | ||
| hidden_states = self.routed_expert_down_proj(hidden_states) | ||
|
|
||
| if False:#if not self.training: |
There was a problem hiding this comment.
The hardcoded if False: condition completely bypasses the moe_infer path, forcing the model to always use moe_train even during inference. This should be changed to check self.training to ensure correct and optimized inference behavior.
| if False:#if not self.training: | |
| if not self.training: |
| caller = self._eval_expr(node.func) | ||
|
|
||
| except Exception: | ||
| except: |
There was a problem hiding this comment.
Avoid using a bare except: clause as it catches all exceptions, including system-exiting ones like SystemExit and KeyboardInterrupt. Use except Exception: instead to only catch standard errors.
| except: | |
| except Exception: |
| # ("meta-llama/Meta-Llama-3-8B-Instruct", AutoModelForCausalLM, None, "text", []), | ||
| # ( | ||
| # "CohereLabs/command-a-vision-07-2025", | ||
| # Cohere2VisionForConditionalGeneration, | ||
| # ["Cohere2DecoderLayer"], | ||
| # "deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct", | ||
| # AutoModelForCausalLM, | ||
| # None, | ||
| # "text", | ||
| # [], | ||
| # ), | ||
| # ( | ||
| # "mistralai/Mixtral-8x7B-Instruct-v0.1", | ||
| # AutoModelForCausalLM, | ||
| # None, | ||
| # "text", | ||
| # [], | ||
| # ), | ||
| # ( | ||
| # "ibm-granite/granite-20b-code-instruct-8k", | ||
| # AutoModelForCausalLM, | ||
| # None, | ||
| # "text", | ||
| # [], | ||
| # ), | ||
| # ("google/gemma-3n-E2B-it", AutoModelForCausalLM, None, "text", ["timm"]), | ||
| # ("unsloth/DeepSeek-R1-0528-BF16", AutoModelForCausalLM, None, "text", []), | ||
| # # --- vision --- | ||
| # ( | ||
| # "HuggingFaceM4/Idefics3-8B-Llama3", | ||
| # Idefics3ForConditionalGeneration, | ||
| # ["LlamaDecoderLayer"], | ||
| # "vision", | ||
| # [], | ||
| # ), | ||
| # ( | ||
| # "llava-hf/llava-1.5-7b-hf", | ||
| # LlavaForConditionalGeneration, | ||
| # ["LlamaDecoderLayer"], | ||
| # "vision", | ||
| # [], | ||
| # ), | ||
| # ( | ||
| # "meta-llama/Llama-3.2-11B-Vision-Instruct", | ||
| # MllamaForConditionalGeneration, | ||
| # ["MllamaSelfAttentionDecoderLayer"], | ||
| # "vision", | ||
| # [], | ||
| # ), | ||
| # # skip phi3_v because of its processor is annoying and requires special code | ||
| # ( | ||
| # "mgoin/pixtral-12b", | ||
| # LlavaForConditionalGeneration, | ||
| # ["MistralDecoderLayer"], | ||
| # "vision", | ||
| # [], | ||
| # ), | ||
| # ( | ||
| # "Qwen/Qwen2.5-VL-7B-Instruct", | ||
| # Qwen2_5_VLForConditionalGeneration, | ||
| # ["Qwen2_5_VLDecoderLayer"], | ||
| # "vision", | ||
| # ["torchvision"], | ||
| # ), | ||
| # # TODO: add gated model command-a to CI runner tokens | ||
| # # ( | ||
| # # "CohereLabs/command-a-vision-07-2025", | ||
| # # Cohere2VisionForConditionalGeneration, | ||
| # # ["Cohere2DecoderLayer"], | ||
| # # "vision", | ||
| # # [], | ||
| # # ), | ||
| # ( | ||
| # "Qwen/Qwen2-VL-2B-Instruct", | ||
| # Qwen2VLForConditionalGeneration, | ||
| # ["Qwen2VLDecoderLayer"], | ||
| # "vision", | ||
| # ["torchvision"], | ||
| # ), | ||
| # ( | ||
| # "mistralai/Mistral-Small-3.1-24B-Instruct-2503", | ||
| # Mistral3ForConditionalGeneration, | ||
| # ["MistralDecoderLayer"], | ||
| # "vision", | ||
| # [], | ||
| # ), | ||
| # ( | ||
| # "google/gemma-3-4b-it", | ||
| # Gemma3ForConditionalGeneration, | ||
| # ["Gemma3DecoderLayer"], | ||
| # "vision", | ||
| # [], | ||
| # ), | ||
| # ( | ||
| # "meta-llama/Llama-4-Scout-17B-16E-Instruct", | ||
| # Llama4ForConditionalGeneration, | ||
| # "Llama4TextDecoderLayer", | ||
| # "vision", | ||
| # [], | ||
| # ), | ||
| # ( | ||
| # "meta-llama/Llama-4-Maverick-17B-128E-Instruct", | ||
| # Llama4ForConditionalGeneration, | ||
| # "Llama4TextDecoderLayer", | ||
| # "vision", | ||
| # [], | ||
| # ), | ||
| ( | ||
| "Qwen/Qwen2-VL-2B-Instruct", | ||
| Qwen2VLForConditionalGeneration, | ||
| ["Qwen2VLDecoderLayer"], | ||
| "vision", | ||
| ["torchvision"], | ||
| ), | ||
| ( | ||
| "mistralai/Mistral-Small-3.1-24B-Instruct-2503", | ||
| Mistral3ForConditionalGeneration, | ||
| ["MistralDecoderLayer"], | ||
| "vision", | ||
| [], | ||
| ), | ||
| ( | ||
| "google/gemma-3-4b-it", | ||
| Gemma3ForConditionalGeneration, | ||
| ["Gemma3DecoderLayer"], | ||
| "vision", | ||
| [], | ||
| ), | ||
| ( | ||
| "meta-llama/Llama-4-Scout-17B-16E-Instruct", | ||
| Llama4ForConditionalGeneration, | ||
| "Llama4TextDecoderLayer", | ||
| "vision", | ||
| [], | ||
| ), | ||
| ( | ||
| "meta-llama/Llama-4-Maverick-17B-128E-Instruct", | ||
| Llama4ForConditionalGeneration, | ||
| "Llama4TextDecoderLayer", | ||
| "vision", | ||
| [], | ||
| ), | ||
| ( | ||
| "google/gemma-3n-E2B-it", | ||
| Gemma3nForConditionalGeneration, | ||
| None, | ||
| "vision", | ||
| ["timm"], | ||
| ), | ||
| # --- audio --- | ||
| ( | ||
| "openai/whisper-large-v3", | ||
| WhisperForConditionalGeneration, | ||
| ["WhisperDecoderLayer"], | ||
| "audio", | ||
| ["librosa", "soundfile", "torchcodec"], | ||
| ), | ||
| # ( | ||
| # "google/gemma-3n-E2B-it", | ||
| # Gemma3nForConditionalGeneration, | ||
| # None, | ||
| # "vision", | ||
| # ["timm"], | ||
| # ), | ||
| # # --- audio --- | ||
| # ( | ||
| # "openai/whisper-large-v3", | ||
| # WhisperForConditionalGeneration, | ||
| # ["WhisperDecoderLayer"], | ||
| # "audio", | ||
| # ["librosa", "soundfile", "torchcodec"], | ||
| # ), |
There was a problem hiding this comment.
Almost all test cases in the test suite have been commented out. While this is common during local development to speed up testing, please ensure these are uncommented before merging to maintain full regression test coverage.
| MODEL_ID, | ||
| device_map="auto", | ||
| max_memory={}, | ||
| offload_folder="/mnt/nvme-data/engine/kylesayrs/offload_folder", |
There was a problem hiding this comment.
Avoid hardcoding absolute paths specific to a local user directory (e.g., /mnt/nvme-data/engine/kylesayrs/...). Use relative paths or configurable arguments to ensure the examples are reproducible and work for other users.
| offload_folder="/mnt/nvme-data/engine/kylesayrs/offload_folder", | |
| offload_folder="offload_folder", |
| ) | ||
|
|
||
| # Save to disk compressed. | ||
| SAVE_DIR = "/mnt/nvme-data/engine/kylesayrs/" + MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4" |
There was a problem hiding this comment.
Avoid hardcoding absolute paths specific to a local user directory. Use relative paths instead.
| SAVE_DIR = "/mnt/nvme-data/engine/kylesayrs/" + MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4" | |
| SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-NVFP4" |
| replace_module_parallel(modules, compress_fn) | ||
|
|
||
| for module in modules: | ||
| assert getattr(module, "quantization_status", None) is None |
There was a problem hiding this comment.
Using assert statements for production state validation is discouraged because they can be globally disabled with Python's -O optimization flag. Consider raising a proper ValueError or using a conditional check instead.
| from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask | ||
| from fla.utils import tensor_cache | ||
| except ImportError: | ||
| raise ImportError("Plese run `pip install -U fla-core`") |
There was a problem hiding this comment.
Typo in user-facing error message: 'Plese' should be 'Please'.
| raise ImportError("Plese run `pip install -U fla-core`") | |
| raise ImportError("Please run 'pip install -U fla-core'") |
|
|
||
| module.quantization_status = QuantizationStatus.FROZEN | ||
| # preserve COMPRESSED status so save_pretrained can skip already-compressed modules | ||
| if module.quantization_status != QuantizationStatus.COMPRESSED: |
There was a problem hiding this comment.
Accessing module.quantization_status directly can raise an AttributeError if the attribute is not set on the module. Use getattr to handle this defensively.
| if module.quantization_status != QuantizationStatus.COMPRESSED: | |
| if getattr(module, "quantization_status", None) != QuantizationStatus.COMPRESSED: |
|
The quality checks have failed. Please run |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
|
|
||
| ## Kimi-K3 (MoE + Multimodal) Example | ||
|
|
||
| `kimi_k3_nvfp4.py` shows how to apply NVFP4 quantization to [Kimi-K3](https://huggingface.co/inference-optimization/Kimi-K3-0.18B), a multimodal Mixture-of-Experts model with a hybrid attention architecture (KDA linear attention + MLA full attention). |
There was a problem hiding this comment.
🚫 [linkspector] reported by reviewdog 🐶
Cannot reach https://huggingface.co/inference-optimization/Kimi-K3-0.18B Status: 401
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
|
The quality checks have failed. Please run |
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
|
The quality checks have failed. Please run |
|
The quality checks have failed. Please run |
|
This pull request has merge conflicts that must be resolved before it can be |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
1 participant
No description provided.