Add Transformer Engine compatibility shim, fallbacks, and modernize TE pins

cosmos_predict1/autoregressive/utils/parallel.py

@@ -21,10 +21,24 @@
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
 from torch.autograd import Function
 from torch.distributed import broadcast, get_process_group_ranks
-from transformer_engine.pytorch.jit import no_torch_dynamo
-from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
-from transformer_engine.pytorch.module.rmsnorm import RMSNorm as RMSNormTE
-from transformer_engine.pytorch.module.rmsnorm import _RMSNorm
+from cosmos_predict1.utils.te_compat import RMSNorm as RMSNormTE, TE_AVAILABLE
+
+try:
+    from transformer_engine.pytorch.jit import no_torch_dynamo
+except Exception:
+    def no_torch_dynamo():
+        def _wrap(fn):
+            return fn
+        return _wrap
+
+try:
+    from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
+except Exception:
+    class TransformerEngineBaseModule:
+        @staticmethod
+        def set_activation_dtype(module, inp):
+            del module
+            return inp.dtype
 
 from cosmos_predict1.utils import log
 
@@ -210,26 +224,11 @@ def __init__(self, hidden_size, process_group, **kwargs):
     def forward(self, inp: torch.Tensor) -> torch.Tensor:
         """RMSNorm FWD"""
 
-        # Set the activation type for AMP.
-        TransformerEngineBaseModule.set_activation_dtype(self, inp)
-
-        if torch.is_grad_enabled():
-            fwd_fn = _RMSNorm.apply
-            args = []
-        else:
-            fwd_fn = _RMSNorm.forward
-            args = [None]
-
-        args += (
-            inp,
-            AllReduceBWD.apply(self.weight, self.process_group),
-            self.eps,
-            self.fwd_rmsnorm_sm_margin,
-            self.bwd_rmsnorm_sm_margin,
-            self.inf_rmsnorm_sm_margin,
-            self.zero_centered_gamma,
-            torch.is_grad_enabled(),
-            self.activation_dtype,
-        )
+        # TE private autograd kernels were removed in newer versions; keep a stable torch-path fallback.
+        if not TE_AVAILABLE or not hasattr(self, "fwd_rmsnorm_sm_margin"):
+            normalized = torch.nn.functional.rms_norm(inp, inp.shape[-1:], self.weight, eps=self.eps)
+            return normalized
 
-        return fwd_fn(*args)
+        TransformerEngineBaseModule.set_activation_dtype(self, inp)
+        out = super().forward(inp)
+        return out

cosmos_predict1/diffusion/module/attention.py

@@ -21,7 +21,7 @@
 from torch import nn
 from torch.utils.checkpoint import checkpoint
 
-from cosmos_predict1.utils.te_compat import DotProductAttention, TE_AVAILABLE, apply_rotary_pos_emb
+from cosmos_predict1.utils.te_compat import DotProductAttention, RMSNorm, TE_AVAILABLE, apply_rotary_pos_emb
 
 # ---------------------- Feed Forward Network -----------------------
 
@@ -128,11 +128,7 @@ def get_normalization(name: str, channels: int):
     if name == "I":
         return nn.Identity()
     elif name == "R":
-        if TE_AVAILABLE:
-            import transformer_engine as te
-
-            return te.pytorch.RMSNorm(channels, eps=1e-6)
-        return nn.RMSNorm(channels, eps=1e-6)
+        return RMSNorm(channels, eps=1e-6)
     else:
         raise ValueError(f"Normalization {name} not found")
 

cosmos_predict1/utils/te_compat.py

@@ -40,6 +40,15 @@
     DotProductAttention = None
 
 
+try:
+    from transformer_engine.pytorch import RMSNorm as _te_rmsnorm_ctor
+except Exception:
+    try:
+        from transformer_engine.pytorch.module.rmsnorm import RMSNorm as _te_rmsnorm_ctor
+    except Exception:
+        _te_rmsnorm_ctor = None
+
+
 try:
     from megatron.core import InferenceParams as InferenceParams
 except Exception:
@@ -66,6 +75,31 @@ def split_along_dim(x: torch.Tensor, split_dim: int, split_sizes: Union[int, Seq
     return torch.split(x, tuple(split_sizes), dim=split_dim)
 
 
+def RMSNorm(hidden_size: int, eps: float = 1e-6, **kwargs):
+    """Create an RMSNorm module using TE when available, else pure torch fallback."""
+    if _te_rmsnorm_ctor is not None:
+        try:
+            return _te_rmsnorm_ctor(hidden_size=hidden_size, eps=eps, **kwargs)
+        except TypeError:
+            return _te_rmsnorm_ctor(hidden_size, eps=eps, **kwargs)
+
+    if hasattr(torch.nn, "RMSNorm"):
+        return torch.nn.RMSNorm(hidden_size, eps=eps, **kwargs)
+
+    class _RMSNormFallback(torch.nn.Module):
+        def __init__(self, hidden_dim: int, eps_: float):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.ones(hidden_dim))
+            self.eps = eps_
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            variance = x.pow(2).mean(dim=-1, keepdim=True)
+            x = x * torch.rsqrt(variance + self.eps)
+            return x * self.weight
+
+    return _RMSNormFallback(hidden_size, eps)
+
+
 def _rotate_half(x: torch.Tensor) -> torch.Tensor:
     x1, x2 = torch.chunk(x, 2, dim=-1)
     return torch.cat((-x2, x1), dim=-1)

requirements.txt

@@ -33,6 +33,6 @@ torch==2.8.0
 torchvision==0.23.0
 tqdm==4.66.5
 transformers==4.49.0
-transformer-engine==2.11.0
-transformer-engine-torch==2.11.0
+transformer-engine>=2.11,<3
+transformer-engine-torch>=2.11,<3
 nvidia-cudnn-cu12>=9.3

scripts/check_te_compat.py

@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+"""Smoke-check Transformer Engine compatibility for inference entrypoints."""
+
+from __future__ import annotations
+
+import importlib
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from cosmos_predict1.utils import te_compat
+
+ENTRYPOINT_MODULES = [
+    "cosmos_predict1.diffusion.inference.inference_inverse_renderer",
+    "cosmos_predict1.diffusion.inference.inference_forward_renderer",
+    "cosmos_predict1.diffusion.inference.text2world",
+]
+
+
+def main() -> int:
+    print(f"TE_AVAILABLE={te_compat.TE_AVAILABLE}")
+    print(f"TE_IMPORT_ERROR={te_compat.TE_IMPORT_ERROR!r}")
+    print(f"TE_DotProductAttention_available={te_compat.DotProductAttention is not None}")
+
+    failures = []
+    for module_name in ENTRYPOINT_MODULES:
+        try:
+            importlib.import_module(module_name)
+            print(f"import_ok: {module_name}")
+        except Exception as exc:  # pragma: no cover - smoke script
+            failures.append((module_name, repr(exc)))
+            print(f"import_failed: {module_name}: {exc!r}")
+
+    fast_path = te_compat.TE_AVAILABLE and te_compat.DotProductAttention is not None
+    print(f"te_fast_path_enabled={fast_path}")
+
+    if failures:
+        print("compat_smoke=WARN")
+        return 0
+
+    print("compat_smoke=PASS")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/install_blackwell_wsl.sh

@@ -36,7 +36,7 @@ PY
 
 # Transformer Engine 2.11.x + cuDNN 9.3+
 pip install "nvidia-cudnn-cu12>=9.3"
-pip install --no-build-isolation transformer-engine==2.11.0 transformer-engine-torch==2.11.0
+pip install --no-build-isolation "transformer-engine>=2.11,<3" "transformer-engine-torch>=2.11,<3"
 
 # Project dependencies.
 pip install -r requirements.txt