Add Pipeline Parallel (and 2D PP+FSDP) support #318

wconstab · 2024-05-09T01:38:00Z

Stack from ghstack (oldest at bottom):

-> Add Pipeline Parallel (and 2D PP+FSDP) support #318

runs PP+DP and PP+TP without issue,
runs PP+TP+DP with decreasing loss, but fails DCP save

Supports only simple schedules currently, gpipe and 1f1b.

Ads cmdline/toml arg for specifiying split points, in a unified
way between tracer or manual frontend.

e.g. user can specifiy "layers.2,layers.4" as split points.

Currently uses manual frontend by default, but allows specifying
tracer frontend. Tracer frontend requires working around additional
compatibility limitations, indicated by raising assertions, and is
not ready for wider use yet.

[ghstack-poisoned]

ghstack-source-id: 318831bda97c4639b745832e7f33f0b616a479ee Pull Request resolved: #318

torchtitan/config_manager.py

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: e49b659e66f4101cef58ad717a80521f5b172347 Pull Request resolved: #318

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: ee49541296ccea6a32a2f17cbc9a6c892216fde4 Pull Request resolved: #318

torchtitan/parallelisms/parallelize_llama.py

wconstab · 2024-05-10T22:05:02Z

torchtitan/parallelisms/parallelize_llama.py

+            model_config.vocab_size, input_shape, dtype=torch.int64, device=device
+        )
+
+        # HACK- can't use shape inference via execution of the PP stage inside ManualPipelineStage API, becuase the


proposal to land this part as-is. We need to add the lazy shape inference part to manual stage api, but i think we can clean this up after it lands.

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: 073ec20f5dfb6f323ce829275f349224bcfff40b Pull Request resolved: #318

train.py

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: ad748435e833ef1e0e4c8e424339532daadf8bee Pull Request resolved: #318

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: 06f8ed615894662b70c37dbb69a0f82c976a5009 Pull Request resolved: #318

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: 9dd9b7ab4c8548757351671dfd63dd338e7cf7ab Pull Request resolved: #318

torchtitan/parallelisms/parallelize_llama.py

train.py

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: f5bd75a194f06879a8fcc3e8ffd34f7824bcfc7f Pull Request resolved: #318

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: d73eb57829322e9909c5e9b541f4c3c65311f4e2 Pull Request resolved: #318

kwen2501

LGTM.
Is there a conclusion with model.reshard()?

torchtitan/parallelisms/parallelize_llama.py

torchtitan/parallelisms/pipelining_utils.py

torchtitan/parallelisms/parallelize_llama.py

[ghstack-poisoned]

wanchaol

This looks reasonablely good to me, have some additional comments inlined.

torchtitan/config_manager.py

wanchaol · 2024-05-21T18:11:00Z

torchtitan/parallelisms/parallelize_llama.py

+            if parallel_dims.dp_enabled
+            else torch.float32,
+            device=device,
+        )


Can we encapsulate L209-L258 into a input_output_shape_inference function and put that into pipeline_utils?

i think those lines are llama-model specific? thats why i kept them here. (The plan is to get rid of these still by providing a runtime shape inference thing, but it might not be the highest priority).

test_runner.py

torchtitan/parallelisms/parallelize_llama.py

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: 765946ac4e0969c08da8d1e722f15fc94eecde34 Pull Request resolved: #318

[ghstack-poisoned]

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: d7e0a1342bc97d6f1bba9e647234d90688ad708f Pull Request resolved: #318

wanchaol · 2024-05-22T00:00:17Z

@wconstab looks like CI is failing now, is it because the APIs for PP not in nightly yet? If so we should probably wait until the nightly is there and then reland this

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: d7e0a1342bc97d6f1bba9e647234d90688ad708f Pull Request resolved: #318

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: d7e0a1342bc97d6f1bba9e647234d90688ad708f Pull Request resolved: pytorch#318

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: 7a1b6ea024726bc7bf2430854c8088b77ff4e29e Pull Request resolved: #318

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: d7e0a1342bc97d6f1bba9e647234d90688ad708f Pull Request resolved: #318

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: d02cc759b34ff53035ac05f6121f0bac255b99fe Pull Request resolved: #318

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: e07b70e49c35445e3f0565f8b9c2038b4c96afd4 Pull Request resolved: #318

runs PP+DP and PP+TP without issue, runs PP+TP+DP with decreasing loss, but fails DCP save Supports only simple schedules currently, gpipe and 1f1b. Ads cmdline/toml arg for specifiying split points, in a unified way between tracer or manual frontend. e.g. user can specifiy "layers.2,layers.4" as split points. Currently uses manual frontend by default, but allows specifying tracer frontend. Tracer frontend requires working around additional compatibility limitations, indicated by raising assertions, and is not ready for wider use yet. ghstack-source-id: d7e0a1342bc97d6f1bba9e647234d90688ad708f Pull Request resolved: pytorch#318

Update

90ce59b

[ghstack-poisoned]

This was referenced May 9, 2024

Add Pipeline Parallel (and 2D PP+FSDP) support #161

Closed

WIP apply PP manually #308

Closed

facebook-github-bot added the CLA Signed This label is managed by the Meta Open Source bot. label May 9, 2024

wconstab added a commit that referenced this pull request May 9, 2024

Mock up a PP config UX

737999a

ghstack-source-id: 318831bda97c4639b745832e7f33f0b616a479ee Pull Request resolved: #318

tianyu-l reviewed May 9, 2024

View reviewed changes

torchtitan/config_manager.py Show resolved Hide resolved

torchtitan/config_manager.py Show resolved Hide resolved

wconstab added 2 commits May 9, 2024 17:03

Update

5b9bf7e

[ghstack-poisoned]

Update

c8b513f

[ghstack-poisoned]

wconstab changed the title ~~Mock up a PP config UX~~ Add Pipeline Parallel (and 2D PP+FSDP) support May 10, 2024

Update

c2b064d

[ghstack-poisoned]

wconstab commented May 10, 2024

View reviewed changes

torchtitan/parallelisms/parallelize_llama.py Outdated Show resolved Hide resolved

wconstab commented May 10, 2024

View reviewed changes

Update

2f0fc1f

[ghstack-poisoned]

This was referenced May 11, 2024

Refactor freqs_cis slice to be safer for PP #321

Merged

Make Transformer tolerate missing layers for PP #322

Merged

tianyu-l reviewed May 11, 2024

View reviewed changes

train.py Outdated Show resolved Hide resolved

train.py Outdated Show resolved Hide resolved

Update

0bf4db7

[ghstack-poisoned]

Update

6adc9bc

[ghstack-poisoned]

Update

7d0b919

[ghstack-poisoned]

wconstab commented May 14, 2024

View reviewed changes

torchtitan/parallelisms/parallelize_llama.py Outdated Show resolved Hide resolved

wconstab commented May 14, 2024

View reviewed changes

train.py Outdated Show resolved Hide resolved

Update

67cfc96

[ghstack-poisoned]

wconstab mentioned this pull request May 17, 2024

Make test_runner.py warn on non-empty output dir #343

Merged

kwen2501 approved these changes May 17, 2024

View reviewed changes

torchtitan/parallelisms/parallelize_llama.py Outdated Show resolved Hide resolved

torchtitan/parallelisms/pipelining_utils.py Outdated Show resolved Hide resolved

kwen2501 reviewed May 17, 2024

View reviewed changes

torchtitan/parallelisms/parallelize_llama.py Outdated Show resolved Hide resolved

torchtitan/parallelisms/parallelize_llama.py Outdated Show resolved Hide resolved

Update

b41f3c8

[ghstack-poisoned]

wconstab mentioned this pull request May 17, 2024

Add 3D support #344

Merged

Update

729f1ea

[ghstack-poisoned]

wconstab mentioned this pull request May 18, 2024

Add a 3-stage PP config #345

Closed

Update

bc814d3

[ghstack-poisoned]

wconstab mentioned this pull request May 20, 2024

Expose mixed_precision dtype arguments #348

Merged

wconstab added 3 commits May 20, 2024 11:33

Update

2690211

[ghstack-poisoned]

Update

8fce38b

[ghstack-poisoned]

Update

6665a52

[ghstack-poisoned]

wanchaol approved these changes May 21, 2024

View reviewed changes

Update

d7463ab

[ghstack-poisoned]

Update

0179b4b

[ghstack-poisoned]

wconstab merged commit 0179b4b into gh/wconstab/12/base May 21, 2024
4 checks passed

wconstab deleted the gh/wconstab/12/head branch May 21, 2024 23:27

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Pipeline Parallel (and 2D PP+FSDP) support #318

Add Pipeline Parallel (and 2D PP+FSDP) support #318

wconstab commented May 9, 2024 •

edited

Loading

wconstab May 10, 2024

kwen2501 left a comment

wanchaol left a comment

wanchaol May 21, 2024

wconstab May 21, 2024

wanchaol commented May 22, 2024

Add Pipeline Parallel (and 2D PP+FSDP) support #318

Add Pipeline Parallel (and 2D PP+FSDP) support #318

Conversation

wconstab commented May 9, 2024 • edited Loading

wconstab May 10, 2024

Choose a reason for hiding this comment

kwen2501 left a comment

Choose a reason for hiding this comment

wanchaol left a comment

Choose a reason for hiding this comment

wanchaol May 21, 2024

Choose a reason for hiding this comment

wconstab May 21, 2024

Choose a reason for hiding this comment

wanchaol commented May 22, 2024

wconstab commented May 9, 2024 •

edited

Loading