Memory Budget¶

Bit-Axon targets the MacBook Air M4 with 16 GB unified memory, of which roughly 8 GB is available for the model. This page breaks down every byte.

Available Memory¶

Weight Memory by Precision¶

FP16 weights alone consume 6.4 GB — nearly the entire budget before any activations or KV cache. Q4 quantization is the default strategy. MLX's nn.QuantizedLinear runs 4-bit matmul as a single fused Metal kernel, delivering 2–3× speedup over FP16 while using 3.6× less memory.

Parameter Count Breakdown¶

SSM State Memory (Fixed, Layers 1–8 and 17–24)¶

The 16 SSM layers maintain a fixed-size hidden state regardless of context length:

• SSM state: \((1, 7680, 16)\) per layer × 16 layers = ~3.7 MB
• Conv cache: \((1, 3, 7680)\) per layer × 16 layers = ~0.7 MB
• Total SSM state: ~4.4 MB — negligible

SWA KV Cache (Layers 9–16, 32 heads, head_dim=80)¶

Only the 8 SWA layers produce a KV cache. Formula:

Total Inference Memory¶

Key conclusions:

• ✅ Q4 alone fits within 8 GB up to ~16K context
• ✅ Q4 + TurboQuant fits within 8 GB at full 64K context with ~5 GB to spare
• ❌ FP16 inference exceeds 8 GB at any context length beyond 4K

QLoRA Training Memory¶

Full fine-tuning (FP16 weights + gradients + Adam optimizer) requires ~38.4 GB — physically impossible on 16 GB. QLoRA (Q4 frozen base + LoRA adapters) compresses training to:

Thermal Management¶

The fanless MacBook Air M4 dissipates all heat passively. Bit-Axon uses a thermal-aware training scheduler that monitors SoC temperature via macOS powermetrics:

CoolingScheduler¶

ThermalPolicy¶

class ThermalPolicy:
    PAUSE_TEMP = 85    # °C — insert micro-pauses
    STOP_TEMP = 95     # °C — halt training entirely
    RESUME_TEMP = 75   # °C — resume normal operation
    CHECK_INTERVAL = 1 # seconds between temperature checks

MoE Thermal Benefit¶

MoE sparsity directly reduces thermal load:

• Only ~1.4B of 3.2B parameters are active per token (~44%)
• 56% of the chip's compute units remain idle on each forward pass
• This is roughly equivalent to running a 1.4B dense model in terms of heat generation
• High likelihood of sustained training on a fanless MacBook Air

Context Length Strategy¶

Weight Tying Savings¶

embed_tokens:  (32000, 2048) = 65.5M params
lm_head:       (32000, 2048) = 65.5M params  ← same tensor
                                ─────────────
Savings:                       65.5M params = ~131 MB (FP16) / ~36 MB (Q4)

By tying embed_tokens.weight = lm_head.weight, the model avoids storing a duplicate copy of the embedding matrix. This is particularly valuable at Q4 where every ~36 MB counts toward the 8 GB ceiling.

See also¶

• ← Architecture Overview
• Quantization Guide — NF4 quantization implementation and QLoRA training
• Training Guide — QLoRA memory usage during training
• TurboQuant Paper — Planned KV cache compression for 64K context
• Inference Guide — KV cache behavior during generation