Fix for when draft model hidden dimension is different from target model hidden dimension

[yilian49]

Motivation

The current draft model files (e.g. llama3_eagle.py) supports draft model config having target_hidden_dimension separate from hidden_dimension. However, in the training scripts, it loads the target model's embedding function by default, which wouldn't work if you have different hidden dimensions between target and draft model.

Modifications

Add a check for if there is a target_hidden_dimension in the draft config and if the two hidden dimensions are different. Don't freeze the embedding function if it doesn't load embedding function from target model.

[gemini-code-assist]

Summary of Changes

Hello @yilian49, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a critical issue in the training scripts where draft models with different hidden dimensions from their target models could lead to incorrect behavior during embedding loading. The changes introduce a robust check to ensure that the target model's embedding function is only loaded and frozen when the hidden dimensions are compatible, thereby preventing potential training failures and ensuring the correct initialization of draft models.

Highlights

• Conditional Embedding Loading: Implemented a conditional check in the training scripts (train_eagle3_offline.py and train_eagle3_online.py) to prevent loading and freezing the target model's embedding function when the draft model's hidden_dimension differs from the target_hidden_dimension.
• Preventing Dimension Mismatch Issues: Ensured that the draft_model.load_embedding and draft_model.freeze_embedding calls are only executed if the target_hidden_size is not explicitly defined or matches the hidden_size of the draft model.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request correctly addresses an issue where the training scripts would fail if the draft and target models have different hidden dimensions. The fix introduces a conditional check to load and freeze embeddings only when the hidden dimensions are compatible, which is a sensible approach. The changes are applied consistently across both offline and online training scripts.

[gemini-code-assist]

This condition can be made more concise and idiomatic by using getattr with a default value. This avoids explicitly checking for the attribute's existence before accessing it.

On a related note, this entire conditional block is duplicated in scripts/train_eagle3_online.py. Consider refactoring the model setup logic into a shared utility function to improve maintainability and avoid future inconsistencies.

Suggested change

	if (
	not hasattr(draft_model_config, "target_hidden_size")
	or draft_model_config.target_hidden_size == draft_model_config.hidden_size
	):
	if (
	getattr(draft_model_config, "target_hidden_size", draft_model_config.hidden_size)
	== draft_model_config.hidden_size
	):

[gemini-code-assist]

This condition can be made more concise and idiomatic by using getattr with a default value. This avoids explicitly checking for the attribute's existence before accessing it.

On a related note, this entire conditional block is duplicated from scripts/train_eagle3_offline.py. Consider refactoring the model setup logic into a shared utility function to improve maintainability and avoid future inconsistencies.

Suggested change

	if (
	not hasattr(draft_model_config, "target_hidden_size")
	or draft_model_config.target_hidden_size == draft_model_config.hidden_size
	):
	if (
	getattr(draft_model_config, "target_hidden_size", draft_model_config.hidden_size)
	== draft_model_config.hidden_size
	):

[zyksir]

this means that the code in inference engine(e.g sglang) also needs to change as well, right?

[b8zhong]

Hi, wonder if during inference, this means that a EAGLE head of a different hidden dim can be inferred with a base model with a different hidden dim, or if they are fundamentally incompatible.

@yilian49