caca-1B-untrained / caca_transformers.py

Update caca_transformers.py

50c5120 verified 3 months ago

106 kB

	from __future__ import annotations

	import math
	import os
	import json
	import logging
	import warnings
	from collections import OrderedDict, defaultdict
	from typing import Any, Dict, List, Optional, Tuple, Union

	import numpy as np
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from transformers import PretrainedConfig, PreTrainedModel
	from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
	from transformers.generation.utils import GenerationMixin

	logger = logging.getLogger(__name__)

	# --- Dependensi opsional ---
	try:
	from flash_attn import flash_attn_func, flash_attn_varlen_func
	HAS_FLASH_ATTN = True
	except ImportError:
	HAS_FLASH_ATTN = False

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

	HAS_SDPA = hasattr(F, "scaled_dot_product_attention")

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

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

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

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

	self._validate_config()

	super().__init__(
	pad_token_id=pad_token_id,
	bos_token_id=bos_token_id,
	eos_token_id=eos_token_id,
	tie_word_embeddings=tie_word_embeddings,
	**kwargs,
	)

	def _validate_config(self) -> None:
	errors: List[str] = []

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	@classmethod
	def from_variant(cls, variant: str, **overrides) -> "CacaConfig":

	if variant not in MODEL_CONFIGS:
	raise ValueError(
	f"Variant '{variant}' tidak dikenal. "
	f"Tersedia: {sorted(MODEL_CONFIGS.keys())}"
	)
	cfg = {MODEL_CONFIGS[variant], overrides}
	return cls(**cfg)

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

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

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

	lm_head = 0 if tie else v * h

	# + h = final RMSNorm
	total = embed + l * (attn_l + ffn_l + norm_l) + h + lm_head

	return {
	"total": total,
	"embeddings": embed,
	"embedding_pct": embed / total * 100,
	"lm_head": lm_head,
	"lm_head_pct": lm_head / total * 100,
	"attention_per_layer": attn_l,
	"ffn_per_layer": ffn_l,
	}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	return inv_freq

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	return top_w, top_idx, aux_loss + ent_loss, z_loss

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

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

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

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

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

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

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

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

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

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

	if self.use_expert_choice:
	weights, aux_loss, z_loss = self.router(hidden_states)
	out = torch.zeros_like(flat)
	for i, expert in enumerate(self.experts):
	mask = weights[:, i] > 1e-6
	if mask.any():
	if not self.training:
	self.expert_usage_count[i] += mask.sum()
	expert_out = self._safe_expert(expert, flat[mask], i, self.config)
	out[mask] += expert_out * weights[mask, i:i+1]
	else:
	top_w, top_idx, aux_loss, z_loss = self.router(hidden_states)
	out = torch.zeros_like(flat)

	for i in range(self.num_experts):
	mask = (top_idx == i).any(-1)
	if not mask.any():
	continue
	if not self.training:
	self.expert_usage_count[i] += mask.sum()
	expert_out = self._safe_expert(self.experts[i], flat[mask], i, self.config)
	token_ids = torch.where(mask)[0]
	expert_pos = (top_idx[mask] == i).float().argmax(dim=-1)
	weights_i = top_w[mask][
	torch.arange(expert_out.shape[0], device=out.device),
	expert_pos
	]
	out.scatter_add_(
	0,
	token_ids.unsqueeze(1).expand_as(expert_out),
	expert_out * weights_i.unsqueeze(1)
	)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	bs, sl, _ = hidden_states.shape

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

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

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

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

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

	present = (k, v) if use_cache else None

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

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

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

	return self.o_proj(out), present

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	aux_loss = z_loss = 0.0

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

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

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

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

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

	return hidden_states, present_kv, aux_loss, z_loss

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

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

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

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

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

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

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

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

	def forward(self, x):

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

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

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

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

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

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

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

	if config.use_spectral_norm:
	self.apply_spectral_norm()

	self.post_init()

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

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

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

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

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

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

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

	hidden_states = self.embed_tokens(input_ids)

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

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

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

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

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

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

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

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

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

	if use_cache: present_kvs.append(pkv_out)

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

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

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

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

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

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

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

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

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

	def forward(
	self,
	input_ids=None, pixel_values=None, audio_features=None,
	attention_mask=None, labels=None, past_key_values=None,
	inputs_embeds=None, use_cache=None, output_attentions=None,
	output_hidden_states=None, return_dict=None, **kwargs,
	):
	if input_ids is not None:
	if input_ids.dtype.is_floating_point:
	raise TypeError(f"input_ids harus int dtype, dapat {input_ids.dtype}")
	if (input_ids < 0).any():
	raise ValueError(f"input_ids memiliki {(input_ids < 0).sum()} nilai negatif")
	max_id = input_ids.max().item()
	if max_id >= self.config.vocab_size:
	raise ValueError(f"input_ids max={max_id} >= vocab_size={self.config.vocab_size}")
	if labels is not None:
	if labels.shape != input_ids.shape:
	raise ValueError(f"labels shape {labels.shape} != input_ids shape {input_ids.shape}")
	valid = labels[labels != -100]
	if valid.numel() > 0 and valid.max().item() >= self.config.vocab_size:
	raise ValueError(f"labels max={valid.max().item()} >= vocab_size={self.config.vocab_size}")
	if valid.numel() == 0:
	logger.warning("Semua label adalah -100. Loss tidak akan dihitung.")
	if attention_mask is not None:
	if attention_mask.shape != input_ids.shape:
	raise ValueError(
	f"attention_mask shape {attention_mask.shape} != input_ids shape {input_ids.shape}"
	)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	# --- AutoClass registration ---
	from transformers import AutoConfig, AutoModel, AutoModelForCausalLM

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