configuration.py

# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# 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.

# ============================================================================
#  ZDTaichu-5.0 — Top-Level Configuration
#
#  Architecture:
#    - LLM backbone: Qwen3 (pure Transformer) → Qwen3.5 (hybrid DeltaNet/Transformer)
#      · 3:1 linear-to-full attention ratio (Gated DeltaNet + full attention)
#      · Custom Qwen3_5DynamicCache for hybrid KV / recurrent states
#      · head_dim=256 (was 128), partial_rotary_factor=0.25
#      · Interleaved M-RoPE with 4D position IDs
#      · Attention output gating (sigmoid gate on q_proj)
#    - Vision encoder: C-RADIOv4-H (unchanged)
#    - Token IDs updated for Qwen3.5 vocabulary (vocab_size=248320)
#      · img_context_token_id: 151655 → 248056  (<|image_pad|>)
#      · video_context_token_id: 151656 → 248057 (<|video_pad|>)
#    - Projector output adapts to Qwen3.5 hidden_size (4096 for 9B variant)
# ============================================================================

from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
from .cradio_config import RADIOConfig

logger = logging.get_logger(__name__)

# ---------------------------------------------------------------------------
# Import Qwen3.5 text config — requires transformers >= 5.3.0
# ---------------------------------------------------------------------------
try:
    from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
except ImportError:
    Qwen3_5TextConfig = None
    logger.warning(
        "Could not import Qwen3_5TextConfig from transformers. "
        "Ensure transformers >= 5.3.0 is installed. "
        "Falling back to PretrainedConfig with manual attributes."
    )


# ---------------------------------------------------------------------------
# Default Qwen3.5 text configuration (9B-class variant)
# ---------------------------------------------------------------------------

_LAYER_TYPES_32 = [
    "linear_attention" if bool((i + 1) % 4) else "full_attention"
    for i in range(32)
]
# Result: [lin, lin, lin, full, lin, lin, lin, full, ... lin, lin, lin, full]
# 24 linear + 8 full attention layers


def _default_qwen3_5_text_dict() -> dict:
    """Return a dict of Qwen3.5 text config values (9B-class)."""
    return dict(
        vocab_size=248320,
        hidden_size=4096,
        intermediate_size=12288,
        num_hidden_layers=32,
        num_attention_heads=16,
        num_key_value_heads=4,
        head_dim=256,
        hidden_act="silu",
        max_position_embeddings=262144,
        rms_norm_eps=1e-6,
        use_cache=True,
        tie_word_embeddings=False,
        attention_bias=False,
        attention_dropout=0.0,
        torch_dtype="bfloat16",
        # --- Hybrid layer architecture ---
        layer_types=list(_LAYER_TYPES_32),  # copy to avoid mutation
        full_attention_interval=4,
        # --- Linear attention (Gated DeltaNet) ---
        linear_conv_kernel_dim=4,
        linear_key_head_dim=128,
        linear_value_head_dim=128,
        linear_num_key_heads=16,
        linear_num_value_heads=32,
        # --- RoPE ---
        rope_parameters={
            "rope_type": "default",
            "rope_theta": 10000000,
            "partial_rotary_factor": 0.25,
            "mrope_interleaved": True,
            "mrope_section": [11, 11, 10],
        },
    )


def _build_llm_config(cfg_dict: dict = None) -> PretrainedConfig:
    """
    Construct the LLM sub-config from a dict or defaults.

    Uses Qwen3_5TextConfig when available (transformers >= 5.3);
    otherwise falls back to a plain PretrainedConfig with the correct
    model_type so that AutoModelForCausalLM can still resolve it.
    """
    if cfg_dict is None:
        cfg_dict = _default_qwen3_5_text_dict()

    if Qwen3_5TextConfig is not None:
        return Qwen3_5TextConfig(**cfg_dict)
    else:
        config = PretrainedConfig(**cfg_dict)
        config.model_type = "qwen3_5_text"
        return config


class ZDTaichu5_0_Config(PretrainedConfig):
    """
    Configuration for ZDTaichu-5.0-9B:
        Vision encoder : C-RADIOv4-H  (ViT-H/16, 653 M params)
        LLM decoder    : Qwen3.5      (hybrid DeltaNet/Transformer)
        Projector      : RMSNorm → Linear(5120→20480) → SquaredReLU → Linear(20480→H)

    The projector input side is unchanged (C-RADIOv4-H ViT-H features at 1280,
    pixel-shuffled to 5120). Only the final projection layer adapts to the
    target LLM hidden_size (4096 for the 9B variant, vs 5120 for Qwen3-14B).

    Qwen3.5 hybrid architecture
    ----------------------------
    The text backbone alternates Gated DeltaNet (linear attention) and standard
    multi-head attention layers in a 3:1 ratio. Linear layers use a causal 1D
    convolution + gated delta rule recurrence for O(1) per-token memory during
    generation, while every 4th layer uses full quadratic attention to preserve
    global context. A custom DynamicCache handles both attention KV states and
    recurrent states.
    """

    model_type = "zdtaichu5_0"
    is_composition = True

    def __init__(
        self,
        vision_config=None,
        llm_config=None,
        force_image_size=None,
        downsample_ratio=0.5,
        template=None,
        ps_version="v2",
        image_tag_type="internvl",
        projector_hidden_size=20480,    # 4 × pixel_shuffle_dim (5120)
        vit_hidden_size=1280,           # ViT-H feature dim — same for C-RADIOv4-H
        attn_implementation="flash_attention_2",
        # Special token IDs for Qwen3.5 vocabulary (vocab_size=248320)
        img_context_token_id: int = 248056,   # <|image_pad|>
        video_context_token_id: int = 248057,  # <|video_pad|>
        **kwargs,
    ):

        # ------------------------------------------------------------------
        # Transformers 5.5.x compatibility:
        # PretrainedConfig.__init__ may call self.get_text_config()
        # during token-id validation. Therefore llm_config must exist
        # before calling super().__init__().
        # ------------------------------------------------------------------

        # ── Vision encoder ───────────────────────────────────────────────────
        if vision_config is not None:
            if isinstance(vision_config, dict):
                self.vision_config = RADIOConfig(**vision_config)
            else:
                self.vision_config = vision_config
        else:
            self.vision_config = RADIOConfig(version="c-radio_v4-h")

        # ── Language model (Qwen3.5 hybrid) ──────────────────────────────────
        if llm_config is not None:
            if isinstance(llm_config, PretrainedConfig):
                self.llm_config = llm_config
            elif isinstance(llm_config, dict):
                self.llm_config = _build_llm_config(llm_config)
            else:
                raise TypeError(
                    f"llm_config must be a dict or PretrainedConfig, got {type(llm_config)}"
                )
        else:
            self.llm_config = _build_llm_config(None)


        # Make tokenizer/generation token ids visible early.
        # Transformers 5.5.x may validate these during super().__init__().
        kwargs.setdefault("bos_token_id", getattr(self.llm_config, "bos_token_id", 248040))
        kwargs.setdefault("eos_token_id", getattr(self.llm_config, "eos_token_id", 248044))
        kwargs.setdefault("pad_token_id", getattr(self.llm_config, "pad_token_id", 248040))
        super().__init__(**kwargs)

        self.tie_word_embeddings = getattr(self.llm_config, "tie_word_embeddings", False)

        # ── VL configuration ─────────────────────────────────────────────────
        self.force_image_size = force_image_size
        self.downsample_ratio = downsample_ratio
        self.template = template
        self.ps_version = ps_version
        self.image_tag_type = image_tag_type
        self.projector_hidden_size = projector_hidden_size
        self.vit_hidden_size = vit_hidden_size

        # Special token IDs
        self.img_context_token_id = img_context_token_id
        self.video_context_token_id = video_context_token_id

        # Attention implementation propagation
        self._attn_implementation = attn_implementation
        self.vision_config.use_flash_attn = (
            self._attn_implementation is not None
            and "flash_attention" in self._attn_implementation
        )
        self.llm_config._attn_implementation = self._attn_implementation

    def get_text_config(self, decoder=False):
        # Robust fallback for Transformers 5.5.x validation.
        if hasattr(self, "llm_config"):
            return self.llm_config
        return _build_llm_config(None)

    @property
    def text_config(self):
        return self.get_text_config(decoder=True)