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