| """Lossless separate-expert layout adapter for pinned official Qwen3.5 MoE. |
| |
| Usage: with native_layout_adapter() as audit: mlx_lm.convert.convert(...) |
| The context changes only source layout handling and restores the original method. |
| All quantization, model math, serialization and reload remain official MLX code. |
| """ |
| import contextlib,importlib.metadata,json,pathlib,re,struct,threading |
| _PATTERN=re.compile(r'^(model\.language_model|language_model\.model)\.layers\.(\d+)\.mlp\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$') |
| _LOCK=threading.Lock() |
|
|
| def _plan(weights,layers,experts,hidden,intermediate): |
| found={};recognized=set() |
| for key,value in weights.items(): |
| match=_PATTERN.fullmatch(key) |
| if match: |
| base,layer,expert,proj=match.groups();layer=int(layer);expert=int(expert) |
| assert key==f'{base}.layers.{layer}.mlp.experts.{expert}.{proj}.weight','Noncanonical expert index' |
| assert 0<=layer<layers and 0<=expert<experts,('Out-of-range expert key',key) |
| shape=tuple(value.shape) if hasattr(value,'shape') else tuple(value['shape']) |
| dtype=str(value.dtype) if hasattr(value,'dtype') else value['dtype'] |
| required=(hidden,intermediate) if proj=='down_proj' else (intermediate,hidden) |
| assert shape==required,('Expert shape mismatch',key,shape,required) |
| found.setdefault((base,layer,proj),{})[expert]=(key,shape,dtype);recognized.add(key) |
| elif re.match(r'^(model\.language_model|language_model\.model)\.layers\.\d+\.mlp\.experts\.',key) and re.search(r'\.experts\.\d+\.',key): |
| raise ValueError('Unexpected independent expert key: '+key) |
| if not found:return [] |
| bases={g[0] for g in found};assert len(bases)==1,'Mixed expert key roots';base=next(iter(bases)) |
| assert len(found)==layers*3,'Missing layer or projection group' |
| plan=[] |
| for layer in range(layers): |
| prefix=f'{base}.layers.{layer}.mlp' |
| assert not any(k.startswith(prefix+'.switch_mlp.') or k in (prefix+'.experts.gate_up_proj',prefix+'.experts.down_proj') for k in weights),'Mixed packed/independent expert layouts' |
| for proj in ('gate_proj','up_proj','down_proj'): |
| group=found[(base,layer,proj)];assert set(group)==set(range(experts)),('Missing/extra experts',layer,proj) |
| assert len({v[2] for v in group.values()})==1,('Mixed expert dtypes',layer,proj) |
| plan.append((prefix+'.switch_mlp.'+proj+'.weight',[group[i][0] for i in range(experts)])) |
| return plan |
|
|
| def stack_independent_experts(weights,args): |
| import mlx.core as mx |
| plan=_plan(weights,args.num_hidden_layers,args.num_experts,args.hidden_size,args.moe_intermediate_size) |
| if not plan:return weights,{'adapted':False,'groups':0,'source_expert_tensors':0} |
| result=dict(weights) |
| for destination,keys in plan: |
| |
| result[destination]=mx.stack([result.pop(k) for k in keys],axis=0) |
| assert len(result)==len(weights)-sum(len(keys)-1 for _,keys in plan) |
| return result,{'adapted':True,'groups':len(plan),'source_expert_tensors':sum(len(k) for _,k in plan),'expert_order':'numeric 0..num_experts-1','value_transform':'stack only; no cast, quantization, norm or conv transform'} |
|
|
| @contextlib.contextmanager |
| def native_layout_adapter(): |
| assert importlib.metadata.version('mlx-lm')=='0.31.3','Adapter audited for mlx-lm 0.31.3 only' |
| from mlx_lm.models.qwen3_5_moe import Model |
| if not _LOCK.acquire(blocking=False):raise RuntimeError('Layout adapter already active in this process') |
| original=Model.sanitize;audit={'calls':[],'native_sanitize_calls':0,'restored':False} |
| def sanitize(self,weights): |
| transformed,row=stack_independent_experts(weights,self.language_model.args) |
| audit['calls'].append(row);audit['native_sanitize_calls']+=1 |
| return original(self,transformed) |
| try:Model.sanitize=sanitize;yield audit |
| finally:Model.sanitize=original;audit['restored']=True;_LOCK.release() |
|
|
| def source_header_preflight(source): |
| source=pathlib.Path(source);config=json.loads((source/'config.json').read_text())['text_config'];index=json.loads((source/'model.safetensors.index.json').read_text())['weight_map'];weights={} |
| for name in sorted(set(index.values())): |
| assert pathlib.Path(name).name==name |
| with (source/name).open('rb') as f: |
| size=struct.unpack('<Q',f.read(8))[0];assert size<100_000_000;header=json.loads(f.read(size)) |
| for key,value in header.items(): |
| if key=='__metadata__':continue |
| assert key not in weights and index.get(key)==name;weights[key]=value |
| assert set(weights)==set(index) |
| plan=_plan(weights,config['num_hidden_layers'],config['num_experts'],config['hidden_size'],config['moe_intermediate_size']) |
| assert len(plan)==120 and sum(len(keys) for _,keys in plan)==30720 |
| return {'passed':True,'source_index_tensors':len(weights),'layers':config['num_hidden_layers'],'experts_per_layer':config['num_experts'],'groups':len(plan),'independent_expert_tensors':30720,'all_shapes_checked':True,'each_group_dtype_uniform':True,'weight_payload_read':False,'model_instantiated':False} |
|
|
