File size: 3,419 Bytes
a4ba82a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | from collections import Counter
import hashlib
from pathlib import Path
import types
def digest(path):
with Path(path).open('rb') as stream:
return hashlib.file_digest(stream, 'sha256').hexdigest()
def load(spec, save):
import torch
from transformers import AutoModelForImageTextToText
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5PreTrainedModel
from safetensors.torch import load_file
from peft import PeftConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict
from production_gdn_fix import _keep_native_gdn_fp32
from verify_gdn import verify
_keep_native_gdn_fp32(Qwen3_5PreTrainedModel)
base, info = AutoModelForImageTextToText.from_pretrained(spec['base'], local_files_only=True, dtype=torch.bfloat16, attn_implementation='sdpa', output_loading_info=True)
assert not any((info.get(k) for k in ['missing_keys', 'unexpected_keys', 'mismatched_keys', 'error_msgs'])), info
path = Path(spec['export']) / 'adapter_model.safetensors'
assert digest(path) == spec['adapter_sha256']
source = load_file(str(path), device='cpu')
assert len(source) == 496 and sum((v.numel() for v in source.values())) == 86556672
assert all((v.dtype == torch.float32 for v in source.values())), 'Expected corrected FP32 checkpoint'
model = get_peft_model(base, PeftConfig.from_pretrained(spec['export'], local_files_only=True), adapter_name='candidate', autocast_adapter_dtype=True)
initial = get_peft_model_state_dict(model, adapter_name='candidate')
assert set(initial) == set(source)
assert all((v.dtype == torch.float32 for v in initial.values())), 'PEFT must create FP32 adapter holders'
del initial
load_info = set_peft_model_state_dict(model, source, adapter_name='candidate')
assert not load_info.unexpected_keys, load_info
actual = get_peft_model_state_dict(model, adapter_name='candidate')
assert set(actual) == set(source)
assert all((actual[k].dtype == v.dtype and torch.equal(actual[k].detach().cpu(), v) for k, v in source.items())), 'Saved FP32 LoRA values must be loaded exactly'
del actual
model.to('cuda:0').eval()
actual = get_peft_model_state_dict(model, adapter_name='candidate')
assert all((actual[k].dtype == torch.float32 and torch.equal(actual[k].detach().cpu(), v) for k, v in source.items())), 'Final device placement must preserve saved FP32 values'
save('ADAPTER_VALUE_VERIFICATION.json', dict(tensors=496, elements=86556672, source_dtypes=dict(Counter((str(v.dtype) for v in source.values()))), holders_fp32_before_assignment=True, all_saved_values_exact_before_and_after_placement=True, adapter_sha256=spec['adapter_sha256'], merged=False))
del source, actual
gdn = verify(model, spec['base'], peft_prefix=True)
save('GDN_VALUE_VERIFICATION.json', gdn)
original_generate = model.generate
def generate_bf16(self, *args, **kwargs):
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
return original_generate(*args, **kwargs)
model.generate = types.MethodType(generate_bf16, model)
save('MODEL_LOADED.json', dict(adapter_persistent_dtype='torch.float32', compute_autocast='torch.bfloat16', gdn_values_exact=True, base_loading_info=info, parameter_dtypes=dict(Counter((str(p.dtype) for p in model.parameters()))), input_output_generation_policy_changed=False))
return model
|