From 4224e1ad91b583e5b46b5bcefe27999f9260bf7e Mon Sep 17 00:00:00 2001 From: Srijan Upadhyay Date: Thu, 9 Jul 2026 11:25:23 +0530 Subject: [PATCH] Add transformer_engine.test() installation smoke check Add a top-level test() entry point so users can verify a Transformer Engine installation without cloning the repository, similar to numpy.test(). It reuses the existing PyPI sanity check and, for each installed framework backend, runs a minimal forward pass on the current device. Checks that need a GPU are skipped when no CUDA device is available, and the PyTorch check reports the device compute capability and which low-precision formats it supports. Closes #170 Signed-off-by: Srijan Upadhyay --- transformer_engine/__init__.py | 79 ++++++++++++++++++++++++++++++++++ 1 file changed, 79 insertions(+) diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index 480a2e9a06..701bb4f2d7 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -109,3 +109,82 @@ def require_nccl_ep() -> None: ) __version__ = str(metadata.version("transformer_engine")) + + +def test(verbose: bool = True) -> bool: + """Run smoke checks to verify the Transformer Engine installation. + + Confirms the package is installed correctly and, when the PyTorch backend is + available, runs a minimal functional check on the current device. Checks that + need a GPU are skipped when no CUDA device is available, so the call is safe + to run anywhere as an installation sanity check. + + Parameters + ---------- + verbose : bool, default = True + Print the result of each check. + + Returns + ------- + bool + ``True`` if every executed check passed, ``False`` otherwise. + """ + # ponytail: smoke test, not the full suite. Covers install integrity and a + # single PyTorch forward pass; deeper coverage stays in tests/ and qa/. + results = [] + + def _record(name: str, passed: bool, detail: str = "") -> None: + results.append(passed) + if verbose: + status = "PASS" if passed else "FAIL" + print(f"[{status}] {name}" + (f": {detail}" if detail else "")) + + # Installation integrity (reuses the PyPI sanity check). + try: + from .common import sanity_checks_for_pypi_installation + + sanity_checks_for_pypi_installation() + _record("installation", True, f"transformer_engine {__version__}") + except Exception as err: # pylint: disable=broad-except + _record("installation", False, str(err)) + + # PyTorch backend, checked only if it is available. A backend that is present + # but fails at runtime is reported, not skipped. + try: + from . import pytorch as te + except (ImportError, FileNotFoundError): + # FileNotFoundError mirrors the module-level guard: torch is installed but + # the Transformer Engine PyTorch extension shared object is missing. + te = None + if te is not None: + try: + import torch + + if not torch.cuda.is_available(): + _record("pytorch", True, "imported; no CUDA device, GPU checks skipped") + else: + major, minor = te.get_device_compute_capability() + dtype = torch.bfloat16 if te.is_bf16_available() else torch.float16 + with torch.no_grad(): + model = te.Linear(16, 16, params_dtype=dtype, device="cuda") + out = model(torch.zeros(4, 16, dtype=dtype, device="cuda")) + if tuple(out.shape) != (4, 16): + raise RuntimeError(f"unexpected output shape {tuple(out.shape)}") + formats = [ + name + for name, available in ( + ("FP8", te.is_fp8_available()), + ("MXFP8", te.is_mxfp8_available()), + ("NVFP4", te.is_nvfp4_available()), + ) + if available + ] + detail = f"sm_{major}{minor}, {dtype}, formats: {', '.join(formats) or 'none'}" + _record("pytorch", True, detail) + except Exception as err: # pylint: disable=broad-except + _record("pytorch", False, str(err)) + + passed = all(results) + if verbose: + print(f"\n{sum(results)}/{len(results)} checks passed.") + return passed