Skip to content

Add Benchmarking and Fine-Tuning Support for ZenFlow #982

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.

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

Open
wants to merge 5 commits into
base: master
Choose a base branch
from

Conversation

Antlera
Copy link

@Antlera Antlera commented Jul 3, 2025

Description:

This PR introduces scripts for benchmarking and fine-tuning with ZenFlow:

  • zf_benchmark.py: Benchmark script for evaluating offloading performance (adapted from offload_states.py by @tohtana ).
  • output_table.py: Parses and summarizes benchmark logs.
  • run_benchmark.sh: Automates benchmark runs with configurable parameters.
  • finetune_llama.py: Fine-tuning script for Llama-2 with DeepSpeed + ZenFlow.
  • finetune_llama.sh: Launch script for fine-tuning with environment setup.
  • zf_config.json: Example DeepSpeed config with ZenFlow optimizations.

Note: This PR is complimentary to PR #7391 on the main repo, and should be merged with (or after) merging PR #7391.

@Antlera Antlera requested a review from tjruwase as a code owner July 3, 2025 01:23
Antlera and others added 2 commits July 2, 2025 21:27
- Introduced `zf_benchmark.py` for model offloading benchmarking with DeepSpeed.
- Added `output_table.py` to parse and display benchmark results in a tabular format.
- Created `run_benchmark.sh` to automate benchmark runs with various configurations.

Signed-off-by: Tingfeng Lan <[email protected]>
- Introduced `finetune_llama.py` for fine-tuning the Llama-2 model using DeepSpeed and ZenFlow.
- Added `finetune_llama.sh` for automated training setup with environment variables and DeepSpeed command.
- Added `zf_config.json` example for DeepSpeed configuration with ZenFlow optimizations.

Signed-off-by: Tingfeng Lan <[email protected]>
Co-authored-by: Yusen Wu <[email protected]>
@Antlera Antlera force-pushed the zenflow_z1_2_example branch from ca441f5 to 0528aed Compare July 3, 2025 01:27
@Antlera
Copy link
Author

Antlera commented Jul 3, 2025

Hi, @tohtana @tjruwase. Could you help review this PR when you have time? This is the example PR for PR #7391. Thanks!

@Antlera
Copy link
Author

Antlera commented Aug 4, 2025

@sfc-gh-truwase Thanks for the great suggestions — I’ve applied them all!

@delock
Copy link
Contributor

delock commented Aug 6, 2025

Hi @Antlera I have a question. I saw ZenFlow has CPU running parameter update. Does DeepSpeed argument --bind_cores_to_rank helping with CPU optimizer efficiency? If true then maybe this switch could be add to launch script.

Here is a link to this switch. This is a switch for CPU backend, but for CPU offload this switch should help as well.
https://github.com/deepspeedai/DeepSpeed/blob/master/docs/_tutorials/accelerator-setup-guide.md#how-to-launch-deepspeed-on-intel-architecture-cpu

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants