62 lines
2.8 KiB
Python
62 lines
2.8 KiB
Python
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import shutil
|
|
import sys
|
|
|
|
import torch
|
|
from peft import PeftModel
|
|
from safetensors import safe_open
|
|
from transformers import AutoModelForCausalLM
|
|
|
|
base, checkpoint, target = (Path(x) for x in sys.argv[1:])
|
|
base_revision = '2c95908929198d6e69af8638f0dbbd9bc6b93f9e'
|
|
pinned = Path.home()/'.cache/huggingface/hub/models--ministral--Ministral-3b-instruct/snapshots'/base_revision
|
|
stage = target.with_name(target.name + '_staging_' + os.environ['SLURM_JOB_ID'])
|
|
|
|
def digest(path):
|
|
with path.open('rb') as handle:
|
|
return hashlib.file_digest(handle, 'sha256').hexdigest()
|
|
|
|
if target.exists():
|
|
raise SystemExit('Merged output already exists')
|
|
index = 'model.safetensors.index.json'
|
|
metadata = ('config.json','generation_config.json','special_tokens_map.json','tokenizer_config.json','tokenizer.json')
|
|
original_keys = json.loads((base/index).read_text())['weight_map']
|
|
for name in (*metadata, index, *sorted(set(original_keys.values()))):
|
|
if digest(base/name) != digest(pinned/name):
|
|
raise RuntimeError(f'Baseline file differs from the pinned training snapshot: {name}')
|
|
print('Baseline model matches the pinned training snapshot', base_revision, flush=True)
|
|
|
|
torch.set_num_threads(8)
|
|
model = AutoModelForCausalLM.from_pretrained(str(base), dtype=torch.float32, local_files_only=True, trust_remote_code=False)
|
|
lora = PeftModel.from_pretrained(model, str(checkpoint), is_trainable=False)
|
|
with torch.no_grad():
|
|
merged = lora.merge_and_unload(safe_merge=True).to(torch.bfloat16)
|
|
stage.mkdir()
|
|
merged.save_pretrained(str(stage), safe_serialization=True, max_shard_size='3GB')
|
|
for name in metadata:
|
|
shutil.copy2(base/name, stage/name)
|
|
|
|
saved_keys = json.loads((stage/index).read_text())['weight_map']
|
|
if set(saved_keys) != set(original_keys) or any('lora_' in key for key in saved_keys):
|
|
raise RuntimeError('Unexpected merged parameter set')
|
|
for key, expected_to_change in (
|
|
('model.layers.0.self_attn.q_proj.weight', True),
|
|
('model.embed_tokens.weight', False),
|
|
):
|
|
with safe_open(str(base/original_keys[key]), framework='pt', device='cpu') as left:
|
|
with safe_open(str(stage/saved_keys[key]), framework='pt', device='cpu') as right:
|
|
original, merged_tensor = left.get_tensor(key), right.get_tensor(key)
|
|
changed = not torch.equal(original, merged_tensor)
|
|
if merged_tensor.dtype != torch.bfloat16 or changed != expected_to_change:
|
|
raise RuntimeError(f'Unexpected merged tensor: {key}')
|
|
(stage/'merge_provenance.json').write_text(json.dumps({
|
|
'checkpoint': str(checkpoint), 'adapter_sha256': digest(checkpoint/'adapter_model.safetensors'),
|
|
'base_revision': base_revision, 'output_dtype': 'bfloat16',
|
|
'tokenizer_and_config': {name: digest(stage/name) for name in metadata},
|
|
}, indent=2) + chr(10))
|
|
stage.rename(target)
|
|
print('MERGED_MODEL_READY', target, flush=True)
|