transformer-engine dependency fix

[xor2k]

Hi everybody,

I have a DGX Spark (aarch64) at home and tried to get cosmos-predict2.5 to run. I haven't tried so with x86, so I can't tell what's the situation there. However, three packages are not provided as aarch64 versions:

transformer-engine
natten
flash-attn

That's how I discovered this repo and I tried to build them myself.

When trying to build transformer-engine, the build fails because in the build environment of transformer-engine includes a non-cuda version of PyTorch and the line

uv pip install "torch==${TORCH_VERSION}.*" --index-url "https://download.pytorch.org/whl/cu${CUDA_NAME}"

in bin/_build.sh will not actually install anything (especially, torch stays non-cuda).

This is the build call (required for cosmos-predict2.5 as mentioned) that fails btw.

just build transformer-engine 2.8 3.12 2.9

The solution is to remove torch from the build dependencies of transformer-engine since it will be installed in a specified version via bin/_build.sh anyway. Since we want other package versions to stay locked in place until we provide an updated, working configuration, I've also removed torch from the uv.lock of transformer-engine. While usually, editing a uv.lock manually is not recommendable, in this particular case (removing a package) it stays identical if you run uv lock in packages/transformer-engine and it's as clean as it gets I guess 😅

Best, Michael