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.

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

Add batch dim idx to support latest deepspeed DistributedAttention #1725

Merged
merged 3 commits into from
Feb 6, 2025

Conversation

bhargaveede
Copy link
Collaborator

@bhargaveede bhargaveede commented Jan 27, 2025

https://github.com/microsoft/DeepSpeed/commits/master/deepspeed/sequence/layer.py
With latest changes in DistributedAttention, we need to add batch_dim_idx and rotary_pos_emb to Distributed Attention.

This change brings changes to add the batch_dim_idx in line with new changes.

deepspeedai/DeepSpeed@ffe0af2

Fixes # (issue)

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you make sure to update the documentation with your changes?
  • Did you write any new necessary tests?

…Attention (#37)

* Temp patch for batch_dim

* Adding batch_dim_idx as per latest deepspeed

* Update modeling_llama.py
@HuggingFaceDocBuilderDev

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@bhargaveede
Copy link
Collaborator Author

@regisss
Distributed Attention implementation has been changed, and function signature is changed in Synapse 1.20.
To use Distributed Attention, we have to provide batch dimension and additional rope parameter from the model side.

Please review it and merge along with other 1.20 Dependent PRs

@libinta libinta changed the title [SW-207148] Add batch dim idx to support latest deepspeed DistributedAttention Add batch dim idx to support latest deepspeed DistributedAttention Feb 3, 2025
@libinta libinta added the run-test Run CI for PRs from external contributors label Feb 3, 2025
Copy link
Collaborator

@regisss regisss left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I left one comments to address + @mlapinskix's comments

optimum/habana/transformers/models/llama/modeling_llama.py Outdated Show resolved Hide resolved
@bhargaveede
Copy link
Collaborator Author

@regisss Addressed review comments. Can you review and merge

@regisss regisss changed the base branch from main to synapse_1_20 February 6, 2025 09:55
Copy link
Collaborator

@regisss regisss left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@regisss regisss merged commit bedc041 into huggingface:synapse_1_20 Feb 6, 2025
4 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
run-test Run CI for PRs from external contributors synapse1.20
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants