Include clean PRIMA code and evaluator; document metrics, initialization, and MIT contributions
91c1e9e verified | """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) | |