from __future__ import annotations import math import os import json import logging import warnings from collections import OrderedDict, defaultdict from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from transformers import PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from transformers.generation.utils import GenerationMixin logger = logging.getLogger(__name__) # --- Dependensi opsional --- 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") 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 --- 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 or (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 or ["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": 2, "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": 2, "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": 2, "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": 2, "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": 2, "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": 4, "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": 2, "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": 2, "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": 2, "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": 2, "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": 4, "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": 4, "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": 4, "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": 4, "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": 2, "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": 4, "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": 8, "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": 2, "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": 2, "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": 4, "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": 8, "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": 8, "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": 4, "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": 16, "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": 8, "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": 8, "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, }, } # --- 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) 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 ffn_l = 3 * h * i norm_l = 2 * h lm_head = 0 if tie else v * h # + h = final RMSNorm total = embed + l * (attn_l + ffn_l + norm_l) + 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_per_layer": attn_l, "ffn_per_layer": ffn_l, } 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" # --- 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] expert_pos = (top_idx[mask] == i).float().argmax(dim=-1) weights_i = top_w[mask][ torch.arange(expert_out.shape[0], device=out.device), expert_pos ] out.scatter_add_( 0, token_ids.unsqueeze(1).expand_as(expert_out), expert_out * weights_i.unsqueeze(1) ) 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 / Audio 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) 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 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 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 int dtype, dapat {input_ids.dtype}") if (input_ids < 0).any(): raise ValueError(f"input_ids memiliki {(input_ids < 0).sum()} nilai negatif") max_id = input_ids.max().item() if max_id >= self.config.vocab_size: raise ValueError(f"input_ids max={max_id} >= vocab_size={self.config.vocab_size}") if labels is not None: if labels.shape != input_ids.shape: raise ValueError(f"labels shape {labels.shape} != input_ids shape {input_ids.shape}") valid = labels[labels != -100] if valid.numel() > 0 and valid.max().item() >= self.config.vocab_size: raise ValueError(f"labels max={valid.max().item()} >= vocab_size={self.config.vocab_size}") if valid.numel() == 0: logger.warning("Semua label adalah -100. Loss tidak akan dihitung.") if attention_mask is not None: if attention_mask.shape != input_ids.shape: raise ValueError( f"attention_mask shape {attention_mask.shape} != input_ids shape {input_ids.shape}" ) 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 } # --- AutoClass registration --- from transformers import AutoConfig, AutoModel, AutoModelForCausalLM AutoConfig.register("caca", CacaConfig) AutoModel.register(CacaConfig, CacaModel) AutoModelForCausalLM.register(CacaConfig, CacaForCausalLM)