fix: project offline state into vocab space

[timmy-feng]

Motivation

PR #293 introduces new behavior which saves only the last layer hidden states from the target model when generating offline hidden states. This saves us from needing to load the entire vocab size of logits into CPU RAM.

This breaks the old OfflineEagle3Model class which passes target to OnlineEagle3Model without projecting the shape from hidden_size to vocab_size.

Modifications

Project the offline hidden state into the correct shape using the target LM head in OfflineEagle3Model.forward().

Related Issues

Accuracy Test

Benchmark & Profiling

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!