Blame for runtime/make_expert_map.py

"""Build a machine-local exact-expert map from GGUF tensor metadata."""
from pathlib import Path
import argparse,struct,sys

def main():
    p=argparse.ArgumentParser();p.add_argument('model',type=Path);p.add_argument('output',type=Path);p.add_argument('--llama-source',type=Path,default=Path('llama.cpp'));a=p.parse_args()
    sys.path.insert(0,str(a.llama_source.resolve()/'gguf-py'))
    import gguf
    model=a.model.resolve();st=model.stat();reader=gguf.GGUFReader(model);ts={t.name:t for t in reader.tensors}
    payload=bytearray(b'WEEPF001'+struct.pack('<QQQQI',st.st_size,st.st_dev,st.st_ino,st.st_mtime_ns,48));count=0
    for layer in range(48):
        experts=[]
        for tier in range(4):
            suffix='' if tier==0 else f'_tier{tier}'
            names=[f'blk.{layer}.ffn_{kind}_exps{suffix}.weight' for kind in ('up','gate','down')]
            if names[0] not in ts:
                assert all(n not in ts for n in names);continue
            tensors=[ts[n] for n in names];n=int(tensors[0].shape[2])
            assert all(len(t.shape)==3 and int(t.shape[2])==n and t.n_bytes%n==0 for t in tensors)
            for e in range(n):
                ranges=[]
                for t in tensors:
                    length=int(t.n_bytes)//n;offset=int(t.data_offset)+e*length
                    assert 0<length<=8*2**20 and 0<=offset<offset+length<=st.st_size
                    ranges.append((offset,length))
                experts.append(ranges)
        assert len(experts)==int(ts[f'blk.{layer}.ffn_gate_inp.weight'].shape[1]);count+=len(experts)
        payload+=struct.pack('<I',len(experts))
        for expert in experts:
            for offset,length in expert:payload+=struct.pack('<QQ',offset,length)
    a.output.write_bytes(payload);print(f'Mapped {count} expert slots into {a.output}')

if __name__=='__main__':main()