Exportable kimi#47096
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Exporter support for Kimi-K2.5 (vision/audio attention registered with the reshaped vision-attention patch) plus: - Kimi vision rotary now emits the standard packed (seq, head_dim) cos/sin instead of the batched (1, seq, head_dim) LLM form, matching every other vision encoder (no exporter-side normalization needed). - ExecuTorch: materialise non-contiguous reshapes via a backend-local _patch_reshape rather than an unconditional clone in the shared vision patch, so ONNX/OpenVINO/dynamo graphs don't carry the copy. - is_multimodal short-circuits to False for non-PreTrainedModel inputs. - vision_utils.get_vision_cu_seqlens gains a merge_temporal option. Co-Authored-By: Claude Opus 4.8 <[email protected]>
c0cd540 to
7f0d32e
Compare
vasqu
left a comment
There was a problem hiding this comment.
Just have a few smaller comments because the changes are honestly a bit hard to verify without integration tests cc @zucchini-nlp whenever you have time to check
| # (n_seg, heads, seg_len, dim) → (n_seg, seg_len, heads, dim) → (seq, heads*dim) | ||
| attn_output = attn_output.transpose(1, 2).reshape(seq_length, -1).contiguous() | ||
| # (n_seg, heads, seg_len, dim) → (n_seg, seg_len, heads, dim) → (seq, heads*dim). | ||
| attn_output = attn_output.transpose(1, 2).reshape(seq_length, -1) |
There was a problem hiding this comment.
not needed in the default (dynamo export) path, if a backend needs it it should handle it on its own (see patch_reshape)
|
|
||
|
|
||
| @register_patch("executorch", "torch.reshape", "torch.Tensor.reshape") | ||
| def _patch_reshape(original): |
There was a problem hiding this comment.
hmm ok so here we have the issue indirectly re contiguous?
There was a problem hiding this comment.
the issue was always in executorch, imo this is more direct 😆
| A non-`PreTrainedModel` (e.g. a bare `nn.Module`) has no canonical `get_encoder`/`get_decoder` | ||
| accessors and is trivially not multi-modal, so it short-circuits to `False`. | ||
| """ | ||
| return isinstance(model, PreTrainedModel) and bool(_find_multimodal_submodules(model)) |
There was a problem hiding this comment.
just fmi where do we suddenly call on a normal nn module?
There was a problem hiding this comment.
not suddenly tbh, this happens during testing and i didn't catch it before
| time_table = torch.cat( | ||
| [self.time_position_embeddings.new_zeros(1, self.dim), self.time_position_embeddings.squeeze(1)] | ||
| ) | ||
| pos = pos + time_table[get_vision_frame_index(grid_thw, kwargs=kwargs)] |
There was a problem hiding this comment.
is this not quite expensive?
There was a problem hiding this comment.
hmm yes could be, it's part of vectorizing the old per-image grid_thw.tolist() loop (untraceable for export)
the zero-pad table replaced the old if t > 1 branch. you are right we can avoid rebuilding it per forward !
moved the pad into the buffer at init, so forward is now just a gather.
| class Kimi_K25VisionRotaryEmbedding(Gemma4VisionRotaryEmbedding): | ||
| def forward(self, x, position_ids): | ||
| position_ids_expanded = position_ids.permute(1, 2, 0)[..., None].float() # shape (bs, positions, 2, 1) | ||
| position_ids_expanded = position_ids.transpose(0, 1)[..., None].float() # (positions, 2, 1) |
There was a problem hiding this comment.
what happened to the batch dim?
There was a problem hiding this comment.
the extra (batch) dim was always size 1 in this packed path, so carrying it through rope was pointless and it was squeezed out
|
I'm off 9-15 July so will do it after I am back (unless someone wants to add tests before that) 😄 |
|
run-slow: kimi_k25 |
|
This comment contains models: ["models/kimi_k25"] |
|
[For maintainers] Suggested jobs to run (before merge) run-slow: kimi_k25 |
CI recapDashboard: View test results in Grafana |
vasqu
left a comment
There was a problem hiding this comment.
Just some smaller comments but overll looks good
| inputs["cu_seqlens"] = get_vision_cu_seqlens(grid_thw, merge_temporal=temporal_encoder) | ||
| inputs["max_seqlen"] = get_max_seqlen(inputs["cu_seqlens"], model.config, kwargs=inputs) |
There was a problem hiding this comment.
Why not use get_vision_attention_seqlens as before?
| self.final_layernorm = nn.LayerNorm(config.hidden_size, eps=1e-05) | ||
| self.post_init() | ||
|
|
||
| def temporal_patch_merger( |
There was a problem hiding this comment.
Technically slightly breaking can we keep BC for a bit at least?
| ) | ||
|
|
||
| cu_seqlens = lengths.cumsum(dim=0, dtype=torch.int32) | ||
| cu_seqlens = get_vision_cu_seqlens(grid_thw, merge_temporal=True, kwargs=kwargs) |
There was a problem hiding this comment.
Can apply our fused version for max seq len together
There was a problem hiding this comment.
General comment and not on you: But we do introduce some very model specific workarounds. I feel like it would make more and more sense to start moving things into a collator / processor that handles these
|
Also cc @zucchini-nlp if you want to take a look |
What does this PR do?
Fixes # (issue)
Code Agent Policy
The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. We are currently bottlenecked by our ability to review and respond to them. As a result,
we ask that new users do not submit pure code agent PRs at this time.
You may use code agents in drafting or to help you diagnose issues. We'd also ask autonomous "OpenClaw"-like agents
not to open any PRs or issues for the moment.
PRs that appear to be fully agent-written will probably be closed without review, and we may block users who do this
repeatedly or maliciously.
This is a rapidly-evolving situation that's causing significant shockwaves in the open-source community. As a result,
this policy is likely to be updated regularly in the near future. For more information, please read
CONTRIBUTING.md.Before submitting
Pull Request checks?
to it if that's the case.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.