Skip to content

[SPMD] Expose apply_backward_optimization_barrier#7477

Merged
alanwaketan merged 2 commits intomasterfrom
alanwaketan/spmd
Jun 26, 2024
Merged

[SPMD] Expose apply_backward_optimization_barrier#7477
alanwaketan merged 2 commits intomasterfrom
alanwaketan/spmd

Conversation

@alanwaketan
Copy link
Copy Markdown
Collaborator

@alanwaketan alanwaketan commented Jun 25, 2024

Summary:
This PR exposes apply_backward_optimization_barrier to the spmd namespace.

Test Plan:
N.A.

@alanwaketan alanwaketan requested review from JackCaoG and jonb377 June 25, 2024 21:06
@alanwaketan alanwaketan self-assigned this Jun 25, 2024
@alanwaketan
Copy link
Copy Markdown
Collaborator Author

alanwaketan commented Jun 25, 2024

Thanks, Jack!

@alanwaketan alanwaketan merged commit 6894a08 into master Jun 26, 2024
@alanwaketan alanwaketan deleted the alanwaketan/spmd branch June 26, 2024 01:15
@bhavya01
Copy link
Copy Markdown
Collaborator

bhavya01 commented Jul 17, 2024 •
edited
Loading

@alanwaketan Should we also merge this change in the 2.4 release? I saw test failure with latest 2.4 wheel

[2024-07-17, 16:47:21 UTC] {logging_mixin.py:150} WARNING - Traceback (most recent call last):
  File "/home/ml-auto-solutions/transformers/examples/pytorch/language-modeling/run_clm.py", line 873, in <module>
    main()
  File "/home/ml-auto-solutions/transformers/examples/pytorch/language-modeling/run_clm.py", line 644, in main
[2024-07-17, 16:47:21 UTC] {logging_mixin.py:150} WARNING -     xs.apply_backward_optimization_barrier(model.model.layers[i])
AttributeError: module 'torch_xla.distributed.spmd' has no attribute 'apply_backward_optimization_barrier'
    ```

@JackCaoG
Copy link
Copy Markdown
Collaborator

JackCaoG commented Jul 17, 2024

which model failed?

@bhavya01
Copy link
Copy Markdown
Collaborator

bhavya01 commented Jul 17, 2024 •
edited
Loading

llama2-train-spmd

Still looking into why this passed in the previous test run
image

@bhavya01
Copy link
Copy Markdown
Collaborator

bhavya01 commented Jul 17, 2024

I realized that this is failing because of pytorch-tpu/transformers@ccf5b15

I replaced torch_xla.experimental.xla_sharding with torch_xla.distributed.spmd in the test and 2.4 release doesn't expose apply_optimization_barrier through the latter. The test passes locally with the fix.

@alanwaketan
Copy link
Copy Markdown
Collaborator Author

alanwaketan commented Jul 17, 2024

Sure, can I still backport things? I might have 2-3 PRs need to be back ported still?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@JackCaoG JackCaoG JackCaoG approved these changes

@jonb377 jonb377 Awaiting requested review from jonb377

Assignees

@alanwaketan alanwaketan

Labels

backport_2.4

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@alanwaketan @bhavya01 @JackCaoG