mirror of
https://github.com/NVIDIA/Model-Optimizer.git
synced 2026-11-08 09:03:20 +08:00
skierat
dd16a96fda
## What does this PR do? **Type of change:** New example **Overview:** Adds examples/speculative_decoding/guides/train_eagle_head_cosmos_reason2.ipynb, a step-by-step Jupyter notebook that walks through the full EAGLE3 draft-head training workflow for nvidia/Cosmos-Reason2-8B. The notebook covers: 1. Installing dependencies 2. Authenticating with Hugging Face 3. Preparing training data from the Nemotron-Post-Training-Dataset-v2 (chat split) using a curated row-selection mapping (guides/nemotron_mapping.csv) 4. Inspecting the bundled EAGLE3 config (guides/CR2_eagle_config.json) tuned for Cosmos-Reason2 (YaRN RoPE, FlexAttention, reduced draft vocabulary) 5. (Optional) Calibrating the draft vocabulary to 32k tokens for faster training and inference 6. Launching training via launch_train.sh with FSDP2 multi-GPU support 7. Exporting the checkpoint to HF format and serving with vLLM Also includes guides/nemotron_mapping.csv and guides/CR2_eagle_config.json as companion files. ## Usage Open and run examples/speculative_decoding/guides/train_eagle_head_cosmos_reason2.ipynb cell by cell. After training, serve the exported checkpoint with: ``` vllm serve nvidia/Cosmos-Reason2-8B \ --host 0.0.0.0 \ --port 8000 \ --speculative-model export/cosmos-reason2-8b-eagle3 \ --num-speculative-tokens 3 \ --dtype bfloat16 ``` ## Testing Tested end-to-end on a 4xB100 GPUs. The exported checkpoint was validated with specdec_bench ## Before your PR is "*Ready for review*" <!-- If you haven't finished some of the above items you can still open `Draft` PR. --> - **Make sure you read and follow [Contributor guidelines](https://github.com/NVIDIA/Model-Optimizer/blob/main/CONTRIBUTING.md)** and your commits are signed. - **Is this change backward compatible?**: Yes - **Did you write any new necessary tests?**: No - **Did you add or update any necessary documentation?**: Yes (the notebook is self-documenting) - **Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?**: No ## Additional Information Cosmos-Reason2-8B requires at least one 80 GB GPU (H100/A100) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added a speculative-decoding configuration for draft model and rotary/attention behavior. * Added an end-to-end training notebook demonstrating training, export, and deployment of a speculative-decoding draft head on Cosmos-Reason2. * Added a data-preparation tool to download, normalize, and convert Nemotron chat conversations into a standardized conversation format. * **Documentation** * Notebook documents environment setup, data prep, training/validation cadence, export, and deployment steps. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Slawek Kierat <skierat@nvidia.com>