File size: 3,531 Bytes
91c1e9e | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Load and export PRIMA trainable state without changing parameter names."""
from pathlib import Path
import torch
PREFIX = 'base_model.model.'
REQUIRED_GROUPS = (
'model.layers.', 'model.grounding_encoder.mask_decoder.',
'model.query_tokens', 'model.residual_query_tokens',
'model.residual_extractor.', 'model.mm_projector.', 'model.text_hidden_fcs.',
)
def load_checkpoint(model, path, *, require_complete=True):
payload = torch.load(path, map_location='cpu')
state = payload.get('state_dict', payload) if isinstance(payload, dict) else payload
if not isinstance(state, dict) or not state:
raise ValueError('Checkpoint must contain a nonempty tensor state dictionary.')
canonical = {}
for key, value in state.items():
name = key[len(PREFIX):] if key.startswith(PREFIX) else key
if name in canonical:
raise ValueError(f'Checkpoint prefix normalization produced a duplicate: {name}')
if not torch.is_tensor(value):
raise ValueError(f'Checkpoint entry is not a tensor: {name}')
canonical[name] = value
current = model.state_dict()
wrapped = any(key.startswith(PREFIX) for key in current)
normalized = {(PREFIX + key if wrapped else key): value for key, value in canonical.items()}
# These are recomputable buffers from supported dependency versions.
ignored = [key for key in normalized if key not in current and
('rotary_emb.inv_freq' in key or key.endswith('position_ids'))]
for key in ignored:
normalized.pop(key)
unexpected = set(normalized) - set(current)
if unexpected:
raise ValueError('Unexpected checkpoint tensors: ' + ', '.join(sorted(unexpected)))
if require_complete:
missing_groups = [group for group in REQUIRED_GROUPS
if not any(key.startswith(group) for key in canonical)]
required = {key for key in current
if any((key[len(PREFIX):] if wrapped else key).startswith(group)
for group in REQUIRED_GROUPS)
and ('lora_' in key or 'model.layers.' not in key)
and 'rotary_emb.inv_freq' not in key}
missing = required - set(normalized)
if missing_groups or missing:
raise ValueError('Incomplete PRIMA checkpoint. Missing groups: '
+ ', '.join(missing_groups) + '; missing tensors: '
+ ', '.join(sorted(missing)))
model.load_state_dict(normalized, strict=False)
return {'loaded_tensors': len(normalized), 'ignored_buffers': ignored}
def export_checkpoint(model, path):
"""Export trainable parameters plus persistent buffers, with compatible keys."""
state = {}
for key, value in model.state_dict().items():
module_path, _, leaf = key.rpartition('.')
parent = model.get_submodule(module_path) if module_path else model
parameter = parent._parameters.get(leaf)
if parameter is not None and not parameter.requires_grad:
continue
name = key[len(PREFIX):] if key.startswith(PREFIX) else key
if name in state:
raise ValueError(f'Export normalization produced a duplicate: {name}')
state[name] = value.detach().cpu().clone()
destination = Path(path)
destination.parent.mkdir(parents=True, exist_ok=True)
temporary = destination.with_name(destination.name + '.tmp')
torch.save(state, temporary)
temporary.replace(destination)
|