Views
No views yet



transformers>=5.0.0from transformers import AutoModel, AutoTokenizer
import torch
repo_name = "nvidia/Nemotron-Labs-Diffusion-14B-Base"
tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
model = AutoModel.from_pretrained(repo_name, trust_remote_code=True)
model = model.cuda().to(torch.bfloat16)
history = []
user_input = input("User: ").strip()
history.append({"role": "user", "content": user_input})
prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
prompt_ids = tokenizer(prompt, return_tensors='pt').input_ids.to(device='cuda')
## Chat in AR Mode
out_ids, nfe = model.ar_generate(inputs.input_ids, max_new_tokens=512)
## Chat in dLM Mode
out_ids, nfe = model.generate(prompt_ids, max_new_tokens=512, block_length=32, threshold=0.9, eos_token_id=tokenizer.eos_token_id)
## Chat in Linear Self-Speculation Mode
out_ids, nfe = model.linear_spec_generate(prompt_ids, max_new_tokens=512, block_length=32, eos_token_id=tokenizer.eos_token_id)
tokenized_out = tokenizer.batch_decode(out_ids[:, prompt_ids.shape[1]:], skip_special_tokens=True)[0]
print(f"Model: {tokenized_out}")
print(f"[Num Function Eval (NFE)={nfe}]")1import torch
2from transformers import AutoModel, AutoTokenizer
3from peft import PeftModel
4
5repo = "nvidia/Nemotron-Labs-Diffusion-14B-Base"
6tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
7model = AutoModel.from_pretrained(repo, trust_remote_code=True)
8model = model.cuda().to(torch.bfloat16)
9
10# Attach the linear_spec LoRA adapter.
11model = PeftModel.from_pretrained(model, repo, subfolder="linear_spec_lora").eval()
12# Unwrap so we can call linear_spec_generate directly (it toggles LoRA internally).
13base = model.model
14
15history = [{"role": "user", "content": "Solve: What is 15% of 240?"}]
16prompt = tokenizer.apply_chat_template(history, tokenize=False, add_generation_prompt=True)
17prompt_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
18
19out_ids, nfe = base.linear_spec_generate(
20 prompt_ids, max_new_tokens=512, block_length=32,
21 eos_token_id=tokenizer.eos_token_id,
22)
23print(tokenizer.decode(out_ids[0, prompt_ids.shape[1]:], skip_special_tokens=True))
24print(f"[NFE={nfe}]")1@techreport{fu2026nemotronlabsdiffusion,
2 title = {Nemotron-Labs-Diffusion: A Tri-Mode Language Model Unifying Autoregressive, Diffusion, and Self-Speculation Decoding},
3 author = {Yonggan Fu and Lexington Whalen and Abhinav Garg and Chengyue Wu and Maksim Khadkevich and Nicolai Oswald and Enze Xie and Daniel Egert and Sharath Turuvekere Sreenivas and Shizhe Diao and Chenhan Yu and Ye Yu and Weijia Chen and Sajad Norouzi and Jingyu Liu and Shiyi Lan and Ligeng Zhu and Jin Wang and Jindong Jiang and Morteza Mardani and Mehran Maghoumi and Song Han and Ante Jukic and Nima Tajbakhsh and Jan Kautz and Pavlo Molchanov},
4 institution = {NVIDIA},
5 year = {2026},
6 note = {Technical report}
7}