popov/neuron-tai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Track Kimi model metadata and cache path

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-01 12:38:31 +02:00
parent 78834e5045
commit bc760c1694
7 changed files with 231 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
packages/node/meshnet_node/model_backend.py
View File

@@ -4,6 +4,7 @@ from __future__ import annotations
import base64
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Literal
Quantization = Literal["bfloat16", "int8", "nf4"]
@@ -65,6 +66,7 @@ class TorchModelShard:
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
cache_dir: Path | None = None,
) -> None:
if shard_start < 0 or shard_end < 0 or shard_start > shard_end:
raise ValueError("shard_start must be <= shard_end and non-negative")
@@ -89,9 +91,9 @@ class TorchModelShard:
model_id,
quantization_config=quant_config,
device_map="auto" if quant_config is not None else None,
torch_dtype=torch.bfloat16,
dtype=torch.bfloat16,
low_cpu_mem_usage=True,
use_safetensors=True,
cache_dir=str(cache_dir) if cache_dir is not None else None,
)
if quant_config is None:
self.model.to(self.device)
@@ -104,7 +106,10 @@ class TorchModelShard:
raise
self.model.eval()
self.tokenizer = AutoTokenizer.from_pretrained(model_id)
self.tokenizer = AutoTokenizer.from_pretrained(
model_id,
cache_dir=str(cache_dir) if cache_dir is not None else None,
)
self.layers = _model_layers(self.model)
self.total_layers = len(self.layers)
# shard_end is INCLUSIVE (last layer index, 0-based), matching the CLI convention.
@@ -336,8 +341,9 @@ def load_torch_shard(
shard_start: int,
shard_end: int,
quantization: Quantization = "bfloat16",
cache_dir: Path | None = None,
) -> TorchModelShard:
return TorchModelShard(model_id, shard_start, shard_end, quantization)
return TorchModelShard(model_id, shard_start, shard_end, quantization, cache_dir)
def _model_layers(model: Any) -> Any: