Skip to content

Tensor parallelism #374

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
Merged
merged 26 commits into from
Jul 22, 2025
Merged

Tensor parallelism #374

merged 26 commits into from
Jul 22, 2025

Conversation

le1nux
Copy link
Member

@le1nux le1nux commented Jun 11, 2025 •
edited
Loading

What does this PR do?

This PR adds support for Tensor Parallelism (including Sequence Parallelism).
Additionally, this PR adds a debugging toolkit to track the input and output tensors during a forward pass, gradients during the backward pass and weight tensors.
Tensors can bei either normal Tensors or DTensors.

Checklist before submitting final PR

  • My PR is minimal and addresses one issue in isolation
  • I have merged the latest version of the target branch into this feature branch
  • I have reviewed my own code w.r.t. correct implementation, missing type hints, proper documentation, etc.
  • I have run a sample config for model training
  • I have checked that all tests run through (python tests/tests.py)
  • I have updated the internal changelog (CHANGELOG_DEV.md)

@le1nux le1nux changed the base branch from main to fsdp2_activation_checkpointing June 11, 2025 21:46
le1nux added 16 commits June 12, 2025 23:22
@le1nux
feat: added TP for attention
19f5298
@le1nux
refactor: TP now runs through
d4b4d18
@le1nux
refactor: changed lm_head TP from rowwise to columnwise
4142ac7
@le1nux
feat: integrated draft for sequence parallelism
7eb6d68
@le1nux
chore: added documentation to TP
efec29f
@le1nux
chore: added documentation for TP attn module input prepraration
a9211bb
@le1nux
feat: added model debugging component
73c81a0
@le1nux
feat: added TP support for absolute positional embeddings
795fb79
@le1nux
feat: improved train step logging
104a68d
@le1nux
feat: addded model step analyzer notebook
157ddd6
@le1nux
refactor: improved logs of debugged model
69c9712
@le1nux
refactor: drafted multi dim TP support in dataloader
89c670e
@le1nux
refactor: minor fix in FQN-based TP intialization
8369226
@le1nux
chore: dp rank now configurable in sampler
4004dac
@le1nux
chore: cleaned up TP implementation
80f51e0
@le1nux
feat: added tensor parallelism tests
fb585aa
@le1nux le1nux marked this pull request as ready for review July 16, 2025 11:57
@le1nux le1nux requested review from flxst and rrutmann July 16, 2025 11:57
flxst
flxst requested changes Jul 16, 2025
View reviewed changes
Copy link
Member

@flxst flxst left a comment

Choose a reason for hiding this comment

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

Great work!

I think there is a problem regarding TP + GELU, see my comment in src/modalities/models/model_factory.py.

rrutmann
rrutmann requested changes Jul 16, 2025
View reviewed changes
Copy link
Collaborator

@rrutmann rrutmann left a comment

Choose a reason for hiding this comment

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

Looks good to me, I just requested some minor changes in the test

@le1nux
Copy link
Member Author

le1nux commented Jul 18, 2025

Great work!

I think there is a problem regarding TP + GELU, see my comment in src/modalities/models/model_factory.py.

I have added GELU support now for TP (including unit tests)

@le1nux le1nux requested review from rrutmann and flxst July 18, 2025 08:41
flxst
flxst requested changes Jul 18, 2025
View reviewed changes
Copy link
Member

@flxst flxst left a comment

Choose a reason for hiding this comment

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

Changes look good! Only some minor issues left, see comments.

rrutmann
rrutmann approved these changes Jul 18, 2025
View reviewed changes
Copy link
Collaborator

@rrutmann rrutmann left a comment

Choose a reason for hiding this comment

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

LGTM :)

@le1nux le1nux requested review from flxst and rrutmann July 18, 2025 13:50
flxst
flxst approved these changes Jul 18, 2025
View reviewed changes
Copy link
Member

@flxst flxst left a comment

Choose a reason for hiding this comment

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

LGTM

@flxst flxst mentioned this pull request Jul 21, 2025
Merged
5 tasks
Base automatically changed from fsdp2_activation_checkpointing to main July 22, 2025 14:19
@le1nux le1nux merged commit 1e4d28e into main Jul 22, 2025
4 of 6 checks passed
@le1nux le1nux deleted the tensor_parallelism branch July 22, 2025 21:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@flxst flxst flxst approved these changes

@rrutmann rrutmann rrutmann approved these changes

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

Successfully merging this pull request may close these issues.
3 participants
@le1nux @flxst @rrutmann