diff --git "a/caca_transformers.py" "b/caca_transformers.py" new file mode 100644--- /dev/null +++ "b/caca_transformers.py" @@ -0,0 +1,2601 @@ +# --- Import Libraries --- +import torch.nn as nn +import torch.nn.functional as F +import math +import os +import json +import logging +from transformers import PretrainedConfig, PreTrainedModel +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.generation.utils import GenerationMixin +from typing import Optional, List, Dict, Tuple, Union, Any +from collections import OrderedDict, defaultdict + +logger = logging.getLogger(__name__) + +try: + from flash_attn import flash_attn_func, flash_attn_varlen_func + HAS_FLASH_ATTN = True +except ImportError: + HAS_FLASH_ATTN = False + +try: + from xformers.ops import memory_efficient_attention + HAS_XFORMERS = True +except ImportError: + HAS_XFORMERS = False + +HAS_SDPA = hasattr(F, "scaled_dot_product_attention") + +def log(msg): + print(f"[LOG] {msg}") + +try: + import bitsandbytes as bnb + HAS_BNB = True +except ImportError: + HAS_BNB = False + +try: + from peft import LoraConfig, get_peft_model, TaskType + HAS_PEFT = True +except ImportError: + HAS_PEFT = False + +# --- CacaConfig & CacaModel--- +class CacaConfig(PretrainedConfig): + model_type = "caca" + + def __init__( + self, + vocab_size: int = 32000, + hidden_size: int = 2048, + intermediate_size: int = 8192, + num_hidden_layers: int = 24, + num_attention_heads: int = 32, + num_key_value_heads: int = 8, + head_dim: Optional[int] = 64, + max_position_embeddings: int = 8192, + rms_norm_eps: float = 1e-6, + qk_norm_eps: float = 1e-6, + initializer_range: float = 0.02, + use_mup: bool = False, + mup_base_width: int = 256, + use_cache: bool = True, + pad_token_id: Optional[int] = None, + bos_token_id: int = 1, + eos_token_id: int = 2, + tie_word_embeddings: bool = False, + rope_theta: float = 10000.0, + rope_scaling: Optional[Dict] = None, + use_rotary_embeddings: bool = True, + rope_type: str = "default", + rope_ntk_alpha: float = 1.0, + use_alibi: bool = False, + attention_bias: bool = False, + attention_dropout: float = 0.0, + attention_temperature: float = 1.0, + use_qk_norm: bool = True, + use_flash_attn: bool = True, + use_grouped_query_attention: bool = False, + use_multi_query_attention: bool = False, + sliding_window: Optional[int] = None, + use_longformer_attention: bool = False, + longformer_attention_window: int = 512, + attn_logit_softcapping: Optional[float] = None, + final_logit_softcapping: Optional[float] = None, + lm_logit_softcapping: Optional[float] = 30.0, + attention_sink_size: int = 4, + attention_sink_window: int = 1024, + use_attention_sink: bool = False, + attention_pattern: str = "all_global", + global_attention_every_n_layers: int = 2, + mlp_bias: bool = False, + hidden_dropout: float = 0.1, + residual_dropout: float = 0.1, + token_dropout: float = 0.0, + use_moe: bool = False, + num_experts: int = 8, + num_experts_per_tok: int = 2, + use_expert_choice: bool = False, + expert_choice_k: float = 0.125, + router_aux_loss_coef: float = 0.01, + router_z_loss_coef: float = 0.001, + moe_layer_frequency: int = 2, + expert_capacity_factor: float = 1.0, + use_grouped_moe: bool = False, + num_expert_groups: int = 1, + expert_load_monitoring: bool = True, + expert_load_warn_threshold: float = 0.3, + use_layer_scale: bool = False, + layer_scale_init: float = 1e-5, + use_stochastic_depth: bool = False, + stochastic_depth_prob: float = 0.1, + use_spectral_norm: bool = False, + label_smoothing: float = 0.0, + use_mixture_of_depths: bool = False, + mod_capacity_factor: float = 0.5, + mod_route_method: str = "learned", + use_cross_attention: bool = False, + cross_attention_frequency: int = 4, + use_multimodal: bool = False, + vision_config: Optional[Dict] = None, + audio_config: Optional[Dict] = None, + projector_hidden_size: Optional[int] = None, + use_soft_merging: bool = False, + merge_threshold: float = 0.5, + use_lora: bool = False, + lora_rank: int = 16, + lora_alpha: float = 32.0, + lora_dropout: float = 0.05, + lora_target_modules: Optional[List[str]] = None, + lora_bias: str = "none", + pretraining_tp: int = 1, + tensor_parallel_size: int = 1, + pipeline_parallel_size: int = 1, + gradient_checkpointing_granularity: str = "full", + nan_recovery_level: int = 1, + chat_template: Optional[str] = None, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + self.head_dim = head_dim if head_dim is not None else (hidden_size // num_attention_heads) + self.max_position_embeddings = max_position_embeddings + self.rms_norm_eps = rms_norm_eps + self.qk_norm_eps = qk_norm_eps + self.initializer_range = initializer_range + self.use_mup = use_mup + self.mup_base_width = mup_base_width + self.use_cache = use_cache + self.tie_word_embeddings = tie_word_embeddings + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.use_rotary_embeddings = use_rotary_embeddings + self.rope_type = rope_type + self.rope_ntk_alpha = rope_ntk_alpha + self.use_alibi = use_alibi + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.attention_temperature = attention_temperature + self.use_qk_norm = use_qk_norm + self.use_flash_attn = use_flash_attn + self.use_grouped_query_attention = use_grouped_query_attention + self.use_multi_query_attention = use_multi_query_attention + self.sliding_window = sliding_window + self.use_longformer_attention = use_longformer_attention + self.longformer_attention_window = longformer_attention_window + self.attn_logit_softcapping = attn_logit_softcapping + self.final_logit_softcapping = final_logit_softcapping + self.lm_logit_softcapping = lm_logit_softcapping + self.attention_sink_size = attention_sink_size + self.attention_sink_window = attention_sink_window + self.use_attention_sink = use_attention_sink + self.attention_pattern = attention_pattern + self.global_attention_every_n_layers = global_attention_every_n_layers + self.mlp_bias = mlp_bias + self.hidden_dropout = hidden_dropout + self.residual_dropout = residual_dropout + self.token_dropout = token_dropout + self.use_moe = use_moe + self.num_experts = num_experts + self.num_experts_per_tok = num_experts_per_tok + self.use_expert_choice = use_expert_choice + self.expert_choice_k = expert_choice_k + self.router_aux_loss_coef = router_aux_loss_coef + self.router_z_loss_coef = router_z_loss_coef + self.moe_layer_frequency = moe_layer_frequency + self.expert_capacity_factor = expert_capacity_factor + self.use_grouped_moe = use_grouped_moe + self.num_expert_groups = num_expert_groups + self.expert_load_monitoring = expert_load_monitoring + self.expert_load_warn_threshold = expert_load_warn_threshold + self.use_layer_scale = use_layer_scale + self.layer_scale_init = layer_scale_init + self.use_stochastic_depth = use_stochastic_depth + self.stochastic_depth_prob = stochastic_depth_prob + self.use_spectral_norm = use_spectral_norm + self.label_smoothing = label_smoothing + self.use_mixture_of_depths = use_mixture_of_depths + self.mod_capacity_factor = mod_capacity_factor + self.mod_route_method = mod_route_method + self.use_cross_attention = use_cross_attention + self.cross_attention_frequency = cross_attention_frequency + self.use_multimodal = use_multimodal + self.vision_config = vision_config or {} + self.audio_config = audio_config or {} + self.projector_hidden_size = projector_hidden_size or hidden_size + self.use_soft_merging = use_soft_merging + self.merge_threshold = merge_threshold + self.use_lora = use_lora + self.lora_rank = lora_rank + self.lora_alpha = lora_alpha + self.lora_dropout = lora_dropout + self.lora_target_modules = ( + lora_target_modules if lora_target_modules is not None else ["q_proj", "v_proj"] + ) + self.lora_bias = lora_bias + self.pretraining_tp = pretraining_tp + self.tensor_parallel_size = tensor_parallel_size + self.pipeline_parallel_size = pipeline_parallel_size + self.gradient_checkpointing_granularity = gradient_checkpointing_granularity + self.nan_recovery_level = nan_recovery_level + self.chat_template = chat_template or ( + "{% for message in messages %}" + "{% if message['role'] == 'system' %}System: {{ message['content'] }}\n" + "{% elif message['role'] == 'user' %}User: {{ message['content'] }}\n" + "{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}\n" + "{% endif %}{% endfor %}" + "{% if add_generation_prompt %}Assistant:{% endif %}" + ) + + self._validate_config() + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + def _validate_config(self) -> None: + errors: List[str] = [] + + if self.num_attention_heads % self.num_key_value_heads != 0: + errors.append( + f"num_attention_heads ({self.num_attention_heads}) harus habis dibagi " + f"num_key_value_heads ({self.num_key_value_heads})" + ) + if self.num_key_value_heads > self.num_attention_heads: + errors.append( + f"num_key_value_heads ({self.num_key_value_heads}) tidak boleh melebihi " + f"num_attention_heads ({self.num_attention_heads})" + ) + if self.hidden_size % self.num_attention_heads != 0: + errors.append( + f"hidden_size ({self.hidden_size}) harus habis dibagi " + f"num_attention_heads ({self.num_attention_heads})" + ) + + expected_head_dim = self.hidden_size // self.num_attention_heads + if self.head_dim != expected_head_dim: + logger.warning( + f"head_dim ({self.head_dim}) != hidden_size // num_heads ({expected_head_dim}). " + "Intentional hanya untuk model dengan custom head projection." + ) + + if self.vocab_size <= 0: + errors.append(f"vocab_size harus > 0, dapat {self.vocab_size}") + if self.vocab_size > 1_000_000: + logger.warning(f"vocab_size ({self.vocab_size:,}) sangat besar — bisa menyebabkan masalah memori.") + + if self.use_flash_attn and not HAS_FLASH_ATTN: + logger.warning("use_flash_attn=True tapi flash-attn belum terinstall. Fallback ke SDPA.") + + if self.sliding_window is not None: + if self.sliding_window > self.max_position_embeddings: + errors.append( + f"sliding_window ({self.sliding_window}) > " + f"max_position_embeddings ({self.max_position_embeddings})" + ) + if self.sliding_window < 128: + logger.warning(f"sliding_window ({self.sliding_window}) sangat kecil.") + + if self.use_moe: + if self.num_experts < self.num_experts_per_tok: + errors.append( + f"num_experts ({self.num_experts}) harus >= " + f"num_experts_per_tok ({self.num_experts_per_tok})" + ) + if self.moe_layer_frequency <= 0: + errors.append("moe_layer_frequency harus > 0") + if self.moe_layer_frequency > self.num_hidden_layers: + logger.warning( + f"moe_layer_frequency ({self.moe_layer_frequency}) > " + f"num_hidden_layers ({self.num_hidden_layers}) — MoE tidak akan aktif." + ) + if self.expert_capacity_factor <= 0: + errors.append("expert_capacity_factor harus > 0") + + if self.use_lora: + if self.lora_rank <= 0: + errors.append(f"lora_rank harus > 0, dapat {self.lora_rank}") + if self.lora_alpha <= 0: + errors.append(f"lora_alpha harus > 0, dapat {self.lora_alpha}") + if not HAS_PEFT: + logger.warning("use_lora=True tapi 'peft' belum terinstall. Jalankan: pip install peft") + + if self.use_mup and self.mup_base_width <= 0: + errors.append(f"mup_base_width harus > 0, dapat {self.mup_base_width}") + + if self.nan_recovery_level not in (0, 1, 2): + errors.append(f"nan_recovery_level harus 0, 1, atau 2, dapat {self.nan_recovery_level}") + + if not (0.0 <= self.label_smoothing < 1.0): + errors.append(f"label_smoothing harus dalam [0.0, 1.0), dapat {self.label_smoothing}") + + if self.final_logit_softcapping is not None and self.lm_logit_softcapping is not None: + if self.final_logit_softcapping != self.lm_logit_softcapping: + logger.warning( + f"final_logit_softcapping ({self.final_logit_softcapping}) dan " + f"lm_logit_softcapping ({self.lm_logit_softcapping}) berbeda. " + "lm_logit_softcapping akan dipakai untuk LM head." + ) + + if errors: + raise ValueError("CacaConfig validation gagal:\n" + "\n".join(f" • {e}" for e in errors)) + + def to_dict(self) -> Dict[str, Any]: + quant_backup = getattr(self, "quantization_config", None) + had_quant = hasattr(self, "quantization_config") + if had_quant and quant_backup is None: + delattr(self, "quantization_config") + try: + output = super().to_dict() + output["auto_map"] = { + "AutoConfig": "caca_transformers.CacaConfig", + "AutoModel": "caca_transformers.CacaModel", + "AutoModelForCausalLM": "caca_transformers.CacaForCausalLM", + } + output["architectures"] = ["CacaForCausalLM"] + finally: + if had_quant: + self.quantization_config = quant_backup + return output + + def get_mup_lr_multiplier(self, param_name: str) -> float: + if not self.use_mup: + return 1.0 + width_ratio = self.mup_base_width / self.hidden_size + if any(k in param_name for k in ("embed_tokens", "lm_head", "norm")): + return 1.0 + return width_ratio + + @classmethod + def from_variant(cls, variant: str, **overrides) -> "CacaConfig": + if variant not in MODEL_CONFIGS: + raise ValueError( + f"Variant '{variant}' tidak dikenal. " + f"Tersedia: {sorted(MODEL_CONFIGS.keys())}" + ) + cfg = {**MODEL_CONFIGS[variant], **overrides} + return cls(**cfg) + +MODEL_CONFIGS = { + "caca-1M": { + "vocab_size": 4000, "hidden_size": 64, "intermediate_size": 64, + "num_hidden_layers": 17, "num_attention_heads": 2, "num_key_value_heads": 2, + "head_dim": 32, "max_position_embeddings": 512, + }, + "caca-1.5M": { + "vocab_size": 4000, "hidden_size": 64, "intermediate_size": 64, + "num_hidden_layers": 40, "num_attention_heads": 2, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 512, + }, + "caca-2M": { + "vocab_size": 4000, "hidden_size": 128, "intermediate_size": 256, + "num_hidden_layers": 7, "num_attention_heads": 4, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 512, + }, + "caca-2.5M": { + "vocab_size": 4000, "hidden_size": 128, "intermediate_size": 512, + "num_hidden_layers": 6, "num_attention_heads": 4, "num_key_value_heads": 2, + "head_dim": 32, "max_position_embeddings": 512, + }, + "caca-3M": { + "vocab_size": 4000, "hidden_size": 160, "intermediate_size": 448, + "num_hidden_layers": 6, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 512, + }, + "caca-3.5M": { + "vocab_size": 4000, "hidden_size": 64, "intermediate_size": 192, + "num_hidden_layers": 56, "num_attention_heads": 2, "num_key_value_heads": 2, + "head_dim": 32, "max_position_embeddings": 512, + }, + "caca-4M": { + "vocab_size": 4000, "hidden_size": 96, "intermediate_size": 448, + "num_hidden_layers": 21, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 512, + }, + "caca-4.5M": { + "vocab_size": 4000, "hidden_size": 64, "intermediate_size": 128, + "num_hidden_layers": 97, "num_attention_heads": 2, "num_key_value_heads": 2, + "head_dim": 32, "max_position_embeddings": 512, + }, + "caca-5M": { + "vocab_size": 4000, "hidden_size": 256, "intermediate_size": 512, + "num_hidden_layers": 5, "num_attention_heads": 8, "num_key_value_heads": 4, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-6M": { + "vocab_size": 8000, "hidden_size": 192, "intermediate_size": 576, + "num_hidden_layers": 7, "num_attention_heads": 6, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-7M": { + "vocab_size": 8000, "hidden_size": 64, "intermediate_size": 256, + "num_hidden_layers": 91, "num_attention_heads": 2, "num_key_value_heads": 2, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-8M": { + "vocab_size": 8000, "hidden_size": 96, "intermediate_size": 64, + "num_hidden_layers": 131, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-9M": { + "vocab_size": 8000, "hidden_size": 448, "intermediate_size": 1024, + "num_hidden_layers": 1, "num_attention_heads": 14, "num_key_value_heads": 2, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-10M": { + "vocab_size": 8000, "hidden_size": 256, "intermediate_size": 512, + "num_hidden_layers": 10, "num_attention_heads": 8, "num_key_value_heads": 4, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-12M": { + "vocab_size": 8000, "hidden_size": 96, "intermediate_size": 256, + "num_hidden_layers": 100, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-14M": { + "vocab_size": 8000, "hidden_size": 160, "intermediate_size": 512, + "num_hidden_layers": 36, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-15M": { + "vocab_size": 8000, "hidden_size": 64, "intermediate_size": 256, + "num_hidden_layers": 227, "num_attention_heads": 2, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-17M": { + "vocab_size": 8000, "hidden_size": 96, "intermediate_size": 256, + "num_hidden_layers": 157, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-20M": { + "vocab_size": 8000, "hidden_size": 160, "intermediate_size": 448, + "num_hidden_layers": 63, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 32, "max_position_embeddings": 1024, + }, + "caca-22M": { + "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 704, + "num_hidden_layers": 30, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 1024, + }, + "caca-25M": { + "vocab_size": 16000, "hidden_size": 576, "intermediate_size": 1344, + "num_hidden_layers": 2, "num_attention_heads": 9, "num_key_value_heads": 3, + "head_dim": 64, "max_position_embeddings": 1024, + }, + "caca-27M": { + "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 512, + "num_hidden_layers": 53, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 1024, + }, + "caca-30M": { + "vocab_size": 16000, "hidden_size": 448, "intermediate_size": 1280, + "num_hidden_layers": 7, "num_attention_heads": 7, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 1024, + }, + "caca-35M": { + "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 512, + "num_hidden_layers": 69, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 1024, + }, + "caca-40M": { + "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 640, + "num_hidden_layers": 33, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 1024, + }, + "caca-45M": { + "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 832, + "num_hidden_layers": 32, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 1024, + }, + "caca-50M": { + "vocab_size": 16000, "hidden_size": 576, "intermediate_size": 1728, + "num_hidden_layers": 8, "num_attention_heads": 9, "num_key_value_heads": 3, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-55M": { + "vocab_size": 16000, "hidden_size": 128, "intermediate_size": 512, + "num_hidden_layers": 194, "num_attention_heads": 2, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-60M": { + "vocab_size": 16000, "hidden_size": 256, "intermediate_size": 640, + "num_hidden_layers": 79, "num_attention_heads": 4, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-65M": { + "vocab_size": 16000, "hidden_size": 256, "intermediate_size": 704, + "num_hidden_layers": 77, "num_attention_heads": 4, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-70M": { + "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 448, + "num_hidden_layers": 179, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-80M": { + "vocab_size": 16000, "hidden_size": 192, "intermediate_size": 448, + "num_hidden_layers": 207, "num_attention_heads": 3, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-90M": { + "vocab_size": 16000, "hidden_size": 768, "intermediate_size": 1984, + "num_hidden_layers": 11, "num_attention_heads": 12, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-100M": { + "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 1024, + "num_hidden_layers": 73, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-110M": { + "vocab_size": 16000, "hidden_size": 256, "intermediate_size": 896, + "num_hidden_layers": 115, "num_attention_heads": 4, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-120M": { + "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 832, + "num_hidden_layers": 105, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-130M": { + "vocab_size": 16000, "hidden_size": 448, "intermediate_size": 1024, + "num_hidden_layers": 63, "num_attention_heads": 7, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-140M": { + "vocab_size": 16000, "hidden_size": 384, "intermediate_size": 832, + "num_hidden_layers": 98, "num_attention_heads": 6, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-150M": { + "vocab_size": 16000, "hidden_size": 640, "intermediate_size": 1728, + "num_hidden_layers": 29, "num_attention_heads": 10, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-160M": { + "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 640, + "num_hidden_layers": 174, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-175M": { + "vocab_size": 16000, "hidden_size": 1536, "intermediate_size": 4096, + "num_hidden_layers": 5, "num_attention_heads": 24, "num_key_value_heads": 8, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-200M": { + "vocab_size": 16000, "hidden_size": 256, "intermediate_size": 896, + "num_hidden_layers": 225, "num_attention_heads": 4, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-225M": { + "vocab_size": 16000, "hidden_size": 512, "intermediate_size": 1408, + "num_hidden_layers": 74, "num_attention_heads": 8, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-250M": { + "vocab_size": 16000, "hidden_size": 448, "intermediate_size": 1216, + "num_hidden_layers": 104, "num_attention_heads": 7, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-275M": { + "vocab_size": 16000, "hidden_size": 576, "intermediate_size": 1536, + "num_hidden_layers": 71, "num_attention_heads": 9, "num_key_value_heads": 3, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-300M": { + "vocab_size": 16000, "hidden_size": 512, "intermediate_size": 1408, + "num_hidden_layers": 103, "num_attention_heads": 8, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-325M": { + "vocab_size": 16000, "hidden_size": 320, "intermediate_size": 1024, + "num_hidden_layers": 256, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-350M": { + "vocab_size": 16000, "hidden_size": 384, "intermediate_size": 1088, + "num_hidden_layers": 205, "num_attention_heads": 6, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-375M": { + "vocab_size": 16000, "hidden_size": 384, "intermediate_size": 960, + "num_hidden_layers": 250, "num_attention_heads": 6, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-400M": { + "vocab_size": 16000, "hidden_size": 512, "intermediate_size": 1344, + "num_hidden_layers": 141, "num_attention_heads": 8, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-450M": { + "vocab_size": 16000, "hidden_size": 512, "intermediate_size": 1536, + "num_hidden_layers": 147, "num_attention_heads": 8, "num_key_value_heads": 1, + "head_dim": 64, "max_position_embeddings": 2048, + }, + "caca-500M": { + "vocab_size": 16000, "hidden_size": 1024, "intermediate_size": 2688, + "num_hidden_layers": 44, "num_attention_heads": 16, "num_key_value_heads": 2, + "head_dim": 64, "max_position_embeddings": 4096, + }, + "caca-550M": { + "vocab_size": 32000, "hidden_size": 640, "intermediate_size": 1920, + "num_hidden_layers": 109, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-600M": { + "vocab_size": 32000, "hidden_size": 5120, "intermediate_size": 13632, + "num_hidden_layers": 1, "num_attention_heads": 40, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-650M": { + "vocab_size": 32000, "hidden_size": 768, "intermediate_size": 1984, + "num_hidden_layers": 101, "num_attention_heads": 6, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-700M": { + "vocab_size": 32000, "hidden_size": 3072, "intermediate_size": 8192, + "num_hidden_layers": 5, "num_attention_heads": 24, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-800M": { + "vocab_size": 32000, "hidden_size": 640, "intermediate_size": 1920, + "num_hidden_layers": 157, "num_attention_heads": 5, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-900M": { + "vocab_size": 32000, "hidden_size": 896, "intermediate_size": 2432, + "num_hidden_layers": 93, "num_attention_heads": 7, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-1B": { + "vocab_size": 32000, "hidden_size": 1024, "intermediate_size": 2688, + "num_hidden_layers": 88, "num_attention_heads": 8, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-1.2B": { + "vocab_size": 32000, "hidden_size": 768, "intermediate_size": 2240, + "num_hidden_layers": 176, "num_attention_heads": 6, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-1.5B": { + "vocab_size": 32000, "hidden_size": 1536, "intermediate_size": 4032, + "num_hidden_layers": 53, "num_attention_heads": 12, "num_key_value_heads": 6, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-1.8B": { + "vocab_size": 32000, "hidden_size": 896, "intermediate_size": 2304, + "num_hidden_layers": 211, "num_attention_heads": 7, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-2B": { + "vocab_size": 32000, "hidden_size": 896, "intermediate_size": 2432, + "num_hidden_layers": 232, "num_attention_heads": 7, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-2.5B": { + "vocab_size": 32000, "hidden_size": 1152, "intermediate_size": 3072, + "num_hidden_layers": 175, "num_attention_heads": 9, "num_key_value_heads": 1, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-3B": { + "vocab_size": 32000, "hidden_size": 1280, "intermediate_size": 3584, + "num_hidden_layers": 159, "num_attention_heads": 10, "num_key_value_heads": 2, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-3.5B": { + "vocab_size": 32000, "hidden_size": 2304, "intermediate_size": 6144, + "num_hidden_layers": 58, "num_attention_heads": 18, "num_key_value_heads": 6, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-4B": { + "vocab_size": 32000, "hidden_size": 3072, "intermediate_size": 8192, + "num_hidden_layers": 39, "num_attention_heads": 24, "num_key_value_heads": 4, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-5B": { + "vocab_size": 32000, "hidden_size": 1536, "intermediate_size": 4288, + "num_hidden_layers": 194, "num_attention_heads": 12, "num_key_value_heads": 2, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-6B": { + "vocab_size": 32000, "hidden_size": 1536, "intermediate_size": 4032, + "num_hidden_layers": 245, "num_attention_heads": 12, "num_key_value_heads": 2, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-7B": { + "vocab_size": 32000, "hidden_size": 12288, "intermediate_size": 32576, + "num_hidden_layers": 4, "num_attention_heads": 96, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-8B": { + "vocab_size": 32000, "hidden_size": 3584, "intermediate_size": 9728, + "num_hidden_layers": 58, "num_attention_heads": 28, "num_key_value_heads": 4, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-9B": { + "vocab_size": 32000, "hidden_size": 2560, "intermediate_size": 7040, + "num_hidden_layers": 122, "num_attention_heads": 20, "num_key_value_heads": 5, + "head_dim": 128, "max_position_embeddings": 4096, + }, + "caca-10B": { + "vocab_size": 32000, "hidden_size": 2560, "intermediate_size": 6784, + "num_hidden_layers": 145, "num_attention_heads": 20, "num_key_value_heads": 4, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-13B": { + "vocab_size": 32000, "hidden_size": 2560, "intermediate_size": 6656, + "num_hidden_layers": 192, "num_attention_heads": 20, "num_key_value_heads": 4, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-14B": { + "vocab_size": 32000, "hidden_size": 2304, "intermediate_size": 6336, + "num_hidden_layers": 244, "num_attention_heads": 18, "num_key_value_heads": 3, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-16B": { + "vocab_size": 32000, "hidden_size": 5120, "intermediate_size": 13440, + "num_hidden_layers": 56, "num_attention_heads": 40, "num_key_value_heads": 10, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-18B": { + "vocab_size": 32000, "hidden_size": 6144, "intermediate_size": 16448, + "num_hidden_layers": 45, "num_attention_heads": 48, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-20B": { + "vocab_size": 32000, "hidden_size": 5120, "intermediate_size": 13696, + "num_hidden_layers": 72, "num_attention_heads": 40, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-25B": { + "vocab_size": 32000, "hidden_size": 3584, "intermediate_size": 9728, + "num_hidden_layers": 180, "num_attention_heads": 28, "num_key_value_heads": 7, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-30B": { + "vocab_size": 32000, "hidden_size": 3584, "intermediate_size": 9344, + "num_hidden_layers": 223, "num_attention_heads": 28, "num_key_value_heads": 7, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-35B": { + "vocab_size": 32000, "hidden_size": 4096, "intermediate_size": 11008, + "num_hidden_layers": 196, "num_attention_heads": 32, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-40B": { + "vocab_size": 32000, "hidden_size": 10240, "intermediate_size": 27392, + "num_hidden_layers": 36, "num_attention_heads": 80, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-45B": { + "vocab_size": 32000, "hidden_size": 6144, "intermediate_size": 16448, + "num_hidden_layers": 114, "num_attention_heads": 48, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-50B": { + "vocab_size": 32000, "hidden_size": 16384, "intermediate_size": 43712, + "num_hidden_layers": 18, "num_attention_heads": 128, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-65B": { + "vocab_size": 32000, "hidden_size": 6144, "intermediate_size": 16448, + "num_hidden_layers": 160, "num_attention_heads": 48, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-70B": { + "vocab_size": 32000, "hidden_size": 6144, "intermediate_size": 16320, + "num_hidden_layers": 179, "num_attention_heads": 48, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-80B": { + "vocab_size": 32000, "hidden_size": 12288, "intermediate_size": 32576, + "num_hidden_layers": 51, "num_attention_heads": 96, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-90B": { + "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54400, + "num_hidden_layers": 21, "num_attention_heads": 160, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-100B": { + "vocab_size": 32000, "hidden_size": 10240, "intermediate_size": 27264, + "num_hidden_layers": 93, "num_attention_heads": 80, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-120B": { + "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54656, + "num_hidden_layers": 28, "num_attention_heads": 160, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-150B": { + "vocab_size": 32000, "hidden_size": 8192, "intermediate_size": 22016, + "num_hidden_layers": 216, "num_attention_heads": 64, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-175B": { + "vocab_size": 32000, "hidden_size": 10240, "intermediate_size": 27520, + "num_hidden_layers": 162, "num_attention_heads": 80, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-200B": { + "vocab_size": 32000, "hidden_size": 10240, "intermediate_size": 27264, + "num_hidden_layers": 183, "num_attention_heads": 80, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-250B": { + "vocab_size": 32000, "hidden_size": 12288, "intermediate_size": 32960, + "num_hidden_layers": 159, "num_attention_heads": 96, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-300B": { + "vocab_size": 32000, "hidden_size": 24576, "intermediate_size": 65536, + "num_hidden_layers": 49, "num_attention_heads": 192, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-350B": { + "vocab_size": 32000, "hidden_size": 14336, "intermediate_size": 38272, + "num_hidden_layers": 165, "num_attention_heads": 112, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-400B": { + "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54784, + "num_hidden_layers": 93, "num_attention_heads": 160, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-450B": { + "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54528, + "num_hidden_layers": 105, "num_attention_heads": 160, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-500B": { + "vocab_size": 32000, "hidden_size": 24576, "intermediate_size": 65728, + "num_hidden_layers": 81, "num_attention_heads": 192, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-600B": { + "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54592, + "num_hidden_layers": 140, "num_attention_heads": 160, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-700B": { + "vocab_size": 32000, "hidden_size": 24576, "intermediate_size": 65344, + "num_hidden_layers": 114, "num_attention_heads": 192, "num_key_value_heads": 16, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-800B": { + "vocab_size": 32000, "hidden_size": 24576, "intermediate_size": 65600, + "num_hidden_layers": 131, "num_attention_heads": 192, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-900B": { + "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54656, + "num_hidden_layers": 212, "num_attention_heads": 160, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, + "caca-1T": { + "vocab_size": 32000, "hidden_size": 20480, "intermediate_size": 54528, + "num_hidden_layers": 236, "num_attention_heads": 160, "num_key_value_heads": 8, + "head_dim": 128, "max_position_embeddings": 8192, + }, +} + +# --- MetricsTracker --- +class MetricsTracker: + def __init__(self, reset_interval: int = 100): + self.metrics: Dict[str, List[float]] = defaultdict(list) + self.reset_interval = reset_interval + self.step_count = 0 + + def log(self, name: str, value: Union[float, torch.Tensor]) -> None: + if isinstance(value, torch.Tensor): + value = value.item() + self.metrics[name].append(value) + + def step(self) -> None: + self.step_count += 1 + if self.step_count % self.reset_interval == 0: + self.clear() + + def get_summary(self) -> Dict[str, Dict[str, float]]: + return { + name: { + "mean": float(np.mean(vals)), + "std": float(np.std(vals)), + "min": float(np.min(vals)), + "max": float(np.max(vals)), + "last": float(vals[-1]), + } + for name, vals in self.metrics.items() if vals + } + + def clear(self) -> None: + self.metrics.clear() + +# --- KV Cache --- +class DynamicCache: + def __init__(self): + self.key_cache: List[torch.Tensor] = [] + self.value_cache: List[torch.Tensor] = [] + self._seen_tokens: int = 0 + + def __len__(self): + return len(self.key_cache) + + def get_seq_length(self, layer_idx: int = 0) -> int: + return self.key_cache[layer_idx].shape[-2] if len(self.key_cache) > layer_idx else 0 + + def update( + self, + key_states: torch.Tensor, + value_states: torch.Tensor, + layer_idx: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + if layer_idx == 0: + self._seen_tokens += key_states.shape[-2] + if len(self.key_cache) == layer_idx: + self.key_cache.append(key_states) + self.value_cache.append(value_states) + else: + self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) + self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) + return self.key_cache[layer_idx], self.value_cache[layer_idx] + + def reorder_cache(self, beam_idx: torch.Tensor) -> None: + for i in range(len(self.key_cache)): + dev = self.key_cache[i].device + self.key_cache[i] = self.key_cache[i].index_select(0, beam_idx.to(dev)) + self.value_cache[i] = self.value_cache[i].index_select(0, beam_idx.to(dev)) + + def crop(self, max_length: int) -> None: + + self._seen_tokens = max(0, max_length) + for i in range(len(self.key_cache)): + self.key_cache[i] = self.key_cache[i][..., :max_length, :] + self.value_cache[i] = self.value_cache[i][..., :max_length, :] + + def to_legacy_cache(self) -> Tuple: + return tuple(zip(self.key_cache, self.value_cache)) + + @classmethod + def from_legacy_cache(cls, past: Optional[Tuple]) -> "DynamicCache": + cache = cls() + if past: + for i, (k, v) in enumerate(past): + cache.key_cache.append(k) + cache.value_cache.append(v) + if i == 0: + cache._seen_tokens = k.shape[-2] + return cache + +class SlidingWindowCache(DynamicCache): + def __init__(self, window_size: int, sink_size: int = 4): + super().__init__() + self.window_size = window_size + self.sink_size = sink_size + + def update(self, key_states, value_states, layer_idx): + k, v = super().update(key_states, value_states, layer_idx) + max_len = self.window_size + self.sink_size + if k.shape[-2] > max_len: + sink_k = k[..., :self.sink_size, :] + sink_v = v[..., :self.sink_size, :] + win_k = k[..., -self.window_size:, :] + win_v = v[..., -self.window_size:, :] + self.key_cache[layer_idx] = torch.cat([sink_k, win_k], dim=-2) + self.value_cache[layer_idx] = torch.cat([sink_v, win_v], dim=-2) + + if layer_idx == 0: + self._seen_tokens = self.key_cache[layer_idx].shape[-2] + return self.key_cache[layer_idx], self.value_cache[layer_idx] + +# --- Normalisasi --- +class CacaRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + var = x.pow(2).mean(-1, keepdim=True) + x = x * torch.rsqrt(var + self.variance_epsilon) + return (self.weight * x).to(dtype) + +class LayerScale(nn.Module): + def __init__(self, dim: int, init_value: float = 1e-5): + super().__init__() + self.gamma = nn.Parameter(init_value * torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.gamma * x + +class StochasticDepth(nn.Module): + def __init__(self, drop_prob: float = 0.0): + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor, training: bool = True) -> torch.Tensor: + if not training or self.drop_prob == 0.0: + return x + keep = 1 - self.drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + mask = keep + torch.rand(shape, dtype=x.dtype, device=x.device) + return x.div(keep) * mask.floor_() + +# --- Positional Encoding --- +def rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + +def apply_rotary_pos_emb(q, k, cos, sin): + cos, sin = cos.to(q.dtype), sin.to(q.dtype) + return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) + +class CacaRotaryEmbedding(nn.Module): + def __init__( + self, + dim: int, + max_position_embeddings: int = 8192, + base: float = 10000.0, + scaling_factor: float = 1.0, + scaling_type: Optional[str] = None, + rope_type: str = "default", + ntk_alpha: float = 1.0, + ): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.scaling_factor = scaling_factor + self.scaling_type = scaling_type + self.rope_type = rope_type + self.ntk_alpha = ntk_alpha + + inv_freq = self._build_inv_freq() + self.register_buffer("inv_freq", inv_freq, persistent=False) + + self._cos_cache: Optional[torch.Tensor] = None + self._sin_cache: Optional[torch.Tensor] = None + self._cached_seq_len: int = 0 + + def _build_inv_freq(self) -> torch.Tensor: + freqs = torch.arange(0, self.dim, 2).float() / self.dim + inv_freq = 1.0 / (self.base ** freqs) + + if self.rope_type == "ntk": + effective_base = self.base * (self.ntk_alpha ** (self.dim / (self.dim - 2))) + inv_freq = 1.0 / (effective_base ** freqs) + elif self.rope_type in ("linear", "yarn"): + inv_freq = inv_freq / self.scaling_factor + # dynamic: ditangani saat forward + + return inv_freq + + def forward(self, x: torch.Tensor, seq_len: int, position_offset: int = 0): + if ( + self._cos_cache is not None + and self._cached_seq_len >= seq_len + position_offset + and position_offset == 0 + ): + return ( + self._cos_cache[:, :, :seq_len, :].to(x.dtype), + self._sin_cache[:, :, :seq_len, :].to(x.dtype), + ) + + t = torch.arange(position_offset, position_offset + seq_len, device=x.device).float() + + if self.rope_type == "dynamic" and seq_len > self.max_position_embeddings: + scale = seq_len / self.max_position_embeddings + t = t / scale + + freqs = torch.outer(t, self.inv_freq) + emb = torch.cat([freqs, freqs], dim=-1) + cos = emb.cos()[None, None] + sin = emb.sin()[None, None] + + if position_offset == 0: + self._cos_cache = cos + self._sin_cache = sin + self._cached_seq_len = seq_len + + return cos.to(x.dtype), sin.to(x.dtype) + +class ALiBiPositionalBias(nn.Module): + def __init__(self, num_heads: int, max_positions: int = 8192): + super().__init__() + slopes = torch.tensor(self._get_slopes(num_heads)) + self.register_buffer("slopes", slopes, persistent=False) + + @staticmethod + def _get_slopes(n: int) -> List[float]: + def _pow2(n): + s = 2 ** (-(2 ** -(math.log2(n) - 3))) + return [s * (s ** i) for i in range(n)] + if math.log2(n).is_integer(): + return _pow2(n) + p = 2 ** math.floor(math.log2(n)) + return _pow2(p) + ALiBiPositionalBias._get_slopes(2 * p)[::2][: n - p] + + def forward(self, seq_len: int, key_len: Optional[int] = None) -> torch.Tensor: + kl = key_len or seq_len + qp = torch.arange(seq_len, device=self.slopes.device).unsqueeze(1) + kp = torch.arange(kl, device=self.slopes.device).unsqueeze(0) + bias = (kp - qp).unsqueeze(0) * self.slopes.unsqueeze(1).unsqueeze(2) + return bias.unsqueeze(0) + +# --- Logit softcap --- +def soft_cap_logits(x: torch.Tensor, cap: Optional[float]) -> torch.Tensor: + if cap is None or cap <= 0: + return x + return cap * torch.tanh(x / cap) + +# --- MoE Router --- +class TopKRouter(nn.Module): + def __init__(self, hidden_size: int, num_experts: int, num_experts_per_tok: int): + super().__init__() + self.num_experts = num_experts + self.num_experts_per_tok = num_experts_per_tok + self.gate = nn.Linear(hidden_size, num_experts, bias=False) + self.gate_norm = nn.LayerNorm(hidden_size) + self.temperature = nn.Parameter(torch.ones(1)) + self.jitter_noise = 0.01 + self.register_buffer("_load_history", torch.zeros(num_experts)) + self._load_steps = 0 + + def forward(self, hidden_states: torch.Tensor): + flat = self.gate_norm(hidden_states.view(-1, hidden_states.shape[-1])) + logits = self.gate(flat) + logits = torch.clamp(logits, -20, 20) + temp = F.softplus(self.temperature).clamp(min=0.1, max=10.0) + logits = logits / temp + + if self.training and self.jitter_noise > 0: + logits = logits + torch.randn_like(logits) * self.jitter_noise + + probs = F.softmax(logits, dim=-1, dtype=torch.float32) + top_w, top_idx = torch.topk(probs, self.num_experts_per_tok, dim=-1) + top_w = top_w / (top_w.sum(-1, keepdim=True) + 1e-9) + + usage = probs.mean(0) + aux_loss = usage.std() / (usage.mean() + 1e-10) + z_loss = torch.logsumexp(logits.float(), -1).pow(2).mean() + ent_loss = -(probs * (probs + 1e-10).log()).sum(-1).mean() * (-0.01) + + if self.training: + self._load_history = 0.9 * self._load_history + 0.1 * usage.detach() + self._load_steps += 1 + + return top_w, top_idx, aux_loss + ent_loss, z_loss + + def get_load_imbalance(self) -> float: + if self._load_steps == 0: + return 0.0 + m = self._load_history.mean() + return (self._load_history.std() / (m + 1e-10)).item() + +class ExpertChoiceRouter(nn.Module): + def __init__(self, hidden_size: int, num_experts: int, expert_choice_k: float): + super().__init__() + self.num_experts = num_experts + self.expert_choice_k = expert_choice_k + self.gate = nn.Linear(hidden_size, num_experts, bias=False) + + def forward(self, hidden_states: torch.Tensor): + bs, sl, hs = hidden_states.shape + flat = hidden_states.view(-1, hs) + logits = self.gate(flat) + probs = F.softmax(logits, -1, dtype=torch.float32) + cap = max(1, int(self.expert_choice_k * bs * sl / self.num_experts)) + _, top_idx = torch.topk(probs.t(), min(cap, bs * sl), dim=-1) + mask = torch.zeros(self.num_experts, bs * sl, device=flat.device) + for e in range(self.num_experts): + mask[e, top_idx[e]] = 1.0 + weights = mask.t() * probs + aux_loss = (probs.mean(0) ** 2).sum() * self.num_experts + z_loss = torch.logsumexp(logits, -1).mean() + return weights, aux_loss, z_loss + +# --- Expert & MoE --- +class Expert(nn.Module): + def __init__(self, config: CacaConfig): + super().__init__() + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias) + self.dropout = nn.Dropout(config.hidden_dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.dropout(F.silu(self.gate_proj(x)) * self.up_proj(x))) + +class MixtureOfExperts(nn.Module): + def __init__(self, config: CacaConfig): + super().__init__() + self.config = config + self.num_experts = config.num_experts + self.use_expert_choice = config.use_expert_choice + self.load_monitoring = config.expert_load_monitoring + + self.experts = nn.ModuleList([Expert(config) for _ in range(config.num_experts)]) + self.router = ( + ExpertChoiceRouter(config.hidden_size, config.num_experts, config.expert_choice_k) + if config.use_expert_choice + else TopKRouter(config.hidden_size, config.num_experts, config.num_experts_per_tok) + ) + self.register_buffer("expert_usage_count", torch.zeros(config.num_experts, dtype=torch.long)) + + def _safe_expert(self, expert_fn, inp: torch.Tensor, expert_idx: int, config: CacaConfig): + try: + out = expert_fn(inp) + if torch.isnan(out).any() or torch.isinf(out).any(): + level = config.nan_recovery_level + if level >= 1: + logger.warning(f"Expert {expert_idx}: output NaN/Inf — di-zero.") + if level == 0: + return out + return torch.zeros_like(out) + return out + except RuntimeError as e: + logger.error(f"Expert {expert_idx} RuntimeError: {e}") + return torch.zeros(inp.shape[0], config.hidden_size, device=inp.device, dtype=inp.dtype) + + def forward(self, hidden_states: torch.Tensor): + bs, sl, hs = hidden_states.shape + flat = hidden_states.view(-1, hs) + + if torch.isnan(flat).any() or torch.isinf(flat).any(): + logger.error("MoE: NaN/Inf pada input. Mengembalikan input tanpa perubahan.") + zero = torch.tensor(0.0, device=flat.device) + return hidden_states, zero, zero + + if self.use_expert_choice: + weights, aux_loss, z_loss = self.router(hidden_states) + out = torch.zeros_like(flat) + for i, expert in enumerate(self.experts): + mask = weights[:, i] > 1e-6 + if mask.any(): + if not self.training: + self.expert_usage_count[i] += mask.sum() + expert_out = self._safe_expert(expert, flat[mask], i, self.config) + out[mask] += expert_out * weights[mask, i:i+1] + else: + top_w, top_idx, aux_loss, z_loss = self.router(hidden_states) + out = torch.zeros_like(flat) + for i in range(self.num_experts): + mask = (top_idx == i).any(-1) + if not mask.any(): + continue + if not self.training: + self.expert_usage_count[i] += mask.sum() + expert_out = self._safe_expert(self.experts[i], flat[mask], i, self.config) + token_ids = torch.where(mask)[0] + for j, tid in enumerate(token_ids): + positions = (top_idx[tid] == i).nonzero(as_tuple=True)[0] + if positions.numel(): + out[tid] += expert_out[j] * top_w[tid, positions[0]] + + if self.load_monitoring and self.training and isinstance(self.router, TopKRouter): + imbalance = self.router.get_load_imbalance() + if imbalance > self.config.expert_load_warn_threshold: + logger.warning( + f"MoE load imbalance CV={imbalance:.3f} " + f"(threshold={self.config.expert_load_warn_threshold})" + ) + + final = out.view(bs, sl, hs) + if torch.isnan(final).any() or torch.isinf(final).any(): + logger.error("MoE: NaN/Inf pada output. Mengembalikan input.") + return hidden_states, aux_loss, z_loss + return final, aux_loss, z_loss + +# --- MLP standar --- +class CacaMLP(nn.Module): + def __init__(self, config: CacaConfig): + super().__init__() + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.mlp_bias) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.mlp_bias) + self.dropout = nn.Dropout(config.hidden_dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.dropout(F.silu(self.gate_proj(x)) * self.up_proj(x))) + +# --- Attention --- +class CacaAttention(nn.Module): + def __init__(self, config: CacaConfig, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.sliding_window = config.sliding_window + self.attn_logit_softcapping = config.attn_logit_softcapping + self.temperature = config.attention_temperature + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) + + if config.use_qk_norm: + self.q_norm = CacaRMSNorm(self.head_dim, config.qk_norm_eps) + self.k_norm = CacaRMSNorm(self.head_dim, config.qk_norm_eps) + else: + self.q_norm = self.k_norm = None + + if config.use_rotary_embeddings: + sf, st = 1.0, None + if config.rope_scaling: + st = config.rope_scaling.get("type", "linear") + sf = config.rope_scaling.get("factor", 1.0) + self.rotary_emb = CacaRotaryEmbedding( + self.head_dim, config.max_position_embeddings, config.rope_theta, + scaling_factor=sf, scaling_type=st, + rope_type=config.rope_type, ntk_alpha=config.rope_ntk_alpha, + ) + else: + self.rotary_emb = None + + self.alibi = ALiBiPositionalBias(self.num_heads) if config.use_alibi else None + self.attn_drop = nn.Dropout(config.attention_dropout) + + self.is_global = self._global_layer(config, layer_idx) + self.has_flash = HAS_FLASH_ATTN and config.use_flash_attn + self.has_xfmrs = HAS_XFORMERS + self.has_sdpa = HAS_SDPA + + self._mask_cache: OrderedDict = OrderedDict() + self._max_cache = 10 + self._cache_hits = self._cache_misses = 0 + + def _global_layer(self, cfg: CacaConfig, idx: Optional[int]) -> bool: + if idx is None: return False + if cfg.attention_pattern == "all_global": return True + if cfg.attention_pattern == "all_local": return False + if cfg.attention_pattern == "interleaved": + return (idx % cfg.global_attention_every_n_layers) == cfg.global_attention_every_n_layers - 1 + return False + + def _causal_mask(self, ql, kl, dtype, device, use_sw): + key = (ql, kl, str(dtype), device.type, use_sw, self.sliding_window if use_sw else None) + if key in self._mask_cache: + self._cache_hits += 1 + self._mask_cache.move_to_end(key) + return self._mask_cache[key].to(dtype=dtype, device=device) + self._cache_misses += 1 + + kl_eff = max(kl, ql) + qp = torch.arange(ql, device=device) + (kl_eff - ql) + kp = torch.arange(kl_eff, device=device) + d = qp[:, None] - kp[None, :] + mask = d < 0 + + if use_sw and self.sliding_window: + if self.config.use_attention_sink and self.config.attention_sink_size > 0: + sink = kp[None, :] < self.config.attention_sink_size + win = (d >= 0) & (d <= self.sliding_window) + mask = (d < 0) | (~sink & ~win) + else: + mask = mask | (d > self.sliding_window) + + float_mask = torch.zeros(1, 1, ql, kl_eff, dtype=dtype, device=device) + float_mask.masked_fill_(mask[None, None], -1e9) + + if len(self._mask_cache) >= self._max_cache: + self._mask_cache.popitem(last=False) + self._mask_cache[key] = float_mask.detach().cpu() + return float_mask + + def get_cache_stats(self) -> Dict: + total = self._cache_hits + self._cache_misses + return { + "hits": self._cache_hits, + "misses": self._cache_misses, + "hit_rate": self._cache_hits / total if total else 0, + "cache_size": len(self._mask_cache), + } + + def forward(self, hidden_states, attention_mask=None, past_key_value=None, use_cache=False): + if hidden_states is None: + raise ValueError("hidden_states tidak boleh None") + if hidden_states.shape[-1] != self.hidden_size: + raise ValueError(f"Ekspektasi hidden_size {self.hidden_size}, dapat {hidden_states.shape[-1]}") + + bs, sl, _ = hidden_states.shape + + q = self.q_proj(hidden_states).view(bs, sl, self.num_heads, self.head_dim).transpose(1, 2) + k = self.k_proj(hidden_states).view(bs, sl, self.num_kv_heads, self.head_dim).transpose(1, 2) + v = self.v_proj(hidden_states).view(bs, sl, self.num_kv_heads, self.head_dim).transpose(1, 2) + + if self.q_norm: + q, k = self.q_norm(q), self.k_norm(k) + + pos_off = 0 + if past_key_value is not None: + try: + if isinstance(past_key_value, (tuple, list)) and past_key_value[0] is not None: + pos_off = past_key_value[0].shape[2] + except Exception: + pos_off = 0 + + if self.rotary_emb: + cos, sin = self.rotary_emb(q, sl, pos_off) + q, k = apply_rotary_pos_emb(q, k, cos, sin) + + if past_key_value is not None and past_key_value[0] is not None: + try: + if past_key_value[0].numel() > 0: + k = torch.cat([past_key_value[0], k], dim=2) + v = torch.cat([past_key_value[1], v], dim=2) + except RuntimeError as e: + logger.error(f"KV concat gagal: {e}") + + present = (k, v) if use_cache else None + + k = k.repeat_interleave(self.num_kv_groups, dim=1) + v = v.repeat_interleave(self.num_kv_groups, dim=1) + + kv_len = k.shape[-2] + use_sw = not self.is_global and self.sliding_window is not None + + if self.has_flash and attention_mask is None and q.device.type == "cuda" and q.dtype in (torch.float16, torch.bfloat16): + try: + out = self._flash_attn(q, k, v, use_sw, kv_len) + except Exception as e: + logger.warning(f"Flash attention gagal ({e}), fallback.") + out = self._fallback(q, k, v, attention_mask, kv_len, use_sw) + else: + out = self._fallback(q, k, v, attention_mask, kv_len, use_sw) + + return self.o_proj(out), present + + def _flash_attn(self, q, k, v, use_sw, kv_len): + bs, _, sl, _ = q.shape + orig_dtype = q.dtype + compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + q = q.transpose(1, 2).contiguous().to(compute_dtype) + k = k.transpose(1, 2).contiguous().to(compute_dtype) + v = v.transpose(1, 2).contiguous().to(compute_dtype) + ws = (self.sliding_window, 0) if use_sw and self.sliding_window < kv_len else (-1, 0) + out = flash_attn_func( + q, k, v, + dropout_p=self.config.attention_dropout if self.training else 0.0, + causal=True, window_size=ws, + ) + return out.to(orig_dtype).reshape(bs, sl, self.hidden_size) + + def _fallback(self, q, k, v, mask, kv_len, use_sw): + if self.has_xfmrs and q.device.type == "cuda" and mask is None: + try: + return self._xformers(q, k, v, kv_len, use_sw) + except Exception: + pass + if self.has_sdpa: + return self._sdpa(q, k, v, mask, kv_len, use_sw) + return self._standard(q, k, v, mask, kv_len, use_sw) + + def _sdpa(self, q, k, v, mask, kv_len, use_sw): + bs, _, sl, _ = q.shape + if mask is None: + mask = self._causal_mask(sl, kv_len, q.dtype, q.device, use_sw) + if self.alibi: + mask = mask + self.alibi(sl, kv_len) + # scale = 1/sqrt(d) lalu dibagi temperature, konsisten dengan standar attention + scale = (1.0 / math.sqrt(self.head_dim)) / self.temperature + out = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, + dropout_p=self.config.attention_dropout if self.training else 0.0, + scale=scale, is_causal=False, + ) + return out.transpose(1, 2).contiguous().reshape(bs, sl, self.hidden_size) + + def _xformers(self, q, k, v, kv_len, use_sw): + bs, _, sl, _ = q.shape + bias = self._causal_mask(sl, kv_len, q.dtype, q.device, use_sw) + out = memory_efficient_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), + attn_bias=bias, + p=self.config.attention_dropout if self.training else 0.0, + ) + return out.reshape(bs, sl, self.hidden_size) + + def _standard(self, q, k, v, mask, kv_len, use_sw): + bs, _, sl, _ = q.shape + + scale = math.sqrt(self.head_dim) * self.temperature + w = torch.matmul(q, k.transpose(-2, -1)) / scale + w = torch.clamp(w, -50, 50) + w = soft_cap_logits(w, self.attn_logit_softcapping) + if mask is None: + mask = self._causal_mask(sl, kv_len, w.dtype, w.device, use_sw) + if self.alibi: + mask = mask + self.alibi(sl, kv_len) + w = F.softmax(w + mask, -1, dtype=torch.float32).to(q.dtype) + w = self.attn_drop(w) + out = torch.matmul(w, v) + return out.transpose(1, 2).contiguous().reshape(bs, sl, self.hidden_size) + +# --- Cross-Attention --- +class CacaCrossAttention(nn.Module): + def __init__(self, config: CacaConfig): + super().__init__() + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = config.head_dim + self.num_kv_groups = self.num_heads // self.num_kv_heads + + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=config.attention_bias) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias) + self.drop = nn.Dropout(config.attention_dropout) + + def forward(self, hidden_states, encoder_hidden_states, attention_mask=None): + bs, sl, _ = hidden_states.shape + esl = encoder_hidden_states.shape[1] + + q = self.q_proj(hidden_states).view(bs, sl, self.num_heads, self.head_dim).transpose(1, 2) + k = self.k_proj(encoder_hidden_states).view(bs, esl, self.num_kv_heads, self.head_dim).transpose(1, 2) + v = self.v_proj(encoder_hidden_states).view(bs, esl, self.num_kv_heads, self.head_dim).transpose(1, 2) + k = k.repeat_interleave(self.num_kv_groups, dim=1) + v = v.repeat_interleave(self.num_kv_groups, dim=1) + + w = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim) + if attention_mask is not None: + w = w + attention_mask + w = F.softmax(w, -1, dtype=torch.float32).to(q.dtype) + w = self.drop(w) + out = torch.matmul(w, v).transpose(1, 2).contiguous().reshape(bs, sl, self.hidden_size) + return self.o_proj(out) + +# --- MoD Router --- +class MixtureOfDepthsRouter(nn.Module): + def __init__(self, hidden_size: int, capacity_factor: float = 0.5, route_method: str = "learned"): + super().__init__() + self.capacity_factor = capacity_factor + self.route_method = route_method + if route_method == "learned": + self.router = nn.Linear(hidden_size, 1) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + bs, sl, _ = hidden_states.shape + if self.route_method == "learned": + logits = self.router(hidden_states).squeeze(-1) + elif self.route_method == "random": + logits = torch.rand(bs, sl, device=hidden_states.device) + else: + logits = torch.zeros(bs, sl, device=hidden_states.device) + cap = max(1, int(sl * self.capacity_factor)) + _, idx = torch.topk(logits, cap, dim=-1) + mask = torch.zeros(bs, sl, dtype=torch.bool, device=hidden_states.device) + mask.scatter_(1, idx, True) + return mask + +# --- Decoder Layer --- +class CacaDecoderLayer(nn.Module): + def __init__(self, config: CacaConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.self_attn = CacaAttention(config, layer_idx) + self.use_moe = config.use_moe and (layer_idx % config.moe_layer_frequency == 0) + self.mlp = MixtureOfExperts(config) if self.use_moe else CacaMLP(config) + self.use_cross = config.use_cross_attention and (layer_idx % config.cross_attention_frequency == 0) + + if self.use_cross: + self.cross_attn = CacaCrossAttention(config) + self.cross_attn_norm = CacaRMSNorm(config.hidden_size, config.rms_norm_eps) + + self.input_layernorm = CacaRMSNorm(config.hidden_size, config.rms_norm_eps) + self.post_attention_layernorm = CacaRMSNorm(config.hidden_size, config.rms_norm_eps) + self.residual_drop = nn.Dropout(config.residual_dropout) + + self.ls1 = LayerScale(config.hidden_size, config.layer_scale_init) if config.use_layer_scale else None + self.ls2 = LayerScale(config.hidden_size, config.layer_scale_init) if config.use_layer_scale else None + self.ls_cross = LayerScale(config.hidden_size, config.layer_scale_init) if (config.use_layer_scale and self.use_cross) else None + + if config.use_stochastic_depth: + dp = config.stochastic_depth_prob * layer_idx / config.num_hidden_layers + self.stoch_depth = StochasticDepth(dp) + else: + self.stoch_depth = None + + self.mod_router = ( + MixtureOfDepthsRouter(config.hidden_size, config.mod_capacity_factor, config.mod_route_method) + if config.use_mixture_of_depths else None + ) + + self._grad_stats = {"max_norm": 0.0, "ema_norm": 0.0, "clip_count": 0} + + def _grad_hook(self, grad: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + if grad is None: + return grad + n = grad.norm().item() + self._grad_stats["max_norm"] = max(self._grad_stats["max_norm"], n) + self._grad_stats["ema_norm"] = 0.9 * self._grad_stats["ema_norm"] + 0.1 * n + if n > 10.0: + self._grad_stats["clip_count"] += 1 + if self._grad_stats["clip_count"] % 100 == 0: + logger.warning( + f"Layer {self.layer_idx}: grad norm {n:.2f} " + f"(clipped {self._grad_stats['clip_count']}x)" + ) + return torch.clamp(grad, -10.0, 10.0) + return grad + + def _recover_nan(self, x: torch.Tensor, name: str) -> torch.Tensor: + level = self.config.nan_recovery_level + if level == 0: + return x + if not (torch.isnan(x).any() or torch.isinf(x).any()): + return x + logger.warning(f"Layer {self.layer_idx}: NaN/Inf pada {name}") + if level >= 1: + return torch.nan_to_num(x, nan=0.0, posinf=1e4, neginf=-1e4) + return x + + def forward( + self, hidden_states, attention_mask=None, + encoder_hidden_states=None, encoder_attention_mask=None, + past_key_value=None, use_cache=False, + ): + if hidden_states is None: + raise ValueError("hidden_states tidak boleh None") + + hidden_states = self._recover_nan(hidden_states, "layer_input") + + if self.training and hidden_states.requires_grad: + hidden_states.register_hook(self._grad_hook) + + aux_loss = z_loss = 0.0 + + if self.mod_router is not None: + proc_mask = self.mod_router(hidden_states) + tokens = hidden_states[proc_mask] + if tokens.numel() == 0: + return hidden_states, past_key_value if use_cache else None, aux_loss, z_loss + else: + proc_mask = None + tokens = hidden_states + + # Self-attention + residual = tokens + attn_out, present_kv = self.self_attn( + self.input_layernorm(tokens), attention_mask, past_key_value, use_cache + ) + attn_out = self._recover_nan(attn_out, "attn_output") + if self.ls1: attn_out = self.ls1(attn_out) + if self.stoch_depth: attn_out = self.stoch_depth(attn_out, self.training) + tokens = residual + self.residual_drop(attn_out) + if self.training: tokens = torch.clamp(tokens, -1e4, 1e4) + + # Cross-attention + if self.use_cross and encoder_hidden_states is not None: + residual = tokens + cross_out = self.cross_attn(self.cross_attn_norm(tokens), encoder_hidden_states, encoder_attention_mask) + cross_out = self._recover_nan(cross_out, "cross_attn") + if self.ls_cross: cross_out = self.ls_cross(cross_out) + if self.stoch_depth: cross_out = self.stoch_depth(cross_out, self.training) + tokens = residual + self.residual_drop(cross_out) + if self.training: tokens = torch.clamp(tokens, -1e4, 1e4) + + # MLP / MoE + residual = tokens + if self.use_moe: + mlp_out, aux_loss, z_loss = self.mlp(self.post_attention_layernorm(tokens)) + else: + mlp_out = self.mlp(self.post_attention_layernorm(tokens)) + mlp_out = self._recover_nan(mlp_out, "mlp_output") + if self.ls2: mlp_out = self.ls2(mlp_out) + if self.stoch_depth: mlp_out = self.stoch_depth(mlp_out, self.training) + tokens = residual + self.residual_drop(mlp_out) + if self.training: tokens = torch.clamp(tokens, -1e4, 1e4) + + if proc_mask is not None: + hidden_states[proc_mask] = tokens + else: + hidden_states = tokens + + return hidden_states, present_kv, aux_loss, z_loss + + def get_gradient_stats(self) -> Dict: + return self._grad_stats.copy() + +# --- PreTrainedModel base --- +class CacaPreTrainedModel(PreTrainedModel): + config_class = CacaConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["CacaDecoderLayer"] + # _skip_keys_device_placement = "past_key_values" + + def _init_weights(self, module): + std = self.config.initializer_range + + if self.config.use_mup and isinstance(module, nn.Linear): + for name, m in self.named_modules(): + if m is module and "o_proj" in name: + std = self.config.initializer_range / math.sqrt(2 * self.config.num_hidden_layers) + break + + if isinstance(module, nn.Linear): + module.weight.data.normal_(0.0, std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(0.0, std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + def _set_gradient_checkpointing(self, module, value=False): + if isinstance(module, CacaModel): + module.gradient_checkpointing = value + + def apply_spectral_norm(self): + if not self.config.use_spectral_norm: + return + for name, module in self.named_modules(): + if isinstance(module, nn.Linear): + try: + nn.utils.spectral_norm(module) + except Exception as e: + logger.warning(f"SpectralNorm gagal pada {name}: {e}") + logger.info("Spectral norm diterapkan ke semua Linear layer.") + +# --- Vision Encoder --- +class VisionTransformerBlock(nn.Module): + def __init__(self, dim, num_heads, mlp_ratio=4.0, dropout=0.0, eps=1e-6): + super().__init__() + self.norm1 = nn.LayerNorm(dim, eps=eps) + self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True) + self.drop1 = nn.Dropout(dropout) + self.norm2 = nn.LayerNorm(dim, eps=eps) + mid = int(dim * mlp_ratio) + self.mlp = nn.Sequential( + nn.Linear(dim, mid), nn.GELU(), nn.Dropout(dropout), + nn.Linear(mid, dim), nn.Dropout(dropout), + ) + self.drop2 = nn.Dropout(dropout) + + def forward(self, x): + + nx = self.norm1(x) + x = x + self.drop1(self.attn(nx, nx, nx, need_weights=False)[0]) + return x + self.drop2(self.mlp(self.norm2(x))) + +class VisionEncoder(nn.Module): + def __init__(self, config: CacaConfig): + super().__init__() + vc = config.vision_config + self.patch_size = vc.get("patch_size", 14) + self.image_size = vc.get("image_size", 224) + self.num_channels = vc.get("num_channels", 3) + self.hidden_size = vc.get("hidden_size", 1024) + self.num_patches = (self.image_size // self.patch_size) ** 2 + self.patch_embed = nn.Conv2d(self.num_channels, self.hidden_size, self.patch_size, self.patch_size, bias=False) + self.cls_token = nn.Parameter(torch.zeros(1, 1, self.hidden_size)) + self.pos_embed = nn.Parameter(torch.zeros(1, self.num_patches + 1, self.hidden_size)) + self.drop = nn.Dropout(vc.get("dropout", 0.0)) + self.blocks = nn.ModuleList([ + VisionTransformerBlock( + self.hidden_size, vc.get("num_heads", 16), + vc.get("intermediate_size", 4096) / self.hidden_size, + vc.get("dropout", 0.0), vc.get("layer_norm_eps", 1e-6), + ) + for _ in range(vc.get("num_layers", 24)) + ]) + self.norm = nn.LayerNorm(self.hidden_size, eps=vc.get("layer_norm_eps", 1e-6)) + nn.init.trunc_normal_(self.pos_embed, std=0.02) + nn.init.trunc_normal_(self.cls_token, std=0.02) + nn.init.trunc_normal_(self.patch_embed.weight, std=0.02) + + def forward(self, pixel_values): + bs = pixel_values.shape[0] + x = self.patch_embed(pixel_values).flatten(2).transpose(1, 2) + x = torch.cat([self.cls_token.expand(bs, -1, -1), x], dim=1) + self.pos_embed + x = self.drop(x) + for blk in self.blocks: + x = blk(x) + return self.norm(x) + +# --- Audio Encoder --- +class AudioTransformerBlock(nn.Module): + def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0, dropout: float = 0.0): + super().__init__() + self.norm1 = nn.LayerNorm(dim) + self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True) + self.drop1 = nn.Dropout(dropout) + self.norm2 = nn.LayerNorm(dim) + mid = int(dim * mlp_ratio) + self.mlp = nn.Sequential( + nn.Linear(dim, mid), nn.GELU(), nn.Dropout(dropout), + nn.Linear(mid, dim), nn.Dropout(dropout), + ) + self.drop2 = nn.Dropout(dropout) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + nx = self.norm1(x) + x = x + self.drop1(self.attn(nx, nx, nx, need_weights=False)[0]) + return x + self.drop2(self.mlp(self.norm2(x))) + +class AudioEncoder(nn.Module): + def __init__(self, config: "CacaConfig"): + super().__init__() + ac = config.audio_config + self.num_mel_bins = ac.get("num_mel_bins", 80) + self.hidden_size = ac.get("hidden_size", 1024) + self.max_audio_length = ac.get("max_audio_length", 3000) + dropout = ac.get("dropout", 0.0) + + self.conv1 = nn.Sequential( + nn.Conv1d(self.num_mel_bins, self.hidden_size, kernel_size=3, padding=1), + nn.GELU(), nn.Dropout(dropout), + ) + self.conv2 = nn.Sequential( + nn.Conv1d(self.hidden_size, self.hidden_size, kernel_size=3, stride=2, padding=1), + nn.GELU(), nn.Dropout(dropout), + ) + self.pos_embed = nn.Parameter(torch.zeros(1, self.max_audio_length // 2, self.hidden_size)) + self.drop = nn.Dropout(dropout) + self.blocks = nn.ModuleList([ + AudioTransformerBlock( + self.hidden_size, ac.get("num_heads", 16), + ac.get("intermediate_size", 4096) / self.hidden_size, dropout, + ) + for _ in range(ac.get("num_layers", 12)) + ]) + self.norm = nn.LayerNorm(self.hidden_size) + nn.init.trunc_normal_(self.pos_embed, std=0.02) + + def forward(self, audio_features: torch.Tensor) -> torch.Tensor: + x = F.gelu(self.conv1(audio_features)) + x = F.gelu(self.conv2(x)) + x = x.transpose(1, 2) + seq_len = x.shape[1] + + if seq_len <= self.pos_embed.shape[1]: + x = x + self.pos_embed[:, :seq_len, :] + else: + pos_interp = F.interpolate( + self.pos_embed.transpose(1, 2), size=seq_len, + mode="linear", align_corners=False, + ).transpose(1, 2) + x = x + pos_interp + + x = self.drop(x) + for blk in self.blocks: + x = blk(x) + return self.norm(x) + +class MultiModalProjector(nn.Module): + def __init__(self, in_size: int, out_size: int, proj_type: str = "mlp", num_layers: int = 2): + super().__init__() + if proj_type == "linear": + self.proj = nn.Linear(in_size, out_size) + else: + layers, cur = [], in_size + for _ in range(num_layers - 1): + layers += [nn.Linear(cur, out_size), nn.GELU(), nn.Dropout(0.1)] + cur = out_size + layers.append(nn.Linear(cur, out_size)) + self.proj = nn.Sequential(*layers) + + def forward(self, x): + return self.proj(x) + +# --- CacaModel --- +class CacaModel(CacaPreTrainedModel): + def __init__(self, config: CacaConfig): + super().__init__(config) + self.config = config + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.token_drop = nn.Dropout(config.token_dropout) if config.token_dropout > 0 else None + self.layers = nn.ModuleList([CacaDecoderLayer(config, i) for i in range(config.num_hidden_layers)]) + self.norm = CacaRMSNorm(config.hidden_size, config.rms_norm_eps) + self.gradient_checkpointing = False + self.metrics = MetricsTracker() + self._last_mem_check = 0 + + if config.use_multimodal: + if config.vision_config: + self.vision_encoder = VisionEncoder(config) + self.vision_projector = MultiModalProjector( + config.vision_config.get("hidden_size", 768), config.hidden_size, + config.vision_config.get("projector_type", "mlp"), + ) + else: + self.vision_encoder = self.vision_projector = None + + if config.audio_config: + self.audio_encoder = AudioEncoder(config) + self.audio_projector = MultiModalProjector( + config.audio_config.get("hidden_size", 768), config.hidden_size, + config.audio_config.get("projector_type", "mlp"), + ) + else: + self.audio_encoder = self.audio_projector = None + + if config.use_spectral_norm: + self.apply_spectral_norm() + + self.post_init() + + def get_input_embeddings(self): return self.embed_tokens + def set_input_embeddings(self, v): self.embed_tokens = v + + def _prep_mask(self, mask, shape, dtype): + if mask is None: return None + if mask.dim() == 2: mask = mask[:, None, None, :] + elif mask.dim() == 3: mask = mask[:, None] + return (1.0 - mask.to(dtype)) * torch.finfo(dtype).min + + def _check_memory(self, device, idx): + if not (torch.cuda.is_available() and device.type == "cuda"): return + if idx - self._last_mem_check < 5: return + alloc = torch.cuda.memory_allocated(device) / 1024**3 + self.metrics.log("gpu_mem_gb", alloc) + if alloc > 12: logger.warning(f"Layer {idx}: GPU mem {alloc:.2f} GB") + if alloc > 14: torch.cuda.empty_cache() + self._last_mem_check = idx + + def forward( + self, input_ids=None, pixel_values=None, audio_features=None, + attention_mask=None, past_key_values=None, use_cache=None, + output_hidden_states=False, return_dict=True, **kwargs, + ): + use_cache = use_cache if use_cache is not None else self.config.use_cache + + if input_ids is None: + raise ValueError("input_ids tidak boleh None") + + bs, sl = input_ids.shape + device = input_ids.device + + self.metrics.log("batch_size", bs) + self.metrics.log("seq_length", sl) + + hidden_states = self.embed_tokens(input_ids) + + if self.token_drop is not None and self.training: + hidden_states = self.token_drop(hidden_states) + + enc_hs = enc_mask = None + if self.config.use_multimodal and pixel_values is not None and self.vision_encoder is not None: + try: + vis_feats = self.vision_encoder(pixel_values.to(device)) + vis_embeds = self.vision_projector(vis_feats) + vis_len = vis_embeds.shape[1] + if self.config.use_cross_attention: + enc_hs = vis_embeds + enc_mask = torch.ones(bs, vis_len, dtype=hidden_states.dtype, device=device) + else: + hidden_states = torch.cat([vis_embeds, hidden_states], dim=1) + sl = hidden_states.shape[1] + + if attention_mask is not None: + vis_mask = torch.ones(bs, vis_len, dtype=attention_mask.dtype, device=device) + attention_mask = torch.cat([vis_mask, attention_mask], dim=1) + + except RuntimeError as e: + logger.error(f"Vision encoder gagal: {e}") + + if self.config.use_multimodal and audio_features is not None and self.audio_encoder is not None: + try: + audio_encoded = self.audio_encoder(audio_features.to(device)) + audio_embeds = self.audio_projector(audio_encoded) + audio_len = audio_embeds.shape[1] + + if self.config.use_cross_attention: + enc_hs = audio_embeds if enc_hs is None else torch.cat([enc_hs, audio_embeds], dim=1) + audio_mask = torch.ones(bs, audio_len, dtype=hidden_states.dtype, device=device) + enc_mask = audio_mask if enc_mask is None else torch.cat([enc_mask, audio_mask], dim=1) + else: + hidden_states = torch.cat([audio_embeds, hidden_states], dim=1) + sl = hidden_states.shape[1] + if attention_mask is not None: + audio_mask_flat = torch.ones(bs, audio_len, dtype=attention_mask.dtype, device=device) + attention_mask = torch.cat([audio_mask_flat, attention_mask], dim=1) + + except RuntimeError as e: + logger.error(f"Audio encoder gagal: {e}") + + if attention_mask is not None: + attention_mask = self._prep_mask(attention_mask, (bs, sl), hidden_states.dtype) + if enc_mask is not None and self.config.use_cross_attention: + enc_mask = self._prep_mask(enc_mask, (bs, enc_hs.shape[1]), hidden_states.dtype) + + if use_cache and past_key_values is None: + past_key_values = tuple([None] * len(self.layers)) + + present_kvs = [] if use_cache else None + all_hs = [hidden_states] if output_hidden_states else None + total_aux = torch.tensor(0.0, device=device) + total_z = torch.tensor(0.0, device=device) + + for idx, layer in enumerate(self.layers): + if self.training: self._check_memory(device, idx) + pkv = past_key_values[idx] if past_key_values else None + + if self.gradient_checkpointing and self.training and not use_cache: + def make_fwd(l): + def fwd(hs): + out, _, al, zl = l(hs, attention_mask, enc_hs, enc_mask, None, False) + return out, al, zl + return fwd + from torch.utils.checkpoint import checkpoint + hidden_states, al, zl = checkpoint(make_fwd(layer), hidden_states, use_reentrant=False) + pkv_out = None + else: + hidden_states, pkv_out, al, zl = layer( + hidden_states, attention_mask, enc_hs, enc_mask, pkv, use_cache + ) + + if use_cache: present_kvs.append(pkv_out) + + if torch.is_tensor(al): + total_aux = total_aux + al + else: + total_aux = total_aux + torch.tensor(float(al), device=device) + if torch.is_tensor(zl): + total_z = total_z + zl + else: + total_z = total_z + torch.tensor(float(zl), device=device) + + if self.training: + self.metrics.log(f"l{idx}_aux", al.item() if torch.is_tensor(al) else al) + self.metrics.log(f"l{idx}_z", zl.item() if torch.is_tensor(zl) else zl) + if output_hidden_states: all_hs.append(hidden_states) + + hidden_states = self.norm(hidden_states) + if output_hidden_states: all_hs.append(hidden_states) + self.metrics.step() + + if not return_dict: + non_none = [hidden_states] + if use_cache: non_none.append(tuple(present_kvs)) + if all_hs: non_none.append(all_hs) + non_none.append(total_aux) + non_none.append(total_z) + return tuple(non_none) + + return ( + BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=tuple(present_kvs) if use_cache else None, + hidden_states=all_hs, + attentions=None, + ), + total_aux, total_z, + ) + + def get_metrics_summary(self): return self.metrics.get_summary() + def get_attention_cache_stats(self): + return {f"layer_{i}": l.self_attn.get_cache_stats() for i, l in enumerate(self.layers)} + +# --- CacaForCausalLM --- +class CacaForCausalLM(CacaPreTrainedModel, GenerationMixin): + # _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config: CacaConfig): + super().__init__(config) + self.model = CacaModel(config) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.post_init() + + def get_input_embeddings(self): return self.model.embed_tokens + def set_input_embeddings(self, v): self.model.embed_tokens = v + def get_output_embeddings(self): return self.lm_head + def set_output_embeddings(self, v): self.lm_head = v + def set_decoder(self, d): self.model = d + def get_decoder(self): return self.model + + def forward( + self, + input_ids=None, pixel_values=None, audio_features=None, + attention_mask=None, labels=None, past_key_values=None, + inputs_embeds=None, use_cache=None, output_attentions=None, + output_hidden_states=None, return_dict=None, **kwargs, + ): + + if input_ids is not None: + if input_ids.dtype.is_floating_point: + raise TypeError( + f"input_ids harus integer dtype, dapat {input_ids.dtype}. " + f"Gunakan input_ids.long() untuk konversi." + ) + + if (input_ids < 0).any(): + neg_count = (input_ids < 0).sum().item() + raise ValueError( + f"input_ids memiliki {neg_count} nilai negatif. " + f"Semua token ID harus >= 0." + ) + + max_id = input_ids.max().item() + if max_id >= self.config.vocab_size: + raise ValueError( + f"input_ids memiliki ID {max_id} yang >= vocab_size " + f"({self.config.vocab_size}). Rentang valid [0, {self.config.vocab_size - 1}]." + ) + + if self.config.pad_token_id is not None: + pad_only = (input_ids == self.config.pad_token_id).all(dim=1) + if pad_only.any(): + logger.warning( + f"{pad_only.sum().item()} sequences hanya berisi padding token. " + f"Ini bisa menyebabkan perilaku tak terduga." + ) + + if labels is not None: + if labels.dtype not in (torch.long, torch.int, torch.int32, torch.int64): + raise TypeError(f"labels harus integer dtype, dapat {labels.dtype}") + + if labels.shape != input_ids.shape: + raise ValueError( + f"labels shape {labels.shape} != input_ids shape {input_ids.shape}" + ) + + valid_labels = labels[labels != -100] + if valid_labels.numel() > 0: + invalid_negs = valid_labels[valid_labels < 0] + if invalid_negs.numel() > 0: + raise ValueError( + f"labels memiliki {invalid_negs.numel()} nilai negatif tak valid (bukan -100)" + ) + max_label = valid_labels.max().item() + if max_label >= self.config.vocab_size: + raise ValueError( + f"labels memiliki ID {max_label} yang >= vocab_size " + f"({self.config.vocab_size})" + ) + else: + logger.warning("Semua label adalah -100 (diabaikan). Loss tidak akan dihitung.") + + if attention_mask is not None: + if attention_mask.shape[0] != input_ids.shape[0]: + raise ValueError( + f"attention_mask batch size ({attention_mask.shape[0]}) != " + f"input_ids batch size ({input_ids.shape[0]})" + ) + if attention_mask.shape[1] != input_ids.shape[1]: + raise ValueError( + f"attention_mask seq length ({attention_mask.shape[1]}) != " + f"input_ids seq length ({input_ids.shape[1]})" + ) + unique_vals = attention_mask.unique() + if not torch.all((unique_vals == 0) | (unique_vals == 1)): + logger.warning( + f"attention_mask memiliki nilai selain 0 dan 1: {unique_vals.tolist()}. " + f"Ini bisa menyebabkan perilaku tak terduga." + ) + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs, aux_loss, z_loss = self.model( + input_ids, pixel_values=pixel_values, audio_features=audio_features, + attention_mask=attention_mask, past_key_values=past_key_values, + use_cache=use_cache, output_hidden_states=output_hidden_states, + return_dict=True, # selalu True agar unpack konsisten + ) + + hidden_states = outputs.last_hidden_state + logits = self.lm_head(hidden_states) + + cap = self.config.lm_logit_softcapping or self.config.final_logit_softcapping + if cap: + logits = soft_cap_logits(logits, cap) + + loss = None + if labels is not None: + sl_logits = logits[..., :-1, :].contiguous() + tl = labels[..., 1:].contiguous() + loss_fct = nn.CrossEntropyLoss(ignore_index=-100, label_smoothing=self.config.label_smoothing) + lm_loss = loss_fct(sl_logits.view(-1, sl_logits.size(-1)), tl.view(-1)) + + if self.training: + with torch.no_grad(): + self.model.metrics.log("perplexity", torch.exp(lm_loss).item()) + self.model.metrics.log("lm_loss", lm_loss.item()) + + if self.config.use_moe: + _aux = aux_loss.item() if torch.is_tensor(aux_loss) else float(aux_loss) + _z = z_loss.item() if torch.is_tensor(z_loss) else float(z_loss) + loss = lm_loss + self.config.router_aux_loss_coef * aux_loss + self.config.router_z_loss_coef * z_loss + if self.training: + self.model.metrics.log("aux_loss", _aux) + self.model.metrics.log("z_loss", _z) + else: + loss = lm_loss + + if not return_dict: + return ((loss, logits) if loss is not None else (logits,)) + + return CausalLMOutputWithPast( + loss=loss, logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=None, + ) + + def prepare_inputs_for_generation( + self, input_ids, past_key_values=None, + attention_mask=None, inputs_embeds=None, + pixel_values=None, audio_features=None, **kwargs, + ): + has_past = False + if past_key_values: + try: + has_past = ( + len(past_key_values) > 0 + and past_key_values[0] is not None + and isinstance(past_key_values[0], (tuple, list)) + and past_key_values[0][0] is not None + and past_key_values[0][0].numel() > 0 + ) + except Exception: + pass + + if has_past: + input_ids = input_ids[:, -1:] + pixel_values = None + audio_features = None + + model_inputs = ( + {"input_ids": input_ids} + if inputs_embeds is None or has_past + else {"inputs_embeds": inputs_embeds} + ) + model_inputs.update({ + "past_key_values": past_key_values if has_past else None, + "use_cache": kwargs.get("use_cache"), + "attention_mask": attention_mask, + "pixel_values": pixel_values, + "audio_features": audio_features, + }) + return model_inputs + + @staticmethod + def _reorder_cache(past_key_values, beam_idx): + reordered = () + for lp in past_key_values: + if lp is not None: + reordered += (tuple( + ps.index_select(0, beam_idx.to(ps.device)) + if ps is not None and ps.numel() > 0 else None + for ps in lp + ),) + else: + reordered += (None,) + return reordered + + def save_pretrained(self, save_directory, **kwargs): + quant_backup = getattr(self.config, "quantization_config", None) + had_quant = hasattr(self.config, "quantization_config") + if had_quant and quant_backup is None: + delattr(self.config, "quantization_config") + try: + super().save_pretrained(save_directory, **kwargs) + if self.training: + stats_path = os.path.join(save_directory, "training_stats.json") + with open(stats_path, "w") as f: + json.dump({ + "metrics": self.model.get_metrics_summary(), + "cache_stats": self.model.get_attention_cache_stats(), + }, f, indent=2) + logger.info(f"Training stats disimpan → {stats_path}") + finally: + if had_quant: + self.config.quantization_config = quant_backup + + def get_model_stats(self) -> Dict: + stats = { + "metrics": self.model.get_metrics_summary(), + "cache_stats": self.model.get_attention_cache_stats(), + "gradient_stats": {f"layer_{i}": l.get_gradient_stats() for i, l in enumerate(self.model.layers)}, + } + if self.config.use_moe: + stats["expert_usage"] = { + f"layer_{i}": l.mlp.expert_usage_count.cpu().tolist() + for i, l in enumerate(self.model.layers) + if hasattr(l.mlp, "expert_usage_count") + } + stats["expert_load_imbalance"] = { + f"layer_{i}": l.mlp.router.get_load_imbalance() + for i, l in enumerate(self.model.layers) + if hasattr(l.mlp, "router") and hasattr(l.mlp.router, "get_load_imbalance") + } + return stats + + def apply_lora(self) -> "CacaForCausalLM": + if not self.config.use_lora: + logger.warning("apply_lora() dipanggil tapi config.use_lora=False. Dilewati.") + return self + if not HAS_PEFT: + raise ImportError("peft diperlukan untuk LoRA. Jalankan: pip install peft") + lora_cfg = LoraConfig( + task_type=TaskType.CAUSAL_LM, + r=self.config.lora_rank, + lora_alpha=self.config.lora_alpha, + lora_dropout=self.config.lora_dropout, + bias=self.config.lora_bias, + target_modules=self.config.lora_target_modules, + ) + peft_model = get_peft_model(self, lora_cfg) + peft_model.print_trainable_parameters() + return peft_model + + def get_mup_param_groups(self, base_lr: float) -> List[Dict]: + if not self.config.use_mup: + return [{"params": self.parameters(), "lr": base_lr}] + groups: Dict[float, List] = defaultdict(list) + for name, param in self.named_parameters(): + if not param.requires_grad: continue + mul = self.config.get_mup_lr_multiplier(name) + groups[base_lr * mul].append(param) + return [{"params": ps, "lr": lr} for lr, ps in groups.items()] + +# --- Quantized variant --- +class CacaForCausalLMQuantized(CacaForCausalLM): + def __init__(self, config: CacaConfig, quantization_config: Optional[Dict] = None): + super().__init__(config) + self.quantization_config = quantization_config + if quantization_config: + self._apply_quantization() + + def _apply_quantization(self): + if not HAS_BNB: + raise ImportError("bitsandbytes diperlukan. pip install bitsandbytes") + if self.quantization_config.get("load_in_8bit"): + self._quantize(bits=8) + elif self.quantization_config.get("load_in_4bit"): + self._quantize(bits=4) + + def _quantize(self, bits: int): + compute_dtype = torch.float16 + if bits == 4 and self.quantization_config.get("bnb_4bit_compute_dtype"): + compute_dtype = getattr(torch, self.quantization_config["bnb_4bit_compute_dtype"]) + for name, module in self.named_modules(): + if not isinstance(module, nn.Linear): continue + has_bias = module.bias is not None + if bits == 8: + new = bnb.nn.Linear8bitLt( + module.in_features, module.out_features, has_bias, + threshold=self.quantization_config.get("llm_int8_threshold", 6.0), + ) + else: + new = bnb.nn.Linear4bit( + module.in_features, module.out_features, bias=has_bias, + compute_dtype=compute_dtype, + quant_type=self.quantization_config.get("bnb_4bit_quant_type", "nf4"), + use_double_quant=self.quantization_config.get("bnb_4bit_use_double_quant", True), + ) + new.weight = module.weight + if has_bias: new.bias = module.bias + parent = ".".join(name.split(".")[:-1]) + child = name.split(".")[-1] + setattr(self.get_submodule(parent) if parent else self, child, new) + logger.info(f"Quantization {bits}-bit diterapkan.") + + @classmethod + def from_pretrained_quantized(cls, model_path: str, quantization_config: Dict): + config = CacaConfig.from_pretrained(model_path) + model = cls(config, quantization_config=quantization_config) + sd = torch.load(os.path.join(model_path, "pytorch_model.bin"), map_location="cpu") + model.load_state_dict(sd, strict=False) + return model + +# --- CacaTrainer --- +class CacaTrainer: + def __init__( + self, model, optimizer, scheduler=None, + gradient_accumulation_steps: int = 1, + max_grad_norm: float = 1.0, + use_amp: bool = False, + log_interval: int = 10, + ): + self.model = model + self.optimizer = optimizer + self.scheduler = scheduler + self.grad_accum = gradient_accumulation_steps + self.max_grad = max_grad_norm + self.use_amp = use_amp + self.log_interval = log_interval + self.global_step = 0 + self.epoch = 0 + self.train_metrics: Dict[str, List[float]] = defaultdict(list) + self.scaler = torch.cuda.amp.GradScaler() if use_amp and torch.cuda.is_available() else None + + def _forward(self, batch) -> torch.Tensor: + + if self.use_amp: + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + with torch.cuda.amp.autocast(dtype=dtype): + out = self.model(**batch) + else: + out = self.model(**batch) + loss = out.loss if hasattr(out, "loss") else out[0] + if loss is None: + raise ValueError("Loss adalah None. Pastikan 'labels' ada di dalam batch.") + return loss + + def training_step(self, batch) -> float: + self.model.train() + loss = self._forward(batch) / self.grad_accum + (self.scaler.scale(loss) if self.scaler else loss).backward() + self.train_metrics["loss"].append(loss.item() * self.grad_accum) + return loss.item() * self.grad_accum + + def optimizer_step(self) -> float: + if self.scaler: + self.scaler.unscale_(self.optimizer) + gn = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad) + if self.scaler: + self.scaler.step(self.optimizer) + self.scaler.update() + else: + self.optimizer.step() + self.optimizer.zero_grad(set_to_none=True) + if self.scheduler: self.scheduler.step() + self.train_metrics["grad_norm"].append(gn.item()) + self.global_step += 1 + return gn.item() + + def train_epoch(self, dataloader) -> float: + self.model.train() + epoch_loss, n = 0.0, 0 + for step, batch in enumerate(dataloader): + epoch_loss += self.training_step(batch) + n += 1 + if (step + 1) % self.grad_accum == 0: + gn = self.optimizer_step() + if self.global_step % self.log_interval == 0: + lr = self.optimizer.param_groups[0]["lr"] + logger.info( + f"Epoch {self.epoch} | Step {self.global_step} | " + f"Loss: {epoch_loss/n:.4f} | GradNorm: {gn:.4f} | LR: {lr:.2e}" + ) + if hasattr(self.model, "get_model_stats"): + stats = self.model.get_model_stats() + ppl = stats.get("metrics", {}).get("perplexity", {}).get("last") + if ppl: logger.info(f" PPL: {ppl:.2f}") + self.epoch += 1 + return epoch_loss / n if n > 0 else 0.0 + + def get_metrics(self) -> Dict: + return { + k: { + "mean": float(np.mean(v)), "std": float(np.std(v)), + "min": float(np.min(v)), "max": float(np.max(v)), + } + for k, v in self.train_metrics.items() if v + } + +print("✓ ARSITEKTUR BERHASIL") + +# --- Inisialisasi Model --- +from transformers import AutoConfig, AutoModel, AutoModelForCausalLM + +AutoConfig.register("caca", CacaConfig) +AutoModel.register(CacaConfig, CacaModel) +AutoModelForCausalLM.register(CacaConfig, CacaForCausalLM) + +# --- Utilitas --- +def calculate_params(config_dict: Dict) -> Dict: + h = config_dict["hidden_size"] + i = config_dict["intermediate_size"] + v = config_dict["vocab_size"] + l = config_dict["num_hidden_layers"] + n_h = config_dict["num_attention_heads"] + n_kv = config_dict["num_key_value_heads"] + hd = config_dict.get("head_dim", h // n_h) + tie = config_dict.get("tie_word_embeddings", False) + use_moe = config_dict.get("use_moe", False) + num_experts = config_dict.get("num_experts", 1) + moe_freq = config_dict.get("moe_layer_frequency", 1) + use_cross = config_dict.get("use_cross_attention", False) + cross_freq = config_dict.get("cross_attention_frequency", 1) + + embed = v * h + q_p = h * (n_h * hd) + kv_p = 2 * h * (n_kv * hd) + o_p = (n_h * hd) * h + attn_l = q_p + kv_p + o_p + norm_l = 2 * h + + lm_head = 0 if tie else v * h + + total_ffn = 0 + total_attn = 0 + total_norm = 0 + total_cross = 0 + + for layer_idx in range(l): + total_attn += attn_l + total_norm += norm_l + + is_moe_layer = use_moe and (layer_idx % moe_freq == 0) + n_experts_this_layer = num_experts if is_moe_layer else 1 + ffn_l = 3 * h * i + total_ffn += ffn_l * n_experts_this_layer + + if is_moe_layer: + total_ffn += h * num_experts + + if use_cross and (layer_idx % cross_freq == 0): + cross_l = attn_l + h + total_cross += cross_l + + total = embed + total_attn + total_ffn + total_norm + total_cross + h + lm_head + + return { + "total": total, + "embeddings": embed, + "embedding_pct": embed / total * 100, + "lm_head": lm_head, + "lm_head_pct": lm_head / total * 100, + "attention_total": total_attn, + "ffn_total": total_ffn, + "cross_attn_total": total_cross, + "attention_per_layer": attn_l, + "ffn_per_layer_dense": 3 * h * i, + } + +def format_params(n: int) -> str: + if n >= 1e12: return f"{n/1e12:.2f}T" + if n >= 1e9: return f"{n/1e9:.2f}B" + if n >= 1e6: return f"{n/1e6:.2f}M" + return f"{n/1e3:.2f}K" + +def count_parameters(model) -> Dict: + total = sum(p.numel() for p in model.parameters()) + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + return { + "total": total, "trainable": trainable, + "total_M": total / 1e6, "trainable_M": trainable / 1e6, + } + +def get_model_size_mb(model) -> Dict: + p_bytes = sum(p.nelement() * p.element_size() for p in model.parameters()) + b_bytes = sum(b.nelement() * b.element_size() for b in model.buffers()) + return { + "param_mb": p_bytes / 1024**2, + "buffer_mb": b_bytes / 1024**2, + "total_mb": (p_bytes + b_bytes) / 1024**2, + } + +def print_model_info(model) -> None: + sep = "=" * 70 + print(f"\n{sep}\nINFO MODEL\n{sep}") + p = count_parameters(model) + print(f"Parameter : {p['total']:,} total ({p['total_M']:.2f}M)") + print(f" {p['trainable']:,} trainable ({p['trainable_M']:.2f}M)") + s = get_model_size_mb(model) + print(f"Ukuran : {s['total_mb']:.1f} MB " + f"(params {s['param_mb']:.1f} MB + buffers {s['buffer_mb']:.1f} MB)") + if hasattr(model, "config"): + c = model.config + print(f"Vocab : {c.vocab_size:,}") + print(f"Hidden : {c.hidden_size} | Layers: {c.num_hidden_layers} | Heads: {c.num_attention_heads}") + print(f"FFN : {c.intermediate_size}") + if c.use_moe: print(f"MoE : {c.num_experts} experts, top-{c.num_experts_per_tok}") + if c.use_lora: print(f"LoRA : rank={c.lora_rank}, alpha={c.lora_alpha}") + if c.use_mup: print(f"MuP : base_width={c.mup_base_width}") + if c.use_spectral_norm: print("SpectralNorm: aktif") + if c.token_dropout > 0: print(f"TokenDrop : {c.token_dropout}") + if c.use_multimodal: print("Multimodal : vision encoder aktif") + print(f"{sep}\n") \ No newline at end of file