| 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 |
|
|
