[PP] Refactor test_schedule_multiproc (#158780)

Browse Source

This refactors the pipelining schedule tests since a lot of them have the same repeated code of:
1. Create pipelined model and reference model
2. Run reference model and pipelined model
3. compare gradients

So this refactors those parts above into helper methods and reduces ~300 LOC. Also adds a better gradient check to resolve flakiness (fixes https://github.com/pytorch/pytorch/issues/154408).

Pull Request resolved: https://github.com/pytorch/pytorch/pull/158780
Approved by: https://github.com/wconstab

...

This commit is contained in:

Howard Huang 2025-07-31 14:35:15 -07:00 committed by PyTorch MergeBot

parent 3967dbedf4

commit b0b3e6e48b

1 changed files with 331 additions and 638 deletions

Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL

Show Stats Download Patch File Download Diff File Expand all files Collapse all files

969

test/distributed/pipelining/test_schedule_multiproc.py

View File

File diff suppressed because it is too large Load Diff