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

dual billing; tracker to node model sharing

Browse Source
This commit is contained in:
Dobromir Popov
2026-07-06 17:31:11 +03:00
parent ccb69c41e3
commit 2e696be80f
14 changed files with 1092 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -85,24 +85,25 @@ class TorchModelShard:
self.torch = torch
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
load_source = str(cache_dir) if cache_dir is not None and (cache_dir / "config.json").exists() else model_id
quant_config, dtype, uses_quantized_weights = _model_load_plan(
AutoConfig,
model_id,
load_source,
quantization,
torch,
cache_dir,
None if load_source != model_id else cache_dir,
)
try:
load_kwargs = {
"device_map": "auto" if uses_quantized_weights else None,
"dtype": dtype,
"low_cpu_mem_usage": True,
"cache_dir": str(cache_dir) if cache_dir is not None else None,
"cache_dir": str(cache_dir) if cache_dir is not None and load_source == model_id else None,
}
if quant_config is not None:
load_kwargs["quantization_config"] = quant_config
self.model = AutoModelForCausalLM.from_pretrained(
model_id,
load_source,
**load_kwargs,
)
if not uses_quantized_weights:
@@ -117,8 +118,8 @@ class TorchModelShard:
self.model.eval()
self.tokenizer = AutoTokenizer.from_pretrained(
model_id,
cache_dir=str(cache_dir) if cache_dir is not None else None,
load_source,
cache_dir=str(cache_dir) if cache_dir is not None and load_source == model_id else None,
)
self.layers = _model_layers(self.model)
self.total_layers = len(self.layers)