feat(adapter/nemo): add fully parallel save/load wrapper option

[devyanic11]

Fixes #48

Add optional use_fully_parallel_wrapper flag to the NeMo wrapper utility. When enabled, save/load strategies are wrapped with FullyParallelSaveStrategyWrapper and FullyParallelLoadStrategyWrapper.

Default behavior remains unchanged.

• Tests pass

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Summary of Changes

Hello @devyanic11, 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 introduces an optional enhancement to the NeMo wrapper utility, allowing users to leverage fully parallel save and load strategies for checkpointing. This provides greater flexibility and potentially improved performance for distributed checkpointing operations, while maintaining backward compatibility by preserving the existing default behavior.

Highlights

• Feature: Fully Parallel Wrapper: An optional use_fully_parallel_wrapper flag has been introduced to the NeMo wrapper utility. When enabled, this flag ensures that save and load strategies are wrapped with FullyParallelSaveStrategyWrapper and FullyParallelLoadStrategyWrapper respectively, enhancing checkpointing capabilities.
• Default Behavior Unchanged: The default behavior of the NeMo wrapper remains unchanged, as the use_fully_parallel_wrapper flag defaults to False, meaning existing implementations will not be affected unless explicitly opted in.
• Testing: New test cases have been added to verify the correct application and non-application of the FullyParallel wrappers based on the use_fully_parallel_wrapper flag, ensuring the feature works as intended.

🧠 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.

Changelog

• src/ml_flashpoint/adapter/nemo/wrapper_util.py
  • Imported FullyParallelLoadStrategyWrapper and FullyParallelSaveStrategyWrapper.
  • Added use_fully_parallel_wrapper parameter to wrap_trainer_and_auto_resume_with_mlflashpoint.
  • Added use_fully_parallel_wrapper parameter to wrap_trainer_checkpoint_io_with_mlflashpoint.
  • Implemented conditional wrapping of save/load strategies with FullyParallel wrappers based on the new flag.
• tests/adapter/nemo/test_wrapper_util.py
  • Imported FullyParallelLoadStrategyWrapper and FullyParallelSaveStrategyWrapper for testing.
  • Updated test_successful_wrap_and_resume_creation to include the new use_fully_parallel_wrapper argument.
  • Added test_fully_parallel_wrapper_enabled to verify correct application of FullyParallel wrappers.
  • Added test_fully_parallel_wrapper_disabled_by_default to verify that FullyParallel wrappers are not applied when the flag is false.

Activity

• No human activity has occurred on this pull request yet.

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 an optional use_fully_parallel_wrapper flag to the NeMo wrapper utility. When enabled, this flag wraps the save and load strategies with FullyParallelSaveStrategyWrapper and FullyParallelLoadStrategyWrapper. The implementation is clean and includes tests for both the enabled and disabled states of the flag. My review identifies a couple of minor documentation improvements to align with the repository's style guide.

[github-actions]

Python Code Coverage Summary

Package	Line Rate	Branch Rate	Health
src.ml_flashpoint	100%	100%	✔
src.ml_flashpoint.adapter	100%	100%	✔
src.ml_flashpoint.adapter.megatron	97%	94%	✔
src.ml_flashpoint.adapter.nemo	98%	94%	✔
src.ml_flashpoint.adapter.pytorch	99%	88%	✔
src.ml_flashpoint.checkpoint_object_manager	94%	92%	➖
src.ml_flashpoint.core	96%	92%	✔
src.ml_flashpoint.replication	81%	81%	❌
Summary	95% (2045 / 2151)	91% (472 / 518)	✔

Minimum allowed line rate is 90%

[github-actions]

C++ Code Coverage Summary

Package	Line Rate	Branch Rate	Health
src.ml_flashpoint.checkpoint_object_manager.buffer_object	94%	54%	✔
src.ml_flashpoint.checkpoint_object_manager.object_manager	70%	37%	❌
src.ml_flashpoint.replication.transfer_service	79%	40%	❌
Summary	81% (888 / 1094)	42% (660 / 1554)	➖

Minimum allowed line rate is 80%

[g-husam]

Looks good, thanks for the contribution! Just need to address the lint failures. I see some unexpected import order changes in buffer_io.py, test_buffer_io.py, and checkpoint_object_manager.py for example