| import collections,json,pathlib |
| import torch |
|
|
|
|
| def audit_loaded_parameters(headers, params, diagnostics_path): |
| counts=collections.Counter();validated_scales=0;native_scales=0;bitcast_scales=0;checked=0 |
| result={'passed':False} |
| try: |
| assert set(params)==set(headers), (set(params)-set(headers),set(headers)-set(params)) |
| dtype_map={'BF16':torch.bfloat16,'F32':torch.float32,'F8_E4M3':torch.float8_e4m3fn,'U8':torch.uint8} |
| for name,h in headers.items(): |
| t=params[name] |
| assert list(t.shape)==h['shape'],(name,list(t.shape),h['shape']) |
| is_scale=h['dtype']=='U8' and name.endswith('.weight_scale_inv') |
| if is_scale: |
| assert t.dtype in (torch.uint8,torch.float8_e8m0fnu),(name,t.dtype) |
| validated_scales+=1 |
| native_scales+=int(t.dtype==torch.uint8) |
| bitcast_scales+=int(t.dtype==torch.float8_e8m0fnu) |
| else: |
| assert t.dtype==dtype_map[h['dtype']],(name,t.dtype,dtype_map[h['dtype']]) |
| counts[str(t.dtype)]+=1;checked+=1 |
| assert len(headers)==19434,len(headers) |
| assert counts['torch.float8_e4m3fn']==9216,dict(counts) |
| assert validated_scales==9216,(validated_scales,native_scales,bitcast_scales) |
| assert native_scales+bitcast_scales==validated_scales |
| result['passed']=True |
| except BaseException as exc: |
| result['error']=repr(exc) |
| raise |
| finally: |
| result.update(checked_tensors=checked,dtype_counts=dict(counts),validated_scale_tensors=validated_scales,native_U8_scale_tensors=native_scales,logical_U8_to_E8M0_scales=bitcast_scales) |
| pathlib.Path(diagnostics_path).write_text(json.dumps(result,indent=2)+chr(10)) |
| return result |
|
|
|
|
| def normalize_loading_info(info): |
| |
| def convert(value): |
| if isinstance(value, dict): |
| return {key:convert(item) for key,item in value.items()} |
| if isinstance(value, (set,frozenset)): |
| return sorted(convert(item) for item in value) |
| if isinstance(value, (list,tuple)): |
| return [convert(item) for item in value] |
| if value is None or isinstance(value,(str,int,float,bool)): |
| return value |
| raise TypeError(f'Unsupported loading metadata type: {type(value).__name__}') |
| return convert(info) |
|
|
