Add custom modeling file

caca_transformers.py

@@ -0,0 +1,2005 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional, Tuple, List
+from transformers import PreTrainedModel, PretrainedConfig
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.generation.utils import GenerationMixin
+from collections import OrderedDict
+import logging
+from functools import lru_cache
+
+logger = logging.getLogger(__name__)
+
+try:
+    from flash_attn import flash_attn_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')
+
+# --- config ---
+class CacaConfig(PretrainedConfig):
+    model_type = "caca"
+    
+    def __init__(
+        self,
+        vocab_size=32000,
+        hidden_size=2048,
+        intermediate_size=8192,
+        num_hidden_layers=24,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        head_dim=64,
+        max_position_embeddings=8192,
+        rms_norm_eps=1e-6,
+        qk_norm_eps=1e-6,
+        initializer_range=0.02,
+        use_cache=True,
+        pad_token_id=None,
+        bos_token_id=1,
+        eos_token_id=2,
+        tie_word_embeddings=False,
+        rope_theta=10000.0,
+        rope_scaling=None,
+        use_rotary_embeddings=True,
+        attention_bias=False,
+        attention_dropout=0.0,
+        use_qk_norm=True,
+        use_alibi=False,
+        use_flash_attn=True,
+        use_grouped_query_attention=False,
+        use_multi_query_attention=False,
+        sliding_window=None,
+        use_longformer_attention=False,
+        longformer_attention_window=512,
+        attn_logit_softcapping=None,
+        final_logit_softcapping=None,
+        attention_sink_size=4,
+        attention_sink_window=1024,
+        use_attention_sink=False,
+        attention_pattern="all_global",
+        global_attention_every_n_layers=2,
+        mlp_bias=False,
+        hidden_dropout=0.1,
+        residual_dropout=0.1,
+        use_moe=False,
+        num_experts=8,
+        num_experts_per_tok=2,
+        use_expert_choice=False,
+        expert_choice_k=0.125,
+        router_aux_loss_coef=0.01,
+        router_z_loss_coef=0.001,
+        moe_layer_frequency=2,
+        expert_capacity_factor=1.0,
+        use_grouped_moe=False,
+        num_expert_groups=1,
+        use_layer_scale=False,
+        layer_scale_init=1e-5,
+        use_stochastic_depth=False,
+        stochastic_depth_prob=0.1,
+        use_mixture_of_depths=False,
+        mod_capacity_factor=0.5,
+        mod_route_method="learned",
+        use_cross_attention=False,
+        cross_attention_frequency=4,
+        use_multimodal=False,
+        vision_config=None,
+        audio_config=None,
+        projector_hidden_size=None,
+        use_soft_merging=False,
+        merge_threshold=0.5,
+        pretraining_tp=1,
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+        chat_template=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 or (hidden_size // num_attention_heads if hidden_size and num_attention_heads else None)
+        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_cache = use_cache
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        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.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.use_qk_norm = use_qk_norm
+        self.use_alibi = use_alibi
+        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.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.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.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_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.pretraining_tp = pretraining_tp
+        self.tensor_parallel_size = tensor_parallel_size
+        self.pipeline_parallel_size = pipeline_parallel_size
+
+        if chat_template is None:
+            self.chat_template = (
+                "{% 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 %}"
+            )
+        else:
+            self.chat_template = chat_template
+
+        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):
+        if self.num_attention_heads % self.num_key_value_heads != 0:
+            raise ValueError(
+                f"num_attention_heads ({self.num_attention_heads}) harus habis dibagi "
+                f"num_key_value_heads ({self.num_key_value_heads})"
+            )
+        
+        if self.use_moe and self.num_experts < self.num_experts_per_tok:
+            raise ValueError(
+                f"num_experts ({self.num_experts}) harus >= "
+                f"num_experts_per_tok ({self.num_experts_per_tok})"
+            )
+        
+        if self.hidden_size % self.num_attention_heads != 0:
+            raise ValueError(
+                f"hidden_size ({self.hidden_size}) harus habis dibagi "
+                f"num_attention_heads ({self.num_attention_heads})"
+            )
+        
+        if self.vocab_size <= 0:
+            raise ValueError(f"vocab_size harus > 0, dapat {self.vocab_size}")
+            
+        if self.use_flash_attn and not HAS_FLASH_ATTN:
+            logger.warning(
+                "use_flash_attn=True tapi flash-attn tidak terinstall. "
+                "Akan fallback ke SDPA/standard attention."
+            )
+            
+        if self.sliding_window is not None:
+            if self.sliding_window > self.max_position_embeddings:
+                raise ValueError(
+                    f"sliding_window ({self.sliding_window}) tidak boleh > "
+                    f"max_position_embeddings ({self.max_position_embeddings})"
+                )
+                
+        if self.use_moe:
+            if self.moe_layer_frequency <= 0:
+                raise ValueError(f"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}). "
+                    f"MoE tidak akan digunakan."
+                )
+
+    def to_dict(self):
+        has_quant_config = hasattr(self, 'quantization_config')
+        quantization_config_backup = getattr(self, 'quantization_config', None)
+        
+        if has_quant_config and quantization_config_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"
+            }
+        finally:
+            if has_quant_config:
+                self.quantization_config = quantization_config_backup
+        
+        return output
+
+# --- Arsitektur Model ---
+class CacaRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x):
+        input_dtype = x.dtype
+        x = x.float()
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return (self.weight * x).to(input_dtype)
+        
+class LayerScale(nn.Module):
+    def __init__(self, dim, init_value=1e-5):
+        super().__init__()
+        self.gamma = nn.Parameter(init_value * torch.ones(dim))
+
+    def forward(self, x):
+        return self.gamma * x
+
+class StochasticDepth(nn.Module):
+    def __init__(self, drop_prob=0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x, training=True):
+        if not training or self.drop_prob == 0.0:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+        random_tensor.floor_()
+        return x.div(keep_prob) * random_tensor
+
+class CacaRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        max_position_embeddings=8192,
+        base=10000.0,
+        scaling_factor=1.0,
+        scaling_type=None,
+    ):
+        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
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)
+        )
+        if scaling_type == "linear":
+            inv_freq = inv_freq / scaling_factor
+        elif scaling_type == "dynamic":
+            inv_freq = inv_freq
+        elif scaling_type == "yarn":
+            inv_freq = self._yarn_get_inv_freq(inv_freq)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def _yarn_get_inv_freq(self, inv_freq):
+        if len(inv_freq) == 0:
+            return inv_freq
+        alpha = self.scaling_factor
+        beta_fast = 32
+        beta_slow = 1
+        freq_threshold = 1 / (self.max_position_embeddings * beta_fast)
+        low_freq_mask = inv_freq > freq_threshold
+        high_freq_mask = ~low_freq_mask
+        low_freq = inv_freq[low_freq_mask]
+        high_freq = inv_freq[high_freq_mask]
+        if len(low_freq) > 0:
+            low_freq = low_freq / alpha
+        if len(high_freq) > 0:
+            smooth_factor = (
+                self.max_position_embeddings * beta_slow / high_freq - beta_fast
+            ) / (beta_slow - beta_fast)
+            smooth_factor = torch.clamp(smooth_factor, 0.0, 1.0)
+            high_freq = (1 - smooth_factor) * (
+                high_freq / alpha
+            ) + smooth_factor * high_freq
+        result = torch.zeros_like(inv_freq)
+        result[low_freq_mask] = low_freq
+        result[high_freq_mask] = high_freq
+        return result
+
+    def forward(self, x, seq_len, position_offset=0):
+        t = torch.arange(
+            position_offset, position_offset + seq_len, device=x.device
+        ).type_as(self.inv_freq)
+        if self.scaling_type == "dynamic":
+            if seq_len > self.max_position_embeddings:
+                dynamic_scale = seq_len / self.max_position_embeddings
+                t = t / dynamic_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, :, :]
+        return cos.to(x.dtype), sin.to(x.dtype)
+
+class ALiBiPositionalBias(nn.Module):
+    def __init__(self, num_heads, max_positions=8192):
+        super().__init__()
+        self.num_heads = num_heads
+        self.max_positions = max_positions
+        slopes = torch.tensor(self._get_slopes(num_heads))
+        self.register_buffer("slopes", slopes, persistent=False)
+
+    def _get_slopes(self, n):
+        def get_slopes_power_of_2(n):
+            start = 2 ** (-(2 ** -(math.log2(n) - 3)))
+            ratio = start
+            return [start * (ratio**i) for i in range(n)]
+
+        if math.log2(n).is_integer():
+            return get_slopes_power_of_2(n)
+        else:
+            closest_power_of_2 = 2 ** math.floor(math.log2(n))
+            return (
+                get_slopes_power_of_2(closest_power_of_2)
+                + self._get_slopes(2 * closest_power_of_2)[0::2][
+                    : n - closest_power_of_2
+                ]
+            )
+
+    def forward(self, seq_len, key_len=None):
+        if key_len is None:
+            key_len = seq_len
+        query_pos = torch.arange(seq_len, device=self.slopes.device).unsqueeze(1)
+        key_pos = torch.arange(key_len, device=self.slopes.device).unsqueeze(0)
+        relative_pos = key_pos - query_pos
+        bias = relative_pos.unsqueeze(0) * self.slopes.unsqueeze(1).unsqueeze(2)
+        return bias.unsqueeze(0)
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    cos = cos.to(q.dtype)
+    sin = sin.to(q.dtype)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+def soft_cap_logits(x, cap):
+    if cap is None or cap <= 0:
+        return x
+    return x.clamp(-cap * 0.99, cap * 0.99)
+
+class TopKRouter(nn.Module):
+    def __init__(self, hidden_size, num_experts, num_experts_per_tok):
+        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)
+
+    def forward(self, hidden_states):
+        batch_size, seq_len, hidden_size = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_size)
+        hidden_states = self.gate_norm(hidden_states)
+        router_logits = self.gate(hidden_states)
+        router_logits = torch.clamp(router_logits, min=-10, max=10)
+        routing_weights = F.softmax(router_logits, dim=-1)
+        top_k_weights, top_k_indices = torch.topk(
+            routing_weights, self.num_experts_per_tok, dim=-1
+        )
+        top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-8)
+        router_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+        expert_usage = router_probs.mean(dim=0)
+        mean_usage = expert_usage.mean()
+        aux_loss = ((expert_usage - mean_usage) ** 2).sum() / (mean_usage + 1e-10)
+        router_logits_for_z = router_logits.to(torch.float32)
+        z_loss = torch.logsumexp(router_logits_for_z, dim=-1).mean()
+        return top_k_weights, top_k_indices, aux_loss, z_loss
+
+class ExpertChoiceRouter(nn.Module):
+    def __init__(self, hidden_size, num_experts, expert_choice_k):
+        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):
+        batch_size, seq_len, hidden_size = hidden_states.shape
+        total_tokens = batch_size * seq_len
+        hidden_states_flat = hidden_states.view(-1, hidden_size)
+        router_logits = self.gate(hidden_states_flat)
+        router_probs = F.softmax(router_logits, dim=-1, dtype=torch.float32)
+        router_probs_t = router_probs.t()
+        capacity = max(1, int(self.expert_choice_k * total_tokens / self.num_experts))
+        top_k_values, top_k_indices = torch.topk(
+            router_probs_t, k=min(capacity, total_tokens), dim=-1
+        )
+        expert_mask = torch.zeros(
+            self.num_experts, total_tokens, device=hidden_states.device
+        )
+        for expert_idx in range(self.num_experts):
+            expert_mask[expert_idx, top_k_indices[expert_idx]] = 1.0
+        routing_weights = expert_mask.t() * router_probs
+        aux_loss = (router_probs.mean(dim=0) ** 2).sum() * self.num_experts
+        z_loss = torch.logsumexp(router_logits, dim=-1).mean()
+        return routing_weights, aux_loss, z_loss
+
+class Expert(nn.Module):
+    def __init__(self, config):
+        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):
+        gate = F.silu(self.gate_proj(x))
+        up = self.up_proj(x)
+        hidden = gate * up
+        hidden = self.dropout(hidden)
+        return self.down_proj(hidden)
+
+class MixtureOfExperts(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.num_experts = config.num_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.use_expert_choice = config.use_expert_choice
+        self.experts = nn.ModuleList([Expert(config) for _ in range(self.num_experts)])
+        if self.use_expert_choice:
+            self.router = ExpertChoiceRouter(
+                config.hidden_size, config.num_experts, config.expert_choice_k
+            )
+        else:
+            self.router = TopKRouter(
+                config.hidden_size, config.num_experts, config.num_experts_per_tok
+            )
+
+    def forward(self, hidden_states):
+        batch_size, seq_len, hidden_size = hidden_states.shape
+        hidden_states_flat = hidden_states.view(-1, hidden_size)
+        if self.use_expert_choice:
+            routing_weights, aux_loss, z_loss = self.router(hidden_states)
+            final_output = torch.zeros_like(hidden_states_flat)
+            for expert_idx, expert in enumerate(self.experts):
+                expert_mask = routing_weights[:, expert_idx] > 0
+                if expert_mask.any():
+                    expert_input = hidden_states_flat[expert_mask]
+                    expert_output = expert(expert_input)
+                    final_output[expert_mask] += (
+                        expert_output
+                        * routing_weights[expert_mask, expert_idx : expert_idx + 1]
+                    )
+        else:
+            top_k_weights, top_k_indices, aux_loss, z_loss = self.router(hidden_states)
+            final_output = torch.zeros_like(hidden_states_flat)
+            for i in range(self.num_experts_per_tok):
+                expert_indices = top_k_indices[:, i]
+                expert_weights = top_k_weights[:, i : i + 1]
+                for expert_idx in range(self.num_experts):
+                    expert_mask = expert_indices == expert_idx
+                    if expert_mask.any():
+                        expert_input = hidden_states_flat[expert_mask]
+                        expert_output = self.experts[expert_idx](expert_input)
+                        final_output[expert_mask] += (
+                            expert_output * expert_weights[expert_mask]
+                        )
+        final_output = final_output.view(batch_size, seq_len, hidden_size)
+        return final_output, aux_loss, z_loss
+
+class MixtureOfDepthsRouter(nn.Module):
+    def __init__(self, hidden_size, capacity_factor=0.5, route_method="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):
+        batch_size, seq_len, hidden_size = hidden_states.shape
+        if self.route_method == "learned":
+            routing_logits = self.router(hidden_states).squeeze(-1)
+        elif self.route_method == "random":
+            routing_logits = torch.rand(
+                batch_size, seq_len, device=hidden_states.device
+            )
+        else:
+            routing_logits = torch.zeros(
+                batch_size, seq_len, device=hidden_states.device
+            )
+        capacity = max(1, int(seq_len * self.capacity_factor))
+        _, top_indices = torch.topk(routing_logits, k=capacity, dim=-1)
+        process_mask = torch.zeros(
+            batch_size, seq_len, dtype=torch.bool, device=hidden_states.device
+        )
+        process_mask.scatter_(1, top_indices, True)
+        return process_mask
+
+class CacaAttention(nn.Module):
+    def __init__(self, config, layer_idx=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_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.sliding_window = config.sliding_window
+        self.attn_logit_softcapping = config.attn_logit_softcapping
+        self.attention_sink_size = config.attention_sink_size
+        self.attention_sink_window = config.attention_sink_window
+        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_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_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, eps=config.qk_norm_eps)
+            self.k_norm = CacaRMSNorm(self.head_dim, eps=config.qk_norm_eps)
+        else:
+            self.q_norm = None
+            self.k_norm = None
+        if config.use_rotary_embeddings:
+            scaling_factor = 1.0
+            scaling_type = None
+            if config.rope_scaling is not None:
+                scaling_type = config.rope_scaling.get("type", "linear")
+                scaling_factor = config.rope_scaling.get("factor", 1.0)
+            self.rotary_emb = CacaRotaryEmbedding(
+                self.head_dim,
+                config.max_position_embeddings,
+                config.rope_theta,
+                scaling_factor=scaling_factor,
+                scaling_type=scaling_type,
+            )
+        else:
+            self.rotary_emb = None
+        if config.use_alibi:
+            self.alibi = ALiBiPositionalBias(
+                self.num_heads, config.max_position_embeddings
+            )
+        else:
+            self.alibi = None
+        self.attention_dropout = nn.Dropout(config.attention_dropout)
+        self.is_global_attention = self._determine_attention_type(config, layer_idx)
+        self.has_flash_attn = HAS_FLASH_ATTN and config.use_flash_attn
+        self.has_xformers = HAS_XFORMERS
+        self.has_sdpa = HAS_SDPA
+        self._mask_cache = {}
+        self._max_cache_size = 10
+
+    def _determine_attention_type(self, config, layer_idx):
+        if layer_idx is None:
+            return False
+        if config.attention_pattern == "all_global":
+            return True
+        elif config.attention_pattern == "all_local":
+            return False
+        elif config.attention_pattern == "interleaved":
+            return (layer_idx % config.global_attention_every_n_layers) == (
+                config.global_attention_every_n_layers - 1
+            )
+        return False
+
+    def forward(
+        self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False
+    ):
+        batch_size, seq_length, _ = hidden_states.size()
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(hidden_states)
+        value_states = self.v_proj(hidden_states)
+        query_states = query_states.view(
+            batch_size, seq_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.view(
+            batch_size, seq_length, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.view(
+            batch_size, seq_length, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        if self.q_norm is not None and self.k_norm is not None:
+            query_states = self.q_norm(query_states)
+            key_states = self.k_norm(key_states)
+        
+        position_offset = 0
+        if past_key_value is not None:
+            try:
+                if isinstance(past_key_value, (tuple, list)) and len(past_key_value) >= 2:
+                    if past_key_value[0] is not None:
+                        position_offset = past_key_value[0].shape[2]
+            except (IndexError, AttributeError, TypeError):
+                position_offset = 0
+        
+        if self.rotary_emb is not None:
+            cos, sin = self.rotary_emb(query_states, seq_length, position_offset)
+            query_states, key_states = apply_rotary_pos_emb(
+                query_states, key_states, cos, sin
+            )
+            
+        if past_key_value is not None and past_key_value[0] is not None:
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        
+        if use_cache:
+            present_key_value = (key_states, value_states)
+        else:
+            present_key_value = None
+        
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        kv_seq_len = key_states.shape[-2]
+
+        use_sliding_window = (not self.is_global_attention) and (
+            self.sliding_window is not None
+        )
+        if self.has_flash_attn and attention_mask is None:
+            if query_states.device.type == "cuda" and query_states.dtype in [
+                torch.float16,
+                torch.bfloat16,
+            ]:
+                try:
+                    attn_output = self._flash_attention(
+                        query_states, key_states, value_states, use_sliding_window
+                    )
+                except Exception as e:
+                    logger.warning(f"Flash Attention gagal, pakai fallback: {e}")
+                    attn_output = self._fallback_attention(
+                        query_states,
+                        key_states,
+                        value_states,
+                        attention_mask,
+                        kv_seq_len,
+                        use_sliding_window,
+                    )
+            else:
+                attn_output = self._fallback_attention(
+                    query_states,
+                    key_states,
+                    value_states,
+                    attention_mask,
+                    kv_seq_len,
+                    use_sliding_window,
+                )
+        else:
+            attn_output = self._fallback_attention(
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                kv_seq_len,
+                use_sliding_window,
+            )
+        attn_output = self.o_proj(attn_output)
+        return attn_output, present_key_value
+
+    def _flash_attention(
+        self, query_states, key_states, value_states, use_sliding_window
+    ):
+        batch_size, num_heads, seq_length, head_dim = query_states.shape
+        kv_seq_len = key_states.shape[-2]
+        original_dtype = query_states.dtype
+        if original_dtype == torch.bfloat16:
+            if not torch.cuda.is_bf16_supported():
+                logger.warning("BF16 not supported on this GPU, falling back to FP16")
+                original_dtype = torch.float16
+        compute_dtype = (
+            torch.bfloat16
+            if original_dtype not in [torch.float16, torch.bfloat16]
+            else original_dtype
+        )
+        query_states = query_states.transpose(1, 2).contiguous().to(compute_dtype)
+        key_states = key_states.transpose(1, 2).contiguous().to(compute_dtype)
+        value_states = value_states.transpose(1, 2).contiguous().to(compute_dtype)
+        if use_sliding_window and self.sliding_window < kv_seq_len:
+            window_size = (self.sliding_window, 0)
+        else:
+            window_size = (-1, 0)
+        attn_output = flash_attn_func(
+            query_states,
+            key_states,
+            value_states,
+            dropout_p=self.config.attention_dropout if self.training else 0.0,
+            softmax_scale=None,
+            causal=True,
+            window_size=window_size,
+        )
+        attn_output = attn_output.to(original_dtype)
+        attn_output = attn_output.reshape(batch_size, seq_length, self.hidden_size)
+        return attn_output
+
+    def _fallback_attention(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        kv_seq_len,
+        use_sliding_window,
+    ):
+        device_type = query_states.device.type
+        if self.has_xformers and device_type == "cuda" and attention_mask is None:
+            try:
+                return self._xformers_attention(
+                    query_states,
+                    key_states,
+                    value_states,
+                    kv_seq_len,
+                    use_sliding_window,
+                )
+            except Exception as e:
+                logger.warning(f"xFormers gagal, pakai SDPA: {e}")
+        if self.has_sdpa:
+            return self._sdpa_attention(
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                kv_seq_len,
+                use_sliding_window,
+            )
+        else:
+            return self._standard_attention(
+                query_states,
+                key_states,
+                value_states,
+                attention_mask,
+                kv_seq_len,
+                use_sliding_window,
+            )
+
+    def _create_causal_mask(
+        self, query_length, key_length, dtype, device, use_sliding_window
+    ):
+        cache_key = (
+            query_length,
+            key_length,
+            str(dtype),
+            use_sliding_window,
+            self.sliding_window if use_sliding_window else None,
+        )
+        if cache_key in self._mask_cache:
+            cached_mask = self._mask_cache[cache_key]
+            return cached_mask.to(device, dtype)
+        if query_length > key_length:
+            key_length = query_length
+        query_pos = torch.arange(query_length, device=device) + (
+            key_length - query_length
+        )
+        key_pos = torch.arange(key_length, device=device)
+        distance = query_pos[:, None] - key_pos[None, :]
+        mask = distance < 0
+
+        if use_sliding_window and self.sliding_window is not None:
+            if self.config.use_attention_sink and self.attention_sink_size > 0:
+                is_sink = key_pos[None, :] < self.attention_sink_size
+                in_window = (distance >= 0) & (distance <= self.sliding_window)
+                mask = (distance < 0) | ((~is_sink) & (~in_window))
+                
+            else:
+                too_far_mask = distance > self.sliding_window
+                mask = mask | too_far_mask
+        float_mask = torch.zeros(
+            1, 1, query_length, key_length, dtype=dtype, device=device
+        )
+        float_mask.masked_fill_(mask.unsqueeze(0).unsqueeze(0), -1e9)
+        if len(self._mask_cache) >= self._max_cache_size:
+            oldest_key = next(iter(self._mask_cache))
+            del self._mask_cache[oldest_key]
+        self._mask_cache[cache_key] = float_mask.detach().cpu()
+        return float_mask
+
+    def _xformers_attention(
+        self, query_states, key_states, value_states, kv_seq_len, use_sliding_window
+    ):
+        batch_size, num_heads, seq_length, head_dim = query_states.shape
+        attn_bias = self._create_causal_mask(
+            seq_length,
+            kv_seq_len,
+            query_states.dtype,
+            query_states.device,
+            use_sliding_window,
+        )
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        attn_output = memory_efficient_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_bias=attn_bias,
+            p=self.config.attention_dropout if self.training else 0.0,
+        )
+        attn_output = attn_output.reshape(batch_size, seq_length, self.hidden_size)
+        return attn_output
+
+    def _sdpa_attention(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        kv_seq_len,
+        use_sliding_window,
+    ):
+        batch_size, num_heads, seq_length, head_dim = query_states.shape
+        if attention_mask is None:
+            attention_mask = self._create_causal_mask(
+                seq_length,
+                kv_seq_len,
+                query_states.dtype,
+                query_states.device,
+                use_sliding_window,
+            )
+        if self.alibi is not None:
+            alibi_bias = self.alibi(seq_length, kv_seq_len)
+            attention_mask = attention_mask + alibi_bias
+        attn_output = F.scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            attn_mask=attention_mask,
+            dropout_p=self.config.attention_dropout if self.training else 0.0,
+            is_causal=False,
+        )
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, seq_length, self.hidden_size)
+        return attn_output
+
+    def _standard_attention(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        kv_seq_len,
+        use_sliding_window,
+    ):
+        batch_size, num_heads, seq_length, head_dim = query_states.shape
+        attn_weights = torch.matmul(
+            query_states, key_states.transpose(2, 3)
+        ) / math.sqrt(head_dim)
+        attn_weights = torch.clamp(attn_weights, min=-50.0, max=50.0)
+        attn_weights = soft_cap_logits(attn_weights, self.attn_logit_softcapping)
+        if attention_mask is None:
+            attention_mask = self._create_causal_mask(
+                seq_length,
+                kv_seq_len,
+                attn_weights.dtype,
+                attn_weights.device,
+                use_sliding_window,
+            )
+        if self.alibi is not None:
+            alibi_bias = self.alibi(seq_length, kv_seq_len)
+            attention_mask = attention_mask + alibi_bias
+        attn_weights = attn_weights + attention_mask
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+            query_states.dtype
+        )
+        attn_weights = self.attention_dropout(attn_weights)
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, seq_length, self.hidden_size)
+        return attn_output
+
+class CacaCrossAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_key_value_groups = self.num_heads // self.num_key_value_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_key_value_heads * self.head_dim,
+            bias=config.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_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.attention_dropout = nn.Dropout(config.attention_dropout)
+
+    def forward(self, hidden_states, encoder_hidden_states, attention_mask=None):
+        batch_size, seq_length, _ = hidden_states.size()
+        encoder_seq_length = encoder_hidden_states.size(1)
+        query_states = self.q_proj(hidden_states)
+        key_states = self.k_proj(encoder_hidden_states)
+        value_states = self.v_proj(encoder_hidden_states)
+        query_states = query_states.view(
+            batch_size, seq_length, self.num_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.view(
+            batch_size, encoder_seq_length, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        value_states = value_states.view(
+            batch_size, encoder_seq_length, self.num_key_value_heads, self.head_dim
+        ).transpose(1, 2)
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        attn_weights = torch.matmul(
+            query_states, key_states.transpose(2, 3)
+        ) / math.sqrt(self.head_dim)
+        if attention_mask is not None:
+            attn_weights = attn_weights + attention_mask
+        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+            query_states.dtype
+        )
+        attn_weights = self.attention_dropout(attn_weights)
+        attn_output = torch.matmul(attn_weights, value_states)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(batch_size, seq_length, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+        return attn_output
+
+class CacaMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=config.mlp_bias
+        )
+        self.up_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=config.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=config.mlp_bias
+        )
+        self.dropout = nn.Dropout(config.hidden_dropout)
+
+    def forward(self, x):
+        gate = F.silu(self.gate_proj(x))
+        up = self.up_proj(x)
+        hidden = gate * up
+        hidden = self.dropout(hidden)
+        output = self.down_proj(hidden)
+        return output
+
+class CacaDecoderLayer(nn.Module):
+    def __init__(self, config, layer_idx):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.self_attn = CacaAttention(config, layer_idx=layer_idx)
+        self.use_moe = config.use_moe and (layer_idx % config.moe_layer_frequency == 0)
+        if self.use_moe:
+            self.mlp = MixtureOfExperts(config)
+        else:
+            self.mlp = CacaMLP(config)
+        self.use_cross_attention = config.use_cross_attention and (
+            layer_idx % config.cross_attention_frequency == 0
+        )
+        if self.use_cross_attention:
+            self.cross_attn = CacaCrossAttention(config)
+            self.cross_attn_layernorm = 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_dropout = nn.Dropout(config.residual_dropout)
+        if config.use_layer_scale:
+            self.layer_scale_1 = LayerScale(config.hidden_size, config.layer_scale_init)
+            self.layer_scale_2 = LayerScale(config.hidden_size, config.layer_scale_init)
+            if self.use_cross_attention:
+                self.layer_scale_cross = LayerScale(
+                    config.hidden_size, config.layer_scale_init
+                )
+        else:
+            self.layer_scale_1 = None
+            self.layer_scale_2 = None
+            self.layer_scale_cross = None
+        if config.use_stochastic_depth:
+            drop_prob = (
+                config.stochastic_depth_prob * layer_idx / config.num_hidden_layers
+            )
+            self.stochastic_depth = StochasticDepth(drop_prob)
+        else:
+            self.stochastic_depth = None
+        if config.use_mixture_of_depths:
+            self.mod_router = MixtureOfDepthsRouter(
+                config.hidden_size, config.mod_capacity_factor, config.mod_route_method
+            )
+        else:
+            self.mod_router = None
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        past_key_value=None,
+        use_cache=False,
+    ):
+        aux_loss = 0.0
+        z_loss = 0.0
+        if self.mod_router is not None:
+            process_mask = self.mod_router(hidden_states)
+            tokens_to_process = hidden_states[process_mask]
+            if tokens_to_process.numel() == 0:
+                present_key_value = past_key_value if use_cache else None
+                return hidden_states, present_key_value, aux_loss, z_loss
+        else:
+            process_mask = None
+            tokens_to_process = hidden_states
+        residual = tokens_to_process
+        tokens_to_process = self.input_layernorm(tokens_to_process)
+        attn_output, present_key_value = self.self_attn(
+            tokens_to_process,
+            attention_mask,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+        )
+        if self.layer_scale_1 is not None:
+            attn_output = self.layer_scale_1(attn_output)
+        if self.stochastic_depth is not None:
+            attn_output = self.stochastic_depth(attn_output, self.training)
+        tokens_to_process = residual + self.residual_dropout(attn_output)
+        if self.use_cross_attention and encoder_hidden_states is not None:
+            residual = tokens_to_process
+            tokens_to_process = self.cross_attn_layernorm(tokens_to_process)
+            cross_attn_output = self.cross_attn(
+                tokens_to_process,
+                encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+            )
+            if self.layer_scale_cross is not None:
+                cross_attn_output = self.layer_scale_cross(cross_attn_output)
+            if self.stochastic_depth is not None:
+                cross_attn_output = self.stochastic_depth(
+                    cross_attn_output, self.training
+                )
+            tokens_to_process = residual + self.residual_dropout(cross_attn_output)
+        residual = tokens_to_process
+        tokens_to_process = self.post_attention_layernorm(tokens_to_process)
+        if self.use_moe:
+            mlp_output, moe_aux_loss, moe_z_loss = self.mlp(tokens_to_process)
+            aux_loss += moe_aux_loss
+            z_loss += moe_z_loss
+        else:
+            mlp_output = self.mlp(tokens_to_process)
+        if self.layer_scale_2 is not None:
+            mlp_output = self.layer_scale_2(mlp_output)
+        if self.stochastic_depth is not None:
+            mlp_output = self.stochastic_depth(mlp_output, self.training)
+        tokens_to_process = residual + self.residual_dropout(mlp_output)
+        if process_mask is not None:
+            hidden_states[process_mask] = tokens_to_process
+        else:
+            hidden_states = tokens_to_process
+        return hidden_states, present_key_value, aux_loss, z_loss
+
+class VisionEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        vision_config = config.vision_config
+        self.patch_size = vision_config.get("patch_size", 14)
+        self.image_size = vision_config.get("image_size", 224)
+        self.num_channels = vision_config.get("num_channels", 3)
+        self.hidden_size = vision_config.get("hidden_size", 1024)
+        self.num_layers = vision_config.get("num_layers", 24)
+        self.num_heads = vision_config.get("num_heads", 16)
+        self.intermediate_size = vision_config.get("intermediate_size", 4096)
+        self.layer_norm_eps = vision_config.get("layer_norm_eps", 1e-6)
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.patch_embed = nn.Sequential(
+            nn.Conv2d(
+                self.num_channels,
+                self.hidden_size,
+                kernel_size=self.patch_size,
+                stride=self.patch_size,
+                bias=False,
+            ),
+            nn.Dropout(p=vision_config.get("dropout", 0.0)),
+        )
+        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.pos_drop = nn.Dropout(p=vision_config.get("dropout", 0.0))
+        self.blocks = nn.ModuleList(
+            [
+                VisionTransformerBlock(
+                    dim=self.hidden_size,
+                    num_heads=self.num_heads,
+                    mlp_ratio=self.intermediate_size / self.hidden_size,
+                    dropout=vision_config.get("dropout", 0.0),
+                    layer_norm_eps=self.layer_norm_eps,
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps)
+        self._init_weights()
+
+    def _init_weights(self):
+        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[0].weight, std=0.02)
+
+    def forward(self, pixel_values):
+        batch_size = pixel_values.shape[0]
+        x = self.patch_embed(pixel_values)
+        x = x.flatten(2).transpose(1, 2)
+        cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+        x = torch.cat([cls_tokens, x], dim=1)
+        x = x + self.pos_embed
+        x = self.pos_drop(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.norm(x)
+        return x
+
+class VisionTransformerBlock(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4.0, dropout=0.0, layer_norm_eps=1e-6):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, eps=layer_norm_eps)
+        self.attn = nn.MultiheadAttention(
+            dim, num_heads, dropout=dropout, batch_first=True
+        )
+        self.drop_path1 = nn.Dropout(dropout)
+        self.norm2 = nn.LayerNorm(dim, eps=layer_norm_eps)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, mlp_hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(mlp_hidden_dim, dim),
+            nn.Dropout(dropout),
+        )
+        self.drop_path2 = nn.Dropout(dropout)
+
+    def forward(self, x):
+        residual = x
+        x = self.norm1(x)
+        x = self.attn(x, x, x, need_weights=False)[0]
+        x = self.drop_path1(x)
+        x = residual + x
+        residual = x
+        x = self.norm2(x)
+        x = self.mlp(x)
+        x = self.drop_path2(x)
+        x = residual + x
+        return x
+
+class AudioEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        audio_config = config.audio_config
+        self.num_mel_bins = audio_config.get("num_mel_bins", 80)
+        self.hidden_size = audio_config.get("hidden_size", 1024)
+        self.num_layers = audio_config.get("num_layers", 12)
+        self.num_heads = audio_config.get("num_heads", 16)
+        self.intermediate_size = audio_config.get("intermediate_size", 4096)
+        self.max_audio_length = audio_config.get("max_audio_length", 3000)
+        self.dropout = audio_config.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(p=self.dropout),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv1d(
+                self.hidden_size, self.hidden_size, kernel_size=3, stride=2, padding=1
+            ),
+            nn.GELU(),
+            nn.Dropout(p=self.dropout),
+        )
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, self.max_audio_length // 2, self.hidden_size)
+        )
+        self.pos_drop = nn.Dropout(p=self.dropout)
+        self.blocks = nn.ModuleList(
+            [
+                AudioTransformerBlock(
+                    dim=self.hidden_size,
+                    num_heads=self.num_heads,
+                    mlp_ratio=self.intermediate_size / self.hidden_size,
+                    dropout=self.dropout,
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(self.hidden_size)
+        self._init_weights()
+
+    def _init_weights(self):
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+
+    def forward(self, audio_features):
+        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_embed_interp = F.interpolate(
+                self.pos_embed.transpose(1, 2),
+                size=seq_len,
+                mode="linear",
+                align_corners=False,
+            ).transpose(1, 2)
+            x = x + pos_embed_interp
+        x = self.pos_drop(x)
+        for block in self.blocks:
+            x = block(x)
+        x = self.norm(x)
+        return x
+
+class AudioTransformerBlock(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4.0, dropout=0.0):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(
+            dim, num_heads, dropout=dropout, batch_first=True
+        )
+        self.drop_path1 = nn.Dropout(dropout)
+        self.norm2 = nn.LayerNorm(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, mlp_hidden_dim),
+            nn.GELU(),
+            nn.Dropout(dropout),
+            nn.Linear(mlp_hidden_dim, dim),
+            nn.Dropout(dropout),
+        )
+        self.drop_path2 = nn.Dropout(dropout)
+
+    def forward(self, x):
+        residual = x
+        x = self.norm1(x)
+        x = self.attn(x, x, x, need_weights=False)[0]
+        x = self.drop_path1(x)
+        x = residual + x
+        residual = x
+        x = self.norm2(x)
+        x = self.mlp(x)
+        x = self.drop_path2(x)
+        x = residual + x
+        return x
+
+class MultiModalProjector(nn.Module):
+    def __init__(self, input_size, output_size, projector_type="mlp", num_layers=2):
+        super().__init__()
+        self.projector_type = projector_type
+        if projector_type == "linear":
+            self.projector = nn.Linear(input_size, output_size)
+        elif projector_type == "mlp":
+            layers = []
+            current_size = input_size
+            for i in range(num_layers - 1):
+                layers.extend(
+                    [nn.Linear(current_size, output_size), nn.GELU(), nn.Dropout(0.1)]
+                )
+                current_size = output_size
+            layers.append(nn.Linear(current_size, output_size))
+            self.projector = nn.Sequential(*layers)
+        elif projector_type == "perceiver":
+            self.projector = PerceiverResampler(
+                input_size, output_size, num_latents=64, num_layers=2
+            )
+        elif projector_type == "qformer":
+            self.projector = QFormerProjector(
+                input_size, output_size, num_queries=32, num_layers=2
+            )
+        else:
+            raise ValueError(f"projector_type tidak dikenal: {projector_type}")
+
+    def forward(self, x):
+        return self.projector(x)
+
+class PerceiverResampler(nn.Module):
+    def __init__(self, input_size, output_size, num_latents=64, num_layers=2):
+        super().__init__()
+        self.num_latents = num_latents
+        self.latents = nn.Parameter(torch.randn(num_latents, output_size))
+        self.layers = nn.ModuleList(
+            [
+                PerceiverLayer(output_size, input_size if i == 0 else output_size)
+                for i in range(num_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(output_size)
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        latents = self.latents.unsqueeze(0).expand(batch_size, -1, -1)
+        for i, layer in enumerate(self.layers):
+            if i == 0:
+                latents = layer(latents, x)
+            else:
+                latents = layer(latents, latents)
+        return self.norm(latents)
+
+class PerceiverLayer(nn.Module):
+    def __init__(self, query_dim, key_dim):
+        super().__init__()
+        self.cross_attn = nn.MultiheadAttention(
+            query_dim, num_heads=8, kdim=key_dim, vdim=key_dim, batch_first=True
+        )
+        self.mlp = nn.Sequential(
+            nn.LayerNorm(query_dim),
+            nn.Linear(query_dim, query_dim * 4),
+            nn.GELU(),
+            nn.Linear(query_dim * 4, query_dim),
+        )
+        self.norm1 = nn.LayerNorm(query_dim)
+        self.norm2 = nn.LayerNorm(query_dim)
+
+    def forward(self, query, key):
+        query = (
+            query + self.cross_attn(self.norm1(query), key, key, need_weights=False)[0]
+        )
+        query = query + self.mlp(self.norm2(query))
+        return query
+
+class QFormerProjector(nn.Module):
+    def __init__(self, input_size, output_size, num_queries=32, num_layers=2):
+        super().__init__()
+        self.num_queries = num_queries
+        self.query_embeds = nn.Parameter(torch.randn(num_queries, output_size))
+        self.query_layers = nn.ModuleList(
+            [
+                nn.TransformerEncoderLayer(
+                    d_model=output_size,
+                    nhead=8,
+                    dim_feedforward=output_size * 4,
+                    batch_first=True,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.cross_attn_layers = nn.ModuleList(
+            [
+                nn.MultiheadAttention(
+                    output_size,
+                    num_heads=8,
+                    kdim=input_size,
+                    vdim=input_size,
+                    batch_first=True,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(output_size)
+
+    def forward(self, x):
+        batch_size = x.shape[0]
+        queries = self.query_embeds.unsqueeze(0).expand(batch_size, -1, -1)
+        for query_layer, cross_attn_layer in zip(
+            self.query_layers, self.cross_attn_layers
+        ):
+            queries = query_layer(queries)
+            queries = queries + cross_attn_layer(queries, x, x, need_weights=False)[0]
+        return self.norm(queries)
+
+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 isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, CacaModel):
+            module.gradient_checkpointing = value
+
+class CacaModel(CacaPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [
+                CacaDecoderLayer(config, layer_idx=idx)
+                for idx in range(config.num_hidden_layers)
+            ]
+        )
+        self.norm = CacaRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.gradient_checkpointing = False
+        if config.use_multimodal:
+            if config.vision_config:
+                self.vision_encoder = VisionEncoder(config)
+                vision_hidden_size = config.vision_config.get("hidden_size", 768)
+                self.vision_projector = MultiModalProjector(
+                    vision_hidden_size,
+                    config.hidden_size,
+                    projector_type=config.vision_config.get("projector_type", "mlp"),
+                )
+            else:
+                self.vision_encoder = None
+                self.vision_projector = None
+            if config.audio_config:
+                self.audio_encoder = AudioEncoder(config)
+                audio_hidden_size = config.audio_config.get("hidden_size", 768)
+                self.audio_projector = MultiModalProjector(
+                    audio_hidden_size,
+                    config.hidden_size,
+                    projector_type=config.audio_config.get("projector_type", "mlp"),
+                )
+            else:
+                self.audio_encoder = None
+                self.audio_projector = None
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def _prepare_attention_mask(self, attention_mask, input_shape, dtype):
+        if attention_mask is None:
+            return None
+        batch_size, seq_length = input_shape
+        if attention_mask.dim() == 2:
+            attention_mask = attention_mask[:, None, None, :]
+        elif attention_mask.dim() == 3:
+            attention_mask = attention_mask[:, None, :, :]
+        attention_mask = attention_mask.to(dtype=dtype)
+        attention_mask = (1.0 - attention_mask) * torch.finfo(dtype).min
+        return attention_mask
+
+    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 not None:
+            batch_size, seq_length = input_ids.shape
+            device = input_ids.device
+            hidden_states = self.embed_tokens(input_ids)
+        else:
+            raise ValueError("input_ids tidak boleh None")
+        
+        if pixel_values is not None:
+            pixel_values = pixel_values.to(device)
+        if audio_features is not None:
+            audio_features = audio_features.to(device)
+        
+        encoder_hidden_states = None
+        encoder_attention_mask = None
+        if self.config.use_multimodal:
+            multimodal_embeds = []
+            if pixel_values is not None and self.vision_encoder is not None:
+                vision_features = self.vision_encoder(pixel_values.to(hidden_states.device))
+                vision_embeds = self.vision_projector(vision_features)
+                multimodal_embeds.append(vision_embeds)
+            if audio_features is not None and self.audio_encoder is not None:
+                audio_encoded = self.audio_encoder(audio_features.to(hidden_states.device))
+                audio_embeds = self.audio_projector(audio_encoded)
+                multimodal_embeds.append(audio_embeds)
+            if multimodal_embeds and self.config.use_cross_attention:
+                encoder_hidden_states = torch.cat(multimodal_embeds, dim=1)
+                encoder_seq_len = encoder_hidden_states.shape[1]
+                encoder_attention_mask = torch.ones(
+                    batch_size,
+                    encoder_seq_len,
+                    dtype=hidden_states.dtype,
+                    device=hidden_states.device,
+                )
+
+            elif multimodal_embeds:
+                multimodal_concat = torch.cat(multimodal_embeds, dim=1)
+                max_multimodal_tokens = self.config.max_position_embeddings // 4
+                if multimodal_concat.shape[1] > max_multimodal_tokens:
+                    logger.warning(
+                        f"Multimodal tokens ({multimodal_concat.shape[1]}) > max ({max_multimodal_tokens}). "
+                        f"Truncating..."
+                    )
+                    multimodal_concat = multimodal_concat[:, :max_multimodal_tokens]
+                hidden_states = torch.cat([multimodal_concat, hidden_states], dim=1)
+                seq_length = hidden_states.shape[1]
+                if attention_mask is not None:
+                    multimodal_mask = torch.ones(
+                        batch_size,
+                        multimodal_concat.shape[1],
+                        dtype=attention_mask.dtype,
+                        device=attention_mask.device,
+                    )
+                    attention_mask = torch.cat([multimodal_mask, attention_mask], dim=1)
+                else:
+                    attention_mask = torch.ones(
+                        batch_size,
+                        seq_length,
+                        dtype=hidden_states.dtype,
+                        device=device,
+                    )
+        if attention_mask is not None:
+            attention_mask = self._prepare_attention_mask(
+                attention_mask, (batch_size, seq_length), hidden_states.dtype
+            )
+        if encoder_attention_mask is not None and self.config.use_cross_attention:
+            encoder_attention_mask = self._prepare_attention_mask(
+                encoder_attention_mask,
+                (batch_size, encoder_hidden_states.shape[1]),
+                hidden_states.dtype,
+            )
+        
+        if use_cache:
+            if past_key_values is None:
+                past_key_values = tuple([None] * len(self.layers))
+
+        present_key_values = [] if use_cache else None
+        all_hidden_states = [] if output_hidden_states else None
+        total_aux_loss = 0.0
+        total_z_loss = 0.0
+        for idx, layer in enumerate(self.layers):
+            if output_hidden_states:
+                all_hidden_states.append(hidden_states)
+            past_key_value = (
+                past_key_values[idx] if past_key_values is not None else None
+            )
+            if self.gradient_checkpointing and self.training and not use_cache:
+                hidden_states, aux_loss, z_loss = self._gradient_checkpointing_forward(
+                    layer,
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+                present_key_value = None
+            else:
+                hidden_states, present_key_value, aux_loss, z_loss = layer(
+                    hidden_states,
+                    attention_mask,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    past_key_value=past_key_value,
+                    use_cache=use_cache,
+                )
+            if use_cache:
+                present_key_values.append(present_key_value)
+            total_aux_loss += aux_loss
+            total_z_loss += z_loss
+            
+            if self.training and torch.cuda.is_available():
+                allocated_gb = torch.cuda.memory_allocated() / 1024**3
+                reserved_gb = torch.cuda.memory_reserved() / 1024**3
+                if allocated_gb > 10:
+                    logger.warning(
+                        f"High GPU memory usage - Allocated: {allocated_gb:.2f}GB, "
+                        f"Reserved: {reserved_gb:.2f}GB"
+                    )
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states:
+            all_hidden_states.append(hidden_states)
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    present_key_values,
+                    all_hidden_states,
+                    total_aux_loss,
+                    total_z_loss,
+                ]
+                if v is not None
+            )
+        return (
+            BaseModelOutputWithPast(
+                last_hidden_state=hidden_states,
+                past_key_values=tuple(present_key_values) if use_cache else None,
+                hidden_states=all_hidden_states,
+                attentions=None,
+            ),
+            total_aux_loss,
+            total_z_loss,
+        )
+
+    def _gradient_checkpointing_forward(
+        self,
+        layer,
+        hidden_states,
+        attention_mask,
+        encoder_hidden_states,
+        encoder_attention_mask,
+    ):
+        from torch.utils.checkpoint import checkpoint
+
+        def custom_forward(hidden_states, attention_mask, encoder_hidden_states, 
+                       encoder_attention_mask):
+            output, _, aux_loss, z_loss = layer(
+                hidden_states, attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                past_key_value=None,
+                use_cache=False,
+            )
+            
+            return output, aux_loss, z_loss
+            
+        hidden_states, aux_loss, z_loss = checkpoint(
+            custom_forward,
+            hidden_states, attention_mask,
+            encoder_hidden_states, encoder_attention_mask,
+            use_reentrant=False,
+        )
+        return hidden_states, aux_loss, z_loss
+        
+class CacaForCausalLM(CacaPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config):
+        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, value):
+        self.model.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    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 convert."
+                )
+            if (input_ids < 0).any():
+                neg_vals = input_ids[input_ids < 0].unique().tolist()
+                raise ValueError(f"input_ids mengandung nilai negatif: {neg_vals}")
+            max_val = input_ids.max().item()
+            if max_val >= self.config.vocab_size:
+                raise ValueError(
+                    f"input_ids mengandung nilai >= vocab_size. "
+                    f"Max value: {max_val}, vocab_size: {self.config.vocab_size:,}"
+                )
+                
+            if labels is not None:
+                if not labels.dtype in [torch.long, torch.int, torch.int32, torch.int64]:
+                    raise TypeError(f"labels harus integer dtype, dapat {labels.dtype}")
+                if (labels[labels != -100] < 0).any():
+                    raise ValueError(f"labels mengandung nilai negatif (selain -100)")
+                max_label = labels[labels != -100].max().item() if (labels != -100).any() else 0
+                if max_label >= self.config.vocab_size:
+                    raise ValueError(
+                        f"labels mengandung nilai >= vocab_size. "
+                        f"Max: {max_label}, vocab_size: {self.config.vocab_size}"
+                    )
+            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]})"
+                    )
+
+        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=return_dict,
+        )
+        if return_dict:
+            hidden_states = outputs.last_hidden_state
+        else:
+            hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states)
+        if self.config.final_logit_softcapping:
+            logits = soft_cap_logits(logits, self.config.final_logit_softcapping)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss_fct = nn.CrossEntropyLoss()
+            lm_loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
+            )
+            if self.config.use_moe:
+                total_loss = (
+                    lm_loss
+                    + (self.config.router_aux_loss_coef * aux_loss)
+                    + (self.config.router_z_loss_coef * z_loss)
+                )
+            else:
+                total_loss = lm_loss
+            loss = total_loss
+        if not return_dict:
+            output = (logits,)
+            if return_dict:
+                output += tuple(
+                    v
+                    for v in [outputs.past_key_values, outputs.hidden_states]
+                    if v is not None
+                )
+            return ((loss,) + output) if loss is not None else output
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values if return_dict else None,
+            hidden_states=outputs.hidden_states if return_dict else None,
+            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 = (
+            past_key_values is not None 
+            and len(past_key_values) > 0 
+            and past_key_values[0] is not None
+        )
+        
+        if has_past:
+            input_ids = input_ids[:, -1:]
+     
+        if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+            attention_mask = attention_mask[:, -input_ids.shape[1]:]
+            
+        if inputs_embeds is not None and not has_past:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        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 if not has_past else None,
+                "audio_features": audio_features if not has_past else None,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            if layer_past is not None and len(layer_past) > 0:
+                reordered_past += (
+                    tuple(
+                        past_state.index_select(0, beam_idx.to(past_state.device))
+                        for past_state in layer_past
+                        if past_state is not None
+                    ),
+                )
+            else:
+                reordered_past += (None,)
+        return reordered_past
+
+    def save_pretrained(self, save_directory, **kwargs):
+        has_quant_config = hasattr(self.config, 'quantization_config')
+        quantization_config_backup = getattr(self.config, 'quantization_config', None)
+        
+        if has_quant_config and quantization_config_backup is None:
+            delattr(self.config, 'quantization_config')
+        
+        try:
+            super().save_pretrained(save_directory, **kwargs)
+        finally:
+            if has_quant_config:
+                self.config.quantization_config = quantization_config_backup
+
+class CacaForCausalLMQuantized(CacaForCausalLM):
+    def __init__(self, config, quantization_config=None):
+        super().__init__(config)
+        self.quantization_config = quantization_config
+        if quantization_config:
+            self._apply_quantization()
+
+    def _apply_quantization(self):
+        if self.quantization_config.get("load_in_8bit"):
+            self._quantize_8bit()
+        elif self.quantization_config.get("load_in_4bit"):
+            self._quantize_4bit()
+
+    def _quantize_8bit(self):
+        try:
+            import bitsandbytes as bnb
+
+            for name, module in self.named_modules():
+                if isinstance(module, nn.Linear):
+                    has_bias = module.bias is not None
+                    new_module = bnb.nn.Linear8bitLt(
+                        module.in_features,
+                        module.out_features,
+                        has_bias,
+                        threshold=self.quantization_config.get(
+                            "llm_int8_threshold", 6.0
+                        ),
+                    )
+                    new_module.weight = module.weight
+                    if has_bias:
+                        new_module.bias = module.bias
+                    parent_name = ".".join(name.split(".")[:-1])
+                    child_name = name.split(".")[-1]
+                    if parent_name:
+                        parent = self.get_submodule(parent_name)
+                        setattr(parent, child_name, new_module)
+                    else:
+                        setattr(self, child_name, new_module)
+            logger.info("Quantisasi 8-bit berhasil diterapkan")
+        except ImportError:
+            logger.error("bitsandbytes tidak terinstall! pip install bitsandbytes")
+
+    def _quantize_4bit(self):
+        try:
+            import bitsandbytes as bnb
+
+            compute_dtype = torch.float16
+            if 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 isinstance(module, nn.Linear):
+                    has_bias = module.bias is not None
+                    new_module = 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_module.weight = module.weight
+                    if has_bias:
+                        new_module.bias = module.bias
+                    parent_name = ".".join(name.split(".")[:-1])
+                    child_name = name.split(".")[-1]
+                    if parent_name:
+                        parent = self.get_submodule(parent_name)
+                        setattr(parent, child_name, new_module)
+                    else:
+                        setattr(self, child_name, new_module)
+            logger.info("Quantisasi 4-bit berhasil diterapkan")
+        except ImportError:
+            logger.error("bitsandbytes tidak terinstall!")
+
+    @classmethod
+    def from_pretrained_quantized(cls, model_path, quantization_config):
+        config = CacaConfig.from_pretrained(model_path)
+        model = cls(config, quantization_config=quantization_config)
+        state_dict = torch.load(f"{model_path}/pytorch_model.bin", map_location="cpu")
+        model.load_state_dict(state_dict, strict=False)
+        return model