Skip to content

Back out "Forward-compat nfp8_dtype fallback for torch.package re-export version blends" #6091

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

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

Back out "Forward-compat nfp8_dtype fallback for torch.package re-export version blends" #6091

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 5 additions & 22 deletions fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,28 +28,11 @@
# @manual=//deeplearning/fbgemm/fbgemm_gpu/codegen:split_embedding_codegen_lookup_invokers
import fbgemm_gpu.split_embedding_codegen_lookup_invokers as invokers
from fbgemm_gpu.config import FeatureGate, FeatureGateName
from fbgemm_gpu.split_embedding_configs import EmbOptimType as OptimType, SparseType

try:
from fbgemm_gpu.split_embedding_configs import nfp8_dtype
except ImportError:
# Forward-compat for torch.package re-export version blends (S685573): a
# model's frozen `split_embedding_configs`, bundled alongside this (newer)
# module at GMPP re-export time, may predate `nfp8_dtype` (added in
# D113263502). Import it defensively and fall back to the pre-arch-aware
# selection so this module finishes initializing regardless of the bundled
# `split_embedding_configs` version -- otherwise the failed top-level import
# leaves a poisoned partial module cached by the torch.package importer and
# a later `from ...ops_training import SplitTableBatchedEmbeddingBagsCodegen`
# fails. Mirrors the D110788115 `sparse_type_to_int` fallback pattern.
def nfp8_dtype() -> torch.dtype:
return (
torch.float8_e4m3fnuz
if torch.version.hip is not None
else torch.float8_e4m3fn
)


from fbgemm_gpu.split_embedding_configs import (
EmbOptimType as OptimType,
nfp8_dtype,
SparseType,
)
from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
BoundsCheckMode,
CacheAlgorithm,
Expand Down
Loading