Skip to content

Implement count-heuristic ZB-H2 pipeline schedule - #95

Open
insujang wants to merge 3 commits into
refactorfrom
codex/zb-h2-pipeline
Open

Implement count-heuristic ZB-H2 pipeline schedule#95
insujang wants to merge 3 commits into
refactorfrom
codex/zb-h2-pipeline

Conversation

@insujang

Copy link
Copy Markdown
Collaborator

Summary

  • split OneForwardOneBackwardSchedule into schedule_1f1b.py while preserving compatibility re-exports
  • add count-heuristic ZB-H2 with the u_s = 2(p-s)-1 warmup and independent count-selected F/B/W operations
  • add genuine deferred linear weight-gradient computation for tied weights, DTensor TP paths, and stacked expert weights
  • add ParallelConfig.pipeline_schedule dispatch and reject inconsistent per-module selections
  • preserve the multimodal encoder-to-LM seam and intentionally omit ZBV/VPP

Validation

  • Ruff checks passed
  • 22 focused schedule, configuration, and pipeline tests passed
  • full distributed parallelism integration matrix: 16 passed
  • existing 1F1B multimodal seam tests plus six-rank ZB-H2 CP seam parity: 3 passed
  • strict two-rank 1F1B versus ZB-H2 loss and local-gradient parity, including consecutive-step cleanup
  • PP+TP and PP+EP Gloo coverage

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.

1 participant