cmedina/agenx-lora-training
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: use torch.load for quantized shards index creation

Browse Source
This commit is contained in:
Christian Medina
2026-07-03 03:58:24 -04:00
parent 5078d7c8dc
commit 1ebad2365c
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
train.py
View File

@@ -36,7 +36,6 @@ def train(config_path):
# Create index file for sharded model # Create index file for sharded model
import glob as glob_mod import glob as glob_mod
import json as json_mod import json as json_mod
from safetensors import safe_open
safetensor_files = glob_mod.glob(f"{config['base_model']}/*.safetensors") safetensor_files = glob_mod.glob(f"{config['base_model']}/*.safetensors")
shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name]) shards = sorted([Path(f).name for f in safetensor_files if "of-" in Path(f).name])
@@ -47,8 +46,9 @@ def train(config_path):
weight_map = {} weight_map = {}
for shard_name in shards: for shard_name in shards:
shard_path = Path(config["base_model"]) / shard_name shard_path = Path(config["base_model"]) / shard_name
with safe_open(str(shard_path), framework="pt", device="cpu") as f: ckpt = torch.load(str(shard_path), map_location="cpu", weights_only=False)
for key in f.keys(): for key in ckpt.keys():
if isinstance(ckpt[key], torch.Tensor):
weight_map[key] = shard_name weight_map[key] = shard_name
index = { index = {