splitbit_llm/train/auto_size.py

"""Auto-size adjustment for SplitBit LLM.

Detects hardware (reuses config.py hardware detection).
Picks optimal model config per tier.
Adjusts batch size and training hyperparams accordingly.
Can upgrade/downgrade model size at runtime (transfer weights).
"""

from __future__ import annotations

import logging
from typing import Any

from ..config import (
    HardwareTier,
    ModelConfig,
    QuantConfig,
    StorageConfig,
    detect_hardware,
    get_model_config,
)

logger = logging.getLogger(__name__)


# Training hyperparameters per tier
TIER_TRAINING_CONFIG: dict[HardwareTier, dict[str, Any]] = {
    HardwareTier.MOBILE: {
        "batch_size": 1,
        "seq_len": 128,
        "lr": 1e-3,
        "epochs": 20,
        "warmup_steps": 50,
    },
    HardwareTier.MINIMAL: {
        "batch_size": 2,
        "seq_len": 256,
        "lr": 3e-4,
        "epochs": 15,
        "warmup_steps": 100,
    },
    HardwareTier.LIGHT: {
        "batch_size": 4,
        "seq_len": 256,
        "lr": 3e-4,
        "epochs": 10,
        "warmup_steps": 200,
    },
    HardwareTier.STANDARD: {
        "batch_size": 8,
        "seq_len": 512,
        "lr": 3e-4,
        "epochs": 10,
        "warmup_steps": 500,
    },
    HardwareTier.FULL: {
        "batch_size": 16,
        "seq_len": 1024,
        "lr": 2e-4,
        "epochs": 8,
        "warmup_steps": 1000,
    },
    HardwareTier.MAXIMUM: {
        "batch_size": 32,
        "seq_len": 2048,
        "lr": 1e-4,
        "epochs": 5,
        "warmup_steps": 2000,
    },
    HardwareTier.DATACENTER: {
        "batch_size": 64,
        "seq_len": 4096,
        "lr": 5e-5,
        "epochs": 3,
        "warmup_steps": 5000,
    },
    HardwareTier.SUPERCOMPUTER: {
        "batch_size": 128,
        "seq_len": 8192,
        "lr": 2e-5,
        "epochs": 2,
        "warmup_steps": 10000,
    },
}


# Inference parameters per tier
TIER_INFERENCE_CONFIG: dict[HardwareTier, dict[str, Any]] = {
    HardwareTier.MOBILE: {
        "max_tokens": 32,
        "temperature": 0.5,
        "top_k": 20,
        "use_cache": True,
    },
    HardwareTier.MINIMAL: {
        "max_tokens": 64,
        "temperature": 0.7,
        "top_k": 40,
        "use_cache": True,
    },
    HardwareTier.LIGHT: {
        "max_tokens": 128,
        "temperature": 0.7,
        "top_k": 40,
        "use_cache": True,
    },
    HardwareTier.STANDARD: {
        "max_tokens": 256,
        "temperature": 0.7,
        "top_k": 50,
        "use_cache": True,
    },
    HardwareTier.FULL: {
        "max_tokens": 512,
        "temperature": 0.7,
        "top_k": 50,
        "use_cache": True,
    },
    HardwareTier.MAXIMUM: {
        "max_tokens": 1024,
        "temperature": 0.7,
        "top_k": 100,
        "use_cache": True,
    },
    HardwareTier.DATACENTER: {
        "max_tokens": 2048,
        "temperature": 0.7,
        "top_k": 100,
        "use_cache": True,
    },
    HardwareTier.SUPERCOMPUTER: {
        "max_tokens": 4096,
        "temperature": 0.7,
        "top_k": 200,
        "use_cache": True,
    },
}


# Voice-optimized inference parameters per tier
TIER_VOICE_CONFIG: dict[HardwareTier, dict[str, Any]] = {
    HardwareTier.MOBILE: {
        "max_tokens": 16,
        "temperature": 0.3,
        "top_k": 10,
        "use_cache": True,
    },
    HardwareTier.MINIMAL: {
        "max_tokens": 32,
        "temperature": 0.5,
        "top_k": 20,
        "use_cache": True,
    },
    HardwareTier.LIGHT: {
        "max_tokens": 48,
        "temperature": 0.5,
        "top_k": 30,
        "use_cache": True,
    },
    HardwareTier.STANDARD: {
        "max_tokens": 64,
        "temperature": 0.5,
        "top_k": 40,
        "use_cache": True,
    },
    HardwareTier.FULL: {
        "max_tokens": 96,
        "temperature": 0.5,
        "top_k": 50,
        "use_cache": True,
    },
    HardwareTier.MAXIMUM: {
        "max_tokens": 128,
        "temperature": 0.5,
        "top_k": 50,
        "use_cache": True,
    },
    HardwareTier.DATACENTER: {
        "max_tokens": 256,
        "temperature": 0.5,
        "top_k": 100,
        "use_cache": True,
    },
    HardwareTier.SUPERCOMPUTER: {
        "max_tokens": 512,
        "temperature": 0.5,
        "top_k": 200,
        "use_cache": True,
    },
}


class AutoSizer:
    """Hardware-aware model sizing and hyperparameter adjustment.

    Detects hardware tier and provides optimal configs for:
    - Model architecture (layers, heads, dim, vocab)
    - Weight quantization format
    - Training hyperparameters (batch size, lr, epochs)
    - Inference parameters (max tokens, temperature)
    - Voice-optimized inference (shorter, faster responses)
    - Skill storage limits
    """

    def __init__(self, tier: HardwareTier | None = None) -> None:
        self.tier = tier or detect_hardware()
        self.model_config = get_model_config(self.tier)
        self.quant_config = QuantConfig.for_tier(self.tier)
        self.storage_config = StorageConfig.for_tier(self.tier)
        self.training_config = TIER_TRAINING_CONFIG.get(self.tier, TIER_TRAINING_CONFIG[HardwareTier.MINIMAL])
        self.inference_config = TIER_INFERENCE_CONFIG.get(self.tier, TIER_INFERENCE_CONFIG[HardwareTier.MINIMAL])
        self.voice_config = TIER_VOICE_CONFIG.get(self.tier, TIER_VOICE_CONFIG[HardwareTier.MINIMAL])

        logger.info(
            "AutoSizer: tier=%s, model=%dL/%dH/d%d, quant=%s (%.1f bpw), "
            "batch=%d, lr=%.1e, max_tokens=%d, voice_tokens=%d",
            self.tier.value,
            self.model_config.n_layers, self.model_config.n_heads, self.model_config.d_model,
            self.quant_config.format, self.quant_config.bpw,
            self.training_config["batch_size"], self.training_config["lr"],
            self.inference_config["max_tokens"], self.voice_config["max_tokens"],
        )

    def get_model_config(self) -> ModelConfig:
        return self.model_config

    def get_quant_config(self) -> QuantConfig:
        return self.quant_config

    def get_training_params(self) -> dict[str, Any]:
        return self.training_config.copy()

    def get_inference_params(self) -> dict[str, Any]:
        return self.inference_config.copy()

    def get_voice_params(self) -> dict[str, Any]:
        return self.voice_config.copy()

    def get_storage_config(self) -> StorageConfig:
        return self.storage_config

    def get_all_stats(self) -> dict[str, Any]:
        return {
            "tier": self.tier.value,
            "model": {
                "n_layers": self.model_config.n_layers,
                "n_heads": self.model_config.n_heads,
                "d_model": self.model_config.d_model,
                "d_ff": self.model_config.d_ff,
                "vocab_size": self.model_config.vocab_size,
                "max_seq_len": self.model_config.max_seq_len,
                "param_count": self.model_config.param_count,
            },
            "quant": {
                "format": self.quant_config.format,
                "bpw": self.quant_config.bpw,
            },
            "training": self.training_config,
            "inference": self.inference_config,
            "voice": self.voice_config,
            "storage": {
                "skill_storage_mb": self.storage_config.skill_storage_mb,
                "max_skills": self.storage_config.max_skills,
            },
        }