Dawn-AI commited on
Commit
a63cbc1
·
verified ·
1 Parent(s): 13f910f

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +0 -140
README.md CHANGED
@@ -4,143 +4,3 @@ license: other
4
  license_name: cc-by-nc-4.0
5
  pipeline_tag: text-generation
6
  ---
7
-
8
- # Nemotron-Flash-3B Base Model
9
-
10
- <p align="center">
11
- 🗞️ <a href="https://arxiv.org/pdf/2511.18890">Paper</a>&nbsp&nbsp| &nbsp&nbsp 🤗 <a href="https://huggingface.co/nvidia/Nemotron-Flash-1B">Nemotron-Flash-1B</a> | &nbsp&nbsp 🤗 <a href="https://huggingface.co/nvidia/Nemotron-Flash-3B">Nemotron-Flash-3B</a> | &nbsp&nbsp 🤗 <a href="https://huggingface.co/nvidia/Nemotron-Flash-3B-Instruct">Nemotron-Flash-3B-Instruct</a> &nbsp
12
- </p>
13
-
14
- ## Model Overview
15
-
16
- Nemotron-Flash is a new hybrid small language model family designed around real-world latency rather than parameter count. It features latency-optimal depth–width ratios, hybrid operators discovered through evolutionary search, and training-time weight normalization. See our <a href="https://arxiv.org/pdf/2511.18890">NeurIPS 2025 paper</a> for more technical details.
17
-
18
- The models achieve SOTA accuracy in math, coding, and commonsense reasoning at the 1B and 3B scales, while delivering decent small-batch latency and large-batch throughput. For example, Nemotron-Flash-1B achieves +5.5% accuracy, 1.9× lower latency, and 45.6× higher throughput compared with Qwen3-0.6B; and Nemotron-Flash-3B achieves +2% / +5.5% accuracy over Qwen2.5-3B / Qwen3-1.7B with 1.3× / 1.7× lower latency and 6.4× / 18.7× higher throughput, respectively.
19
-
20
-
21
- <div align="center">
22
- <img src="https://huggingface.co/nvidia/Nemotron-Flash-3B/resolve/main/images/nemotron_flash_result.png" alt="Compare with SOTA SLMs" width="800">
23
- </div>
24
-
25
-
26
- ## Environment
27
- ```bash
28
- torch<=2.9.1
29
- transformers<=4.56.2
30
- causal-conv1d
31
- flash-attn<=2.7.3
32
- mamba-ssm
33
- flash-linear-attention
34
- ```
35
- We provide a <a href="https://huggingface.co/nvidia/Nemotron-Flash-3B/resolve/main/setup.sh">script</a> to build the conda environment: `bash setup.sh`.
36
-
37
-
38
-
39
- ## Chat with Nemotron-Flash
40
-
41
- We integrated the attention kernel from <a href="https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html">TRT-LLM AutoDeploy</a> to enable generation with CUDA Graph:
42
-
43
- ```
44
- from transformers import AutoModelForCausalLM, AutoTokenizer
45
- import torch
46
-
47
- repo_name = "nvidia/Nemotron-Flash-3B"
48
-
49
- tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
50
- model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
51
- model = model.cuda().to(torch.bfloat16)
52
-
53
- print('Initializing generation state...')
54
- generation_state = model.init_cuda_graph_generation(
55
- max_new_tokens=max_new_tokens,
56
- batch_size=1,
57
- device='cuda',
58
- )
59
-
60
- prompt = input("User:")
61
- inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
62
-
63
- print(f"Generating with CUDA graph acceleration...")
64
- outputs = model.generate_with_cuda_graph(
65
- input_ids=inputs["input_ids"],
66
- generation_state=generation_state,
67
- max_new_tokens=256,
68
- temperature=0,
69
- eos_token_id=tokenizer.eos_token_id,
70
- )
71
-
72
- response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
73
- print(f"Response: {response}")
74
- ```
75
-
76
- Another option is to perform generation w/o CUDA Graph:
77
-
78
- ```
79
- outputs = model.generate_with_cache(
80
- input_ids=inputs["input_ids"],
81
- max_new_tokens=256,
82
- temperature=0,
83
- eos_token_id=tokenizer.eos_token_id,
84
- )
85
- ```
86
-
87
- ## Finetune Nemotron-Flash
88
-
89
- To finetune Nemotron-Flash models, switch the attention kernel to FlashAttention2 when loading the model:
90
-
91
- ```
92
- from transformers import AutoConfig, AutoModelForCausalLM
93
- repo_name = "nvidia/Nemotron-Flash-3B"
94
-
95
- config = AutoConfig.from_pretrained(repo_name, trust_remote_code=True)
96
- setattr(config, "attn_implementation_new", "flash_attention_2")
97
- model = AutoModelForCausalLM.from_pretrained(repo_name, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True)
98
- ```
99
-
100
- ## Running Nemotron-Flash with TensorRT-LLM
101
-
102
- ### Setup
103
- Installation + quick start for TensorRT-LLM: <a href="https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html">Tutorial</a>.
104
-
105
- ### Quick example
106
-
107
- An example script for running through the generation workflow:
108
- ```
109
- cd examples/auto_deploy
110
- python build_and_run_ad.py --model nvidia/Nemotron-Flash-3B-Instruct --args.yaml-extra nemotron_flash.yaml
111
- ```
112
-
113
- ### Serving with trtllm-serve
114
-
115
- - Spin up a trtllm server (more details are in this <a href="https://nvidia.github.io/TensorRT-LLM/commands/trtllm-serve/trtllm-serve.html#starting-a-server">doc</a>):
116
- ```
117
- trtllm-serve serve nvidia/Nemotron-Flash-3B-Instruct \
118
- --backend _autodeploy \
119
- --trust_remote_code \
120
- --extra_llm_api_options examples/auto_deploy/nemotron_flash.yaml
121
- ```
122
-
123
- - Send a request (more details are in this <a href="https://nvidia.github.io/TensorRT-LLM/examples/curl_chat_client.html">doc</a>):
124
- ```
125
- curl http://localhost:8000/v1/chat/completions \
126
- -H "Content-Type: application/json" \
127
- -d '{
128
- "model": "nvidia/Nemotron-Flash-3B-Instruct",
129
- "messages":[{"role": "user", "content": "Where is New York?"}],
130
- "max_tokens": 16,
131
- "temperature": 0
132
- }'
133
- ```
134
-
135
-
136
- ## Citation
137
- ```
138
- @misc{fu2025nemotronflash,
139
- title={Nemotron-Flash: Towards Latency-Optimal Hybrid Small Language Models},
140
- author={Yonggan Fu and Xin Dong and Shizhe Diao and Matthijs Van keirsbilck and Hanrong Ye and Wonmin Byeon and Yashaswi Karnati and Lucas Liebenwein and Hannah Zhang and Nikolaus Binder and Maksim Khadkevich and Alexander Keller and Jan Kautz and Yingyan Celine Lin and Pavlo Molchanov},
141
- year={2025},
142
- eprint={2511.18890},
143
- archivePrefix={arXiv},
144
- primaryClass={cs.LG},
145
- url={https://arxiv.org/abs/2511.18890},
146
- }
 
4
  license_name: cc-by-nc-4.0
5
  pipeline_tag: text-generation
6
  ---