diff --git a/cosmos_predict1/autoregressive/utils/parallel.py b/cosmos_predict1/autoregressive/utils/parallel.py index 05f7733..f0302e0 100644 --- a/cosmos_predict1/autoregressive/utils/parallel.py +++ b/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 diff --git a/cosmos_predict1/diffusion/module/attention.py b/cosmos_predict1/diffusion/module/attention.py index 1feb6ea..19fad85 100644 --- a/cosmos_predict1/diffusion/module/attention.py +++ b/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") diff --git a/cosmos_predict1/utils/te_compat.py b/cosmos_predict1/utils/te_compat.py index 0c11e29..aba2ddd 100644 --- a/cosmos_predict1/utils/te_compat.py +++ b/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) diff --git a/requirements.txt b/requirements.txt index ecf45e4..183ff92 100644 --- a/requirements.txt +++ b/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 diff --git a/scripts/check_te_compat.py b/scripts/check_te_compat.py new file mode 100644 index 0000000..66649b8 --- /dev/null +++ b/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()) diff --git a/scripts/install_blackwell_wsl.sh b/scripts/install_blackwell_wsl.sh index cb341ac..55698c2 100755 --- a/scripts/install_blackwell_wsl.sh +++ b/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