fix data processing + resume training

[Dogacel]

Motivation

• Fix training pipeline to resume from the last checkpoint properly.
• Loss tracking on wandb causes gradient accumulation and thus memory leak as wandb commits periodically.
• Track draft configuration on WandB.
• Remove bad data during data preparation to fix zero loss dat points.

Modifications

• Allow loading trainer optimizer states, global step and epoch from the checkpoint.
• Use detached losses for tracking metrics.
• Skip rows that start with assistant rows.

Related Issues

Accuracy Test

• I have additional changes on my private branch to make it work on our HPC cluster and I am using those changes to train our models. However I would be glad if someone can double-check my patch also works on your setup.
• I've run processing and I see ShareGPT dataset has equal number of train/test data poitns with post-processing (regenerate train data) and with no post-processing.

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]

Summary of Changes

Hello @Dogacel, 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 significantly enhances the robustness and reliability of the training pipeline by addressing several critical issues. It introduces comprehensive checkpoint loading capabilities, ensuring seamless resumption of training from any saved state. Furthermore, it resolves a potential memory leak during metric logging to Weights & Biases and improves data quality by filtering out problematic conversational data points. These changes collectively aim to provide a more stable and efficient training experience.

Highlights

• Training Resume: The training pipeline now correctly resumes from the last checkpoint, loading optimizer states, global step, and epoch to ensure continuity.
• Memory Leak Fix: Loss and accuracy tracking for Weights & Biases (WandB) now uses detached tensors, preventing memory leaks caused by gradient accumulation during logging.
• WandB Configuration Tracking: The draft model configuration is now explicitly tracked and saved on Weights & Biases, providing more comprehensive experiment logging.
• Data Quality Improvement: The data preparation script has been updated to filter out 'bad data' (conversations starting with an assistant message) to prevent issues like zero loss data points during training.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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 introduces significant changes to the training pipeline, primarily focusing on enabling robust training resumption from checkpoints in scripts/train_eagle3.py. A new load_checkpoint function was added to retrieve the starting epoch, global step, and optimizer state, which are then used to initialize the training loop and load the optimizer's state. The build_draft_model function was refactored to accept the last checkpoint path, and the main training loop was updated to correctly resume iteration within an epoch by calculating target_batch_idx. Additionally, changes were made to scripts/prepare_data.py to skip entire conversations that start with an 'assistant' message, and specforge/tracker.py was updated to log draft model configuration and save the config file to WandB. The checkpoint naming convention in specforge/utils.py was also updated to include step information. Review comments pointed out a potential IndexError in scripts/prepare_data.py when accessing conversation data without checking for an empty list, noted a regression in scripts/train_eagle3.py where the --ckpt-dir functionality for fine-tuning from a base model was removed, and raised a security concern regarding weights_only=False when loading checkpoints, though acknowledging its necessity for optimizer state.

I am having trouble creating individual review comments. Click here to see my feedback.

scripts/prepare_data.py (196)

Accessing row["conversations"][0] without checking if row["conversations"] is empty can lead to an IndexError, which would crash the data preparation script. It's safer to check for an empty list first.

            if row["conversations"] and row["conversations"][0]["role"] == "assistant":

scripts/prepare_data.py (210)

Accessing row["conversations"][0] without checking if row["conversations"] is empty can lead to an IndexError, which would crash the data preparation script. It's safer to check for an empty list first. This is the same issue as in the loop for the training set.

                if row["conversations"] and row["conversations"][0]["role"] == "assistant":

scripts/train_eagle3.py (346-363)

The logic to handle --ckpt-dir for starting fine-tuning from a specific base model checkpoint has been removed. This appears to be a regression, as the command-line argument still exists but is no longer used. Was this removal intentional? If not, this functionality should be restored. The previous logic could conflict with --resume, so you might want to rethink how they interact (e.g., make them mutually exclusive or have a clear precedence).

scripts/train_eagle3.py (337)

Loading checkpoints with weights_only=False can be a security vulnerability, as it may lead to arbitrary code execution if the checkpoint file is compromised. Since you are loading optimizer state which requires this setting, it's acceptable, but it's crucial to ensure that you only load checkpoints from a trusted source.

[FrankLeeeee]

Hi @Dogacel , can you apply code formatting via pre-commit?

[Dogacel]

Hi @Dogacel , can you apply code formatting via pre-commit?

Got it!

The link in PR description is dead, but I assume it is this one -> https://pre-commit.com/#install

I've run pre-commit run --all-file and committed.

[Dogacel]

Gentle ping @FrankLeeeee