Skip to content
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.

Sign up for GitHub

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

Jump to bottom

Fori loop simple case test test #7028

Closed
wants to merge 53 commits into from
Closed

Fori loop simple case test test #7028

wants to merge 53 commits into from

Conversation

ManfeiBai
Copy link
Collaborator

No description provided.

qihqi and others added 30 commits April 17, 2024 15:31
@qihqi
Add doc explaining how it works (pytorch#6937)
f44e1df
@vanbasten23
Fix profiling in benchmark script (pytorch#6934)
310f08d
@alanwaketan
[Pallas] Integrate FlashAttention with SPMD (pytorch#6935)
9f2b82d 
Summary:
This pull request integrating FlashAttention with SPMD. The way it works is to create a manual sharding region for the kernel which means we wraps all the inputs with enable_manual_sharding and all the outputs with disable_manual_sharding.

Added a new test file because the original test file is not SPMD aware.

Test Plan:
PJRT_DEVICE=TPU python test/test_pallas_spmd.py
@will-cromar
Use TPU build for CPU and GPU Python tests (pytorch#6921)
b2556d6
@ManfeiBai
[Fori_loop|While_loop] Create fori_loop.md (pytorch#6942)
0417d4d 
Add doc for fori_loop/while_loop and add simple user guide for simple test case
@yeounoh
Update XLA pin, 04/19/2024 (pytorch#6944)
2ec7706
@yeounoh
Update jinja and sphinx versions to address the vulnearbility (pytorc…
9ba844a 
…h#6946)
@ysiraichi
Make nms fallback by default. (pytorch#6933)
b06c9c7
@bdhirsh
revert expand test with dynamo (pytorch#6950)
62a2b11
@bhavya01
Lower embedding bag forward only (pytorch#6951)
46919a4
@dewitt
Cleanup the code example in the torch_xla2 README. (pytorch#6939)
6fd448d
@will-cromar @qihqi
[torch_xla2] Simplify developer setup steps (pytorch#6905)
89efd17 
Co-authored-by: qihqi <[email protected]>
@lsy323
Create rc13 trigger (pytorch#6956)
fa090a2
@qihqi
update torch deps to 2.3 (pytorch#6959)
b5574d8
@lsy323
update rc14 (pytorch#6962)
69eeace
@alanwaketan
[Doc] Update Pallas user guide (pytorch#6961)
a7749fa 
Summary:
This PR adds a preliminary user guide for Pallas.

Test Plan:
Skip CI.
@lsy323
Add final 2.3 trigger (pytorch#6963)
5369e7d
@will-cromar
Temporarily ignore torch commit in CI test (pytorch#6964)
76f7dd0
@alanwaketan
[Doc] Improve docker instructions (pytorch#6969)
af74c34 
Summary:
Clarify some caveats for using the TPU docker images in the landing page.

Test Plan:
Skip CI.
@wonjoolee95
Enable PagedAttention through Pallas (pytorch#6912)
6ed2026
@lsy323
Update readme for 2.3 releae (pytorch#6967)
0a204a6
@lsy323
Update GPU readme (pytorch#6968)
abe090a
@will-cromar
Write torch_xla.version.__torch_gitrev__ to file directly (pytorch#…
0054ec0 
…6966)
@bdhirsh
fix pytorch CI after pin update, change test to use assertLessEqual (p…
2a204e9 
…ytorch#6973)
@ManfeiBai
Update Openxla-pin to 04/24 (pytorch#6975)
4b48134
@JackCaoG
Move test_grad_checkpoint.py to tpu test list (pytorch#6976)
2bf59e0
@ManfeiBai
Revert "Update Openxla-pin to 04/24" (pytorch#6980)
023e2c8
@will-cromar
Update CODEOWNERS for build infrastructure (pytorch#6953)
3f5ff0f
@will-cromar
Move .torch_pin and handle in ansible (pytorch#6920)
b3be775
@JackCaoG
Update dynamo test to be less constrain (pytorch#6981)
b9a9449
will-cromar and others added 23 commits April 26, 2024 15:12
@will-cromar
Build CPP tests in new CI workflow (pytorch#6947)
b834e49
@will-cromar
Run TPU CI when label is on PR (pytorch#6984)
174f407
@qihqi
Add readme to call a model (lost due to merge conflicts) (pytorch#6986)
6443e59
@will-cromar
Fix permission issues during CI checkout (pytorch#6985)
6d01bb6
@ManfeiBai
[Revert Revert] Update OpenXLA-pin update to Apr25 (pytorch#6982)
971ebe1
@avikchaudhuri
Update test_export_fx_passes.py (pytorch#6972)
7527816
@will-cromar
Change name of CI documentation (pytorch#6994)
7cc78a6
@will-cromar
Rework docs push (pytorch#6954)
c91171d
@will-cromar
sudo rm leftover files in GHA (pytorch#6995)
0e032b1
@avikchaudhuri
Fixes to dynamic_shapes args in test_unbounded_dynamism.py (pytorch#6999
73b915b 
)
@will-cromar
Manually push to gh-pages branch (pytorch#6996)
d25f475
@ysiraichi
Re-land: dynamo expand test with view-replay. (pytorch#6958)
42db709
@JackCaoG
Move the nightly whl instruction out of the hide area (pytorch#7000)
87329ce
@will-cromar
Don't fail docs push if there's nothing to commit (pytorch#7001)
5f75290
@will-cromar
Complain when TensorFlow is installed (pytorch#7004)
4a5e238
@will-cromar
Clean up workspace before test (pytorch#7005)
b8f8fa9
@will-cromar
Tag CI build with git hash (pytorch#7003)
77bbf7f
@qihqi
fix addbmm opinfo (pytorch#6993)
2399e10
@qihqi
Fix more opinfo tests (pytorch#7008)
2907ab3
@wonjoolee95
Fix q dtype in paged attention kernel (pytorch#7011)
865836a
@will-cromar
Build upstream CI image on push to master (pytorch#6952)
4883f6f
@alanwaketan
[Pallas] Support segment ids in flash attention (pytorch#6943)
400bd0c 
Summary:
This PR is to add segment ids to the flash attention wrapper. The segment ids are a way to create an attention mask where each token can only attend to other tokens within the same segment. The mask is therefore a block diagonal matrix.

To support it, we further split the flash attention forward into tracing and execution part, and implement all the shape operations to make it compatible with the kernel.

Test Plan:
PJRT_DEVICE=TPU python test/test_pallas.py
@JackCaoG @will-cromar
Support pin pr number in new .torch_pin (pytorch#6998)
cbbefa2 
Co-authored-by: Will Cromar <[email protected]>
@ManfeiBai ManfeiBai closed this May 6, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@will-cromar will-cromar Awaiting requested review from will-cromar will-cromar will be requested when the pull request is marked ready for review will-cromar is a code owner

@JackCaoG JackCaoG Awaiting requested review from JackCaoG JackCaoG will be requested when the pull request is marked ready for review JackCaoG is a code owner

@mateuszlewko mateuszlewko Awaiting requested review from mateuszlewko mateuszlewko will be requested when the pull request is marked ready for review mateuszlewko is a code owner

@stgpetrovic stgpetrovic Awaiting requested review from stgpetrovic stgpetrovic will be requested when the pull request is marked ready for review stgpetrovic is a code owner

Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.