GLM-5.3 is live. Run Z.ai's latest model on Friendli Model APIs. Try it today ➜
Jiarui-Wang
Available on FriendliAI
Run this model inference on single tenant GPU with unmatched speed and reliability at scale.
Run this model inference with full control and performance in your environment.
Model Details
Model Provider
Jiarui-Wang
Model Tree
Qwen/Qwen3-8B
Input Modalities
Output Modalities
Supported Functionality
MemSFT specializes modern large language models with an external parametric memory. This checkpoint contains the Qwen3-8B memory trained on DISC-Law-SFT. The memory learns to approximate retrieval-based teacher distributions over domain SFT data. At each decoding step, a learned token-level router combines the next-token distributions of the frozen base model and memory. This checkpoint is an auxiliary memory, not a standalone chat model. Its key advantages are:
The evaluated 14B + 8B configuration is intended for a CUDA GPU with sufficient memory to load both models in BF16.
bash
git clone https://github.com/LUMIA-Group/MemSFT.gitcd MemSFTconda create -n memsft-generate python=3.10 pip -yconda activate memsft-generatepython -m pip install -e .python -m pip install \"torch>=2.4,<2.7" \"transformers==4.51.3" \"huggingface-hub==0.35.3" \"accelerate>=0.34,<2"
python
from pathlib import Pathimport torchfrom huggingface_hub import snapshot_downloadfrom transformers import AutoModelForCausalLM, AutoTokenizer, set_seedfrom memsft.router.adaptive_memdec import AdaptiveMemoryDecoderdevice = torch.device("cuda:0")base_id = "Qwen/Qwen3-14B"memory_id = "Jiarui-Wang/MemSFT-Qwen3-Law-Memory-8B"router_repo = "Jiarui-Wang/MemSFT-Qwen3-Routers"router_subdir = "Qwen3-14B-Law-M8B-Router"router_root = snapshot_download(repo_id=router_repo,revision="v1.0.0",allow_patterns=[f"{router_subdir}/*"],)router_path = str(Path(router_root) / router_subdir)tokenizer = AutoTokenizer.from_pretrained(base_id,revision="40c069824f4251a91eefaf281ebe4c544efd3e18",)base = AutoModelForCausalLM.from_pretrained(base_id,revision="40c069824f4251a91eefaf281ebe4c544efd3e18",torch_dtype=torch.bfloat16,low_cpu_mem_usage=True,).to(device).eval()memory = AutoModelForCausalLM.from_pretrained(memory_id,revision="v1.0.0",torch_dtype=torch.bfloat16,low_cpu_mem_usage=True,).to(device).eval()vocab_size = len(tokenizer)base.resize_token_embeddings(vocab_size)memory.resize_token_embeddings(vocab_size)base.requires_grad_(False)memory.requires_grad_(False)model = AdaptiveMemoryDecoder(base_lm=base,knn_generator=memory,router_path=router_path,router_device=device,).eval()model.set_tokenizer(tokenizer)
python
instruction = ("依据给出的实体类型提取句子的实体信息，实体类型包括:犯罪嫌疑人、受害人、""被盗货币、物品价值、盗窃获利、被盗物品、作案工具、时间、地点、组织机构。""逐个列出实体信息。")question = "句子:被告人周某甲被归案。"prompt = f"{instruction}\n{question}"messages = [{"role": "user", "content": prompt}]prompt_text = tokenizer.apply_chat_template(messages,tokenize=False,add_generation_prompt=True,enable_thinking=False,)inputs = tokenizer(prompt_text, return_tensors="pt").to(device)set_seed(42)with torch.inference_mode():output_ids = model.generate(**inputs,do_sample=True,temperature=0.6,top_p=0.95,top_k=20,max_new_tokens=256,eos_token_id=tokenizer.eos_token_id,pad_token_id=tokenizer.eos_token_id,)answer = tokenizer.decode(output_ids[0, inputs["input_ids"].shape[1]:],skip_special_tokens=True,)print(answer)
Example output:
text
犯罪嫌疑人：周某甲
The output exactly matches the reference entity set. With the same prompt,
seed, and generation settings, Qwen3-14B alone instead enumerates all ten
entity types and assigns the nine unmentioned types the value 无. Both runs
ended naturally at EOS in BF16 on NVIDIA A800 80GB GPUs.
LawBench is distributed under the Apache License 2.0.
The released Law memory is evaluated with Qwen3-14B and its matching router. LawBench Avg. is computed over the 19 evaluated legal subtasks on a 0–100 scale, where higher is better. General is the average over MATH-500, C-Eval, IFEval, MMLU-Redux, and INCLUDE.
| Base model | LawBench Avg. ↑ | General ↑ |
|---|---|---|
| Qwen3-14B | 49.83 | 83.22 |
| Qwen3-14B + MemSFT | 56.47 | 83.79 |
The evaluated pairing is:
Qwen/Qwen3-14BJiarui-Wang/MemSFT-Qwen3-Law-Memory-8BJiarui-Wang/MemSFT-Qwen3-Routers/Qwen3-14B-Law-M8B-RouterMemSFT prefers tensor-only .safetensors router checkpoints. Legacy .pt
checkpoints should be loaded only from trusted sources; the MemSFT loader uses
PyTorch's restricted weights_only=True mode for compatibility.
This checkpoint is intended for reproducing MemSFT and for augmenting Qwen3-14B on the evaluated legal-domain tasks. It should be used with the matching Qwen3-14B/M8B router. Performance outside the evaluated model combination and domain has not been established. Model outputs may be incorrect or incomplete and should not be treated as legal advice. Qualified professionals should review outputs before use in consequential legal applications.
This MemSFT checkpoint is released under the Apache License 2.0. Upstream models, software, and datasets remain subject to their respective licenses and terms.
If you find MemSFT helpful in your research, please consider citing:
bibtex
@misc{wang2026memsftmitigatingalignmenttax,title={MemSFT: Mitigating Alignment Tax with an External Parametric Memory},author={Jiarui Wang and Xiang Shi and Jiaqi Cao and Rubin Wei and Xiquan Wang and Hao Sun and Jingzhi Wang and Zhiqi Yang and Qipeng Guo and Bowen Zhou and Zhouhan Lin},year={2026},eprint={2607.25614},archivePrefix={arXiv},primaryClass={cs.LG},url={https://arxiv.org/abs/2607.25614},}
For questions and discussions, feel free to email wangjiarui1@sjtu.edu.cn.