fix: include medusa in data_module assignment in main.py

[yeyu-nvidia]

Problem

When training.mode == "medusa" is used in main.py, the data_module variable is never assigned because line 344 only covered eagle3 and dflash modes. This causes an UnboundLocalError when the trainer is constructed with **data_module.

Fixes OMNIML-4147

Fix

Add "medusa" to the training_args.mode in ("eagle3", "dflash") condition so data_module is correctly populated for medusa training.

Summary by CodeRabbit

• Bug Fixes
  • Fixed speculative decoding example to properly handle "medusa" mode alongside existing "eagle3" and "dflash" modes.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: d0a67b92-0065-4a5f-86e9-5760321ec053

📥 Commits

Reviewing files that changed from the base of the PR and between 9bb917d and 4475873.

📒 Files selected for processing (1)

• examples/speculative_decoding/main.py

---

📝 Walkthrough

Walkthrough

The train() function in the speculative decoding example is modified to include "medusa" mode in the conditional check for using make_speculative_data_module(). Previously, only "eagle3" and "dflash" modes triggered speculative data module construction; "medusa" now qualifies for the same treatment.

Changes

Cohort / File(s)	Summary
Speculative Decoding Mode Support examples/speculative_decoding/main.py	Added "medusa" to the set of training modes that use speculative data module construction, alongside existing "eagle3" and "dflash" modes.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately describes the main change: adding 'medusa' to the data_module assignment condition in main.py. It is specific, concise, and directly summarizes the primary purpose of the fix.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	The pull request is a minimal bug fix that only adds 'medusa' to an existing conditional statement on line 344. The code does not introduce any security anti-patterns outlined in SECURITY.md: it uses torch.load() with weights_only=True (safe), exposes trust_remote_code as a configurable parameter with a secure default of False, avoids eval/exec calls, contains no # nosec comments, and adds no new dependencies. The fix correctly prevents an UnboundLocalError when training_args.mode == 'medusa'.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/medusa-unbound-data-module

---

_{Review rate limit: 9/10 reviews remaining, refill in 6 minutes.}

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.49%. Comparing base (077e29a) to head (4475873).
⚠️ Report is 13 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1370      +/-   ##
==========================================
+ Coverage   76.48%   77.49%   +1.00%     
==========================================
  Files         471      471              
  Lines       50487    50487              
==========================================
+ Hits        38617    39124     +507     
+ Misses      11870    11363     -507     

Flag	Coverage Δ
examples	41.58% <ø> (+10.83%)	⬆️
regression	14.91% <ø> (+0.20%)	⬆️
unit	52.78% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

PR Preview Action v1.8.1
Preview removed because the pull request was closed.
2026-05-04 09:48 UTC