Pin transformer-engine[pytorch]==2.5.0

[nhayesroth]

Summary

Pins our transformer-engine dependency to 2.5.0.

Version 2.6.0 was released yesterday and appears to include a dependency mismatch.

Testing

Built and pushed a new image version: nvcr.io/ebplj5vljyh4/dev/cosmos-transfer1:ci-32978582

Current Dockerfile (unpinned dep => 2.6.0) => OSError

root@before:/workspace# nvidia-smi | grep "CUDA Version"
| NVIDIA-SMI 570.133.20             Driver Version: 570.133.20     CUDA Version: 12.9     |

root@before:/workspace# ls -la /usr/local/cuda*/lib64/libcublas.so* 
lrwxrwxrwx 1 root root        15 Apr  8 04:16 /usr/local/cuda-12.9/lib64/libcublas.so -> libcublas.so.12
lrwxrwxrwx 1 root root        21 Apr  8 04:16 /usr/local/cuda-12.9/lib64/libcublas.so.12 -> libcublas.so.12.9.0.2
-rw-r--r-- 1 root root 105136880 Apr  8 04:16 /usr/local/cuda-12.9/lib64/libcublas.so.12.9.0.2
lrwxrwxrwx 1 root root        15 Apr  8 04:16 /usr/local/cuda-12/lib64/libcublas.so -> libcublas.so.12
lrwxrwxrwx 1 root root        21 Apr  8 04:16 /usr/local/cuda-12/lib64/libcublas.so.12 -> libcublas.so.12.9.0.2
-rw-r--r-- 1 root root 105136880 Apr  8 04:16 /usr/local/cuda-12/lib64/libcublas.so.12.9.0.2
lrwxrwxrwx 1 root root        15 Apr  8 04:16 /usr/local/cuda/lib64/libcublas.so -> libcublas.so.12
lrwxrwxrwx 1 root root        21 Apr  8 04:16 /usr/local/cuda/lib64/libcublas.so.12 -> libcublas.so.12.9.0.2
-rw-r--r-- 1 root root 105136880 Apr  8 04:16 /usr/local/cuda/lib64/libcublas.so.12.9.0.2

root@before:/workspace# pip list | grep -i transformer
transformer_engine                2.6.0
transformer_engine_cu12           2.6.0
transformer_engine_torch          2.6.0
transformers                      4.49.0

root@before:/workspace# python3 -c "import transformer_engine" && echo 'it works'
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/__init__.py", line 11, in <module>
    import transformer_engine.common
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/common/__init__.py", line 308, in <module>
    _TE_LIB_CTYPES = _load_core_library()
                     ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/transformer_engine/common/__init__.py", line 300, in _load_core_library
    return ctypes.CDLL(_get_shared_object_file("core"), mode=ctypes.RTLD_GLOBAL)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/lib/python3.12/ctypes/__init__.py", line 379, in __init__
    self._handle = _dlopen(self._name, mode)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^
OSError: libcublas.so.13: cannot open shared object file: No such file or directory

After pinning to 2.5.0 => no error

root@pinned-2.5:/mnt/pvc# nvidia-smi | grep "CUDA Version"
| NVIDIA-SMI 570.133.20             Driver Version: 570.133.20     CUDA Version: 12.9     |

root@pinned-2.5:/mnt/pvc# ls -la /usr/local/cuda*/lib64/libcublas.so*
lrwxrwxrwx 1 root root        15 Apr  8 04:16 /usr/local/cuda-12.9/lib64/libcublas.so -> libcublas.so.12
lrwxrwxrwx 1 root root        21 Apr  8 04:16 /usr/local/cuda-12.9/lib64/libcublas.so.12 -> libcublas.so.12.9.0.2
-rw-r--r-- 1 root root 105136880 Apr  8 04:16 /usr/local/cuda-12.9/lib64/libcublas.so.12.9.0.2
lrwxrwxrwx 1 root root        15 Apr  8 04:16 /usr/local/cuda-12/lib64/libcublas.so -> libcublas.so.12
lrwxrwxrwx 1 root root        21 Apr  8 04:16 /usr/local/cuda-12/lib64/libcublas.so.12 -> libcublas.so.12.9.0.2
-rw-r--r-- 1 root root 105136880 Apr  8 04:16 /usr/local/cuda-12/lib64/libcublas.so.12.9.0.2
lrwxrwxrwx 1 root root        15 Apr  8 04:16 /usr/local/cuda/lib64/libcublas.so -> libcublas.so.12
lrwxrwxrwx 1 root root        21 Apr  8 04:16 /usr/local/cuda/lib64/libcublas.so.12 -> libcublas.so.12.9.0.2
-rw-r--r-- 1 root root 105136880 Apr  8 04:16 /usr/local/cuda/lib64/libcublas.so.12.9.0.2

root@pinned-2.5:/mnt/pvc# pip list | grep -i transformer
transformer_engine                2.5.0
transformer_engine_cu12           2.5.0
transformer_engine_torch          2.5.0
transformers                      4.49.0

root@pinned-2.5:/mnt/pvc# python3 -c "import transformer_engine" && echo 'it works'
it works

[grossanchez]

There is now 2.6.0-post1 available that seems to work fine.