Migrate internal tbe/* and torchrec callers to canonical type packages#5780
Closed
q10 wants to merge 1 commit into
Closed
Migrate internal tbe/* and torchrec callers to canonical type packages#5780q10 wants to merge 1 commit into
q10 wants to merge 1 commit into
Conversation
Summary: ...canonical packages where those types now live (`fbgemm_gpu.tbe.config`, `fbgemm_gpu.tbe.cache.cache_config`, `fbgemm_gpu.tbe.ssd.ssd_config`). Also remove the eager `KVEmbeddingInference` re-export from `tbe/cache/__init__.py` (and update its 1 in-tree caller `bench/tbe/tbe_kv_benchmark.py`) so that loading `tbe.cache` no longer transitively imports the heavy `kv_embedding_ops_inference` module. Fold the one stray external caller (`torchrec/distributed/batched_embedding_kernel.py`) that D103493980 missed: `BackendType` is consolidated into the existing `from fbgemm_gpu.tbe.ssd import (...)` block. After this diff, no `fbgemm_gpu/` source file depends on the shell for any type that has a canonical home elsewhere (only `construct_cache_state` and `get_new_embedding_location` still come from the shell since they have not yet been canonicalised). This is the prerequisite for D103477971: it breaks the latent circular import (ops_common -> tbe.ssd.ssd_config -> tbe/ssd/__init__.py -> tbe/ssd/inference.py -> ops_common) at its source, allowing D103477971 to convert ops_common.py to a re-export shell without relocating ssd_config.py outside the tbe.ssd package. Files migrated (internal, fbgemm_gpu/): - fbgemm_gpu/tbe/ssd/inference.py, inference_serving.py, training.py - fbgemm_gpu/tbe/cache/__init__.py, kv_embedding_ops_inference.py - fbgemm_gpu/tbe/bench/benchmark_click_interface.py, embedding_ops_common_config.py - fbgemm_gpu/split_embedding_configs.py, split_embedding_inference_converter.py - fbgemm_gpu/split_table_batched_embeddings_ops.py - fbgemm_gpu/split_table_batched_embeddings_ops_inference.py - fbgemm_gpu/split_table_batched_embeddings_ops_training.py - fbgemm_gpu/split_table_batched_embeddings_ops_training_common.py - fbgemm_gpu/sparse_ops.py - fbgemm_gpu/bench/tbe/tbe_kv_benchmark.py Files migrated (external): - torchrec/distributed/batched_embedding_kernel.py: `BackendType` line 37 folded into the existing `tbe.ssd` package umbrella block. Backward-compat fallback paths inside torchrec (D105254047's try/except ImportError shims in embedding_lookup.py / sharding_config.py / types.py) are intentionally LEFT UNTOUCHED - they protect old base-layer fbgemm_gpu consumers and remain dead code on master. No public-API change. The legacy shell still re-exports every type, so external callers (and the pyper_models snapshot copies of older torchrec) are unaffected. Reviewed By: cthi Differential Revision: D105789345
Contributor
|
@q10 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D105789345. |
Contributor
|
This pull request has been merged in 8e5780d. |
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.
Summary:
...canonical packages where those types now live
(
fbgemm_gpu.tbe.config,fbgemm_gpu.tbe.cache.cache_config,fbgemm_gpu.tbe.ssd.ssd_config). Also remove the eagerKVEmbeddingInferencere-export fromtbe/cache/__init__.py(andupdate its 1 in-tree caller
bench/tbe/tbe_kv_benchmark.py) so thatloading
tbe.cacheno longer transitively imports the heavykv_embedding_ops_inferencemodule. Fold the one stray externalcaller (
torchrec/distributed/batched_embedding_kernel.py) thatD103493980 missed:
BackendTypeis consolidated into the existingfrom fbgemm_gpu.tbe.ssd import (...)block.After this diff, no
fbgemm_gpu/source file depends on the shell forany type that has a canonical home elsewhere (only
construct_cache_stateand
get_new_embedding_locationstill come from the shell since theyhave not yet been canonicalised). This is the prerequisite for D103477971:
it breaks the latent circular import (ops_common -> tbe.ssd.ssd_config ->
tbe/ssd/init.py -> tbe/ssd/inference.py -> ops_common) at its source,
allowing D103477971 to convert ops_common.py to a re-export shell without
relocating ssd_config.py outside the tbe.ssd package.
Files migrated (internal, fbgemm_gpu/):
Files migrated (external):
BackendTypeline 37folded into the existing
tbe.ssdpackage umbrella block.Backward-compat fallback paths inside torchrec (D105254047's
try/except ImportError shims in embedding_lookup.py / sharding_config.py /
types.py) are intentionally LEFT UNTOUCHED - they protect old base-layer
fbgemm_gpu consumers and remain dead code on master.
No public-API change. The legacy shell still re-exports every type, so
external callers (and the pyper_models snapshot copies of older torchrec)
are unaffected.
Reviewed By: cthi
Differential Revision: D105789345