Skip to content

ZeroBubble Pipeline Parallelism#7

Open
pgusto34 wants to merge 26 commits intomainfrom
ZeroBubble
Open

ZeroBubble Pipeline Parallelism#7
pgusto34 wants to merge 26 commits intomainfrom
ZeroBubble

Conversation

@pgusto34
Copy link
Collaborator

High-level Description

Added split backward pass support to Piper, which enabled the system to support a wider array of schedules (e.g., ZeroBubble).

Description of changes:

  • Refactored piper into a piper directory/project and renamed files accordingly by removing the "piper" prefix (if you prefer to keep src as the directory name, that's totally cool too)
  • Added the backward_weight() and backward_input() functions to actor.py, which are the functions that implement the split backward pass
  • Added a backward_utils file containing helper functions for the split_backward_pass. These functions are essentially identically to the ones that PiPPy uses outside of a few minor tweaks.
  • Updated the Task class in exec.py to store a TaskType enum instead of booleans when determining what type of task type (f, b, u, i, w) each object represents
  • Updated the validate_schedule() function in exec.py to use the updated Task class and check for correctness when backward_input/backward_weight tasks could be present.
  • Updated piper_exec() in exec.py to include cases for handling backward_input and backward_weight tasks
  • Added zb1p as an option to the test_llama.py --schedule argument
  • Modified the schedule_helpers.py to use the updated Task class
  • Created a Dockerfile and run.sh script to deploy Piper in a Docker container

@pgusto34 pgusto34 self-assigned this Jan 28, 2026
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.

2 participants