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.

[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