Fix VBE batch sizes not passed to request builder (#5653)#5653
Closed
gregmacnamara wants to merge 1 commit into
Closed
Fix VBE batch sizes not passed to request builder (#5653)#5653gregmacnamara wants to merge 1 commit into
gregmacnamara wants to merge 1 commit into
Conversation
Contributor
|
@gregmacnamara has exported this pull request. If you are a Meta employee, you can view the originating Diff in D94380814. |
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
06cf316 to
af6b699
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
36ac52f to
0d9ad95
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: Pull Request resolved: pytorch#5653 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
0d9ad95 to
e997bf5
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: Pull Request resolved: pytorch#5653 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
a42d0fc to
9c9b46b
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
9c9b46b to
2b9a287
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: Pull Request resolved: pytorch#5653 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
2b9a287 to
14ec6d9
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 17, 2026
Summary: Pull Request resolved: pytorch#5653 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Differential Revision: D94380814
14ec6d9 to
106e7ba
Compare
106e7ba to
ec004c8
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 27, 2026
Summary: X-link: facebookresearch/FBGEMM#2626 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Reviewed By: q10 Differential Revision: D94380814
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 27, 2026
Summary: X-link: facebookresearch/FBGEMM#2626 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Reviewed By: q10 Differential Revision: D94380814
ec004c8 to
082474f
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 27, 2026
Summary: X-link: facebookresearch/FBGEMM#2626 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Reviewed By: q10 Differential Revision: D94380814
082474f to
18c0c8b
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 27, 2026
Summary: X-link: facebookresearch/FBGEMM#2626 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Reviewed By: q10 Differential Revision: D94380814
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 27, 2026
Summary: X-link: facebookresearch/FBGEMM#2626 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Reviewed By: q10 Differential Revision: D94380814
18c0c8b to
c3f2f8a
Compare
Summary: X-link: facebookresearch/FBGEMM#2626 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Reviewed By: q10 Differential Revision: D94380814
c3f2f8a to
06f17b7
Compare
gregmacnamara
added a commit
to gregmacnamara/FBGEMM
that referenced
this pull request
Apr 27, 2026
Summary: X-link: facebookresearch/FBGEMM#2626 When `batch_size_per_feature_per_rank` is provided externally (e.g. from a sharding plan), it was not being forwarded to `_build_requests_jagged`. Both branches of the VBE/non-VBE conditional were passing the raw `batch_size_per_feature_per_rank` parameter (which could be None in the else branch) instead of the resolved `Bs_feature_rank` from `_generate_batch_sizes`. This fix: 1. Captures the externally-provided batch sizes into `Bs_feature_rank` when `batch_size_per_feature_per_rank` is set 2. Uses `_generate_batch_sizes` return value otherwise 3. Removes the duplicate if/else branches that were doing the same thing 4. Passes the resolved `Bs_feature_rank` to `_build_requests_jagged` Reviewed By: q10 Differential Revision: D94380814
Contributor
|
This pull request has been merged in b8dc9a1. |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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:
X-link: https://github.com/facebookresearch/FBGEMM/pull/2626
When
batch_size_per_feature_per_rankis provided externally (e.g. from asharding plan), it was not being forwarded to
_build_requests_jagged.Both branches of the VBE/non-VBE conditional were passing the raw
batch_size_per_feature_per_rankparameter (which could be None in the elsebranch) instead of the resolved
Bs_feature_rankfrom_generate_batch_sizes.This fix:
Bs_feature_rankwhenbatch_size_per_feature_per_rankis set_generate_batch_sizesreturn value otherwiseBs_feature_rankto_build_requests_jaggedReviewed By: q10
Differential Revision: D94380814