File size: 2,426 Bytes
ff7679d | 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 | 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):
# Preserve loader sets as sorted lists before report/progress serialization.
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)
|