caca_transformers.py

# --- Import Libraries ---
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from transformers import PretrainedConfig, PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
from transformers.generation.utils import GenerationMixin
from safetensors.torch import save_file
from collections import OrderedDict, defaultdict
from typing import Any, Dict, List, Optional, Tuple, Union
import tempfile
import math
import os
import json
import logging
import warnings

logger = logging.getLogger(__name__)

try:
    from flash_attn import flash_attn_func, flash_attn_varlen_func
    HAS_FLASH_ATTN = True
except ImportError:
    HAS_FLASH_ATTN = False

try:
    from xformers.ops import memory_efficient_attention
    HAS_XFORMERS = True
except ImportError:
    HAS_XFORMERS = False

HAS_SDPA = hasattr(F, "scaled_dot_product_attention")

def log(msg):
    print(f"[LOG] {msg}")

try:
    import bitsandbytes as bnb
    HAS_BNB = True
except ImportError:
    HAS_BNB = False

try:
    from peft import LoraConfig, get_peft_model, TaskType
    HAS_PEFT = True
except ImportError:
    HAS_PEFT = False

# --- CacaConfig & CacaModel---
class CacaConfig(PretrainedConfig):
    model_type = "caca"

    def __init__(
        self,
        vocab_size: int = 32000,
        hidden_size: int = 2048,
        intermediate_size: int = 8192,
        num_hidden_layers: int = 24,
        num_attention_heads: int = 32,
        num_key_value_heads: int = 8,
        head_dim: Optional[int] = 64,
        max_position_embeddings: int = 8192,
        rms_norm_eps: float = 1e-6,
        qk_norm_eps: float = 1e-6,
        initializer_range: float = 0.02,
        use_mup: bool = False,
        mup_base_width: int = 256,
        use_cache: bool = True,
        pad_token_id: Optional[int] = None,
        bos_token_id: int = 1,
        eos_token_id: int = 2,
        tie_word_embeddings: bool = False,
        rope_theta: float = 10000.0,
        rope_scaling: Optional[Dict] = None,
        use_rotary_embeddings: bool = True,
        rope_type: str = "default",
        rope_ntk_alpha: float = 1.0,
        use_alibi: bool = False,
        attention_bias: bool = False,
        attention_dropout: float = 0.0,
        attention_temperature: float = 1.0,
        use_qk_norm: bool = True,
        use_flash_attn: bool = True,
        use_grouped_query_attention: bool = False,
        use_multi_query_attention: bool = False,
        sliding_window: Optional[int] = None,
        use_longformer_attention: bool = False,
        longformer_attention_window: int = 512,
        attn_logit_softcapping: Optional[float] = None,
        final_logit_softcapping: Optional[float] = None,
        lm_logit_softcapping: Optional[float] = 30.0,
        attention_sink_size: int = 4,
        attention_sink_window: int = 1024,
        use_attention_sink: bool = False,
        attention_pattern: str = "all_global",
        global_attention_every_n_layers: int = 2,
        mlp_bias: bool = False,
        hidden_dropout: float = 0.1,
        residual_dropout: float = 0.1,
        token_dropout: float = 0.0,
        use_moe: bool = False,
        num_experts: int = 8,
        num_experts_per_tok: int = 2,
        use_expert_choice: bool = False,
        expert_choice_k: float = 0.125,
        router_aux_loss_coef: float = 0.01,
        router_z_loss_coef: float = 0.001,
        moe_layer_frequency: int = 2,
        expert_capacity_factor: float = 1.0,
        use_grouped_moe: bool = False,
        num_expert_groups: int = 1,
        expert_load_monitoring: bool = True,
        expert_load_warn_threshold: float = 0.3,
        use_layer_scale: bool = False,
        layer_scale_init: float = 1e-5,
        use_stochastic_depth: bool = False,
        stochastic_depth_prob: float = 0.1,
        use_spectral_norm: bool = False,
        label_smoothing: float = 0.0,
        use_mixture_of_depths: bool = False,
        mod_capacity_factor: float = 0.5,
        mod_route_method: str = "learned",
        use_cross_attention: bool = False,
        cross_attention_frequency: int = 4,
        use_multimodal: bool = False,
        vision_config: Optional[Dict] = None,
        audio_config: Optional[Dict] = None,
        projector_hidden_size: Optional[int] = None,
        use_soft_merging: bool = False,
        merge_threshold: float = 0.5,
        use_lora: bool = False,
        lora_rank: int = 16,
        lora_alpha: float = 32.0,
        lora_dropout: float = 0.05,
        lora_target_modules: Optional[List[str]] = None,
        lora_bias: str = "none",
        pretraining_tp: int = 1,
        tensor_parallel_size: int = 1,
        pipeline_parallel_size: int = 1,
        gradient_checkpointing_granularity: str = "full",
        nan_recovery_level: int = 1,
        chat_template: Optional[str] = None,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        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.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim if head_dim is not None else (hidden_size // num_attention_heads)
        self.max_position_embeddings = max_position_embeddings
        self.rms_norm_eps = rms_norm_eps
        self.qk_norm_eps = qk_norm_eps
        self.initializer_range = initializer_range
        self.use_mup = use_mup
        self.mup_base_width = mup_base_width
        self.use_cache = use_cache
        self.tie_word_embeddings = tie_word_embeddings
        self.rope_theta = rope_theta
        self.rope_scaling = rope_scaling
        self.use_rotary_embeddings = use_rotary_embeddings
        self.rope_type = rope_type
        self.rope_ntk_alpha = rope_ntk_alpha
        self.use_alibi = use_alibi
        self.attention_bias = attention_bias
        self.attention_dropout = attention_dropout
        self.attention_temperature = attention_temperature
        self.use_qk_norm = use_qk_norm
        self.use_flash_attn = use_flash_attn
        self.use_grouped_query_attention = use_grouped_query_attention
        self.use_multi_query_attention = use_multi_query_attention
        self.sliding_window = sliding_window
        self.use_longformer_attention = use_longformer_attention
        self.longformer_attention_window = longformer_attention_window
        self.attn_logit_softcapping = attn_logit_softcapping
        self.final_logit_softcapping = final_logit_softcapping
        self.lm_logit_softcapping = lm_logit_softcapping
        self.attention_sink_size = attention_sink_size
        self.attention_sink_window = attention_sink_window
        self.use_attention_sink = use_attention_sink
        self.attention_pattern = attention_pattern
        self.global_attention_every_n_layers = global_attention_every_n_layers
        self.mlp_bias = mlp_bias
        self.hidden_dropout = hidden_dropout
        self.residual_dropout = residual_dropout
        self.token_dropout = token_dropout
        self.use_moe = use_moe
        self.num_experts = num_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.use_expert_choice = use_expert_choice
        self.expert_choice_k = expert_choice_k
        self.router_aux_loss_coef = router_aux_loss_coef
        self.router_z_loss_coef = router_z_loss_coef
        self.moe_layer_frequency = moe_layer_frequency
        self.expert_capacity_factor = expert_capacity_factor
        self.use_grouped_moe = use_grouped_moe
        self.num_expert_groups = num_expert_groups
        self.expert_load_monitoring = expert_load_monitoring
        self.expert_load_warn_threshold = expert_load_warn_threshold
        self.use_layer_scale = use_layer_scale
        self.layer_scale_init = layer_scale_init
        self.use_stochastic_depth = use_stochastic_depth
        self.stochastic_depth_prob = stochastic_depth_prob
        self.use_spectral_norm = use_spectral_norm
        self.label_smoothing = label_smoothing
        self.use_mixture_of_depths = use_mixture_of_depths
        self.mod_capacity_factor = mod_capacity_factor
        self.mod_route_method = mod_route_method
        self.use_cross_attention = use_cross_attention
        self.cross_attention_frequency = cross_attention_frequency
        self.use_multimodal = use_multimodal
        self.vision_config = vision_config or {}
        self.audio_config = audio_config or {}
        self.projector_hidden_size = projector_hidden_size or hidden_size
        self.use_soft_merging = use_soft_merging
        self.merge_threshold = merge_threshold
        self.use_lora = use_lora
        self.lora_rank = lora_rank
        self.lora_alpha = lora_alpha
        self.lora_dropout = lora_dropout
        self.lora_target_modules = (
            lora_target_modules if lora_target_modules is not None else ["q_proj", "v_proj"]
        )
        self.lora_bias = lora_bias
        self.pretraining_tp = pretraining_tp
        self.tensor_parallel_size = tensor_parallel_size
        self.pipeline_parallel_size = pipeline_parallel_size
        self.gradient_checkpointing_granularity = gradient_checkpointing_granularity
        self.nan_recovery_level = nan_recovery_level
        self.chat_template = chat_template or (
            "{% for message in messages %}"
            "{% if message['role'] == 'system' %}System: {{ message['content'] }}\n"
            "{% elif message['role'] == 'user' %}User: {{ message['content'] }}\n"
            "{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}\n"
            "{% endif %}{% endfor %}"
            "{% if add_generation_prompt %}Assistant:{% endif %}"
        )

        self._validate_config()

        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 _validate_config(self) -> None:
        errors: List[str] = []

        if self.num_attention_heads % self.num_key_value_heads != 0:
            errors.append(
                f"num_attention_heads ({self.num_attention_heads}) harus habis dibagi "
                f"num_key_value_heads ({self.num_key_value_heads})"
            )
        if self.num_key_value_heads > self.num_attention_heads:
            errors.append(
                f"num_key_value_heads ({self.num_key_value_heads}) tidak boleh melebihi "
                f"num_attention_heads ({self.num_attention_heads})"
            )
        if self.hidden_size % self.num_attention_heads != 0:
            errors.append(
                f"hidden_size ({self.hidden_size}) harus habis dibagi "
                f"num_attention_heads ({self.num_attention_heads})"
            )

        expected_head_dim = self.hidden_size // self.num_attention_heads
        if self.head_dim != expected_head_dim:
            logger.warning(
                f"head_dim ({self.head_dim}) != hidden_size // num_heads ({expected_head_dim}). "
                "Intentional hanya untuk model dengan custom head projection."
            )

        if self.vocab_size <= 0:
            errors.append(f"vocab_size harus > 0, dapat {self.vocab_size}")
        if self.vocab_size > 1_000_000:
            logger.warning(f"vocab_size ({self.vocab_size:,}) sangat besar — bisa menyebabkan masalah memori.")

        if self.use_flash_attn and not HAS_FLASH_ATTN:
            logger.warning("use_flash_attn=True tapi flash-attn belum terinstall. Fallback ke SDPA.")

        if self.sliding_window is not None:
            if self.sliding_window > self.max_position_embeddings:
                errors.append(
                    f"sliding_window ({self.sliding_window}) > "
                    f"max_position_embeddings ({self.max_position_embeddings})"
                )
            if self.sliding_window < 128:
                logger.warning(f"sliding_window ({self.sliding_window}) sangat kecil.")

        if self.use_moe:
            if self.num_experts < self.num_experts_per_tok:
                errors.append(
                    f"num_experts ({self.num_experts}) harus >= "
                    f"num_experts_per_tok ({self.num_experts_per_tok})"
                )
            if self.moe_layer_frequency <= 0:
                errors.append("moe_layer_frequency harus > 0")
            if self.moe_layer_frequency > self.num_hidden_layers:
                logger.warning(
                    f"moe_layer_frequency ({self.moe_layer_frequency}) > "
                    f"num_hidden_layers ({self.num_hidden_layers}) — MoE tidak akan aktif."
                )
            if self.expert_capacity_factor <= 0:
                errors.append("expert_capacity_factor harus > 0")

        if self.use_lora:
            if self.lora_rank <= 0:
                errors.append(f"lora_rank harus > 0, dapat {self.lora_rank}")
            if self.lora_alpha <= 0:
                errors.append(f"lora_alpha harus > 0, dapat {self.lora_alpha}")
            if not HAS_PEFT:
                logger.warning("use_lora=True tapi 'peft' belum terinstall. Jalankan: pip install peft")

        if self.use_mup and self.mup_base_width <= 0:
            errors.append(f"mup_base_width harus > 0, dapat {self.mup_base_width}")

        if self.nan_recovery_level not in (0, 1, 2):
            errors.append(f"nan_recovery_level harus 0, 1, atau 2, dapat {self.nan_recovery_level}")

        if not (0.0 <= self.label_smoothing < 1.0):
            errors.append(f"label_smoothing harus dalam [0.0, 1.0), dapat {self.label_smoothing}")

        if self.final_logit_softcapping is not None and self.lm_logit_softcapping is not None:
            if self.final_logit_softcapping != self.lm_logit_softcapping:
                logger.warning(
                    f"final_logit_softcapping ({self.final_logit_softcapping}) dan "
                    f"lm_logit_softcapping ({self.lm_logit_softcapping}) berbeda. "
                    "lm_logit_softcapping akan dipakai untuk LM head."
                )

        if errors:
            raise ValueError("CacaConfig validation gagal:\n" + "\n".join(f"  • {e}" for e in errors))

    def to_dict(self) -> Dict[str, Any]:
        quant_backup = getattr(self, "quantization_config", None)
        had_quant = hasattr(self, "quantization_config")
        if had_quant and quant_backup is None:
            delattr(self, "quantization_config")
        try:
            output = super().to_dict()
            output["auto_map"] = {
                "AutoConfig": "caca_transformers.CacaConfig",
                "AutoModel": "caca_transformers.CacaModel",
                "AutoModelForCausalLM": "caca_transformers.CacaForCausalLM",
            }
            output["architectures"] = ["CacaForCausalLM"]
        finally:
            if had_quant:
                self.quantization_config = quant_backup
        return output

    def get_mup_lr_multiplier(self, param_name: str) -> float:
        if not self.use_mup:
            return 1.0
        width_ratio = self.mup_base_width / self.hidden_size
        if any(k in param_name for k in ("embed_tokens", "lm_head", "norm")):
            return 1.0
        return width_ratio

    @classmethod
    def from_variant(cls, variant: str, **overrides) -> "CacaConfig":
        if variant not in MODEL_CONFIGS:
            raise ValueError(
                f"Variant '{variant}' tidak dikenal. "
                f"Tersedia: {sorted(MODEL_CONFIGS.keys())}"
            )
        cfg = {**MODEL_CONFIGS[variant], **overrides}
        return cls(**cfg)

MODEL_CONFIGS = {
    "caca-1M": {
        "vocab_size": 4000, "hidden_size": 64, "intermediate_size": 64,
        "num_hidden_layers": 17, "num_attention_heads": 2, "num_key_value_heads": 2,
        "head_dim": 32, "max_position_embeddings": 512,
    },
    "caca-1.5M": {
        "vocab_size": 4000, "hidden_size": 64, "intermediate_size": 64,
        "num_hidden_layers": 40, "num_attention_heads": 2, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 512,
    },
    "caca-2M": {
        "vocab_size": 4000, "hidden_size": 128, "intermediate_size": 256,
        "num_hidden_layers": 7, "num_attention_heads": 4, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 512,
    },
    "caca-2.5M": {
        "vocab_size": 4000, "hidden_size": 128, "intermediate_size": 512,
        "num_hidden_layers": 6, "num_attention_heads": 4, "num_key_value_heads": 2,
        "head_dim": 32, "max_position_embeddings": 512,
    },
    "caca-3M": {
        "vocab_size": 4000, "hidden_size": 160, "intermediate_size": 448,
        "num_hidden_layers": 6, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 512,
    },
    "caca-3.5M": {
        "vocab_size": 4000, "hidden_size": 64, "intermediate_size": 192,
        "num_hidden_layers": 56, "num_attention_heads": 2, "num_key_value_heads": 2,
        "head_dim": 32, "max_position_embeddings": 512,
    },
    "caca-4M": {
        "vocab_size": 4000, "hidden_size": 96, "intermediate_size": 448,
        "num_hidden_layers": 21, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 512,
    },
    "caca-4.5M": {
        "vocab_size": 4000, "hidden_size": 64, "intermediate_size": 128,
        "num_hidden_layers": 97, "num_attention_heads": 2, "num_key_value_heads": 2,
        "head_dim": 32, "max_position_embeddings": 512,
    },
    "caca-5M": {
        "vocab_size": 4000, "hidden_size": 256, "intermediate_size": 512,
        "num_hidden_layers": 5, "num_attention_heads": 8, "num_key_value_heads": 4,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-6M": {
        "vocab_size": 8000, "hidden_size": 192, "intermediate_size": 576,
        "num_hidden_layers": 7, "num_attention_heads": 6, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-7M": {
        "vocab_size": 8000, "hidden_size": 64, "intermediate_size": 256,
        "num_hidden_layers": 91, "num_attention_heads": 2, "num_key_value_heads": 2,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-8M": {
        "vocab_size": 8000, "hidden_size": 96, "intermediate_size": 64,
        "num_hidden_layers": 131, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-9M": {
        "vocab_size": 8000, "hidden_size": 448, "intermediate_size": 1024,
        "num_hidden_layers": 1, "num_attention_heads": 14, "num_key_value_heads": 2,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-10M": {
        "vocab_size": 8000, "hidden_size": 256, "intermediate_size": 512,
        "num_hidden_layers": 10, "num_attention_heads": 8, "num_key_value_heads": 4,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-12M": {
        "vocab_size": 8000, "hidden_size": 96, "intermediate_size": 256,
        "num_hidden_layers": 100, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-14M": {
        "vocab_size": 8000, "hidden_size": 160, "intermediate_size": 512,
        "num_hidden_layers": 36, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-15M": {
        "vocab_size": 8000, "hidden_size": 64, "intermediate_size": 256,
        "num_hidden_layers": 227, "num_attention_heads": 2, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-17M": {
        "vocab_size": 8000, "hidden_size": 96, "intermediate_size": 256,
        "num_hidden_layers": 157, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-20M": {
        "vocab_size": 8000, "hidden_size": 160, "intermediate_size": 448,
        "num_hidden_layers": 63, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 32, "max_position_embeddings": 1024,
    },
    "caca-22M": {
        "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 704,
        "num_hidden_layers": 30, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 1024,
    },
    "caca-25M": {
        "vocab_size": 16000, "hidden_size": 576, "intermediate_size": 1344,
        "num_hidden_layers": 2, "num_attention_heads": 9, "num_key_value_heads": 3,
        "head_dim": 64, "max_position_embeddings": 1024,
    },
    "caca-27M": {
        "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 512,
        "num_hidden_layers": 53, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 1024,
    },
    "caca-30M": {
        "vocab_size": 16000, "hidden_size": 448, "intermediate_size": 1280,
        "num_hidden_layers": 7, "num_attention_heads": 7, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 1024,
    },
    "caca-35M": {
        "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 512,
        "num_hidden_layers": 69, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 1024,
    },
    "caca-40M": {
        "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 640,
        "num_hidden_layers": 33, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 1024,
    },
    "caca-45M": {
        "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 832,
        "num_hidden_layers": 32, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 1024,
    },
    "caca-50M": {
        "vocab_size": 16000, "hidden_size": 576, "intermediate_size": 1728,
        "num_hidden_layers": 8, "num_attention_heads": 9, "num_key_value_heads": 3,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-55M": {
        "vocab_size": 16000, "hidden_size": 128, "intermediate_size": 512,
        "num_hidden_layers": 194, "num_attention_heads": 2, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-60M": {
        "vocab_size": 16000, "hidden_size": 256, "intermediate_size": 640,
        "num_hidden_layers": 79, "num_attention_heads": 4, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-65M": {
        "vocab_size": 16000, "hidden_size": 256, "intermediate_size": 704,
        "num_hidden_layers": 77, "num_attention_heads": 4, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-70M": {
        "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 448,
        "num_hidden_layers": 179, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-80M": {
        "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 448,
        "num_hidden_layers": 207, "num_attention_heads": 3, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-90M": {
        "vocab_size": 16000, "hidden_size": 768, "intermediate_size": 1984,
        "num_hidden_layers": 11, "num_attention_heads": 12, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-100M": {
        "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 1024,
        "num_hidden_layers": 73, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-110M": {
        "vocab_size": 16000, "hidden_size": 256, "intermediate_size": 896,
        "num_hidden_layers": 115, "num_attention_heads": 4, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-120M": {
        "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 832,
        "num_hidden_layers": 105, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-130M": {
        "vocab_size": 16000, "hidden_size": 448, "intermediate_size": 1024,
        "num_hidden_layers": 63, "num_attention_heads": 7, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-140M": {
        "vocab_size": 16000, "hidden_size": 384, "intermediate_size": 832,
        "num_hidden_layers": 98, "num_attention_heads": 6, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-150M": {
        "vocab_size": 16000, "hidden_size": 640, "intermediate_size": 1728,
        "num_hidden_layers": 29, "num_attention_heads": 10, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-160M": {
        "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 640,
        "num_hidden_layers": 174, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-175M": {
        "vocab_size": 16000, "hidden_size": 1536, "intermediate_size": 4096,
        "num_hidden_layers": 5, "num_attention_heads": 24, "num_key_value_heads": 8,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-200M": {
        "vocab_size": 16000, "hidden_size": 256, "intermediate_size": 896,
        "num_hidden_layers": 225, "num_attention_heads": 4, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-225M": {
        "vocab_size": 16000, "hidden_size": 512, "intermediate_size": 1408,
        "num_hidden_layers": 74, "num_attention_heads": 8, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-250M": {
        "vocab_size": 16000, "hidden_size": 448, "intermediate_size": 1216,
        "num_hidden_layers": 104, "num_attention_heads": 7, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-275M": {
        "vocab_size": 16000, "hidden_size": 576, "intermediate_size": 1536,
        "num_hidden_layers": 71, "num_attention_heads": 9, "num_key_value_heads": 3,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-300M": {
        "vocab_size": 16000, "hidden_size": 512, "intermediate_size": 1408,
        "num_hidden_layers": 103, "num_attention_heads": 8, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-325M": {
        "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 1024,
        "num_hidden_layers": 256, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-350M": {
        "vocab_size": 16000, "hidden_size": 384, "intermediate_size": 1088,
        "num_hidden_layers": 205, "num_attention_heads": 6, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-375M": {
        "vocab_size": 16000, "hidden_size": 384, "intermediate_size": 960,
        "num_hidden_layers": 250, "num_attention_heads": 6, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-400M": {
        "vocab_size": 16000, "hidden_size": 512, "intermediate_size": 1344,
        "num_hidden_layers": 141, "num_attention_heads": 8, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-450M": {
        "vocab_size": 16000, "hidden_size": 512, "intermediate_size": 1536,
        "num_hidden_layers": 147, "num_attention_heads": 8, "num_key_value_heads": 1,
        "head_dim": 64, "max_position_embeddings": 2048,
    },
    "caca-500M": {
        "vocab_size": 16000, "hidden_size": 1024, "intermediate_size": 2688,
        "num_hidden_layers": 44, "num_attention_heads": 16, "num_key_value_heads": 2,
        "head_dim": 64, "max_position_embeddings": 4096,
    },
    "caca-550M": {
        "vocab_size": 32000, "hidden_size": 640, "intermediate_size": 1920,
        "num_hidden_layers": 109, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-600M": {
        "vocab_size": 32000, "hidden_size": 5120, "intermediate_size": 13632,
        "num_hidden_layers": 1, "num_attention_heads": 40, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-650M": {
        "vocab_size": 32000, "hidden_size": 768, "intermediate_size": 1984,
        "num_hidden_layers": 101, "num_attention_heads": 6, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-700M": {
        "vocab_size": 32000, "hidden_size": 3072, "intermediate_size": 8192,
        "num_hidden_layers": 5, "num_attention_heads": 24, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-800M": {
        "vocab_size": 32000, "hidden_size": 640, "intermediate_size": 1920,
        "num_hidden_layers": 157, "num_attention_heads": 5, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-900M": {
        "vocab_size": 32000, "hidden_size": 896, "intermediate_size": 2432,
        "num_hidden_layers": 93, "num_attention_heads": 7, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-1B": {
        "vocab_size": 32000, "hidden_size": 1024, "intermediate_size": 2688,
        "num_hidden_layers": 88, "num_attention_heads": 8, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-1.2B": {
        "vocab_size": 32000, "hidden_size": 768, "intermediate_size": 2240,
        "num_hidden_layers": 176, "num_attention_heads": 6, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-1.5B": {
        "vocab_size": 32000, "hidden_size": 1536, "intermediate_size": 4032,
        "num_hidden_layers": 53, "num_attention_heads": 12, "num_key_value_heads": 6,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-1.8B": {
        "vocab_size": 32000, "hidden_size": 896, "intermediate_size": 2304,
        "num_hidden_layers": 211, "num_attention_heads": 7, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-2B": {
        "vocab_size": 32000, "hidden_size": 896, "intermediate_size": 2432,
        "num_hidden_layers": 232, "num_attention_heads": 7, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-2.5B": {
        "vocab_size": 32000, "hidden_size": 1152, "intermediate_size": 3072,
        "num_hidden_layers": 175, "num_attention_heads": 9, "num_key_value_heads": 1,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-3B": {
        "vocab_size": 32000, "hidden_size": 1280, "intermediate_size": 3584,
        "num_hidden_layers": 159, "num_attention_heads": 10, "num_key_value_heads": 2,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-3.5B": {
        "vocab_size": 32000, "hidden_size": 2304, "intermediate_size": 6144,
        "num_hidden_layers": 58, "num_attention_heads": 18, "num_key_value_heads": 6,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-4B": {
        "vocab_size": 32000, "hidden_size": 3072, "intermediate_size": 8192,
        "num_hidden_layers": 39, "num_attention_heads": 24, "num_key_value_heads": 4,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-5B": {
        "vocab_size": 32000, "hidden_size": 1536, "intermediate_size": 4288,
        "num_hidden_layers": 194, "num_attention_heads": 12, "num_key_value_heads": 2,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-6B": {
        "vocab_size": 32000, "hidden_size": 1536, "intermediate_size": 4032,
        "num_hidden_layers": 245, "num_attention_heads": 12, "num_key_value_heads": 2,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-7B": {
        "vocab_size": 32000, "hidden_size": 12288, "intermediate_size": 32576,
        "num_hidden_layers": 4, "num_attention_heads": 96, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-8B": {
        "vocab_size": 32000, "hidden_size": 3584, "intermediate_size": 9728,
        "num_hidden_layers": 58, "num_attention_heads": 28, "num_key_value_heads": 4,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-9B": {
        "vocab_size": 32000, "hidden_size": 2560, "intermediate_size": 7040,
        "num_hidden_layers": 122, "num_attention_heads": 20, "num_key_value_heads": 5,
        "head_dim": 128, "max_position_embeddings": 4096,
    },
    "caca-10B": {
        "vocab_size": 32000, "hidden_size": 2560, "intermediate_size": 6784,
        "num_hidden_layers": 145, "num_attention_heads": 20, "num_key_value_heads": 4,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-13B": {
        "vocab_size": 32000, "hidden_size": 2560, "intermediate_size": 6656,
        "num_hidden_layers": 192, "num_attention_heads": 20, "num_key_value_heads": 4,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-14B": {
        "vocab_size": 32000, "hidden_size": 2304, "intermediate_size": 6336,
        "num_hidden_layers": 244, "num_attention_heads": 18, "num_key_value_heads": 3,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-16B": {
        "vocab_size": 32000, "hidden_size": 5120, "intermediate_size": 13440,
        "num_hidden_layers": 56, "num_attention_heads": 40, "num_key_value_heads": 10,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-18B": {
        "vocab_size": 32000, "hidden_size": 6144, "intermediate_size": 16448,
        "num_hidden_layers": 45, "num_attention_heads": 48, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-20B": {
        "vocab_size": 32000, "hidden_size": 5120, "intermediate_size": 13696,
        "num_hidden_layers": 72, "num_attention_heads": 40, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-25B": {
        "vocab_size": 32000, "hidden_size": 3584, "intermediate_size": 9728,
        "num_hidden_layers": 180, "num_attention_heads": 28, "num_key_value_heads": 7,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-30B": {
        "vocab_size": 32000, "hidden_size": 3584, "intermediate_size": 9344,
        "num_hidden_layers": 223, "num_attention_heads": 28, "num_key_value_heads": 7,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-35B": {
        "vocab_size": 32000, "hidden_size": 4096, "intermediate_size": 11008,
        "num_hidden_layers": 196, "num_attention_heads": 32, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-40B": {
        "vocab_size": 32000, "hidden_size": 10240, "intermediate_size": 27392,
        "num_hidden_layers": 36, "num_attention_heads": 80, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-45B": {
        "vocab_size": 32000, "hidden_size": 6144, "intermediate_size": 16448,
        "num_hidden_layers": 114, "num_attention_heads": 48, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-50B": {
        "vocab_size": 32000, "hidden_size": 16384, "intermediate_size": 43712,
        "num_hidden_layers": 18, "num_attention_heads": 128, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-65B": {
        "vocab_size": 32000, "hidden_size": 6144, "intermediate_size": 16448,
        "num_hidden_layers": 160, "num_attention_heads": 48, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-70B": {
        "vocab_size": 32000, "hidden_size": 6144, "intermediate_size": 16320,
        "num_hidden_layers": 179, "num_attention_heads": 48, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-80B": {
        "vocab_size": 32000, "hidden_size": 12288, "intermediate_size": 32576,
        "num_hidden_layers": 51, "num_attention_heads": 96, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-90B": {
        "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54400,
        "num_hidden_layers": 21, "num_attention_heads": 160, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-100B": {
        "vocab_size": 32000, "hidden_size": 10240, "intermediate_size": 27264,
        "num_hidden_layers": 93, "num_attention_heads": 80, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-120B": {
        "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54656,
        "num_hidden_layers": 28, "num_attention_heads": 160, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-150B": {
        "vocab_size": 32000, "hidden_size": 8192, "intermediate_size": 22016,
        "num_hidden_layers": 216, "num_attention_heads": 64, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-175B": {
        "vocab_size": 32000, "hidden_size": 10240, "intermediate_size": 27520,
        "num_hidden_layers": 162, "num_attention_heads": 80, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-200B": {
        "vocab_size": 32000, "hidden_size": 10240, "intermediate_size": 27264,
        "num_hidden_layers": 183, "num_attention_heads": 80, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-250B": {
        "vocab_size": 32000, "hidden_size": 12288, "intermediate_size": 32960,
        "num_hidden_layers": 159, "num_attention_heads": 96, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-300B": {
        "vocab_size": 32000, "hidden_size": 24576, "intermediate_size": 65536,
        "num_hidden_layers": 49, "num_attention_heads": 192, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-350B": {
        "vocab_size": 32000, "hidden_size": 14336, "intermediate_size": 38272,
        "num_hidden_layers": 165, "num_attention_heads": 112, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-400B": {
        "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54784,
        "num_hidden_layers": 93, "num_attention_heads": 160, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-450B": {
        "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54528,
        "num_hidden_layers": 105, "num_attention_heads": 160, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-500B": {
        "vocab_size": 32000, "hidden_size": 24576, "intermediate_size": 65728,
        "num_hidden_layers": 81, "num_attention_heads": 192, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-600B": {
        "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54592,
        "num_hidden_layers": 140, "num_attention_heads": 160, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-700B": {
        "vocab_size": 32000, "hidden_size": 24576, "intermediate_size": 65344,
        "num_hidden_layers": 114, "num_attention_heads": 192, "num_key_value_heads": 16,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-800B": {
        "vocab_size": 32000, "hidden_size": 24576, "intermediate_size": 65600,
        "num_hidden_layers": 131, "num_attention_heads": 192, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-900B": {
        "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54656,
        "num_hidden_layers": 212, "num_attention_heads": 160, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
    "caca-1T": {
        "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54528,
        "num_hidden_layers": 236, "num_attention_heads": 160, "num_key_value_heads": 8,
        "head_dim": 128, "max_position_embeddings": 8192,
    },
}

# --- MetricsTracker ---
class MetricsTracker:
    def __init__(self, reset_interval: int = 100):
        self.metrics: Dict[str, List[float]] = defaultdict(list)
        self.reset_interval = reset_interval
        self.step_count = 0

    def log(self, name: str, value: Union[float, torch.Tensor]) -> None:
        if isinstance(value, torch.Tensor):
            value = value.item()
        self.metrics[name].append(value)

    def step(self) -> None:
        self.step_count += 1
        if self.step_count % self.reset_interval == 0:
            self.clear()

    def get_summary(self) -> Dict[str, Dict[str, float]]:
        return {
            name: {
                "mean": float(np.mean(vals)),
                "std":  float(np.std(vals)),
                "min":  float(np.min(vals)),
                "max":  float(np.max(vals)),
                "last": float(vals[-1]),
            }
            for name, vals in self.metrics.items() if vals
        }

    def clear(self) -> None:
        self.metrics.clear()

# --- KV Cache ---
class DynamicCache:
    def __init__(self):
        self.key_cache:   List[torch.Tensor] = []
        self.value_cache: List[torch.Tensor] = []
        self._seen_tokens: int = 0

    def __len__(self):
        return len(self.key_cache)

    def get_seq_length(self, layer_idx: int = 0) -> int:
        return self.key_cache[layer_idx].shape[-2] if len(self.key_cache) > layer_idx else 0

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if layer_idx == 0:
            self._seen_tokens += key_states.shape[-2]
        if len(self.key_cache) == layer_idx:
            self.key_cache.append(key_states)
            self.value_cache.append(value_states)
        else:
            self.key_cache[layer_idx]   = torch.cat([self.key_cache[layer_idx],   key_states],   dim=-2)
            self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)
        return self.key_cache[layer_idx], self.value_cache[layer_idx]

    def reorder_cache(self, beam_idx: torch.Tensor) -> None:
        for i in range(len(self.key_cache)):
            dev = self.key_cache[i].device
            self.key_cache[i]   = self.key_cache[i].index_select(0,   beam_idx.to(dev))
            self.value_cache[i] = self.value_cache[i].index_select(0, beam_idx.to(dev))

    def crop(self, max_length: int) -> None:

        self._seen_tokens = max(0, max_length)
        for i in range(len(self.key_cache)):
            self.key_cache[i]   = self.key_cache[i][...,   :max_length, :]
            self.value_cache[i] = self.value_cache[i][..., :max_length, :]

    def to_legacy_cache(self) -> Tuple:
        return tuple(zip(self.key_cache, self.value_cache))

    @classmethod
    def from_legacy_cache(cls, past: Optional[Tuple]) -> "DynamicCache":
        cache = cls()
        if past:
            for i, (k, v) in enumerate(past):
                cache.key_cache.append(k)
                cache.value_cache.append(v)
                if i == 0:
                    cache._seen_tokens = k.shape[-2]
        return cache

class SlidingWindowCache(DynamicCache):
    def __init__(self, window_size: int, sink_size: int = 4):
        super().__init__()
        self.window_size = window_size
        self.sink_size   = sink_size

    def update(self, key_states, value_states, layer_idx):
        k, v = super().update(key_states, value_states, layer_idx)
        max_len = self.window_size + self.sink_size
        if k.shape[-2] > max_len:
            sink_k = k[..., :self.sink_size, :]
            sink_v = v[..., :self.sink_size, :]
            win_k  = k[..., -self.window_size:, :]
            win_v  = v[..., -self.window_size:, :]
            self.key_cache[layer_idx]   = torch.cat([sink_k, win_k], dim=-2)
            self.value_cache[layer_idx] = torch.cat([sink_v, win_v], dim=-2)

            if layer_idx == 0:
                self._seen_tokens = self.key_cache[layer_idx].shape[-2]
        return self.key_cache[layer_idx], self.value_cache[layer_idx]

# --- Normalisasi ---
class CacaRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        x = x.float()
        var = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(var + self.variance_epsilon)
        return (self.weight * x).to(dtype)

class LayerScale(nn.Module):
    def __init__(self, dim: int, init_value: float = 1e-5):
        super().__init__()
        self.gamma = nn.Parameter(init_value * torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.gamma * x

class StochasticDepth(nn.Module):
    def __init__(self, drop_prob: float = 0.0):
        super().__init__()
        self.drop_prob = drop_prob

    def forward(self, x: torch.Tensor, training: bool = True) -> torch.Tensor:
        if not training or self.drop_prob == 0.0:
            return x
        keep  = 1 - self.drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
        mask  = keep + torch.rand(shape, dtype=x.dtype, device=x.device)
        return x.div(keep) * mask.floor_()

# --- Positional Encoding ---
def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin):
    cos, sin = cos.to(q.dtype), sin.to(q.dtype)
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)

class CacaRotaryEmbedding(nn.Module):
    def __init__(
        self,
        dim: int,
        max_position_embeddings: int = 8192,
        base: float = 10000.0,
        scaling_factor: float = 1.0,
        scaling_type: Optional[str] = None,
        rope_type: str = "default",
        ntk_alpha: float = 1.0,
    ):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        self.scaling_factor = scaling_factor
        self.scaling_type   = scaling_type
        self.rope_type      = rope_type
        self.ntk_alpha      = ntk_alpha

        inv_freq = self._build_inv_freq()
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        self._cos_cache: Optional[torch.Tensor] = None
        self._sin_cache: Optional[torch.Tensor] = None
        self._cached_seq_len: int = 0

    def _build_inv_freq(self) -> torch.Tensor:
        freqs    = torch.arange(0, self.dim, 2).float() / self.dim
        inv_freq = 1.0 / (self.base ** freqs)

        if self.rope_type == "ntk":
            effective_base = self.base * (self.ntk_alpha ** (self.dim / (self.dim - 2)))
            inv_freq = 1.0 / (effective_base ** freqs)
        elif self.rope_type in ("linear", "yarn"):
            inv_freq = inv_freq / self.scaling_factor
        # dynamic: ditangani saat forward

        return inv_freq

    def forward(self, x: torch.Tensor, seq_len: int, position_offset: int = 0):
        if (
            self._cos_cache is not None
            and self._cached_seq_len >= seq_len + position_offset
            and position_offset == 0
        ):
            return (
                self._cos_cache[:, :, :seq_len, :].to(x.dtype),
                self._sin_cache[:, :, :seq_len, :].to(x.dtype),
            )

        t = torch.arange(position_offset, position_offset + seq_len, device=x.device).float()

        if self.rope_type == "dynamic" and seq_len > self.max_position_embeddings:
            scale = seq_len / self.max_position_embeddings
            t = t / scale

        freqs = torch.outer(t, self.inv_freq)
        emb   = torch.cat([freqs, freqs], dim=-1)
        cos   = emb.cos()[None, None]
        sin   = emb.sin()[None, None]

        if position_offset == 0:
            self._cos_cache      = cos
            self._sin_cache      = sin
            self._cached_seq_len = seq_len

        return cos.to(x.dtype), sin.to(x.dtype)

class ALiBiPositionalBias(nn.Module):
    def __init__(self, num_heads: int, max_positions: int = 8192):
        super().__init__()
        slopes = torch.tensor(self._get_slopes(num_heads))
        self.register_buffer("slopes", slopes, persistent=False)

    @staticmethod
    def _get_slopes(n: int) -> List[float]:
        def _pow2(n):
            s = 2 ** (-(2 ** -(math.log2(n) - 3)))
            return [s * (s ** i) for i in range(n)]
        if math.log2(n).is_integer():
            return _pow2(n)
        p = 2 ** math.floor(math.log2(n))
        return _pow2(p) + ALiBiPositionalBias._get_slopes(2 * p)[::2][: n - p]

    def forward(self, seq_len: int, key_len: Optional[int] = None) -> torch.Tensor:
        kl = key_len or seq_len
        qp = torch.arange(seq_len, device=self.slopes.device).unsqueeze(1)
        kp = torch.arange(kl,      device=self.slopes.device).unsqueeze(0)
        bias = (kp - qp).unsqueeze(0) * self.slopes.unsqueeze(1).unsqueeze(2)
        return bias.unsqueeze(0)

# --- Logit softcap ---
def soft_cap_logits(x: torch.Tensor, cap: Optional[float]) -> torch.Tensor:
    if cap is None or cap <= 0:
        return x
    return cap * torch.tanh(x / cap)

# --- MoE Router ---
class TopKRouter(nn.Module):
    def __init__(self, hidden_size: int, num_experts: int, num_experts_per_tok: int):
        super().__init__()
        self.num_experts         = num_experts
        self.num_experts_per_tok = num_experts_per_tok
        self.gate      = nn.Linear(hidden_size, num_experts, bias=False)
        self.gate_norm = nn.LayerNorm(hidden_size)
        self.temperature   = nn.Parameter(torch.ones(1))
        self.jitter_noise  = 0.01
        self.register_buffer("_load_history", torch.zeros(num_experts))
        self._load_steps = 0

    def forward(self, hidden_states: torch.Tensor):
        flat   = self.gate_norm(hidden_states.view(-1, hidden_states.shape[-1]))
        logits = self.gate(flat)
        logits = torch.clamp(logits, -20, 20)
        temp = F.softplus(self.temperature).clamp(min=0.1, max=10.0)
        logits = logits / temp

        if self.training and self.jitter_noise > 0:
            logits = logits + torch.randn_like(logits) * self.jitter_noise

        probs  = F.softmax(logits, dim=-1, dtype=torch.float32)
        top_w, top_idx = torch.topk(probs, self.num_experts_per_tok, dim=-1)
        top_w  = top_w / (top_w.sum(-1, keepdim=True) + 1e-9)

        usage    = probs.mean(0)
        aux_loss = usage.std() / (usage.mean() + 1e-10)
        z_loss   = torch.logsumexp(logits.float(), -1).pow(2).mean()
        ent_loss = -(probs * (probs + 1e-10).log()).sum(-1).mean() * (-0.01)

        if self.training:
            self._load_history = 0.9 * self._load_history + 0.1 * usage.detach()
            self._load_steps  += 1

        return top_w, top_idx, aux_loss + ent_loss, z_loss

    def get_load_imbalance(self) -> float:
        if self._load_steps == 0:
            return 0.0
        m = self._load_history.mean()
        return (self._load_history.std() / (m + 1e-10)).item()

class ExpertChoiceRouter(nn.Module):
    def __init__(self, hidden_size: int, num_experts: int, expert_choice_k: float):
        super().__init__()
        self.num_experts     = num_experts
        self.expert_choice_k = expert_choice_k
        self.gate = nn.Linear(hidden_size, num_experts, bias=False)

    def forward(self, hidden_states: torch.Tensor):
        bs, sl, hs = hidden_states.shape
        flat   = hidden_states.view(-1, hs)
        logits = self.gate(flat)
        probs  = F.softmax(logits, -1, dtype=torch.float32)
        cap    = max(1, int(self.expert_choice_k * bs * sl / self.num_experts))
        _, top_idx = torch.topk(probs.t(), min(cap, bs * sl), dim=-1)
        mask = torch.zeros(self.num_experts, bs * sl, device=flat.device)
        for e in range(self.num_experts):
            mask[e, top_idx[e]] = 1.0
        weights  = mask.t() * probs
        aux_loss = (probs.mean(0) ** 2).sum() * self.num_experts
        z_loss   = torch.logsumexp(logits, -1).mean()
        return weights, aux_loss, z_loss

# --- Expert & MoE ---
class Expert(nn.Module):
    def __init__(self, config: CacaConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)
        self.dropout   = nn.Dropout(config.hidden_dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.dropout(F.silu(self.gate_proj(x)) * self.up_proj(x)))

class MixtureOfExperts(nn.Module):
    def __init__(self, config: CacaConfig):
        super().__init__()
        self.config            = config
        self.num_experts       = config.num_experts
        self.use_expert_choice = config.use_expert_choice
        self.load_monitoring   = config.expert_load_monitoring

        self.experts = nn.ModuleList([Expert(config) for _ in range(config.num_experts)])
        self.router  = (
            ExpertChoiceRouter(config.hidden_size, config.num_experts, config.expert_choice_k)
            if config.use_expert_choice
            else TopKRouter(config.hidden_size, config.num_experts, config.num_experts_per_tok)
        )
        self.register_buffer("expert_usage_count", torch.zeros(config.num_experts, dtype=torch.long))

    def _safe_expert(self, expert_fn, inp: torch.Tensor, expert_idx: int, config: CacaConfig):
        try:
            out = expert_fn(inp)
            if torch.isnan(out).any() or torch.isinf(out).any():
                level = config.nan_recovery_level
                if level >= 1:
                    logger.warning(f"Expert {expert_idx}: output NaN/Inf — di-zero.")
                if level == 0:
                    return out
                return torch.zeros_like(out)
            return out
        except RuntimeError as e:
            logger.error(f"Expert {expert_idx} RuntimeError: {e}")
            return torch.zeros(inp.shape[0], config.hidden_size, device=inp.device, dtype=inp.dtype)

    def forward(self, hidden_states: torch.Tensor):
        bs, sl, hs = hidden_states.shape
        flat = hidden_states.view(-1, hs)

        if torch.isnan(flat).any() or torch.isinf(flat).any():
            logger.error("MoE: NaN/Inf pada input. Mengembalikan input tanpa perubahan.")
            zero = torch.tensor(0.0, device=flat.device)
            return hidden_states, zero, zero

        if self.use_expert_choice:
            weights, aux_loss, z_loss = self.router(hidden_states)
            out = torch.zeros_like(flat)
            for i, expert in enumerate(self.experts):
                mask = weights[:, i] > 1e-6
                if mask.any():
                    if not self.training:
                        self.expert_usage_count[i] += mask.sum()
                    expert_out = self._safe_expert(expert, flat[mask], i, self.config)
                    out[mask] += expert_out * weights[mask, i:i+1]
        else:
            top_w, top_idx, aux_loss, z_loss = self.router(hidden_states)
            out = torch.zeros_like(flat)
            for i in range(self.num_experts):
                mask = (top_idx == i).any(-1)
                if not mask.any():
                    continue
                if not self.training:
                    self.expert_usage_count[i] += mask.sum()
                expert_out = self._safe_expert(self.experts[i], flat[mask], i, self.config)
                token_ids  = torch.where(mask)[0]
                for j, tid in enumerate(token_ids):
                    positions = (top_idx[tid] == i).nonzero(as_tuple=True)[0]
                    if positions.numel():
                        out[tid] += expert_out[j] * top_w[tid, positions[0]]

        if self.load_monitoring and self.training and isinstance(self.router, TopKRouter):
            imbalance = self.router.get_load_imbalance()
            if imbalance > self.config.expert_load_warn_threshold:
                logger.warning(
                    f"MoE load imbalance CV={imbalance:.3f} "
                    f"(threshold={self.config.expert_load_warn_threshold})"
                )

        final = out.view(bs, sl, hs)
        if torch.isnan(final).any() or torch.isinf(final).any():
            logger.error("MoE: NaN/Inf pada output. Mengembalikan input.")
            return hidden_states, aux_loss, z_loss
        return final, aux_loss, z_loss

# --- MLP standar ---
class CacaMLP(nn.Module):
    def __init__(self, config: CacaConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias)
        self.dropout   = nn.Dropout(config.hidden_dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(self.dropout(F.silu(self.gate_proj(x)) * self.up_proj(x)))

# --- Attention ---
class CacaAttention(nn.Module):
    def __init__(self, config: CacaConfig, layer_idx: Optional[int] = None):
        super().__init__()
        self.config      = config
        self.layer_idx   = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads   = config.num_attention_heads
        self.num_kv_heads = config.num_key_value_heads
        self.head_dim    = config.head_dim
        self.num_kv_groups = self.num_heads // self.num_kv_heads
        self.sliding_window         = config.sliding_window
        self.attn_logit_softcapping = config.attn_logit_softcapping
        self.temperature = config.attention_temperature

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads    * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(self.num_heads    * self.head_dim, self.hidden_size,  bias=config.attention_bias)

        if config.use_qk_norm:
            self.q_norm = CacaRMSNorm(self.head_dim, config.qk_norm_eps)
            self.k_norm = CacaRMSNorm(self.head_dim, config.qk_norm_eps)
        else:
            self.q_norm = self.k_norm = None

        if config.use_rotary_embeddings:
            sf, st = 1.0, None
            if config.rope_scaling:
                st = config.rope_scaling.get("type", "linear")
                sf = config.rope_scaling.get("factor", 1.0)
            self.rotary_emb = CacaRotaryEmbedding(
                self.head_dim, config.max_position_embeddings, config.rope_theta,
                scaling_factor=sf, scaling_type=st,
                rope_type=config.rope_type, ntk_alpha=config.rope_ntk_alpha,
            )
        else:
            self.rotary_emb = None

        self.alibi     = ALiBiPositionalBias(self.num_heads) if config.use_alibi else None
        self.attn_drop = nn.Dropout(config.attention_dropout)

        self.is_global = self._global_layer(config, layer_idx)
        self.has_flash = HAS_FLASH_ATTN and config.use_flash_attn
        self.has_xfmrs = HAS_XFORMERS
        self.has_sdpa  = HAS_SDPA

        self._mask_cache: OrderedDict = OrderedDict()
        self._max_cache  = 10
        self._cache_hits = self._cache_misses = 0

    def _global_layer(self, cfg: CacaConfig, idx: Optional[int]) -> bool:
        if idx is None: return False
        if cfg.attention_pattern == "all_global":  return True
        if cfg.attention_pattern == "all_local":   return False
        if cfg.attention_pattern == "interleaved":
            return (idx % cfg.global_attention_every_n_layers) == cfg.global_attention_every_n_layers - 1
        return False

    def _causal_mask(self, ql, kl, dtype, device, use_sw):
        key = (ql, kl, str(dtype), device.type, use_sw, self.sliding_window if use_sw else None)
        if key in self._mask_cache:
            self._cache_hits += 1
            self._mask_cache.move_to_end(key)
            return self._mask_cache[key].to(dtype=dtype, device=device)
        self._cache_misses += 1

        kl_eff = max(kl, ql)
        qp = torch.arange(ql, device=device) + (kl_eff - ql)
        kp = torch.arange(kl_eff, device=device)
        d  = qp[:, None] - kp[None, :]
        mask = d < 0

        if use_sw and self.sliding_window:
            if self.config.use_attention_sink and self.config.attention_sink_size > 0:
                sink = kp[None, :] < self.config.attention_sink_size
                win  = (d >= 0) & (d <= self.sliding_window)
                mask = (d < 0) | (~sink & ~win)
            else:
                mask = mask | (d > self.sliding_window)

        float_mask = torch.zeros(1, 1, ql, kl_eff, dtype=dtype, device=device)
        float_mask.masked_fill_(mask[None, None], -1e9)

        if len(self._mask_cache) >= self._max_cache:
            self._mask_cache.popitem(last=False)
        self._mask_cache[key] = float_mask.detach().cpu()
        return float_mask

    def get_cache_stats(self) -> Dict:
        total = self._cache_hits + self._cache_misses
        return {
            "hits":       self._cache_hits,
            "misses":     self._cache_misses,
            "hit_rate":   self._cache_hits / total if total else 0,
            "cache_size": len(self._mask_cache),
        }

    def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False):
        if hidden_states is None:
            raise ValueError("hidden_states tidak boleh None")
        if hidden_states.shape[-1] != self.hidden_size:
            raise ValueError(f"Ekspektasi hidden_size {self.hidden_size}, dapat {hidden_states.shape[-1]}")

        bs, sl, _ = hidden_states.shape

        q = self.q_proj(hidden_states).view(bs, sl, self.num_heads,    self.head_dim).transpose(1, 2)
        k = self.k_proj(hidden_states).view(bs, sl, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(hidden_states).view(bs, sl, self.num_kv_heads, self.head_dim).transpose(1, 2)

        if self.q_norm:
            q, k = self.q_norm(q), self.k_norm(k)

        pos_off = 0
        if past_key_value is not None:
            try:
                if isinstance(past_key_value, (tuple, list)) and past_key_value[0] is not None:
                    pos_off = past_key_value[0].shape[2]
            except Exception:
                pos_off = 0

        if self.rotary_emb:
            cos, sin = self.rotary_emb(q, sl, pos_off)
            q, k = apply_rotary_pos_emb(q, k, cos, sin)

        if past_key_value is not None and past_key_value[0] is not None:
            try:
                if past_key_value[0].numel() > 0:
                    k = torch.cat([past_key_value[0], k], dim=2)
                    v = torch.cat([past_key_value[1], v], dim=2)
            except RuntimeError as e:
                logger.error(f"KV concat gagal: {e}")

        present = (k, v) if use_cache else None

        k = k.repeat_interleave(self.num_kv_groups, dim=1)
        v = v.repeat_interleave(self.num_kv_groups, dim=1)

        kv_len = k.shape[-2]
        use_sw = not self.is_global and self.sliding_window is not None

        if self.has_flash and attention_mask is None and q.device.type == "cuda" and q.dtype in (torch.float16, torch.bfloat16):
            try:
                out = self._flash_attn(q, k, v, use_sw, kv_len)
            except Exception as e:
                logger.warning(f"Flash attention gagal ({e}), fallback.")
                out = self._fallback(q, k, v, attention_mask, kv_len, use_sw)
        else:
            out = self._fallback(q, k, v, attention_mask, kv_len, use_sw)

        return self.o_proj(out), present

    def _flash_attn(self, q, k, v, use_sw, kv_len):
        bs, _, sl, _ = q.shape
        orig_dtype    = q.dtype
        compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
        q = q.transpose(1, 2).contiguous().to(compute_dtype)
        k = k.transpose(1, 2).contiguous().to(compute_dtype)
        v = v.transpose(1, 2).contiguous().to(compute_dtype)
        ws  = (self.sliding_window, 0) if use_sw and self.sliding_window < kv_len else (-1, 0)
        out = flash_attn_func(
            q, k, v,
            dropout_p=self.config.attention_dropout if self.training else 0.0,
            causal=True, window_size=ws,
        )
        return out.to(orig_dtype).reshape(bs, sl, self.hidden_size)

    def _fallback(self, q, k, v, mask, kv_len, use_sw):
        if self.has_xfmrs and q.device.type == "cuda" and mask is None:
            try:
                return self._xformers(q, k, v, kv_len, use_sw)
            except Exception:
                pass
        if self.has_sdpa:
            return self._sdpa(q, k, v, mask, kv_len, use_sw)
        return self._standard(q, k, v, mask, kv_len, use_sw)

    def _sdpa(self, q, k, v, mask, kv_len, use_sw):
        bs, _, sl, _ = q.shape
        if mask is None:
            mask = self._causal_mask(sl, kv_len, q.dtype, q.device, use_sw)
        if self.alibi:
            mask = mask + self.alibi(sl, kv_len)
        # scale = 1/sqrt(d) lalu dibagi temperature, konsisten dengan standar attention
        scale = (1.0 / math.sqrt(self.head_dim)) / self.temperature
        out = F.scaled_dot_product_attention(
            q, k, v, attn_mask=mask,
            dropout_p=self.config.attention_dropout if self.training else 0.0,
            scale=scale, is_causal=False,
        )
        return out.transpose(1, 2).contiguous().reshape(bs, sl, self.hidden_size)

    def _xformers(self, q, k, v, kv_len, use_sw):
        bs, _, sl, _ = q.shape
        bias = self._causal_mask(sl, kv_len, q.dtype, q.device, use_sw)
        out  = memory_efficient_attention(
            q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2),
            attn_bias=bias,
            p=self.config.attention_dropout if self.training else 0.0,
        )
        return out.reshape(bs, sl, self.hidden_size)

    def _standard(self, q, k, v, mask, kv_len, use_sw):
        bs, _, sl, _ = q.shape

        scale = math.sqrt(self.head_dim) * self.temperature
        w = torch.matmul(q, k.transpose(-2, -1)) / scale
        w = torch.clamp(w, -50, 50)
        w = soft_cap_logits(w, self.attn_logit_softcapping)
        if mask is None:
            mask = self._causal_mask(sl, kv_len, w.dtype, w.device, use_sw)
        if self.alibi:
            mask = mask + self.alibi(sl, kv_len)
        w = F.softmax(w + mask, -1, dtype=torch.float32).to(q.dtype)
        w = self.attn_drop(w)
        out = torch.matmul(w, v)
        return out.transpose(1, 2).contiguous().reshape(bs, sl, self.hidden_size)

# --- Cross-Attention ---
class CacaCrossAttention(nn.Module):
    def __init__(self, config: CacaConfig):
        super().__init__()
        self.hidden_size   = config.hidden_size
        self.num_heads     = config.num_attention_heads
        self.num_kv_heads  = config.num_key_value_heads
        self.head_dim      = config.head_dim
        self.num_kv_groups = self.num_heads // self.num_kv_heads

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads    * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(self.num_heads    * self.head_dim, self.hidden_size,  bias=config.attention_bias)
        self.drop   = nn.Dropout(config.attention_dropout)

    def forward(self, hidden_states, encoder_hidden_states, attention_mask=None):
        bs, sl, _ = hidden_states.shape
        esl = encoder_hidden_states.shape[1]

        q = self.q_proj(hidden_states).view(bs, sl,  self.num_heads,    self.head_dim).transpose(1, 2)
        k = self.k_proj(encoder_hidden_states).view(bs, esl, self.num_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(encoder_hidden_states).view(bs, esl, self.num_kv_heads, self.head_dim).transpose(1, 2)
        k = k.repeat_interleave(self.num_kv_groups, dim=1)
        v = v.repeat_interleave(self.num_kv_groups, dim=1)

        w = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        if attention_mask is not None:
            w = w + attention_mask
        w   = F.softmax(w, -1, dtype=torch.float32).to(q.dtype)
        w   = self.drop(w)
        out = torch.matmul(w, v).transpose(1, 2).contiguous().reshape(bs, sl, self.hidden_size)
        return self.o_proj(out)

# --- MoD Router ---
class MixtureOfDepthsRouter(nn.Module):
    def __init__(self, hidden_size: int, capacity_factor: float = 0.5, route_method: str = "learned"):
        super().__init__()
        self.capacity_factor = capacity_factor
        self.route_method    = route_method
        if route_method == "learned":
            self.router = nn.Linear(hidden_size, 1)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        bs, sl, _ = hidden_states.shape
        if self.route_method == "learned":
            logits = self.router(hidden_states).squeeze(-1)
        elif self.route_method == "random":
            logits = torch.rand(bs, sl, device=hidden_states.device)
        else:
            logits = torch.zeros(bs, sl, device=hidden_states.device)
        cap = max(1, int(sl * self.capacity_factor))
        _, idx = torch.topk(logits, cap, dim=-1)
        mask = torch.zeros(bs, sl, dtype=torch.bool, device=hidden_states.device)
        mask.scatter_(1, idx, True)
        return mask

# --- Decoder Layer ---
class CacaDecoderLayer(nn.Module):
    def __init__(self, config: CacaConfig, layer_idx: int):
        super().__init__()
        self.config    = config
        self.layer_idx = layer_idx

        self.self_attn = CacaAttention(config, layer_idx)
        self.use_moe   = config.use_moe and (layer_idx % config.moe_layer_frequency == 0)
        self.mlp       = MixtureOfExperts(config) if self.use_moe else CacaMLP(config)
        self.use_cross = config.use_cross_attention and (layer_idx % config.cross_attention_frequency == 0)

        if self.use_cross:
            self.cross_attn      = CacaCrossAttention(config)
            self.cross_attn_norm = CacaRMSNorm(config.hidden_size, config.rms_norm_eps)

        self.input_layernorm          = CacaRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.post_attention_layernorm = CacaRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.residual_drop            = nn.Dropout(config.residual_dropout)

        self.ls1    = LayerScale(config.hidden_size, config.layer_scale_init) if config.use_layer_scale else None
        self.ls2    = LayerScale(config.hidden_size, config.layer_scale_init) if config.use_layer_scale else None
        self.ls_cross = LayerScale(config.hidden_size, config.layer_scale_init) if (config.use_layer_scale and self.use_cross) else None

        if config.use_stochastic_depth:
            dp = config.stochastic_depth_prob * layer_idx / config.num_hidden_layers
            self.stoch_depth = StochasticDepth(dp)
        else:
            self.stoch_depth = None

        self.mod_router = (
            MixtureOfDepthsRouter(config.hidden_size, config.mod_capacity_factor, config.mod_route_method)
            if config.use_mixture_of_depths else None
        )

        self._grad_stats = {"max_norm": 0.0, "ema_norm": 0.0, "clip_count": 0}

    def _grad_hook(self, grad: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
        if grad is None:
            return grad
        n = grad.norm().item()
        self._grad_stats["max_norm"] = max(self._grad_stats["max_norm"], n)
        self._grad_stats["ema_norm"] = 0.9 * self._grad_stats["ema_norm"] + 0.1 * n
        if n > 10.0:
            self._grad_stats["clip_count"] += 1
            if self._grad_stats["clip_count"] % 100 == 0:
                logger.warning(
                    f"Layer {self.layer_idx}: grad norm {n:.2f} "
                    f"(clipped {self._grad_stats['clip_count']}x)"
                )
            return torch.clamp(grad, -10.0, 10.0)
        return grad

    def _recover_nan(self, x: torch.Tensor, name: str) -> torch.Tensor:
        level = self.config.nan_recovery_level
        if level == 0:
            return x
        if not (torch.isnan(x).any() or torch.isinf(x).any()):
            return x
        logger.warning(f"Layer {self.layer_idx}: NaN/Inf pada {name}")
        if level >= 1:
            return torch.nan_to_num(x, nan=0.0, posinf=1e4, neginf=-1e4)
        return x

    def forward(
        self, hidden_states, attention_mask=None,
        encoder_hidden_states=None, encoder_attention_mask=None,
        past_key_value=None, use_cache=False,
    ):
        if hidden_states is None:
            raise ValueError("hidden_states tidak boleh None")

        hidden_states = self._recover_nan(hidden_states, "layer_input")

        if self.training and hidden_states.requires_grad:
            hidden_states.register_hook(self._grad_hook)

        aux_loss = z_loss = 0.0

        if self.mod_router is not None:
            proc_mask = self.mod_router(hidden_states)
            tokens    = hidden_states[proc_mask]
            if tokens.numel() == 0:
                return hidden_states, past_key_value if use_cache else None, aux_loss, z_loss
        else:
            proc_mask = None
            tokens    = hidden_states

        # Self-attention
        residual = tokens
        attn_out, present_kv = self.self_attn(
            self.input_layernorm(tokens), attention_mask, past_key_value, use_cache
        )
        attn_out = self._recover_nan(attn_out, "attn_output")
        if self.ls1:         attn_out = self.ls1(attn_out)
        if self.stoch_depth: attn_out = self.stoch_depth(attn_out, self.training)
        tokens = residual + self.residual_drop(attn_out)
        if self.training: tokens = torch.clamp(tokens, -1e4, 1e4)

        # Cross-attention
        if self.use_cross and encoder_hidden_states is not None:
            residual  = tokens
            cross_out = self.cross_attn(self.cross_attn_norm(tokens), encoder_hidden_states, encoder_attention_mask)
            cross_out = self._recover_nan(cross_out, "cross_attn")
            if self.ls_cross:    cross_out = self.ls_cross(cross_out)
            if self.stoch_depth: cross_out = self.stoch_depth(cross_out, self.training)
            tokens = residual + self.residual_drop(cross_out)
            if self.training: tokens = torch.clamp(tokens, -1e4, 1e4)

        # MLP / MoE
        residual = tokens
        if self.use_moe:
            mlp_out, aux_loss, z_loss = self.mlp(self.post_attention_layernorm(tokens))
        else:
            mlp_out = self.mlp(self.post_attention_layernorm(tokens))
        mlp_out = self._recover_nan(mlp_out, "mlp_output")
        if self.ls2:         mlp_out = self.ls2(mlp_out)
        if self.stoch_depth: mlp_out = self.stoch_depth(mlp_out, self.training)
        tokens = residual + self.residual_drop(mlp_out)
        if self.training: tokens = torch.clamp(tokens, -1e4, 1e4)

        if proc_mask is not None:
            hidden_states[proc_mask] = tokens
        else:
            hidden_states = tokens

        return hidden_states, present_kv, aux_loss, z_loss

    def get_gradient_stats(self) -> Dict:
        return self._grad_stats.copy()

# --- PreTrainedModel base ---
class CacaPreTrainedModel(PreTrainedModel):
    config_class = CacaConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["CacaDecoderLayer"]
   # _skip_keys_device_placement = "past_key_values"

    def _init_weights(self, module):
        std = self.config.initializer_range

        if self.config.use_mup and isinstance(module, nn.Linear):
            for name, m in self.named_modules():
                if m is module and "o_proj" in name:
                    std = self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers)
                    break

        if isinstance(module, nn.Linear):
            module.weight.data.normal_(0.0, std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(0.0, std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

    def _set_gradient_checkpointing(self, module, value=False):
        if isinstance(module, CacaModel):
            module.gradient_checkpointing = value

    def apply_spectral_norm(self):
        if not self.config.use_spectral_norm:
            return
        for name, module in self.named_modules():
            if isinstance(module, nn.Linear):
                try:
                    nn.utils.spectral_norm(module)
                except Exception as e:
                    logger.warning(f"SpectralNorm gagal pada {name}: {e}")
        logger.info("Spectral norm diterapkan ke semua Linear layer.")

# --- Vision Encoder ---
class VisionTransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_ratio=4.0, dropout=0.0, eps=1e-6):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim, eps=eps)
        self.attn  = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)
        self.drop1 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(dim, eps=eps)
        mid = int(dim * mlp_ratio)
        self.mlp   = nn.Sequential(
            nn.Linear(dim, mid), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(mid, dim), nn.Dropout(dropout),
        )
        self.drop2 = nn.Dropout(dropout)

    def forward(self, x):

        nx = self.norm1(x)
        x  = x + self.drop1(self.attn(nx, nx, nx, need_weights=False)[0])
        return x + self.drop2(self.mlp(self.norm2(x)))

class VisionEncoder(nn.Module):
    def __init__(self, config: CacaConfig):
        super().__init__()
        vc = config.vision_config
        self.patch_size   = vc.get("patch_size", 14)
        self.image_size   = vc.get("image_size", 224)
        self.num_channels = vc.get("num_channels", 3)
        self.hidden_size  = vc.get("hidden_size", 1024)
        self.num_patches  = (self.image_size // self.patch_size) ** 2
        self.patch_embed  = nn.Conv2d(self.num_channels, self.hidden_size, self.patch_size, self.patch_size, bias=False)
        self.cls_token    = nn.Parameter(torch.zeros(1, 1, self.hidden_size))
        self.pos_embed    = nn.Parameter(torch.zeros(1, self.num_patches + 1, self.hidden_size))
        self.drop         = nn.Dropout(vc.get("dropout", 0.0))
        self.blocks       = nn.ModuleList([
            VisionTransformerBlock(
                self.hidden_size, vc.get("num_heads", 16),
                vc.get("intermediate_size", 4096) / self.hidden_size,
                vc.get("dropout", 0.0), vc.get("layer_norm_eps", 1e-6),
            )
            for _ in range(vc.get("num_layers", 24))
        ])
        self.norm = nn.LayerNorm(self.hidden_size, eps=vc.get("layer_norm_eps", 1e-6))
        nn.init.trunc_normal_(self.pos_embed,        std=0.02)
        nn.init.trunc_normal_(self.cls_token,        std=0.02)
        nn.init.trunc_normal_(self.patch_embed.weight, std=0.02)

    def forward(self, pixel_values):
        bs = pixel_values.shape[0]
        x  = self.patch_embed(pixel_values).flatten(2).transpose(1, 2)
        x  = torch.cat([self.cls_token.expand(bs, -1, -1), x], dim=1) + self.pos_embed
        x  = self.drop(x)
        for blk in self.blocks:
            x = blk(x)
        return self.norm(x)

# --- Audio Encoder ---
class AudioTransformerBlock(nn.Module):
    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn  = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)
        self.drop1 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(dim)
        mid = int(dim * mlp_ratio)
        self.mlp   = nn.Sequential(
            nn.Linear(dim, mid), nn.GELU(), nn.Dropout(dropout),
            nn.Linear(mid, dim), nn.Dropout(dropout),
        )
        self.drop2 = nn.Dropout(dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        nx = self.norm1(x)
        x  = x + self.drop1(self.attn(nx, nx, nx, need_weights=False)[0])
        return x + self.drop2(self.mlp(self.norm2(x)))

class AudioEncoder(nn.Module):
    def __init__(self, config: "CacaConfig"):
        super().__init__()
        ac = config.audio_config
        self.num_mel_bins    = ac.get("num_mel_bins", 80)
        self.hidden_size     = ac.get("hidden_size", 1024)
        self.max_audio_length = ac.get("max_audio_length", 3000)
        dropout = ac.get("dropout", 0.0)

        self.conv1 = nn.Sequential(
            nn.Conv1d(self.num_mel_bins, self.hidden_size, kernel_size=3, padding=1),
            nn.GELU(), nn.Dropout(dropout),
        )
        self.conv2 = nn.Sequential(
            nn.Conv1d(self.hidden_size, self.hidden_size, kernel_size=3, stride=2, padding=1),
            nn.GELU(), nn.Dropout(dropout),
        )
        self.pos_embed = nn.Parameter(torch.zeros(1, self.max_audio_length // 2, self.hidden_size))
        self.drop      = nn.Dropout(dropout)
        self.blocks    = nn.ModuleList([
            AudioTransformerBlock(
                self.hidden_size, ac.get("num_heads", 16),
                ac.get("intermediate_size", 4096) / self.hidden_size, dropout,
            )
            for _ in range(ac.get("num_layers", 12))
        ])
        self.norm = nn.LayerNorm(self.hidden_size)
        nn.init.trunc_normal_(self.pos_embed, std=0.02)

    def forward(self, audio_features: torch.Tensor) -> torch.Tensor:
        x = F.gelu(self.conv1(audio_features))
        x = F.gelu(self.conv2(x))
        x = x.transpose(1, 2)
        seq_len = x.shape[1]

        if seq_len <= self.pos_embed.shape[1]:
            x = x + self.pos_embed[:, :seq_len, :]
        else:
            pos_interp = F.interpolate(
                self.pos_embed.transpose(1, 2), size=seq_len,
                mode="linear", align_corners=False,
            ).transpose(1, 2)
            x = x + pos_interp

        x = self.drop(x)
        for blk in self.blocks:
            x = blk(x)
        return self.norm(x)

class MultiModalProjector(nn.Module):
    def __init__(self, in_size: int, out_size: int, proj_type: str = "mlp", num_layers: int = 2):
        super().__init__()
        if proj_type == "linear":
            self.proj = nn.Linear(in_size, out_size)
        else:
            layers, cur = [], in_size
            for _ in range(num_layers - 1):
                layers += [nn.Linear(cur, out_size), nn.GELU(), nn.Dropout(0.1)]
                cur = out_size
            layers.append(nn.Linear(cur, out_size))
            self.proj = nn.Sequential(*layers)

    def forward(self, x):
        return self.proj(x)

# --- CacaModel ---
class CacaModel(CacaPreTrainedModel):
    def __init__(self, config: CacaConfig):
        super().__init__(config)
        self.config       = config
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.token_drop   = nn.Dropout(config.token_dropout) if config.token_dropout > 0 else None
        self.layers       = nn.ModuleList([CacaDecoderLayer(config, i) for i in range(config.num_hidden_layers)])
        self.norm         = CacaRMSNorm(config.hidden_size, config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.metrics      = MetricsTracker()
        self._last_mem_check = 0

        if config.use_multimodal:
            if config.vision_config:
                self.vision_encoder   = VisionEncoder(config)
                self.vision_projector = MultiModalProjector(
                    config.vision_config.get("hidden_size", 768), config.hidden_size,
                    config.vision_config.get("projector_type", "mlp"),
                )
            else:
                self.vision_encoder = self.vision_projector = None

        if config.audio_config:
            self.audio_encoder   = AudioEncoder(config)
            self.audio_projector = MultiModalProjector(
                config.audio_config.get("hidden_size", 768), config.hidden_size,
                config.audio_config.get("projector_type", "mlp"),
            )
        else:
            self.audio_encoder = self.audio_projector = None

        if config.use_spectral_norm:
            self.apply_spectral_norm()

        self.post_init()

    def get_input_embeddings(self): return self.embed_tokens
    def set_input_embeddings(self, v): self.embed_tokens = v

    def _prep_mask(self, mask, shape, dtype):
        if mask is None: return None
        if mask.dim() == 2:   mask = mask[:, None, None, :]
        elif mask.dim() == 3: mask = mask[:, None]
        return (1.0 - mask.to(dtype)) * torch.finfo(dtype).min

    def _check_memory(self, device, idx):
        if not (torch.cuda.is_available() and device.type == "cuda"): return
        if idx - self._last_mem_check < 5: return
        alloc = torch.cuda.memory_allocated(device) / 1024**3
        self.metrics.log("gpu_mem_gb", alloc)
        if alloc > 12: logger.warning(f"Layer {idx}: GPU mem {alloc:.2f} GB")
        if alloc > 14: torch.cuda.empty_cache()
        self._last_mem_check = idx

    def forward(
        self, input_ids=None, pixel_values=None, audio_features=None,
        attention_mask=None, past_key_values=None, use_cache=None,
        output_hidden_states=False, return_dict=True, **kwargs,
    ):
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        if input_ids is None:
            raise ValueError("input_ids tidak boleh None")

        bs, sl = input_ids.shape
        device = input_ids.device

        self.metrics.log("batch_size", bs)
        self.metrics.log("seq_length", sl)

        hidden_states = self.embed_tokens(input_ids)

        if self.token_drop is not None and self.training:
            hidden_states = self.token_drop(hidden_states)

        enc_hs = enc_mask = None
        if self.config.use_multimodal and pixel_values is not None and self.vision_encoder is not None:
            try:
                vis_feats  = self.vision_encoder(pixel_values.to(device))
                vis_embeds = self.vision_projector(vis_feats)
                vis_len    = vis_embeds.shape[1]
                if self.config.use_cross_attention:
                    enc_hs   = vis_embeds
                    enc_mask = torch.ones(bs, vis_len, dtype=hidden_states.dtype, device=device)
                else:
                    hidden_states = torch.cat([vis_embeds, hidden_states], dim=1)
                    sl = hidden_states.shape[1]

                    if attention_mask is not None:
                        vis_mask = torch.ones(bs, vis_len, dtype=attention_mask.dtype, device=device)
                        attention_mask = torch.cat([vis_mask, attention_mask], dim=1)

            except RuntimeError as e:
                logger.error(f"Vision encoder gagal: {e}")

        if self.config.use_multimodal and audio_features is not None and self.audio_encoder is not None:
            try:
                audio_encoded = self.audio_encoder(audio_features.to(device))
                audio_embeds  = self.audio_projector(audio_encoded)
                audio_len     = audio_embeds.shape[1]

                if self.config.use_cross_attention:
                    enc_hs = audio_embeds if enc_hs is None else torch.cat([enc_hs, audio_embeds], dim=1)
                    audio_mask = torch.ones(bs, audio_len, dtype=hidden_states.dtype, device=device)
                    enc_mask = audio_mask if enc_mask is None else torch.cat([enc_mask, audio_mask], dim=1)
                else:
                    hidden_states = torch.cat([audio_embeds, hidden_states], dim=1)
                    sl = hidden_states.shape[1]
                    if attention_mask is not None:
                        audio_mask_flat = torch.ones(bs, audio_len, dtype=attention_mask.dtype, device=device)
                        attention_mask = torch.cat([audio_mask_flat, attention_mask], dim=1)

            except RuntimeError as e:
                logger.error(f"Audio encoder gagal: {e}")

        if attention_mask is not None:
            attention_mask = self._prep_mask(attention_mask, (bs, sl), hidden_states.dtype)
        if enc_mask is not None and self.config.use_cross_attention:
            enc_mask = self._prep_mask(enc_mask, (bs, enc_hs.shape[1]), hidden_states.dtype)

        if use_cache and past_key_values is None:
            past_key_values = tuple([None] * len(self.layers))

        present_kvs = [] if use_cache else None
        all_hs      = [hidden_states] if output_hidden_states else None
        total_aux   = torch.tensor(0.0, device=device)
        total_z     = torch.tensor(0.0, device=device)

        for idx, layer in enumerate(self.layers):
            if self.training: self._check_memory(device, idx)
            pkv = past_key_values[idx] if past_key_values else None

            if self.gradient_checkpointing and self.training and not use_cache:
                def make_fwd(l):
                    def fwd(hs):
                        out, _, al, zl = l(hs, attention_mask, enc_hs, enc_mask, None, False)
                        return out, al, zl
                    return fwd
                from torch.utils.checkpoint import checkpoint
                hidden_states, al, zl = checkpoint(make_fwd(layer), hidden_states, use_reentrant=False)
                pkv_out = None
            else:
                hidden_states, pkv_out, al, zl = layer(
                    hidden_states, attention_mask, enc_hs, enc_mask, pkv, use_cache
                )

            if use_cache: present_kvs.append(pkv_out)

            if torch.is_tensor(al):
                total_aux = total_aux + al
            else:
                total_aux = total_aux + torch.tensor(float(al), device=device)
            if torch.is_tensor(zl):
                total_z = total_z + zl
            else:
                total_z = total_z + torch.tensor(float(zl), device=device)

            if self.training:
                self.metrics.log(f"l{idx}_aux", al.item() if torch.is_tensor(al) else al)
                self.metrics.log(f"l{idx}_z",   zl.item() if torch.is_tensor(zl) else zl)
            if output_hidden_states: all_hs.append(hidden_states)

        hidden_states = self.norm(hidden_states)
        if output_hidden_states: all_hs.append(hidden_states)
        self.metrics.step()

        if not return_dict:
            non_none = [hidden_states]
            if use_cache:    non_none.append(tuple(present_kvs))
            if all_hs:       non_none.append(all_hs)
            non_none.append(total_aux)
            non_none.append(total_z)
            return tuple(non_none)

        return (
            BaseModelOutputWithPast(
                last_hidden_state=hidden_states,
                past_key_values=tuple(present_kvs) if use_cache else None,
                hidden_states=all_hs,
                attentions=None,
            ),
            total_aux, total_z,
        )

    def get_metrics_summary(self): return self.metrics.get_summary()
    def get_attention_cache_stats(self):
        return {f"layer_{i}": l.self_attn.get_cache_stats() for i, l in enumerate(self.layers)}

# --- CacaForCausalLM ---
class CacaForCausalLM(CacaPreTrainedModel, GenerationMixin):
   # _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: CacaConfig):
        super().__init__(config)
        self.model   = CacaModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):     return self.model.embed_tokens
    def set_input_embeddings(self, v):  self.model.embed_tokens = v
    def get_output_embeddings(self):    return self.lm_head
    def set_output_embeddings(self, v): self.lm_head = v
    def set_decoder(self, d):           self.model = d
    def get_decoder(self):              return self.model

    def forward(
        self,
        input_ids=None, pixel_values=None, audio_features=None,
        attention_mask=None, labels=None, past_key_values=None,
        inputs_embeds=None, use_cache=None, output_attentions=None,
        output_hidden_states=None, return_dict=None, **kwargs,
    ):
        
        if input_ids is not None:
            if input_ids.dtype.is_floating_point:
                raise TypeError(
                    f"input_ids harus integer dtype, dapat {input_ids.dtype}. "
                    f"Gunakan input_ids.long() untuk konversi."
                )

            if (input_ids < 0).any():
                neg_count = (input_ids < 0).sum().item()
                raise ValueError(
                    f"input_ids memiliki {neg_count} nilai negatif. "
                    f"Semua token ID harus >= 0."
                )

            max_id = input_ids.max().item()
            if max_id >= self.config.vocab_size:
                raise ValueError(
                    f"input_ids memiliki ID {max_id} yang >= vocab_size "
                    f"({self.config.vocab_size}). Rentang valid [0, {self.config.vocab_size - 1}]."
                )

            if self.config.pad_token_id is not None:
                pad_only = (input_ids == self.config.pad_token_id).all(dim=1)
                if pad_only.any():
                    logger.warning(
                        f"{pad_only.sum().item()} sequences hanya berisi padding token. "
                        f"Ini bisa menyebabkan perilaku tak terduga."
                    )

            if labels is not None:
                if labels.dtype not in (torch.long, torch.int, torch.int32, torch.int64):
                    raise TypeError(f"labels harus integer dtype, dapat {labels.dtype}")

                if labels.shape != input_ids.shape:
                    raise ValueError(
                        f"labels shape {labels.shape} != input_ids shape {input_ids.shape}"
                    )

                valid_labels = labels[labels != -100]
                if valid_labels.numel() > 0:
                    invalid_negs = valid_labels[valid_labels < 0]
                    if invalid_negs.numel() > 0:
                        raise ValueError(
                            f"labels memiliki {invalid_negs.numel()} nilai negatif tak valid (bukan -100)"
                        )
                    max_label = valid_labels.max().item()
                    if max_label >= self.config.vocab_size:
                        raise ValueError(
                            f"labels memiliki ID {max_label} yang >= vocab_size "
                            f"({self.config.vocab_size})"
                        )
                else:
                    logger.warning("Semua label adalah -100 (diabaikan). Loss tidak akan dihitung.")

            if attention_mask is not None:
                if attention_mask.shape[0] != input_ids.shape[0]:
                    raise ValueError(
                        f"attention_mask batch size ({attention_mask.shape[0]}) != "
                        f"input_ids batch size ({input_ids.shape[0]})"
                    )
                if attention_mask.shape[1] != input_ids.shape[1]:
                    raise ValueError(
                        f"attention_mask seq length ({attention_mask.shape[1]}) != "
                        f"input_ids seq length ({input_ids.shape[1]})"
                    )
                unique_vals = attention_mask.unique()
                if not torch.all((unique_vals == 0) | (unique_vals == 1)):
                    logger.warning(
                        f"attention_mask memiliki nilai selain 0 dan 1: {unique_vals.tolist()}. "
                        f"Ini bisa menyebabkan perilaku tak terduga."
                    )

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs, aux_loss, z_loss = self.model(
            input_ids, pixel_values=pixel_values, audio_features=audio_features,
            attention_mask=attention_mask, past_key_values=past_key_values,
            use_cache=use_cache, output_hidden_states=output_hidden_states,
            return_dict=True,  # selalu True agar unpack konsisten
        )

        hidden_states = outputs.last_hidden_state
        logits = self.lm_head(hidden_states)

        cap = self.config.lm_logit_softcapping or self.config.final_logit_softcapping
        if cap:
            logits = soft_cap_logits(logits, cap)

        loss = None
        if labels is not None:
            sl_logits = logits[..., :-1, :].contiguous()
            tl        = labels[..., 1:].contiguous()
            loss_fct  = nn.CrossEntropyLoss(ignore_index=-100, label_smoothing=self.config.label_smoothing)
            lm_loss   = loss_fct(sl_logits.view(-1, sl_logits.size(-1)), tl.view(-1))

            if self.training:
                with torch.no_grad():
                    self.model.metrics.log("perplexity", torch.exp(lm_loss).item())
                    self.model.metrics.log("lm_loss",    lm_loss.item())

            if self.config.use_moe:
                _aux = aux_loss.item() if torch.is_tensor(aux_loss) else float(aux_loss)
                _z   = z_loss.item()   if torch.is_tensor(z_loss)   else float(z_loss)
                loss = lm_loss + self.config.router_aux_loss_coef * aux_loss + self.config.router_z_loss_coef * z_loss
                if self.training:
                    self.model.metrics.log("aux_loss", _aux)
                    self.model.metrics.log("z_loss",   _z)
            else:
                loss = lm_loss

        if not return_dict:
            return ((loss, logits) if loss is not None else (logits,))

        return CausalLMOutputWithPast(
            loss=loss, logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=None,
        )

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None,
        attention_mask=None, inputs_embeds=None,
        pixel_values=None, audio_features=None, **kwargs,
    ):
        has_past = False
        if past_key_values:
            try:
                has_past = (
                    len(past_key_values) > 0
                    and past_key_values[0] is not None
                    and isinstance(past_key_values[0], (tuple, list))
                    and past_key_values[0][0] is not None
                    and past_key_values[0][0].numel() > 0
                )
            except Exception:
                pass

        if has_past:
            input_ids      = input_ids[:, -1:]
            pixel_values   = None
            audio_features = None

        model_inputs = (
            {"input_ids": input_ids}
            if inputs_embeds is None or has_past
            else {"inputs_embeds": inputs_embeds}
        )
        model_inputs.update({
            "past_key_values": past_key_values if has_past else None,
            "use_cache":       kwargs.get("use_cache"),
            "attention_mask":  attention_mask,
            "pixel_values":    pixel_values,
            "audio_features":  audio_features,
        })
        return model_inputs

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        reordered = ()
        for lp in past_key_values:
            if lp is not None:
                reordered += (tuple(
                    ps.index_select(0, beam_idx.to(ps.device))
                    if ps is not None and ps.numel() > 0 else None
                    for ps in lp
                ),)
            else:
                reordered += (None,)
        return reordered

    def save_pretrained(self, save_directory, **kwargs):
        quant_backup = getattr(self.config, "quantization_config", None)
        had_quant    = hasattr(self.config, "quantization_config")
        if had_quant and quant_backup is None:
            delattr(self.config, "quantization_config")
        try:
            super().save_pretrained(save_directory, **kwargs)
            if self.training:
                stats_path = os.path.join(save_directory, "training_stats.json")
                with open(stats_path, "w") as f:
                    json.dump({
                        "metrics":     self.model.get_metrics_summary(),
                        "cache_stats": self.model.get_attention_cache_stats(),
                    }, f, indent=2)
                logger.info(f"Training stats disimpan → {stats_path}")
        finally:
            if had_quant:
                self.config.quantization_config = quant_backup

    def get_model_stats(self) -> Dict:
        stats = {
            "metrics":        self.model.get_metrics_summary(),
            "cache_stats":    self.model.get_attention_cache_stats(),
            "gradient_stats": {f"layer_{i}": l.get_gradient_stats() for i, l in enumerate(self.model.layers)},
        }
        if self.config.use_moe:
            stats["expert_usage"] = {
                f"layer_{i}": l.mlp.expert_usage_count.cpu().tolist()
                for i, l in enumerate(self.model.layers)
                if hasattr(l.mlp, "expert_usage_count")
            }
            stats["expert_load_imbalance"] = {
                f"layer_{i}": l.mlp.router.get_load_imbalance()
                for i, l in enumerate(self.model.layers)
                if hasattr(l.mlp, "router") and hasattr(l.mlp.router, "get_load_imbalance")
            }
        return stats

    def apply_lora(self) -> "CacaForCausalLM":
        if not self.config.use_lora:
            logger.warning("apply_lora() dipanggil tapi config.use_lora=False. Dilewati.")
            return self
        if not HAS_PEFT:
            raise ImportError("peft diperlukan untuk LoRA. Jalankan: pip install peft")
        lora_cfg = LoraConfig(
            task_type=TaskType.CAUSAL_LM,
            r=self.config.lora_rank,
            lora_alpha=self.config.lora_alpha,
            lora_dropout=self.config.lora_dropout,
            bias=self.config.lora_bias,
            target_modules=self.config.lora_target_modules,
        )
        peft_model = get_peft_model(self, lora_cfg)
        peft_model.print_trainable_parameters()
        return peft_model

    def get_mup_param_groups(self, base_lr: float) -> List[Dict]:
        if not self.config.use_mup:
            return [{"params": self.parameters(), "lr": base_lr}]
        groups: Dict[float, List] = defaultdict(list)
        for name, param in self.named_parameters():
            if not param.requires_grad: continue
            mul = self.config.get_mup_lr_multiplier(name)
            groups[base_lr * mul].append(param)
        return [{"params": ps, "lr": lr} for lr, ps in groups.items()]

# --- Quantized variant ---
class CacaForCausalLMQuantized(CacaForCausalLM):
    def __init__(self, config: CacaConfig, quantization_config: Optional[Dict] = None):
        super().__init__(config)
        self.quantization_config = quantization_config
        if quantization_config:
            self._apply_quantization()

    def _apply_quantization(self):
        if not HAS_BNB:
            raise ImportError("bitsandbytes diperlukan. pip install bitsandbytes")
        if self.quantization_config.get("load_in_8bit"):
            self._quantize(bits=8)
        elif self.quantization_config.get("load_in_4bit"):
            self._quantize(bits=4)

    def _quantize(self, bits: int):
        compute_dtype = torch.float16
        if bits == 4 and self.quantization_config.get("bnb_4bit_compute_dtype"):
            compute_dtype = getattr(torch, self.quantization_config["bnb_4bit_compute_dtype"])
        for name, module in self.named_modules():
            if not isinstance(module, nn.Linear): continue
            has_bias = module.bias is not None
            if bits == 8:
                new = bnb.nn.Linear8bitLt(
                    module.in_features, module.out_features, has_bias,
                    threshold=self.quantization_config.get("llm_int8_threshold", 6.0),
                )
            else:
                new = bnb.nn.Linear4bit(
                    module.in_features, module.out_features, bias=has_bias,
                    compute_dtype=compute_dtype,
                    quant_type=self.quantization_config.get("bnb_4bit_quant_type", "nf4"),
                    use_double_quant=self.quantization_config.get("bnb_4bit_use_double_quant", True),
                )
            new.weight = module.weight
            if has_bias: new.bias = module.bias
            parent = ".".join(name.split(".")[:-1])
            child  = name.split(".")[-1]
            setattr(self.get_submodule(parent) if parent else self, child, new)
        logger.info(f"Quantization {bits}-bit diterapkan.")

    @classmethod
    def from_pretrained_quantized(cls, model_path: str, quantization_config: Dict):
        config = CacaConfig.from_pretrained(model_path)
        model  = cls(config, quantization_config=quantization_config)
        sd     = torch.load(os.path.join(model_path, "pytorch_model.bin"), map_location="cpu")
        model.load_state_dict(sd, strict=False)
        return model

# --- CacaTrainer ---
class CacaTrainer:
    def __init__(
        self, model, optimizer, scheduler=None,
        gradient_accumulation_steps: int = 1,
        max_grad_norm: float = 1.0,
        use_amp: bool = False,
        log_interval: int = 10,
    ):
        self.model       = model
        self.optimizer   = optimizer
        self.scheduler   = scheduler
        self.grad_accum  = gradient_accumulation_steps
        self.max_grad    = max_grad_norm
        self.use_amp     = use_amp
        self.log_interval = log_interval
        self.global_step  = 0
        self.epoch        = 0
        self.train_metrics: Dict[str, List[float]] = defaultdict(list)
        self.scaler = torch.cuda.amp.GradScaler() if use_amp and torch.cuda.is_available() else None

    def _forward(self, batch) -> torch.Tensor:

        if self.use_amp:
            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
            with torch.cuda.amp.autocast(dtype=dtype):
                out = self.model(**batch)
        else:
            out = self.model(**batch)
        loss = out.loss if hasattr(out, "loss") else out[0]
        if loss is None:
            raise ValueError("Loss adalah None. Pastikan 'labels' ada di dalam batch.")
        return loss

    def training_step(self, batch) -> float:
        self.model.train()
        loss = self._forward(batch) / self.grad_accum
        (self.scaler.scale(loss) if self.scaler else loss).backward()
        self.train_metrics["loss"].append(loss.item() * self.grad_accum)
        return loss.item() * self.grad_accum

    def optimizer_step(self) -> float:
        if self.scaler:
            self.scaler.unscale_(self.optimizer)
        gn = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad)
        if self.scaler:
            self.scaler.step(self.optimizer)
            self.scaler.update()
        else:
            self.optimizer.step()
        self.optimizer.zero_grad(set_to_none=True)
        if self.scheduler: self.scheduler.step()
        self.train_metrics["grad_norm"].append(gn.item())
        self.global_step += 1
        return gn.item()

    def train_epoch(self, dataloader) -> float:
        self.model.train()
        epoch_loss, n = 0.0, 0
        for step, batch in enumerate(dataloader):
            epoch_loss += self.training_step(batch)
            n += 1
            if (step + 1) % self.grad_accum == 0:
                gn = self.optimizer_step()
                if self.global_step % self.log_interval == 0:
                    lr = self.optimizer.param_groups[0]["lr"]
                    logger.info(
                        f"Epoch {self.epoch} | Step {self.global_step} | "
                        f"Loss: {epoch_loss/n:.4f} | GradNorm: {gn:.4f} | LR: {lr:.2e}"
                    )
                    if hasattr(self.model, "get_model_stats"):
                        stats = self.model.get_model_stats()
                        ppl   = stats.get("metrics", {}).get("perplexity", {}).get("last")
                        if ppl: logger.info(f"  PPL: {ppl:.2f}")
        self.epoch += 1
        return epoch_loss / n if n > 0 else 0.0

    def get_metrics(self) -> Dict:
        return {
            k: {
                "mean": float(np.mean(v)), "std": float(np.std(v)),
                "min":  float(np.min(v)),  "max": float(np.max(v)),
            }
            for k, v in self.train_metrics.items() if v
        }

print("✓ ARSITEKTUR BERHASIL")

# --- Inisialisasi Model ---
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM

AutoConfig.register("caca", CacaConfig)
AutoModel.register(CacaConfig, CacaModel)
AutoModelForCausalLM.register(CacaConfig, CacaForCausalLM)

# --- Utilitas ---
def calculate_params(config_dict: Dict) -> Dict:
    h    = config_dict["hidden_size"]
    i    = config_dict["intermediate_size"]
    v    = config_dict["vocab_size"]
    l    = config_dict["num_hidden_layers"]
    n_h  = config_dict["num_attention_heads"]
    n_kv = config_dict["num_key_value_heads"]
    hd   = config_dict.get("head_dim", h // n_h)
    tie  = config_dict.get("tie_word_embeddings", False)
    use_moe      = config_dict.get("use_moe", False)
    num_experts  = config_dict.get("num_experts", 1)
    moe_freq     = config_dict.get("moe_layer_frequency", 1)
    use_cross    = config_dict.get("use_cross_attention", False)
    cross_freq   = config_dict.get("cross_attention_frequency", 1)

    embed  = v * h
    q_p    = h * (n_h * hd)
    kv_p   = 2 * h * (n_kv * hd)
    o_p    = (n_h * hd) * h
    attn_l = q_p + kv_p + o_p
    norm_l = 2 * h

    lm_head = 0 if tie else v * h

    total_ffn = 0
    total_attn = 0
    total_norm = 0
    total_cross = 0

    for layer_idx in range(l):
        total_attn += attn_l
        total_norm += norm_l

        is_moe_layer = use_moe and (layer_idx % moe_freq == 0)
        n_experts_this_layer = num_experts if is_moe_layer else 1
        ffn_l = 3 * h * i
        total_ffn += ffn_l * n_experts_this_layer

        if is_moe_layer:
            total_ffn += h * num_experts
            
        if use_cross and (layer_idx % cross_freq == 0):
            cross_l = attn_l + h
            total_cross += cross_l

    total = embed + total_attn + total_ffn + total_norm + total_cross + h + lm_head

    return {
        "total": total,
        "embeddings": embed,
        "embedding_pct": embed / total * 100,
        "lm_head": lm_head,
        "lm_head_pct": lm_head / total * 100,
        "attention_total": total_attn,
        "ffn_total": total_ffn,
        "cross_attn_total": total_cross,
        "attention_per_layer": attn_l,
        "ffn_per_layer_dense": 3 * h * i,
    }

def format_params(n: int) -> str:
    if n >= 1e12: return f"{n/1e12:.2f}T"
    if n >= 1e9:  return f"{n/1e9:.2f}B"
    if n >= 1e6:  return f"{n/1e6:.2f}M"
    return f"{n/1e3:.2f}K"

def count_parameters(model) -> Dict:
    total     = sum(p.numel() for p in model.parameters())
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    return {
        "total": total, "trainable": trainable,
        "total_M": total / 1e6, "trainable_M": trainable / 1e6,
    }

def get_model_size_mb(model) -> Dict:
    p_bytes = sum(p.nelement() * p.element_size() for p in model.parameters())
    b_bytes = sum(b.nelement() * b.element_size() for b in model.buffers())
    return {
        "param_mb":  p_bytes / 1024**2,
        "buffer_mb": b_bytes / 1024**2,
        "total_mb":  (p_bytes + b_bytes) / 1024**2,
    }

def print_model_info(model) -> None:
    sep = "=" * 70
    print(f"\n{sep}\nINFO MODEL\n{sep}")
    p = count_parameters(model)
    print(f"Parameter   : {p['total']:,} total  ({p['total_M']:.2f}M)")
    print(f"              {p['trainable']:,} trainable ({p['trainable_M']:.2f}M)")
    s = get_model_size_mb(model)
    print(f"Ukuran      : {s['total_mb']:.1f} MB  "
          f"(params {s['param_mb']:.1f} MB + buffers {s['buffer_mb']:.1f} MB)")
    if hasattr(model, "config"):
        c = model.config
        print(f"Vocab       : {c.vocab_size:,}")
        print(f"Hidden      : {c.hidden_size}  |  Layers: {c.num_hidden_layers}  |  Heads: {c.num_attention_heads}")
        print(f"FFN         : {c.intermediate_size}")
        if c.use_moe:            print(f"MoE         : {c.num_experts} experts, top-{c.num_experts_per_tok}")
        if c.use_lora:           print(f"LoRA        : rank={c.lora_rank}, alpha={c.lora_alpha}")
        if c.use_mup:            print(f"MuP         : base_width={c.mup_base_width}")
        if c.use_spectral_norm:  print("SpectralNorm: aktif")
        if c.token_dropout > 0:  print(f"TokenDrop   : {c.token_dropout}")
        if c.use_multimodal:     print("Multimodal  : vision encoder aktif")
    print(f"{sep}\n")