Update README.md

README.md

@@ -1,18 +1,44 @@
 ---
 library_name: transformers
-tags: []
+license: other
+license_name: cc-by-nc-4.0
+pipeline_tag: text-generation
 ---
 
-# Nemotron-Flash-3B-Instruct
+# Nemotron-Flash-3B Instruct Model
 
-Nemotron-Flash is a new hybrid SLM model family that outperforms Qwen models in accuracy (math, coding, and commonsense), batch-size-1 latency, and throughput. More details are in our NeurIPS 2025 [paper](https://drive.google.com/drive/folders/17vOGktwUfUpRAJPGJUV6oX8XwLSczZtv?usp=sharing).
+<p align="center">
+  🗞️ <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
+</p>
 
-Docker path: `/lustre/fsw/portfolios/nvr/users/yongganf/docker/megatron_py25_fast_slm.sqsh` on NRT.
+## Model Overview
 
+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.
 
-## Chat with Nemotron-Flash-3B-Instruct
+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.
 
-We wrap the model into CUDA Graph for fast generation:
+
+<div align="center">
+<img src="https://huggingface.co/nvidia/Nemotron-Flash-3B/resolve/main/images/nemotron_flash_result.png" alt="Compare with SOTA SLMs" width="800">
+</div>
+
+
+## Environment
+```bash
+torch<=2.9.1
+transformers<=4.56.2
+causal-conv1d
+flash-attn<=2.7.3
+mamba-ssm
+flash-linear-attention
+```
+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`.
+
+
+
+## Chat with Nemotron-Flash
+
+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:
 
 ```
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -24,9 +50,6 @@ tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)
 model = model.cuda().to(torch.bfloat16)
 
-
-max_new_tokens = 256
-
 print('Initializing generation state...')
 generation_state = model.init_cuda_graph_generation(
     max_new_tokens=max_new_tokens,
@@ -34,25 +57,55 @@ generation_state = model.init_cuda_graph_generation(
     device='cuda',
 )
 
-while True:
-    prompt = input("User:")
-    if prompt.lower() == "exit":
-        break
-
-    inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
-
-    print(f"Generating with CUDA graph acceleration...")
-    outputs = model.generate_with_cuda_graph(
-        input_ids=inputs["input_ids"],
-        generation_state=generation_state,
-        max_new_tokens=max_new_tokens,
-        temperature=0,
-        top_k=50,
-        eos_token_id=tokenizer.eos_token_id,
-        profiling=False,
+prompt = input("User:")
+prompt = "User: " + prompt + "\nAssistant:"
+inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
+
+print(f"Generating with CUDA graph acceleration...")
+outputs = model.generate_with_cuda_graph(
+    input_ids=inputs["input_ids"],
+    generation_state=generation_state,
+    max_new_tokens=256,
+    temperature=0,
+    eos_token_id=tokenizer.eos_token_id,
+)
+
+response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+print(f"Response: {response}")
+```
+
+Another option is to perform generation w/o CUDA Graph:
+
+```
+outputs = model.generate_with_cache(
+    input_ids=inputs["input_ids"],
+    max_new_tokens=256,
+    temperature=0,
+    eos_token_id=tokenizer.eos_token_id,
     )
+```
+
+## Finetune Nemotron-Flash
+
+To finetune Nemotron-Flash models, switch the attention kernel to FlashAttention2 when loading the model:
 
-    response = tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
-    
-    print(f"Response: {response}")
-```
+```
+from transformers import AutoConfig, AutoModelForCausalLM
+repo_name = "nvidia/Nemotron-Flash-3B-Instruct"
+
+config = AutoConfig.from_pretrained(repo_name, trust_remote_code=True)
+setattr(config, "attention_implementation_new", "flash_attention_2")
+model = AutoModelForCausalLM.from_pretrained(repo_name, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True)
+```
+
+## Citation
+```
+@misc{fu2025nemotronflash,
+      title={Nemotron-Flash: Towards Latency-Optimal Hybrid Small Language Models}, 
+      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},
+      year={2025},
+      eprint={2511.18890},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2511.18890}, 
+}