[feature]:Add adaptation of all models to zero3

[Feng0w0]

1.Based on the current hot models of the five major series, zero 3 is adapted and aligned with zero 2 for accuracy, with a loss error within 0.001.
2.Add a zero 3 YAML file for each series of models.
3.Add some training script files using zero3 on NPU.

[gemini-code-assist]

Summary of Changes

Hello @Feng0w0, 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 training capabilities by integrating DeepSpeed Zero3 across various model architectures. The changes aim to optimize memory usage and training efficiency, especially for large models and NPU environments, while meticulously preserving model accuracy. It provides the necessary infrastructure, including configuration files and training scripts, to leverage these advanced optimization techniques.

Highlights

• DeepSpeed Zero3 Integration: Adapted five major model series to utilize DeepSpeed Zero3 for enhanced memory efficiency during training.
• Accuracy Preservation: Ensured that the Zero3 adaptation maintains accuracy, with a reported loss error within 0.001% compared to Zero2.
• Configuration and Scripting: Added dedicated Zero3 YAML configuration files for each model series and new training scripts optimized for NPU environments.
• Model Loading and Checkpointing Enhancements: Improved model loading mechanisms to support Zero3 initialization and updated gradient checkpointing with determinism checks.

🧠 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 adapt the models for DeepSpeed Zero3 training, particularly for NPU environments. Key updates include modifications to gradient checkpointing, a new model initialization context for Zero3, and adjusted state dictionary loading. The changes also involve updating several forward methods in wan_video_vae.py to correctly handle feature caching, and adding new configuration files and training scripts for NPU. Overall, the changes align with the goal of enabling Zero3 for various models.

[gemini-code-assist]

The change from torch.inference_mode() to torch.no_grad() might have minor performance implications during inference. torch.inference_mode() is generally preferred for inference tasks as it can offer better performance by disabling more autograd-related features than torch.no_grad(). While torch.no_grad() is functionally correct for disabling gradient computation, inference_mode is specifically optimized for inference workflows.

Suggested change

	with torch.no_grad():
	with torch.inference_mode():