| """Verified packed projection loading and lossless vLLM shard assembly.""" |
| from __future__ import annotations |
| import hashlib |
| import json |
| from pathlib import Path |
| import re |
| import torch |
| from safetensors import safe_open |
| from adapters.ouro import PAPER_GROUPS |
| from .export import load_export_artifact |
| from .packed_weight import PackedLoopQWeight |
| from .paths import pinned_snapshot, REVISIONS |
|
|
|
|
| def _sha256(path): |
| digest = hashlib.sha256() |
| with Path(path).open('rb') as stream: |
| for block in iter(lambda: stream.read(1024*1024), b''): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def concatenate_packed_rows(shards): |
| """Preserve per-row codes/scales while assembling Q,K,V or gate,up order.""" |
| if not shards: |
| raise ValueError('cannot concatenate an empty shard list') |
| for shard in shards: |
| shard.validate() |
| first = shards[0] |
| if any(s.shape[1] != first.shape[1] or s.output_dtype != first.output_dtype |
| or s.codes.device != first.codes.device or s.scales.device != first.scales.device for s in shards): |
| raise ValueError('packed shards must share input width, dtype and device') |
| return PackedLoopQWeight(torch.cat([s.codes for s in shards], dim=0), |
| torch.cat([s.scales for s in shards], dim=0), |
| (sum(s.shape[0] for s in shards), first.shape[1]), first.output_dtype) |
|
|
|
|
| def load_packed_ouro_bundle(directory, component_artifact, *, allow_diagnostic=False): |
| """Return group/loop packed matrices; do not allocate dense model weights. |
| |
| Keys are (canonical group, None) for shared weights and (group, loop) for |
| selected variants. Files and component identity are verified before use. |
| """ |
| root = Path(directory).resolve() |
| manifest = json.loads((root/'manifest.json').read_text()) |
| if manifest.get('format') != 'loopq_packed_ouro_projections' or manifest.get('format_version') != 1: |
| raise ValueError('unsupported packed Ouro manifest') |
| diagnostic = manifest.get('diagnostic') |
| if type(diagnostic) is not bool or manifest.get('status') != ('diagnostic_exported' if diagnostic else 'exported'): |
| raise ValueError('packed export is incomplete or mislabeled') |
| if diagnostic and not allow_diagnostic: |
| raise ValueError('diagnostic packed bundle requires allow_diagnostic=True') |
| if _sha256(component_artifact) != manifest.get('component_sha256'): |
| raise ValueError('packed bundle component artifact hash mismatch') |
| component = load_export_artifact(component_artifact) |
| if component['model'] != manifest.get('model') or component['model'].get('revision') != REVISIONS['ouro'][1]: |
| raise ValueError('packed bundle model mismatch') |
| calibration = component.get('calibration', {}) |
| if calibration.get('completed') is not True or (not diagnostic and calibration.get('paper_calibration') is not True): |
| raise ValueError('packed bundle calibration eligibility mismatch') |
| source = pinned_snapshot('ouro')/'model.safetensors' |
| if _sha256(source) != manifest.get('backbone_sha256'): |
| raise ValueError('packed bundle backbone hash mismatch') |
| groups = {f'model.layers.{layer}.{group}' for layer in range(24) for group in PAPER_GROUPS} |
| components = component['components'] |
| if set(components['shared_transforms']) != groups: |
| raise ValueError('component must cover all Ouro groups') |
| variants = [(key, None) for key in sorted(groups)] |
| for key, loops in components['selected_loop_transforms'].items(): |
| if key not in groups or set(map(int, loops)) != set(range(4)): |
| raise ValueError('invalid selected loop coverage') |
| variants.extend((key, loop) for loop in range(4)) |
| expected = {} |
| for key, loop in variants: |
| prefix, group = key.rsplit('.',1) |
| for projection in PAPER_GROUPS[group]['hf_weights']: |
| expected[(key, loop, f'{prefix}.{projection}.weight')] = None |
| loaded = {} |
| with safe_open(source, framework='pt', device='cpu') as reader: |
| for row in manifest['weights']: |
| loop = row['loop'] |
| if loop is not None and type(loop) is not int: |
| raise ValueError('loop index must be an integer or null') |
| key = (row['group'], loop, row['source_name']) |
| if key not in expected or key in loaded: |
| raise ValueError('unexpected or duplicate packed projection') |
| if not re.fullmatch(r'weight_\d{4}\.pt', row['path']): |
| raise ValueError('invalid packed file name') |
| path = (root/row['path']).resolve() |
| if path.parent != root or _sha256(path) != row['sha256']: |
| raise ValueError('packed file path or hash mismatch') |
| packed = PackedLoopQWeight.from_state_dict(torch.load(path, weights_only=True, map_location='cpu')) |
| if (list(packed.shape) != reader.get_slice(row['source_name']).get_shape() |
| or list(packed.shape) != row['shape'] or packed.output_dtype != row['dtype'] |
| or packed.payload_bytes != row['payload_bytes']): |
| raise ValueError('packed projection metadata mismatch') |
| loaded[key] = packed |
| if set(loaded) != set(expected): |
| raise ValueError('packed bundle is missing projections') |
| if sum(x.payload_bytes for x in loaded.values()) != manifest['tensor_payload_bytes']: |
| raise ValueError('packed payload total mismatch') |
| assembled = {} |
| for key, loop in variants: |
| prefix, group = key.rsplit('.',1) |
| assembled[(key, loop)] = concatenate_packed_rows([ |
| loaded[(key, loop, f'{prefix}.{projection}.weight')] |
| for projection in PAPER_GROUPS[group]['hf_weights']]) |
| return assembled |
|
|
|
|
| class PackedProjectionDispatch: |
| """Inference-only packed residence with a transient dense GEMM weight. |
| |
| This uses the existing BF16 linear operation, not a native INT4 kernel. |
| Installation checks everything before releasing any dense parameters. |
| """ |
| def __init__(self, groups, parameters_by_group): |
| if len({id(p) for p in parameters_by_group.values()}) != len(parameters_by_group): |
| raise ValueError('packed groups must own distinct projection parameters') |
| if {key for key,loop in groups if loop is None} != set(parameters_by_group): |
| raise ValueError('packed shared group coverage mismatch') |
| prepared = {} |
| for (key,loop),packed in groups.items(): |
| if key not in parameters_by_group or (loop is not None and (type(loop) is not int or not 0 <= loop < 4)): |
| raise ValueError('invalid packed group or loop') |
| parameter = parameters_by_group[key] |
| if parameter.requires_grad: |
| raise ValueError('packed dispatch requires frozen parameters') |
| packed.validate() |
| if tuple(parameter.shape) != packed.shape or str(parameter.dtype) != packed.output_dtype: |
| raise ValueError('packed dispatch shape/dtype mismatch (requires unsharded compatible weights)') |
| prepared[(key,loop)] = PackedLoopQWeight(packed.codes.to(parameter.device), |
| packed.scales.to(parameter.device), packed.shape, packed.output_dtype) |
| for key in parameters_by_group: |
| loops = {loop for group,loop in prepared if group == key and loop is not None} |
| if loops and loops != set(range(4)): |
| raise ValueError('selected packed group must cover all four loops') |
| self.parameter_ids = {key:id(p) for key,p in parameters_by_group.items()} |
| self.groups = prepared |
| self.report = {'packed_payload_bytes':sum(p.payload_bytes for p in prepared.values()), |
| 'shared_dense_bytes_released':sum(p.numel()*p.element_size() for p in parameters_by_group.values()), |
| 'groups':len(prepared), 'native_int4_gemm':False} |
| for parameter in parameters_by_group.values(): |
| parameter.data = parameter.data.new_empty(0) |
|
|
| def __call__(self, projection, key, loop, value): |
| if torch.is_grad_enabled(): |
| raise RuntimeError('packed projection dispatch is inference-only') |
| if type(loop) is not int or not 0 <= loop < 4: |
| raise ValueError('invalid packed recurrence index') |
| if id(projection.weight) != self.parameter_ids[key]: |
| raise ValueError('projection parameter does not match packed group') |
| packed = self.groups.get((key,loop), self.groups[(key,None)]) |
| original = projection.weight.data |
| projection.weight.data = packed._dequantize_validated(device=value.device) |
| try: |
| return projection(value) |
| finally: |
| projection.weight.data = original |
|
|
