1
2
3
4
5
6
7
8
9

@staticmethod
def _get_layer_type(layer_idx: int, total_layers: int) -> str:
    third = total_layers // 3  # 각 8개 레이어
    if layer_idx < third:           # Layer 0-7: 순수 SSM
        return "ssm"
    elif layer_idx < 2 * third:     # Layer 8-15: SWA + MoE
        return "swa_moe"
    else:                           # Layer 16-23: SSM + MoE
        return "ssm_moe"

1
2
3
4
5
6
7
8
9

class AxonSSM(nn.Module):
    def __init__(self, config: BitAxonConfig):
        self.in_proj = nn.Linear(D, 2 * E, bias=False)          # 입력을 x와 z 브랜치로 분할
        self.conv1d = nn.Conv1d(E, E, kernel_size=d_conv, groups=E)  # 깊이별 인과 컨볼루션
        self.x_proj = nn.Linear(E, d_state * 2 + 1, bias=False)   # B, C, dt 파라미터로 투영
        self.dt_proj = nn.Linear(1, E, bias=True)                # 채널별 스텝 사이즈
        self.out_proj = nn.Linear(E, D, bias=False)              # 출력 투영
        self.A_log = mx.log(mx.arange(1, d_state + 1))            # 대각선 SSM 상태 행렬
        self.D = mx.ones((E,))                                    # 스킵 연결 파라미터

1
2
3
4
5
6
7
8
9

def _ssm_scan_parallel(self, x, dt, B_in, C_in):
    step = config.ssm_scan_step  # 기본값 64
    for j in range(d_state):     # 상태 차원별로 독립 처리
        for i in range(0, L, step):
            S = min(step, L - i)
            dtA_chunk = dtA[:, i : i + S, :]
            dtx_chunk = dtx[:, i : i + S, :]
            B_chunk = B_in[:, i : i + S, j]
            C_chunk = C_in[:, i : i + S, j]

1
2
3
4
5
6
7

def segsum(x: mx.array) -> mx.array:
    """하드웨어 효율적인 병렬 세그먼트 합"""
    seq_len = x.shape[-1]
    cs = mx.cumsum(x, axis=-1)
    diff = cs[..., :, None] - cs[..., None, :]
    mask = mx.tril(mx.ones((seq_len, seq_len), dtype=diff.dtype), -1)
    return diff * mask

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

def _make_sliding_window_mask(self, seq_len: int, kv_len: int, q_offset: int = 0):
    q_pos = mx.arange(q_offset, q_offset + seq_len)
    k_pos = mx.arange(kv_len)

    # 인과 제약: 미래 토큰을 볼 수 없음
    causal_mask = k_pos[None, :] <= (q_pos[:, None] + causal_offset)

    # 윈도우 제약: 제한된 어텐션 범위
    window_mask = (q_pos[:, None] + causal_offset) - k_pos[None, :] < self.window_size

    # 결합: 인과+윈도우 밖의 위치는 -inf
    mask = mx.where(causal_mask & window_mask, 0.0, -mx.inf)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

class KVCache:
    def __init__(self, window_size: int | None = None):
        self.window_size = window_size

    def update_and_fetch(self, xk: mx.array, xv: mx.array) -> tuple[mx.array, mx.array]:
        self.k = mx.concatenate([self.k, xk], axis=2)
        self.v = mx.concatenate([self.v, xv], axis=2)
        if self.window_size is not None:
            self.k = self.k[:, :, -self.window_size:]  # 윈도우로 트리밍
            self.v = self.v[:, :, -self.window_size:]

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12

class SharedExpertMoE(nn.Module):
    def __call__(self, x: mx.array) -> mx.array:
        gates = self.gate(x)                      # (batch, seq_len, num_experts)
        gates = mx.softmax(gates, axis=-1)        # 전문가별 소프트맥스

        # Top-k 전문가 선택
        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k-1, axis=-1)[..., :k])
        scores = mx.take_along_axis(gates, inds, axis=-1)

        # 전문가 처리
        y = self.switch_mlp(x, inds)
        y = (y * scores[..., None]).sum(axis=-2)  # 가중치 합산

1
2
3
4
5
6
7
8
9

class SwitchLinear(nn.Module):
    def __call__(self, x: mx.array, indices: mx.array) -> mx.array:
        B, L, K = indices.shape
        flat_idx = indices.reshape(-1)
        x_flat = x.reshape(-1, 1, D)

        w_t = self.weight.swapaxes(-1, -2)
        out = mx.gather_mm(x_flat, w_t, rhs_indices=flat_idx,
                          sorted_indices=sorted_indices)

1
2
3
4

# 공유 전문가 게이팅
shared_out = self.shared_expert(x)
gate = sigmoid(shared_expert_gate(x))
output = gated_expert_output + gate * shared_out

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class AxonSSMMoEBlock(nn.Module):
    def __call__(self, x, cache=None):
        # SSM with residual
        residual = x
        x = self.input_norm(x)
        ssm_out, ssm_cache = self.ssm(x, cache=cache)
        x = residual + ssm_out

        # MoE with residual
        residual = x
        x = self.post_ssm_norm(x)
        x = residual + self.moe(x)
        return x, ssm_cache

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13

class QuantizedSwitchLinear(nn.Module):
    def __init__(self, input_dims, output_dims, num_experts,
                 group_size=64, bits=4):
        # 그룹별 양자화: 64개 원소마다 스케일 팩터와 바이어스
        self.weight, self.scales, self.biases_quant = \
            mx.quantize(weight, group_size=group_size, bits=bits)

    def __call__(self, x: mx.array, indices: mx.array):
        # 양자화된 가중치 + gather를 결합한 단일 연산
        out = mx.gather_qmm(
            x_flat, self.weight, self.scales, self.biases_quant,
            rhs_indices=flat_idx, group_size=self.group_size, bits=self.bits
        )

1
2
3
4
5
6

@dataclass
class ThermalPolicy:
    max_speed_temp: float = 75.0    # 이 온도 이하: 전속 학습
    pause_temp: float = 85.0        # 이 온도 이상: 학습 일시 정지
    stop_temp: float = 95.0         # 이 온도 이상: 학습 중단
    pause_duration: float = 0.5      # 냉각 대기 시간 (초)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

class CoolingScheduler:
    def __init__(self, monitor, policy: ThermalPolicy = None):
        self._monitor = monitor
        self._policy = policy or ThermalPolicy()
        self._total_pause_time: float = 0.0

    def check_before_step(self, step: int) -> None:
        temp = self._monitor.temperature
        if temp >= self._policy.stop_temp:  # 95°C 임계
            raise ThermalShutdownError(
                f"SoC temperature {temp:.1f}C exceeds stop threshold")
        while temp >= self._policy.pause_temp:  # 85°C 임계
            time.sleep(self._policy.pause_duration)  # 0.5초 대기
            self._total_pause_time += self._policy.pause_duration

    def should_reduce_batch(self) -> bool:
        temp = self._monitor.temperature
        return self._policy.max_speed_temp <= temp < self._policy.pause_temp

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17

class SequencePacker:
    def __init__(self, max_seq_len: int = 2048, eos_token_id: int = 151645):
        self.max_seq_len = max_seq_len
        self.eos_token_id = eos_token_id

    def add_example(self, token_ids: list[int], loss_mask: list[int]):
        # 버퍼가 비어있지 않으면 EOS 구분자 삽입
        if self._buffer_ids:
            self._buffer_ids.append(self.eos_token_id)
            self._buffer_mask.append(0)  # 구분자에서는 손실 계산 안 함

        # 버퍼가 가득 차면 배치 반환
        while len(self._buffer_ids) >= self.max_seq_len:
            yield PackedBatch(
                token_ids=buffer[:self.max_seq_len],
                loss_mask=mask[:self.max_seq_len]
            )

1
2
3
4
5
6
7

def orpo_loss(chosen_logps, rejected_logps, beta=0.1):
    # 오즈비(odds ratio) 계산
    log_odds = (chosen_logps - rejected_logps) - \
               (log1mexp(chosen_logps) - log1mexp(rejected_logps))
    # 시그모이드 페널티
    loss = -mx.mean(nn.log_sigmoid(beta * log_odds))
    return loss

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16

def compute_orpo_loss(model, chosen_ids, chosen_labels,
                      rejected_ids, rejected_labels, beta=0.1):
    # 순방향 패스 (2회 — 참조 모델 불필요)
    logits_chosen = model(chosen_ids)
    logits_rejected = model(rejected_ids)

    # 선택된 시퀀스에서 NLL 손실
    nll_loss = cross_entropy_loss(logits_chosen, chosen_labels)

    # 선호도 비교
    chosen_logps = get_logps(logits_chosen, chosen_labels)
    rejected_logps = get_logps(logits_rejected, rejected_labels)

    # 결합 목적 함수
    orpo_penalty = orpo_loss(chosen_logps, rejected_logps, beta)
    total_loss = nll_loss + orpo_penalty

1
2
3
4
5
6
7
8

def log1mexp(x: mx.array) -> mx.array:
    threshold = mx.array(-_LN2)  # -ln(2)
    use_branch1 = x < threshold
    x_branch1 = mx.where(use_branch1, x, mx.zeros_like(x))
    x_branch2 = mx.where(~use_branch1, x, mx.zeros_like(x))
    branch1 = mx.log(-mx.expm1(x_branch1))         # x < -ln(2)인 경우
    branch2 = mx.log1p(-mx.exp(x_branch2))          # x >= -ln(2)인 경우
    return mx.where(use_branch1, branch1, branch2)

1
2
3
4
5
6
7
8

@dataclass
class TrainingConfig:
    quantize_bits: int = 4
    quantize_group_size: int = 64
    lora_rank: int = 8
    lora_dropout: float = 0.0
    lora_scale: float = 20.0
    use_dora: bool = True  # Weight-Decomposed LoRA

1
2
3
4
5
6
7
8

def __call__(self, x):
    y = self.linear(x)
    z = (self.dropout(x) @ self.lora_a) @ self.lora_b
    out = y + (self.scale * z).astype(x.dtype)

    # 원래 크기 보존 (DoRA 핵심)
    denom = mx.sqrt(self._dora_w_sq_norm + cross + d_sq)
    out = (self.m / denom).astype(x.dtype) * out