Add files using upload-large-folder tool

NOTICE.txt

@@ -0,0 +1,4 @@
+CVRR applies a learned recurrent transition to OpenGVLab/InternVL3-9B.
+Native dense weights are preserved; recurrent projection weights are modified
+by merging the trained CVRR LoRA. CVRR-specific code is additional code, not
+part of the original backbone distribution. See the included upstream notices.

README.md

@@ -0,0 +1,121 @@
+---
+base_model: OpenGVLab/InternVL3-9B
+base_model_relation: finetune
+library_name: transformers
+license: mit
+tags:
+  - cvrr
+  - custom_code
+  - latent-reasoning
+  - visual-question-answering
+  - arxiv:2609.06746
+---
+
+# InternVL3-9B-CVRR
+
+<p align="center">
+  <a href="https://arxiv.org/abs/2609.06746"><strong>📄 Paper</strong></a> ·
+  <a href="https://github.com/dmis-lab/CVRR"><strong>💻 Code</strong></a>
+</p>
+
+## Introduction
+
+Latent visual reasoning is not just about storing visual information in hidden
+states. That information should matter for the model's answer. When an answer
+decoder can access the original multimodal context through a parallel path,
+the presence of informative latent states alone does not establish their use.
+
+Our work, **Reason Through the Latent! Making Latent Visual Reasoning Necessary**,
+introduces **Causal Visual Recurrent Reasoning (CVRR)**. CVRR separates visual
+access during reasoning from visual access during answering, focusing on three
+design principles:
+
+**👁️ Preserving pretrained visual competence**  
+Reasoning starts from the native image-conditioned question representation,
+preserving visual information already integrated by the pretrained VLM.
+
+**🔁 Refining reasoning with persistent visual evidence**  
+A shared native decoder layer updates the question state while retaining access
+to fixed visual evidence. The recurrent transition is adapted with LoRA and
+answer-token supervision.
+
+**🔒 Making the reasoning state the visual answer interface**  
+The upper answer decoder receives the final recurrent question state, without
+direct access to the original visual rows or original multimodal prefix caches.
+The lower-layer prefix context used for answer generation follows a text-only path.
+
+## Model Checkpoints
+
+### Final Models
+
+We release CVRR checkpoints based on the following pretrained vision-language
+backbones. Each repository contains the native backbone weights, the merged
+recurrent transition, and the custom inference code; no separate LoRA download
+is required.
+
+- 🤗 [Qwen2.5-VL-7B-CVRR](https://huggingface.co/dmis-lab/Qwen2.5-VL-7B-CVRR)
+- 🤗 [Qwen3-VL-8B-CVRR](https://huggingface.co/dmis-lab/Qwen3-VL-8B-CVRR)
+- 🤗 [Gemma3-4B-CVRR](https://huggingface.co/dmis-lab/Gemma3-4B-CVRR)
+- 🤗 [Gemma3-12B-CVRR](https://huggingface.co/dmis-lab/Gemma3-12B-CVRR)
+- 🤗 [Gemma4-12B-CVRR](https://huggingface.co/dmis-lab/Gemma4-12B-CVRR)
+- 🤗 [InternVL3-9B-CVRR](https://huggingface.co/dmis-lab/InternVL3-9B-CVRR)
+
+This checkpoint is based on **`OpenGVLab/InternVL3-9B`**.
+
+## Loading
+
+Download the complete repository and install its model-specific dependencies:
+
+```bash
+hf download dmis-lab/InternVL3-9B-CVRR --local-dir ./InternVL3-9B-CVRR
+pip install -r ./InternVL3-9B-CVRR/requirements.txt
+```
+
+```python
+import torch
+from PIL import Image
+from transformers import AutoModelForImageTextToText
+
+model = AutoModelForImageTextToText.from_pretrained(
+    "./InternVL3-9B-CVRR",
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+    device_map="cuda:0",
+).eval()
+
+image = Image.open("example.jpg").convert("RGB")
+inputs = model.prepare_inputs(
+    image,
+    "What is the dominant color? A. Red B. Blue C. Green D. Yellow. Answer with the option letter.",
+)
+logits = model.next_token_logits(**inputs)
+token_id = logits.argmax(dim=-1).item()
+print(model.tokenizer.decode([token_id]))
+```
+
+This model exposes the strict first-answer-token readout used in our analyses. Multi-token generation is not currently supported by this loader.
+
+Use one complete model replica per GPU. This custom Transformers loader does
+not support automatic model sharding, generic `save_pretrained()` reserialization,
+or direct loading through vLLM/SGLang. Preserve the downloaded directory and
+pin the Hub revision for reproducible use.
+
+## License
+
+Please follow the upstream model's terms and the included
+[`NOTICE.txt`](NOTICE.txt) and [`licenses/`](licenses/) attribution files.
+The component licenses remain applicable to their respective materials.
+
+## Citation
+
+```bibtex
+@misc{park2026reasonlatentmakinglatent,
+      title={Reason Through the Latent! Making Latent Visual Reasoning Necessary},
+      author={Suhyeong Park and Junha Jung and Jaewoo Kang},
+      year={2026},
+      eprint={2609.06746},
+      archivePrefix={arXiv},
+      primaryClass={cs.AI},
+      url={https://arxiv.org/abs/2609.06746},
+}
+```

config.json

@@ -0,0 +1,30 @@
+{
+  "model_type": "cvrr_merged",
+  "architectures": [
+    "CVRRMergedModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_cvrr_merged.CVRRMergedConfig",
+    "AutoModel": "modeling_cvrr_merged.CVRRMergedModel",
+    "AutoModelForImageTextToText": "modeling_cvrr_merged.CVRRMergedModel"
+  },
+  "release": {
+    "name": "CVRR-InternVL3-9B",
+    "base_model": "OpenGVLab/InternVL3-9B",
+    "ell_star": 35,
+    "recurrent_layer": 36,
+    "upper_decoder_start": 37,
+    "inference_T": 4,
+    "inference_beta": 0.33,
+    "lora_rank": 32,
+    "lora_alpha": 12.0,
+    "lora_dropout": 0.01,
+    "checkpoint_step": 500,
+    "format": "cvrr_native_plus_merged_transition_v1",
+    "merged_weight_dtype": "float32",
+    "native_weight_bytes": 18277586944,
+    "status": "gpu_vstar_comparison_completed",
+    "upload_ready": true,
+    "hub_model_id": "dmis-lab/InternVL3-9B-CVRR"
+  }
+}

configuration_cvrr_merged.py

@@ -0,0 +1,18 @@
+from transformers import PretrainedConfig
+
+
+class CVRRMergedConfig(PretrainedConfig):
+    model_type = 'cvrr_merged'
+
+    def __init__(self, release=None, **kwargs):
+        super().__init__(**kwargs)
+        self.release = {} if release is None else release
+        self.architectures = ['CVRRMergedModel']
+        self.auto_map = {
+            'AutoConfig': 'configuration_cvrr_merged.CVRRMergedConfig',
+            'AutoModel': 'modeling_cvrr_merged.CVRRMergedModel',
+            'AutoModelForImageTextToText': 'modeling_cvrr_merged.CVRRMergedModel',
+        }
+
+
+CVRRMergedConfig.register_for_auto_class()

configuration_source_qwen25.py

@@ -0,0 +1,1142 @@
+"""Configuration for CLOSE latent reasoning on Qwen2.5-VL.
+
+The current ``raw_evidence_question`` path keeps the complete layer-``ell_star``
+visual memory as immutable evidence and recurrently updates the text-only
+question state by cross-attending to it. Its selected interface averages all
+four residual recurrent states and may use a training-only restoration target
+from the normal multimodal question state. Historical slot-workspace fields are
+retained so old checkpoints remain loadable, but they are mutually exclusive
+with the raw-evidence path.
+
+Every architectural hyperparameter of Method §3 lives here -- nothing is
+hardcoded in ``modeling_close_qwen2_5_vl.py``.
+
+Layer-index convention (used consistently across the repo)
+----------------------------------------------------------
+``ell_star`` is **0-indexed and inclusive**: it is the index of the last decoder
+layer belonging to the lower branch. With ``num_hidden_layers == 28``::
+
+    F_{<=l*} = language_model.layers[: ell_star + 1]     # lower_slice
+    F_{>l*}  = language_model.layers[ell_star + 1 :]     # upper_slice
+
+so ``ell_star=18`` means layers 0..18 run before the split and 19..27 after it.
+``ell_star`` is measured once by ``scripts/localize_read_point.py`` (§3.1) and
+frozen before any workspace training.
+
+Careful with ``Qwen2_5_VLConfig``: its ``__setattr__`` forwards any attribute
+whose name already exists in ``text_config.__dict__`` down to the sub-config, and
+its ``__init__`` builds ``text_config`` from ``**kwargs``. Workspace fields are
+therefore declared as explicit named parameters so they never enter ``kwargs``
+and never shadow a text-config field. ``tests/test_config_roundtrip.py`` guards
+this.
+"""
+
+from __future__ import annotations
+
+import math
+
+from transformers.models.qwen2_5_vl.configuration_qwen2_5_vl import Qwen2_5_VLConfig
+
+#: How the ``S`` workspace slots are assigned M-RoPE positions once they are
+#: spliced into the replacement cache above ``ell_star``. Qwen2.5-VL uses 3D
+#: M-RoPE (t, h, w) with ``mrope_section=[16, 24, 24]``, so slots -- which have
+#: no spatial extent -- need an explicit convention.
+WORKSPACE_ROPE_MODES = (
+    # Slots continue the 1-D text position sequence directly after Q*, with
+    # t == h == w. Treats the workspace as "more text".
+    "continue",
+    # Slots reuse the (t, h, w) span the image tokens occupied in the original
+    # multimodal prefill, subsampled to S positions. Preserves the read point's
+    # positional geometry but reintroduces image-derived indices.
+    "image_span",
+    # All slots share one position (the first position after Q*). Makes the
+    # workspace order-free, matching the set semantics of Perceiver-style slots.
+    "shared",
+)
+
+
+class CloseQwen2_5_VLConfig(Qwen2_5_VLConfig):
+    r"""Configuration for :class:`CloseQwen2_5_VLForConditionalGeneration`.
+
+    Args:
+        ell_star: 0-indexed inclusive last layer of the lower branch. ``None``
+            until §3.1 localization has run; the model refuses to build with
+            ``None`` rather than falling back to a guessed depth.
+        num_workspace_steps: the fixed recurrent horizon (paper notation ``T``;
+            per-item segment counts ``n`` exist only in the losses, never here).
+            The trainable parameter count is independent of it because
+            ``f_theta`` is shared across steps, so it is a free experimental
+            axis. ``0`` is legal and is the no-recurrence floor: ``Z^(0)`` goes
+            straight to the decoder.
+        num_workspace_slots: ``S``, slots in ``Z``. Free capacity knob.
+        workspace_width: ``d_w``, the slot width. Free capacity knob -- the
+            method's claim is path exclusivity (all visual information routed
+            through ``Z``), NOT an information bottleneck, so ``S`` and ``d_w``
+            carry no methodological commitment. Kept below ``d`` (3584 for the
+            7B) since ``P_Z`` projects up at the splice.
+        workspace_num_heads: attention heads inside ``r_theta`` / ``f_theta``.
+        read_num_blocks: cross-attention blocks in ``r_theta``. The *read* is
+            once regardless of depth: all blocks see the same ``V*`` in a single
+            forward, and ``V*`` is discarded afterwards.
+        transition_num_blocks: blocks in the shared recurrent ``f_theta``.
+        workspace_ffn_mult: FFN expansion inside workspace blocks.
+        workspace_dropout: dropout inside workspace blocks.
+        workspace_rope_mode: one of :data:`WORKSPACE_ROPE_MODES`.
+        adapter_rank: rank of the low-rank adapter on the post-``ell_star``
+            backbone layers. ``0`` disables it (frozen backbone only).
+        adapter_alpha: LoRA-style scaling; effective scale is
+            ``adapter_alpha / adapter_rank``.
+        adapter_dropout: dropout probability applied only to the LoRA branch.
+        lambda_traj: legacy trajectory-loss weight. The current method sets it
+            to zero and trains with final-answer CE only; it remains serialized
+            solely so historical checkpoints can still be inspected.
+    """
+
+    model_type = "close_qwen2_5_vl"
+
+    def __init__(
+        self,
+        ell_star: int | None = None,
+        num_workspace_steps: int = 4,
+        num_workspace_slots: int = 32,
+        workspace_width: int = 512,
+        workspace_num_heads: int = 8,
+        read_num_blocks: int = 1,
+        transition_num_blocks: int = 2,
+        workspace_ffn_mult: int = 4,
+        workspace_dropout: float = 0.0,
+        workspace_rope_mode: str = "continue",
+        adapter_rank: int = 16,
+        adapter_alpha: int = 32,
+        adapter_dropout: float = 0.0,
+        adapter_exclude_layers: tuple = (),
+        lambda_traj: float = 0.0,
+        read_shortcut: bool = False,
+        splice_scale_match: bool = False,
+        read_gating: bool = False,
+        gate_init_eps: float = 0.02,
+        read_residual_pure: bool = False,
+        residual_norm_cap: float = 0.0,
+        pure_warm_start: bool = False,
+        final_verify: bool = False,
+        mid_read_step: int = 0,
+        persistent_slot_id: bool = False,
+        mid_read_alpha: float = 1.0,
+        evidence_reasoning: bool = False,
+        raw_evidence_question: bool = False,
+        raw_attention_source_layer: int = -1,
+        raw_transition_ffn: bool | None = None,
+        raw_state_replacement: bool = False,
+        raw_question_anchor: bool = False,
+        raw_mean_aggregation: bool = False,
+        hierarchical_visual_layers: tuple = (),
+        counterfactual_residual_recurrence: bool = False,
+        visual_counterfactual_recurrence: bool = False,
+        visual_cumulative_recurrence: bool = False,
+        visual_full_state_recurrence: bool = False,
+        visual_boundary_reentry: bool = False,
+        visual_source_centered_reentry: bool = False,
+        recurrence_only_adapter: bool = False,
+        visual_preserving_adapter_correction: bool = False,
+        perceive_deliberate_chain: bool = False,
+        local_visual_latent_chain: bool = False,
+        pdlc_transition_conditioned: bool = False,
+        pdlc_question_visible_through_pair: int = 1,
+        pdlc_question_curriculum: bool = False,
+        latent_chain_init_token_id: int = -1,
+        spatial_visual_recurrence: bool = False,
+        spatial_recurrence_steps: int = 8,
+        spatial_recurrence_inner_width: int = 256,
+        spatial_recurrence_heads: int = 8,
+        spatial_step_override: int = -1,
+        counterfactual_beta: float = 1.0,
+        crr_concat_aggregation: bool = False,
+        crr_policy_rank: int = 0,
+        crr_policy_alpha: float = 32.0,
+        counterfactual_reliance_loss: bool = False,
+        transition_aware_reliance_loss: bool = False,
+        corrective_transition_reliance_loss: bool = False,
+        counterfactual_step_retention_loss: bool = False,
+        functional_state_loss: bool = False,
+        per_example_answer_loss: bool = False,
+        interface_loss: bool = False,
+        read_anchor: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.ell_star = ell_star
+        self.num_workspace_steps = num_workspace_steps
+        self.num_workspace_slots = num_workspace_slots
+        self.workspace_width = workspace_width
+        self.workspace_num_heads = workspace_num_heads
+        self.read_num_blocks = read_num_blocks
+        self.transition_num_blocks = transition_num_blocks
+        self.workspace_ffn_mult = workspace_ffn_mult
+        self.workspace_dropout = workspace_dropout
+        self.workspace_rope_mode = workspace_rope_mode
+        self.adapter_rank = adapter_rank
+        self.adapter_alpha = adapter_alpha
+        self.adapter_dropout = adapter_dropout
+        # E run: dense full-FT 층은 LoRA 대상에서 제외 (이중 파라미터화 방지)
+        self.adapter_exclude_layers = list(adapter_exclude_layers or ())
+        self.lambda_traj = lambda_traj
+        # v2 anti-collapse switches (defaults False so v1 checkpoints reload).
+        # read_shortcut: r_theta adds an orthogonally-projected pooled-V* term to
+        # every slot, so Z^(0) is image-dependent from step 0 instead of relying
+        # on a random cross-attention finding signal.
+        # splice_scale_match: P_Z(Z) rows are RMS-matched to the question rows at
+        # the splice. Unmatched, slot states are ~O(1) against layer-20 hiddens
+        # in the hundreds -- the decoder could ignore them for free, and did
+        # (probe 20143: visual sensitivity 0.94 -> 1.0000 over training).
+        self.read_shortcut = read_shortcut
+        self.splice_scale_match = splice_scale_match
+        # v6 candidate: gated re-reads DURING the recurrence. Every re-read
+        # goes through r_theta into Z, so the mediation invariant (decoder sees
+        # only [Q*; P_Z(Z)]) is untouched; what changes is that the workspace
+        # may consult V* again mid-reasoning, scaled by a learned scalar gate.
+        # The gate is free at EVERY step -- no hand-coded schedule. Degeneration
+        # into answer-time retrieval is disciplined by the objective itself:
+        # L_traj pins each intermediate state to its teacher segment, so a
+        # wait-then-fetch trajectory misses its intermediate targets. The gate
+        # bias initialises negative, i.e. training starts from the proven
+        # no-reread regime and departs only where gradients demand it.
+        # Default False: v4/v5 checkpoints reload bit-identically.
+        self.read_gating = read_gating
+        # Initial gate open-rate epsilon_gamma: b_gamma = logit(eps). Sigmoid
+        # never reaches exactly 0 at finite logits, so this is near-identity
+        # (not bitwise) initialisation w.r.t. the read-once path.
+        self.gate_init_eps = gate_init_eps
+        # v7: always-on PURE visual residual rereads. Every step adds
+        # Z = U + W_o XAttn(LN(U) -> queries, V* -> keys/values): no self-attn,
+        # no FFN, no output norm, no biases in the branch, W_o zero-init.
+        # Motivation (battery, ep1 ckpt 20752): read_once hurt badly while
+        # null_step (blank memory at one step) was harmless -- the gated
+        # reader's value was extra recurrent DEPTH, not vision. This branch
+        # removes every path that can change Z without V* content, so any
+        # benefit it shows IS visual by construction.
+        self.read_residual_pure = read_residual_pure
+        # v7.1 (advisor): relative-norm trust region on the visual residual.
+        #   s_k = min(1, rho * RMS(U) / (RMS(Delta) + eps));  Z = U + s_k * Delta
+        # Unlike a fixed alpha, the optimizer cannot undo this by rescaling
+        # W_o (s_k adapts); ||Delta_used|| <~ rho * ||U|| structurally, and
+        # V* = 0 => Delta = 0 keeps the zero-memory identity. 0.0 disables.
+        self.residual_norm_cap = residual_norm_cap
+        # v7.1: initialise pure_attn K/V projections from the initial visual
+        # reader's v_in (same [d_w, d] shape) and Q from its first read
+        # block's in-proj slice, instead of random xavier -- the ep1 finding
+        # was that Q/K/V stayed ~random while only W_o trained.
+        self.pure_warm_start = pure_warm_start
+        # MAIN METHOD (advisor-confirmed): Final-Verified Recurrent Workspace.
+        # Initial read -> PURE latent recurrence (no mid-step visual anything)
+        # -> exactly ONE pure visual verification at the end:
+        #   Z_out = U^(T) + W_o XAttn(LN(U^(T)), V*),  alpha = 1, W_o = 0 init.
+        # Uses the same bias-free verifier module as read_residual_pure; the
+        # two flags differ only in WHERE the branch fires (every step vs once).
+        self.final_verify = final_verify
+        # MAIN (advisor 최종 확정): Mid-Read Recurrent Workspace. 고정 상수
+        # k* = floor(T/2) = 4 (스텝 sweep 산물이 아니라 판독 전후 동수 전이의
+        # midpoint)에서 pure visual read 1회:
+        #   Z^(4) = U^(4) + W_o XAttn(LN(U^(4)), V*),  alpha=1, W_o=0 init.
+        # 논리적 step 4 == python 루프 index 3 (0-based) -- off-by-one은
+        # tests/test_mid_read.py가 상태 분기 위치로 고정한다. 0이면 비활성.
+        self.mid_read_step = mid_read_step
+        # advisor(2026-08-06): E_slot(=workspace_read.slots)을 매 recurrent
+        # step 입력과 mid-read query에 일시 제공. state에는 누적하지 않는다.
+        self.persistent_slot_id = persistent_slot_id
+        # 고정 read 스케일 Z = U + alpha*W_o A (advisor scale-mismatch 검사).
+        # 1.0 = 기존과 비트동일. eval의 gate_overrides alpha는 절대값으로 이를 대체.
+        self.mid_read_alpha = mid_read_alpha
+        # advisor 최종안: E(persistent evidence) / R(recurrent reasoning) 분리.
+        # states = [E^(0), R^(1..T)], 결합 Z^(T)=E+R^(T)는 prefill에서.
+        self.evidence_reasoning = evidence_reasoning
+        # Current replacement candidate: no learned evidence/reasoning slots.
+        # E is the complete raw V* sequence and R0 is the text-only Q* sequence.
+        # A single shared CrossAttn transition produces R1..RT. Residual mode
+        # decodes RT; the replacement candidate decodes Q*+HT. Cross-attention
+        # uses the backbone's native GQA shape and
+        # is initialized from the first upper layer by default. Checkpoints
+        # created before the CrossAttn-only decision did not serialize
+        # raw_transition_ffn; None therefore preserves their historical FFN
+        # for faithful evaluation, while every new training launch passes False.
+        self.raw_evidence_question = bool(raw_evidence_question)
+        self.raw_attention_source_layer = int(raw_attention_source_layer)
+        self.raw_transition_ffn = (
+            True if raw_transition_ffn is None else bool(raw_transition_ffn)
+        )
+        # Experimental simplification of the same raw-E/Q path.  The residual
+        # carrier is replaced, not augmented:
+        #   H0=Q*, Hk=CrossAttn(H{k-1}, E), decoder_latent=Q*+HT.
+        # False preserves all residual raw-E/Q checkpoints and active runs.
+        self.raw_state_replacement = bool(raw_state_replacement)
+        # Question-anchored replacement keeps the fixed text scaffold in every
+        # recurrent query without restoring the old recurrent identity path:
+        #   Z0=Q*, Zk=Q*+CrossAttn(Z{k-1}, E), decoder_latent=ZT.
+        # This replaces the one-time Q*+HT decoder interface above.
+        self.raw_question_anchor = bool(raw_question_anchor)
+        # Q-shaped recurrent candidate selected by the V* interface oracle:
+        #   R0=Q*, Rk=R{k-1}+CrossAttn(R{k-1}, V*)
+        #   R_agg=mean(R1..RT)
+        # This is parameter-free and replaces both question-anchored state
+        # replacement and final-state-only decoding.
+        self.raw_mean_aggregation = bool(raw_mean_aggregation)
+        # A non-empty sequence maps the first T-1 recurrent steps to frozen
+        # visual-token states from the listed language layers.  The final step
+        # uses Q* as memory to integrate the staged reads without another image
+        # access.  One CrossAttention module is shared by all T steps.
+        self.hierarchical_visual_layers = [
+            int(layer) for layer in (hierarchical_visual_layers or ())
+        ]
+        # Counterfactual Residual Recurrence (CRR): the image-conditioned and
+        # text-only question states are propagated by the same native layer
+        # F=L(ell_star+1).  Only their counterfactual residual is recurrently
+        # refined; there are no slots, V* memory, custom attention modules, or
+        # inference-time auxiliary modules.  beta is the fixed relaxed-update
+        # coefficient; an optional batch-swap reliance loss is training-only.
+        self.counterfactual_residual_recurrence = bool(
+            counterfactual_residual_recurrence
+        )
+        # Causal Visual Residual Recurrence (CVRR), implemented as the
+        # vision-native CRR transition.  At every recurrent step the same
+        # native layer is evaluated with and without only Q->vision attention;
+        # their difference is the image-caused update.  It deliberately
+        # reuses CRR's strict Q-shaped decoder/cache interface and mean
+        # aggregation, hence the parent CRR flag remains explicit.
+        self.visual_counterfactual_recurrence = bool(
+            visual_counterfactual_recurrence
+        )
+        # Minimal cumulative CVRR variant.  Unlike the historical anchored
+        # mean path, visual updates are accumulated into the running state and
+        # the final state is decoded directly:
+        #   R_k = R_{k-1} + beta * Delta_k^V,  decode R_T.
+        # One explicit flag controls both choices so a checkpoint cannot
+        # silently combine cumulative dynamics with the old mean interface.
+        self.visual_cumulative_recurrence = bool(
+            visual_cumulative_recurrence
+        )
+        # Full-state CVRR keeps the native recurrent computation instead of
+        # treating the visual on/off intervention difference as the next
+        # hidden state:
+        #   R_tilde_k = Pi_Q F([E; R_{k-1}])
+        #   R_k = R_{k-1} + beta * (R_tilde_k - R_{k-1}).
+        # The visual-off branch remains available to evaluation code, but is
+        # not part of this model's train/inference transition.
+        self.visual_full_state_recurrence = bool(
+            visual_full_state_recurrence
+        )
+        # Input-grounded closure control.  Later recurrent calls rebuild the
+        # native post-L20 scaffold and inject only the carried question-state
+        # deviation, rather than feeding post-L21 rows back into L21.
+        self.visual_boundary_reentry = bool(visual_boundary_reentry)
+        # Source-centered boundary recurrence.  The text-only L21 state is a
+        # fixed anchor, while the native multimodal residual initializes a
+        # nonzero, image-dependent deviation.  The recurrent cell evolves
+        # only that deviation after subtracting its zero-deviation response.
+        self.visual_source_centered_reentry = bool(
+            visual_source_centered_reentry
+        )
+        # Restrict LoRA to the single shared recurrent cell L(ell_star+1).
+        # Its execution is disabled for R1, the text-only cache, answer-token
+        # decoding, and L22+; only recurrent transitions k>=2 enable it.
+        self.recurrence_only_adapter = bool(recurrence_only_adapter)
+        # Preserve the pretrained image-caused transition while removing the
+        # frozen layer's non-visual recurrent drift:
+        #   Delta_k = F_{base+LoRA,on}(E,R_{k-1}) - F_{base,off}(E,R_{k-1})
+        #           = (F_{base,on} - F_{base,off})
+        #             + (F_{base+LoRA,on} - F_{base,on}).
+        # This is a full-state recurrence subtype and introduces no new
+        # parameters or inference-time module.
+        self.visual_preserving_adapter_correction = bool(
+            visual_preserving_adapter_correction
+        )
+        # Local Visual Latent Chain (LVLC): K homogeneous latent-token
+        # positions are evaluated in one native Transformer forward.  The
+        # first latent sees clean Q plus visual source rows; every later latent
+        # sees only the immediately preceding latent plus the same visual
+        # rows; answer rows see clean Q plus the final latent.  This is a local
+        # attention graph, not a recurrent forward or temporal aggregation.
+        self.local_visual_latent_chain = bool(local_visual_latent_chain)
+        # Perceive--Deliberate Latent Chain (PDLC).  K must be even and is
+        # interpreted as K/2 native LOOK/THINK pairs.  The complete VLM stack
+        # runs once over a block-sparse graph: LOOK rows may inspect the
+        # multimodal prompt, THINK rows may inspect only clean question rows
+        # and their paired LOOK, and answer rows may inspect only clean
+        # question rows and THINK states.  There is no recurrent cell, visual
+        # state update, split-state splice, or temporal aggregator.
+        # LVLC reuses the same one-pass latent-sequence execution plumbing but
+        # replaces the LOOK/THINK token types and visibility graph entirely.
+        self.perceive_deliberate_chain = bool(
+            perceive_deliberate_chain or self.local_visual_latent_chain
+        )
+        # Tight PDLC graph: pair 1 bootstraps from prompt/question, later
+        # LOOK rows receive visual placeholders plus the preceding latent
+        # prefix, later THINK rows receive only that causal latent prefix, and
+        # the answer consumes final THINK.  It changes visibility only and
+        # allocates no additional trainable module.
+        self.pdlc_transition_conditioned = bool(
+            pdlc_transition_conditioned
+        )
+        # Strict inference uses 1: only the bootstrap pair reads the clean
+        # question.  Larger values are an explicit evaluation intervention.
+        # The training-only curriculum changes a runtime override and always
+        # returns to 1 for validation/inference.
+        self.pdlc_question_visible_through_pair = int(
+            pdlc_question_visible_through_pair
+        )
+        self.pdlc_question_curriculum = bool(pdlc_question_curriculum)
+        # Both learned latent token types start from one native text embedding
+        # (the launcher supplies the tokenizer's newline id).  -1 falls back
+        # to token 0 for synthetic/unit-test configs only.
+        self.latent_chain_init_token_id = int(latent_chain_init_token_id)
+        # Spatial Recurrent Visual Field (SRVF).  Unlike all historical
+        # ell_star methods, the recurrent state is the complete native visual
+        # token grid immediately after the vision merger.  The final field is
+        # inserted into the ordinary Qwen prefix and the full decoder runs from
+        # scratch, making the zero-output initialization exactly base-equivalent.
+        self.spatial_visual_recurrence = bool(spatial_visual_recurrence)
+        self.spatial_recurrence_steps = int(spatial_recurrence_steps)
+        self.spatial_recurrence_inner_width = int(
+            spatial_recurrence_inner_width
+        )
+        self.spatial_recurrence_heads = int(spatial_recurrence_heads)
+        # -1 uses the trained/default horizon.  Non-negative values are an
+        # evaluation-only prefix sweep, with 0 being the exact base path.
+        self.spatial_step_override = int(spatial_step_override)
+        self.counterfactual_beta = float(counterfactual_beta)
+        # Learn one feature-wise mixture of the complete recurrent residual
+        # trajectory.  The bias-free T*d -> d map starts from
+        # [I/T, ..., I/T], so enabling it is exactly the historical CRR mean
+        # interface before the first optimizer step.  A dedicated flag keeps
+        # old CRR checkpoints and the historical slot-ER aggregator distinct.
+        self.crr_concat_aggregation = bool(crr_concat_aggregation)
+        # Optional latent-only policy adapter.  It corrects the deterministic
+        # C2..CT means but never touches B, C1, or answer-token decoding.  Rank 0
+        # is the exact historical CRR path; RL checkpoints serialize rank > 0.
+        self.crr_policy_rank = int(crr_policy_rank)
+        self.crr_policy_alpha = float(crr_policy_alpha)
+        # Training-only CRR intervention.  A batch-matched residual from a
+        # different-answer example replaces the factual residual before the
+        # upper decoder.  The model returns the factual-minus-swapped answer
+        # score gap; the Trainer applies the margin loss.  Evaluation and
+        # generation never execute the extra branch.
+        self.counterfactual_reliance_loss = bool(
+            counterfactual_reliance_loss
+        )
+        # Replace the aggregate residual swap with one intervention inside the
+        # recurrent chain.  C_k is swapped and C_{k+1:T} is rolled out again,
+        # so the answer gap is attributable to a selected transition rather
+        # than to an undifferentiated full-trajectory replacement.
+        self.transition_aware_reliance_loss = bool(
+            transition_aware_reliance_loss
+        )
+        # Training-only hard-case corrective ranking.  At one selected
+        # transition k>=2, the actual decoder scores Z_{k-1}, Z_k, and a
+        # donor-swapped Z_k.  The Trainer requires the factual transition to
+        # improve over both detached comparators.  This replaces, rather than
+        # stacks with, the historical transition-aware reliance objective.
+        self.corrective_transition_reliance_loss = bool(
+            corrective_transition_reliance_loss
+        )
+        # Training-only functional no-regression objective.  Gold-answer NLL
+        # is measured through the actual upper decoder at every recurrent CRR
+        # prefix; inference remains the unchanged T-step mean interface.
+        self.counterfactual_step_retention_loss = bool(
+            counterfactual_step_retention_loss
+        )
+        # Training-only self-distillation target.  The multimodal lower pass
+        # already computed for V* also provides its image-conditioned question
+        # rows.  They supervise R_agg at the split boundary, but are never
+        # exposed to the inference decoder.
+        self.functional_state_loss = bool(functional_state_loss)
+        # Token CE lets long caption answers dominate one-token visual-choice
+        # examples.  The per-example variant first averages valid answer tokens
+        # within each sample, then averages samples.  It changes training/eval
+        # loss reduction only and has no inference-time effect.
+        self.per_example_answer_loss = bool(per_example_answer_loss)
+        # ER v2 (advisor 2026-08-11): explicit evidence read —
+        # R^(k) = f_theta(R^(k-1), Q*; E), E는 별도 memory로 cross-attn 읽기.
+        # (E+R 합산-감산 제거. evidence_reasoning=True 전제.)
+        self.explicit_evidence_read = bool(kwargs.pop("explicit_evidence_read", False))
+        # Historical field name retained for checkpoint compatibility. In ER
+        # v2.2 this means R^(0) is a batch-broadcast set of learnable slots; it
+        # is deliberately NOT initialized from Q*. Q* conditions every shared
+        # transition instead.
+        self.r_init_question = bool(kwargs.pop("r_init_question", False))
+        # ER v2.1 (advisor 2026-08-11): evidence read의 competitive normalization —
+        # slot 축 softmax 후 evidence 축 재정규화 (Slot Attention식 경쟁).
+        # 파라미터/모듈/loss 무추가; evidence_read의 W_qkv/W_o 재사용.
+        self.competitive_evidence_read = bool(
+            kwargs.pop("competitive_evidence_read", False))
+        # ER v2.2 C: k>=1 traj 감독 부재 시 twin 체인 제거 (states = answer 체인).
+        self.er_single_chain = bool(kwargs.pop("er_single_chain", False))
+        # ER v2.2 C (§5-7): decoder latent = softmax(temporal_logits)로 가중한
+        # R1..RT 결합 (전역 스칼라 T개). er_single_chain 전제. (legacy)
+        self.temporal_aggregation = bool(kwargs.pop("temporal_aggregation", False))
+        # ER v2.2 C 확정판: R_agg = Linear(Concat(R1..RT), bias=False),
+        # 초기 W=[I/T,...,I/T] (R_agg=mean(R1..RT)). er_single_chain 전제.
+        self.concat_aggregation = bool(kwargs.pop("concat_aggregation", False))
+        # ER v2: 워크스페이스 블록의 rank-r 부분공간 attention/FFN (0 = dense).
+        # d_w=d에서 state width 유지와 파라미터 폭증을 분리 (advisor).
+        self.workspace_low_rank = int(kwargs.pop("workspace_low_rank", 0))
+        # advisor: single-layer text-side interface alignment (bridge-only).
+        # teacher = frozen-base mm 경로 q_end@l*+1 (adapter off, sg);
+        # student = detach(R^T) splice의 l*+1 한 층 재계산. 추론 시 무존재.
+        self.interface_loss = interface_loss
+        # Training-only visual anchor switch. False in the current answer-only
+        # method; when false the model does not even materialize pooled V* for
+        # the Trainer.
+        self.read_anchor = bool(read_anchor)
+        if evidence_reasoning and persistent_slot_id:
+            raise ValueError("evidence_reasoning은 persistent_slot_id와 동시 사용 불가")
+        if evidence_reasoning and not mid_read_step:
+            raise ValueError("evidence_reasoning은 mid_read_step > 0 필요")
+        self.validate()
+
+    # -- derived views -----------------------------------------------------
+
+    @property
+    def num_decoder_layers(self) -> int:
+        return self.text_config.num_hidden_layers
+
+    @property
+    def backbone_width(self) -> int:
+        """``d`` -- backbone hidden width (3584 for the 7B)."""
+        return self.text_config.hidden_size
+
+    @property
+    def lower_slice(self) -> slice:
+        """Layers forming ``F_{<=l*}``."""
+        self._require_ell_star()
+        return slice(0, self.ell_star + 1)
+
+    @property
+    def upper_slice(self) -> slice:
+        """Layers forming ``F_{>l*}``."""
+        self._require_ell_star()
+        return slice(self.ell_star + 1, self.num_decoder_layers)
+
+    # -- validation --------------------------------------------------------
+
+    def _require_ell_star(self) -> None:
+        if self.ell_star is None:
+            raise ValueError(
+                "ell_star is unset. Run scripts/localize_read_point.py (Method 3.1) "
+                "and pass the measured layer explicitly; it must not be guessed."
+            )
+
+    def validate(self) -> None:
+        """Reject configurations that silently break the §3.2 contract."""
+        # ``PretrainedConfig.from_pretrained`` may apply user overrides with
+        # setattr *after* __init__.  In that path, setting LVLC=True would not
+        # re-run the constructor's promotion of the shared one-pass execution
+        # flag.  Re-establish the invariant before any model reads the config.
+        if getattr(self, "local_visual_latent_chain", False):
+            self.perceive_deliberate_chain = True
+        if self.ell_star is not None:
+            n = self.num_decoder_layers
+            # Both branches must be non-empty: an empty lower branch means there
+            # is no V* to read, an empty upper branch means the workspace never
+            # reaches the decoder.
+            if not 0 <= self.ell_star <= n - 2:
+                raise ValueError(
+                    f"ell_star={self.ell_star} out of range for {n} decoder layers; "
+                    f"expected 0 <= ell_star <= {n - 2} so both branches are non-empty."
+                )
+        # K=0 is the no-recurrence control (read only), not a misconfiguration.
+        if self.num_workspace_steps < 0:
+            raise ValueError("num_workspace_steps (K) must be >= 0.")
+        if self.num_workspace_slots < 1:
+            raise ValueError("num_workspace_slots (S) must be >= 1.")
+        if self.workspace_num_heads < 1:
+            raise ValueError("workspace_num_heads must be >= 1.")
+        if self.workspace_width % self.workspace_num_heads != 0:
+            raise ValueError(
+                f"workspace_width={self.workspace_width} must be divisible by "
+                f"workspace_num_heads={self.workspace_num_heads}."
+            )
+        if self.workspace_width > self.backbone_width:
+            raise ValueError(
+                f"workspace_width (d_w={self.workspace_width}) cannot exceed "
+                f"backbone width (d={self.backbone_width}). Native-width d_w=d is "
+                "the current method and uses an Identity decoder interface."
+            )
+        if self.workspace_rope_mode not in WORKSPACE_ROPE_MODES:
+            raise ValueError(
+                f"workspace_rope_mode={self.workspace_rope_mode!r} not in "
+                f"{WORKSPACE_ROPE_MODES}."
+            )
+        if self.adapter_rank < 0:
+            raise ValueError("adapter_rank must be >= 0 (0 disables the adapter).")
+        if not math.isfinite(self.adapter_dropout) or not 0.0 <= self.adapter_dropout < 1.0:
+            raise ValueError("adapter_dropout must be finite and in [0, 1).")
+        if getattr(self, "crr_policy_rank", 0) < 0:
+            raise ValueError("crr_policy_rank must be >= 0.")
+        if not math.isfinite(getattr(self, "crr_policy_alpha", 0.0)) or getattr(
+            self, "crr_policy_alpha", 0.0
+        ) <= 0.0:
+            raise ValueError("crr_policy_alpha must be finite and > 0.")
+        if self.read_num_blocks < 1 or self.transition_num_blocks < 1:
+            raise ValueError("read_num_blocks and transition_num_blocks must be >= 1.")
+        if self.workspace_ffn_mult < 1:
+            raise ValueError("workspace_ffn_mult must be >= 1.")
+        mid = int(getattr(self, "mid_read_step", 0) or 0)
+        if mid < 0 or mid > self.num_workspace_steps:
+            raise ValueError(
+                f"mid_read_step={mid} must be in [0, num_workspace_steps="
+                f"{self.num_workspace_steps}]."
+            )
+        low_rank = int(getattr(self, "workspace_low_rank", 0) or 0)
+        if low_rank < 0:
+            raise ValueError("workspace_low_rank must be >= 0.")
+        if low_rank and low_rank % self.workspace_num_heads != 0:
+            raise ValueError(
+                f"workspace_low_rank={low_rank} must be divisible by "
+                f"workspace_num_heads={self.workspace_num_heads}."
+            )
+        if getattr(self, "competitive_evidence_read", False) and not getattr(
+            self, "explicit_evidence_read", False
+        ):
+            raise ValueError(
+                "competitive_evidence_read requires explicit_evidence_read=True; "
+                "otherwise the flag has no execution path."
+            )
+        if getattr(self, "explicit_evidence_read", False) and not getattr(
+            self, "evidence_reasoning", False
+        ):
+            raise ValueError(
+                "explicit_evidence_read requires evidence_reasoning=True; otherwise "
+                "the evidence-read module is allocated but never called."
+            )
+        if getattr(self, "r_init_question", False) and not getattr(
+            self, "evidence_reasoning", False
+        ):
+            raise ValueError(
+                "r_init_question (legacy name for learnable R0 slots) requires "
+                "evidence_reasoning=True; otherwise reasoning_slots are unused."
+            )
+        if (getattr(self, "concat_aggregation", False)
+                or getattr(self, "temporal_aggregation", False)):
+            if self.num_workspace_steps < 1:
+                raise ValueError("reasoning aggregation requires num_workspace_steps >= 1.")
+            if not getattr(self, "evidence_reasoning", False):
+                raise ValueError(
+                    "reasoning aggregation requires evidence_reasoning=True; "
+                    "otherwise the aggregation module is unused."
+                )
+        if getattr(self, "concat_aggregation", False) and not getattr(
+            self, "er_single_chain", False
+        ):
+            raise ValueError("concat_aggregation requires er_single_chain=True.")
+        if getattr(self, "raw_evidence_question", False):
+            if self.workspace_width != self.backbone_width:
+                raise ValueError(
+                    "raw_evidence_question requires native workspace_width == "
+                    "backbone_width: E=V* and R0=Q* are used without projections."
+                )
+            if self.num_workspace_steps < 1:
+                raise ValueError(
+                    "raw_evidence_question requires num_workspace_steps >= 1."
+                )
+            source_layer = (
+                self.ell_star + 1
+                if self.raw_attention_source_layer < 0
+                else self.raw_attention_source_layer
+            )
+            if not self.ell_star < source_layer < self.text_config.num_hidden_layers:
+                raise ValueError(
+                    "raw_attention_source_layer must be an upper-backbone layer, "
+                    f"got ell_star={self.ell_star}, source={source_layer}."
+                )
+            incompatible = {
+                "evidence_reasoning": bool(getattr(self, "evidence_reasoning", False)),
+                "explicit_evidence_read": bool(getattr(self, "explicit_evidence_read", False)),
+                "competitive_evidence_read": bool(getattr(self, "competitive_evidence_read", False)),
+                "r_init_question": bool(getattr(self, "r_init_question", False)),
+                "er_single_chain": bool(getattr(self, "er_single_chain", False)),
+                "temporal_aggregation": bool(getattr(self, "temporal_aggregation", False)),
+                "concat_aggregation": bool(getattr(self, "concat_aggregation", False)),
+                "read_shortcut": bool(getattr(self, "read_shortcut", False)),
+                "splice_scale_match": bool(getattr(self, "splice_scale_match", False)),
+                "read_gating": bool(getattr(self, "read_gating", False)),
+                "read_residual_pure": bool(getattr(self, "read_residual_pure", False)),
+                "final_verify": bool(getattr(self, "final_verify", False)),
+                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
+                "persistent_slot_id": bool(getattr(self, "persistent_slot_id", False)),
+                "interface_loss": bool(getattr(self, "interface_loss", False)),
+                "read_anchor": bool(getattr(self, "read_anchor", False)),
+            }
+            active = [name for name, enabled in incompatible.items() if enabled]
+            if active:
+                raise ValueError(
+                    "raw_evidence_question replaces the slot workspace and is "
+                    f"incompatible with: {', '.join(active)}"
+                )
+            if self.raw_state_replacement and self.raw_transition_ffn:
+                raise ValueError(
+                    "raw_state_replacement is CrossAttn-only and is incompatible "
+                    "with raw_transition_ffn=True."
+                )
+            if self.raw_question_anchor and not self.raw_state_replacement:
+                raise ValueError(
+                    "raw_question_anchor requires raw_state_replacement=True."
+                )
+            if self.raw_mean_aggregation and (
+                self.raw_state_replacement or self.raw_question_anchor
+            ):
+                raise ValueError(
+                    "raw_mean_aggregation replaces raw_state_replacement and "
+                    "raw_question_anchor; use the residual recurrence."
+                )
+            hierarchy = list(
+                getattr(self, "hierarchical_visual_layers", ()) or ()
+            )
+            if hierarchy:
+                if not self.raw_mean_aggregation:
+                    raise ValueError(
+                        "hierarchical_visual_layers requires "
+                        "raw_mean_aggregation=True."
+                    )
+                if self.raw_transition_ffn:
+                    raise ValueError(
+                        "hierarchical visual recurrence is CrossAttn-only and "
+                        "requires raw_transition_ffn=False."
+                    )
+                expected_visual_reads = self.num_workspace_steps - 1
+                if len(hierarchy) != expected_visual_reads:
+                    raise ValueError(
+                        "hierarchical_visual_layers must contain exactly one "
+                        "layer per visual-read step, followed by one Q* "
+                        "integration step: "
+                        f"got {len(hierarchy)} layers for "
+                        f"T={self.num_workspace_steps} "
+                        f"(expected {expected_visual_reads})."
+                    )
+                if hierarchy != sorted(set(hierarchy)):
+                    raise ValueError(
+                        "hierarchical_visual_layers must be strictly increasing."
+                    )
+                invalid = [
+                    layer
+                    for layer in hierarchy
+                    if not 0 <= layer <= int(self.ell_star)
+                ]
+                if invalid:
+                    raise ValueError(
+                        "hierarchical_visual_layers must lie in the frozen lower "
+                        f"branch [0,{self.ell_star}], got {invalid}."
+                    )
+                if self.functional_state_loss:
+                    raise ValueError(
+                        "functional_state_loss is not defined for hierarchical "
+                        "visual recurrence; use final-answer CE only."
+                    )
+            if self.functional_state_loss and not self.raw_mean_aggregation:
+                raise ValueError(
+                    "functional_state_loss requires raw_mean_aggregation=True "
+                    "so its target is the exact decoder-facing recurrence."
+                )
+        elif getattr(self, "raw_state_replacement", False):
+            raise ValueError(
+                "raw_state_replacement requires raw_evidence_question=True."
+            )
+        elif getattr(self, "raw_question_anchor", False):
+            raise ValueError(
+                "raw_question_anchor requires raw_evidence_question=True."
+            )
+        elif getattr(self, "raw_mean_aggregation", False):
+            raise ValueError(
+                "raw_mean_aggregation requires raw_evidence_question=True."
+            )
+        elif getattr(self, "hierarchical_visual_layers", ()):
+            raise ValueError(
+                "hierarchical_visual_layers requires "
+                "raw_evidence_question=True."
+            )
+        elif getattr(self, "functional_state_loss", False):
+            raise ValueError(
+                "functional_state_loss requires raw_evidence_question=True."
+            )
+        if (
+            getattr(self, "counterfactual_reliance_loss", False)
+            and not getattr(self, "counterfactual_residual_recurrence", False)
+        ):
+            raise ValueError(
+                "counterfactual_reliance_loss requires "
+                "counterfactual_residual_recurrence=True."
+            )
+        if (
+            getattr(self, "transition_aware_reliance_loss", False)
+            and not getattr(self, "counterfactual_reliance_loss", False)
+        ):
+            raise ValueError(
+                "transition_aware_reliance_loss requires "
+                "counterfactual_reliance_loss=True."
+            )
+        if (
+            getattr(self, "transition_aware_reliance_loss", False)
+            and self.num_workspace_steps < 2
+        ):
+            raise ValueError(
+                "transition_aware_reliance_loss requires at least two steps."
+            )
+        if (
+            getattr(self, "corrective_transition_reliance_loss", False)
+            and not getattr(self, "counterfactual_reliance_loss", False)
+        ):
+            raise ValueError(
+                "corrective_transition_reliance_loss requires "
+                "counterfactual_reliance_loss=True."
+            )
+        if (
+            getattr(self, "corrective_transition_reliance_loss", False)
+            and self.num_workspace_steps < 2
+        ):
+            raise ValueError(
+                "corrective_transition_reliance_loss requires at least two steps."
+            )
+        if (
+            getattr(self, "corrective_transition_reliance_loss", False)
+            and getattr(self, "transition_aware_reliance_loss", False)
+        ):
+            raise ValueError(
+                "corrective_transition_reliance_loss replaces, and cannot be "
+                "combined with, transition_aware_reliance_loss."
+            )
+        if (
+            getattr(self, "corrective_transition_reliance_loss", False)
+            and getattr(self, "counterfactual_step_retention_loss", False)
+        ):
+            raise ValueError(
+                "corrective_transition_reliance_loss replaces, and cannot be "
+                "combined with, counterfactual_step_retention_loss."
+            )
+        if (
+            getattr(self, "counterfactual_step_retention_loss", False)
+            and not getattr(self, "counterfactual_residual_recurrence", False)
+        ):
+            raise ValueError(
+                "counterfactual_step_retention_loss requires "
+                "counterfactual_residual_recurrence=True."
+            )
+        if (
+            getattr(self, "counterfactual_step_retention_loss", False)
+            and self.num_workspace_steps < 2
+        ):
+            raise ValueError(
+                "counterfactual_step_retention_loss requires at least two steps."
+            )
+        if getattr(self, "crr_concat_aggregation", False) and not getattr(
+            self, "counterfactual_residual_recurrence", False
+        ):
+            raise ValueError(
+                "crr_concat_aggregation requires "
+                "counterfactual_residual_recurrence=True."
+            )
+        if getattr(self, "counterfactual_residual_recurrence", False):
+            if getattr(self, "raw_evidence_question", False):
+                raise ValueError(
+                    "counterfactual_residual_recurrence replaces "
+                    "raw_evidence_question; enable exactly one q-state method."
+                )
+            if self.num_workspace_steps < 1:
+                raise ValueError(
+                    "counterfactual_residual_recurrence requires "
+                    "num_workspace_steps >= 1."
+                )
+            if self.crr_policy_rank > 0 and self.num_workspace_steps < 2:
+                raise ValueError(
+                    "crr_policy_rank > 0 requires num_workspace_steps >= 2."
+                )
+            if not 0.0 <= self.counterfactual_beta <= 1.0:
+                raise ValueError(
+                    "counterfactual_beta must lie in [0, 1], got "
+                    f"{self.counterfactual_beta}."
+                )
+            incompatible = {
+                "evidence_reasoning": bool(getattr(self, "evidence_reasoning", False)),
+                "explicit_evidence_read": bool(getattr(self, "explicit_evidence_read", False)),
+                "competitive_evidence_read": bool(getattr(self, "competitive_evidence_read", False)),
+                "r_init_question": bool(getattr(self, "r_init_question", False)),
+                "er_single_chain": bool(getattr(self, "er_single_chain", False)),
+                "temporal_aggregation": bool(getattr(self, "temporal_aggregation", False)),
+                "concat_aggregation": bool(getattr(self, "concat_aggregation", False)),
+                "read_shortcut": bool(getattr(self, "read_shortcut", False)),
+                "splice_scale_match": bool(getattr(self, "splice_scale_match", False)),
+                "read_gating": bool(getattr(self, "read_gating", False)),
+                "read_residual_pure": bool(getattr(self, "read_residual_pure", False)),
+                "final_verify": bool(getattr(self, "final_verify", False)),
+                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
+                "persistent_slot_id": bool(getattr(self, "persistent_slot_id", False)),
+                "interface_loss": bool(getattr(self, "interface_loss", False)),
+                "read_anchor": bool(getattr(self, "read_anchor", False)),
+                "functional_state_loss": bool(getattr(self, "functional_state_loss", False)),
+            }
+            active = [name for name, enabled in incompatible.items() if enabled]
+            if active:
+                raise ValueError(
+                    "counterfactual_residual_recurrence replaces the slot/raw "
+                    f"workspace and is incompatible with: {', '.join(active)}"
+                )
+        elif getattr(self, "crr_policy_rank", 0) > 0:
+            raise ValueError(
+                "crr_policy_rank > 0 requires "
+                "counterfactual_residual_recurrence=True."
+            )
+        if getattr(self, "visual_counterfactual_recurrence", False):
+            if not getattr(
+                self, "counterfactual_residual_recurrence", False
+            ):
+                raise ValueError(
+                    "visual_counterfactual_recurrence requires "
+                    "counterfactual_residual_recurrence=True."
+                )
+            # CVRR is the intentionally minimal CE-only candidate.  The old
+            # CRR tail-rollout losses/policy implement a different transition
+            # equation and must not be silently applied to this trajectory.
+            incompatible = {
+                "counterfactual_reliance_loss": bool(
+                    getattr(self, "counterfactual_reliance_loss", False)
+                ),
+                "transition_aware_reliance_loss": bool(
+                    getattr(self, "transition_aware_reliance_loss", False)
+                ),
+                "corrective_transition_reliance_loss": bool(
+                    getattr(self, "corrective_transition_reliance_loss", False)
+                ),
+                "counterfactual_step_retention_loss": bool(
+                    getattr(self, "counterfactual_step_retention_loss", False)
+                ),
+                "crr_concat_aggregation": bool(
+                    getattr(self, "crr_concat_aggregation", False)
+                ),
+                "crr_policy_rank": bool(getattr(self, "crr_policy_rank", 0)),
+            }
+            active = [name for name, enabled in incompatible.items() if enabled]
+            if active:
+                raise ValueError(
+                    "visual_counterfactual_recurrence uses final-answer CE only "
+                    "and a parameter-free decoder interface; incompatible with: "
+                    + ", ".join(active)
+                )
+        elif getattr(self, "visual_cumulative_recurrence", False):
+            raise ValueError(
+                "visual_cumulative_recurrence requires "
+                "visual_counterfactual_recurrence=True."
+            )
+        if getattr(self, "visual_full_state_recurrence", False):
+            if not getattr(self, "visual_counterfactual_recurrence", False):
+                raise ValueError(
+                    "visual_full_state_recurrence requires "
+                    "visual_counterfactual_recurrence=True."
+                )
+            if not getattr(self, "visual_cumulative_recurrence", False):
+                raise ValueError(
+                    "visual_full_state_recurrence requires "
+                    "visual_cumulative_recurrence=True."
+                )
+        if getattr(self, "visual_boundary_reentry", False) and not getattr(
+            self, "visual_full_state_recurrence", False
+        ):
+            raise ValueError(
+                "visual_boundary_reentry requires "
+                "visual_full_state_recurrence=True."
+            )
+        if getattr(self, "visual_source_centered_reentry", False):
+            if not getattr(self, "visual_boundary_reentry", False):
+                raise ValueError(
+                    "visual_source_centered_reentry requires "
+                    "visual_boundary_reentry=True."
+                )
+            if getattr(self, "visual_preserving_adapter_correction", False):
+                raise ValueError(
+                    "visual_source_centered_reentry cannot be combined with "
+                    "visual_preserving_adapter_correction."
+                )
+        if getattr(self, "recurrence_only_adapter", False):
+            if not getattr(self, "visual_full_state_recurrence", False):
+                raise ValueError(
+                    "recurrence_only_adapter requires "
+                    "visual_full_state_recurrence=True."
+                )
+            if self.adapter_rank <= 0:
+                raise ValueError(
+                    "recurrence_only_adapter requires adapter_rank > 0."
+                )
+            if self.ell_star is None:
+                raise ValueError(
+                    "recurrence_only_adapter requires an explicit ell_star."
+                )
+            recurrent_layer = int(self.ell_star) + 1
+            if recurrent_layer in set(self.adapter_exclude_layers):
+                raise ValueError(
+                    "recurrence_only_adapter cannot exclude its recurrent "
+                    f"layer {recurrent_layer} from LoRA."
+                )
+        if getattr(self, "visual_preserving_adapter_correction", False):
+            if not getattr(self, "visual_full_state_recurrence", False):
+                raise ValueError(
+                    "visual_preserving_adapter_correction requires "
+                    "visual_full_state_recurrence=True."
+                )
+            if not getattr(self, "recurrence_only_adapter", False):
+                raise ValueError(
+                    "visual_preserving_adapter_correction requires "
+                    "recurrence_only_adapter=True."
+                )
+        if getattr(self, "perceive_deliberate_chain", False):
+            if self.num_workspace_steps < 2 or self.num_workspace_steps % 2:
+                raise ValueError(
+                    "one-pass latent chains require an even "
+                    "num_workspace_steps >= 2."
+                )
+            vocab_size = int(self.text_config.vocab_size)
+            if not -1 <= self.latent_chain_init_token_id < vocab_size:
+                raise ValueError(
+                    "latent_chain_init_token_id must be -1 or a valid text "
+                    f"token id below {vocab_size}."
+                )
+            num_pairs = self.num_workspace_steps // 2
+            if self.local_visual_latent_chain and (
+                self.pdlc_transition_conditioned
+                or self.pdlc_question_curriculum
+                or self.pdlc_question_visible_through_pair != 1
+            ):
+                raise ValueError(
+                    "local_visual_latent_chain replaces PDLC transition and "
+                    "question-visibility controls"
+                )
+            if not 1 <= self.pdlc_question_visible_through_pair <= num_pairs:
+                raise ValueError(
+                    "pdlc_question_visible_through_pair must lie in "
+                    f"[1,{num_pairs}]"
+                )
+            if self.pdlc_question_curriculum and not (
+                getattr(self, "pdlc_transition_conditioned", False)
+            ):
+                raise ValueError(
+                    "pdlc_question_curriculum requires "
+                    "pdlc_transition_conditioned=True"
+                )
+            incompatible = {
+                "raw_evidence_question": bool(
+                    getattr(self, "raw_evidence_question", False)
+                ),
+                "counterfactual_residual_recurrence": bool(
+                    getattr(self, "counterfactual_residual_recurrence", False)
+                ),
+                "spatial_visual_recurrence": bool(
+                    getattr(self, "spatial_visual_recurrence", False)
+                ),
+                "evidence_reasoning": bool(
+                    getattr(self, "evidence_reasoning", False)
+                ),
+                "read_gating": bool(getattr(self, "read_gating", False)),
+                "read_residual_pure": bool(
+                    getattr(self, "read_residual_pure", False)
+                ),
+                "final_verify": bool(getattr(self, "final_verify", False)),
+                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
+                "interface_loss": bool(getattr(self, "interface_loss", False)),
+                "read_anchor": bool(getattr(self, "read_anchor", False)),
+                "counterfactual_reliance_loss": bool(
+                    getattr(self, "counterfactual_reliance_loss", False)
+                ),
+                "counterfactual_step_retention_loss": bool(
+                    getattr(self, "counterfactual_step_retention_loss", False)
+                ),
+                "functional_state_loss": bool(
+                    getattr(self, "functional_state_loss", False)
+                ),
+            }
+            active = [name for name, enabled in incompatible.items() if enabled]
+            if active:
+                raise ValueError(
+                    "perceive_deliberate_chain replaces every recurrent/"
+                    "workspace path; incompatible with: " + ", ".join(active)
+                )
+        elif getattr(self, "pdlc_transition_conditioned", False):
+            raise ValueError(
+                "pdlc_transition_conditioned requires "
+                "perceive_deliberate_chain=True."
+            )
+        elif getattr(self, "pdlc_question_curriculum", False) or (
+            getattr(self, "pdlc_question_visible_through_pair", 1) != 1
+        ):
+            raise ValueError(
+                "PDLC question visibility controls require "
+                "perceive_deliberate_chain=True."
+            )
+        if getattr(self, "spatial_visual_recurrence", False):
+            if not 0.0 <= self.counterfactual_beta <= 1.0:
+                raise ValueError(
+                    "counterfactual_beta must lie in [0, 1], got "
+                    f"{self.counterfactual_beta}."
+                )
+            if self.spatial_recurrence_steps < 1:
+                raise ValueError(
+                    "spatial_recurrence_steps must be >= 1 for training"
+                )
+            if self.spatial_recurrence_inner_width < 1:
+                raise ValueError(
+                    "spatial_recurrence_inner_width must be positive"
+                )
+            if self.spatial_recurrence_heads < 1 or (
+                self.spatial_recurrence_inner_width
+                % self.spatial_recurrence_heads
+            ):
+                raise ValueError(
+                    "spatial_recurrence_inner_width must be divisible by "
+                    "spatial_recurrence_heads"
+                )
+            if self.spatial_step_override < -1:
+                raise ValueError("spatial_step_override must be >= -1")
+            if self.adapter_rank != 0:
+                raise ValueError(
+                    "spatial_visual_recurrence replaces decoder LoRA; set "
+                    "adapter_rank=0 so the T=0 path remains the frozen base model"
+                )
+            incompatible = {
+                "raw_evidence_question": bool(
+                    getattr(self, "raw_evidence_question", False)
+                ),
+                "counterfactual_residual_recurrence": bool(
+                    getattr(self, "counterfactual_residual_recurrence", False)
+                ),
+                "evidence_reasoning": bool(
+                    getattr(self, "evidence_reasoning", False)
+                ),
+                "read_gating": bool(getattr(self, "read_gating", False)),
+                "read_residual_pure": bool(
+                    getattr(self, "read_residual_pure", False)
+                ),
+                "final_verify": bool(getattr(self, "final_verify", False)),
+                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
+                "interface_loss": bool(getattr(self, "interface_loss", False)),
+                "read_anchor": bool(getattr(self, "read_anchor", False)),
+                "counterfactual_reliance_loss": bool(
+                    getattr(self, "counterfactual_reliance_loss", False)
+                ),
+                "counterfactual_step_retention_loss": bool(
+                    getattr(self, "counterfactual_step_retention_loss", False)
+                ),
+                "functional_state_loss": bool(
+                    getattr(self, "functional_state_loss", False)
+                ),
+            }
+            active = [name for name, enabled in incompatible.items() if enabled]
+            if active:
+                raise ValueError(
+                    "spatial_visual_recurrence replaces every historical "
+                    "workspace/reliance path; incompatible with: "
+                    + ", ".join(active)
+                )
+        modes = [bool(getattr(self, m, False)) for m in
+                 ("read_gating", "read_residual_pure", "final_verify")]
+        modes.append(bool(getattr(self, "mid_read_step", 0)))
+        if sum(modes) > 1:
+            raise ValueError(
+                "read_gating / read_residual_pure / final_verify are mutually "
+                "exclusive reread modes."
+            )
+
+
+__all__ = ["CloseQwen2_5_VLConfig", "WORKSPACE_ROPE_MODES"]

configuration_source_qwen3.py

@@ -0,0 +1,740 @@
+"""Configuration for CLOSE latent reasoning on Qwen3-VL.
+
+The current ``raw_evidence_question`` path keeps the complete layer-``ell_star``
+visual memory as immutable evidence and recurrently updates the text-only
+question state by cross-attending to it. Its selected interface averages all
+four residual recurrent states and may use a training-only restoration target
+from the normal multimodal question state. Historical slot-workspace fields are
+retained so old checkpoints remain loadable, but they are mutually exclusive
+with the raw-evidence path.
+
+Every architectural hyperparameter of Method §3 lives here -- nothing is
+hardcoded in ``modeling_close_qwen3_vl.py``.
+
+Layer-index convention (used consistently across the repo)
+----------------------------------------------------------
+``ell_star`` is **0-indexed and inclusive**: it is the index of the last decoder
+layer belonging to the lower branch. With ``num_hidden_layers == 28``::
+
+    F_{<=l*} = language_model.layers[: ell_star + 1]     # lower_slice
+    F_{>l*}  = language_model.layers[ell_star + 1 :]     # upper_slice
+
+so ``ell_star=18`` means layers 0..18 run before the split and 19..27 after it.
+``ell_star`` is measured once by ``scripts/localize_read_point.py`` (§3.1) and
+frozen before any workspace training.
+
+Careful with ``Qwen3VLConfig``: its ``__setattr__`` forwards any attribute
+whose name already exists in ``text_config.__dict__`` down to the sub-config, and
+its ``__init__`` builds ``text_config`` from ``**kwargs``. Workspace fields are
+therefore declared as explicit named parameters so they never enter ``kwargs``
+and never shadow a text-config field. ``tests/test_config_roundtrip.py`` guards
+this.
+"""
+
+from __future__ import annotations
+
+import math
+
+from transformers.models.qwen3_vl.configuration_qwen3_vl import Qwen3VLConfig
+
+#: How the ``S`` workspace slots are assigned M-RoPE positions once they are
+#: spliced into the replacement cache above ``ell_star``. Qwen3-VL uses 3D
+#: M-RoPE (t, h, w) with ``mrope_section=[16, 24, 24]``, so slots -- which have
+#: no spatial extent -- need an explicit convention.
+WORKSPACE_ROPE_MODES = (
+    # Slots continue the 1-D text position sequence directly after Q*, with
+    # t == h == w. Treats the workspace as "more text".
+    "continue",
+    # Slots reuse the (t, h, w) span the image tokens occupied in the original
+    # multimodal prefill, subsampled to S positions. Preserves the read point's
+    # positional geometry but reintroduces image-derived indices.
+    "image_span",
+    # All slots share one position (the first position after Q*). Makes the
+    # workspace order-free, matching the set semantics of Perceiver-style slots.
+    "shared",
+)
+
+
+class CloseQwen3VLConfig(Qwen3VLConfig):
+    r"""Configuration for :class:`CloseQwen3VLForConditionalGeneration`.
+
+    Args:
+        ell_star: 0-indexed inclusive last layer of the lower branch. ``None``
+            until §3.1 localization has run; the model refuses to build with
+            ``None`` rather than falling back to a guessed depth.
+        num_workspace_steps: the fixed recurrent horizon (paper notation ``T``;
+            per-item segment counts ``n`` exist only in the losses, never here).
+            The trainable parameter count is independent of it because
+            ``f_theta`` is shared across steps, so it is a free experimental
+            axis. ``0`` is legal and is the no-recurrence floor: ``Z^(0)`` goes
+            straight to the decoder.
+        num_workspace_slots: ``S``, slots in ``Z``. Free capacity knob.
+        workspace_width: ``d_w``, the slot width. Free capacity knob -- the
+            method's claim is path exclusivity (all visual information routed
+            through ``Z``), NOT an information bottleneck, so ``S`` and ``d_w``
+            carry no methodological commitment. Kept below ``d`` (4096 for the
+            7B) since ``P_Z`` projects up at the splice.
+        workspace_num_heads: attention heads inside ``r_theta`` / ``f_theta``.
+        read_num_blocks: cross-attention blocks in ``r_theta``. The *read* is
+            once regardless of depth: all blocks see the same ``V*`` in a single
+            forward, and ``V*`` is discarded afterwards.
+        transition_num_blocks: blocks in the shared recurrent ``f_theta``.
+        workspace_ffn_mult: FFN expansion inside workspace blocks.
+        workspace_dropout: dropout inside workspace blocks.
+        workspace_rope_mode: one of :data:`WORKSPACE_ROPE_MODES`.
+        adapter_rank: rank of the low-rank adapter on the post-``ell_star``
+            backbone layers. ``0`` disables it (frozen backbone only).
+        adapter_alpha: LoRA-style scaling; effective scale is
+            ``adapter_alpha / adapter_rank``.
+        adapter_dropout: dropout probability applied only to the LoRA branch.
+        lambda_traj: legacy trajectory-loss weight. The current method sets it
+            to zero and trains with final-answer CE only; it remains serialized
+            solely so historical checkpoints can still be inspected.
+    """
+
+    model_type = "close_qwen3_vl"
+
+    def __init__(
+        self,
+        ell_star: int | None = None,
+        num_workspace_steps: int = 4,
+        num_workspace_slots: int = 32,
+        workspace_width: int = 512,
+        workspace_num_heads: int = 8,
+        read_num_blocks: int = 1,
+        transition_num_blocks: int = 2,
+        workspace_ffn_mult: int = 4,
+        workspace_dropout: float = 0.0,
+        workspace_rope_mode: str = "continue",
+        adapter_rank: int = 16,
+        adapter_alpha: int = 32,
+        adapter_dropout: float = 0.0,
+        adapter_exclude_layers: tuple = (),
+        lambda_traj: float = 0.0,
+        read_shortcut: bool = False,
+        splice_scale_match: bool = False,
+        read_gating: bool = False,
+        gate_init_eps: float = 0.02,
+        read_residual_pure: bool = False,
+        residual_norm_cap: float = 0.0,
+        pure_warm_start: bool = False,
+        final_verify: bool = False,
+        mid_read_step: int = 0,
+        persistent_slot_id: bool = False,
+        mid_read_alpha: float = 1.0,
+        evidence_reasoning: bool = False,
+        raw_evidence_question: bool = False,
+        raw_attention_source_layer: int = -1,
+        raw_transition_ffn: bool | None = None,
+        raw_state_replacement: bool = False,
+        raw_question_anchor: bool = False,
+        raw_mean_aggregation: bool = False,
+        counterfactual_residual_recurrence: bool = False,
+        visual_counterfactual_recurrence: bool = False,
+        visual_cumulative_recurrence: bool = False,
+        visual_full_state_recurrence: bool = False,
+        recurrence_only_adapter: bool = False,
+        counterfactual_beta: float = 1.0,
+        crr_concat_aggregation: bool = False,
+        crr_policy_rank: int = 0,
+        crr_policy_alpha: float = 32.0,
+        counterfactual_reliance_loss: bool = False,
+        transition_aware_reliance_loss: bool = False,
+        counterfactual_step_retention_loss: bool = False,
+        functional_state_loss: bool = False,
+        per_example_answer_loss: bool = False,
+        interface_loss: bool = False,
+        read_anchor: bool = False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.ell_star = ell_star
+        self.num_workspace_steps = num_workspace_steps
+        self.num_workspace_slots = num_workspace_slots
+        self.workspace_width = workspace_width
+        self.workspace_num_heads = workspace_num_heads
+        self.read_num_blocks = read_num_blocks
+        self.transition_num_blocks = transition_num_blocks
+        self.workspace_ffn_mult = workspace_ffn_mult
+        self.workspace_dropout = workspace_dropout
+        self.workspace_rope_mode = workspace_rope_mode
+        self.adapter_rank = adapter_rank
+        self.adapter_alpha = adapter_alpha
+        self.adapter_dropout = adapter_dropout
+        # E run: dense full-FT 층은 LoRA 대상에서 제외 (이중 파라미터화 방지)
+        self.adapter_exclude_layers = list(adapter_exclude_layers or ())
+        self.lambda_traj = lambda_traj
+        # v2 anti-collapse switches (defaults False so v1 checkpoints reload).
+        # read_shortcut: r_theta adds an orthogonally-projected pooled-V* term to
+        # every slot, so Z^(0) is image-dependent from step 0 instead of relying
+        # on a random cross-attention finding signal.
+        # splice_scale_match: P_Z(Z) rows are RMS-matched to the question rows at
+        # the splice. Unmatched, slot states are ~O(1) against layer-20 hiddens
+        # in the hundreds -- the decoder could ignore them for free, and did
+        # (probe 20143: visual sensitivity 0.94 -> 1.0000 over training).
+        self.read_shortcut = read_shortcut
+        self.splice_scale_match = splice_scale_match
+        # v6 candidate: gated re-reads DURING the recurrence. Every re-read
+        # goes through r_theta into Z, so the mediation invariant (decoder sees
+        # only [Q*; P_Z(Z)]) is untouched; what changes is that the workspace
+        # may consult V* again mid-reasoning, scaled by a learned scalar gate.
+        # The gate is free at EVERY step -- no hand-coded schedule. Degeneration
+        # into answer-time retrieval is disciplined by the objective itself:
+        # L_traj pins each intermediate state to its teacher segment, so a
+        # wait-then-fetch trajectory misses its intermediate targets. The gate
+        # bias initialises negative, i.e. training starts from the proven
+        # no-reread regime and departs only where gradients demand it.
+        # Default False: v4/v5 checkpoints reload bit-identically.
+        self.read_gating = read_gating
+        # Initial gate open-rate epsilon_gamma: b_gamma = logit(eps). Sigmoid
+        # never reaches exactly 0 at finite logits, so this is near-identity
+        # (not bitwise) initialisation w.r.t. the read-once path.
+        self.gate_init_eps = gate_init_eps
+        # v7: always-on PURE visual residual rereads. Every step adds
+        # Z = U + W_o XAttn(LN(U) -> queries, V* -> keys/values): no self-attn,
+        # no FFN, no output norm, no biases in the branch, W_o zero-init.
+        # Motivation (battery, ep1 ckpt 20752): read_once hurt badly while
+        # null_step (blank memory at one step) was harmless -- the gated
+        # reader's value was extra recurrent DEPTH, not vision. This branch
+        # removes every path that can change Z without V* content, so any
+        # benefit it shows IS visual by construction.
+        self.read_residual_pure = read_residual_pure
+        # v7.1 (advisor): relative-norm trust region on the visual residual.
+        #   s_k = min(1, rho * RMS(U) / (RMS(Delta) + eps));  Z = U + s_k * Delta
+        # Unlike a fixed alpha, the optimizer cannot undo this by rescaling
+        # W_o (s_k adapts); ||Delta_used|| <~ rho * ||U|| structurally, and
+        # V* = 0 => Delta = 0 keeps the zero-memory identity. 0.0 disables.
+        self.residual_norm_cap = residual_norm_cap
+        # v7.1: initialise pure_attn K/V projections from the initial visual
+        # reader's v_in (same [d_w, d] shape) and Q from its first read
+        # block's in-proj slice, instead of random xavier -- the ep1 finding
+        # was that Q/K/V stayed ~random while only W_o trained.
+        self.pure_warm_start = pure_warm_start
+        # MAIN METHOD (advisor-confirmed): Final-Verified Recurrent Workspace.
+        # Initial read -> PURE latent recurrence (no mid-step visual anything)
+        # -> exactly ONE pure visual verification at the end:
+        #   Z_out = U^(T) + W_o XAttn(LN(U^(T)), V*),  alpha = 1, W_o = 0 init.
+        # Uses the same bias-free verifier module as read_residual_pure; the
+        # two flags differ only in WHERE the branch fires (every step vs once).
+        self.final_verify = final_verify
+        # MAIN (advisor 최종 확정): Mid-Read Recurrent Workspace. 고정 상수
+        # k* = floor(T/2) = 4 (스텝 sweep 산물이 아니라 판독 전후 동수 전이의
+        # midpoint)에서 pure visual read 1회:
+        #   Z^(4) = U^(4) + W_o XAttn(LN(U^(4)), V*),  alpha=1, W_o=0 init.
+        # 논리적 step 4 == python 루프 index 3 (0-based) -- off-by-one은
+        # tests/test_mid_read.py가 상태 분기 위치로 고정한다. 0이면 비활성.
+        self.mid_read_step = mid_read_step
+        # advisor(2026-08-06): E_slot(=workspace_read.slots)을 매 recurrent
+        # step 입력과 mid-read query에 일시 제공. state에는 누적하지 않는다.
+        self.persistent_slot_id = persistent_slot_id
+        # 고정 read 스케일 Z = U + alpha*W_o A (advisor scale-mismatch 검사).
+        # 1.0 = 기존과 비트동일. eval의 gate_overrides alpha는 절대값으로 이를 대체.
+        self.mid_read_alpha = mid_read_alpha
+        # advisor 최종안: E(persistent evidence) / R(recurrent reasoning) 분리.
+        # states = [E^(0), R^(1..T)], 결합 Z^(T)=E+R^(T)는 prefill에서.
+        self.evidence_reasoning = evidence_reasoning
+        # Current replacement candidate: no learned evidence/reasoning slots.
+        # E is the complete raw V* sequence and R0 is the text-only Q* sequence.
+        # A single shared CrossAttn transition produces R1..RT. Residual mode
+        # decodes RT; the replacement candidate decodes Q*+HT. Cross-attention
+        # uses the backbone's native GQA shape and
+        # is initialized from the first upper layer by default. Checkpoints
+        # created before the CrossAttn-only decision did not serialize
+        # raw_transition_ffn; None therefore preserves their historical FFN
+        # for faithful evaluation, while every new training launch passes False.
+        self.raw_evidence_question = bool(raw_evidence_question)
+        self.raw_attention_source_layer = int(raw_attention_source_layer)
+        self.raw_transition_ffn = (
+            True if raw_transition_ffn is None else bool(raw_transition_ffn)
+        )
+        # Experimental simplification of the same raw-E/Q path.  The residual
+        # carrier is replaced, not augmented:
+        #   H0=Q*, Hk=CrossAttn(H{k-1}, E), decoder_latent=Q*+HT.
+        # False preserves all residual raw-E/Q checkpoints and active runs.
+        self.raw_state_replacement = bool(raw_state_replacement)
+        # Question-anchored replacement keeps the fixed text scaffold in every
+        # recurrent query without restoring the old recurrent identity path:
+        #   Z0=Q*, Zk=Q*+CrossAttn(Z{k-1}, E), decoder_latent=ZT.
+        # This replaces the one-time Q*+HT decoder interface above.
+        self.raw_question_anchor = bool(raw_question_anchor)
+        # Q-shaped recurrent candidate selected by the V* interface oracle:
+        #   R0=Q*, Rk=R{k-1}+CrossAttn(R{k-1}, V*)
+        #   R_agg=mean(R1..RT)
+        # This is parameter-free and replaces both question-anchored state
+        # replacement and final-state-only decoding.
+        self.raw_mean_aggregation = bool(raw_mean_aggregation)
+        # Counterfactual Residual Recurrence (CRR): the image-conditioned and
+        # text-only question states are propagated by the same native layer
+        # F=L(ell_star+1).  Only their counterfactual residual is recurrently
+        # refined; there are no slots, V* memory, custom attention modules, or
+        # inference-time auxiliary modules.  beta is the fixed relaxed-update
+        # coefficient; an optional batch-swap reliance loss is training-only.
+        self.counterfactual_residual_recurrence = bool(
+            counterfactual_residual_recurrence
+        )
+        # Qwen3 counterpart of the promoted Qwen2.5 full-state CRR path.  The
+        # multimodal scaffold is used only inside the shared recurrent cell;
+        # answer decoding still receives a question-shaped state and a
+        # text-only prefix cache.
+        self.visual_counterfactual_recurrence = bool(
+            visual_counterfactual_recurrence
+        )
+        self.visual_cumulative_recurrence = bool(
+            visual_cumulative_recurrence
+        )
+        self.visual_full_state_recurrence = bool(
+            visual_full_state_recurrence
+        )
+        self.recurrence_only_adapter = bool(recurrence_only_adapter)
+        self.counterfactual_beta = float(counterfactual_beta)
+        # Bias-free Concat(C1..CT) -> C_agg, initialized to the exact
+        # historical mean interface for checkpoint-compatible warm starts.
+        self.crr_concat_aggregation = bool(crr_concat_aggregation)
+        # Optional latent-only policy adapter.  It corrects the deterministic
+        # C2..CT means but never touches B, C1, or answer-token decoding.  Rank 0
+        # is the exact historical CRR path; RL checkpoints serialize rank > 0.
+        self.crr_policy_rank = int(crr_policy_rank)
+        self.crr_policy_alpha = float(crr_policy_alpha)
+        # Training-only CRR intervention.  A batch-matched residual from a
+        # different-answer example replaces the factual residual before the
+        # upper decoder.  The model returns the factual-minus-swapped answer
+        # score gap; the Trainer applies the margin loss.  Evaluation and
+        # generation never execute the extra branch.
+        self.counterfactual_reliance_loss = bool(
+            counterfactual_reliance_loss
+        )
+        self.transition_aware_reliance_loss = bool(
+            transition_aware_reliance_loss
+        )
+        # Training-only functional no-regression objective.  Gold-answer NLL
+        # is measured through the actual upper decoder at every recurrent CRR
+        # prefix; inference remains the unchanged T-step mean interface.
+        self.counterfactual_step_retention_loss = bool(
+            counterfactual_step_retention_loss
+        )
+        # Training-only self-distillation target.  The multimodal lower pass
+        # already computed for V* also provides its image-conditioned question
+        # rows.  They supervise R_agg at the split boundary, but are never
+        # exposed to the inference decoder.
+        self.functional_state_loss = bool(functional_state_loss)
+        # Token CE lets long caption answers dominate one-token visual-choice
+        # examples.  The per-example variant first averages valid answer tokens
+        # within each sample, then averages samples.  It changes training/eval
+        # loss reduction only and has no inference-time effect.
+        self.per_example_answer_loss = bool(per_example_answer_loss)
+        # ER v2 (advisor 2026-08-11): explicit evidence read —
+        # R^(k) = f_theta(R^(k-1), Q*; E), E는 별도 memory로 cross-attn 읽기.
+        # (E+R 합산-감산 제거. evidence_reasoning=True 전제.)
+        self.explicit_evidence_read = bool(kwargs.pop("explicit_evidence_read", False))
+        # Historical field name retained for checkpoint compatibility. In ER
+        # v2.2 this means R^(0) is a batch-broadcast set of learnable slots; it
+        # is deliberately NOT initialized from Q*. Q* conditions every shared
+        # transition instead.
+        self.r_init_question = bool(kwargs.pop("r_init_question", False))
+        # ER v2.1 (advisor 2026-08-11): evidence read의 competitive normalization —
+        # slot 축 softmax 후 evidence 축 재정규화 (Slot Attention식 경쟁).
+        # 파라미터/모듈/loss 무추가; evidence_read의 W_qkv/W_o 재사용.
+        self.competitive_evidence_read = bool(
+            kwargs.pop("competitive_evidence_read", False))
+        # ER v2.2 C: k>=1 traj 감독 부재 시 twin 체인 제거 (states = answer 체인).
+        self.er_single_chain = bool(kwargs.pop("er_single_chain", False))
+        # ER v2.2 C (§5-7): decoder latent = softmax(temporal_logits)로 가중한
+        # R1..RT 결합 (전역 스칼라 T개). er_single_chain 전제. (legacy)
+        self.temporal_aggregation = bool(kwargs.pop("temporal_aggregation", False))
+        # ER v2.2 C 확정판: R_agg = Linear(Concat(R1..RT), bias=False),
+        # 초기 W=[I/T,...,I/T] (R_agg=mean(R1..RT)). er_single_chain 전제.
+        self.concat_aggregation = bool(kwargs.pop("concat_aggregation", False))
+        # ER v2: 워크스페이스 블록의 rank-r 부분공간 attention/FFN (0 = dense).
+        # d_w=d에서 state width 유지와 파라미터 폭증을 분리 (advisor).
+        self.workspace_low_rank = int(kwargs.pop("workspace_low_rank", 0))
+        # advisor: single-layer text-side interface alignment (bridge-only).
+        # teacher = frozen-base mm 경로 q_end@l*+1 (adapter off, sg);
+        # student = detach(R^T) splice의 l*+1 한 층 재계산. 추론 시 무존재.
+        self.interface_loss = interface_loss
+        # Training-only visual anchor switch. False in the current answer-only
+        # method; when false the model does not even materialize pooled V* for
+        # the Trainer.
+        self.read_anchor = bool(read_anchor)
+        if evidence_reasoning and persistent_slot_id:
+            raise ValueError("evidence_reasoning은 persistent_slot_id와 동시 사용 불가")
+        if evidence_reasoning and not mid_read_step:
+            raise ValueError("evidence_reasoning은 mid_read_step > 0 필요")
+        self.validate()
+
+    # -- derived views -----------------------------------------------------
+
+    @property
+    def num_decoder_layers(self) -> int:
+        return self.text_config.num_hidden_layers
+
+    @property
+    def backbone_width(self) -> int:
+        """``d`` -- backbone hidden width (4096 for the 7B)."""
+        return self.text_config.hidden_size
+
+    @property
+    def lower_slice(self) -> slice:
+        """Layers forming ``F_{<=l*}``."""
+        self._require_ell_star()
+        return slice(0, self.ell_star + 1)
+
+    @property
+    def upper_slice(self) -> slice:
+        """Layers forming ``F_{>l*}``."""
+        self._require_ell_star()
+        return slice(self.ell_star + 1, self.num_decoder_layers)
+
+    # -- validation --------------------------------------------------------
+
+    def _require_ell_star(self) -> None:
+        if self.ell_star is None:
+            raise ValueError(
+                "ell_star is unset. Run scripts/localize_read_point.py (Method 3.1) "
+                "and pass the measured layer explicitly; it must not be guessed."
+            )
+
+    def validate(self) -> None:
+        """Reject configurations that silently break the §3.2 contract."""
+        if self.ell_star is not None:
+            n = self.num_decoder_layers
+            # Both branches must be non-empty: an empty lower branch means there
+            # is no V* to read, an empty upper branch means the workspace never
+            # reaches the decoder.
+            if not 0 <= self.ell_star <= n - 2:
+                raise ValueError(
+                    f"ell_star={self.ell_star} out of range for {n} decoder layers; "
+                    f"expected 0 <= ell_star <= {n - 2} so both branches are non-empty."
+                )
+        # K=0 is the no-recurrence control (read only), not a misconfiguration.
+        if self.num_workspace_steps < 0:
+            raise ValueError("num_workspace_steps (K) must be >= 0.")
+        if self.num_workspace_slots < 1:
+            raise ValueError("num_workspace_slots (S) must be >= 1.")
+        if self.workspace_num_heads < 1:
+            raise ValueError("workspace_num_heads must be >= 1.")
+        if self.workspace_width % self.workspace_num_heads != 0:
+            raise ValueError(
+                f"workspace_width={self.workspace_width} must be divisible by "
+                f"workspace_num_heads={self.workspace_num_heads}."
+            )
+        if self.workspace_width > self.backbone_width:
+            raise ValueError(
+                f"workspace_width (d_w={self.workspace_width}) cannot exceed "
+                f"backbone width (d={self.backbone_width}). Native-width d_w=d is "
+                "the current method and uses an Identity decoder interface."
+            )
+        if self.workspace_rope_mode not in WORKSPACE_ROPE_MODES:
+            raise ValueError(
+                f"workspace_rope_mode={self.workspace_rope_mode!r} not in "
+                f"{WORKSPACE_ROPE_MODES}."
+            )
+        if self.adapter_rank < 0:
+            raise ValueError("adapter_rank must be >= 0 (0 disables the adapter).")
+        if not math.isfinite(self.adapter_dropout) or not 0.0 <= self.adapter_dropout < 1.0:
+            raise ValueError("adapter_dropout must be finite and in [0, 1).")
+        if getattr(self, "crr_policy_rank", 0) < 0:
+            raise ValueError("crr_policy_rank must be >= 0.")
+        if not math.isfinite(getattr(self, "crr_policy_alpha", 0.0)) or getattr(
+            self, "crr_policy_alpha", 0.0
+        ) <= 0.0:
+            raise ValueError("crr_policy_alpha must be finite and > 0.")
+        if self.read_num_blocks < 1 or self.transition_num_blocks < 1:
+            raise ValueError("read_num_blocks and transition_num_blocks must be >= 1.")
+        if self.workspace_ffn_mult < 1:
+            raise ValueError("workspace_ffn_mult must be >= 1.")
+        mid = int(getattr(self, "mid_read_step", 0) or 0)
+        if mid < 0 or mid > self.num_workspace_steps:
+            raise ValueError(
+                f"mid_read_step={mid} must be in [0, num_workspace_steps="
+                f"{self.num_workspace_steps}]."
+            )
+        low_rank = int(getattr(self, "workspace_low_rank", 0) or 0)
+        if low_rank < 0:
+            raise ValueError("workspace_low_rank must be >= 0.")
+        if low_rank and low_rank % self.workspace_num_heads != 0:
+            raise ValueError(
+                f"workspace_low_rank={low_rank} must be divisible by "
+                f"workspace_num_heads={self.workspace_num_heads}."
+            )
+        if getattr(self, "competitive_evidence_read", False) and not getattr(
+            self, "explicit_evidence_read", False
+        ):
+            raise ValueError(
+                "competitive_evidence_read requires explicit_evidence_read=True; "
+                "otherwise the flag has no execution path."
+            )
+        if getattr(self, "explicit_evidence_read", False) and not getattr(
+            self, "evidence_reasoning", False
+        ):
+            raise ValueError(
+                "explicit_evidence_read requires evidence_reasoning=True; otherwise "
+                "the evidence-read module is allocated but never called."
+            )
+        if getattr(self, "r_init_question", False) and not getattr(
+            self, "evidence_reasoning", False
+        ):
+            raise ValueError(
+                "r_init_question (legacy name for learnable R0 slots) requires "
+                "evidence_reasoning=True; otherwise reasoning_slots are unused."
+            )
+        if (getattr(self, "concat_aggregation", False)
+                or getattr(self, "temporal_aggregation", False)):
+            if self.num_workspace_steps < 1:
+                raise ValueError("reasoning aggregation requires num_workspace_steps >= 1.")
+            if not getattr(self, "evidence_reasoning", False):
+                raise ValueError(
+                    "reasoning aggregation requires evidence_reasoning=True; "
+                    "otherwise the aggregation module is unused."
+                )
+        if getattr(self, "concat_aggregation", False) and not getattr(
+            self, "er_single_chain", False
+        ):
+            raise ValueError("concat_aggregation requires er_single_chain=True.")
+        if getattr(self, "raw_evidence_question", False):
+            if self.workspace_width != self.backbone_width:
+                raise ValueError(
+                    "raw_evidence_question requires native workspace_width == "
+                    "backbone_width: E=V* and R0=Q* are used without projections."
+                )
+            if self.num_workspace_steps < 1:
+                raise ValueError(
+                    "raw_evidence_question requires num_workspace_steps >= 1."
+                )
+            source_layer = (
+                self.ell_star + 1
+                if self.raw_attention_source_layer < 0
+                else self.raw_attention_source_layer
+            )
+            if not self.ell_star < source_layer < self.text_config.num_hidden_layers:
+                raise ValueError(
+                    "raw_attention_source_layer must be an upper-backbone layer, "
+                    f"got ell_star={self.ell_star}, source={source_layer}."
+                )
+            incompatible = {
+                "evidence_reasoning": bool(getattr(self, "evidence_reasoning", False)),
+                "explicit_evidence_read": bool(getattr(self, "explicit_evidence_read", False)),
+                "competitive_evidence_read": bool(getattr(self, "competitive_evidence_read", False)),
+                "r_init_question": bool(getattr(self, "r_init_question", False)),
+                "er_single_chain": bool(getattr(self, "er_single_chain", False)),
+                "temporal_aggregation": bool(getattr(self, "temporal_aggregation", False)),
+                "concat_aggregation": bool(getattr(self, "concat_aggregation", False)),
+                "read_shortcut": bool(getattr(self, "read_shortcut", False)),
+                "splice_scale_match": bool(getattr(self, "splice_scale_match", False)),
+                "read_gating": bool(getattr(self, "read_gating", False)),
+                "read_residual_pure": bool(getattr(self, "read_residual_pure", False)),
+                "final_verify": bool(getattr(self, "final_verify", False)),
+                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
+                "persistent_slot_id": bool(getattr(self, "persistent_slot_id", False)),
+                "interface_loss": bool(getattr(self, "interface_loss", False)),
+                "read_anchor": bool(getattr(self, "read_anchor", False)),
+            }
+            active = [name for name, enabled in incompatible.items() if enabled]
+            if active:
+                raise ValueError(
+                    "raw_evidence_question replaces the slot workspace and is "
+                    f"incompatible with: {', '.join(active)}"
+                )
+            if self.raw_state_replacement and self.raw_transition_ffn:
+                raise ValueError(
+                    "raw_state_replacement is CrossAttn-only and is incompatible "
+                    "with raw_transition_ffn=True."
+                )
+            if self.raw_question_anchor and not self.raw_state_replacement:
+                raise ValueError(
+                    "raw_question_anchor requires raw_state_replacement=True."
+                )
+            if self.raw_mean_aggregation and (
+                self.raw_state_replacement or self.raw_question_anchor
+            ):
+                raise ValueError(
+                    "raw_mean_aggregation replaces raw_state_replacement and "
+                    "raw_question_anchor; use the residual recurrence."
+                )
+            if self.functional_state_loss and not self.raw_mean_aggregation:
+                raise ValueError(
+                    "functional_state_loss requires raw_mean_aggregation=True "
+                    "so its target is the exact decoder-facing recurrence."
+                )
+        elif getattr(self, "raw_state_replacement", False):
+            raise ValueError(
+                "raw_state_replacement requires raw_evidence_question=True."
+            )
+        elif getattr(self, "raw_question_anchor", False):
+            raise ValueError(
+                "raw_question_anchor requires raw_evidence_question=True."
+            )
+        elif getattr(self, "raw_mean_aggregation", False):
+            raise ValueError(
+                "raw_mean_aggregation requires raw_evidence_question=True."
+            )
+        elif getattr(self, "functional_state_loss", False):
+            raise ValueError(
+                "functional_state_loss requires raw_evidence_question=True."
+            )
+        if (
+            getattr(self, "counterfactual_reliance_loss", False)
+            and not getattr(self, "counterfactual_residual_recurrence", False)
+        ):
+            raise ValueError(
+                "counterfactual_reliance_loss requires "
+                "counterfactual_residual_recurrence=True."
+            )
+        if (
+            getattr(self, "transition_aware_reliance_loss", False)
+            and not getattr(self, "counterfactual_reliance_loss", False)
+        ):
+            raise ValueError(
+                "transition_aware_reliance_loss requires "
+                "counterfactual_reliance_loss=True."
+            )
+        if (
+            getattr(self, "transition_aware_reliance_loss", False)
+            and self.num_workspace_steps < 2
+        ):
+            raise ValueError(
+                "transition_aware_reliance_loss requires at least two steps."
+            )
+        if (
+            getattr(self, "counterfactual_step_retention_loss", False)
+            and not getattr(self, "counterfactual_residual_recurrence", False)
+        ):
+            raise ValueError(
+                "counterfactual_step_retention_loss requires "
+                "counterfactual_residual_recurrence=True."
+            )
+        if (
+            getattr(self, "counterfactual_step_retention_loss", False)
+            and self.num_workspace_steps < 2
+        ):
+            raise ValueError(
+                "counterfactual_step_retention_loss requires at least two steps."
+            )
+        if getattr(self, "crr_concat_aggregation", False) and not getattr(
+            self, "counterfactual_residual_recurrence", False
+        ):
+            raise ValueError(
+                "crr_concat_aggregation requires "
+                "counterfactual_residual_recurrence=True."
+            )
+        if getattr(self, "counterfactual_residual_recurrence", False):
+            if getattr(self, "raw_evidence_question", False):
+                raise ValueError(
+                    "counterfactual_residual_recurrence replaces "
+                    "raw_evidence_question; enable exactly one q-state method."
+                )
+            if self.num_workspace_steps < 1:
+                raise ValueError(
+                    "counterfactual_residual_recurrence requires "
+                    "num_workspace_steps >= 1."
+                )
+            if self.crr_policy_rank > 0 and self.num_workspace_steps < 2:
+                raise ValueError(
+                    "crr_policy_rank > 0 requires num_workspace_steps >= 2."
+                )
+            if not 0.0 <= self.counterfactual_beta <= 1.0:
+                raise ValueError(
+                    "counterfactual_beta must lie in [0, 1], got "
+                    f"{self.counterfactual_beta}."
+                )
+            incompatible = {
+                "evidence_reasoning": bool(getattr(self, "evidence_reasoning", False)),
+                "explicit_evidence_read": bool(getattr(self, "explicit_evidence_read", False)),
+                "competitive_evidence_read": bool(getattr(self, "competitive_evidence_read", False)),
+                "r_init_question": bool(getattr(self, "r_init_question", False)),
+                "er_single_chain": bool(getattr(self, "er_single_chain", False)),
+                "temporal_aggregation": bool(getattr(self, "temporal_aggregation", False)),
+                "concat_aggregation": bool(getattr(self, "concat_aggregation", False)),
+                "read_shortcut": bool(getattr(self, "read_shortcut", False)),
+                "splice_scale_match": bool(getattr(self, "splice_scale_match", False)),
+                "read_gating": bool(getattr(self, "read_gating", False)),
+                "read_residual_pure": bool(getattr(self, "read_residual_pure", False)),
+                "final_verify": bool(getattr(self, "final_verify", False)),
+                "mid_read_step": bool(getattr(self, "mid_read_step", 0)),
+                "persistent_slot_id": bool(getattr(self, "persistent_slot_id", False)),
+                "interface_loss": bool(getattr(self, "interface_loss", False)),
+                "read_anchor": bool(getattr(self, "read_anchor", False)),
+                "functional_state_loss": bool(getattr(self, "functional_state_loss", False)),
+            }
+            active = [name for name, enabled in incompatible.items() if enabled]
+            if active:
+                raise ValueError(
+                    "counterfactual_residual_recurrence replaces the slot/raw "
+                    f"workspace and is incompatible with: {', '.join(active)}"
+                )
+        elif getattr(self, "crr_policy_rank", 0) > 0:
+            raise ValueError(
+                "crr_policy_rank > 0 requires "
+                "counterfactual_residual_recurrence=True."
+            )
+        if getattr(self, "visual_counterfactual_recurrence", False):
+            if not getattr(self, "counterfactual_residual_recurrence", False):
+                raise ValueError(
+                    "visual_counterfactual_recurrence requires "
+                    "counterfactual_residual_recurrence=True."
+                )
+            if not getattr(self, "visual_cumulative_recurrence", False):
+                raise ValueError(
+                    "Qwen3 visual recurrence currently requires the promoted "
+                    "cumulative full-state path."
+                )
+            if not getattr(self, "visual_full_state_recurrence", False):
+                raise ValueError(
+                    "Qwen3 visual recurrence currently requires "
+                    "visual_full_state_recurrence=True."
+                )
+            incompatible = {
+                "counterfactual_reliance_loss": bool(
+                    getattr(self, "counterfactual_reliance_loss", False)
+                ),
+                "transition_aware_reliance_loss": bool(
+                    getattr(self, "transition_aware_reliance_loss", False)
+                ),
+                "counterfactual_step_retention_loss": bool(
+                    getattr(self, "counterfactual_step_retention_loss", False)
+                ),
+                "crr_concat_aggregation": bool(
+                    getattr(self, "crr_concat_aggregation", False)
+                ),
+                "crr_policy_rank": bool(getattr(self, "crr_policy_rank", 0)),
+            }
+            active = [name for name, enabled in incompatible.items() if enabled]
+            if active:
+                raise ValueError(
+                    "Qwen3 full-state visual recurrence uses final-answer CE "
+                    "and final-state decoding only; incompatible with: "
+                    + ", ".join(active)
+                )
+        elif getattr(self, "visual_cumulative_recurrence", False) or getattr(
+            self, "visual_full_state_recurrence", False
+        ):
+            raise ValueError(
+                "visual cumulative/full-state flags require "
+                "visual_counterfactual_recurrence=True."
+            )
+        if getattr(self, "recurrence_only_adapter", False):
+            if not getattr(self, "visual_full_state_recurrence", False):
+                raise ValueError(
+                    "recurrence_only_adapter requires "
+                    "visual_full_state_recurrence=True."
+                )
+            if self.adapter_rank <= 0:
+                raise ValueError(
+                    "recurrence_only_adapter requires adapter_rank > 0."
+                )
+            recurrent_layer = int(self.ell_star) + 1
+            if recurrent_layer in set(self.adapter_exclude_layers):
+                raise ValueError(
+                    "recurrence_only_adapter cannot exclude its recurrent "
+                    f"layer {recurrent_layer} from LoRA."
+                )
+        modes = [bool(getattr(self, m, False)) for m in
+                 ("read_gating", "read_residual_pure", "final_verify")]
+        modes.append(bool(getattr(self, "mid_read_step", 0)))
+        if sum(modes) > 1:
+            raise ValueError(
+                "read_gating / read_residual_pure / final_verify are mutually "
+                "exclusive reread modes."
+            )
+
+
+__all__ = ["CloseQwen3VLConfig", "WORKSPACE_ROPE_MODES"]

cvrr_release_config.json

@@ -0,0 +1,19 @@
+{
+  "name": "CVRR-InternVL3-9B",
+  "base_model": "OpenGVLab/InternVL3-9B",
+  "ell_star": 35,
+  "recurrent_layer": 36,
+  "upper_decoder_start": 37,
+  "inference_T": 4,
+  "inference_beta": 0.33,
+  "lora_rank": 32,
+  "lora_alpha": 12.0,
+  "lora_dropout": 0.01,
+  "checkpoint_step": 500,
+  "format": "cvrr_native_plus_merged_transition_v1",
+  "merged_weight_dtype": "float32",
+  "native_weight_bytes": 18277586944,
+  "status": "gpu_vstar_comparison_completed",
+  "upload_ready": true,
+  "hub_model_id": "dmis-lab/InternVL3-9B-CVRR"
+}

licenses/Apache-2.0.txt

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

licenses/InternVL-MIT.txt

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 OpenGVLab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

licenses/UPSTREAM_MODEL_CARD.md

@@ -0,0 +1,701 @@
+---
+license: mit
+pipeline_tag: image-text-to-text
+library_name: transformers
+base_model:
+- OpenGVLab/InternVL3-9B-Instruct
+base_model_relation: finetune
+datasets:
+- OpenGVLab/MMPR-v1.2
+language:
+- multilingual
+tags:
+- internvl
+- custom_code
+---
+
+# InternVL3-9B
+
+[\[📂 GitHub\]](https://github.com/OpenGVLab/InternVL)  [\[📜 InternVL 1.0\]](https://huggingface.co/papers/2312.14238)  [\[📜 InternVL 1.5\]](https://huggingface.co/papers/2404.16821)  [\[📜 InternVL 2.5\]](https://huggingface.co/papers/2412.05271)  [\[📜 InternVL2.5-MPO\]](https://huggingface.co/papers/2411.10442)  [\[📜 InternVL3\]](https://huggingface.co/papers/2504.10479)
+
+[\[🆕 Blog\]](https://internvl.github.io/blog/)  [\[🗨️ Chat Demo\]](https://internvl.opengvlab.com/)  [\[🤗 HF Demo\]](https://huggingface.co/spaces/OpenGVLab/InternVL)  [\[🚀 Quick Start\]](#quick-start)  [\[📖 Documents\]](https://internvl.readthedocs.io/en/latest/)
+
+<div align="center">
+  <img width="500" alt="image" src="https://cdn-uploads.huggingface.co/production/uploads/64006c09330a45b03605bba3/zJsd2hqd3EevgXo6fNgC-.png">
+</div>
+
+## Introduction
+
+We introduce InternVL3, an advanced multimodal large language model (MLLM) series that demonstrates superior overall performance.
+Compared to InternVL 2.5, InternVL3 exhibits superior multimodal perception and reasoning capabilities, while further extending its multimodal capabilities to encompass tool usage, GUI agents, industrial image analysis, 3D vision perception, and more.
+Additionally, we compare InternVL3 with  Qwen2.5 Chat models, whose corresponding pre-trained base models are employed as the initialization of the langauge component in InternVL3. Benefitting from Native Multimodal Pre-Training, the InternVL3 series achieves even better overall text performance than the Qwen2.5 series.
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/overall.png)
+
+## InternVL3 Family
+
+In the following table, we provide an overview of the InternVL3 series.
+
+|  Model Name   |                                       Vision Part                                       |                                 Language Part                                  |                         HF Link                          |
+| :-----------: | :-------------------------------------------------------------------------------------: | :----------------------------------------------------------------------------: | :------------------------------------------------------: |
+| InternVL3-1B  | [InternViT-300M-448px-V2_5](https://huggingface.co/OpenGVLab/InternViT-300M-448px-V2_5) |            [Qwen2.5-0.5B](https://huggingface.co/Qwen/Qwen2.5-0.5B)            | [🤗 link](https://huggingface.co/OpenGVLab/InternVL3-1B)  |
+| InternVL3-2B  | [InternViT-300M-448px-V2_5](https://huggingface.co/OpenGVLab/InternViT-300M-448px-V2_5) |            [Qwen2.5-1.5B](https://huggingface.co/Qwen/Qwen2.5-1.5B)            | [🤗 link](https://huggingface.co/OpenGVLab/InternVL3-2B)  |
+| InternVL3-8B  | [InternViT-300M-448px-V2_5](https://huggingface.co/OpenGVLab/InternViT-300M-448px-V2_5) |              [Qwen2.5-7B](https://huggingface.co/Qwen/Qwen2.5-7B)              | [🤗 link](https://huggingface.co/OpenGVLab/InternVL3-8B)  |
+| InternVL3-9B  | [InternViT-300M-448px-V2_5](https://huggingface.co/OpenGVLab/InternViT-300M-448px-V2_5) | [internlm3-8b-instruct](https://huggingface.co/internlm/internlm3-8b-instruct) | [🤗 link](https://huggingface.co/OpenGVLab/InternVL3-9B)  |
+| InternVL3-14B | [InternViT-300M-448px-V2_5](https://huggingface.co/OpenGVLab/InternViT-300M-448px-V2_5) |             [Qwen2.5-14B](https://huggingface.co/Qwen/Qwen2.5-14B)             | [🤗 link](https://huggingface.co/OpenGVLab/InternVL3-14B) |
+| InternVL3-38B |   [InternViT-6B-448px-V2_5](https://huggingface.co/OpenGVLab/InternViT-6B-448px-V2_5)   |             [Qwen2.5-32B](https://huggingface.co/Qwen/Qwen2.5-32B)             | [🤗 link](https://huggingface.co/OpenGVLab/InternVL3-38B) |
+| InternVL3-78B |   [InternViT-6B-448px-V2_5](https://huggingface.co/OpenGVLab/InternViT-6B-448px-V2_5)   |             [Qwen2.5-72B](https://huggingface.co/Qwen/Qwen2.5-72B)             | [🤗 link](https://huggingface.co/OpenGVLab/InternVL3-78B) |
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/overall-table.png)
+
+## Model Architecture
+
+As shown in the following figure, [InternVL3](https://internvl.github.io/blog/2025-04-11-InternVL-3/) retains the same model architecture as [InternVL 2.5](https://internvl.github.io/blog/2024-12-05-InternVL-2.5/) and its predecessors, InternVL 1.5 and 2.0, following the "ViT-MLP-LLM" paradigm. In this new version, we integrate a newly incrementally pre-trained InternViT with various pre-trained LLMs, including InternLM 3 and Qwen 2.5, using a randomly initialized MLP projector.
+
+
+![image/png](https://cdn-uploads.huggingface.co/production/uploads/64119264f0f81eb569e0d569/BiiyXN6NOk0p-3rl3ueyL.png)
+
+As in the previous version, we applied a pixel unshuffle operation, reducing the number of visual tokens to one-quarter of the original. Besides, we adopted a similar dynamic resolution strategy as InternVL 1.5, dividing images into tiles of 448×448 pixels. The key difference, starting from InternVL 2.0, is that we additionally introduced support for multi-image and video data.
+
+Notably, in InternVL3, we integrate the [Variable Visual Position Encoding (V2PE)](https://arxiv.org/abs/2412.09616), which utilizes smaller, more flexible position increments for visual tokens. Benefiting from V2PE, InternVL3 exhibits better long context understanding capabilities compared to its predecessors.
+
+## Training Strategy
+
+### Native Multimodal Pre-Training
+
+We propose a [Native Multimodal Pre-Training](https://huggingface.co/papers/2504.10479) approach that consolidates language and vision learning into a single pre-training stage.
+In contrast to standard paradigms that first train a language-only model and subsequently adapt it to handle additional modalities, our method interleaves multimodal data (e.g., image-text, video-text, or image-text interleaved sequences) with large-scale textual corpora. This unified training scheme allows the model to learn both linguistic and multimodal representations simultaneously, ultimately enhancing its capability to handle vision-language tasks without the need for separate alignment or bridging modules.
+Please see [our paper](https://huggingface.co/papers/2504.10479) for more details.
+
+### Supervised Fine-Tuning
+
+In this phase, the techniques of random JPEG compression, square loss re-weighting, and multimodal data packing proposed in [InternVL2.5](https://arxiv.org/abs/2412.05271) are also employed in the InternVL3 series.
+The main advancement of the SFT phase in InternVL3 compared to InternVL2.5 lies in the use of higher-quality and more diverse training data.
+Specifically, we further extend  training samples for tool use, 3D scene understanding, GUI operations, long context tasks, video understanding, scientific diagrams, creative writing, and multimodal reasoning.
+
+### Mixed Preference Optimization
+
+During Pre-training and SFT, the model is trained to predict the next token conditioned on previous ground-truth tokens.
+However, during inference, the model predicts each token based on its own prior outputs. 
+This discrepancy between ground-truth tokens and model-predicted tokens introduces a distribution shift, which can impair the model’s Chain-of-Thought (CoT) reasoning capabilities.
+To mitigate this issue, we employ [MPO](https://arxiv.org/abs/2411.10442), which introduces additional supervision from both positive and negative samples to align the model response distribution with the ground-truth distribution, thereby improving reasoning performance.
+Specifically, the training objective of MPO is a combination of
+preference loss \\(\mathcal{L}_{\text{p}}\\),
+quality loss \\(\mathcal{L}_{\text{q}}\\),
+and generation loss \\(\mathcal{L}_{\text{g}}\\),
+which can be formulated as follows:
+
+
+$$
+\mathcal{L}=w_{p}\cdot\mathcal{L}_{\text{p}} + w_{q}\cdot\mathcal{L}_{\text{q}} + w_{g}\cdot\mathcal{L}_{\text{g}},
+$$
+
+
+where \\(w_{*}\\) represents the weight assigned to each loss component. Please see [our paper](https://arxiv.org/abs/2411.10442) for more details about MPO.
+
+
+### Test-Time Scaling
+
+Test-Time Scaling has been shown to be an effective method to enhance the reasoning abilities of LLMs and MLLMs.
+In this work, we use the Best-of-N evaluation strategy and employ [VisualPRM-8B](https://huggingface.co/OpenGVLab/VisualPRM-8B) as the critic model to select the best response for reasoning and mathematics evaluation.
+
+## Evaluation on Multimodal Capability
+
+### Multimodal Reasoning and Mathematics
+
+![image/png](https://huggingface.co/OpenGVLab/VisualPRM-8B-v1_1/resolve/main/visualprm-performance.png)
+
+### OCR, Chart, and Document Understanding
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/ocr.png)
+
+### Multi-Image & Real-World Comprehension
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/multi-images.png)
+
+### Comprehensive Multimodal & Hallucination Evaluation
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/comprehensive.png)
+
+### Visual Grounding
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/grounding.png)
+
+### Multimodal Multilingual Understanding
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/multilingual.png)
+
+### Video Understanding
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/video.png)
+
+### GUI Grounding
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/gui.png)
+
+### Spatial Reasoning
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/vsi.png)
+
+## Evaluation on Language Capability
+
+We compare InternVL3 with  Qwen2.5 Chat models, whose corresponding pre-trained base models are employed as the initialization of the langauge component in InternVL3.
+Benefitting from Native Multimodal Pre-Training, the InternVL3 series achieves even better overall text performance than the Qwen2.5 series.
+Please note that the evaluation scores of Qwen2.5 series  may differ from those officially reported, as we have adopted the prompt versions provided in the table across all datasets for OpenCompass evaluation.
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/text.png)
+
+## Ablation Study
+
+### Native Multimodal Pre-Training
+
+We conduct experiments on the InternVL2-8B model while keeping its architecture, initialization parameters, and training data entirely unchanged. Traditionally, InternVL2-8B employs a training pipeline that begins with an MLP warmup phase for feature alignment followed by an Instruction Tuning stage. In our experiments, we substitute the conventional MLP warmup phase with a native multimodal pre-training process. This modification isolates the contribution of native multimodal pre-training to the overall multimodal capability of the model.
+
+The evaluation results in the Figure below shows that the model with native multimodal pre-training exhibits performance on most benchmarks that is comparable to the fully multi-stage-trained InternVL2-8B baseline. Furthermore, when followed by instruction tuning on higher-quality data, the model demonstrates further performance gains across evaluated multimodal tasks. These findings underscore the efficiency of native multimodal pre-training in imparting powerful multimodal capabilities to MLLMs.
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/ablation-native.png)
+
+### Mixed Preference Optimization
+
+As shown in the table below, models fine-tuned with MPO demonstrate superior reasoning performance across seven multimodal reasoning benchmarks compared to their counterparts without MPO. Specifically, InternVL3-78B and InternVL3-38B outperform their counterparts by 4.1 and 4.5 points, respectively. Notably, the training data used for MPO is a subset of that used for SFT, indicating that the performance improvements primarily stem from the training algorithm rather than the training data.
+
+![image/png](https://huggingface.co/datasets/OpenGVLab/MMPR-v1.2-prompts/resolve/main/ablation-mpo.png)
+
+### Variable Visual Position Encoding
+
+As reported in the table below, the introduction of V2PE leads to significant performance gains across most evaluation metrics. In addition, our ablation studies—by varying the positional increment \\( \delta \\)—reveal that even for tasks primarily involving conventional contexts, relatively small \\( \delta \\) values can achieve optimal performance. These findings provide important insights for future efforts aimed at refining position encoding strategies for visual tokens in MLLMs.
+
+![image/png](https://huggingface.co/datasets/Weiyun1025/InternVL-Performance/resolve/main/internvl3/ablation-v2pe.png)
+
+## Quick Start
+
+We provide an example code to run `InternVL3-9B` using `transformers`.
+
+> Please use transformers>=4.37.2 to ensure the model works normally.
+
+### Model Loading
+
+#### 16-bit (bf16 / fp16)
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+path = "OpenGVLab/InternVL3-9B"
+model = AutoModel.from_pretrained(
+    path,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
+    use_flash_attn=True,
+    trust_remote_code=True).eval().cuda()
+```
+
+#### BNB 8-bit Quantization
+
+```python
+import torch
+from transformers import AutoTokenizer, AutoModel
+path = "OpenGVLab/InternVL3-9B"
+model = AutoModel.from_pretrained(
+    path,
+    torch_dtype=torch.bfloat16,
+    load_in_8bit=True,
+    low_cpu_mem_usage=True,
+    use_flash_attn=True,
+    trust_remote_code=True).eval()
+```
+
+#### Multiple GPUs
+
+The reason for writing the code this way is to avoid errors that occur during multi-GPU inference due to tensors not being on the same device. By ensuring that the first and last layers of the large language model (LLM) are on the same device, we prevent such errors.
+
+```python
+import math
+import torch
+from transformers import AutoTokenizer, AutoModel
+
+def split_model(model_name):
+    device_map = {}
+    world_size = torch.cuda.device_count()
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    num_layers = config.llm_config.num_hidden_layers
+    # Since the first GPU will be used for ViT, treat it as half a GPU.
+    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
+    num_layers_per_gpu = [num_layers_per_gpu] * world_size
+    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
+    layer_cnt = 0
+    for i, num_layer in enumerate(num_layers_per_gpu):
+        for j in range(num_layer):
+            device_map[f'language_model.model.layers.{layer_cnt}'] = i
+            layer_cnt += 1
+    device_map['vision_model'] = 0
+    device_map['mlp1'] = 0
+    device_map['language_model.model.tok_embeddings'] = 0
+    device_map['language_model.model.embed_tokens'] = 0
+    device_map['language_model.output'] = 0
+    device_map['language_model.model.norm'] = 0
+    device_map['language_model.model.rotary_emb'] = 0
+    device_map['language_model.lm_head'] = 0
+    device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
+
+    return device_map
+
+path = "OpenGVLab/InternVL3-9B"
+device_map = split_model('InternVL3-9B')
+model = AutoModel.from_pretrained(
+    path,
+    torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=True,
+    use_flash_attn=True,
+    trust_remote_code=True,
+    device_map=device_map).eval()
+```
+
+### Inference with Transformers
+
+```python
+import math
+import numpy as np
+import torch
+import torchvision.transforms as T
+from decord import VideoReader, cpu
+from PIL import Image
+from torchvision.transforms.functional import InterpolationMode
+from transformers import AutoModel, AutoTokenizer
+
+IMAGENET_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_STD = (0.229, 0.224, 0.225)
+
+def build_transform(input_size):
+    MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
+    transform = T.Compose([
+        T.Lambda(lambda img: img.convert('RGB') if img.mode != 'RGB' else img),
+        T.Resize((input_size, input_size), interpolation=InterpolationMode.BICUBIC),
+        T.ToTensor(),
+        T.Normalize(mean=MEAN, std=STD)
+    ])
+    return transform
+
+def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
+    best_ratio_diff = float('inf')
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+def dynamic_preprocess(image, min_num=1, max_num=12, image_size=448, use_thumbnail=False):
+    orig_width, orig_height = image.size
+    aspect_ratio = orig_width / orig_height
+
+    # calculate the existing image aspect ratio
+    target_ratios = set(
+        (i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
+        i * j <= max_num and i * j >= min_num)
+    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(
+        aspect_ratio, target_ratios, orig_width, orig_height, image_size)
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    # resize the image
+    resized_img = image.resize((target_width, target_height))
+    processed_images = []
+    for i in range(blocks):
+        box = (
+            (i % (target_width // image_size)) * image_size,
+            (i // (target_width // image_size)) * image_size,
+            ((i % (target_width // image_size)) + 1) * image_size,
+            ((i // (target_width // image_size)) + 1) * image_size
+        )
+        # split the image
+        split_img = resized_img.crop(box)
+        processed_images.append(split_img)
+    assert len(processed_images) == blocks
+    if use_thumbnail and len(processed_images) != 1:
+        thumbnail_img = image.resize((image_size, image_size))
+        processed_images.append(thumbnail_img)
+    return processed_images
+
+def load_image(image_file, input_size=448, max_num=12):
+    image = Image.open(image_file).convert('RGB')
+    transform = build_transform(input_size=input_size)
+    images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
+    pixel_values = [transform(image) for image in images]
+    pixel_values = torch.stack(pixel_values)
+    return pixel_values
+
+def split_model(model_name):
+    device_map = {}
+    world_size = torch.cuda.device_count()
+    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    num_layers = config.llm_config.num_hidden_layers
+    # Since the first GPU will be used for ViT, treat it as half a GPU.
+    num_layers_per_gpu = math.ceil(num_layers / (world_size - 0.5))
+    num_layers_per_gpu = [num_layers_per_gpu] * world_size
+    num_layers_per_gpu[0] = math.ceil(num_layers_per_gpu[0] * 0.5)
+    layer_cnt = 0
+    for i, num_layer in enumerate(num_layers_per_gpu):
+        for j in range(num_layer):
+            device_map[f'language_model.model.layers.{layer_cnt}'] = i
+            layer_cnt += 1
+    device_map['vision_model'] = 0
+    device_map['mlp1'] = 0
+    device_map['language_model.model.tok_embeddings'] = 0
+    device_map['language_model.model.embed_tokens'] = 0
+    device_map['language_model.output'] = 0
+    device_map['language_model.model.norm'] = 0
+    device_map['language_model.model.rotary_emb'] = 0
+    device_map['language_model.lm_head'] = 0
+    device_map[f'language_model.model.layers.{num_layers - 1}'] = 0
+
+    return device_map
+
+# If you set `load_in_8bit=True`, you will need two 80GB GPUs.
+# If you set `load_in_8bit=False`, you will need at least three 80GB GPUs.
+path = 'OpenGVLab/InternVL3-9B'
+device_map = split_model('InternVL3-9B')
+model = AutoModel.from_pretrained(
+    path,
+    torch_dtype=torch.bfloat16,
+    load_in_8bit=False,
+    low_cpu_mem_usage=True,
+    use_flash_attn=True,
+    trust_remote_code=True,
+    device_map=device_map).eval()
+tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True, use_fast=False)
+
+# set the max number of tiles in `max_num`
+pixel_values = load_image('./examples/image1.jpg', max_num=12).to(torch.bfloat16).cuda()
+generation_config = dict(max_new_tokens=1024, do_sample=True)
+
+# pure-text conversation (纯文本对话)
+question = 'Hello, who are you?'
+response, history = model.chat(tokenizer, None, question, generation_config, history=None, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+question = 'Can you tell me a story?'
+response, history = model.chat(tokenizer, None, question, generation_config, history=history, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+# single-image single-round conversation (单图单轮对话)
+question = '<image>\nPlease describe the image shortly.'
+response = model.chat(tokenizer, pixel_values, question, generation_config)
+print(f'User: {question}\nAssistant: {response}')
+
+# single-image multi-round conversation (单图多轮对话)
+question = '<image>\nPlease describe the image in detail.'
+response, history = model.chat(tokenizer, pixel_values, question, generation_config, history=None, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+question = 'Please write a poem according to the image.'
+response, history = model.chat(tokenizer, pixel_values, question, generation_config, history=history, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+# multi-image multi-round conversation, combined images (多图多轮对话，拼接图像)
+pixel_values1 = load_image('./examples/image1.jpg', max_num=12).to(torch.bfloat16).cuda()
+pixel_values2 = load_image('./examples/image2.jpg', max_num=12).to(torch.bfloat16).cuda()
+pixel_values = torch.cat((pixel_values1, pixel_values2), dim=0)
+
+question = '<image>\nDescribe the two images in detail.'
+response, history = model.chat(tokenizer, pixel_values, question, generation_config,
+                               history=None, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+question = 'What are the similarities and differences between these two images.'
+response, history = model.chat(tokenizer, pixel_values, question, generation_config,
+                               history=history, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+# multi-image multi-round conversation, separate images (多图多轮对话，独立图像)
+pixel_values1 = load_image('./examples/image1.jpg', max_num=12).to(torch.bfloat16).cuda()
+pixel_values2 = load_image('./examples/image2.jpg', max_num=12).to(torch.bfloat16).cuda()
+pixel_values = torch.cat((pixel_values1, pixel_values2), dim=0)
+num_patches_list = [pixel_values1.size(0), pixel_values2.size(0)]
+
+question = 'Image-1: <image>\nImage-2: <image>\nDescribe the two images in detail.'
+response, history = model.chat(tokenizer, pixel_values, question, generation_config,
+                               num_patches_list=num_patches_list,
+                               history=None, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+question = 'What are the similarities and differences between these two images.'
+response, history = model.chat(tokenizer, pixel_values, question, generation_config,
+                               num_patches_list=num_patches_list,
+                               history=history, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+# batch inference, single image per sample (单图批处理)
+pixel_values1 = load_image('./examples/image1.jpg', max_num=12).to(torch.bfloat16).cuda()
+pixel_values2 = load_image('./examples/image2.jpg', max_num=12).to(torch.bfloat16).cuda()
+num_patches_list = [pixel_values1.size(0), pixel_values2.size(0)]
+pixel_values = torch.cat((pixel_values1, pixel_values2), dim=0)
+
+questions = ['<image>\nDescribe the image in detail.'] * len(num_patches_list)
+responses = model.batch_chat(tokenizer, pixel_values,
+                             num_patches_list=num_patches_list,
+                             questions=questions,
+                             generation_config=generation_config)
+for question, response in zip(questions, responses):
+    print(f'User: {question}\nAssistant: {response}')
+
+# video multi-round conversation (视频多轮对话)
+def get_index(bound, fps, max_frame, first_idx=0, num_segments=32):
+    if bound:
+        start, end = bound[0], bound[1]
+    else:
+        start, end = -100000, 100000
+    start_idx = max(first_idx, round(start * fps))
+    end_idx = min(round(end * fps), max_frame)
+    seg_size = float(end_idx - start_idx) / num_segments
+    frame_indices = np.array([
+        int(start_idx + (seg_size / 2) + np.round(seg_size * idx))
+        for idx in range(num_segments)
+    ])
+    return frame_indices
+
+def load_video(video_path, bound=None, input_size=448, max_num=1, num_segments=32):
+    vr = VideoReader(video_path, ctx=cpu(0), num_threads=1)
+    max_frame = len(vr) - 1
+    fps = float(vr.get_avg_fps())
+
+    pixel_values_list, num_patches_list = [], []
+    transform = build_transform(input_size=input_size)
+    frame_indices = get_index(bound, fps, max_frame, first_idx=0, num_segments=num_segments)
+    for frame_index in frame_indices:
+        img = Image.fromarray(vr[frame_index].asnumpy()).convert('RGB')
+        img = dynamic_preprocess(img, image_size=input_size, use_thumbnail=True, max_num=max_num)
+        pixel_values = [transform(tile) for tile in img]
+        pixel_values = torch.stack(pixel_values)
+        num_patches_list.append(pixel_values.shape[0])
+        pixel_values_list.append(pixel_values)
+    pixel_values = torch.cat(pixel_values_list)
+    return pixel_values, num_patches_list
+
+video_path = './examples/red-panda.mp4'
+pixel_values, num_patches_list = load_video(video_path, num_segments=8, max_num=1)
+pixel_values = pixel_values.to(torch.bfloat16).cuda()
+video_prefix = ''.join([f'Frame{i+1}: <image>\n' for i in range(len(num_patches_list))])
+question = video_prefix + 'What is the red panda doing?'
+# Frame1: <image>\nFrame2: <image>\n...\nFrame8: <image>\n{question}
+response, history = model.chat(tokenizer, pixel_values, question, generation_config,
+                               num_patches_list=num_patches_list, history=None, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+
+question = 'Describe this video in detail.'
+response, history = model.chat(tokenizer, pixel_values, question, generation_config,
+                               num_patches_list=num_patches_list, history=history, return_history=True)
+print(f'User: {question}\nAssistant: {response}')
+```
+
+#### Streaming Output
+
+Besides this method, you can also use the following code to get streamed output.
+
+```python
+from transformers import TextIteratorStreamer
+from threading import Thread
+
+# Initialize the streamer
+streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=10)
+# Define the generation configuration
+generation_config = dict(max_new_tokens=1024, do_sample=False, streamer=streamer)
+# Start the model chat in a separate thread
+thread = Thread(target=model.chat, kwargs=dict(
+    tokenizer=tokenizer, pixel_values=pixel_values, question=question,
+    history=None, return_history=False, generation_config=generation_config,
+))
+thread.start()
+
+# Initialize an empty string to store the generated text
+generated_text = ''
+# Loop through the streamer to get the new text as it is generated
+for new_text in streamer:
+    if new_text == model.conv_template.sep:
+        break
+    generated_text += new_text
+    print(new_text, end='', flush=True)  # Print each new chunk of generated text on the same line
+```
+
+## Finetune
+
+Many repositories now support fine-tuning of the InternVL series models, including [InternVL](https://github.com/OpenGVLab/InternVL), [SWIFT](https://github.com/modelscope/ms-swift), [XTurner](https://github.com/InternLM/xtuner), and others. Please refer to their documentation for more details on fine-tuning.
+
+## Deployment
+
+### LMDeploy
+
+LMDeploy is a toolkit for compressing, deploying, and serving LLMs & VLMs.
+
+```sh
+# if lmdeploy<0.7.3, you need to explicitly set chat_template_config=ChatTemplateConfig(model_name='internvl2_5')
+pip install lmdeploy>=0.7.3
+```
+
+LMDeploy abstracts the complex inference process of multi-modal Vision-Language Models (VLM) into an easy-to-use pipeline, similar to the Large Language Model (LLM) inference pipeline.
+
+#### A 'Hello, world' Example
+
+```python
+from lmdeploy import pipeline, TurbomindEngineConfig, ChatTemplateConfig
+from lmdeploy.vl import load_image
+
+model = 'OpenGVLab/InternVL3-9B'
+image = load_image('https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/tests/data/tiger.jpeg')
+pipe = pipeline(model, backend_config=TurbomindEngineConfig(session_len=16384, tp=1), chat_template_config=ChatTemplateConfig(model_name='internvl2_5'))
+response = pipe(('describe this image', image))
+print(response.text)
+```
+
+If `ImportError` occurs while executing this case, please install the required dependency packages as prompted.
+
+#### Multi-images Inference
+
+When dealing with multiple images, you can put them all in one list. Keep in mind that multiple images will lead to a higher number of input tokens, and as a result, the size of the context window typically needs to be increased.
+
+```python
+from lmdeploy import pipeline, TurbomindEngineConfig, ChatTemplateConfig
+from lmdeploy.vl import load_image
+from lmdeploy.vl.constants import IMAGE_TOKEN
+
+model = 'OpenGVLab/InternVL3-9B'
+pipe = pipeline(model, backend_config=TurbomindEngineConfig(session_len=16384, tp=1), chat_template_config=ChatTemplateConfig(model_name='internvl2_5'))
+
+image_urls=[
+    'https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/demo/resources/human-pose.jpg',
+    'https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/demo/resources/det.jpg'
+]
+
+images = [load_image(img_url) for img_url in image_urls]
+# Numbering images improves multi-image conversations
+response = pipe((f'Image-1: {IMAGE_TOKEN}\nImage-2: {IMAGE_TOKEN}\ndescribe these two images', images))
+print(response.text)
+```
+
+#### Batch Prompts Inference
+
+Conducting inference with batch prompts is quite straightforward; just place them within a list structure:
+
+```python
+from lmdeploy import pipeline, TurbomindEngineConfig, ChatTemplateConfig
+from lmdeploy.vl import load_image
+
+model = 'OpenGVLab/InternVL3-9B'
+pipe = pipeline(model, backend_config=TurbomindEngineConfig(session_len=16384, tp=1), chat_template_config=ChatTemplateConfig(model_name='internvl2_5'))
+
+image_urls=[
+    "https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/demo/resources/human-pose.jpg",
+    "https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/demo/resources/det.jpg"
+]
+prompts = [('describe this image', load_image(img_url)) for img_url in image_urls]
+response = pipe(prompts)
+print(response)
+```
+
+#### Multi-turn Conversation
+
+There are two ways to do the multi-turn conversations with the pipeline. One is to construct messages according to the format of OpenAI and use above introduced method, the other is to use the `pipeline.chat` interface.
+
+```python
+from lmdeploy import pipeline, TurbomindEngineConfig, GenerationConfig, ChatTemplateConfig
+from lmdeploy.vl import load_image
+
+model = 'OpenGVLab/InternVL3-9B'
+pipe = pipeline(model, backend_config=TurbomindEngineConfig(session_len=16384, tp=1), chat_template_config=ChatTemplateConfig(model_name='internvl2_5'))
+
+image = load_image('https://raw.githubusercontent.com/open-mmlab/mmdeploy/main/demo/resources/human-pose.jpg')
+gen_config = GenerationConfig(top_k=40, top_p=0.8, temperature=0.8)
+sess = pipe.chat(('describe this image', image), gen_config=gen_config)
+print(sess.response.text)
+sess = pipe.chat('What is the woman doing?', session=sess, gen_config=gen_config)
+print(sess.response.text)
+```
+
+#### Service
+
+LMDeploy's `api_server` enables models to be easily packed into services with a single command. The provided RESTful APIs are compatible with OpenAI's interfaces. Below are an example of service startup:
+
+```shell
+lmdeploy serve api_server OpenGVLab/InternVL3-9B --chat-template internvl2_5 --server-port 23333 --tp 1
+```
+
+To use the OpenAI-style interface, you need to install OpenAI:
+
+```shell
+pip install openai
+```
+
+Then, use the code below to make the API call:
+
+```python
+from openai import OpenAI
+
+client = OpenAI(api_key='YOUR_API_KEY', base_url='http://0.0.0.0:23333/v1')
+model_name = client.models.list().data[0].id
+response = client.chat.completions.create(
+    model=model_name,
+    messages=[{
+        'role':
+        'user',
+        'content': [{
+            'type': 'text',
+            'text': 'describe this image',
+        }, {
+            'type': 'image_url',
+            'image_url': {
+                'url':
+                'https://modelscope.oss-cn-beijing.aliyuncs.com/resource/tiger.jpeg',
+            },
+        }],
+    }],
+    temperature=0.8,
+    top_p=0.8)
+print(response)
+```
+
+## License
+
+This project is released under the MIT License.
+
+## Citation
+
+If you find this project useful in your research, please consider citing:
+
+```BibTeX
+@article{chen2024expanding,
+  title={Expanding Performance Boundaries of Open-Source Multimodal Models with Model, Data, and Test-Time Scaling},
+  author={Chen, Zhe and Wang, Weiyun and Cao, Yue and Liu, Yangzhou and Gao, Zhangwei and Cui, Erfei and Zhu, Jinguo and Ye, Shenglong and Tian, Hao and Liu, Zhaoyang and others},
+  journal={arXiv preprint arXiv:2412.05271},
+  year={2024}
+}
+@article{wang2024mpo,
+  title={Enhancing the Reasoning Ability of Multimodal Large Language Models via Mixed Preference Optimization},
+  author={Wang, Weiyun and Chen, Zhe and Wang, Wenhai and Cao, Yue and Liu, Yangzhou and Gao, Zhangwei and Zhu, Jinguo and Zhu, Xizhou and Lu, Lewei and Qiao, Yu and Dai, Jifeng},
+  journal={arXiv preprint arXiv:2411.10442},
+  year={2024}
+}
+@article{chen2024far,
+  title={How Far Are We to GPT-4V? Closing the Gap to Commercial Multimodal Models with Open-Source Suites},
+  author={Chen, Zhe and Wang, Weiyun and Tian, Hao and Ye, Shenglong and Gao, Zhangwei and Cui, Erfei and Tong, Wenwen and Hu, Kongzhi and Luo, Jiapeng and Ma, Zheng and others},
+  journal={arXiv preprint arXiv:2404.16821},
+  year={2024}
+}
+@inproceedings{chen2024internvl,
+  title={Internvl: Scaling up vision foundation models and aligning for generic visual-linguistic tasks},
+  author={Chen, Zhe and Wu, Jiannan and Wang, Wenhai and Su, Weijie and Chen, Guo and Xing, Sen and Zhong, Muyan and Zhang, Qinglong and Zhu, Xizhou and Lu, Lewei and others},
+  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},
+  pages={24185--24198},
+  year={2024}
+}
+```

merged_transition.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4189bcb7300b89ff02c1631ec8644b6fce8d119e1dee409336b54404cb203b27
+size 645923336

modeling_cvrr_merged.py

@@ -0,0 +1,279 @@
+"""Inference-only merged CVRR, preserving the reference strict decoding path.
+
+Native projections and dense learned projections are separate. No A/B adapter
+tensors or external backbone download is required after loading the package.
+"""
+from __future__ import annotations
+
+import contextlib
+import json
+import types
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+from safetensors.torch import load_file
+from transformers import PreTrainedModel
+
+from .configuration_cvrr_merged import CVRRMergedConfig
+
+# Transformers 4.57's local-directory loader copies direct relative imports;
+# enumerate transitive helpers without eagerly importing every backbone.
+if TYPE_CHECKING:
+    from .source_splitting import SplitContext
+    from .source_perceive import PerceiveDeliberateLayout
+    from .source_spatial import SpatialVisualRecurrentCell
+    from .source_helpers import _layer_hidden
+    from .source_gemma import GemmaCVRR
+    from .source_internvl import InternVLCVRR
+    from .configuration_source_qwen25 import CloseQwen2_5_VLConfig
+    from .configuration_source_qwen3 import CloseQwen3VLConfig
+    from .modeling_source_qwen25 import CloseQwen2_5_VLForConditionalGeneration
+    from .modeling_source_qwen3 import CloseQwen3VLForConditionalGeneration
+
+
+class NativeMergedLinear(nn.Module):
+    def __init__(self, native, merged):
+        super().__init__()
+        if tuple(native.weight.shape) != tuple(merged.shape):
+            raise ValueError('Native/merged projection shape mismatch')
+        self.base = native
+        self.register_buffer('merged_weight', merged.to(device=native.weight.device, dtype=torch.float32))
+        self.enabled = False
+        self.dropout = nn.Identity()
+        self.requires_grad_(False)
+
+    def _apply(self, fn, recurse=True):
+        # .to(device) is allowed; a caller's dtype cast must not truncate the
+        # FP32 merged weights and then merely upcast already-lost precision.
+        original = self.merged_weight
+        super()._apply(fn, recurse=recurse)
+        self.merged_weight = original.to(device=self.merged_weight.device, dtype=torch.float32)
+        return self
+
+    def forward(self, x):
+        if not self.enabled:
+            return self.base(x)
+        bias = None if self.base.bias is None else self.base.bias.float()
+        return torch.nn.functional.linear(x.float(), self.merged_weight, bias).to(x.dtype)
+
+
+@contextlib.contextmanager
+def merged_execution(model, enabled):
+    modules = model._cvrr_merged_projections
+    previous = [m.enabled for m in modules]
+    for m in modules:
+        m.enabled = bool(enabled)
+    try:
+        yield
+    finally:
+        for m, old in zip(modules, previous):
+            m.enabled = old
+
+
+def set_merged_execution(model, enabled):
+    for m in model._cvrr_merged_projections:
+        m.enabled = bool(enabled)
+
+
+def install_merged(runtime, weights, *, generic):
+    cell = runtime.layers[runtime.cell_index] if generic else runtime.text_model.layers[int(runtime.config.ell_star) + 1]
+    modules = []
+    for key, weight in sorted(weights.items()):
+        if not key.endswith('.weight'):
+            raise ValueError(f'Unexpected merged state key: {key}')
+        parts = key.removesuffix('.weight').split('.')
+        parent = cell
+        for part in parts[:-1]:
+            parent = getattr(parent, part)
+        old = getattr(parent, parts[-1])
+        native = old.base if generic else old.base_layer
+        wrapped = NativeMergedLinear(native, weight)
+        setattr(parent, parts[-1], wrapped)
+        modules.append(wrapped)
+    # Deliberately a tuple, not another ModuleList alias in state_dict.
+    runtime._cvrr_merged_projections = tuple(modules)
+    if generic:
+        runtime.lora = {str(i): m for i, m in enumerate(modules)}
+        runtime.adapters = types.MethodType(merged_execution, runtime)
+    else:
+        runtime._recurrence_adapter_execution = types.MethodType(merged_execution, runtime)
+        runtime._set_adapters_enabled = types.MethodType(set_merged_execution, runtime)
+    if any('lora_' in n for n, _ in runtime.named_parameters()):
+        raise RuntimeError('Unmerged LoRA parameters remain')
+    runtime.requires_grad_(False)
+    return runtime.eval()
+
+
+class CVRRMergedModel(PreTrainedModel):
+    config_class = CVRRMergedConfig
+    main_input_name = 'input_ids'
+
+    def __init__(self, config, runtime=None):
+        super().__init__(config)
+        if runtime is None:
+            raise ValueError('Load a packaged release using from_pretrained')
+        self.runtime = runtime
+        self.is_generic = 'Qwen' not in config.release['name']
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
+        if args:
+            raise TypeError('Positional loader overrides are not supported')
+        config = kwargs.pop('config', None)
+        device = kwargs.pop('device_map', kwargs.pop('device', 'cpu'))
+        if device is None:
+            device = 'cpu'
+        if isinstance(device, dict) or str(device) == 'auto':
+            raise ValueError('Use one complete replica per device; pass device_map="cuda:0" or "cpu"')
+        dtype = kwargs.pop('dtype', kwargs.pop('torch_dtype', torch.bfloat16))
+        if dtype not in (torch.bfloat16, 'bfloat16', 'auto', None):
+            raise ValueError('This release preserves BF16 native weights and FP32 merged projections')
+        offline = kwargs.pop('local_files_only', False)
+        revision = kwargs.pop('revision', None)
+        token = kwargs.pop('token', None)
+        kwargs.pop('trust_remote_code', None)
+        kwargs.pop('_from_auto', None)
+        kwargs.pop('adapter_kwargs', None)
+        kwargs.pop('name_or_path', None)
+        kwargs.pop('_commit_hash', None)
+        if kwargs:
+            raise TypeError(f'Unsupported loader options: {sorted(kwargs)}')
+        path = Path(pretrained_model_name_or_path)
+        if not path.is_dir():
+            from huggingface_hub import snapshot_download
+            path = Path(snapshot_download(str(pretrained_model_name_or_path), revision=revision,
+                                          token=token, local_files_only=offline))
+        if config is None:
+            config = CVRRMergedConfig.from_pretrained(path, local_files_only=True)
+        settings = json.loads((path / 'cvrr_release_config.json').read_text())
+        if config.release != settings:
+            raise ValueError('Root configuration and release metadata disagree')
+        if settings['format'] != 'cvrr_native_plus_merged_transition_v1':
+            raise ValueError('Unsupported release layout')
+        native_path = str(path / 'native_backbone')
+        generic = 'Qwen' not in settings['name']
+        common = dict(ell_star=settings['ell_star'], steps=settings['inference_T'],
+                      beta=settings['inference_beta'], rank=settings['lora_rank'],
+                      alpha=settings['lora_alpha'], dropout=settings['lora_dropout'],
+                      device=device, offline=True)
+        if generic:
+            if 'InternVL' in settings['name']:
+                from .source_internvl import InternVLCVRR
+                runtime = InternVLCVRR(native_path, **common)
+            else:
+                from .source_gemma import GemmaCVRR
+                runtime = GemmaCVRR(native_path, **common)
+        else:
+            if 'Qwen2.5' in settings['name']:
+                from transformers import Qwen2_5_VLForConditionalGeneration as Native
+                from .configuration_source_qwen25 import CloseQwen2_5_VLConfig as SourceConfig
+                from .modeling_source_qwen25 import CloseQwen2_5_VLForConditionalGeneration as Source
+            else:
+                from transformers import Qwen3VLForConditionalGeneration as Native
+                from .configuration_source_qwen3 import CloseQwen3VLConfig as SourceConfig
+                from .modeling_source_qwen3 import CloseQwen3VLForConditionalGeneration as Source
+            source_config = SourceConfig.from_pretrained(path / 'source_config', local_files_only=True)
+            source_config.num_workspace_steps = settings['inference_T']
+            source_config.counterfactual_beta = settings['inference_beta']
+            for flag in ('counterfactual_reliance_loss', 'counterfactual_step_retention_loss',
+                         'functional_state_loss', 'interface_loss'):
+                setattr(source_config, flag, False)
+            native = Native.from_pretrained(native_path, dtype=torch.bfloat16,
+                                           device_map=str(device), local_files_only=True,
+                                           attn_implementation='sdpa')
+            runtime = Source(source_config, backbone=native).to(device)
+        weights = load_file(path / 'merged_transition.safetensors', device='cpu')
+        runtime = install_merged(runtime, weights, generic=generic)
+        model = cls(config, runtime).eval()
+        model.release_path = path
+        return model
+
+    def prepare_inputs(self, image, question, *, max_visual_tokens=8192, max_tiles=12):
+        """Single-image, batch-one preparation using the reference model prompts."""
+        if max_visual_tokens < 1:
+            raise ValueError('max_visual_tokens must be positive')
+        device = next(self.runtime.parameters()).device
+        if 'InternVL' in self.config.release['name']:
+            from .source_internvl import InternVLArrowCollator
+            from .source_helpers import dynamic_tiles, _normalize_tiles
+            helper = InternVLArrowCollator(self.runtime, max_tiles=max_tiles)
+            tiles = dynamic_tiles(image, image_size=helper.image_size, max_tiles=max_tiles,
+                                  thumbnail=helper.use_thumbnail)
+            prompt = helper._query(question, '', len(tiles))
+            self.tokenizer = self.runtime.tokenizer
+            encoded = dict(self.tokenizer(prompt, return_tensors='pt'))
+            encoded.update(pixel_values=_normalize_tiles(tiles).to(torch.bfloat16),
+                           image_flags=torch.ones(len(tiles), 1, dtype=torch.long))
+        else:
+            from transformers import AutoProcessor
+            processor = AutoProcessor.from_pretrained(self.release_path / 'native_backbone',
+                           local_files_only=True, use_fast=('Qwen3' in self.config.release['name']))
+            self.tokenizer = processor.tokenizer
+            if not self.is_generic:
+                cfg = self.runtime.config
+                patch = int(cfg.vision_config.patch_size) * int(cfg.vision_config.spatial_merge_size)
+                pixels = max_visual_tokens * patch * patch
+                ip = processor.image_processor
+                if isinstance(getattr(ip, 'size', None), dict):
+                    ip.size = dict(ip.size, longest_edge=pixels)
+                ip.max_pixels = pixels
+            prompt = processor.apply_chat_template([{'role':'user', 'content':[
+                {'type':'image'}, {'type':'text','text':question}]}],
+                tokenize=False, add_generation_prompt=True)
+            encoded = dict(processor(text=[prompt], images=[image], return_tensors='pt'))
+            if not self.is_generic:
+                txt = processor.apply_chat_template([{'role':'user','content':[
+                    {'type':'text','text':question}]}], tokenize=False, add_generation_prompt=True)
+                q = processor.tokenizer(txt, return_tensors='pt', add_special_tokens=False)
+                encoded['question_ids'] = q['input_ids']
+                encoded['question_attention_mask'] = q['attention_mask']
+        result = {k: v.to(device) for k,v in encoded.items() if isinstance(v, torch.Tensor)}
+        if 'pixel_values' in result:
+            result['pixel_values'] = result['pixel_values'].to(torch.bfloat16)
+        if self.is_generic:
+            count = int((self.runtime._modality(result).eq(1) & result['attention_mask'].bool()).sum())
+        else:
+            count = int(result['input_ids'].eq(self.runtime.config.image_token_id).sum())
+            allowed = {'input_ids','attention_mask','pixel_values','image_grid_thw',
+                       'question_ids','question_attention_mask'}
+            result = {k:v for k,v in result.items() if k in allowed}
+        if not 0 < count <= max_visual_tokens:
+            raise ValueError(f'Visual token count {count} exceeds cap {max_visual_tokens} or is empty')
+        return result
+
+    def get_input_embeddings(self):
+        if self.is_generic:
+            return self.runtime.base_model.get_input_embeddings()
+        return self.runtime.backbone.get_input_embeddings()
+
+    def save_pretrained(self, *args, **kwargs):
+        raise NotImplementedError('Keep the exported directory intact; generic HF reserialization would change the release layout')
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError('Inference-only release: use next_token_logits or generate')
+
+    @torch.inference_mode()
+    def next_token_logits(self, **inputs):
+        if self.is_generic:
+            trace = self.runtime.extract(inputs)
+            state = self.runtime.rollout(trace)[-1]
+            return self.runtime.next_token_logits(trace, state)
+        upper, _cache, _states, _evidence = self.runtime.prefill(**inputs)
+        mask = inputs.get('question_attention_mask')
+        last = (mask.long().sum(-1) - 1 if mask is not None else
+                torch.full((upper.shape[0],), upper.shape[1]-1, device=upper.device))
+        hidden = upper[torch.arange(upper.shape[0], device=upper.device), last.long()]
+        return self.runtime.backbone.lm_head(self.runtime.text_model.norm(hidden)).float()
+
+    @torch.inference_mode()
+    def generate(self, *, do_sample=False, max_new_tokens=32, **inputs):
+        if do_sample:
+            raise ValueError('Only deterministic greedy decoding is supported by this release')
+        if self.is_generic:
+            raise NotImplementedError('Gemma/InternVL release validation currently covers strict next-token readout only')
+        return self.runtime.generate_answer(**inputs, max_new_tokens=max_new_tokens)
+
+
+CVRRMergedModel.register_for_auto_class('AutoModelForImageTextToText')

native_backbone/added_tokens.json

@@ -0,0 +1,11 @@
+{
+  "</box>": 128141,
+  "</img>": 128134,
+  "</quad>": 128137,
+  "</ref>": 128139,
+  "<IMG_CONTEXT>": 128135,
+  "<box>": 128140,
+  "<img>": 128133,
+  "<quad>": 128136,
+  "<ref>": 128138
+}

native_backbone/config.json

@@ -0,0 +1,225 @@
+{
+  "_commit_hash": null,
+  "architectures": [
+    "InternVLChatModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_internvl_chat.InternVLChatConfig",
+    "AutoModel": "modeling_internvl_chat.InternVLChatModel",
+    "AutoModelForCausalLM": "modeling_internvl_chat.InternVLChatModel"
+  },
+  "downsample_ratio": 0.5,
+  "dynamic_image_size": true,
+  "force_image_size": 448,
+  "hidden_size": 4096,
+  "image_fold": null,
+  "llm_config": {
+    "_attn_implementation_autoset": true,
+    "add_cross_attention": false,
+    "architectures": [
+      "InternLM2ForCausalLM"
+    ],
+    "attn_implementation": "flash_attention_2",
+    "auto_map": {
+      "AutoConfig": "configuration_internlm2.InternLM2Config",
+      "AutoModel": "modeling_internlm2.InternLM2ForCausalLM",
+      "AutoModelForCausalLM": "modeling_internlm2.InternLM2ForCausalLM"
+    },
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bias": false,
+    "bos_token_id": 1,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 2,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "silu",
+    "hidden_size": 4096,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 10240,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 32768,
+    "min_length": 0,
+    "model_type": "internlm2",
+    "moe_config": null,
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 32,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 48,
+    "num_key_value_heads": 2,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 2,
+    "prefix": null,
+    "pretraining_tp": 1,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": {
+      "factor": 2.0,
+      "type": "dynamic"
+    },
+    "rope_theta": 50000000,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "transformers_version": "4.48.3",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": false,
+    "vocab_size": 128142
+  },
+  "max_dynamic_patch": 12,
+  "min_dynamic_patch": 1,
+  "model_type": "internvl_chat",
+  "pad2square": false,
+  "ps_version": "v2",
+  "select_layer": -1,
+  "system_message": null,
+  "template": "internvl2_5",
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": null,
+  "use_backbone_lora": 0,
+  "use_img_start_end_token": true,
+  "use_llm_lora": 0,
+  "use_thumbnail": true,
+  "vision_config": {
+    "_attn_implementation_autoset": true,
+    "add_cross_attention": false,
+    "architectures": [
+      "InternVisionModel"
+    ],
+    "attention_dropout": 0.0,
+    "auto_map": {
+      "AutoConfig": "configuration_intern_vit.InternVisionConfig",
+      "AutoModel": "modeling_intern_vit.InternVisionModel"
+    },
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "capacity_factor": 1.2,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "drop_path_rate": 0.1,
+    "dropout": 0.0,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "eval_capacity_factor": 1.4,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 448,
+    "initializer_factor": 0.1,
+    "initializer_range": 1e-10,
+    "intermediate_size": 4096,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "laux_allreduce": "all_nodes",
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "intern_vit_6b",
+    "moe_coeff_ratio": 0.5,
+    "moe_intermediate_size": 768,
+    "moe_output_scale": 4.0,
+    "no_repeat_ngram_size": 0,
+    "noisy_gate_policy": "RSample_before",
+    "norm_type": "layer_norm",
+    "num_attention_heads": 16,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_experts": 8,
+    "num_hidden_layers": 24,
+    "num_return_sequences": 1,
+    "num_routed_experts": 4,
+    "num_shared_experts": 4,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "qk_normalization": false,
+    "qkv_bias": true,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "shared_expert_intermediate_size": 3072,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "transformers_version": "4.48.3",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_flash_attn": true,
+    "use_moe": false,
+    "use_residual": true,
+    "use_rts": false,
+    "use_weighted_residual": false
+  }
+}

native_backbone/configuration_intern_vit.py

@@ -0,0 +1,119 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+import os
+from typing import Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class InternVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`InternVisionModel`]. It is used to
+    instantiate a vision encoder according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        num_channels (`int`, *optional*, defaults to 3):
+            Number of color channels in the input images (e.g., 3 for RGB).
+        patch_size (`int`, *optional*, defaults to 14):
+            The size (resolution) of each patch.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        qkv_bias (`bool`, *optional*, defaults to `False`):
+            Whether to add a bias to the queries and values in the self-attention layers.
+        hidden_size (`int`, *optional*, defaults to 3200):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_attention_heads (`int`, *optional*, defaults to 25):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (`int`, *optional*, defaults to 12800):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        qk_normalization (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the queries and keys in the self-attention layers.
+        num_hidden_layers (`int`, *optional*, defaults to 48):
+            Number of hidden layers in the Transformer encoder.
+        use_flash_attn (`bool`, *optional*, defaults to `True`):
+            Whether to use flash attention mechanism.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` ``"gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-6):
+            The epsilon used by the layer normalization layers.
+        dropout (`float`, *optional*, defaults to 0.0):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        drop_path_rate (`float`, *optional*, defaults to 0.0):
+            Dropout rate for stochastic depth.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        initializer_factor (`float`, *optional*, defaults to 0.1):
+            A factor for layer scale.
+    """
+
+    model_type = 'intern_vit_6b'
+
+    def __init__(
+            self,
+            num_channels=3,
+            patch_size=14,
+            image_size=224,
+            qkv_bias=False,
+            hidden_size=3200,
+            num_attention_heads=25,
+            intermediate_size=12800,
+            qk_normalization=True,
+            num_hidden_layers=48,
+            use_flash_attn=True,
+            hidden_act='gelu',
+            norm_type='rms_norm',
+            layer_norm_eps=1e-6,
+            dropout=0.0,
+            drop_path_rate=0.0,
+            attention_dropout=0.0,
+            initializer_range=0.02,
+            initializer_factor=0.1,
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.dropout = dropout
+        self.drop_path_rate = drop_path_rate
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.initializer_range = initializer_range
+        self.initializer_factor = initializer_factor
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.norm_type = norm_type
+        self.qkv_bias = qkv_bias
+        self.qk_normalization = qk_normalization
+        self.use_flash_attn = use_flash_attn
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> 'PretrainedConfig':
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        if 'vision_config' in config_dict:
+            config_dict = config_dict['vision_config']
+
+        if 'model_type' in config_dict and hasattr(cls, 'model_type') and config_dict['model_type'] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
+            )
+
+        return cls.from_dict(config_dict, **kwargs)

native_backbone/configuration_internlm2.py

@@ -0,0 +1,150 @@
+# Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on transformers/src/transformers/models/llama/configuration_llama.py
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" InternLM2 model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+INTERNLM2_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
+
+
+# Modified from transformers.model.llama.configuration_llama.LlamaConfig
+class InternLM2Config(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`InternLM2Model`]. It is used to instantiate
+    an InternLM2 model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the InternLM2-7B.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 32000):
+            Vocabulary size of the InternLM2 model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`InternLM2Model`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 11008):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 32):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details checkout [this
+            paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
+            `num_attention_heads`.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 2048):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-12):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings(`bool`, *optional*, defaults to `False`):
+            Whether to tie weight embeddings
+        Example:
+
+    """
+    model_type = 'internlm2'
+    _auto_class = 'AutoConfig'
+
+    def __init__(  # pylint: disable=W0102
+        self,
+        vocab_size=103168,
+        hidden_size=4096,
+        intermediate_size=11008,
+        num_hidden_layers=32,
+        num_attention_heads=32,
+        num_key_value_heads=None,
+        hidden_act='silu',
+        max_position_embeddings=2048,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        use_cache=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        bias=True,
+        rope_theta=10000,
+        rope_scaling=None,
+        attn_implementation='eager',
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.bias = bias
+
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self._rope_scaling_validation()
+
+        self.attn_implementation = attn_implementation
+        if self.attn_implementation is None:
+            self.attn_implementation = 'eager'
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    def _rope_scaling_validation(self):
+        """
+        Validate the `rope_scaling` configuration.
+        """
+        if self.rope_scaling is None:
+            return
+
+        if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
+            raise ValueError(
+                '`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, '
+                f'got {self.rope_scaling}'
+            )
+        rope_scaling_type = self.rope_scaling.get('type', None)
+        rope_scaling_factor = self.rope_scaling.get('factor', None)
+        if rope_scaling_type is None or rope_scaling_type not in ['linear', 'dynamic']:
+            raise ValueError(
+                f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
+            )
+        if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor < 1.0:
+            raise ValueError(f"`rope_scaling`'s factor field must be a float >= 1, got {rope_scaling_factor}")

native_backbone/configuration_internvl_chat.py

@@ -0,0 +1,101 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import copy
+
+from transformers import AutoConfig, LlamaConfig, Qwen2Config
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+from .configuration_intern_vit import InternVisionConfig
+from .configuration_internlm2 import InternLM2Config
+
+
+logger = logging.get_logger(__name__)
+
+
+class InternVLChatConfig(PretrainedConfig):
+    model_type = 'internvl_chat'
+    is_composition = True
+
+    def __init__(
+            self,
+            vision_config=None,
+            llm_config=None,
+            use_backbone_lora=0,
+            use_llm_lora=0,
+            select_layer=-1,
+            force_image_size=None,
+            downsample_ratio=0.5,
+            template=None,
+            dynamic_image_size=False,
+            use_thumbnail=False,
+            ps_version='v1',
+            min_dynamic_patch=1,
+            max_dynamic_patch=6,
+            **kwargs):
+        super().__init__(**kwargs)
+
+        if vision_config is None:
+            vision_config = {'architectures': ['InternVisionModel']}
+            logger.info('vision_config is None. Initializing the InternVisionConfig with default values.')
+
+        if llm_config is None:
+            llm_config = {'architectures': ['InternLM2ForCausalLM']}
+            logger.info('llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`).')
+
+        self.vision_config = InternVisionConfig(**vision_config)
+        if llm_config['architectures'][0] == 'LlamaForCausalLM':
+            self.llm_config = LlamaConfig(**llm_config)
+        elif llm_config['architectures'][0] == 'InternLM2ForCausalLM':
+            self.llm_config = InternLM2Config(**llm_config)
+        elif llm_config['architectures'][0] == 'Qwen2ForCausalLM':
+            self.llm_config = Qwen2Config(**llm_config)
+        else:
+            raise ValueError('Unsupported architecture: {}'.format(llm_config['architectures'][0]))
+        self.use_backbone_lora = use_backbone_lora
+        self.use_llm_lora = use_llm_lora
+        self.select_layer = select_layer
+        self.force_image_size = force_image_size
+        self.downsample_ratio = downsample_ratio
+        self.template = template
+        self.dynamic_image_size = dynamic_image_size
+        self.use_thumbnail = use_thumbnail
+        self.ps_version = ps_version  # pixel shuffle version
+        self.min_dynamic_patch = min_dynamic_patch
+        self.max_dynamic_patch = max_dynamic_patch
+        # By default, we use tie_word_embeddings=False for models of all sizes.
+        self.tie_word_embeddings = self.llm_config.tie_word_embeddings
+
+        logger.info(f'vision_select_layer: {self.select_layer}')
+        logger.info(f'ps_version: {self.ps_version}')
+        logger.info(f'min_dynamic_patch: {self.min_dynamic_patch}')
+        logger.info(f'max_dynamic_patch: {self.max_dynamic_patch}')
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output['vision_config'] = self.vision_config.to_dict()
+        output['llm_config'] = self.llm_config.to_dict()
+        output['model_type'] = self.__class__.model_type
+        output['use_backbone_lora'] = self.use_backbone_lora
+        output['use_llm_lora'] = self.use_llm_lora
+        output['select_layer'] = self.select_layer
+        output['force_image_size'] = self.force_image_size
+        output['downsample_ratio'] = self.downsample_ratio
+        output['template'] = self.template
+        output['dynamic_image_size'] = self.dynamic_image_size
+        output['use_thumbnail'] = self.use_thumbnail
+        output['ps_version'] = self.ps_version
+        output['min_dynamic_patch'] = self.min_dynamic_patch
+        output['max_dynamic_patch'] = self.max_dynamic_patch
+
+        return output

native_backbone/conversation.py

@@ -0,0 +1,391 @@
+"""
+Conversation prompt templates.
+
+We kindly request that you import fastchat instead of copying this file if you wish to use it.
+If you have changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
+
+Modified from https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
+"""
+
+import dataclasses
+from enum import IntEnum, auto
+from typing import Dict, List, Tuple, Union
+
+
+class SeparatorStyle(IntEnum):
+    """Separator styles."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    ADD_COLON_SPACE_SINGLE = auto()
+    NO_COLON_SINGLE = auto()
+    NO_COLON_TWO = auto()
+    ADD_NEW_LINE_SINGLE = auto()
+    LLAMA2 = auto()
+    CHATGLM = auto()
+    CHATML = auto()
+    CHATINTERN = auto()
+    DOLLY = auto()
+    RWKV = auto()
+    PHOENIX = auto()
+    ROBIN = auto()
+    FALCON_CHAT = auto()
+    CHATGLM3 = auto()
+    INTERNVL_ZH = auto()
+    MPT = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The template of the system prompt
+    system_template: str = '{system_message}'
+    # The system message
+    system_message: str = ''
+    # The names of two roles
+    roles: Tuple[str] = ('USER', 'ASSISTANT')
+    # All messages. Each item is (role, message).
+    messages: List[List[str]] = ()
+    # The number of few shot examples
+    offset: int = 0
+    # The separator style and configurations
+    sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
+    sep: str = '\n'
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: Union[str, List[str]] = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ': '  # must be end with a space
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + message + seps[i % 2]
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                        role
+                        + ': '
+                        + message.replace('\r\n', '\n').replace('\n\n', '\n')
+                    )
+                    ret += '\n\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA2:
+            seps = [self.sep, self.sep2]
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = '[INST] '
+            for i, (role, message) in enumerate(self.messages):
+                tag = self.roles[i % 2]
+                if message:
+                    if i == 0:
+                        ret += message + ' '
+                    else:
+                        ret += tag + ' ' + message + seps[i % 2]
+                else:
+                    ret += tag
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM:
+            # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+            # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+            round_add_n = 1 if self.name == 'chatglm2' else 0
+            if system_prompt:
+                ret = system_prompt + self.sep
+            else:
+                ret = ''
+
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += f'[Round {i//2 + round_add_n}]{self.sep}'
+
+                if message:
+                    ret += f'{role}：{message}{self.sep}'
+                else:
+                    ret += f'{role}：'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = '' if system_prompt == '' else system_prompt + self.sep + '\n'
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep + '\n'
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM3:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + ' ' + message
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATINTERN:
+            # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                # if i % 2 == 0:
+                #     ret += "<s>"
+                if message:
+                    ret += role + ':' + message + seps[i % 2] + '\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ':\n' + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += '\n\n'
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + '<s>' + message + '</s>'
+                else:
+                    ret += role + ': ' + '<s>'
+            return ret
+        elif self.sep_style == SeparatorStyle.ROBIN:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ':\n' + message + self.sep
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.FALCON_CHAT:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+
+            return ret
+        elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
+            seps = [self.sep, self.sep2]
+            ret = self.system_message + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        else:
+            raise ValueError(f'Invalid style: {self.sep_style}')
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        ret = [{'role': 'system', 'content': self.system_message}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({'role': 'user', 'content': msg})
+            else:
+                if msg is not None:
+                    ret.append({'role': 'assistant', 'content': msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+    def dict(self):
+        return {
+            'template_name': self.name,
+            'system_message': self.system_message,
+            'roles': self.roles,
+            'messages': self.messages,
+            'offset': self.offset,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f'{template.name} has been registered.'
+
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+# Both Hermes-2 and internlm2-chat are chatml-format conversation templates. The difference
+# is that during training, the preprocessing function for the Hermes-2 template doesn't add
+# <s> at the beginning of the tokenized sequence, while the internlm2-chat template does.
+# Therefore, they are completely equivalent during inference.
+register_conv_template(
+    Conversation(
+        name='Hermes-2',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_str='<|endoftext|>',
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='internlm2-chat',
+        system_template='<|im_start|>system\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='phi3-chat',
+        system_template='<|system|>\n{system_message}',
+        # note: The new system prompt was not used here to avoid changes in benchmark performance.
+        # system_message='我是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        system_message='你是由上海人工智能实验室联合商汤科技开发的书生多模态大模型，英文名叫InternVL, 是一个有用无害的人工智能助手。',
+        roles=('<|user|>\n', '<|assistant|>\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|end|>',
+    )
+)
+
+
+register_conv_template(
+    Conversation(
+        name='internvl2_5',
+        system_template='<|im_start|>system\n{system_message}',
+        system_message='你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>\n',
+    )
+)

native_backbone/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.48.3"
+}

native_backbone/model.safetensors.index.json

@@ -0,0 +1,692 @@
+{
+  "metadata": {
+    "total_size": 18277586944
+  },
+  "weight_map": {
+    "language_model.model.layers.0.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.0.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.0.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.0.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.0.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.0.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.0.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.1.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.1.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.1.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.1.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.1.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.1.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.1.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.10.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.10.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.10.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.10.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.10.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.10.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.10.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.11.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.11.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.11.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.11.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.11.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.11.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.11.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.12.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.12.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.12.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.12.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.12.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.12.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.12.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.13.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.13.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.13.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.13.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.13.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.13.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.13.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.14.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.14.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.14.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.14.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.14.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.14.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.14.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.15.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.15.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.15.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.15.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.15.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.15.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.15.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.16.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.16.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.16.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.16.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.16.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.16.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.16.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.17.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.17.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.17.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.17.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.17.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.17.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.17.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.18.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.18.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.18.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.18.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.18.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.18.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.18.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.19.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.19.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.19.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.19.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.19.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.19.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.19.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.2.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.2.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.2.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.2.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.2.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.2.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.2.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.20.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.20.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.20.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.20.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.20.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.20.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.20.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.21.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.21.attention.wqkv.weight": "native-00001.safetensors",
+    "language_model.model.layers.21.attention_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.21.feed_forward.w1.weight": "native-00001.safetensors",
+    "language_model.model.layers.21.feed_forward.w2.weight": "native-00001.safetensors",
+    "language_model.model.layers.21.feed_forward.w3.weight": "native-00001.safetensors",
+    "language_model.model.layers.21.ffn_norm.weight": "native-00001.safetensors",
+    "language_model.model.layers.22.attention.wo.weight": "native-00001.safetensors",
+    "language_model.model.layers.22.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.22.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.22.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.22.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.22.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.22.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.23.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.23.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.23.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.23.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.23.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.23.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.23.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.24.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.24.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.24.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.24.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.24.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.24.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.24.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.25.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.25.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.25.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.25.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.25.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.25.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.25.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.26.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.26.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.26.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.26.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.26.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.26.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.26.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.27.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.27.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.27.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.27.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.27.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.27.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.27.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.28.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.28.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.28.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.28.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.28.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.28.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.28.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.29.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.29.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.29.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.29.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.29.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.29.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.29.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.3.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.3.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.3.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.3.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.3.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.3.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.3.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.30.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.30.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.30.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.30.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.30.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.30.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.30.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.31.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.31.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.31.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.31.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.31.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.31.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.31.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.32.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.32.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.32.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.32.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.32.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.32.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.32.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.33.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.33.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.33.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.33.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.33.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.33.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.33.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.34.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.34.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.34.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.34.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.34.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.34.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.34.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.35.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.35.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.35.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.35.feed_forward.w1.weight": "native-00002.safetensors",
+    "language_model.model.layers.35.feed_forward.w2.weight": "native-00002.safetensors",
+    "language_model.model.layers.35.feed_forward.w3.weight": "native-00002.safetensors",
+    "language_model.model.layers.35.ffn_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.36.attention.wo.weight": "native-00002.safetensors",
+    "language_model.model.layers.36.attention.wqkv.weight": "native-00002.safetensors",
+    "language_model.model.layers.36.attention_norm.weight": "native-00002.safetensors",
+    "language_model.model.layers.36.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.36.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.36.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.36.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.37.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.37.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.37.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.37.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.37.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.37.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.37.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.38.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.38.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.38.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.38.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.38.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.38.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.38.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.39.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.39.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.39.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.39.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.39.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.39.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.39.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.4.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.4.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.4.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.4.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.4.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.4.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.4.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.40.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.40.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.40.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.40.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.40.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.40.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.40.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.41.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.41.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.41.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.41.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.41.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.41.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.41.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.42.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.42.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.42.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.42.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.42.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.42.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.42.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.43.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.43.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.43.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.43.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.43.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.43.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.43.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.44.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.44.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.44.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.44.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.44.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.44.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.44.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.45.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.45.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.45.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.45.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.45.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.45.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.45.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.46.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.46.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.46.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.46.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.46.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.46.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.46.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.47.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.47.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.47.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.47.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.47.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.47.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.47.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.5.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.5.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.5.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.5.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.5.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.5.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.5.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.6.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.6.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.6.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.6.feed_forward.w1.weight": "native-00003.safetensors",
+    "language_model.model.layers.6.feed_forward.w2.weight": "native-00003.safetensors",
+    "language_model.model.layers.6.feed_forward.w3.weight": "native-00003.safetensors",
+    "language_model.model.layers.6.ffn_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.7.attention.wo.weight": "native-00003.safetensors",
+    "language_model.model.layers.7.attention.wqkv.weight": "native-00003.safetensors",
+    "language_model.model.layers.7.attention_norm.weight": "native-00003.safetensors",
+    "language_model.model.layers.7.feed_forward.w1.weight": "native-00004.safetensors",
+    "language_model.model.layers.7.feed_forward.w2.weight": "native-00004.safetensors",
+    "language_model.model.layers.7.feed_forward.w3.weight": "native-00004.safetensors",
+    "language_model.model.layers.7.ffn_norm.weight": "native-00004.safetensors",
+    "language_model.model.layers.8.attention.wo.weight": "native-00004.safetensors",
+    "language_model.model.layers.8.attention.wqkv.weight": "native-00004.safetensors",
+    "language_model.model.layers.8.attention_norm.weight": "native-00004.safetensors",
+    "language_model.model.layers.8.feed_forward.w1.weight": "native-00004.safetensors",
+    "language_model.model.layers.8.feed_forward.w2.weight": "native-00004.safetensors",
+    "language_model.model.layers.8.feed_forward.w3.weight": "native-00004.safetensors",
+    "language_model.model.layers.8.ffn_norm.weight": "native-00004.safetensors",
+    "language_model.model.layers.9.attention.wo.weight": "native-00004.safetensors",
+    "language_model.model.layers.9.attention.wqkv.weight": "native-00004.safetensors",
+    "language_model.model.layers.9.attention_norm.weight": "native-00004.safetensors",
+    "language_model.model.layers.9.feed_forward.w1.weight": "native-00004.safetensors",
+    "language_model.model.layers.9.feed_forward.w2.weight": "native-00004.safetensors",
+    "language_model.model.layers.9.feed_forward.w3.weight": "native-00004.safetensors",
+    "language_model.model.layers.9.ffn_norm.weight": "native-00004.safetensors",
+    "language_model.model.norm.weight": "native-00004.safetensors",
+    "language_model.model.tok_embeddings.weight": "native-00004.safetensors",
+    "language_model.output.weight": "native-00004.safetensors",
+    "mlp1.0.bias": "native-00004.safetensors",
+    "mlp1.0.weight": "native-00004.safetensors",
+    "mlp1.1.bias": "native-00004.safetensors",
+    "mlp1.1.weight": "native-00004.safetensors",
+    "mlp1.3.bias": "native-00004.safetensors",
+    "mlp1.3.weight": "native-00004.safetensors",
+    "vision_model.embeddings.class_embedding": "native-00004.safetensors",
+    "vision_model.embeddings.patch_embedding.bias": "native-00004.safetensors",
+    "vision_model.embeddings.patch_embedding.weight": "native-00004.safetensors",
+    "vision_model.embeddings.position_embedding": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.0.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.1.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.10.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.11.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.12.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.13.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.14.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.15.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.16.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.17.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.18.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.19.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.2.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.20.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.21.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.22.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.23.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.3.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.4.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.5.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.6.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.7.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.8.norm2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.attn.proj.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.attn.proj.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.attn.qkv.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.attn.qkv.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.ls1": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.ls2": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.mlp.fc1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.mlp.fc1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.mlp.fc2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.mlp.fc2.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.norm1.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.norm1.weight": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.norm2.bias": "native-00004.safetensors",
+    "vision_model.encoder.layers.9.norm2.weight": "native-00004.safetensors"
+  }
+}

native_backbone/modeling_intern_vit.py

@@ -0,0 +1,429 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from timm.layers import DropPath
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutput,
+                                           BaseModelOutputWithPooling)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration_intern_vit import InternVisionConfig
+
+try:
+    from flash_attn.bert_padding import pad_input, unpad_input
+    from flash_attn.flash_attn_interface import \
+        flash_attn_varlen_qkvpacked_func
+    has_flash_attn = True
+except:
+    print('FlashAttention2 is not installed.')
+    has_flash_attn = False
+
+logger = logging.get_logger(__name__)
+
+
+class FlashAttention(nn.Module):
+    """Implement the scaled dot product attention with softmax.
+    Arguments
+    ---------
+        softmax_scale: The temperature to use for the softmax attention.
+                      (default: 1/sqrt(d_keys) where d_keys is computed at
+                      runtime)
+        attention_dropout: The dropout rate to apply to the attention
+                           (default: 0.0)
+    """
+
+    def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):
+        super().__init__()
+        self.softmax_scale = softmax_scale
+        self.dropout_p = attention_dropout
+
+    def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens=None,
+                max_s=None, need_weights=False):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None
+                if unpadded: (nnz, 3, h, d)
+            key_padding_mask: a bool tensor of shape (B, S)
+        """
+        assert not need_weights
+        assert qkv.dtype in [torch.float16, torch.bfloat16]
+        assert qkv.is_cuda
+
+        if cu_seqlens is None:
+            batch_size = qkv.shape[0]
+            seqlen = qkv.shape[1]
+            if key_padding_mask is None:
+                qkv = rearrange(qkv, 'b s ... -> (b s) ...')
+                max_s = seqlen
+                cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32,
+                                          device=qkv.device)
+                output = flash_attn_varlen_qkvpacked_func(
+                    qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
+            else:
+                nheads = qkv.shape[-2]
+                x = rearrange(qkv, 'b s three h d -> b s (three h d)')
+                x_unpad, indices, cu_seqlens, max_s = unpad_input(x, key_padding_mask)
+                x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)
+                output_unpad = flash_attn_varlen_qkvpacked_func(
+                    x_unpad, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),
+                                             indices, batch_size, seqlen),
+                                   'b s (h d) -> b s h d', h=nheads)
+        else:
+            assert max_s is not None
+            output = flash_attn_varlen_qkvpacked_func(
+                qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                softmax_scale=self.softmax_scale, causal=causal
+            )
+
+        return output, None
+
+
+class InternRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+try:
+    from apex.normalization import FusedRMSNorm
+
+    InternRMSNorm = FusedRMSNorm  # noqa
+
+    logger.info('Discovered apex.normalization.FusedRMSNorm - will use it instead of InternRMSNorm')
+except ImportError:
+    # using the normal InternRMSNorm
+    pass
+except Exception:
+    logger.warning('discovered apex but it failed to load, falling back to InternRMSNorm')
+    pass
+
+
+NORM2FN = {
+    'rms_norm': InternRMSNorm,
+    'layer_norm': nn.LayerNorm,
+}
+
+
+class InternVisionEmbeddings(nn.Module):
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.class_embedding = nn.Parameter(
+            torch.randn(1, 1, self.embed_dim),
+        )
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches + 1
+
+        self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim))
+
+    def _get_pos_embed(self, pos_embed, H, W):
+        target_dtype = pos_embed.dtype
+        pos_embed = pos_embed.float().reshape(
+            1, self.image_size // self.patch_size, self.image_size // self.patch_size, -1).permute(0, 3, 1, 2)
+        pos_embed = F.interpolate(pos_embed, size=(H, W), mode='bicubic', align_corners=False). \
+            reshape(1, -1, H * W).permute(0, 2, 1).to(target_dtype)
+        return pos_embed
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, channel, width, height]
+        batch_size, _, height, width = patch_embeds.shape
+        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+        class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype)
+        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+        position_embedding = torch.cat([
+            self.position_embedding[:, :1, :],
+            self._get_pos_embed(self.position_embedding[:, 1:, :], height, width)
+        ], dim=1)
+        embeddings = embeddings + position_embedding.to(target_dtype)
+        return embeddings
+
+
+class InternAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.use_flash_attn = config.use_flash_attn and has_flash_attn
+        if config.use_flash_attn and not has_flash_attn:
+            print('Warning: Flash Attention is not available, use_flash_attn is set to False.')
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f'embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:'
+                f' {self.num_heads}).'
+            )
+
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias)
+        self.attn_drop = nn.Dropout(config.attention_dropout)
+        self.proj_drop = nn.Dropout(config.dropout)
+
+        self.qk_normalization = config.qk_normalization
+
+        if self.qk_normalization:
+            self.q_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+            self.k_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+        if self.use_flash_attn:
+            self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def _naive_attn(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
+
+        if self.qk_normalization:
+            B_, H_, N_, D_ = q.shape
+            q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
+            k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
+
+        attn = ((q * self.scale) @ k.transpose(-2, -1))
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def _flash_attn(self, x, key_padding_mask=None, need_weights=False):
+        qkv = self.qkv(x)
+        qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)
+
+        if self.qk_normalization:
+            q, k, v = qkv.unbind(2)
+            q = self.q_norm(q.flatten(-2, -1)).view(q.shape)
+            k = self.k_norm(k.flatten(-2, -1)).view(k.shape)
+            qkv = torch.stack([q, k, v], dim=2)
+
+        context, _ = self.inner_attn(
+            qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=False
+        )
+        outs = self.proj(rearrange(context, 'b s h d -> b s (h d)'))
+        outs = self.proj_drop(outs)
+        return outs
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        x = self._naive_attn(hidden_states) if not self.use_flash_attn else self._flash_attn(hidden_states)
+        return x
+
+
+class InternMLP(nn.Module):
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.act = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class InternVisionEncoderLayer(nn.Module):
+    def __init__(self, config: InternVisionConfig, drop_path_rate: float):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.norm_type = config.norm_type
+
+        self.attn = InternAttention(config)
+        self.mlp = InternMLP(config)
+        self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
+        self.norm2 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
+
+        self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
+        self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
+        self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+        self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+    ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)`
+        """
+        hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states).to(hidden_states.dtype)) * self.ls1)
+
+        hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states).to(hidden_states.dtype)) * self.ls2)
+
+        return hidden_states
+
+
+class InternVisionEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`InternEncoderLayer`].
+
+    Args:
+        config (`InternConfig`):
+            The corresponding vision configuration for the `InternEncoder`.
+    """
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]
+        self.layers = nn.ModuleList([
+            InternVisionEncoderLayer(config, dpr[idx]) for idx in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = True
+
+    def forward(
+            self,
+            inputs_embeds,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Embedded representation of the inputs. Should be float, not int tokens.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        encoder_states = () if output_hidden_states else None
+        hidden_states = inputs_embeds
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    encoder_layer,
+                    hidden_states)
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                )
+            hidden_states = layer_outputs
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states
+        )
+
+
+class InternVisionModel(PreTrainedModel):
+    main_input_name = 'pixel_values'
+    _supports_flash_attn_2 = True
+    config_class = InternVisionConfig
+    _no_split_modules = ['InternVisionEncoderLayer']
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = InternVisionEmbeddings(config)
+        self.encoder = InternVisionEncoder(config)
+
+    def resize_pos_embeddings(self, old_size, new_size, patch_size):
+        pos_emb = self.embeddings.position_embedding
+        _, num_positions, embed_dim = pos_emb.shape
+        cls_emb = pos_emb[:, :1, :]
+        pos_emb = pos_emb[:, 1:, :].reshape(1, old_size // patch_size, old_size // patch_size, -1).permute(0, 3, 1, 2)
+        pos_emb = F.interpolate(pos_emb.float(), size=new_size // patch_size, mode='bicubic', align_corners=False)
+        pos_emb = pos_emb.to(cls_emb.dtype).reshape(1, embed_dim, -1).permute(0, 2, 1)
+        pos_emb = torch.cat([cls_emb, pos_emb], dim=1)
+        self.embeddings.position_embedding = nn.Parameter(pos_emb)
+        self.embeddings.image_size = new_size
+        logger.info('Resized position embeddings from {} to {}'.format(old_size, new_size))
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def forward(
+            self,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            pixel_embeds: Optional[torch.FloatTensor] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if pixel_values is None and pixel_embeds is None:
+            raise ValueError('You have to specify pixel_values or pixel_embeds')
+
+        if pixel_embeds is not None:
+            hidden_states = pixel_embeds
+        else:
+            if len(pixel_values.shape) == 4:
+                hidden_states = self.embeddings(pixel_values)
+            else:
+                raise ValueError(f'wrong pixel_values size: {pixel_values.shape}')
+        encoder_outputs = self.encoder(
+            inputs_embeds=hidden_states,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        last_hidden_state = encoder_outputs.last_hidden_state
+        pooled_output = last_hidden_state[:, 0, :]
+
+        if not return_dict:
+            return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=last_hidden_state,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )

native_backbone/modeling_internlm2.py

@@ -0,0 +1,1456 @@
+# Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on transformers/src/transformers/models/llama/modeling_llama.py
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch InternLM2 model."""
+import math
+import queue
+import threading
+import warnings
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutputWithPast,
+                                           CausalLMOutputWithPast,
+                                           SequenceClassifierOutputWithPast)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import (add_start_docstrings,
+                                add_start_docstrings_to_model_forward, logging,
+                                replace_return_docstrings)
+
+try:
+    from transformers.generation.streamers import BaseStreamer
+except:  # noqa # pylint: disable=bare-except
+    BaseStreamer = None
+
+from .configuration_internlm2 import InternLM2Config
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = 'InternLM2Config'
+
+flash_attn_func, flash_attn_varlen_func = None, None
+pad_input, index_first_axis, unpad_input = None, None, None
+try:
+    from flash_attn import flash_attn_func as _flash_attn_func
+    from flash_attn import flash_attn_varlen_func as _flash_attn_varlen_func
+    from flash_attn.bert_padding import index_first_axis as _index_first_axis
+    from flash_attn.bert_padding import pad_input as _pad_input
+    from flash_attn.bert_padding import unpad_input as _unpad_input
+
+    flash_attn_func, flash_attn_varlen_func = _flash_attn_func, _flash_attn_varlen_func
+    pad_input, index_first_axis, unpad_input = _pad_input, _index_first_axis, _unpad_input
+    has_flash_attn = True
+except:
+    has_flash_attn = False
+
+
+def _import_flash_attn():
+    global flash_attn_func, flash_attn_varlen_func
+    global pad_input, index_first_axis, unpad_input
+    try:
+        from flash_attn import flash_attn_func as _flash_attn_func
+        from flash_attn import \
+            flash_attn_varlen_func as _flash_attn_varlen_func
+        from flash_attn.bert_padding import \
+            index_first_axis as _index_first_axis
+        from flash_attn.bert_padding import pad_input as _pad_input
+        from flash_attn.bert_padding import unpad_input as _unpad_input
+        flash_attn_func, flash_attn_varlen_func = _flash_attn_func, _flash_attn_varlen_func
+        pad_input, index_first_axis, unpad_input = _pad_input, _index_first_axis, _unpad_input
+    except ImportError:
+        raise ImportError('flash_attn is not installed.')
+
+
+# Copied from transformers.models.llama.modeling_llama._get_unpad_data
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+# Copied from transformers.models.bart.modeling_bart._make_causal_mask
+def _make_causal_mask(
+        input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
+):
+    """
+    Make causal mask used for bi-directional self-attention.
+    """
+    bsz, tgt_len = input_ids_shape
+    mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min, device=device), device=device)
+    mask_cond = torch.arange(mask.size(-1), device=device)
+    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
+    mask = mask.to(dtype)
+
+    if past_key_values_length > 0:
+        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
+    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
+
+
+# Copied from transformers.models.bart.modeling_bart._expand_mask
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
+    """
+    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
+    """
+    bsz, src_len = mask.size()
+    tgt_len = tgt_len if tgt_len is not None else src_len
+
+    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+
+    inverted_mask = 1.0 - expanded_mask
+
+    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->InternLM2
+class InternLM2RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        InternLM2RMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+try:
+    from functools import partial
+
+    from apex.normalization import FusedRMSNorm
+    InternLM2RMSNorm = partial(FusedRMSNorm, eps=1e-6)   # noqa
+    print('Discovered apex.normalization.FusedRMSNorm - will use it instead of InternLM2RMSNorm')
+except ImportError:
+    # using the normal LlamaRMSNorm
+    pass
+except Exception:
+    print('discovered apex but it failed to load, falling back to InternLM2RMSNorm')
+    pass
+
+
+# Copied from transformers.model.llama.modeling_llama.LlamaRotaryEmbedding with Llama->InternLM2
+class InternLM2RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        self.inv_freq = None
+        # inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+        # self.register_buffer('inv_freq', inv_freq, persistent=False)
+
+        self.max_seq_len_cached = -1 
+        # Build here to make `torch.jit.trace` work.
+        # self._set_cos_sin_cache(
+        #     seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
+        # )
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        if self.inv_freq is None:
+            inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+            del self.inv_freq
+            self.register_buffer('inv_freq', inv_freq, persistent=False)
+
+
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device).to(dtype=self.inv_freq.dtype)
+
+        # freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        freqs = torch.outer(t, self.inv_freq.to(device=t.device))
+
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)
+        self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
+
+        return (
+            self.cos_cached[:seq_len].to(dtype=x.dtype),
+            self.sin_cached[:seq_len].to(dtype=x.dtype),
+        )
+
+
+# Copied from transformers.model.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding with Llama->InternLM2
+class InternLM2LinearScalingRotaryEmbedding(InternLM2RotaryEmbedding):
+    """InternLM2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        if self.inv_freq is None:
+            inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+            del self.inv_freq
+            self.register_buffer('inv_freq', inv_freq, persistent=False)
+        
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device).to(dtype=self.inv_freq.dtype)
+        t = t / self.scaling_factor
+
+        # freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        freqs = torch.outer(t, self.inv_freq.to(device=t.device))
+
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)
+        self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)
+
+
+# Copied from transformers.model.llama.modeling_llama.LlamaDynamicNTKScalingRotaryEmbedding with Llama->InternLM2
+class InternLM2DynamicNTKScalingRotaryEmbedding(InternLM2RotaryEmbedding):
+    """InternLM2RotaryEmbedding extended with Dynamic NTK scaling.
+    Credits to the Reddit users /u/bloc97 and /u/emozilla.
+    """
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        if self.inv_freq is None:
+            inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim))
+            del self.inv_freq
+            self.register_buffer('inv_freq', inv_freq, persistent=False)
+
+
+        self.max_seq_len_cached = seq_len
+
+        if seq_len > self.max_position_embeddings:
+            base = self.base * (
+                    (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
+            ) ** (self.dim / (self.dim - 2))
+            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+            self.register_buffer('inv_freq', inv_freq, persistent=False)
+
+        t = torch.arange(self.max_seq_len_cached, device=device).to(dtype=self.inv_freq.dtype)
+
+        # freqs = torch.einsum('i,j->ij', t, self.inv_freq)
+        freqs = torch.outer(t, self.inv_freq.to(device=t.device))
+        
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)
+        self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)
+
+
+# Copied from transformers.model.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2:]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+
+# Copied from transformers.model.llama.modeling_llama.apply_rotary_pos_emb; float 
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors."""
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim).float()
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim).float()
+    q_dtype, k_dtype = q.dtype, k.dtype
+    q, k = q.float(), k.float()
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed.to(dtype=q_dtype), k_embed.to(dtype=k_dtype)
+
+
+class InternLM2MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.w2(self.act_fn(self.w1(x)) * self.w3(x))
+
+        return down_proj
+
+
+# Copied from transformers.model.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+# Modified from transformers.model.llama.modeling_llama.LlamaAttention
+class InternLM2Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: InternLM2Config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.is_causal = True
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f'hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}'
+                f' and `num_heads`: {self.num_heads}).'
+            )
+
+        self.wqkv = nn.Linear(
+            self.hidden_size,
+            (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
+            bias=config.bias,
+        )
+
+        self.wo = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias)
+        self._init_rope()
+
+    def _init_rope(self):
+        if self.config.rope_scaling is None:
+            self.rotary_emb = InternLM2RotaryEmbedding(
+                self.head_dim,
+                max_position_embeddings=self.max_position_embeddings,
+                base=self.config.rope_theta,
+            )
+        else:
+            scaling_type = self.config.rope_scaling['type']
+            scaling_factor = self.config.rope_scaling['factor']
+            if scaling_type == 'dynamic':
+                self.rotary_emb = InternLM2DynamicNTKScalingRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    base=self.config.rope_theta,
+                    scaling_factor=scaling_factor,
+                )
+            elif scaling_type == 'linear':
+                self.rotary_emb = InternLM2LinearScalingRotaryEmbedding(
+                    self.head_dim,
+                    max_position_embeddings=self.max_position_embeddings,
+                    base=self.config.rope_theta,
+                    scaling_factor=scaling_factor,
+                )
+            else:
+                raise ValueError("Currently we only support rotary embedding's type being 'dynamic' or 'linear'.")
+        return self.rotary_emb
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_value: Optional[Tuple[torch.Tensor]] = None,
+            output_attentions: bool = False,
+            use_cache: bool = False,
+            **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        if 'padding_mask' in kwargs:
+            warnings.warn(
+                'Passing `padding_mask` is deprecated and will be removed in v4.37. '
+                'Please make sure use `attention_mask` instead.`'
+            )
+
+        bsz, q_len, _ = hidden_states.size()
+
+        qkv_states = self.wqkv(hidden_states)
+
+        qkv_states = rearrange(
+            qkv_states,
+            'b q (h gs d) -> b q h gs d',
+            gs=2 + self.num_key_value_groups,
+            d=self.head_dim,
+        )
+
+        query_states = qkv_states[..., : self.num_key_value_groups, :]
+        query_states = rearrange(query_states, 'b q h gs d -> b q (h gs) d')
+        key_states = qkv_states[..., -2, :]
+        value_states = qkv_states[..., -1, :]
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        key_states = repeat_kv(key_states, self.num_key_value_groups)
+        value_states = repeat_kv(value_states, self.num_key_value_groups)
+
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f'Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is'
+                f' {attn_weights.size()}'
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f'Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}'
+                )
+            attn_weights = attn_weights + attention_mask
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f'`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is'
+                f' {attn_output.size()}'
+            )
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+        attn_output = self.wo(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+
+# Modified from transformers.model.llama.modeling_llama.InternLM2FlashAttention2
+class InternLM2FlashAttention2(InternLM2Attention):
+    """
+    InternLM2 flash attention module. This module inherits from `InternLM2Attention` as the weights of the module stays
+    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
+    flash attention and deal with padding tokens in case the input contains any of them.
+    """
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.LongTensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_value: Optional[Tuple[torch.Tensor]] = None,
+            output_attentions: bool = False,
+            use_cache: bool = False,
+            **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        # InternLM2FlashAttention2 attention does not support output_attentions
+        if 'padding_mask' in kwargs:
+            warnings.warn(
+                'Passing `padding_mask` is deprecated and will be removed in v4.37. '
+                'Please make sure use `attention_mask` instead.`'
+            )
+
+            # overwrite attention_mask with padding_mask
+            attention_mask = kwargs.pop('padding_mask')
+
+        output_attentions = False
+
+        bsz, q_len, _ = hidden_states.size()
+
+        qkv_states = self.wqkv(hidden_states)
+
+        qkv_states = rearrange(
+            qkv_states,
+            'b q (h gs d) -> b q h gs d',
+            gs=2 + self.num_key_value_groups,
+            d=self.head_dim,
+        )
+
+        query_states = qkv_states[..., : self.num_key_value_groups, :]
+        query_states = rearrange(query_states, 'b q h gs d -> b q (h gs) d')
+        key_states = qkv_states[..., -2, :]
+        value_states = qkv_states[..., -1, :]
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        kv_seq_len = key_states.shape[-2]
+        if past_key_value is not None:
+            kv_seq_len += past_key_value[0].shape[-2]
+
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        attn_output = self._flash_attention_forward(
+            query_states, key_states, value_states, attention_mask, q_len
+        )
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+        attn_output = self.wo(attn_output)
+
+        if not output_attentions:
+            attn_weights = None
+
+        return attn_output, attn_weights, past_key_value
+
+    def _flash_attention_forward(
+            self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
+    ):
+        """
+        Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
+        first unpad the input, then computes the attention scores and pad the final attention scores.
+
+        Args:
+            query_states (`torch.Tensor`):
+                Input query states to be passed to Flash Attention API
+            key_states (`torch.Tensor`):
+                Input key states to be passed to Flash Attention API
+            value_states (`torch.Tensor`):
+                Input value states to be passed to Flash Attention API
+            attention_mask (`torch.Tensor`):
+                The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
+                position of padding tokens and 1 for the position of non-padding tokens.
+            dropout (`int`, *optional*):
+                Attention dropout
+            softmax_scale (`float`, *optional*):
+                The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
+        """
+        # Contains at least one padding token in the sequence
+        causal = self.is_causal and query_length != 1
+        if attention_mask is not None:
+            batch_size = query_states.shape[0]
+            query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._unpad_input(
+                query_states, key_states, value_states, attention_mask, query_length
+            )
+
+            cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+            max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+
+            attn_output_unpad = flash_attn_varlen_func(
+                query_states,
+                key_states,
+                value_states,
+                cu_seqlens_q=cu_seqlens_q,
+                cu_seqlens_k=cu_seqlens_k,
+                max_seqlen_q=max_seqlen_in_batch_q,
+                max_seqlen_k=max_seqlen_in_batch_k,
+                dropout_p=dropout,
+                softmax_scale=softmax_scale,
+                causal=causal,
+            )
+
+            attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)
+        else:
+            attn_output = flash_attn_func(
+                query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
+            )
+
+        return attn_output
+
+    def _unpad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
+        indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
+        batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
+
+        key_layer = index_first_axis(
+            key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+        value_layer = index_first_axis(
+            value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
+        )
+
+        if query_length == kv_seq_len:
+            query_layer = index_first_axis(
+                query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
+            )
+            cu_seqlens_q = cu_seqlens_k
+            max_seqlen_in_batch_q = max_seqlen_in_batch_k
+            indices_q = indices_k
+        elif query_length == 1:
+            max_seqlen_in_batch_q = 1
+            cu_seqlens_q = torch.arange(
+                batch_size + 1, dtype=torch.int32, device=query_layer.device
+            )  # There is a memcpy here, that is very bad.
+            indices_q = cu_seqlens_q[:-1]
+            query_layer = query_layer.squeeze(1)
+        else:
+            # The -q_len: slice assumes left padding.
+            attention_mask = attention_mask[:, -query_length:]
+            query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask)
+
+        return (
+            query_layer,
+            key_layer,
+            value_layer,
+            indices_q.to(torch.int64),
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        )
+
+
+INTERNLM2_ATTENTION_CLASSES = {
+    'eager': InternLM2Attention,
+    'flash_attention_2': InternLM2FlashAttention2,
+}
+
+
+# Modified from transformers.model.llama.modeling_llama.LlamaDecoderLayer
+class InternLM2DecoderLayer(nn.Module):
+    def __init__(self, config: InternLM2Config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.attention = INTERNLM2_ATTENTION_CLASSES[config.attn_implementation](config=config)
+
+        self.feed_forward = InternLM2MLP(config)
+        self.attention_norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.ffn_norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_value: Optional[Tuple[torch.Tensor]] = None,
+            output_attentions: Optional[bool] = False,
+            use_cache: Optional[bool] = False,
+            **kwargs,
+    ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*):
+                attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+                query_sequence_length, key_sequence_length)` if default attention is used.
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            use_cache (`bool`, *optional*):
+                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
+                (see `past_key_values`).
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
+        """
+        if 'padding_mask' in kwargs:
+            warnings.warn(
+                'Passing `padding_mask` is deprecated and will be removed in v4.37. '
+                'Please make sure use `attention_mask` instead.`'
+            )
+
+        residual = hidden_states
+
+        hidden_states = self.attention_norm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            use_cache=use_cache,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.ffn_norm(hidden_states)
+        hidden_states = self.feed_forward(hidden_states)
+        hidden_states = residual + hidden_states
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        return outputs
+
+
+InternLM2_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`InternLM2Config`]):
+            Model configuration class with all the parameters of the model. Initializing with a config file does not
+            load the weights associated with the model, only the configuration. Check out the
+            [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaPreTrainedModel with Llama->InternLM2
+@add_start_docstrings(
+    'The bare InternLM2 Model outputting raw hidden-states without any specific head on top.',
+    InternLM2_START_DOCSTRING,
+)
+class InternLM2PreTrainedModel(PreTrainedModel):
+    config_class = InternLM2Config
+    base_model_prefix = 'model'
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['InternLM2DecoderLayer']
+    _skip_keys_device_placement = 'past_key_values'
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+InternLM2_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or
+            when `config.use_cache=True`):
+            Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape
+            `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape
+            `(batch_size, num_heads, decoder_sequence_length, embed_size_per_head)`.
+
+            Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+# Modified from transformers.model.llama.modeling_llama.LlamaModel
+@add_start_docstrings(
+    'The bare InternLM2 Model outputting raw hidden-states without any specific head on top.',
+    InternLM2_START_DOCSTRING,
+)
+class InternLM2Model(InternLM2PreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`InternLM2DecoderLayer`]
+
+    Args:
+        config: InternLM2Config
+    """
+
+    _auto_class = 'AutoModel'
+
+    def __init__(self, config: InternLM2Config):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.config = config
+        if not has_flash_attn:
+            self.config.attn_implementation = 'eager'
+            print('Warning: Flash attention is not available, using eager attention instead.')
+
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+
+        self.layers = nn.ModuleList([InternLM2DecoderLayer(config) for _ in range(config.num_hidden_layers)])
+        self.norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.tok_embeddings
+
+    def set_input_embeddings(self, value):
+        self.tok_embeddings = value
+
+    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
+        # create causal mask
+        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+        combined_attention_mask = None
+        if input_shape[-1] > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape,
+                inputs_embeds.dtype,
+                device=inputs_embeds.device,
+                past_key_values_length=past_key_values_length,
+            )
+
+        if attention_mask is not None:
+            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
+            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
+                inputs_embeds.device
+            )
+            combined_attention_mask = (
+                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
+            )
+
+        return combined_attention_mask
+
+    @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPast]:
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if self.config.attn_implementation == 'flash_attention_2':
+            _import_flash_attn()
+
+        # retrieve input_ids and inputs_embeds
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time')
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError('You have to specify either input_ids or inputs_embeds')
+
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+        if past_key_values is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+
+        if position_ids is None:
+            device = input_ids.device if input_ids is not None else inputs_embeds.device
+            position_ids = torch.arange(
+                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
+            )
+            position_ids = position_ids.unsqueeze(0)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.tok_embeddings(input_ids)
+
+        if self.config.attn_implementation == 'flash_attention_2':
+            # 2d mask is passed through the layers
+            attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
+        else:
+            if attention_mask is None:
+                attention_mask = torch.ones(
+                    (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+                )
+            attention_mask = self._prepare_decoder_attention_mask(
+                attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
+            )
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    '`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...'
+                )
+                use_cache = False
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = () if use_cache else None
+
+        for idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            past_key_value = past_key_values[idx] if past_key_values is not None else None
+
+            if self.gradient_checkpointing and self.training:
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # None for past_key_value
+                        return module(*inputs, output_attentions, None)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(decoder_layer),
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    None,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_value,
+                    output_attentions=output_attentions,
+                    use_cache=use_cache,
+                )
+
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = next_decoder_cache if use_cache else None
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+        )
+
+
+# Modified from transformers.model.llama.modeling_llama.LlamaForCausalLM
+class InternLM2ForCausalLM(InternLM2PreTrainedModel):
+    _auto_class = 'AutoModelForCausalLM'
+
+    _tied_weights_keys = ['output.weight']
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = InternLM2Model(config)
+        self.vocab_size = config.vocab_size
+        self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.tok_embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.tok_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.output
+
+    def set_output_embeddings(self, new_embeddings):
+        self.output = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
+    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, InternLM2ForCausalLM
+
+        >>> model = InternLM2ForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = outputs[0]
+        logits = self.output(hidden_states)
+        logits = logits.float()
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+        output = CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+        output['logits'] = output['logits'].to(device)
+        return output
+
+    def prepare_inputs_for_generation(
+            self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+    ):
+        if past_key_values is not None:
+            past_length = past_key_values[0][0].shape[2]
+
+            # Some generation methods already pass only the last input ID
+            if input_ids.shape[1] > past_length:
+                remove_prefix_length = past_length
+            else:
+                # Default to old behavior: keep only final ID
+                remove_prefix_length = input_ids.shape[1] - 1
+
+            input_ids = input_ids[:, remove_prefix_length:]
+
+        position_ids = kwargs.get('position_ids', None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1]:]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {'inputs_embeds': inputs_embeds}
+        else:
+            model_inputs = {'input_ids': input_ids}
+
+        model_inputs.update(
+            {
+                'position_ids': position_ids,
+                'past_key_values': past_key_values,
+                'use_cache': kwargs.get('use_cache'),
+                'attention_mask': attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
+            )
+        return reordered_past
+
+    def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = [], meta_instruction=''):
+        if tokenizer.add_bos_token:
+            prompt = ''
+        else:
+            prompt = tokenizer.bos_token
+        if meta_instruction:
+            prompt += f"""<|im_start|>system\n{meta_instruction}<|im_end|>\n"""
+        for record in history:
+            prompt += f"""<|im_start|>user\n{record[0]}<|im_end|>\n<|im_start|>assistant\n{record[1]}<|im_end|>\n"""
+        prompt += f"""<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"""
+        return tokenizer([prompt], return_tensors='pt')
+
+    @torch.no_grad()
+    def chat(
+            self,
+            tokenizer,
+            query: str,
+            history: List[Tuple[str, str]] = [],
+            streamer: Optional[BaseStreamer] = None,
+            max_new_tokens: int = 1024,
+            do_sample: bool = True,
+            temperature: float = 0.8,
+            top_p: float = 0.8,
+            meta_instruction: str = 'You are an AI assistant whose name is InternLM (书生·浦语).\n'
+                                    '- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory (上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n'
+                                    '- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such as English and 中文.',
+            **kwargs,
+    ):
+        inputs = self.build_inputs(tokenizer, query, history, meta_instruction)
+        inputs = {k: v.to(self.device) for k, v in inputs.items() if torch.is_tensor(v)}
+        # also add end-of-assistant token in eos token id to avoid unnecessary generation
+        eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(['<|im_end|>'])[0]]
+        outputs = self.generate(
+            **inputs,
+            streamer=streamer,
+            max_new_tokens=max_new_tokens,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_p=top_p,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+        outputs = outputs[0].cpu().tolist()[len(inputs['input_ids'][0]):]
+        response = tokenizer.decode(outputs, skip_special_tokens=True)
+        response = response.split('<|im_end|>')[0]
+        history = history + [(query, response)]
+        return response, history
+
+    @torch.no_grad()
+    def stream_chat(
+            self,
+            tokenizer,
+            query: str,
+            history: List[Tuple[str, str]] = [],
+            max_new_tokens: int = 1024,
+            do_sample: bool = True,
+            temperature: float = 0.8,
+            top_p: float = 0.8,
+            **kwargs,
+    ):
+        """
+        Return a generator in format: (response, history)
+        Eg.
+        ('你好，有什么可以帮助您的吗', [('你好', '你好，有什么可以帮助您的吗')])
+        ('你好，有什么可以帮助您的吗？', [('你好', '你好，有什么可以帮助您的吗？')])
+        """
+        if BaseStreamer is None:
+            raise ModuleNotFoundError(
+                'The version of `transformers` is too low. Please make sure '
+                'that you have installed `transformers>=4.28.0`.'
+            )
+
+        response_queue = queue.Queue(maxsize=20)
+
+        class ChatStreamer(BaseStreamer):
+            def __init__(self, tokenizer) -> None:
+                super().__init__()
+                self.tokenizer = tokenizer
+                self.queue = response_queue
+                self.query = query
+                self.history = history
+                self.response = ''
+                self.cache = []
+                self.received_inputs = False
+                self.queue.put((self.response, history + [(self.query, self.response)]))
+
+            def put(self, value):
+                if len(value.shape) > 1 and value.shape[0] > 1:
+                    raise ValueError('ChatStreamer only supports batch size 1')
+                elif len(value.shape) > 1:
+                    value = value[0]
+
+                if not self.received_inputs:
+                    # The first received value is input_ids, ignore here
+                    self.received_inputs = True
+                    return
+
+                self.cache.extend(value.tolist())
+                token = self.tokenizer.decode(self.cache, skip_special_tokens=True)
+                if token.strip() != '<|im_end|>':
+                    self.response = self.response + token
+                    history = self.history + [(self.query, self.response)]
+                    self.queue.put((self.response, history))
+                    self.cache = []
+                else:
+                    self.end()
+
+            def end(self):
+                self.queue.put(None)
+
+        def stream_producer():
+            return self.chat(
+                tokenizer=tokenizer,
+                query=query,
+                streamer=ChatStreamer(tokenizer=tokenizer),
+                history=history,
+                max_new_tokens=max_new_tokens,
+                do_sample=do_sample,
+                temperature=temperature,
+                top_p=top_p,
+                **kwargs,
+            )
+
+        def consumer():
+            producer = threading.Thread(target=stream_producer)
+            producer.start()
+            while True:
+                res = response_queue.get()
+                if res is None:
+                    return
+                yield res
+
+        return consumer()
+
+
+# Copied from transformers.model.llama.modeling_llama.LlamaForSequenceClassification with Llama->InternLM2
+@add_start_docstrings(
+    """
+    The InternLM2 Model transformer with a sequence classification head on top (linear layer).
+
+    [`InternLM2ForSequenceClassification`] uses the last token in order to do the classification,
+    as other causal models (e.g. GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """,
+    InternLM2_START_DOCSTRING,
+)
+class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = InternLM2Model(config)
+        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.tok_embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.tok_embeddings = value
+
+    @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
+    def forward(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        transformer_outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError('Cannot handle batch sizes > 1 if no padding token is defined.')
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1).to(
+                    logits.device
+                )
+            else:
+                sequence_lengths = -1
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = 'regression'
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = 'single_label_classification'
+                else:
+                    self.config.problem_type = 'multi_label_classification'
+
+            if self.config.problem_type == 'regression':
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == 'single_label_classification':
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == 'multi_label_classification':
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )

native_backbone/modeling_internvl_chat.py

@@ -0,0 +1,363 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2024 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+import warnings
+from typing import Any, List, Optional, Tuple, Union
+
+import torch.utils.checkpoint
+import transformers
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import (AutoModel, GenerationConfig, LlamaForCausalLM,
+                          LlamaTokenizer, Qwen2ForCausalLM)
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import ModelOutput, logging
+
+from .configuration_internvl_chat import InternVLChatConfig
+from .conversation import get_conv_template
+from .modeling_intern_vit import InternVisionModel, has_flash_attn
+from .modeling_internlm2 import InternLM2ForCausalLM
+
+logger = logging.get_logger(__name__)
+
+
+def version_cmp(v1, v2, op='eq'):
+    import operator
+
+    from packaging import version
+    op_func = getattr(operator, op)
+    return op_func(version.parse(v1), version.parse(v2))
+
+
+class InternVLChatModel(PreTrainedModel):
+    config_class = InternVLChatConfig
+    main_input_name = 'pixel_values'
+    _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
+    _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'InternLM2DecoderLayer',
+                         'Qwen2DecoderLayer']
+
+    def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
+        super().__init__(config)
+
+        assert version_cmp(transformers.__version__, '4.36.2', 'ge')
+        image_size = config.force_image_size or config.vision_config.image_size
+        patch_size = config.vision_config.patch_size
+        self.patch_size = patch_size
+        self.select_layer = config.select_layer
+        self.template = config.template
+        self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
+        self.downsample_ratio = config.downsample_ratio
+        self.ps_version = config.ps_version
+        use_flash_attn = use_flash_attn if has_flash_attn else False
+        config.vision_config.use_flash_attn = True if use_flash_attn else False
+        config.llm_config.attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'
+
+        logger.info(f'num_image_token: {self.num_image_token}')
+        logger.info(f'ps_version: {self.ps_version}')
+        if vision_model is not None:
+            self.vision_model = vision_model
+        else:
+            self.vision_model = InternVisionModel(config.vision_config)
+        if language_model is not None:
+            self.language_model = language_model
+        else:
+            if config.llm_config.architectures[0] == 'LlamaForCausalLM':
+                self.language_model = LlamaForCausalLM(config.llm_config)
+            elif config.llm_config.architectures[0] == 'InternLM2ForCausalLM':
+                self.language_model = InternLM2ForCausalLM(config.llm_config)
+            elif config.llm_config.architectures[0] == 'Qwen2ForCausalLM':
+                self.language_model = Qwen2ForCausalLM(config.llm_config)
+            else:
+                raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')
+
+        vit_hidden_size = config.vision_config.hidden_size
+        llm_hidden_size = config.llm_config.hidden_size
+
+        self.mlp1 = nn.Sequential(
+            nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2),
+            nn.Linear(vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size),
+            nn.GELU(),
+            nn.Linear(llm_hidden_size, llm_hidden_size)
+        )
+
+        self.img_context_token_id = None
+        self.conv_template = get_conv_template(self.template)
+        self.system_message = self.conv_template.system_message
+
+    def forward(
+            self,
+            pixel_values: torch.FloatTensor,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            image_flags: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            labels: Optional[torch.LongTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        image_flags = image_flags.squeeze(-1)
+        input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        vit_embeds = self.extract_feature(pixel_values)
+        vit_embeds = vit_embeds[image_flags == 1]
+        vit_batch_size = pixel_values.shape[0]
+
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+
+        if torch.distributed.get_rank() == 0:
+            print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
+
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+        try:
+            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+        except Exception as e:
+            vit_embeds = vit_embeds.reshape(-1, C)
+            print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
+                  f'vit_embeds.shape={vit_embeds.shape}')
+            n_token = min(selected.sum(), vit_embeds.size(0))
+            input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]
+
+        input_embeds = input_embeds.reshape(B, N, C)
+
+        outputs = self.language_model(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        logits = outputs.logits
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.language_model.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def pixel_shuffle(self, x, scale_factor=0.5):
+        n, w, h, c = x.size()
+        # N, W, H, C --> N, W, H * scale, C // scale
+        x = x.view(n, w, int(h * scale_factor), int(c / scale_factor))
+        # N, W, H * scale, C // scale --> N, H * scale, W, C // scale
+        x = x.permute(0, 2, 1, 3).contiguous()
+        # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2)
+        x = x.view(n, int(h * scale_factor), int(w * scale_factor),
+                   int(c / (scale_factor * scale_factor)))
+        if self.ps_version == 'v1':
+            warnings.warn("In ps_version 'v1', the height and width have not been swapped back, "
+                          'which results in a transposed image.')
+        else:
+            x = x.permute(0, 2, 1, 3).contiguous()
+        return x
+
+    def extract_feature(self, pixel_values):
+        if self.select_layer == -1:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values,
+                output_hidden_states=False,
+                return_dict=True).last_hidden_state
+        else:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values,
+                output_hidden_states=True,
+                return_dict=True).hidden_states[self.select_layer]
+        vit_embeds = vit_embeds[:, 1:, :]
+
+        h = w = int(vit_embeds.shape[1] ** 0.5)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
+        vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio)
+        vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1])
+        vit_embeds = self.mlp1(vit_embeds)
+        return vit_embeds
+
+    def batch_chat(self, tokenizer, pixel_values, questions, generation_config, num_patches_list=None,
+                   history=None, return_history=False, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',
+                   IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', verbose=False, image_counts=None):
+        if history is not None or return_history:
+            print('Now multi-turn chat is not supported in batch_chat.')
+            raise NotImplementedError
+
+        if image_counts is not None:
+            num_patches_list = image_counts
+            print('Warning: `image_counts` is deprecated. Please use `num_patches_list` instead.')
+
+        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+
+        if verbose and pixel_values is not None:
+            image_bs = pixel_values.shape[0]
+            print(f'dynamic ViT batch size: {image_bs}')
+
+        queries = []
+        for idx, num_patches in enumerate(num_patches_list):
+            question = questions[idx]
+            if pixel_values is not None and '<image>' not in question:
+                question = '<image>\n' + question
+            template = get_conv_template(self.template)
+            template.system_message = self.system_message
+            template.append_message(template.roles[0], question)
+            template.append_message(template.roles[1], None)
+            query = template.get_prompt()
+
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+            queries.append(query)
+
+        tokenizer.padding_side = 'left'
+        model_inputs = tokenizer(queries, return_tensors='pt', padding=True)
+        input_ids = model_inputs['input_ids'].cuda()
+        attention_mask = model_inputs['attention_mask'].cuda()
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
+        generation_config['eos_token_id'] = eos_token_id
+        generation_output = self.generate(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **generation_config
+        )
+        responses = tokenizer.batch_decode(generation_output, skip_special_tokens=True)
+        responses = [response.split(template.sep)[0].strip() for response in responses]
+        return responses
+
+    def chat(self, tokenizer, pixel_values, question, generation_config, history=None, return_history=False,
+             num_patches_list=None, IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>',
+             verbose=False):
+
+        if history is None and pixel_values is not None and '<image>' not in question:
+            question = '<image>\n' + question
+
+        if num_patches_list is None:
+            num_patches_list = [pixel_values.shape[0]] if pixel_values is not None else []
+        assert pixel_values is None or len(pixel_values) == sum(num_patches_list)
+
+        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+
+        template = get_conv_template(self.template)
+        template.system_message = self.system_message
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
+
+        history = [] if history is None else history
+        for (old_question, old_answer) in history:
+            template.append_message(template.roles[0], old_question)
+            template.append_message(template.roles[1], old_answer)
+        template.append_message(template.roles[0], question)
+        template.append_message(template.roles[1], None)
+        query = template.get_prompt()
+
+        if verbose and pixel_values is not None:
+            image_bs = pixel_values.shape[0]
+            print(f'dynamic ViT batch size: {image_bs}')
+
+        for num_patches in num_patches_list:
+            image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
+            query = query.replace('<image>', image_tokens, 1)
+
+        model_inputs = tokenizer(query, return_tensors='pt')
+        input_ids = model_inputs['input_ids'].cuda()
+        attention_mask = model_inputs['attention_mask'].cuda()
+        generation_config['eos_token_id'] = eos_token_id
+        generation_output = self.generate(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **generation_config
+        )
+        response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
+        response = response.split(template.sep.strip())[0].strip()
+        history.append((question, response))
+        if return_history:
+            return response, history
+        else:
+            query_to_print = query.replace(IMG_CONTEXT_TOKEN, '')
+            query_to_print = query_to_print.replace(f'{IMG_START_TOKEN}{IMG_END_TOKEN}', '<image>')
+            if verbose:
+                print(query_to_print, response)
+            return response
+
+    @torch.no_grad()
+    def generate(
+            self,
+            pixel_values: Optional[torch.FloatTensor] = None,
+            input_ids: Optional[torch.FloatTensor] = None,
+            attention_mask: Optional[torch.LongTensor] = None,
+            visual_features: Optional[torch.FloatTensor] = None,
+            generation_config: Optional[GenerationConfig] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            **generate_kwargs,
+    ) -> torch.LongTensor:
+
+        assert self.img_context_token_id is not None
+        if pixel_values is not None:
+            if visual_features is not None:
+                vit_embeds = visual_features
+            else:
+                vit_embeds = self.extract_feature(pixel_values)
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+            B, N, C = input_embeds.shape
+            input_embeds = input_embeds.reshape(B * N, C)
+
+            input_ids = input_ids.reshape(B * N)
+            selected = (input_ids == self.img_context_token_id)
+            assert selected.sum() != 0
+            input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
+
+            input_embeds = input_embeds.reshape(B, N, C)
+        else:
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        outputs = self.language_model.generate(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            output_hidden_states=output_hidden_states,
+            # return_dict=return_dict, # return_dict is not supported in transformers 4.44
+            use_cache=True,
+            **generate_kwargs,
+        )
+
+        return outputs
+
+    @property
+    def lm_head(self):
+        return self.language_model.get_output_embeddings()
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()

native_backbone/native-00001.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8595a8061a1eacd66c8e8872754ff18b36113c660d01ebfdf91cf9701477543
+size 4878234984

native_backbone/native-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66ebb3671f03e25fc350750513958e39078d851ac332aedb92cbaa5aba071bc6
+size 4882437624

native_backbone/native-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:186bf3a374834e1e573e9d8ff373e1dec200eca2b5dc23e2ddcdcc75cf4ecb5d
+size 4844680424

native_backbone/native-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a97f493300043c8b238f70fceaaab6b490f527ea2330cfabfe4a1c24aa23d185
+size 3672318056

native_backbone/special_tokens_map.json

@@ -0,0 +1,63 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|action_start|>",
+    "<|action_end|>",
+    "<|interpreter|>",
+    "<|plugin|>",
+    "<restate>",
+    "</restate>",
+    "<planning>",
+    "</planning>",
+    "<recollect>",
+    "</recollect>",
+    "<execution>",
+    "</execution>",
+    "<review>",
+    "</review>",
+    "<summarize>",
+    "</summarize>",
+    "<retry>",
+    "</retry>",
+    "<conclude>",
+    "</conclude>",
+    "<img>",
+    "</img>",
+    "<IMG_CONTEXT>",
+    "<quad>",
+    "</quad>",
+    "<ref>",
+    "</ref>",
+    "<box>",
+    "</box>"
+  ],
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

native_backbone/tokenization_internlm3.py

@@ -0,0 +1,294 @@
+import os
+from shutil import copyfile
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
+
+import sentencepiece as spm
+from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
+from transformers.utils import logging
+
+if TYPE_CHECKING:
+    from transformers.tokenization_utils_base import TextInput
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "tokenizer.model"}
+
+SPIECE_UNDERLINE = "▁"
+
+
+class InternLM3Tokenizer(PreTrainedTokenizer):
+    """
+    Construct a InternLM3 tokenizer. Based on byte-level Byte-Pair-Encoding. The default padding token is unset as there is
+    no padding token in the original model.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+        unk_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<unk>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        bos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"<s>"`):
+            The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+        eos_token (`str` or `tokenizers.AddedToken`, *optional*, defaults to `"</s>"`):
+            The end of sequence token.
+        pad_token (`str` or `tokenizers.AddedToken`, *optional*):
+            A special token used to make arrays of tokens the same size for batching purpose. Will then be ignored by
+            attention mechanisms or loss computation.
+        sp_model_kwargs (`Dict[str, Any]`, `Optional`, *optional*):
+            Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
+            SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
+            to set:
+
+            - `enable_sampling`: Enable subword regularization.
+            - `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
+
+              - `nbest_size = {0,1}`: No sampling is performed.
+              - `nbest_size > 1`: samples from the nbest_size results.
+              - `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
+                using forward-filtering-and-backward-sampling algorithm.
+
+            - `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
+              BPE-dropout.
+
+        add_bos_token (`bool`, *optional*, defaults to `True`):
+            Whether or not to add an `bos_token` at the start of sequences.
+        add_eos_token (`bool`, *optional*, defaults to `False`):
+            Whether or not to add an `eos_token` at the end of sequences.
+        clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+            Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
+            extra spaces.
+        use_default_system_prompt (`bool`, *optional*, defaults to `False`):
+            Whether or not the default system prompt for InternLM3 should be used.
+        spaces_between_special_tokens (`bool`, *optional*, defaults to `False`):
+            Whether or not to add spaces between special tokens.
+        spaces_for_interleaved_special_tokens (`bool`, *optional*, defaults to `False`):
+           Whether or not to add spaces between special tokens that are interleaved with normal tokens.
+        add_prefix_space (`bool`, *optional*, defaults to `True`):
+            Whether or not to add an initial space to the input. This allows to treat the leading word just as any
+            other word. Again, this should be set with `from_slow=True` to make sure it's taken into account.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        unk_token="<unk>",
+        bos_token="<s>",
+        eos_token="</s>",
+        pad_token=None,
+        sp_model_kwargs: Optional[Dict[str, Any]] = None,
+        add_bos_token=True,
+        add_eos_token=False,
+        clean_up_tokenization_spaces=False,
+        use_default_system_prompt=False,
+        spaces_between_special_tokens=False,
+        spaces_for_interleaved_special_tokens=False,
+        add_prefix_space=True,
+        **kwargs,
+    ):
+        self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
+        bos_token = AddedToken(bos_token, normalized=False, special=True) if isinstance(bos_token, str) else bos_token
+        eos_token = AddedToken(eos_token, normalized=False, special=True) if isinstance(eos_token, str) else eos_token
+        unk_token = AddedToken(unk_token, normalized=False, special=True) if isinstance(unk_token, str) else unk_token
+        pad_token = AddedToken(pad_token, normalized=False, special=True) if isinstance(pad_token, str) else pad_token
+
+        self.vocab_file = vocab_file
+        self.add_bos_token = add_bos_token
+        self.add_eos_token = add_eos_token
+        self.use_default_system_prompt = use_default_system_prompt
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.Load(vocab_file)
+        self.add_prefix_space = add_prefix_space
+        self.spaces_for_interleaved_special_tokens = spaces_for_interleaved_special_tokens
+
+        vocab_size = self.sp_model.get_piece_size()
+        self.decoder = {i: self.sp_model.id_to_piece(i) for i in range(vocab_size)}
+
+        super().__init__(
+            bos_token=bos_token,
+            eos_token=eos_token,
+            unk_token=unk_token,
+            pad_token=pad_token,
+            add_bos_token=add_bos_token,
+            add_eos_token=add_eos_token,
+            sp_model_kwargs=sp_model_kwargs,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            use_default_system_prompt=use_default_system_prompt,
+            spaces_between_special_tokens=spaces_between_special_tokens,
+            add_prefix_space=add_prefix_space,
+            **kwargs,
+        )
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        state["sp_model"] = None
+        state["sp_model_proto"] = self.sp_model.serialized_model_proto()
+        return state
+
+    def __setstate__(self, d):
+        self.__dict__.update(d)
+        self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
+        self.sp_model.LoadFromSerializedProto(self.sp_model_proto)
+
+    @property
+    def vocab_size(self):
+        """Returns vocab size"""
+        return self.sp_model.get_piece_size()
+
+    def get_vocab(self):
+        """Returns vocab as a dict"""
+        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
+        vocab.update(self.added_tokens_encoder)
+        return vocab
+
+    def tokenize(self, text: "TextInput", **kwargs) -> List[str]:
+        """
+        Args:
+            text: TextInput
+        Simply calls PreTrainedTokenizer's method
+        """
+        return super().tokenize(text, **kwargs)
+
+    def _tokenize(self, text, **kwargs):
+        """
+        Args:
+            text: TextInput
+        Returns a tokenized string. The Gemma tokenizer never adds a prefix space.
+        """
+        return self.sp_model.encode(text, out_type=str)
+
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.sp_model.piece_to_id(token)
+
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.decoder.get(index, "")
+
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        # since we manually add the prefix space, we have to remove it when decoding
+        if tokens[0].startswith(SPIECE_UNDERLINE) and self.add_prefix_space:
+            tokens[0] = tokens[0][1:]
+
+        current_sub_tokens = []
+        out_string = ""
+        prev_is_special = False
+        for i, token in enumerate(tokens):
+            # make sure that special tokens are not decoded using sentencepiece model
+            if token in self.all_special_tokens:
+                if not prev_is_special and i != 0 and self.spaces_for_interleaved_special_tokens:
+                    out_string += " "
+                out_string += self.sp_model.decode(current_sub_tokens) + token
+                prev_is_special = True
+                current_sub_tokens = []
+            else:
+                if (
+                    prev_is_special
+                    and i == 1
+                    and self.add_prefix_space
+                    and not token.startswith(SPIECE_UNDERLINE)
+                    and self.spaces_for_interleaved_special_tokens
+                ):
+                    out_string += " "
+                current_sub_tokens.append(token)
+                prev_is_special = False
+        out_string += self.sp_model.decode(current_sub_tokens)
+        return out_string
+
+    def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        """
+        Save the vocabulary and special tokens file to a directory.
+
+        Args:
+            save_directory (`str`):
+                The directory in which to save the vocabulary.
+
+        Returns:
+            `Tuple(str)`: Paths to the files saved.
+        """
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        out_vocab_file = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"])
+
+        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
+            copyfile(self.vocab_file, out_vocab_file)
+        elif not os.path.isfile(self.vocab_file):
+            with open(out_vocab_file, "wb") as fi:
+                content_spiece_model = self.sp_model.serialized_model_proto()
+                fi.write(content_spiece_model)
+
+        return (out_vocab_file,)
+
+    def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = bos_token_id + token_ids_0 + eos_token_id
+
+        if token_ids_1 is not None:
+            output = output + bos_token_id + token_ids_1 + eos_token_id
+
+        return output
+
+    def get_special_tokens_mask(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
+    ) -> List[int]:
+        """
+        Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+        special tokens using the tokenizer `prepare_for_model` method.
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of IDs.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+            already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether or not the token list is already formatted with special tokens for the model.
+
+        Returns:
+            `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+        """
+        if already_has_special_tokens:
+            return super().get_special_tokens_mask(token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True)
+
+        bos_token_id = [1] if self.add_bos_token else []
+        eos_token_id = [1] if self.add_eos_token else []
+
+        if token_ids_1 is None:
+            return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id
+        return bos_token_id + ([0] * len(token_ids_0)) + eos_token_id + bos_token_id + ([0] * len(token_ids_1)) + eos_token_id
+
+    def create_token_type_ids_from_sequences(self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None) -> List[int]:
+        """
+        Creates a mask from the two sequences passed to be used in a sequence-pair classification task. An ALBERT
+        sequence pair mask has the following format:
+
+        ```
+        0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1
+        | first sequence    | second sequence |
+        ```
+
+        if token_ids_1 is None, only returns the first portion of the mask (0s).
+
+        Args:
+            token_ids_0 (`List[int]`):
+                List of ids.
+            token_ids_1 (`List[int]`, *optional*):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            `List[int]`: List of [token type IDs](../glossary#token-type-ids) according to the given sequence(s).
+        """
+        bos_token_id = [self.bos_token_id] if self.add_bos_token else []
+        eos_token_id = [self.eos_token_id] if self.add_eos_token else []
+
+        output = [0] * len(bos_token_id + token_ids_0 + eos_token_id)
+
+        if token_ids_1 is not None:
+            output += [1] * len(bos_token_id + token_ids_1 + eos_token_id)
+
+        return output

native_backbone/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bcacff3229854f5103ee7a85473a30ca9a8b3a68f3aae9b7479574b23ac2256b
+size 2475075

native_backbone/tokenizer_config.json

@@ -0,0 +1,330 @@
+{
+  "add_bos_token": true,
+  "add_eos_token": false,
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128111": {
+      "content": "<restate>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128112": {
+      "content": "</restate>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128113": {
+      "content": "<planning>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128114": {
+      "content": "</planning>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128115": {
+      "content": "<recollect>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128116": {
+      "content": "</recollect>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128117": {
+      "content": "<execution>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128118": {
+      "content": "</execution>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128119": {
+      "content": "<review>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128120": {
+      "content": "</review>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128121": {
+      "content": "<summarize>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128122": {
+      "content": "</summarize>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128123": {
+      "content": "<retry>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128124": {
+      "content": "</retry>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128125": {
+      "content": "<conclude>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128126": {
+      "content": "</conclude>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128127": {
+      "content": "<|plugin|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128128": {
+      "content": "<|interpreter|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128129": {
+      "content": "<|action_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128130": {
+      "content": "<|action_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128131": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128132": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128133": {
+      "content": "<img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128134": {
+      "content": "</img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128135": {
+      "content": "<IMG_CONTEXT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128136": {
+      "content": "<quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128137": {
+      "content": "</quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128138": {
+      "content": "<ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128139": {
+      "content": "</ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128140": {
+      "content": "<box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "128141": {
+      "content": "</box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|action_start|>",
+    "<|action_end|>",
+    "<|interpreter|>",
+    "<|plugin|>",
+    "<restate>",
+    "</restate>",
+    "<planning>",
+    "</planning>",
+    "<recollect>",
+    "</recollect>",
+    "<execution>",
+    "</execution>",
+    "<review>",
+    "</review>",
+    "<summarize>",
+    "</summarize>",
+    "<retry>",
+    "</retry>",
+    "<conclude>",
+    "</conclude>",
+    "<img>",
+    "</img>",
+    "<IMG_CONTEXT>",
+    "<quad>",
+    "</quad>",
+    "<ref>",
+    "</ref>",
+    "<box>",
+    "</box>"
+  ],
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_internlm3.InternLM3Tokenizer",
+      null
+    ]
+  },
+  "bos_token": "<s>",
+  "chat_template": "{{ bos_token }}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "extra_special_tokens": {},
+  "model_max_length": 8192,
+  "pad_token": "</s>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "InternLM3Tokenizer",
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false
+}

requirements.txt

@@ -0,0 +1,11 @@
+torch==2.9.1
+transformers==4.57.6
+peft==0.17.1
+timm
+einops
+accelerate
+safetensors
+huggingface_hub
+Pillow
+numpy
+sentencepiece

source_gemma.py

@@ -0,0 +1,685 @@
+"""Strict persistent-visual CVRR wrapper for Gemma 3 and Gemma 4.
+
+The implementation intentionally wraps released Transformers models instead
+of copying their decoder source.  Native layer-call arguments are captured
+during frozen prefix passes and replayed for the shared recurrent cell and
+upper text-only continuation.  This preserves each release's masks, rotary
+geometry, and attention type while enforcing a visibly auditable interface:
+the upper decoder receives only non-visual recurrent rows.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import json
+import math
+import pathlib
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .source_helpers import _layer_hidden
+
+
+class LoRALinear(nn.Module):
+    def __init__(
+        self,
+        base: nn.Linear,
+        *,
+        rank: int,
+        alpha: float,
+        dropout: float,
+    ):
+        super().__init__()
+        if rank <= 0:
+            raise ValueError("LoRA rank must be positive")
+        self.base = base
+        for parameter in self.base.parameters():
+            parameter.requires_grad_(False)
+        self.rank = int(rank)
+        self.scale = float(alpha) / float(rank)
+        self.dropout = nn.Dropout(float(dropout))
+        self.lora_A = nn.Parameter(torch.empty(rank, base.in_features))
+        self.lora_B = nn.Parameter(torch.zeros(base.out_features, rank))
+        nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
+        self.enabled = True
+
+    def forward(self, inputs):
+        result = self.base(inputs)
+        if not self.enabled:
+            return result
+        update = F.linear(self.dropout(inputs).float(), self.lora_A.float())
+        update = F.linear(update, self.lora_B.float())
+        return result + (update * self.scale).to(result.dtype)
+
+
+def _module_parent(root: nn.Module, path: str):
+    parts = path.split(".")
+    parent = root
+    for part in parts[:-1]:
+        parent = getattr(parent, part)
+    return parent, parts[-1]
+
+
+def inject_cell_lora(
+    cell: nn.Module,
+    *,
+    rank: int,
+    alpha: float,
+    dropout: float,
+    suffixes: set[str] | None = None,
+) -> dict[str, LoRALinear]:
+    if suffixes is None:
+        suffixes = {
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        }
+    selected = {
+        name: module
+        for name, module in cell.named_modules()
+        if isinstance(module, nn.Linear) and name.rsplit(".", 1)[-1] in suffixes
+    }
+    if not selected:
+        raise RuntimeError("no attention/MLP projections found in recurrent cell")
+    wrappers = {}
+    # Materialize the list before replacing children during traversal.
+    for name, module in selected.items():
+        parent, child = _module_parent(cell, name)
+        wrapper = LoRALinear(
+            module, rank=rank, alpha=alpha, dropout=dropout
+        )
+        setattr(parent, child, wrapper)
+        wrappers[name] = wrapper
+    return wrappers
+
+
+@dataclass
+class LayerCall:
+    positional_tail: tuple[Any, ...]
+    keywords: dict[str, Any]
+
+
+@dataclass
+class GemmaCVRRTrace:
+    """Frozen native states needed for one strict recurrent rollout.
+
+    ``scaffold`` is multimodal and is consumed only by the shared recurrent
+    cell.  ``text_calls`` comes from an image-free sequence and is the sole
+    context replayed by the upper decoder.  Keeping those two objects separate
+    makes the no-bypass contract directly inspectable.
+    """
+
+    scaffold: torch.Tensor
+    text_rows: torch.Tensor
+    visual_rows: torch.Tensor
+    question_valid: torch.Tensor
+    base_anchor: torch.Tensor
+    r1: torch.Tensor
+    mm_cell_call: LayerCall
+    text_calls: dict[int, LayerCall]
+
+    @property
+    def question_lengths(self) -> torch.Tensor:
+        return self.question_valid.sum(dim=1)
+
+    @property
+    def visual_lengths(self) -> torch.Tensor:
+        return self.visual_rows.sum(dim=1)
+
+
+class _StopAfterCell(RuntimeError):
+    pass
+
+
+def _capture_call(storage: dict[int, LayerCall], index: int):
+    def hook(_module, args, kwargs):
+        storage[index] = LayerCall(tuple(args[1:]), dict(kwargs))
+
+    return hook
+
+
+def _gather_rows(full, mask):
+    batch, _, width = full.shape
+    lengths = mask.sum(dim=1).tolist()
+    result = full.new_zeros((batch, max(lengths), width))
+    for row, length in enumerate(lengths):
+        result[row, :length] = full[row, mask[row]]
+    return result
+
+
+def _gather_ids(full, mask, *, pad_value: int):
+    lengths = mask.sum(dim=1).tolist()
+    result = full.new_full((full.shape[0], max(lengths)), pad_value)
+    valid = torch.zeros_like(result, dtype=torch.bool)
+    for row, length in enumerate(lengths):
+        result[row, :length] = full[row, mask[row]]
+        valid[row, :length] = True
+    return result, valid
+
+
+def _replace_rows(full, mask, rows):
+    result = full.clone()
+    for batch_index in range(full.shape[0]):
+        count = int(mask[batch_index].sum())
+        result[batch_index, mask[batch_index]] = rows[batch_index, :count]
+    return result
+
+
+class GemmaCVRR(nn.Module):
+    """One localized Gemma backbone with one shared recurrent-cell LoRA."""
+
+    def __init__(
+        self,
+        model_path: str,
+        *,
+        ell_star: int,
+        steps: int = 4,
+        beta: float = 0.33,
+        rank: int = 32,
+        alpha: float = 12.0,
+        dropout: float = 0.01,
+        device: str | torch.device = "cuda:0",
+        offline: bool = True,
+    ):
+        super().__init__()
+        from transformers import AutoConfig
+
+        self.model_path = str(model_path)
+        self.device_ref = torch.device(device)
+        config = AutoConfig.from_pretrained(
+            model_path, local_files_only=offline
+        )
+        if config.model_type == "gemma3":
+            from transformers import Gemma3ForConditionalGeneration as ModelClass
+        elif config.model_type == "gemma4_unified":
+            from transformers import (
+                Gemma4UnifiedForConditionalGeneration as ModelClass,
+            )
+            if int(getattr(config.text_config, "num_kv_shared_layers", 0)):
+                raise NotImplementedError(
+                    "Gemma4 cross-layer shared KV would be an upper visual bypass"
+                )
+        else:
+            raise ValueError(f"unsupported Gemma model_type={config.model_type!r}")
+        self.base_model = ModelClass.from_pretrained(
+            model_path,
+            dtype=torch.bfloat16,
+            device_map=str(self.device_ref),
+            local_files_only=offline,
+            attn_implementation="sdpa",
+        )
+        self.model_type = config.model_type
+        self.layers = self.base_model.model.language_model.layers
+        self.ell_star = int(ell_star)
+        self.cell_index = self.ell_star + 1
+        self.upper_start = self.cell_index + 1
+        if not 0 <= self.ell_star <= len(self.layers) - 2:
+            raise ValueError(
+                f"ell_star={ell_star} must leave a recurrent cell and upper decoder"
+            )
+        if steps < 2:
+            raise ValueError("CVRR training requires at least two recurrent states")
+        if not 0.0 <= beta <= 1.0:
+            raise ValueError("beta must lie in [0,1]")
+        self.steps = int(steps)
+        self.beta = float(beta)
+        for parameter in self.base_model.parameters():
+            parameter.requires_grad_(False)
+        self.lora = inject_cell_lora(
+            self.layers[self.cell_index],
+            rank=rank,
+            alpha=alpha,
+            dropout=dropout,
+        )
+        # The dense checkpoint was placed before adapters were constructed;
+        # newly allocated A/B tensors otherwise remain on CPU until an outer
+        # training entrypoint happens to call ``model.to(device)``.
+        for module in self.lora.values():
+            module.to(self.device_ref)
+        self.rank = int(rank)
+        self.alpha = float(alpha)
+        self.adapter_dropout = float(dropout)
+        self.base_model.eval()
+
+    @contextlib.contextmanager
+    def adapters(self, enabled: bool):
+        previous = [module.enabled for module in self.lora.values()]
+        for module in self.lora.values():
+            module.enabled = bool(enabled)
+        try:
+            yield
+        finally:
+            for module, value in zip(self.lora.values(), previous):
+                module.enabled = value
+
+    def train(self, mode: bool = True):
+        super().train(mode)
+        # Frozen dense modules stay deterministic; only LoRA dropout follows
+        # training mode.
+        self.base_model.eval()
+        for module in self.lora.values():
+            module.dropout.train(mode)
+        return self
+
+    def _modality(self, mm_inputs):
+        if "token_type_ids" in mm_inputs:
+            return mm_inputs["token_type_ids"]
+        if "mm_token_type_ids" in mm_inputs:
+            return mm_inputs["mm_token_type_ids"]
+        raise ValueError("Gemma multimodal inputs have no modality IDs")
+
+    def _pad_token_id(self) -> int:
+        return int(self.base_model.config.text_config.pad_token_id)
+
+    def _initial_multimodal(self, mm_inputs):
+        calls: dict[int, LayerCall] = {}
+        captured = {}
+        cell = self.layers[self.cell_index]
+
+        def stop(_module, _args, output):
+            captured["hidden"] = _layer_hidden(output).detach()
+            raise _StopAfterCell
+
+        pre = cell.register_forward_pre_hook(
+            _capture_call(calls, self.cell_index), with_kwargs=True
+        )
+        post = cell.register_forward_hook(stop)
+        try:
+            with torch.no_grad(), self.adapters(False):
+                try:
+                    self.base_model.model(
+                        **mm_inputs, use_cache=False, return_dict=True
+                    )
+                except _StopAfterCell:
+                    pass
+        finally:
+            pre.remove()
+            post.remove()
+        if "hidden" not in captured or self.cell_index not in calls:
+            raise RuntimeError("failed to capture native multimodal recurrent cell")
+        return captured["hidden"], calls[self.cell_index]
+
+    def _text_context(self, question_ids, question_mask):
+        calls: dict[int, LayerCall] = {}
+        captured = {}
+        handles = []
+        for index in range(self.cell_index, len(self.layers)):
+            handles.append(
+                self.layers[index].register_forward_pre_hook(
+                    _capture_call(calls, index), with_kwargs=True
+                )
+            )
+
+        def capture_cell(_module, _args, output):
+            captured["anchor"] = _layer_hidden(output).detach()
+
+        handles.append(self.layers[self.cell_index].register_forward_hook(capture_cell))
+        try:
+            with torch.no_grad(), self.adapters(False):
+                self.base_model.model(
+                    input_ids=question_ids,
+                    attention_mask=question_mask,
+                    use_cache=False,
+                    return_dict=True,
+                )
+        finally:
+            for handle in handles:
+                handle.remove()
+        missing = [
+            index
+            for index in range(self.cell_index, len(self.layers))
+            if index not in calls
+        ]
+        if missing or "anchor" not in captured:
+            raise RuntimeError(f"failed to capture text context; missing={missing}")
+        return captured["anchor"], calls
+
+    @staticmethod
+    def _call_layer(layer, hidden, call: LayerCall):
+        return _layer_hidden(
+            layer(hidden, *call.positional_tail, **call.keywords)
+        )
+
+    def _upper(self, state, text_calls):
+        hidden = state
+        with self.adapters(False):
+            for index in range(self.upper_start, len(self.layers)):
+                hidden = self._call_layer(
+                    self.layers[index], hidden, text_calls[index]
+                )
+        hidden = self.base_model.model.language_model.norm(hidden)
+        logits = self.base_model.lm_head(hidden)
+        if self.model_type == "gemma4_unified":
+            cap = self.base_model.config.text_config.final_logit_softcapping
+            if cap is not None:
+                logits = torch.tanh(logits / cap) * cap
+        return logits
+
+    def extract(self, mm_inputs: dict[str, torch.Tensor]) -> GemmaCVRRTrace:
+        """Extract the frozen native first read and text-only upper context."""
+
+        attention = mm_inputs["attention_mask"].bool()
+        visual = self._modality(mm_inputs).eq(1) & attention
+        text_rows = (~visual) & attention
+        question_ids, question_valid = _gather_ids(
+            mm_inputs["input_ids"],
+            text_rows,
+            pad_value=self._pad_token_id(),
+        )
+        question_mask = question_valid.long()
+        first_full, mm_cell_call = self._initial_multimodal(mm_inputs)
+        base_anchor, text_calls = self._text_context(question_ids, question_mask)
+        r1 = _gather_rows(first_full, text_rows)
+        if r1.shape != base_anchor.shape:
+            raise RuntimeError(
+                "native multimodal and image-free question states are misaligned: "
+                f"R1={tuple(r1.shape)}, B={tuple(base_anchor.shape)}"
+            )
+        return GemmaCVRRTrace(
+            scaffold=first_full.detach(),
+            text_rows=text_rows,
+            visual_rows=visual,
+            question_valid=question_valid,
+            base_anchor=base_anchor.detach(),
+            r1=r1.detach(),
+            mm_cell_call=mm_cell_call,
+            text_calls=text_calls,
+        )
+
+    def rollout(
+        self,
+        trace: GemmaCVRRTrace,
+        *,
+        steps: int | None = None,
+        initial_state: torch.Tensor | None = None,
+    ) -> list[torch.Tensor]:
+        """Run the shared native cell and return ``[R1, ..., R_T]``."""
+
+        horizon = self.steps if steps is None else int(steps)
+        if horizon < 1:
+            raise ValueError("rollout steps must be positive")
+        state = trace.r1 if initial_state is None else initial_state
+        if state.shape != trace.r1.shape:
+            raise ValueError(
+                f"initial state shape {tuple(state.shape)} != {tuple(trace.r1.shape)}"
+            )
+        states = [state]
+        for _ in range(1, horizon):
+            recurrent_input = _replace_rows(
+                trace.scaffold, trace.text_rows, state
+            )
+            with self.adapters(True):
+                proposal_full = self._call_layer(
+                    self.layers[self.cell_index],
+                    recurrent_input,
+                    trace.mm_cell_call,
+                )
+            proposal = _gather_rows(proposal_full, trace.text_rows)
+            state = state + self.beta * (proposal - state)
+            states.append(state)
+        return states
+
+    def decode_logits(
+        self,
+        trace: GemmaCVRRTrace,
+        state: torch.Tensor,
+    ) -> torch.Tensor:
+        """Decode one question-shaped state through the strict text-only path."""
+
+        if state.shape != trace.base_anchor.shape:
+            raise ValueError(
+                f"decoder state shape {tuple(state.shape)} != "
+                f"text anchor {tuple(trace.base_anchor.shape)}"
+            )
+        # Written explicitly as B + C_T to mirror the method definition.  No
+        # multimodal row or multimodal cache is passed to `_upper`.
+        decoder_state = trace.base_anchor + (state - trace.base_anchor)
+        return self._upper(decoder_state, trace.text_calls).float()
+
+    def next_token_logits(
+        self,
+        trace: GemmaCVRRTrace,
+        state: torch.Tensor,
+    ) -> torch.Tensor:
+        """Return the distribution after each sample's final valid prompt row."""
+
+        logits = self.decode_logits(trace, state)
+        row = trace.question_lengths.to(logits.device) - 1
+        if bool((row < 0).any()):
+            raise ValueError("empty question sequence")
+        batch = torch.arange(logits.shape[0], device=logits.device)
+        return logits[batch, row]
+
+    def residual(self, trace: GemmaCVRRTrace, state: torch.Tensor) -> torch.Tensor:
+        return state - trace.base_anchor
+
+    def state_from_residual(
+        self,
+        trace: GemmaCVRRTrace,
+        residual: torch.Tensor,
+    ) -> torch.Tensor:
+        if residual.shape != trace.base_anchor.shape:
+            raise ValueError(
+                f"residual shape {tuple(residual.shape)} != "
+                f"text anchor {tuple(trace.base_anchor.shape)}"
+            )
+        return trace.base_anchor + residual
+
+    def forward(self, mm_inputs: dict[str, torch.Tensor], mm_labels):
+        question_labels, _ = _gather_ids(
+            mm_labels,
+            ((~self._modality(mm_inputs).eq(1)) & mm_inputs["attention_mask"].bool()),
+            pad_value=-100,
+        )
+        trace = self.extract(mm_inputs)
+        state = self.rollout(trace)[-1]
+        logits = self.decode_logits(trace, state)
+        shift_logits = logits[:, :-1]
+        shift_labels = question_labels[:, 1:]
+        token_loss = F.cross_entropy(
+            shift_logits.reshape(-1, shift_logits.shape[-1]),
+            shift_labels.reshape(-1),
+            ignore_index=-100,
+            reduction="none",
+        ).reshape(shift_labels.shape)
+        valid = shift_labels.ne(-100)
+        counts = valid.sum(dim=1).clamp_min(1)
+        per_example = (token_loss * valid).sum(dim=1) / counts
+        return {
+            "loss": per_example.mean(),
+            "logits": logits,
+            "labels": question_labels,
+            "r1": trace.r1.detach(),
+            "rT": state.detach(),
+            "base_anchor": trace.base_anchor.detach(),
+            "visual_rows": (
+                self._modality(mm_inputs).eq(1)
+                & mm_inputs["attention_mask"].bool()
+            ).sum(dim=1).detach(),
+        }
+
+    def adapter_state_dict(self):
+        return {
+            name: tensor.detach().cpu()
+            for name, tensor in self.state_dict().items()
+            if ".lora_A" in name or ".lora_B" in name
+        }
+
+    def save_adapter(self, output_dir: str | pathlib.Path, *, step: int):
+        output = pathlib.Path(output_dir)
+        output.mkdir(parents=True, exist_ok=True)
+        torch.save(self.adapter_state_dict(), output / "adapter_model.pt")
+        metadata = {
+            "format": "gemma_cvrr_lora_v1",
+            "base_model": self.model_path,
+            "model_type": self.model_type,
+            "ell_star": self.ell_star,
+            "cell_layer": self.cell_index,
+            "num_workspace_steps": self.steps,
+            "counterfactual_beta": self.beta,
+            "adapter_rank": self.rank,
+            "adapter_alpha": self.alpha,
+            "adapter_dropout": self.adapter_dropout,
+            "step": int(step),
+            "strict_path": True,
+        }
+        (output / "cvrr_config.json").write_text(json.dumps(metadata, indent=2))
+
+    def load_adapter(self, adapter_dir: str | pathlib.Path):
+        """Load an adapter exactly; missing or surplus LoRA tensors are fatal."""
+
+        adapter_dir = pathlib.Path(adapter_dir).expanduser().resolve()
+        metadata = json.loads((adapter_dir / "cvrr_config.json").read_text())
+        checks = {
+            "model_type": self.model_type,
+            "ell_star": self.ell_star,
+            "cell_layer": self.cell_index,
+            "num_workspace_steps": self.steps,
+            "adapter_rank": self.rank,
+        }
+        mismatches = {
+            name: (metadata.get(name), expected)
+            for name, expected in checks.items()
+            if metadata.get(name) != expected
+        }
+        if mismatches:
+            raise ValueError(f"adapter metadata mismatch: {mismatches}")
+        expected_keys = set(self.adapter_state_dict())
+        try:
+            payload = torch.load(
+                adapter_dir / "adapter_model.pt",
+                map_location="cpu",
+                weights_only=True,
+            )
+        except TypeError:
+            payload = torch.load(adapter_dir / "adapter_model.pt", map_location="cpu")
+        actual_keys = set(payload)
+        if actual_keys != expected_keys:
+            raise RuntimeError(
+                "adapter tensor mismatch: "
+                f"missing={sorted(expected_keys - actual_keys)[:8]}, "
+                f"unexpected={sorted(actual_keys - expected_keys)[:8]}"
+            )
+        incompatible = self.load_state_dict(payload, strict=False)
+        unexpected = list(incompatible.unexpected_keys)
+        missing_lora = [key for key in incompatible.missing_keys if key in expected_keys]
+        if unexpected or missing_lora:
+            raise RuntimeError(
+                f"adapter load failed: missing={missing_lora}, unexpected={unexpected}"
+            )
+        return metadata
+
+
+def _prompt(processor, question: str, hint: str):
+    content = [
+        {"type": "image"},
+        {"type": "text", "text": question.strip() + str(hint)},
+    ]
+    return processor.apply_chat_template(
+        [{"role": "user", "content": content}],
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+
+class GemmaArrowCollator:
+    def __init__(self, processor):
+        self.processor = processor
+        self.tokenizer = processor.tokenizer
+
+    @staticmethod
+    def _pad_1d(values, pad):
+        width = max(item.shape[0] for item in values)
+        result = values[0].new_full((len(values), width), pad)
+        for index, item in enumerate(values):
+            result[index, : item.shape[0]] = item
+        return result
+
+    def __call__(self, features):
+        from PIL import Image
+
+        examples = []
+        for feature in features:
+            raw = feature["image_bytes"]
+            if isinstance(raw, memoryview):
+                raw = raw.tobytes()
+            with Image.open(io.BytesIO(raw)) as opened:
+                image = opened.convert("RGB")
+            prompt = _prompt(
+                self.processor,
+                str(feature["fixed_question"]),
+                str(feature["fixed_hint"]),
+            )
+            prompt_item = self.processor(
+                text=prompt, images=[image], return_tensors="pt"
+            )
+            full_item = self.processor(
+                text=prompt + str(feature["fixed_answer"]).strip(),
+                images=[image],
+                return_tensors="pt",
+            )
+            prompt_ids = prompt_item["input_ids"][0]
+            full_ids = full_item["input_ids"][0]
+            if not torch.equal(full_ids[: prompt_ids.numel()], prompt_ids):
+                raise RuntimeError("Gemma answer serialization changed the prompt prefix")
+            eos = full_ids.new_tensor([self.tokenizer.eos_token_id])
+            item = {name: value for name, value in full_item.items()}
+            item["input_ids"] = torch.cat((full_ids, eos))
+            item["attention_mask"] = torch.cat(
+                (item["attention_mask"][0], torch.ones_like(eos))
+            )
+            modality_name = (
+                "token_type_ids"
+                if "token_type_ids" in item
+                else "mm_token_type_ids"
+            )
+            item[modality_name] = torch.cat(
+                (item[modality_name][0], torch.zeros_like(eos))
+            )
+            answer_ids = item["input_ids"][prompt_ids.numel() :]
+            item["labels"] = torch.cat(
+                (torch.full_like(prompt_ids, -100), answer_ids)
+            )
+            examples.append(item)
+
+        sequence_names = {
+            "input_ids": int(self.tokenizer.pad_token_id),
+            "attention_mask": 0,
+            "labels": -100,
+        }
+        modality_name = (
+            "token_type_ids"
+            if "token_type_ids" in examples[0]
+            else "mm_token_type_ids"
+        )
+        sequence_names[modality_name] = 0
+        batch = {
+            name: self._pad_1d([item[name] for item in examples], pad)
+            for name, pad in sequence_names.items()
+        }
+        for name in examples[0]:
+            if name in sequence_names or name == "labels":
+                continue
+            values = [item[name] for item in examples]
+            batch[name] = torch.cat(values, dim=0)
+        labels = batch.pop("labels")
+        return {"mm_inputs": batch, "mm_labels": labels}
+
+
+def move_batch(batch, device):
+    return {
+        "mm_inputs": {
+            name: value.to(device, non_blocking=True)
+            for name, value in batch["mm_inputs"].items()
+        },
+        "mm_labels": batch["mm_labels"].to(device, non_blocking=True),
+    }

source_helpers.py

@@ -0,0 +1,69 @@
+def _layer_hidden(output):
+    """Return the hidden tensor from a decoder-layer output."""
+
+    import torch
+
+    if torch.is_tensor(output):
+        return output
+    if isinstance(output, (tuple, list)) and output and torch.is_tensor(output[0]):
+        return output[0]
+    raise TypeError(f"unsupported decoder-layer output type: {type(output)!r}")
+
+def _closest_ratio(aspect, ratios, width, height, image_size):
+    best = (1, 1)
+    difference = float("inf")
+    area = width * height
+    for ratio in ratios:
+        candidate = ratio[0] / ratio[1]
+        current = abs(aspect - candidate)
+        if current < difference or (
+            current == difference
+            and area > 0.5 * image_size * image_size * ratio[0] * ratio[1]
+        ):
+            difference = current
+            best = ratio
+    return best
+
+def dynamic_tiles(image, *, image_size: int, max_tiles: int, thumbnail: bool):
+    """Official InternVL dynamic tiling, kept local for reproducibility."""
+
+    ratios = sorted(
+        {
+            (i, j)
+            for n in range(1, max_tiles + 1)
+            for i in range(1, n + 1)
+            for j in range(1, n + 1)
+            if 1 <= i * j <= max_tiles
+        },
+        key=lambda item: item[0] * item[1],
+    )
+    width, height = image.size
+    columns, rows = _closest_ratio(
+        width / height, ratios, width, height, image_size
+    )
+    resized = image.convert("RGB").resize(
+        (image_size * columns, image_size * rows), resample=3
+    )
+    tiles = []
+    for index in range(columns * rows):
+        left = (index % columns) * image_size
+        top = (index // columns) * image_size
+        tiles.append(
+            resized.crop((left, top, left + image_size, top + image_size))
+        )
+    if thumbnail and len(tiles) != 1:
+        tiles.append(image.convert("RGB").resize((image_size, image_size), 3))
+    return tiles
+
+def _normalize_tiles(tiles):
+    import numpy as np
+    import torch
+
+    mean = torch.tensor((0.485, 0.456, 0.406)).view(3, 1, 1)
+    std = torch.tensor((0.229, 0.224, 0.225)).view(3, 1, 1)
+    tensors = []
+    for tile in tiles:
+        array = np.asarray(tile, dtype=np.float32) / 255.0
+        tensor = torch.from_numpy(array).permute(2, 0, 1)
+        tensors.append((tensor - mean) / std)
+    return torch.stack(tensors)

source_internvl.py

@@ -0,0 +1,286 @@
+"""Strict persistent-visual CVRR wrapper for released InternVL3 chat models."""
+
+from __future__ import annotations
+
+import copy
+import io
+import pathlib
+
+import torch
+
+from .source_gemma import (
+    GemmaCVRR,
+    LayerCall,
+    _StopAfterCell,
+    _capture_call,
+    _gather_ids,
+    _gather_rows,
+    _replace_rows,
+    inject_cell_lora,
+)
+from .source_helpers import _layer_hidden
+from .source_helpers import _normalize_tiles, dynamic_tiles
+
+
+class InternVLCVRR(GemmaCVRR):
+    def __init__(
+        self,
+        model_path: str,
+        *,
+        ell_star: int,
+        steps: int = 4,
+        beta: float = 0.33,
+        rank: int = 32,
+        alpha: float = 12.0,
+        dropout: float = 0.01,
+        device: str | torch.device = "cuda:0",
+        offline: bool = True,
+    ):
+        # Bypass GemmaCVRR.__init__, retaining its audited recurrence, adapter
+        # toggling, loss, and serialization methods.
+        torch.nn.Module.__init__(self)
+        from transformers import AutoModel, AutoTokenizer
+
+        import torch.distributed as dist
+
+        self._owns_process_group = False
+        if dist.is_available() and not dist.is_initialized():
+            import os
+            import tempfile
+
+            rendezvous = pathlib.Path(tempfile.gettempdir()) / f"cvrr_iv_train_{os.getpid()}"
+            dist.init_process_group(
+                "gloo", init_method=f"file://{rendezvous}", rank=0, world_size=1
+            )
+            self._owns_process_group = True
+
+        self.model_path = str(model_path)
+        self.device_ref = torch.device(device)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+            use_fast=False,
+            local_files_only=offline,
+        )
+        self.base_model = AutoModel.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+            local_files_only=offline,
+            low_cpu_mem_usage=True,
+            use_flash_attn=False,
+            dtype=torch.bfloat16,
+            device_map=str(self.device_ref),
+        )
+        self.model_type = "internvl_chat"
+        self.layers = self.base_model.language_model.model.layers
+        self.ell_star = int(ell_star)
+        self.cell_index = self.ell_star + 1
+        self.upper_start = self.cell_index + 1
+        if not 0 <= self.ell_star <= len(self.layers) - 2:
+            raise ValueError("ell_star must leave a cell and upper decoder")
+        if steps < 2 or not 0.0 <= beta <= 1.0:
+            raise ValueError("invalid recurrence depth or beta")
+        self.steps = int(steps)
+        self.beta = float(beta)
+        for parameter in self.base_model.parameters():
+            parameter.requires_grad_(False)
+        self.lora = inject_cell_lora(
+            self.layers[self.cell_index],
+            rank=rank,
+            alpha=alpha,
+            dropout=dropout,
+            suffixes={"wqkv", "wo", "w1", "w2", "w3"},
+        )
+        for module in self.lora.values():
+            module.to(self.device_ref)
+        self.rank = int(rank)
+        self.alpha = float(alpha)
+        self.adapter_dropout = float(dropout)
+        self.image_token_id = int(
+            self.tokenizer.convert_tokens_to_ids("<IMG_CONTEXT>")
+        )
+        self.base_model.img_context_token_id = self.image_token_id
+        self.base_model.eval()
+
+    def _pad_token_id(self) -> int:
+        return int(self.base_model.config.llm_config.pad_token_id)
+
+    def _modality(self, mm_inputs):
+        return mm_inputs["input_ids"].eq(self.image_token_id).long()
+
+    def _initial_multimodal(self, mm_inputs):
+        calls: dict[int, LayerCall] = {}
+        captured = {}
+        cell = self.layers[self.cell_index]
+
+        def stop(_module, _args, output):
+            captured["hidden"] = _layer_hidden(output).detach()
+            raise _StopAfterCell
+
+        pre = cell.register_forward_pre_hook(
+            _capture_call(calls, self.cell_index), with_kwargs=True
+        )
+        post = cell.register_forward_hook(stop)
+        try:
+            with torch.no_grad(), self.adapters(False):
+                try:
+                    self.base_model(
+                        **mm_inputs,
+                        use_cache=False,
+                        output_hidden_states=False,
+                        return_dict=True,
+                    )
+                except _StopAfterCell:
+                    pass
+        finally:
+            pre.remove()
+            post.remove()
+        if "hidden" not in captured or self.cell_index not in calls:
+            raise RuntimeError("failed to capture InternVL recurrent cell")
+        return captured["hidden"], calls[self.cell_index]
+
+    def _text_context(self, question_ids, question_mask):
+        calls: dict[int, LayerCall] = {}
+        captured = {}
+        handles = []
+        for index in range(self.cell_index, len(self.layers)):
+            handles.append(
+                self.layers[index].register_forward_pre_hook(
+                    _capture_call(calls, index), with_kwargs=True
+                )
+            )
+
+        def capture(_module, _args, output):
+            captured["anchor"] = _layer_hidden(output).detach()
+
+        handles.append(self.layers[self.cell_index].register_forward_hook(capture))
+        try:
+            with torch.no_grad(), self.adapters(False):
+                self.base_model.language_model(
+                    input_ids=question_ids,
+                    attention_mask=question_mask,
+                    use_cache=False,
+                    output_hidden_states=False,
+                    return_dict=True,
+                )
+        finally:
+            for handle in handles:
+                handle.remove()
+        missing = [
+            index
+            for index in range(self.cell_index, len(self.layers))
+            if index not in calls
+        ]
+        if missing or "anchor" not in captured:
+            raise RuntimeError(f"failed to capture InternVL text path: {missing}")
+        return captured["anchor"], calls
+
+    def _upper(self, state, text_calls):
+        hidden = state
+        with self.adapters(False):
+            for index in range(self.upper_start, len(self.layers)):
+                hidden = self._call_layer(
+                    self.layers[index], hidden, text_calls[index]
+                )
+        hidden = self.base_model.language_model.model.norm(hidden)
+        return self.base_model.language_model.output(hidden).float()
+
+
+class InternVLArrowCollator:
+    def __init__(self, model: InternVLCVRR, *, max_tiles: int = 12):
+        self.tokenizer = model.tokenizer
+        self.template = copy.deepcopy(model.base_model.conv_template)
+        self.system_message = model.base_model.system_message
+        self.num_image_token = int(model.base_model.num_image_token)
+        self.image_size = int(
+            model.base_model.config.force_image_size
+            or model.base_model.config.vision_config.image_size
+        )
+        self.use_thumbnail = bool(model.base_model.config.use_thumbnail)
+        self.max_tiles = int(max_tiles)
+
+    @staticmethod
+    def _pad(values, pad):
+        width = max(value.shape[0] for value in values)
+        output = values[0].new_full((len(values), width), pad)
+        for index, value in enumerate(values):
+            output[index, : value.shape[0]] = value
+        return output
+
+    def _query(self, question, hint, num_tiles):
+        template = copy.deepcopy(self.template)
+        template.system_message = self.system_message
+        template.append_message(
+            template.roles[0],
+            "<image>\n" + str(question).strip() + str(hint),
+        )
+        template.append_message(template.roles[1], None)
+        query = template.get_prompt()
+        visual = (
+            "<img>"
+            + "<IMG_CONTEXT>" * self.num_image_token * num_tiles
+            + "</img>"
+        )
+        return query.replace("<image>", visual, 1)
+
+    def __call__(self, features):
+        from PIL import Image
+
+        rows = []
+        for feature in features:
+            raw = feature["image_bytes"]
+            if isinstance(raw, memoryview):
+                raw = raw.tobytes()
+            with Image.open(io.BytesIO(raw)) as opened:
+                tiles = dynamic_tiles(
+                    opened.convert("RGB"),
+                    image_size=self.image_size,
+                    max_tiles=self.max_tiles,
+                    thumbnail=self.use_thumbnail,
+                )
+            # The released InternViT does not cast inputs internally; its model
+            # card explicitly converts pixel_values to bfloat16 before forward.
+            pixels = _normalize_tiles(tiles).to(torch.bfloat16)
+            query = self._query(
+                feature["fixed_question"], feature["fixed_hint"], len(tiles)
+            )
+            tokenized = self.tokenizer(query, return_tensors="pt")
+            prompt = tokenized.input_ids[0]
+            full = self.tokenizer(
+                query + str(feature["fixed_answer"]).strip(),
+                return_tensors="pt",
+            ).input_ids[0]
+            if not torch.equal(full[: prompt.numel()], prompt):
+                raise RuntimeError("InternVL answer serialization changed the prompt prefix")
+            full = torch.cat(
+                (full, full.new_tensor([self.tokenizer.eos_token_id]))
+            )
+            answer = full[prompt.numel() :]
+            rows.append(
+                {
+                    "input_ids": full,
+                    "attention_mask": torch.cat(
+                        (tokenized.attention_mask[0], torch.ones_like(answer))
+                    ),
+                    "labels": torch.cat((torch.full_like(prompt, -100), answer)),
+                    "pixel_values": pixels,
+                    "image_flags": torch.ones(len(tiles), 1, dtype=torch.long),
+                }
+            )
+        return {
+            "mm_inputs": {
+                "input_ids": self._pad(
+                    [row["input_ids"] for row in rows], self.tokenizer.pad_token_id
+                ),
+                "attention_mask": self._pad(
+                    [row["attention_mask"] for row in rows], 0
+                ),
+                "pixel_values": torch.cat(
+                    [row["pixel_values"] for row in rows], dim=0
+                ),
+                "image_flags": torch.cat(
+                    [row["image_flags"] for row in rows], dim=0
+                ),
+            },
+            "mm_labels": self._pad([row["labels"] for row in rows], -100),
+        }

source_perceive.py

@@ -0,0 +1,414 @@
+"""Attention layouts for non-recurrent visual latent reasoning chains.
+
+The sequence is laid out as::
+
+    [multimodal prompt ; clean question ; LOOK_1 ; THINK_1 ; ... ; answer]
+
+All rows are processed once by the native VLM decoder.  A block-sparse causal
+graph makes LOOK rows the only latent rows with access to the multimodal
+prefix, while THINK rows integrate one visual read without directly seeing
+that prefix.  Answer rows can see the clean question and THINK rows, but never
+the multimodal prefix or LOOK rows.  Consequently the image-dependent answer
+path is
+
+    image -> LOOK_k -> THINK_k -> answer,
+
+without a tied recurrent cell, a visual-state update, or an aggregation
+module.
+
+The transition-conditioned variant tightens the graph without adding a
+module.  The first LOOK/THINK pair bootstraps from the complete prompt and the
+clean question.  Every later pair sees the causal latent prefix, while its
+LOOK row receives only visual placeholder rows from the multimodal prefix.
+Thus later pairs cannot independently re-solve the original image/question
+prompt; their task semantics must arrive through earlier latent states.  The
+last THINK is the sole latent exposed to answer rows and therefore acts as the
+native-transformer aggregation state.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass(frozen=True)
+class PerceiveDeliberateLayout:
+    """Physical row layout of one padded training/generation sequence."""
+
+    multimodal_length: int
+    question_length: int
+    num_pairs: int
+    answer_input_length: int = 0
+
+    def __post_init__(self) -> None:
+        if self.multimodal_length < 1:
+            raise ValueError("multimodal_length must be positive")
+        if self.question_length < 1:
+            raise ValueError("question_length must be positive")
+        if self.num_pairs < 1:
+            raise ValueError("num_pairs must be positive")
+        if self.answer_input_length < 0:
+            raise ValueError("answer_input_length must be non-negative")
+
+    @property
+    def num_latent_tokens(self) -> int:
+        return 2 * self.num_pairs
+
+    @property
+    def question_start(self) -> int:
+        return self.multimodal_length
+
+    @property
+    def latent_start(self) -> int:
+        return self.multimodal_length + self.question_length
+
+    @property
+    def answer_start(self) -> int:
+        return self.latent_start + self.num_latent_tokens
+
+    @property
+    def sequence_length(self) -> int:
+        return self.answer_start + self.answer_input_length
+
+    @property
+    def question_slice(self) -> slice:
+        return slice(self.question_start, self.latent_start)
+
+    @property
+    def answer_slice(self) -> slice:
+        return slice(self.answer_start, self.sequence_length)
+
+    @property
+    def look_indices(self) -> tuple[int, ...]:
+        return tuple(self.latent_start + 2 * index for index in range(self.num_pairs))
+
+    @property
+    def think_indices(self) -> tuple[int, ...]:
+        return tuple(index + 1 for index in self.look_indices)
+
+
+def build_perceive_deliberate_mask(
+    layout: PerceiveDeliberateLayout,
+    multimodal_attention_mask: torch.Tensor,
+    question_attention_mask: torch.Tensor,
+    answer_input_attention_mask: torch.Tensor | None = None,
+    *,
+    multimodal_visual_mask: torch.Tensor | None = None,
+    local_visual_chain: bool = False,
+    transition_conditioned: bool = False,
+    question_visible_through_pair: int = 1,
+    disable_chain_links: bool = False,
+    disable_late_visual: bool = False,
+    final_pair_only: bool = False,
+) -> torch.BoolTensor:
+    """Build the strict LOOK/THINK/answer visibility graph.
+
+    The returned boolean mask has shape ``[B, 1, L, L]`` and follows PyTorch
+    SDPA semantics: ``True`` means that the query may attend to the key.
+
+    In the original graph, LOOK_k sees the valid multimodal prompt, the clean
+    question, and (for k>1) only THINK_{k-1}; THINK_k sees the clean question
+    and LOOK_k.  In the transition-conditioned graph, pair 1 keeps those
+    bootstrap inputs, but every later latent sees its complete causal latent
+    prefix, later LOOK rows see only visual placeholder keys, and later THINK
+    rows normally receive no independent question shortcut.  For a controlled
+    curriculum/intervention, ``question_visible_through_pair`` may temporarily
+    retain the clean-question edge through a later pair; ``1`` is the strict
+    inference graph.  Answer rows see only the final THINK in that variant.
+    Every row also sees itself so padded query
+    rows remain numerically defined; padded rows are never exposed as keys to
+    semantic rows.
+    """
+
+    if multimodal_attention_mask.ndim != 2:
+        raise ValueError("multimodal_attention_mask must have shape [B, L_mm]")
+    if question_attention_mask.ndim != 2:
+        raise ValueError("question_attention_mask must have shape [B, L_q]")
+    batch_size = multimodal_attention_mask.shape[0]
+    expected_mm = (batch_size, layout.multimodal_length)
+    expected_q = (batch_size, layout.question_length)
+    if tuple(multimodal_attention_mask.shape) != expected_mm:
+        raise ValueError(
+            "multimodal mask/layout mismatch: "
+            f"expected {expected_mm}, got {tuple(multimodal_attention_mask.shape)}"
+        )
+    if tuple(question_attention_mask.shape) != expected_q:
+        raise ValueError(
+            "question mask/layout mismatch: "
+            f"expected {expected_q}, got {tuple(question_attention_mask.shape)}"
+        )
+    if layout.answer_input_length:
+        if answer_input_attention_mask is None:
+            raise ValueError("answer_input_attention_mask is required")
+        expected_answer = (batch_size, layout.answer_input_length)
+        if tuple(answer_input_attention_mask.shape) != expected_answer:
+            raise ValueError(
+                "answer mask/layout mismatch: "
+                f"expected {expected_answer}, got "
+                f"{tuple(answer_input_attention_mask.shape)}"
+            )
+    elif answer_input_attention_mask is not None and answer_input_attention_mask.numel():
+        raise ValueError("received answer mask for an empty answer-input segment")
+
+    device = multimodal_attention_mask.device
+    length = layout.sequence_length
+    visible = torch.zeros(
+        batch_size, length, length, dtype=torch.bool, device=device
+    )
+    mm_valid = multimodal_attention_mask.bool()
+    q_valid = question_attention_mask.bool()
+    visual_valid = None
+    if transition_conditioned or local_visual_chain:
+        if not 1 <= question_visible_through_pair <= layout.num_pairs:
+            raise ValueError(
+                "question_visible_through_pair must lie in "
+                f"[1,{layout.num_pairs}]"
+            )
+        if multimodal_visual_mask is None:
+            raise ValueError(
+                "multimodal_visual_mask is required for the "
+                "transition-conditioned graph"
+            )
+        if tuple(multimodal_visual_mask.shape) != expected_mm:
+            raise ValueError(
+                "visual mask/layout mismatch: "
+                f"expected {expected_mm}, got "
+                f"{tuple(multimodal_visual_mask.shape)}"
+            )
+        visual_valid = multimodal_visual_mask.bool()
+        if bool((visual_valid & ~mm_valid).any()):
+            raise ValueError("visual keys must be a subset of valid multimodal keys")
+
+    # The source multimodal prompt keeps the native causal graph.  Padded query
+    # rows are harmless and receive a self edge below.
+    mm_causal = torch.ones(
+        layout.multimodal_length,
+        layout.multimodal_length,
+        dtype=torch.bool,
+        device=device,
+    ).tril()
+    visible[:, : layout.multimodal_length, : layout.multimodal_length] = (
+        mm_causal.unsqueeze(0) & mm_valid[:, None, :]
+    )
+
+    # The duplicated question is deliberately text-only: it is causally
+    # connected only to preceding valid rows of its own segment.
+    q_causal = torch.ones(
+        layout.question_length,
+        layout.question_length,
+        dtype=torch.bool,
+        device=device,
+    ).tril()
+    visible[
+        :, layout.question_slice, layout.question_slice
+    ] = q_causal.unsqueeze(0) & q_valid[:, None, :]
+
+    if local_visual_chain:
+        # Homogeneous one-pass latent chain.  Visual memory is persistent and
+        # available at every position, while task state has exactly one local
+        # predecessor edge.  No latent may consume the complete causal prefix.
+        latent_indices = range(layout.latent_start, layout.answer_start)
+        for latent_offset, latent in enumerate(latent_indices):
+            visible[:, latent, : layout.multimodal_length] = visual_valid
+            if latent_offset == 0:
+                visible[:, latent, layout.question_slice] = q_valid
+            else:
+                visible[:, latent, latent - 1] = True
+
+        if layout.answer_input_length:
+            answer_valid = answer_input_attention_mask.bool()
+            answer_causal = torch.ones(
+                layout.answer_input_length,
+                layout.answer_input_length,
+                dtype=torch.bool,
+                device=device,
+            ).tril()
+            visible[:, layout.answer_slice, layout.question_slice] = q_valid[:, None, :]
+            visible[:, layout.answer_slice, layout.answer_start - 1] = True
+            visible[:, layout.answer_slice, layout.answer_slice] = (
+                answer_causal.unsqueeze(0) & answer_valid[:, None, :]
+            )
+
+        diagonal = torch.arange(length, device=device)
+        visible[:, diagonal, diagonal] = True
+        return visible.unsqueeze(1)
+
+    for pair_index, (look, think) in enumerate(
+        zip(layout.look_indices, layout.think_indices)
+    ):
+        if final_pair_only and pair_index != layout.num_pairs - 1:
+            continue
+
+        if transition_conditioned and not final_pair_only:
+            # Pair 1 bootstraps the causal latent prefix.  Later pairs cannot
+            # recover the task from the original prompt: only image placeholder
+            # rows remain visible outside the latent prefix.
+            if pair_index == 0:
+                visible[:, look, : layout.multimodal_length] = mm_valid
+                visible[:, look, layout.question_slice] = q_valid
+                visible[:, think, layout.question_slice] = q_valid
+            else:
+                if not disable_late_visual:
+                    visible[:, look, : layout.multimodal_length] = visual_valid
+                if not disable_chain_links:
+                    visible[:, look, layout.latent_start:look] = True
+                    visible[:, think, layout.latent_start:look] = True
+                if pair_index < question_visible_through_pair:
+                    visible[:, look, layout.question_slice] = q_valid
+                    visible[:, think, layout.question_slice] = q_valid
+
+            # The local LOOK -> THINK edge is never an inter-pair chain-link
+            # intervention and therefore remains present in every arm.
+            visible[:, think, look] = True
+        else:
+            # Original graph, also used by the final-pair-only capacity
+            # control: each active pair may independently consume prompt + Q.
+            if not disable_late_visual or pair_index == 0 or final_pair_only:
+                visible[:, look, : layout.multimodal_length] = mm_valid
+            visible[:, look, layout.question_slice] = q_valid
+            if pair_index and not disable_chain_links and not final_pair_only:
+                visible[:, look, layout.think_indices[pair_index - 1]] = True
+
+            # Deliberation: no original multimodal key can be consumed here.
+            visible[:, think, layout.question_slice] = q_valid
+            visible[:, think, look] = True
+
+    if layout.answer_input_length:
+        answer_valid = answer_input_attention_mask.bool()
+        answer_causal = torch.ones(
+            layout.answer_input_length,
+            layout.answer_input_length,
+            dtype=torch.bool,
+            device=device,
+        ).tril()
+        visible[:, layout.answer_slice, layout.question_slice] = q_valid[:, None, :]
+        answer_thinks = (
+            [layout.think_indices[-1]]
+            if final_pair_only or transition_conditioned
+            else list(layout.think_indices)
+        )
+        visible[:, layout.answer_slice, answer_thinks] = True
+        visible[:, layout.answer_slice, layout.answer_slice] = (
+            answer_causal.unsqueeze(0) & answer_valid[:, None, :]
+        )
+
+    # Avoid all-masked softmax rows for physical padding.  Semantic queries do
+    # not receive padded keys because every segment assignment above uses its
+    # validity mask.
+    diagonal = torch.arange(length, device=device)
+    visible[:, diagonal, diagonal] = True
+    return visible.unsqueeze(1)
+
+
+def build_perceive_deliberate_positions(
+    layout: PerceiveDeliberateLayout,
+    multimodal_position_ids: torch.LongTensor,
+    multimodal_attention_mask: torch.Tensor,
+    question_attention_mask: torch.Tensor,
+    answer_input_attention_mask: torch.Tensor | None = None,
+) -> torch.LongTensor:
+    """Extend native multimodal M-RoPE with logical text positions.
+
+    The clean question keeps ordinary relative text positions but is shifted
+    after the largest valid multimodal M-RoPE coordinate.  Latent and answer
+    rows then continue from each item's *logical* question length, independent
+    of right padding.  All three M-RoPE coordinates are equal for these
+    non-spatial rows.
+    """
+
+    if multimodal_position_ids.ndim != 3 or multimodal_position_ids.shape[0] != 3:
+        raise ValueError("multimodal_position_ids must have shape [3, B, L_mm]")
+    batch_size = multimodal_attention_mask.shape[0]
+    if tuple(multimodal_position_ids.shape[1:]) != (
+        batch_size,
+        layout.multimodal_length,
+    ):
+        raise ValueError("multimodal position/layout mismatch")
+
+    mm_valid = multimodal_attention_mask.bool()
+    masked_mm = multimodal_position_ids.masked_fill(~mm_valid.unsqueeze(0), -1)
+    continuation = masked_mm.amax(dim=(0, 2)).clamp_min(-1) + 1  # [B]
+
+    q_valid = question_attention_mask.bool()
+    q_relative = q_valid.long().cumsum(dim=-1) - 1
+    q_relative = q_relative.masked_fill(~q_valid, 0)
+    q_positions = continuation[:, None] + q_relative
+    q_lengths = q_valid.sum(dim=-1)
+
+    latent_relative = torch.arange(
+        layout.num_latent_tokens,
+        device=multimodal_position_ids.device,
+    )
+    latent_positions = (
+        continuation[:, None] + q_lengths[:, None] + latent_relative[None, :]
+    )
+
+    segments = [multimodal_position_ids, q_positions.unsqueeze(0).expand(3, -1, -1)]
+    segments.append(latent_positions.unsqueeze(0).expand(3, -1, -1))
+    if layout.answer_input_length:
+        if answer_input_attention_mask is None:
+            raise ValueError("answer_input_attention_mask is required")
+        answer_relative = torch.arange(
+            layout.answer_input_length,
+            device=multimodal_position_ids.device,
+        )
+        answer_positions = (
+            continuation[:, None]
+            + q_lengths[:, None]
+            + layout.num_latent_tokens
+            + answer_relative[None, :]
+        )
+        segments.append(answer_positions.unsqueeze(0).expand(3, -1, -1))
+    return torch.cat(segments, dim=-1)
+
+
+def build_answer_decode_mask(
+    layout: PerceiveDeliberateLayout,
+    question_attention_mask: torch.Tensor,
+    generated_attention_mask: torch.Tensor,
+    *,
+    local_visual_chain: bool = False,
+    final_pair_only: bool = False,
+    transition_conditioned: bool = False,
+) -> torch.BoolTensor:
+    """Visibility for one cached answer query after the latent prefill.
+
+    ``generated_attention_mask`` includes the current answer token.  Physical
+    multimodal and LOOK cache rows remain present but are invisible.
+    """
+
+    if generated_attention_mask.ndim != 2:
+        raise ValueError("generated_attention_mask must have shape [B, N]")
+    batch_size, generated_length = generated_attention_mask.shape
+    if tuple(question_attention_mask.shape) != (
+        batch_size,
+        layout.question_length,
+    ):
+        raise ValueError("question mask/layout mismatch")
+    total_keys = layout.answer_start + generated_length
+    visible = torch.zeros(
+        batch_size, 1, 1, total_keys, dtype=torch.bool,
+        device=question_attention_mask.device,
+    )
+    visible[:, 0, 0, layout.question_slice] = question_attention_mask.bool()
+    if local_visual_chain:
+        visible[:, 0, 0, layout.answer_start - 1] = True
+    else:
+        answer_thinks = (
+            [layout.think_indices[-1]]
+            if final_pair_only or transition_conditioned
+            else list(layout.think_indices)
+        )
+        visible[:, 0, 0, answer_thinks] = True
+    visible[:, 0, 0, layout.answer_start:] = generated_attention_mask.bool()
+    return visible
+
+
+__all__ = [
+    "PerceiveDeliberateLayout",
+    "build_answer_decode_mask",
+    "build_perceive_deliberate_mask",
+    "build_perceive_deliberate_positions",
+]

source_spatial.py

@@ -0,0 +1,349 @@
+"""Question-conditioned recurrence over the complete native visual field.
+
+The state is the dynamic-resolution sequence emitted by Qwen2.5-VL's vision
+merger, before those image embeddings enter the language model.  Tokens keep
+their native two-dimensional grid throughout the recurrence.  One shared cell
+combines directional local messages with visual-to-question cross-attention::
+
+    V_0 = VisionMerger(image)
+    V_{k+1} = V_k + Cell(V_k, question, grid),  k = 0, ..., T - 1
+
+Only ``V_T`` is inserted into the ordinary Qwen multimodal prefix.  The cell's
+output projection is exactly zero at initialization, so every horizon,
+including the default T=8, is initially bitwise identical to the base visual
+embedding path.  The full language model is then run from scratch; no cache
+created while constructing the visual field is available to answer decoding.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.models.qwen2.modeling_qwen2 import Qwen2RMSNorm
+
+
+@dataclass
+class SpatialVisualTelemetry:
+    """Detached recurrence diagnostics; no visual trajectory is retained."""
+
+    update_rms: torch.Tensor  # [T]
+    update_relative: torch.Tensor  # [T]
+    step_drift: torch.Tensor  # [T], 1 - cos(V_k, V_{k+1})
+    final_relative: torch.Tensor  # scalar, RMS(V_T - V_0) / RMS(V_0)
+
+
+class SpatialVisualRecurrentCell(nn.Module):
+    """Shared O(N) spatial cell for packed, variable-resolution image tokens.
+
+    Global N-by-N visual attention is deliberately absent: a sample may carry
+    8192 merged visual tokens.  Four directional grid messages preserve local
+    geometry in linear time, while low-rank visual-to-question attention makes
+    every patch update task dependent.  The recurrent *state* remains at the
+    native language-model width; ``inner_width`` limits only the update
+    operator's compute and is not a state bottleneck.
+    """
+
+    def __init__(
+        self,
+        width: int,
+        inner_width: int,
+        num_heads: int,
+        *,
+        rms_norm_eps: float,
+        residual_scale: float,
+    ) -> None:
+        super().__init__()
+        if width < 1 or inner_width < 1:
+            raise ValueError("width and inner_width must be positive")
+        if num_heads < 1 or inner_width % num_heads:
+            raise ValueError(
+                "inner_width must be divisible by the positive num_heads"
+            )
+        if not 0.0 <= residual_scale <= 1.0:
+            raise ValueError("residual_scale must lie in [0, 1]")
+        self.width = int(width)
+        self.inner_width = int(inner_width)
+        self.num_heads = int(num_heads)
+        self.head_dim = self.inner_width // self.num_heads
+        self.residual_scale = float(residual_scale)
+
+        self.state_norm = Qwen2RMSNorm(self.width, eps=rms_norm_eps)
+        self.question_norm = Qwen2RMSNorm(self.width, eps=rms_norm_eps)
+
+        # One projection contains distinct north/south/west/east maps.  This
+        # keeps orientation identifiable without allocating a dense 3x3 dxd
+        # convolution over the 3584-wide native visual state.
+        self.center_proj = nn.Linear(self.width, self.inner_width, bias=False)
+        self.neighbor_proj = nn.Linear(
+            self.width, 4 * self.inner_width, bias=False
+        )
+
+        self.query_proj = nn.Linear(self.width, self.inner_width, bias=False)
+        self.key_proj = nn.Linear(self.width, self.inner_width, bias=False)
+        self.value_proj = nn.Linear(self.width, self.inner_width, bias=False)
+        self.output_proj = nn.Linear(self.inner_width, self.width, bias=False)
+
+    def zero_output(self) -> None:
+        """Restore the exact base-model identity after generic HF init."""
+        with torch.no_grad():
+            self.output_proj.weight.zero_()
+
+    @staticmethod
+    def _token_layout(
+        merged_grid_thw: torch.LongTensor,
+        image_to_batch: torch.LongTensor,
+        batch_size: int,
+        expected_tokens: int,
+    ) -> tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]:
+        """Return token-to-batch ids, within-sample positions and counts."""
+        if merged_grid_thw.ndim != 2 or merged_grid_thw.shape[-1] != 3:
+            raise ValueError("merged_grid_thw must have shape [num_images, 3]")
+        if image_to_batch.ndim != 1 or image_to_batch.shape[0] != len(
+            merged_grid_thw
+        ):
+            raise ValueError("image_to_batch must contain one id per image")
+        per_image = merged_grid_thw.prod(dim=-1).long()
+        if int(per_image.sum()) != int(expected_tokens):
+            raise ValueError(
+                "merged grids and visual field disagree: "
+                f"grid tokens={int(per_image.sum())}, state tokens={expected_tokens}"
+            )
+        token_to_batch = torch.repeat_interleave(image_to_batch, per_image)
+        if token_to_batch.numel() and (
+            int(token_to_batch.min()) < 0
+            or int(token_to_batch.max()) >= batch_size
+        ):
+            raise ValueError("image_to_batch contains an out-of-range batch id")
+        if token_to_batch.numel() > 1 and bool(
+            (token_to_batch[1:] < token_to_batch[:-1]).any()
+        ):
+            raise ValueError(
+                "images must follow the same batch-major order as Qwen inputs"
+            )
+        counts = torch.bincount(token_to_batch, minlength=batch_size)
+        starts = counts.cumsum(0) - counts
+        within = torch.arange(
+            expected_tokens, device=token_to_batch.device
+        ) - torch.repeat_interleave(starts, counts)
+        return token_to_batch, within, counts
+
+    @staticmethod
+    def _neighbor_edges(
+        merged_grid_thw: torch.LongTensor,
+        expected_tokens: int,
+    ) -> tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]:
+        """Build packed directed 4-neighbour edges once for all recurrent steps.
+
+        Direction ids are 0/1/2/3 = north/south/west/east *as seen by the
+        destination token*.  Temporal slices are separate 2-D fields.
+        """
+        device = merged_grid_thw.device
+        destinations: list[torch.Tensor] = []
+        sources: list[torch.Tensor] = []
+        directions: list[torch.Tensor] = []
+        offset = 0
+        for t_value, h_value, w_value in merged_grid_thw.detach().cpu().tolist():
+            t, h, w = int(t_value), int(h_value), int(w_value)
+            count = t * h * w
+            if min(t, h, w) < 1:
+                raise ValueError(f"invalid merged visual grid {(t, h, w)}")
+            index = torch.arange(
+                offset, offset + count, device=device, dtype=torch.long
+            ).view(t, h, w)
+
+            def append(dst: torch.Tensor, src: torch.Tensor, direction: int) -> None:
+                if dst.numel() == 0:
+                    return
+                destinations.append(dst.reshape(-1))
+                sources.append(src.reshape(-1))
+                directions.append(
+                    torch.full(
+                        (dst.numel(),), direction, device=device, dtype=torch.long
+                    )
+                )
+
+            append(index[:, 1:, :], index[:, :-1, :], 0)  # north
+            append(index[:, :-1, :], index[:, 1:, :], 1)  # south
+            append(index[:, :, 1:], index[:, :, :-1], 2)  # west
+            append(index[:, :, :-1], index[:, :, 1:], 3)  # east
+            offset += count
+        if offset != expected_tokens:
+            raise ValueError(
+                f"edge grids contain {offset} tokens, expected {expected_tokens}"
+            )
+        if not destinations:
+            empty = torch.empty(0, device=device, dtype=torch.long)
+            return empty, empty, empty
+        return (
+            torch.cat(destinations),
+            torch.cat(sources),
+            torch.cat(directions),
+        )
+
+    def _question_attention(
+        self,
+        normalized_state: torch.Tensor,
+        question_keys: torch.Tensor,
+        question_values: torch.Tensor,
+        question_valid: torch.BoolTensor,
+        token_to_batch: torch.LongTensor,
+        within_sample: torch.LongTensor,
+        token_counts: torch.LongTensor,
+    ) -> torch.Tensor:
+        """Visual-query/text-memory SDPA for a packed visual sequence."""
+        batch_size = question_keys.shape[0]
+        max_tokens = int(token_counts.max()) if token_counts.numel() else 0
+        if max_tokens == 0:
+            return normalized_state.new_empty(0, self.inner_width)
+
+        query_flat = self.query_proj(normalized_state)
+        query = query_flat.new_zeros(
+            batch_size, max_tokens, self.inner_width
+        )
+        query[token_to_batch, within_sample] = query_flat
+
+        bsz, question_length, _ = question_keys.shape
+        query = query.view(
+            bsz, max_tokens, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key = question_keys.view(
+            bsz, question_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        value = question_values.view(
+            bsz, question_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        # Boolean SDPA masks use True for entries that are allowed to attend.
+        allowed = question_valid[:, None, None, :]
+        attended = F.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=allowed,
+            dropout_p=0.0,
+            is_causal=False,
+        )
+        attended = attended.transpose(1, 2).reshape(
+            bsz, max_tokens, self.inner_width
+        )
+        return attended[token_to_batch, within_sample]
+
+    def forward(
+        self,
+        visual_state: torch.Tensor,
+        merged_grid_thw: torch.LongTensor,
+        image_to_batch: torch.LongTensor,
+        question_embeddings: torch.Tensor,
+        question_attention_mask: torch.Tensor | None,
+        *,
+        steps: int,
+    ) -> tuple[torch.Tensor, SpatialVisualTelemetry]:
+        if visual_state.ndim != 2 or visual_state.shape[-1] != self.width:
+            raise ValueError(
+                f"visual_state must be [N,{self.width}], got "
+                f"{tuple(visual_state.shape)}"
+            )
+        if question_embeddings.ndim != 3 or question_embeddings.shape[-1] != self.width:
+            raise ValueError(
+                f"question_embeddings must be [B,Q,{self.width}]"
+            )
+        if steps < 0:
+            raise ValueError("steps must be non-negative")
+        batch_size = question_embeddings.shape[0]
+        token_to_batch, within_sample, token_counts = self._token_layout(
+            merged_grid_thw,
+            image_to_batch,
+            batch_size,
+            visual_state.shape[0],
+        )
+        edge_dst, edge_src, edge_direction = self._neighbor_edges(
+            merged_grid_thw, visual_state.shape[0]
+        )
+
+        if question_attention_mask is None:
+            question_valid = torch.ones(
+                question_embeddings.shape[:2],
+                dtype=torch.bool,
+                device=question_embeddings.device,
+            )
+        else:
+            if question_attention_mask.shape != question_embeddings.shape[:2]:
+                raise ValueError("question attention mask shape mismatch")
+            question_valid = question_attention_mask > 0
+        if not bool(question_valid.any(dim=-1).all()):
+            raise ValueError("every image-bearing sample needs a question token")
+
+        normalized_question = self.question_norm(question_embeddings)
+        question_keys = self.key_proj(normalized_question)
+        question_values = self.value_proj(normalized_question)
+
+        initial = visual_state
+        state = visual_state
+        update_rms: list[torch.Tensor] = []
+        update_relative: list[torch.Tensor] = []
+        step_drift: list[torch.Tensor] = []
+        for _ in range(steps):
+            normalized = self.state_norm(state)
+            center = self.center_proj(normalized)
+            directional = self.neighbor_proj(normalized).view(
+                state.shape[0], 4, self.inner_width
+            )
+            spatial = center.new_zeros(center.shape)
+            degree = center.new_zeros(center.shape[0], 1)
+            if edge_dst.numel():
+                messages = directional[edge_src, edge_direction]
+                spatial.index_add_(0, edge_dst, messages)
+                degree.index_add_(
+                    0,
+                    edge_dst,
+                    torch.ones(
+                        edge_dst.shape[0], 1, dtype=center.dtype, device=center.device
+                    ),
+                )
+                spatial = spatial / degree.clamp_min(1.0)
+            question = self._question_attention(
+                normalized,
+                question_keys,
+                question_values,
+                question_valid,
+                token_to_batch,
+                within_sample,
+                token_counts,
+            )
+            mixed = F.silu((center + spatial + question) / (3.0**0.5))
+            update = self.output_proj(mixed) * self.residual_scale
+            next_state = state + update
+
+            with torch.no_grad():
+                state_rms = state.float().square().mean().sqrt().clamp_min(1e-8)
+                update_norm = update.float().square().mean().sqrt()
+                cosine = F.cosine_similarity(
+                    state.float(), next_state.float(), dim=-1
+                ).mean()
+                update_rms.append(update_norm.detach())
+                update_relative.append((update_norm / state_rms).detach())
+                step_drift.append((1.0 - cosine).detach())
+            state = next_state
+
+        with torch.no_grad():
+            initial_rms = initial.float().square().mean().sqrt().clamp_min(1e-8)
+            final_relative = (
+                (state.float() - initial.float()).square().mean().sqrt()
+                / initial_rms
+            ).detach()
+            empty = initial.new_empty(0, dtype=torch.float32)
+            telemetry = SpatialVisualTelemetry(
+                update_rms=(torch.stack(update_rms) if update_rms else empty),
+                update_relative=(
+                    torch.stack(update_relative) if update_relative else empty
+                ),
+                step_drift=(torch.stack(step_drift) if step_drift else empty),
+                final_relative=final_relative,
+            )
+        return state, telemetry
+
+
+__all__ = ["SpatialVisualRecurrentCell", "SpatialVisualTelemetry"]

source_splitting.py

@@ -0,0 +1,473 @@
+"""Layer-split forward for Qwen2.5-VL and Qwen3-VL.
+
+Method §3.2 needs ``F = F_{>l*} o F_{<=l*}`` as two separately runnable halves so
+that the multimodal branch can be cut off at ``l*`` and replaced by the
+workspace. §3.1 needs the same split to read and patch activations at a chosen
+depth.
+
+This module reimplements the prologue of the native Qwen VL text-model forward
+-- embedding merge, M-RoPE index, causal mask, rotary embeddings -- as a
+reusable :class:`SplitContext`, then exposes the decoder-layer loop as a range
+you can run piecewise.  Qwen3-VL additionally injects three ``DeepStack``
+vision features after language layers 0--2; those tensors are carried in the
+split context and applied at the identical layer boundaries.  It deliberately
+mirrors transformers 4.57.6 rather than monkeypatching it;
+``tests/test_split_equivalence.py`` asserts the composed halves reproduce the
+stock forward bit-for-bit, which is what makes the mirroring safe to rely on.
+
+Shapes use ``B`` batch, ``L`` sequence, ``d`` backbone width (3584 on the 7B),
+``N_v`` visual tokens, ``N_q`` question tokens.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from typing import Any
+
+import torch
+from transformers.cache_utils import Cache
+from transformers.masking_utils import (
+    create_causal_mask,
+    create_sliding_window_causal_mask,
+)
+
+
+@dataclass
+class SplitContext:
+    """Per-forward state shared by every decoder layer.
+
+    Computed once by :func:`make_split_context` so that layer ranges can be run
+    independently without recomputing masks or rotary tables.
+    """
+
+    hidden_states: torch.Tensor  # [B, L, d] - mutated as layers run
+    position_ids: torch.Tensor  # [3, B, L] - M-RoPE (t, h, w)
+    position_embeddings: tuple[torch.Tensor, torch.Tensor]  # (cos, sin) [B, L, head_dim]
+    causal_mask_mapping: dict[str, torch.Tensor | None]
+    cache_position: torch.Tensor  # [L]
+    text_position_ids: torch.Tensor | None  # [B, L] only when packed
+    past_key_values: Cache | None
+    # Original 2-D key-padding mask.  ``create_causal_mask`` is allowed to
+    # return ``None`` for SDPA and delegate causality to ``is_causal``; keeping
+    # this tensor lets counterfactual branches materialize the equivalent mask
+    # before removing a precisely selected set of attention edges.
+    attention_mask: torch.Tensor | None = None  # [B, L_kv]
+    # Qwen3-VL only.  DeepStack adds one visual feature tensor after each of
+    # the first three language layers.  They stay ``None`` for Qwen2.5-VL and
+    # for every text-only branch.
+    visual_pos_masks: torch.Tensor | None = None  # [B, L] bool
+    deepstack_visual_embeds: list[torch.Tensor] | None = None
+
+    def clone_at(self, hidden_states: torch.Tensor) -> "SplitContext":
+        """Same context, different hidden states (for patched re-runs)."""
+        return SplitContext(
+            hidden_states=hidden_states,
+            position_ids=self.position_ids,
+            position_embeddings=self.position_embeddings,
+            causal_mask_mapping=self.causal_mask_mapping,
+            cache_position=self.cache_position,
+            text_position_ids=self.text_position_ids,
+            past_key_values=self.past_key_values,
+            attention_mask=self.attention_mask,
+            visual_pos_masks=self.visual_pos_masks,
+            deepstack_visual_embeds=self.deepstack_visual_embeds,
+        )
+
+
+# ---------------------------------------------------------------------------
+# embedding / position construction
+# ---------------------------------------------------------------------------
+
+
+def embed_multimodal(
+    vl_model,
+    input_ids: torch.LongTensor,  # [B, L]
+    pixel_values: torch.Tensor | None = None,
+    image_grid_thw: torch.LongTensor | None = None,
+    attention_mask: torch.Tensor | None = None,
+    *,
+    return_deepstack: bool = False,
+) -> (
+    tuple[torch.Tensor, torch.Tensor]
+    | tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | None,
+        list[torch.Tensor] | None,
+    ]
+):
+    """Token embeddings with image features scattered in, plus M-RoPE indices.
+
+    Mirrors the prefill path of ``Qwen2_5_VLModel.forward``. Pass
+    ``pixel_values=None`` to get the text-only branch used for ``Q*``.
+
+    Args:
+        vl_model: a native ``Qwen2_5_VLModel`` or ``Qwen3VLModel`` (i.e.
+            ``model.model``, not the ``...ForConditionalGeneration`` wrapper).
+        return_deepstack: also return Qwen3-VL's visual-position mask and
+            DeepStack features.  The default two-value return keeps all
+            Qwen2.5 callers backward compatible.
+
+    Returns:
+        ``(inputs_embeds [B, L, d], position_ids [3, B, L])`` and optionally
+        ``(visual_pos_masks, deepstack_visual_embeds)``.
+    """
+    inputs_embeds = vl_model.get_input_embeddings()(input_ids)  # [B, L, d]
+    model_type = str(getattr(vl_model.config, "model_type", ""))
+    # A freshly wrapped model exposes the native qwen3_vl config here.
+    # Reloading a fully saved CLOSE checkpoint reconstructs the nested native
+    # backbone from the wrapper config, so its model view carries
+    # close_qwen3_vl instead. Both use the tuple-returning Qwen3 feature API.
+    is_qwen3_vl = model_type in {"qwen3_vl", "close_qwen3_vl"}
+    visual_pos_masks = None
+    deepstack_visual_embeds = None
+
+    if pixel_values is not None:
+        image_features = vl_model.get_image_features(pixel_values, image_grid_thw)
+        if is_qwen3_vl:
+            image_embeds, deepstack_visual_embeds = image_features
+        else:
+            image_embeds = image_features
+        image_embeds = torch.cat(image_embeds, dim=0).to(
+            inputs_embeds.device, inputs_embeds.dtype
+        )  # [N_v_total, d]
+        image_mask, _ = vl_model.get_placeholder_mask(
+            input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds
+        )
+        inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+        if is_qwen3_vl:
+            visual_pos_masks = image_mask[..., 0]
+
+    if is_qwen3_vl:
+        position_ids, _ = vl_model.get_rope_index(
+            input_ids,
+            image_grid_thw,
+            None,  # video_grid_thw
+            attention_mask=attention_mask,
+        )
+    else:
+        position_ids, _ = vl_model.get_rope_index(
+            input_ids,
+            image_grid_thw,
+            None,  # video_grid_thw
+            second_per_grid_ts=None,
+            attention_mask=attention_mask,
+        )
+    if return_deepstack:
+        return (
+            inputs_embeds,
+            position_ids,
+            visual_pos_masks,
+            deepstack_visual_embeds,
+        )
+    return inputs_embeds, position_ids
+
+
+def make_split_context(
+    text_model,
+    inputs_embeds: torch.Tensor,  # [B, L, d]
+    position_ids: torch.Tensor,  # [3, B, L]
+    attention_mask: torch.Tensor | None = None,
+    past_key_values: Cache | None = None,
+    cache_position: torch.Tensor | None = None,
+    visual_pos_masks: torch.Tensor | None = None,
+    deepstack_visual_embeds: list[torch.Tensor] | None = None,
+) -> SplitContext:
+    """Build masks and rotary embeddings once, as the stock forward does.
+
+    Args:
+        text_model: ``Qwen2_5_VLTextModel`` (``vl_model.language_model``).
+    """
+    if cache_position is None:
+        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+        cache_position = torch.arange(
+            past_seen, past_seen + inputs_embeds.shape[1], device=inputs_embeds.device
+        )  # [L]
+
+    if position_ids.ndim == 2:
+        position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+    model_type = str(getattr(text_model.config, "model_type", ""))
+
+    # Packed-sequence convention: a leading text-only row makes it [4, B, L].
+    if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+        text_position_ids = position_ids[0]  # [B, L]
+        position_ids = position_ids[1:]  # [3, B, L]
+    elif model_type == "qwen3_vl_text":
+        # Qwen3-VL always passes the temporal M-RoPE row to both the causal-mask
+        # builder and decoder layers, even for ordinary (non-packed) inputs.
+        text_position_ids = position_ids[0]
+    else:
+        text_position_ids = None
+
+    mask_kwargs: dict[str, Any] = {
+        "config": text_model.config,
+        "input_embeds": inputs_embeds,
+        "attention_mask": attention_mask,
+        "cache_position": cache_position,
+        "past_key_values": past_key_values,
+        "position_ids": text_position_ids,
+    }
+    causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)}
+    if getattr(text_model, "has_sliding_layers", False):
+        causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(
+            **mask_kwargs
+        )
+
+    position_embeddings = text_model.rotary_emb(inputs_embeds, position_ids)
+
+    return SplitContext(
+        hidden_states=inputs_embeds,
+        position_ids=position_ids,
+        position_embeddings=position_embeddings,
+        causal_mask_mapping=causal_mask_mapping,
+        cache_position=cache_position,
+        text_position_ids=text_position_ids,
+        past_key_values=past_key_values,
+        attention_mask=attention_mask,
+        visual_pos_masks=visual_pos_masks,
+        deepstack_visual_embeds=deepstack_visual_embeds,
+    )
+
+
+def block_attention_edges(
+    ctx: SplitContext,
+    query_mask: torch.Tensor,
+    key_mask: torch.Tensor,
+) -> SplitContext:
+    """Return ``ctx`` with selected query-to-key attention edges removed.
+
+    ``query_mask`` and ``key_mask`` are boolean ``[B, L]`` supports in the
+    current (cache-free) sequence.  Every ordinary causal/padding constraint is
+    preserved; only their Cartesian product is additionally masked.  Both
+    boolean SDPA masks (``True`` means visible) and additive eager masks
+    (``0``/negative infinity) are supported.
+
+    The helper deliberately rejects cached contexts.  Its intended use is a
+    counterfactual recurrent layer evaluation, never autoregressive decoding,
+    and silently guessing the key offset of a populated cache would invalidate
+    the causal comparison.
+    """
+    if ctx.past_key_values is not None and ctx.past_key_values.get_seq_length() > 0:
+        raise ValueError("block_attention_edges requires a cache-free context")
+    if query_mask.dtype != torch.bool or key_mask.dtype != torch.bool:
+        raise TypeError("query_mask and key_mask must be boolean tensors")
+    if query_mask.shape != key_mask.shape or query_mask.ndim != 2:
+        raise ValueError(
+            "query_mask and key_mask must have the same [B, L] shape, got "
+            f"{tuple(query_mask.shape)} and {tuple(key_mask.shape)}"
+        )
+    batch_size, seq_len = query_mask.shape
+    if ctx.hidden_states.shape[:2] != (batch_size, seq_len):
+        raise ValueError(
+            "edge masks must match the SplitContext sequence, got "
+            f"{tuple(query_mask.shape)} for {tuple(ctx.hidden_states.shape[:2])}"
+        )
+
+    blocked = query_mask[:, None, :, None] & key_mask[:, None, None, :]
+    updated: dict[str, torch.Tensor] = {}
+    for attention_type, base_mask in ctx.causal_mask_mapping.items():
+        if base_mask is None:
+            # SDPA may omit an all-valid causal mask.  Materialize exactly that
+            # lower triangle, then reapply key padding before deleting edges.
+            q_positions = ctx.cache_position
+            if q_positions.numel() != seq_len:
+                raise ValueError(
+                    "cache-free context must have one cache position per row"
+                )
+            key_positions = torch.arange(seq_len, device=query_mask.device)
+            visible = key_positions[None, :] <= q_positions[:, None]
+            visible = visible[None, None].expand(batch_size, 1, -1, -1)
+            if ctx.attention_mask is not None:
+                if ctx.attention_mask.shape != (batch_size, seq_len):
+                    raise ValueError(
+                        "counterfactual edge masking expects a 2-D [B, L] "
+                        "attention mask"
+                    )
+                visible = visible & ctx.attention_mask[:, None, None, :].bool()
+            updated[attention_type] = visible & ~blocked
+            continue
+
+        if not isinstance(base_mask, torch.Tensor) or base_mask.ndim != 4:
+            raise TypeError(
+                "counterfactual edge masking supports tensor 4-D attention "
+                f"masks, got {type(base_mask)!r}"
+            )
+        if base_mask.shape[0] not in (1, batch_size):
+            raise ValueError("attention-mask batch dimension is incompatible")
+        if base_mask.shape[-2:] != (seq_len, seq_len):
+            raise ValueError(
+                "counterfactual edge masking expects a square cache-free mask, "
+                f"got {tuple(base_mask.shape)}"
+            )
+        if base_mask.dtype == torch.bool:
+            updated[attention_type] = base_mask & ~blocked
+        elif base_mask.is_floating_point():
+            updated[attention_type] = base_mask.masked_fill(
+                blocked, torch.finfo(base_mask.dtype).min
+            )
+        else:
+            raise TypeError(
+                f"unsupported attention mask dtype {base_mask.dtype}"
+            )
+
+    return replace(ctx, causal_mask_mapping=updated)
+
+
+# ---------------------------------------------------------------------------
+# running layer ranges
+# ---------------------------------------------------------------------------
+
+
+def run_layer_range(
+    text_model,
+    ctx: SplitContext,
+    start: int,
+    stop: int | None = None,
+    use_cache: bool = False,
+    hidden_states: torch.Tensor | None = None,
+    collect: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
+    """Run ``text_model.layers[start:stop]`` on ``ctx``.
+
+    ``self.norm`` is *not* applied -- it belongs to the very top of the stack.
+    Call :func:`final_norm` after the last range.
+
+    Args:
+        hidden_states: override the context's states (leave ``None`` to chain).
+        collect: also return the input hidden states of every layer in the range
+            plus the range output, i.e. ``stop - start + 1`` tensors.
+
+    Returns:
+        ``[B, L, d]``, or ``(output, collected)`` when ``collect``.
+    """
+    layers = text_model.layers
+    stop = len(layers) if stop is None else stop
+    h = ctx.hidden_states if hidden_states is None else hidden_states
+
+    collected: list[torch.Tensor] = []
+    for layer_index, layer in enumerate(layers[start:stop], start=start):
+        if collect:
+            collected.append(h)
+        attention_type = getattr(layer, "attention_type", "full_attention")
+        h = layer(
+            h,
+            attention_mask=ctx.causal_mask_mapping[attention_type],
+            position_ids=ctx.text_position_ids,
+            past_key_values=ctx.past_key_values,
+            use_cache=use_cache,
+            cache_position=ctx.cache_position,
+            position_embeddings=ctx.position_embeddings,
+        )
+        # 4.57 decoder layers return a bare tensor; older ones returned a tuple.
+        if isinstance(h, tuple):
+            h = h[0]
+        if (
+            ctx.deepstack_visual_embeds is not None
+            and layer_index < len(ctx.deepstack_visual_embeds)
+        ):
+            if ctx.visual_pos_masks is None:
+                raise ValueError("DeepStack features require visual_pos_masks")
+            h = text_model._deepstack_process(
+                h,
+                ctx.visual_pos_masks,
+                ctx.deepstack_visual_embeds[layer_index],
+            )
+
+    if collect:
+        collected.append(h)
+        return h, collected
+    return h
+
+
+def final_norm(text_model, hidden_states: torch.Tensor) -> torch.Tensor:
+    """Apply the stack's final RMSNorm. ``[B, L, d] -> [B, L, d]``."""
+    return text_model.norm(hidden_states)
+
+
+# ---------------------------------------------------------------------------
+# token selection operators (Pi_img / Pi_q in §3.2)
+# ---------------------------------------------------------------------------
+
+
+def image_token_mask(input_ids: torch.LongTensor, image_token_id: int) -> torch.Tensor:
+    """``Pi_img`` support: ``[B, L]`` bool, True at image placeholder positions."""
+    return input_ids == image_token_id
+
+
+def vision_span_mask(input_ids: torch.LongTensor, config) -> torch.Tensor:
+    """``[B, L]`` bool covering ``<|vision_start|>``, image pads, ``<|vision_end|>``.
+
+    Use this (not :func:`image_token_mask`) when *removing* the visual segment to
+    build the text-only branch, so the delimiters do not survive as orphans.
+    """
+    ids = {
+        config.vision_start_token_id,
+        config.vision_end_token_id,
+        config.image_token_id,
+        config.video_token_id,
+    }
+    mask = torch.zeros_like(input_ids, dtype=torch.bool)
+    for tid in ids:
+        mask |= input_ids == tid
+    return mask
+
+
+def select_tokens(
+    hidden_states: torch.Tensor,  # [B, L, d]
+    mask: torch.Tensor,  # [B, L] bool
+) -> torch.Tensor:
+    """Gather masked positions. Requires an equal count per batch element.
+
+    Returns ``[B, N, d]`` where ``N`` is that per-element count.
+    """
+    counts = mask.sum(dim=1)
+    if counts.numel() > 1 and not bool((counts == counts[0]).all()):
+        raise ValueError(
+            f"select_tokens needs the same number of selected tokens per batch "
+            f"element, got {counts.tolist()}. Bucket by visual-token count or "
+            f"gather per-example instead."
+        )
+    n = int(counts[0])
+    b, _, d = hidden_states.shape
+    return hidden_states[mask].view(b, n, d)
+
+
+def select_tokens_padded(
+    hidden_states: torch.Tensor,  # [B, L, d]
+    mask: torch.Tensor,  # [B, L] bool
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Gather masked positions, right-padded to the batch maximum.
+
+    Qwen2.5-VL uses dynamic resolution, so ``N_v`` differs across a batch. §3.1's
+    patching genuinely needs equal counts (it transplants position by position),
+    but ``r_theta`` only cross-attends over ``V*`` -- a variable-length memory is
+    exactly what a key-padding mask is for.
+
+    Returns ``(padded [B, N_max, d], key_padding_mask [B, N_max])`` where the
+    mask is ``True`` at padding, matching ``nn.MultiheadAttention``.
+    """
+    counts = mask.sum(dim=1)
+    n_max = int(counts.max())
+    b, _, d = hidden_states.shape
+    out = hidden_states.new_zeros((b, n_max, d))
+    pad = torch.ones((b, n_max), dtype=torch.bool, device=hidden_states.device)
+    for i in range(b):
+        n = int(counts[i])
+        out[i, :n] = hidden_states[i][mask[i]]
+        pad[i, :n] = False
+    return out, pad
+
+
+__all__ = [
+    "SplitContext",
+    "embed_multimodal",
+    "make_split_context",
+    "block_attention_edges",
+    "run_layer_range",
+    "final_norm",
+    "image_token_mask",
+    "vision_span_mask",
+    "select_tokens",
+    "select_tokens_padded",
+]