Fails to run due to transformer_engine.pytorch StopIteration error

[Fritingo]

Issue Description

When trying to run examples.text2image, the script fails with a StopIteration error related to transformer_engine.pytorch. It seems that the shared library for Transformer Engine is not found or loaded properly.

python -m examples.text2image \
    --prompt "${PROMPT_}" \
    --model_size 2B \
    --save_path output/text2image_2b.jpg

It seems transformer_engine cannot find the required .so file for the PyTorch extension. Even though transformer-engine[pytorch] was installed via uv pip, the shared object library is missing.

Traceback (most recent call last):
  File ".../transformer_engine/pytorch/__init__.py", line 52, in _load_library
    so_path = next(so_dir.glob(f"{module_name}.*.{extension}"))
StopIteration
...
  File ".../transformer_engine/pytorch/__init__.py", line 55, in _load_library
    so_path = next(so_dir.glob(f"{module_name}.*.{extension}"))
StopIteration

System Information

• OS: Ubuntu 22.04
• GPU: NVIDIA Blackwell 6000 Pro
• Driver version: 575
• Python version: 3.10 (in virtual environment .venv)
• Installation method: using uv as described in the setup guide:
  https://github.com/nvidia-cosmos/cosmos-predict2/blob/main/documentations/setup.md

[Fritingo]

The full error message

(cosmos-predict2) user@vm01:~/Documents/cosmos-predict2$ python -m examples.text2image     --prompt "${PROMPT_}"     --model_size 2B     --save_path output/text2image_2b.jpg
Traceback (most recent call last):
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/transformer_engine/pytorch/__init__.py", line 52, in _load_library
    so_path = next(so_dir.glob(f"{module_name}.*.{extension}"))
StopIteration

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/user/Documents/cosmos-predict2/examples/text2image.py", line 20, in <module>
    from imaginaire.auxiliary.text_encoder import CosmosTextEncoder
  File "/home/user/Documents/cosmos-predict2/imaginaire/auxiliary/text_encoder.py", line 35, in <module>
    from imaginaire.models.vlm_qwen import build_tokenizer
  File "/home/user/Documents/cosmos-predict2/imaginaire/models/vlm_qwen.py", line 39, in <module>
    from imaginaire.utils.checkpointer import _IncompatibleKeys
  File "/home/user/Documents/cosmos-predict2/imaginaire/utils/checkpointer.py", line 27, in <module>
    from imaginaire.utils import callback, distributed, log, misc
  File "/home/user/Documents/cosmos-predict2/imaginaire/utils/callback.py", line 29, in <module>
    from imaginaire.utils import distributed, log
  File "/home/user/Documents/cosmos-predict2/imaginaire/utils/distributed.py", line 45, in <module>
    from megatron.core import parallel_state
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/__init__.py", line 2, in <module>
    import megatron.core.tensor_parallel
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/tensor_parallel/__init__.py", line 2, in <module>
    from .cross_entropy import vocab_parallel_cross_entropy
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/tensor_parallel/cross_entropy.py", line 7, in <module>
    from megatron.core.parallel_state import (
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/parallel_state.py", line 14, in <module>
    from .utils import GlobalMemoryBuffer
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/utils.py", line 33, in <module>
    from megatron.core.dist_checkpointing.mapping import ShardedTensor
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/dist_checkpointing/__init__.py", line 5, in <module>
    from .serialization import (
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/dist_checkpointing/serialization.py", line 28, in <module>
    from .state_dict_transformation import load_preprocess, save_preprocess
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/dist_checkpointing/state_dict_transformation.py", line 12, in <module>
    from .exchange_utils import determine_main_replica_uniform_distribution, exchange_by_distribution
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/megatron/core/dist_checkpointing/exchange_utils.py", line 24, in <module>
    from transformer_engine.pytorch.float8_tensor import Float8Tensor
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/transformer_engine/pytorch/__init__.py", line 63, in <module>
    _load_library()
  File "/home/user/Documents/cosmos-predict2/.venv/lib/python3.10/site-packages/transformer_engine/pytorch/__init__.py", line 55, in _load_library
    so_path = next(so_dir.glob(f"{module_name}.*.{extension}"))
StopIteration