diff --git a/pyproject.toml b/pyproject.toml index ff2ffd9a280..22fde66b520 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -69,7 +69,7 @@ automodel = [ "mamba-ssm", "causal-conv1d", "nv-grouped-gemm", - "transformer-engine[pytorch]>=2.9.0a0,<2.12.0", + "transformer-engine[pytorch,core_cu12] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1", "deep_ep @ git+https://github.com/deepseek-ai/DeepEP.git@bfded34800dfec415b71503f8205181de90b2480", ] vllm = [ @@ -99,7 +99,7 @@ mcore = [ # This dependency also needs to be compatible with the spec in Megatron-Bridge/pyproject.toml. # It is specified here since we don't directly use Megatron-Bridge/pyproject.toml, but a proxy setup.py+pyproject.toml combo # outside to allow "optionally" installing the megatron path. It's simpler to deal with transformer-engine here in the NeMo RL pyproject.toml - "transformer-engine[pytorch,core_cu12] @ git+https://github.com/NVIDIA/TransformerEngine.git@71bbefbf153418f943640df0f7373625dc93fa46", + "transformer-engine[pytorch,core_cu12] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1", "megatron-core", "megatron-bridge", # Flash-attn version should be selected to satisfy both TE + vLLM requirements (xformers in particular) @@ -242,7 +242,7 @@ link-mode = "copy" # The timm override is needed because current automodel pins to 1.0.16. This can be removed once we move ToT automodel # The nvidia-modelopt override is needed because mcore is still on 0.33 override-dependencies = [ - "transformer-engine[pytorch,core_cu12] @ git+https://github.com/NVIDIA/TransformerEngine.git@71bbefbf153418f943640df0f7373625dc93fa46", + "transformer-engine[pytorch,core_cu12] @ git+https://github.com/NVIDIA/TransformerEngine.git@v2.14.1", "nvidia-cudnn-cu12==9.19.0.56; sys_platform != 'darwin'", "opencv-python-headless>=4.11.0", "timm<=1.0.22", @@ -375,12 +375,12 @@ requires-dist = ["torch", "packaging", "ninja"] [[tool.uv.dependency-metadata]] name = "transformer-engine" -version = "2.14.0+71bbefb" +version = "2.14.1+366798e" requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"] [[tool.uv.dependency-metadata]] name = "transformer-engine-torch" -version = "2.14.0+71bbefb" +version = "2.14.1+366798e" requires-dist = ["torch", "transformer-engine"] [[tool.uv.dependency-metadata]] diff --git a/uv.lock b/uv.lock index 59723314766..f72b112a598 100644 --- a/uv.lock +++ b/uv.lock @@ -110,7 +110,7 @@ overrides = [ { name = "torch", marker = "sys_platform == 'darwin'", specifier = "==2.10.0", index = "https://pypi.org/simple" }, { name = "torchaudio", marker = "sys_platform != 'darwin'", specifier = "==2.10.0", index = "https://download.pytorch.org/whl/cu129" }, { name = "torchaudio", marker = "sys_platform == 'darwin'", specifier = "==2.10.0", index = "https://pypi.org/simple" }, - { name = "transformer-engine", extras = ["pytorch", "core-cu12"], git = "https://github.com/NVIDIA/TransformerEngine.git?rev=71bbefbf153418f943640df0f7373625dc93fa46" }, + { name = "transformer-engine", extras = ["pytorch", "core-cu12"], git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" }, { name = "transformers", specifier = "==5.3.0" }, { name = "xgrammar", specifier = "==0.1.33" }, ] @@ -171,12 +171,12 @@ requires-dist = ["torch", "scikit-build-core", "wheel"] [[manifest.dependency-metadata]] name = "transformer-engine" -version = "2.14.0+71bbefb" +version = "2.14.1+366798e" requires-dist = ["torch", "pydantic", "importlib-metadata>=1.0", "packaging"] [[manifest.dependency-metadata]] name = "transformer-engine-torch" -version = "2.14.0+71bbefb" +version = "2.14.1+366798e" requires-dist = ["torch", "transformer-engine"] [[package]] @@ -4398,8 +4398,8 @@ requires-dist = [ { name = "torchdata" }, { name = "torchvision", marker = "sys_platform != 'darwin'", specifier = "==0.25.0", index = "https://download.pytorch.org/whl/cu129" }, { name = "torchvision", marker = "sys_platform == 'darwin'", specifier = "==0.25.0", index = "https://pypi.org/simple" }, - { name = "transformer-engine", extras = ["core-cu12", "pytorch"], marker = "extra == 'mcore'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=71bbefbf153418f943640df0f7373625dc93fa46" }, - { name = "transformer-engine", extras = ["pytorch"], marker = "extra == 'automodel'", specifier = ">=2.9.0a0,<2.12.0" }, + { name = "transformer-engine", extras = ["core-cu12", "pytorch"], marker = "extra == 'automodel'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" }, + { name = "transformer-engine", extras = ["core-cu12", "pytorch"], marker = "extra == 'mcore'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1" }, { name = "transformers", specifier = "==5.3.0" }, { name = "transformers", marker = "extra == 'automodel'", specifier = ">=5.3.0" }, { name = "triton", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')", index = "https://download.pytorch.org/whl/cu129" }, @@ -8136,8 +8136,8 @@ wheels = [ [[package]] name = "transformer-engine" -version = "2.14.0+71bbefb" -source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=71bbefbf153418f943640df0f7373625dc93fa46#71bbefbf153418f943640df0f7373625dc93fa46" } +version = "2.14.1+366798e" +source = { git = "https://github.com/NVIDIA/TransformerEngine.git?rev=v2.14.1#366798ef8a0a00d8f2c1650d11e7e623d7c33e26" } dependencies = [ { name = "importlib-metadata" }, { name = "packaging" },