[None][feat] Support multi-modal part of K3 - #17050
Conversation
|
In order to test it, I renamed |
| delta | K2.5 | K3 | ||
| ---------------------+----------------------------+--------------------------- | ||
| vision norms | LayerNorm | RMSNorm (torch.nn.RMSNorm) | ||
| attention head_dim | vt_hidden_size // heads | qkv_hidden_size // heads |
There was a problem hiding this comment.
Nit: might as well align the | here?
| # Vision tower is not quantized (checkpoint quant ignore list covers | ||
| # vision_tower.* / mm_projector.*): keep only the kv-cache quant algo. | ||
| self.model_config.quant_config = QuantConfig( | ||
| kv_cache_quant_algo=model_config.quant_config.kv_cache_quant_algo |
There was a problem hiding this comment.
Should this be kept? AFAIK we don't even have the notion of KV cache for the vision encoder's transformer layers, and from issues we saw with past models, the value here may influence the choice of kernels.
See #12851 as an example.
| self.model_config.pretrained_config = copy.copy(model_config.pretrained_config) | ||
| pretrained_config = self.model_config.pretrained_config | ||
| model_dtype = ( | ||
| getattr(pretrained_config, "torch_dtype", None) |
There was a problem hiding this comment.
Some lines (e.g. 238) access model_config.torch_dtype. Is there any meaningful difference between that and pretrained_config.torch_dtype? If not, shall we settle on model_config.torch_dtype to get away from needing getattr etc.?
| model_dtype = getattr(torch, model_dtype, torch.bfloat16) | ||
| pretrained_config.torch_dtype = model_dtype | ||
|
|
||
| vision_cfg = getattr(pretrained_config, "vision_config", {}) |
There was a problem hiding this comment.
Can't we assume that, if we're mapped to KimiK3ForConditionalGeneration, these fields / nested fields all exist?
| **kwargs, | ||
| ) -> None: | ||
| config = model_config.pretrained_config | ||
| self._supports_sdpa = True |
There was a problem hiding this comment.
Nit: maybe leave a comment for this?
| # K3 components below. | ||
| PreTrainedModel.__init__(self, config) | ||
|
|
||
| if hasattr(self, "llm"): |
There was a problem hiding this comment.
Could you leave a comment for why we have this early return? It's not quite apparent for less familiar readers such as myself.
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Solid bring-up — the delta table in the module docstring and the head_dim / replication comments make this much easier to follow than the usual VLM port, and the MMMU channel extractor is a real scoring bug worth fixing. Given this targets feat/kimi_k3 and the branch is racing to functional completeness, I've split the items below into before merge (cheap, seconds each) and follow-up PR (everything else — none of it should hold this up).
Before merge (cheap)
Ticket. Title is [None][feat] on a ~1100-line model bring-up. This should carry TRTLLM-14704 per the repo convention — [TRTLLM-14704][feat] Support multi-modal part of K3. [None] is for chores/docs/waives.
Stale path reference. run_gsm8k_kimi_k3.sbatch → run_eval_kimi_k3.sbatch is a user-facing rename; docs are updated, but tests/integration/defs/kimi_k3_disagg_parity.py:42 still points at the old name. One-line fix.
Description vs. diff. Two things a future reader (and the merge-back reviewer) will need, and they're just description lines:
- The
modeling_kimi_k25.pychanges aren't wire-up — they change behavior for the shipped K2.5 model:_vision_requires_replicationnow silently replicates the vision tower for anytp_sizethat doesn't divide the head count, andload_weightsnow materializes every checkpoint slice up front. Both look right, but please say so explicitly and note whether K2.5 was re-validated. - TEP16 support in the sbatch isn't mentioned at all.
Follow-up PR (file under TRTLLM-14674, not blocking)
Test coverage. tests/unittest/others/test_lm_eval.py isn't referenced by any file under tests/integration/test_lists/, so the tests listed as this PR's coverage don't run in pre-merge CI. The 650 lines of new model + config code have no tests at all — is_kimi_k3_multimodal_config is a pure dict predicate and KimiK3Config.from_dict round-trip is a few lines; both are cheap to pin and both are exactly the kind of routing logic that breaks silently when a checkpoint field is renamed. This should land before the branch merges back, but it doesn't need to be in this PR.
Inheritance structure. KimiK3VisionModel.__init__ and K3MoonViT3dEncoder.__init__ deliberately skip their parents and re-implement them — ~90 lines duplicated from the K2.5 versions, already diverged (K3 adds qkv_hidden_size, drops layer_norm_eps). A shared _parse_vision_cfg(vision_cfg) helper plus a _build_blocks() hook would let both share the config parsing and keep future K2.5 fixes reaching K3. Worth doing while the K2.5/K3 pairing is still fresh in someone's head, but it's a refactor of working code — follow-up.
| self.norm1 = nn.RMSNorm(hidden_dim) | ||
| # head_dim is taken from model_config.pretrained_config.head_dim, which | ||
| # KimiK3VisionModel sets to qkv_hidden_size // num_heads (128). | ||
| self.attn = KimiK25VisionAttention( |
There was a problem hiding this comment.
Follow-up, not blocking.
Attention.__init__ accepts an explicit head_dim= kwarg, added specifically for "sub-modules (e.g. VLM vision encoders) whose head_dim does not match the top-level config's head_dim" (tensorrt_llm/_torch/modules/attention.py:450). Routing it through pretrained_config.head_dim (set on a copy.copy of the composite config at line 330) works, but it's action-at-a-distance: anything else that reads head_dim off that config copy now sees 128 instead of the text head_dim, and the coupling is invisible from here.
In a follow-up, add a head_dim: Optional[int] = None param to KimiK25VisionAttention, forward it to Attention, and drop the pretrained_config.head_dim mutation.
There was a problem hiding this comment.
Agreed — deferring to the follow-up as scoped: add head_dim: Optional[int] = None to KimiK25VisionAttention, forward it to Attention, and drop the pretrained_config.head_dim mutation on the config copy.
| # Reference uses torch.nn.RMSNorm(hidden_dim) with default eps for the | ||
| # per-layer norms; match it exactly (created in fp32, cast with the rest | ||
| # of the vision tower to the model dtype in load_weights()). | ||
| self.norm0 = nn.RMSNorm(hidden_dim) |
There was a problem hiding this comment.
Follow-up, not blocking — but worth a comment now if it's zero-cost.
nn.RMSNorm(hidden_dim) leaves eps=None, which torch resolves at runtime to torch.finfo(input.dtype).eps — ~7.8e-3 for bf16 vs ~1.2e-7 for fp32. So "match [the reference] exactly" holds only while both sides run the tower in the same dtype, and the value silently changes if the vision tower dtype ever changes. Pinning it explicitly (nn.RMSNorm(hidden_dim, eps=<value the reference resolves to>)) with the reference's dtype noted in the comment would make that assumption durable. Same for final_layernorm at line 207.
There was a problem hiding this comment.
Documented in a97b9c2 at both sites (norm0/norm1 + final_layernorm): eps=None resolves at runtime to torch.finfo(input.dtype).eps, so exact parity with the reference holds only while both towers run the same dtype (bf16 today). Left the value unpinned since the reference uses the default too — pinning would have to happen on both sides together to keep bit-parity, which the comment now spells out.
| if not isinstance(vision_cfg, dict): | ||
| vision_cfg = (vision_cfg.to_dict() | ||
| if hasattr(vision_cfg, "to_dict") else vars(vision_cfg)) | ||
| num_heads = vision_cfg.get("vt_num_attention_heads", |
There was a problem hiding this comment.
The 12 fallback here is the K3 default, but this helper is also on the K2.5 path (KimiK25VisionModel defaults to 16 heads at line 828). For a K2.5 config that omits both vt_num_attention_heads and num_attention_heads, this computes 12 % tp_size while the model builds a 16-head tower — the two disagree about whether replication is needed. It stays self-consistent because the mapping is applied unconditionally, so no correctness bug today; the effect is that a shardable K2.5 tower can be silently forced to tp=1.
Cleaner: take num_heads as an argument from the caller (both KimiK*VisionModel.__init__ already resolve it with the right per-model default) rather than re-deriving it with a hardcoded default. This is the only item here that touches shipped K2.5 behavior, so if you want it in this PR it's a ~5-line change; otherwise fold it into the follow-up.
There was a problem hiding this comment.
Done in a97b9c2 — _vision_requires_replication / _get_vision_tp_mapping now take num_heads as a required argument and never re-derive it; both KimiK*VisionModel.__init__s resolve it first (16- resp. 12-head defaults) and the two projector MLPs receive it through their constructors, so every replication decision uses the head count the tower is actually built with.
| and required_multimodal_token_ids.issubset(config_dict)) | ||
|
|
||
|
|
||
| def is_kimi_k3_multimodal_config(config_dict: dict) -> bool: |
There was a problem hiding this comment.
Follow-up, not blocking.
This is a pure dict → bool predicate and the routing decision for the whole K3 checkpoint family; it's worth a handful of unit-test cases (composite → True; language_model_only: true → False; missing vision_config → False; empty-dict sub-configs → False; text-only kimi_linear → False). Cheap, and it's the thing that will break silently if the released config renames a field — please get it in before the branch merges back.
Separately (also follow-up): language_model_only is a checkpoint field, so today the only way to run a K3 VL checkpoint text-only is to edit config.json. If text-only serving of a VL checkpoint is a use case (e.g. perf comparison against the DEP16 references), consider an env/arg escape hatch.
There was a problem hiding this comment.
Added tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py in a97b9c2 with the five cases you listed (composite → True; language_model_only: true → False; missing vision_config → False; empty-dict sub-configs → False; text-only kimi_linear → False) plus non-dict sub-configs, missing text_config, explicit language_model_only: false, and a wrong-model_type guard.
On the text-only escape hatch for VL checkpoints: agreed it's worth having (checkpoint-only fields are awkward to override) — deferring to the follow-up; an explicit LLM arg / model_kwargs override seems better than an env var.
| # default mode the batch bump also rewrites the YAML's own max_batch_size and | ||
| # the nested cuda_graph_config max_batch_size so graph coverage follows. | ||
| if [[ "$PARALLEL" == tep ]]; then | ||
| TEP_CONFIG=$REPO/examples/kimi_k3/.eval_tep_runtime.$SLURM_JOB_ID.yaml |
There was a problem hiding this comment.
Two issues with the TEP rewrite. Both are follow-up material — the script works as written — but (1) is a one-line fix if you want it now.
-
$REPO/examples/kimi_k3/.eval_tep_runtime.$SLURM_JOB_ID.yamlwrites a per-job file into the source tree and never removes it, so repeated TEP runs accumulate dotfiles in a checked-in examples dir.$SLURM_SUBMIT_DIR(where the log already goes) ormktemp+ atrap ... EXITwould be cleaner. If it must stay under$REPO, at minimum add it to.gitignore— that part is a one-liner worth doing here. -
sed -e "s/^ max_batch_size: .*/..."keys on exactly two spaces of indentation and will rewrite any singly-nestedmax_batch_size, not justcuda_graph_config's — silently a no-op if the YAML is ever reindented, silently wrong if another nested block gains the key. A short Pythonyaml.safe_load/safe_dumprewrite would be both correct and more readable, and matches the repo's preference for Python over shell once there's real logic. Follow-up.
There was a problem hiding this comment.
(1) Done in a97b9c2: the per-job yaml is now removed by a trap ... EXIT in the batch script, and both .eval_tep_runtime.*.yaml / .eval_dflash_runtime.*.yaml are gitignored (the dflash mode shared the same wart). It stays under $REPO deliberately — that's the only path guaranteed container-mounted at the same location on every rank; $SLURM_SUBMIT_DIR only coincides with $REPO by default, and a node-local mktemp file would be invisible to the other three nodes. The script now carries a comment saying exactly that.
(2) Agreed on the indent-keyed sed — taking the yaml.safe_load/safe_dump rewrite as the follow-up. One wrinkle to solve there: the rewrite currently runs in the host batch shell where PyYAML isn't guaranteed, so it likely moves inside the container command (rank-0-gated) or gains a python-with-fallback path.
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: Michal Guzek <mguzek@nvidia.com>
- Pass the vision tower head count from both KimiK*VisionModel inits and the projector MLPs into _vision_requires_replication / _get_vision_tp_mapping instead of re-deriving it there with a hardcoded K3 default (a K2.5 config omitting both head-count keys could silently force a shardable 16-head tower to tp=1). - Document the runtime-dtype-dependent eps of the reference-matching default-eps RMSNorms in the K3 vision tower (norm0/norm1 and final_layernorm). - Clean up the per-job TEP runtime yaml: remove it via trap on batch script exit and gitignore the .eval_tep_runtime/.eval_dflash_runtime patterns. - Unit-test is_kimi_k3_multimodal_config (composite, language_model_only opt-out, missing/empty/non-dict sub-configs, text-only kimi_linear, wrong model_type) so released-config field renames fail loudly. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
183b018 to
a97b9c2
Compare
4 participants
@coderabbitai summary
Description
Adds the multi-modal (vision-language) portion of Kimi-K3 to the PyTorch backend.
tensorrt_llm/_torch/configs/kimi_k3.py— K3 multimodal configtensorrt_llm/_torch/models/modeling_kimi_k3_vl.py— K3 VL modelconfig_utils.py, model/config__init__.py,modeling_kimi_k25.py,modeling_kimi_linear.pyevaluate/lm_eval.py+evaluate/post_processing.pytests/unittest/others/test_lm_eval.pyTest Coverage
tests/unittest/others/test_lm_eval.py