File size: 2,862 Bytes
d53adc9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | # Generated by ml.integrations.export.runtime_packager.write_remote_code_bundle.
# Exported for HuggingFace trust_remote_code loading.
# This file is intentionally self-contained.
from __future__ import annotations
from dataclasses import dataclass
from .semantics import canonicalize_model_values, validate_model_values
@dataclass
class ModelArgs:
"""Torch-free runtime configuration for Sophia Hybrid."""
vocab_size: int = 65536
dim: int = 1536
n_layers: int = 28
num_heads: int = 16
head_dim: int = 128
ffn_hidden: int = 3968
kda_decay_rank: int = 128
kda_output_gate_rank: int = 128
kda_output_gate_full_rank: bool = True
kda_decay_lower_bound: float = -5.0
kda_dt_min: float = 1e-3
kda_dt_max: float = 1e-1
kda_dt_floor: float = 1e-4
kda_a_log_init: float = 0.0
mla_q_rank: int = 384
mla_kv_rank: int = 128
short_conv_kernel: int = 4
attn_res_block_size: int = 4
situ_gate_softcap: float = 4.0
situ_up_softcap: float = 25.0
norm_eps: float = 1e-5
max_seq_len: int = 4096
max_batch_size: int = 4
dropout: float = 0.0
initializer_range: float = 0.02
kda_backend: str = "auto"
use_cache: bool = True
def __post_init__(self) -> None:
normalized = canonicalize_model_values(vars(self))
for name, value in normalized.items():
if hasattr(self, name):
setattr(self, name, value)
validate_model_values(normalized)
@property
def attention_inner_dim(self) -> int:
return int(self.num_heads) * int(self.head_dim)
def layer_type(self, layer_idx: int) -> str:
return "mla" if int(layer_idx) % 4 == 3 else "kda"
def ensure_runtime_batch_capacity(self, batch_size: int) -> int:
required = int(batch_size)
if required <= 0:
raise ValueError(f"batch_size must be > 0, got {required}")
self.max_batch_size = max(int(self.max_batch_size), required)
return int(self.max_batch_size)
def runtime_batch_capacity(self) -> int:
return int(self.max_batch_size)
def ensure_runtime_sequence_capacity(self, max_seq_len: int) -> int:
required = int(max_seq_len)
if required <= 0:
raise ValueError(f"max_seq_len must be > 0, got {required}")
configured = int(self.max_seq_len)
if required > configured:
raise ValueError(
"runtime sequence capacity cannot exceed configured max_seq_len: "
f"required={required} configured={configured}"
)
return configured
def runtime_sequence_capacity(self) -> int:
return int(self.max_seq_len)
def runtime_max_seq_len(self) -> int:
return self.runtime_sequence_capacity()
def runtime_layer_count(self) -> int:
return int(self.n_layers)
|