Skip to content
Merged
Show file tree
Hide file tree
Changes from 28 commits
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
f4cadcc
remove detachedeagle and handle all offline mode in the eagle model
yeyu-nvidia Sep 15, 2025
dd2cef6
apply coderabbit suggestion
yeyu-nvidia Sep 15, 2025
6cf677e
remove OfflineEagleDMRegistry
yeyu-nvidia Sep 15, 2025
e959ae1
apply suggestion to cover eagle1 case
yeyu-nvidia Sep 15, 2025
9ba9935
debug
yeyu-nvidia Sep 15, 2025
0f8493a
debug
yeyu-nvidia Sep 15, 2025
0670b8b
minor
yeyu-nvidia Sep 16, 2025
86bc66e
fix the bug in megatron import
yeyu-nvidia Sep 16, 2025
a26a5a6
Added support for qwen3-next quantization and export (#323)
kinjalpatel27 Sep 16, 2025
742429d
Update distill Megatron plugin (#319)
AAnoosheh Sep 16, 2025
62c1c99
Remove unused utilities for ModelOpt <0.29 MCore checkpoints (#322)
realAsma Sep 16, 2025
6a3edec
Upgrade TensorRT-LLM docker to 1.1.0RC2 (#327)
cjluo-nv Sep 16, 2025
68abf3e
[1/N] QATTrainer training workflow fixes and clean up; Added backend …
realAsma Sep 17, 2025
e40fb07
import fix for torch 2.9 (#315)
i-riyad Sep 17, 2025
8f25717
checkout main branch's fix for megatron+importer
yeyu-nvidia Sep 18, 2025
912e427
pad labels if it's 1 token shorter than input_ids
yeyu-nvidia Sep 18, 2025
0582079
Revert "pad labels if it's 1 token shorter than input_ids"
yeyu-nvidia Sep 18, 2025
ca7ec41
pad labels if it's 1 token shorter than input_ids
yeyu-nvidia Sep 18, 2025
a6a0805
Fix issues of attention.core_attention.softmax_offset is None (#330)
yueshen2016 Sep 17, 2025
0e01a5d
update eagle example notebook (#314)
h-guo18 Sep 17, 2025
5f6bedc
Allow KD loss in val mode for MLM plugin (#331)
AAnoosheh Sep 17, 2025
b504296
Deprecate TRTLLM-build in examples (#297)
cjluo-nv Sep 17, 2025
3303f35
Fix issue of attention.core_attention is None (#334)
yueshen2016 Sep 18, 2025
0f7b7ea
Reinstate int8_sq support for vlm_example. (#333)
cjluo-nv Sep 18, 2025
89c65b9
Update News README.md (#336)
omrialmog Sep 18, 2025
f7425fc
[NVBug: 5525758] Update VLM-PTQ readme (#339)
cjluo-nv Sep 18, 2025
2ec5654
Slurm support for QAT Simplified Flow + Qwen3-8B recipe (#285)
jenchen13 Sep 18, 2025
0c203e0
Merge branch 'main' into yeyu/move_offline_eagle_to_online
yeyu-nvidia Sep 18, 2025
fc9cca8
Merge branch 'main' into yeyu/move_offline_eagle_to_online
yeyu-nvidia Sep 18, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 4 additions & 7 deletions modelopt/torch/speculative/eagle/conversion.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,24 +24,21 @@
from ..config import EagleConfig

EagleDMRegistry = _DMRegistryCls(prefix="Eagle") # global instance for the registry
OfflineEagleDMRegistry = _DMRegistryCls(prefix="DetachedEagle") # global instance for the registry


def convert_to_eagle_model(model: nn.Module, config: EagleConfig) -> ConvertReturnType:
"""Convert the model to a eagle model as per `config`."""
# initialize the true module if necessary
model = model.init_modellike() if isinstance(model, ModelLikeModule) else model

registry = OfflineEagleDMRegistry if config.eagle_offline else EagleDMRegistry

original_cls = type(model)
if original_cls not in registry:
for cls in registry._registry:
if original_cls not in EagleDMRegistry:
for cls in EagleDMRegistry._registry:
if issubclass(original_cls, cls):
registry.register({original_cls: "base_model_class"})(registry[cls])
EagleDMRegistry.register({original_cls: "base_model_class"})(EagleDMRegistry[cls])
break

eagle_model = registry.convert(model)
eagle_model = EagleDMRegistry.convert(model)
eagle_model.modify(
eagle_offline=config.eagle_offline,
eagle_hidden_state_distillation=config.eagle_hidden_state_distillation,
Expand Down
Loading
Loading