Add option to use gather to select indices in EC #3479
Closed
+26
−3
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
meta-cla
bot
added
the
CLA Signed
|
@elokrainz has exported this pull request. If you are a Meta employee, you can view the originating Diff in D85309309. |
elokrainz
pushed a commit
to elokrainz/torchrec
that referenced
this pull request
Oct 24, 2025
Summary:
Due to atomic add in torch.index_select, the backward performance sometimes is bad comparing with gather. In this diff, it provides users with control over the indexing process and select the suitable operator based on specific cases.
Perf comparison on pure operators(forward+backward)
2D Embedding, No Repetition
Config: shape=(1000000, 256), dim=0, indices=100000, unique=95300 (95.3%)
Method Time (s) Speedup Status
torch.gather 0.9439 1.00 x 🏆
torch.index_select 1.0509 0.90 x
2D Embedding, Low Repetition
Config: shape=(1000000, 256), dim=0, indices=100000, unique=48732 (48.7%)
Method Time (s) Speedup Status
torch.gather 0.9076 1.00 x 🏆
torch.index_select 1.0415 0.87 x
2D Embedding, High Repetition
Config: shape=(1000000, 256), dim=0, indices=250000, unique=9957 (4.0%)
Method Time (s) Speedup Status
torch.gather 1.2385 1.00 x 🏆
torch.index_select 1.6225 0.76 x
Small Vocab, Low Repetition
Config: shape=(1000, 256), dim=0, indices=2000, unique=635 (31.8%)
Method Time (s) Speedup Status
torch.gather 0.1502 1.00 x 🏆
torch.index_select 0.1763 0.85 x
Small Vocab, Very High Repetition
Config: shape=(1000, 256), dim=0, indices=100000, unique=625 (0.6%)
Method Time (s) Speedup Status
torch.gather 0.2626 1.00 x 🏆
torch.index_select 0.4126 0.64 x
Large Vocab, No Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=9996 (100.0%)
Method Time (s) Speedup Status
torch.gather 5.8014 1.00 x 🏆
torch.index_select 5.8184 1.00 x
Large Vocab, Low Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=5000 (50.0%)
Method Time (s) Speedup Status
torch.gather 5.7912 1.00 x 🏆
torch.index_select 5.8137 1.00 x
Large Vocab, High Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=400 (4.0%)
Method Time (s) Speedup Status
torch.gather 5.7784 1.00 x 🏆
torch.index_select 5.8100 0.99 x
Mast Job Test:
baseline: fire-jingchang-f816557933
torch.index_select backward takes ~37ms
{F1982939713}
exp: fire-jingchang-f816355728
torch.gather backward takes ~10ms
{F1982939742}
Reviewed By: TroyGarden
Differential Revision: D85309309
elokrainz
force-pushed
the
export-D85309309
branch
from
f8a0b12 to
987f44c
Compare
elokrainz
pushed a commit
to elokrainz/torchrec
that referenced
this pull request
Oct 24, 2025
Summary: Pull Request resolved: meta-pytorch#3479 Due to atomic add in torch.index_select, the backward performance sometimes is bad comparing with gather. In this diff, it provides users with control over the indexing process and select the suitable operator based on specific cases. Perf comparison on pure operators(forward+backward) 2D Embedding, No Repetition Config: shape=(1000000, 256), dim=0, indices=100000, unique=95300 (95.3%) Method Time (s) Speedup Status torch.gather 0.9439 1.00 x 🏆 torch.index_select 1.0509 0.90 x 2D Embedding, Low Repetition Config: shape=(1000000, 256), dim=0, indices=100000, unique=48732 (48.7%) Method Time (s) Speedup Status torch.gather 0.9076 1.00 x 🏆 torch.index_select 1.0415 0.87 x 2D Embedding, High Repetition Config: shape=(1000000, 256), dim=0, indices=250000, unique=9957 (4.0%) Method Time (s) Speedup Status torch.gather 1.2385 1.00 x 🏆 torch.index_select 1.6225 0.76 x Small Vocab, Low Repetition Config: shape=(1000, 256), dim=0, indices=2000, unique=635 (31.8%) Method Time (s) Speedup Status torch.gather 0.1502 1.00 x 🏆 torch.index_select 0.1763 0.85 x Small Vocab, Very High Repetition Config: shape=(1000, 256), dim=0, indices=100000, unique=625 (0.6%) Method Time (s) Speedup Status torch.gather 0.2626 1.00 x 🏆 torch.index_select 0.4126 0.64 x Large Vocab, No Repetition Config: shape=(10000000, 256), dim=0, indices=10000, unique=9996 (100.0%) Method Time (s) Speedup Status torch.gather 5.8014 1.00 x 🏆 torch.index_select 5.8184 1.00 x Large Vocab, Low Repetition Config: shape=(10000000, 256), dim=0, indices=10000, unique=5000 (50.0%) Method Time (s) Speedup Status torch.gather 5.7912 1.00 x 🏆 torch.index_select 5.8137 1.00 x Large Vocab, High Repetition Config: shape=(10000000, 256), dim=0, indices=10000, unique=400 (4.0%) Method Time (s) Speedup Status torch.gather 5.7784 1.00 x 🏆 torch.index_select 5.8100 0.99 x Mast Job Test: baseline: fire-jingchang-f816557933 torch.index_select backward takes ~37ms {F1982939713} exp: fire-jingchang-f816355728 torch.gather backward takes ~10ms {F1982939742} Reviewed By: TroyGarden Differential Revision: D85309309
elokrainz
force-pushed
the
export-D85309309
branch
from
987f44c to
5458b03
Compare
Summary:
Due to atomic add in torch.index_select, the backward performance sometimes is bad comparing with gather. In this diff, it provides users with control over the indexing process and select the suitable operator based on specific cases.
Perf comparison on pure operators(forward+backward)
2D Embedding, No Repetition
Config: shape=(1000000, 256), dim=0, indices=100000, unique=95300 (95.3%)
Method Time (s) Speedup Status
torch.gather 0.9439 1.00 x 🏆
torch.index_select 1.0509 0.90 x
2D Embedding, Low Repetition
Config: shape=(1000000, 256), dim=0, indices=100000, unique=48732 (48.7%)
Method Time (s) Speedup Status
torch.gather 0.9076 1.00 x 🏆
torch.index_select 1.0415 0.87 x
2D Embedding, High Repetition
Config: shape=(1000000, 256), dim=0, indices=250000, unique=9957 (4.0%)
Method Time (s) Speedup Status
torch.gather 1.2385 1.00 x 🏆
torch.index_select 1.6225 0.76 x
Small Vocab, Low Repetition
Config: shape=(1000, 256), dim=0, indices=2000, unique=635 (31.8%)
Method Time (s) Speedup Status
torch.gather 0.1502 1.00 x 🏆
torch.index_select 0.1763 0.85 x
Small Vocab, Very High Repetition
Config: shape=(1000, 256), dim=0, indices=100000, unique=625 (0.6%)
Method Time (s) Speedup Status
torch.gather 0.2626 1.00 x 🏆
torch.index_select 0.4126 0.64 x
Large Vocab, No Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=9996 (100.0%)
Method Time (s) Speedup Status
torch.gather 5.8014 1.00 x 🏆
torch.index_select 5.8184 1.00 x
Large Vocab, Low Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=5000 (50.0%)
Method Time (s) Speedup Status
torch.gather 5.7912 1.00 x 🏆
torch.index_select 5.8137 1.00 x
Large Vocab, High Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=400 (4.0%)
Method Time (s) Speedup Status
torch.gather 5.7784 1.00 x 🏆
torch.index_select 5.8100 0.99 x
Mast Job Test:
baseline: fire-jingchang-f816557933
torch.index_select backward takes ~37ms
{F1982939713}
exp: fire-jingchang-f816355728
torch.gather backward takes ~10ms
{F1982939742}
Reviewed By: TroyGarden
Differential Revision: D85309309
elokrainz
force-pushed
the
export-D85309309
branch
from
5458b03 to
885c63d
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Summary:
Due to atomic add in torch.index_select, the backward performance sometimes is bad comparing with gather. In this diff, it provides users with control over the indexing process and select the suitable operator based on specific cases.
Perf comparison on pure operators(forward+backward)
2D Embedding, No Repetition
Config: shape=(1000000, 256), dim=0, indices=100000, unique=95300 (95.3%)
Method Time (s) Speedup Status
torch.gather 0.9439 1.00 x 🏆
torch.index_select 1.0509 0.90 x
2D Embedding, Low Repetition
Config: shape=(1000000, 256), dim=0, indices=100000, unique=48732 (48.7%)
Method Time (s) Speedup Status
torch.gather 0.9076 1.00 x 🏆
torch.index_select 1.0415 0.87 x
2D Embedding, High Repetition
Config: shape=(1000000, 256), dim=0, indices=250000, unique=9957 (4.0%)
Method Time (s) Speedup Status
torch.gather 1.2385 1.00 x 🏆
torch.index_select 1.6225 0.76 x
Small Vocab, Low Repetition
Config: shape=(1000, 256), dim=0, indices=2000, unique=635 (31.8%)
Method Time (s) Speedup Status
torch.gather 0.1502 1.00 x 🏆
torch.index_select 0.1763 0.85 x
Small Vocab, Very High Repetition
Config: shape=(1000, 256), dim=0, indices=100000, unique=625 (0.6%)
Method Time (s) Speedup Status
torch.gather 0.2626 1.00 x 🏆
torch.index_select 0.4126 0.64 x
Large Vocab, No Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=9996 (100.0%)
Method Time (s) Speedup Status
torch.gather 5.8014 1.00 x 🏆
torch.index_select 5.8184 1.00 x
Large Vocab, Low Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=5000 (50.0%)
Method Time (s) Speedup Status
torch.gather 5.7912 1.00 x 🏆
torch.index_select 5.8137 1.00 x
Large Vocab, High Repetition
Config: shape=(10000000, 256), dim=0, indices=10000, unique=400 (4.0%)
Method Time (s) Speedup Status
torch.gather 5.7784 1.00 x 🏆
torch.index_select 5.8100 0.99 x
Differential Revision: D85309309