Skip codegen of empty PT2 wrappers for optimizers without backend support - #6043
Draft
jithunnair-amd wants to merge 1 commit into
Draft
Skip codegen of empty PT2 wrappers for optimizers without backend support#6043jithunnair-amd wants to merge 1 commit into
jithunnair-amd wants to merge 1 commit into
Conversation
…port
TBE PT2 codegen emitted a CPU and CUDA backend wrapper TU for every entry in
ALL_OPTIMIZERS, including the six deprecated optimizers (approx_sgd,
approx_rowwise_adagrad, approx_rowwise_adagrad_with_counter,
approx_rowwise_adagrad_with_weight_decay, rowwise_adagrad_with_weight_decay,
rowwise_weighted_adagrad), whose has_cpu_support and has_gpu_support are both
False. The wrapper templates guard their entire body on those flags, so for
these optimizers they render to byte-identical, effectively empty translation
units that only differ by filename.
Besides being wasted compilation, the six identical CUDA wrappers are a
liability under a preprocessor-off, path-insensitive compiler cache (e.g.
sccache classic mode): the cache key ignores the input/output path, so all six
collapse to one cached object that is copied to every output. On HIP each copy
then carries the same __hip_cuid symbol and the link fails with
"multiple definition of __hip_cuid_...".
Gate wrapper generation on the corresponding support flag (autograd, which
carries the deprecation stub needed for backward compatibility, is still
generated for every optimizer), and build the CMake wrapper source lists from
CPU_OPTIMIZERS / GPU_OPTIMIZERS instead of ALL_OPTIMIZERS. At current head those
groups exactly equal the sets with has_cpu_support / has_gpu_support True, so
the generated file set matches the listed sources.
Test Plan:
```
python -c "import tbe_sources as t; \
print(len([f for f in t.gen_gpu_files_training_pt2 if 'pt2_cuda_wrapper' in f]), \
len([f for f in t.gen_cpu_files_training_pt2 if 'pt2_cpu_wrapper' in f]))"
# 12 (11 GPU + 1 SSD) and 5, none for deprecated optimizers; 18 autograd TUs kept
```
Authored with assistance from Cursor
Co-authored-by: Cursor <cursoragent@cursor.com>
This was referenced Jul 22, 2026
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
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.
This PR tries to address an issue discovered in pytorch/pytorch#190601.
TBE PT2 codegen emitted a CPU and CUDA backend wrapper TU for every entry in
ALL_OPTIMIZERS, including the six deprecated optimizers (approx_sgd,approx_rowwise_adagrad,approx_rowwise_adagrad_with_counter,approx_rowwise_adagrad_with_weight_decay,rowwise_adagrad_with_weight_decay,rowwise_weighted_adagrad), whosehas_cpu_supportandhas_gpu_supportare both False. The wrapper templates guard their entire body on those flags, so for these optimizers they render to byte-identical, effectively empty translation units that only differ by filename.Besides being wasted compilation, the six identical CUDA wrappers are a liability under a preprocessor-off, path-insensitive compiler cache (e.g. sccache classic mode): the cache key ignores the input/output path, so all six collapse to one cached object that is copied to every output. On HIP each copy then carries the same
__hip_cuidsymbol and the link fails with "multiple definition of__hip_cuid_...".Gate wrapper generation on the corresponding support flag (autograd, which carries the deprecation stub needed for backward compatibility, is still generated for every optimizer), and build the CMake wrapper source lists from
CPU_OPTIMIZERS/GPU_OPTIMIZERSinstead ofALL_OPTIMIZERS. At current head those groups exactly equal the sets withhas_cpu_support/has_gpu_supportTrue, so the generated file set matches the listed sources.Note: this is one of two alternative fixes for the same issue. The other alternative is #6044
Authored with assistance from Cursor
Made with Cursor