Skip to content

fix save checkpoint bug #406

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
FrankLeeeee merged 1 commit into from
Jan 12, 2026
Merged

fix save checkpoint bug #406

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions scripts/train_eagle3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -862,6 +862,13 @@ def main():
if args.max_num_steps is not None and global_step >= args.max_num_steps:
break

# Save final checkpoint if training ended without saving
if global_step % args.save_interval != 0:
print_on_rank0(
f"Training completed at step {global_step}, saving final checkpoint..."
)
save_checkpoints(args, epoch, global_step, eagle3_model, optimizer)
Comment on lines +866 to +870
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot Jan 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

To prevent a potential ZeroDivisionError if args.save_interval is set to 0, it's good practice to add a check to ensure args.save_interval is positive before performing the modulo operation. This would make the script more robust against invalid inputs. A similar check should also be considered for the checkpoint logic inside the training loop on line 855 for consistency.

Suggested change
if global_step % args.save_interval != 0:
print_on_rank0(
f"Training completed at step {global_step}, saving final checkpoint..."
)
save_checkpoints(args, epoch, global_step, eagle3_model, optimizer)
if args.save_interval > 0 and global_step % args.save_interval != 0:
print_on_rank0(
f"Training completed at step {global_step}, saving final checkpoint..."
)
save_checkpoints(args, epoch, global_step, eagle3_model, optimizer)


# Close the tracker
tracker.close()
destroy_distributed()
Expand Down