Skip to content
New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

WIP: TensorParallel with new strategy #1421

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Conversation

awaelchli
Copy link
Contributor

Shows how the new ModelParallelStrategy could be applied in generate/tp.py.

Some caveats:

  • Need to apply quantization in the parallelize function (to be done)
  • DTensors don't work with torch.inference_mode

There is room to apply more parallelism. For example, the RMSNorm could leverage SequenceParallel.

# for free to join this conversation on GitHub. Already have an account? # to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant