| import json |
| import os |
| import re |
| import shutil |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Dict, Optional, Tuple |
| from urllib.request import Request, urlopen |
|
|
| import torch |
|
|
| from .config import TextConfig |
|
|
|
|
| class AdapterLoadError(RuntimeError): |
| pass |
|
|
|
|
| def _cache_root() -> Path: |
| hf_hub_cache = os.environ.get("HF_HUB_CACHE") |
| if hf_hub_cache: |
| return Path(hf_hub_cache) |
|
|
| hf_home = os.environ.get("HF_HOME") |
| if hf_home: |
| return Path(hf_home) / "hub" |
|
|
| return Path("~/.cache/huggingface/hub").expanduser() |
|
|
|
|
| def adapter_cache_dir() -> Path: |
| return _cache_root() / "md_finetunes" |
|
|
|
|
| def normalize_adapter_id(value: Optional[str]) -> Optional[str]: |
| if not value: |
| return None |
| tail = value.split("/")[-1].strip() |
| if "@" not in tail: |
| return None |
| return tail |
|
|
|
|
| def parse_adapter_id(adapter_id: str) -> Tuple[str, str]: |
| if not adapter_id or "@" not in adapter_id: |
| raise AdapterLoadError( |
| f"Invalid adapter id '{adapter_id}'. Expected 'finetune_id@step'." |
| ) |
| finetune_id, step = adapter_id.split("@", 1) |
| if not finetune_id or not step: |
| raise AdapterLoadError( |
| f"Invalid adapter id '{adapter_id}'. Expected 'finetune_id@step'." |
| ) |
| return finetune_id, step |
|
|
|
|
| def _fetch_presigned_url(finetune_id: str, step: str) -> str: |
| endpoint = os.getenv("MOONDREAM_ENDPOINT", "https://api.moondream.ai").rstrip("/") |
| api_key = os.getenv("MOONDREAM_API_KEY") |
| if not api_key: |
| raise AdapterLoadError("MOONDREAM_API_KEY is required to load finetune adapters.") |
|
|
| headers = {"User-Agent": "moondream-torch", "X-Moondream-Auth": api_key} |
| url = f"{endpoint}/v1/tuning/finetunes/{finetune_id}/checkpoints/{step}/download" |
| req = Request(url, headers=headers) |
| try: |
| with urlopen(req) as r: |
| payload = json.loads(r.read().decode("utf-8")) |
| except Exception as e: |
| raise AdapterLoadError(f"Failed to fetch adapter URL: {e}") from e |
|
|
| presigned = payload.get("url") |
| if not presigned: |
| raise AdapterLoadError("Adapter URL response missing 'url' field.") |
| return presigned |
|
|
|
|
| def cached_adapter_path(adapter_id: str) -> Path: |
| finetune_id, step = parse_adapter_id(adapter_id) |
|
|
| cache_dir = adapter_cache_dir() / finetune_id / step |
| cache_dir.mkdir(parents=True, exist_ok=True) |
|
|
| for name in ("adapter.pt", "adapter.safetensors"): |
| path = cache_dir / name |
| if path.exists() and path.stat().st_size > 0: |
| return path |
|
|
| presigned_url = _fetch_presigned_url(finetune_id, step) |
| dest = cache_dir / "adapter.pt" |
|
|
| try: |
| with urlopen(presigned_url) as r, open(dest, "wb") as f: |
| shutil.copyfileobj(r, f) |
| except Exception as e: |
| raise AdapterLoadError(f"Failed to download adapter: {e}") from e |
| return dest |
|
|
|
|
| def _load_state_dict(path: Path, device: torch.device) -> Dict[str, Any]: |
| if path.suffix == ".safetensors": |
| try: |
| from safetensors.torch import safe_open |
| except Exception as e: |
| raise AdapterLoadError( |
| "safetensors is required to load .safetensors adapters." |
| ) from e |
| data = {} |
| with safe_open(str(path), framework="pt") as f: |
| for key in f.keys(): |
| data[key] = f.get_tensor(key).to(device=device) |
| return data |
|
|
| try: |
| return torch.load(path, map_location=device, weights_only=True) |
| except TypeError: |
| return torch.load(path, map_location=device) |
|
|
|
|
| @dataclass |
| class DenseLoRALayer: |
| up_a: torch.Tensor |
| up_b: torch.Tensor |
| down_a: torch.Tensor |
| down_b: torch.Tensor |
|
|
|
|
| @dataclass |
| class MoELoRALayer: |
| up_a: torch.Tensor |
| up_b: torch.Tensor |
| down_a: torch.Tensor |
| down_b: torch.Tensor |
|
|
|
|
| class TextLoRA: |
| def __init__( |
| self, |
| text_config: TextConfig, |
| *, |
| rank: int, |
| max_rank: int, |
| dtype: torch.dtype, |
| device: torch.device, |
| adapter_id: Optional[str] = None, |
| ) -> None: |
| if rank <= 0: |
| raise AdapterLoadError("LoRA rank must be positive.") |
| if max_rank < rank: |
| raise AdapterLoadError("max_rank must be >= rank.") |
|
|
| self.text_config = text_config |
| self.rank = rank |
| self.max_rank = max_rank |
| self.adapter_id = adapter_id |
|
|
| moe_cfg = text_config.moe |
| self.start_layer = moe_cfg.start_layer if moe_cfg else text_config.n_layers |
|
|
| if moe_cfg is not None: |
| self.rank_per_expert = rank // moe_cfg.experts_per_token |
| if self.rank_per_expert < 1: |
| raise AdapterLoadError( |
| f"rank ({rank}) must be >= experts_per_token ({moe_cfg.experts_per_token})" |
| ) |
| self.max_rank_per_expert = max_rank // moe_cfg.experts_per_token |
| if self.max_rank_per_expert < 1: |
| raise AdapterLoadError( |
| f"max_rank ({max_rank}) must be >= experts_per_token ({moe_cfg.experts_per_token})" |
| ) |
| else: |
| self.rank_per_expert = 0 |
| self.max_rank_per_expert = 0 |
|
|
| d_model = text_config.dim |
| d_ffn = text_config.ff_dim |
|
|
| self.dense: list[DenseLoRALayer] = [] |
| for _ in range(self.start_layer): |
| self.dense.append( |
| DenseLoRALayer( |
| up_a=torch.zeros((max_rank, d_model), device=device, dtype=dtype), |
| up_b=torch.zeros((d_ffn, max_rank), device=device, dtype=dtype), |
| down_a=torch.zeros((max_rank, d_ffn), device=device, dtype=dtype), |
| down_b=torch.zeros((d_model, max_rank), device=device, dtype=dtype), |
| ) |
| ) |
|
|
| self.moe: list[MoELoRALayer] = [] |
| if moe_cfg is not None: |
| num_experts = moe_cfg.num_experts |
| d_expert = moe_cfg.expert_inner_dim |
| for _ in range(text_config.n_layers - self.start_layer): |
| self.moe.append( |
| MoELoRALayer( |
| up_a=torch.zeros( |
| (num_experts, self.max_rank_per_expert, d_model), |
| device=device, |
| dtype=dtype, |
| ), |
| up_b=torch.zeros( |
| (num_experts, d_expert * 2, self.max_rank_per_expert), |
| device=device, |
| dtype=dtype, |
| ), |
| down_a=torch.zeros( |
| (num_experts, self.max_rank_per_expert, d_expert), |
| device=device, |
| dtype=dtype, |
| ), |
| down_b=torch.zeros( |
| (num_experts, d_model, self.max_rank_per_expert), |
| device=device, |
| dtype=dtype, |
| ), |
| ) |
| ) |
|
|
| def dense_layer(self, layer_idx: int) -> Optional[DenseLoRALayer]: |
| if layer_idx < len(self.dense): |
| return self.dense[layer_idx] |
| return None |
|
|
| def moe_layer(self, layer_idx: int) -> Optional[MoELoRALayer]: |
| moe_idx = layer_idx - self.start_layer |
| if 0 <= moe_idx < len(self.moe): |
| return self.moe[moe_idx] |
| return None |
|
|
| @staticmethod |
| def _pad_axis(tensor: torch.Tensor, target: int, axis: int) -> torch.Tensor: |
| if tensor.shape[axis] == target: |
| return tensor |
| if tensor.shape[axis] > target: |
| raise AdapterLoadError( |
| f"LoRA tensor rank {tensor.shape[axis]} exceeds max {target}" |
| ) |
| pad_shape = list(tensor.shape) |
| pad_shape[axis] = target - tensor.shape[axis] |
| pad = torch.zeros(pad_shape, device=tensor.device, dtype=tensor.dtype) |
| return torch.cat([tensor, pad], dim=axis) |
|
|
| @staticmethod |
| def detect_rank(state_dict: Dict[str, Any], text_config: TextConfig) -> int: |
| for key, tensor in state_dict.items(): |
| if "dense" in key and "up_a" in key: |
| return int(tensor.shape[0]) |
| for key, tensor in state_dict.items(): |
| if "moe" in key and "up_a" in key: |
| rank_per_expert = int(tensor.shape[1]) |
| moe_cfg = text_config.moe |
| if moe_cfg: |
| return rank_per_expert * moe_cfg.experts_per_token |
| return rank_per_expert |
| raise AdapterLoadError("Could not detect LoRA rank from state dict.") |
|
|
| @classmethod |
| def from_state_dict( |
| cls, |
| state_dict: Dict[str, Any], |
| *, |
| text_config: TextConfig, |
| max_rank: int, |
| dtype: torch.dtype, |
| device: torch.device, |
| adapter_id: Optional[str] = None, |
| ) -> "TextLoRA": |
| rank = cls.detect_rank(state_dict, text_config) |
| if rank > max_rank: |
| raise AdapterLoadError( |
| f"Adapter rank ({rank}) exceeds max_rank ({max_rank})." |
| ) |
|
|
| lora = cls( |
| text_config, |
| rank=rank, |
| max_rank=max_rank, |
| dtype=dtype, |
| device=device, |
| adapter_id=adapter_id, |
| ) |
|
|
| dense_seen = set() |
| moe_seen = set() |
|
|
| pattern = re.compile(r"(dense|moe)\.(\d+)\.(up_a|up_b|down_a|down_b)$") |
| for key, tensor in state_dict.items(): |
| match = pattern.search(key) |
| if not match: |
| continue |
| kind, idx_str, name = match.group(1), match.group(2), match.group(3) |
| idx = int(idx_str) |
| arr = tensor.to(device=device, dtype=dtype) |
|
|
| if kind == "dense": |
| if idx >= len(lora.dense): |
| raise AdapterLoadError(f"Dense LoRA layer index {idx} out of range.") |
| layer = lora.dense[idx] |
| if name in ("up_a", "down_a"): |
| arr = cls._pad_axis(arr, lora.max_rank, axis=0) |
| else: |
| arr = cls._pad_axis(arr, lora.max_rank, axis=1) |
| setattr(layer, name, arr) |
| dense_seen.add((idx, name)) |
| else: |
| if idx >= len(lora.moe): |
| raise AdapterLoadError(f"MoE LoRA layer index {idx} out of range.") |
| layer = lora.moe[idx] |
| if name in ("up_a", "down_a"): |
| arr = cls._pad_axis(arr, lora.max_rank_per_expert, axis=1) |
| else: |
| arr = cls._pad_axis(arr, lora.max_rank_per_expert, axis=2) |
| setattr(layer, name, arr) |
| moe_seen.add((idx, name)) |
|
|
| for layer_idx in range(len(lora.dense)): |
| for name in ("up_a", "up_b", "down_a", "down_b"): |
| if (layer_idx, name) not in dense_seen: |
| raise AdapterLoadError( |
| f"Adapter missing dense LoRA for layer {layer_idx} ({name})." |
| ) |
| for layer_idx in range(len(lora.moe)): |
| for name in ("up_a", "up_b", "down_a", "down_b"): |
| if (layer_idx, name) not in moe_seen: |
| raise AdapterLoadError( |
| f"Adapter missing MoE LoRA for layer {layer_idx} ({name})." |
| ) |
|
|
| return lora |
|
|
|
|
| def select_layer_lora( |
| lora: Optional[TextLoRA], layer_idx: int, *, is_moe: bool |
| ) -> Optional[object]: |
| if lora is None: |
| return None |
| return lora.moe_layer(layer_idx) if is_moe else lora.dense_layer(layer_idx) |
|
|
|
|
| def apply_dense_lora( |
| x: torch.Tensor, lora_a: torch.Tensor, lora_b: torch.Tensor |
| ) -> torch.Tensor: |
| b, t, c = x.shape |
| x_flat = x.reshape(-1, c) |
| lora_mid = torch.matmul(x_flat, lora_a.t()) |
| lora_out = torch.matmul(lora_mid, lora_b.t()) |
| return lora_out.reshape(b, t, -1) |
|
|
|
|
| def apply_moe_lora_fc1_flat( |
| x_expanded: torch.Tensor, lora: MoELoRALayer, flat_idxs: torch.Tensor |
| ) -> torch.Tensor: |
| lora_up_a = lora.up_a[flat_idxs] |
| lora_up_b = lora.up_b[flat_idxs] |
| lora_mid = torch.bmm(lora_up_a, x_expanded.unsqueeze(-1)).squeeze(-1) |
| lora_up = torch.bmm(lora_up_b, lora_mid.unsqueeze(-1)).squeeze(-1) |
| return lora_up |
|
|
|
|
| def apply_moe_lora_fc2_flat( |
| h: torch.Tensor, lora: MoELoRALayer, flat_idxs: torch.Tensor |
| ) -> torch.Tensor: |
| lora_down_a = lora.down_a[flat_idxs] |
| lora_down_b = lora.down_b[flat_idxs] |
| lora_mid = torch.bmm(lora_down_a, h.unsqueeze(-1)).squeeze(-1) |
| lora_down = torch.bmm(lora_down_b, lora_mid.unsqueeze(-1)).squeeze(-1) |
| return lora_down |
|
|
|
|
| _ADAPTER_CACHE: Dict[Tuple[str, str, str, Tuple], TextLoRA] = {} |
| _CACHE_ORDER: list[Tuple[str, str, str, Tuple]] = [] |
| _CACHE_SIZE = 8 |
|
|
|
|
| def _config_key(text_config: TextConfig) -> Tuple: |
| moe = text_config.moe |
| moe_key = None |
| if moe is not None: |
| moe_key = ( |
| moe.num_experts, |
| moe.start_layer, |
| moe.experts_per_token, |
| moe.expert_inner_dim, |
| ) |
| return ( |
| text_config.dim, |
| text_config.ff_dim, |
| text_config.n_layers, |
| moe_key, |
| ) |
|
|
|
|
| def load_adapter( |
| adapter_id: Optional[str], |
| *, |
| text_config: TextConfig, |
| device: torch.device, |
| dtype: torch.dtype, |
| max_rank: int = 16, |
| ) -> Optional[TextLoRA]: |
| if adapter_id is None: |
| return None |
|
|
| adapter_id = normalize_adapter_id(adapter_id) |
| if adapter_id is None: |
| return None |
|
|
| key = (adapter_id, str(device), str(dtype), _config_key(text_config)) |
| cached = _ADAPTER_CACHE.get(key) |
| if cached is not None: |
| return cached |
|
|
| path = cached_adapter_path(adapter_id) |
| checkpoint = _load_state_dict(path, device) |
| if not isinstance(checkpoint, dict): |
| raise AdapterLoadError("Invalid adapter checkpoint format.") |
|
|
| state_dict = checkpoint.get("lora_state_dict", checkpoint) |
| if not isinstance(state_dict, dict): |
| raise AdapterLoadError("Adapter checkpoint missing lora_state_dict.") |
|
|
| lora = TextLoRA.from_state_dict( |
| state_dict, |
| text_config=text_config, |
| max_rank=max_rank, |
| dtype=dtype, |
| device=device, |
| adapter_id=adapter_id, |
| ) |
|
|
| _ADAPTER_CACHE[key] = lora |
| _CACHE_ORDER.append(key) |
| if len(_CACHE_ORDER) > _CACHE_SIZE: |
| old = _CACHE_ORDER.pop(0) |
| _ADAPTER_CACHE.pop(old, None) |
|
|
| return lora |
|
|
