NVIDIA Cosmos (Predict2.5) on DGX Spark (ARM / aarch64)

[amirtd]

I'm not able to run NVIDIA Cosmos (Predict2.5 / text2world) on DGX Spark (ARM / aarch64)?
I’m hitting dependency blockers around Transformer Engine / flash-attn and some configs that seem x86_64-centric. Curious if anyone has a confirmed working setup, workarounds, or if Cosmos is currently x86-only in practice.

Environment

• Hardware: NVIDIA DGX Spark
• Architecture: aarch64 (ARM)
• OS: Ubuntu (Spark image)
• Kernel: 6.14.0-1015-nvidia
• GPU: NVIDIA GB10
• Driver: 580.126.09
• CUDA: 13.0
• Python: 3.10 (Conda)
• Cosmos Version: cosmos-predict2.5

Key Issues Encountered

1. Transformer Engine Is a Hard Dependency (and Not Available on ARM)
Multiple Cosmos modules unconditionally import transformer_engine, including:

• minimal_v4_dit.py
• wan2pt1.py
• fused_adam_dtensor.py
• optimizer and attention layers
  Example failure:

ModuleNotFoundError: No module named 'transformer_engine'

Even after:

• Rewriting attention layers
• Introducing fallback implementations
• Removing explicit class inheritance from DotProductAttention

…the import graph still eventually requires transformer_engine.

Root problem:
transformer-engine wheels are not available for aarch64, and building from source is not supported on DGX Spark.

2. flash-attn / flash-attn-2 Is x86_64-Centric
Cosmos enforces availability checks such as:

assert is_flash_attn_2_available()

On DGX Spark:

• flash-attn is not available
• flash-attn-2 cannot be installed
• CUDA + ARM + GB10 has no supported build
  This causes early termination even when the feature is not strictly required for inference.

3. Config System Assumes Missing Experiment Modules
Text2world config attempts to auto-import:

cosmos_predict2._src.predict2.configs.text2world.experiment
But the directory does not exist in the repo, leading to:
ModuleNotFoundError:
No module named 'cosmos_predict2._src.predict2.configs.text2world.experiment'

This suggests:

• Incomplete packaging
• Internal-only config assumptions
• Or missing repo components

4. Inference Path Is Hardwired to Video2World
Even with:

• --inference-type=text2world
• Environment overrides
• Patched model_loader.py

The system still instantiates:

Video2WorldInference(...)

This drags in:

• Vision-language models
• Qwen-VL
• Transformer Engine
• Flash attention paths

There is no clean text-only inference path that avoids these dependencies.

5. Cascading Failures from Partial Patching

Attempts to patch around missing dependencies led to secondary failures:

• NoneType takes no arguments (class inheritance on disabled components)
• NameError: _SACConfig not defined
• SyntaxError: return outside function (after fallback edits)

This indicates the codebase is not designed for partial feature removal.

What Would Be Needed to Support DGX Spark

1. Optional Transformer Engine

• Clean fallbacks
• No hard imports
• No class inheritance from TE types

2. flash-attn Optionality

• Runtime feature flags
• Safe PyTorch attention fallback

3. True Text-Only Inference Path

• No Video2World initialization
• No vision or VLM imports

4. ARM-Compatible Build Artifacts

• TE or equivalent kernels for aarch64
• Or official statement of non-support

[robmurrer]

did you use uv sync --extra=cu130 ?

[amirtd]

did you use uv sync --extra=cu130 ?

uv sync --extra=cu130 won’t fix this on DGX Spark (ARM).
--extra=cu130 only selects CUDA 13 wheels but it does not change CPU architecture. I did try CUDA 13 (cu130), but the issue isn’t the CUDA version. On DGX Spark the blocker is ARM/aarch64: Cosmos hard-depends on transformer-engine and flash-attn, which don’t have ARM wheels. uv sync --extra=cu130 still resolves x86-only packages, so runtime imports fail.

[amirtd]

Update on Cosmos + DGX Spark:

• video2world runs end-to-end on DGX Spark (ARM) and successfully generates MP4 output using the Cosmos-Predict 2.5 video models.
• text2world, however, currently fails on Spark due to hard dependencies on Transformer Engine / flash-attn, which don’t have ARM builds yet.

Same Cosmos version, different dependency paths!

[sbhavani]

Related context from Transformer Engine: support on DGX Spark (SM120) is still being actively brought up (see NVIDIA/TransformerEngine#2668)

[jcl2023]

@sbhavani

Any update on this?

Thanks,

[xor2k]

I got cosmos-predict2.5 working on my DGX Spark. There are three packages that don't provide aarch64 versions

transformer-engine
natten
flash-attn

As I have discovered, there is a cosmos-dependencies repository that can build those packages. However, the build environment of transformer-engine has a bug that needs to be fixed first, I've handed in a pull request: nvidia-cosmos/cosmos-dependencies#55

So as long as no aarch64 versions for the packages mentioned are provided, one needs to build the required .whl files oneself. Ensure to

• Have cosmos-dependencies and cosmos-predict2.5 in the same directory/on the same directory level
• Apply the changes from the PR mentioned above
• Use Python-3.12 (as is default in DGX OS right now)

Then, in cosmos-dependencies, with docker (details see README.md there) or possibly without, run (e.g. put it into build.sh and run sh build.sh)

export CUDA_VERSION=130
TORCH_VERSION=2.9
PYTHON_VERSION=3.12

just build transformer-engine 2.8 ${PYTHON_VERSION} ${TORCH_VERSION} &&
just build natten 0.21.0 ${PYTHON_VERSION} ${TORCH_VERSION} &&
just build flash-attn 2.7.4.post1 ${PYTHON_VERSION} ${TORCH_VERSION} &&
mkdir -p wheels &&
for file in $(find build -name "*.whl"); do
    cp -a $file wheels
done

This will put mentioned three packages in

cosmos-dependencies/wheels

Next, we need to patch cosmos-predict2.5 to use those self-built packages with git apply patch.diff and patch.diff containing

diff --git a/packages/cosmos-oss/pyproject.toml b/packages/cosmos-oss/pyproject.toml
index 352dc34..0b25882 100644
--- a/packages/cosmos-oss/pyproject.toml
+++ b/packages/cosmos-oss/pyproject.toml
@@ -130,7 +130,7 @@ cu130_torch29 = [
   "natten==0.21.0+cu130.torch29",
   "transformer-engine==2.8+cu130.torch29",
   # Torch
-  "torch==2.9.1",
+  "torch==2.9.1+cu130",
   # Dependencies constrained by torch
   "torchvision",
   "triton",
@@ -176,17 +176,17 @@ decord = [
 flash-attn = [
   { index = "cosmos-cu128-torch27", extra = "cu128_torch27" },
   { index = "cosmos", extra = "cu130_torch29", marker = "python_version >= '3.13'" },
-  { index = "cosmos-cu130-torch29", extra = "cu130_torch29", marker = "python_version < '3.13'" },
+  { index = "cosmos-cu130-torch29-local", extra = "cu130_torch29", marker = "python_version < '3.13'" },
 ]
 natten = [
   { index = "cosmos-cu128-torch27", extra = "cu128_torch27" },
   { index = "cosmos", extra = "cu130_torch29", marker = "python_version >= '3.13'" },
-  { index = "cosmos-cu130-torch29", extra = "cu130_torch29", marker = "python_version < '3.13'" },
+  { index = "cosmos-cu130-torch29-local", extra = "cu130_torch29", marker = "python_version < '3.13'" },
 ]
 transformer-engine = [
   { index = "cosmos-cu128-torch27", extra = "cu128_torch27" },
   { index = "cosmos", extra = "cu130_torch29", marker = "python_version >= '3.13'" },
-  { index = "cosmos-cu130-torch29", extra = "cu130_torch29", marker = "python_version < '3.13'" },
+  { index = "cosmos-cu130-torch29-local", extra = "cu130_torch29", marker = "python_version < '3.13'" },
 ]
 torch = [
   { index = "cosmos-cu128-torch27", extra = "cu128_torch27" },
@@ -210,6 +210,12 @@ name = "cosmos-cu128-torch27"
 url = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.2.0/cu128_torch27/simple"
 explicit = true
 
+[[tool.uv.index]]
+name = "cosmos-cu130-torch29-local"
+url = "../cosmos-dependencies/wheels"
+format = "flat"
+explicit = true
+
 [[tool.uv.index]]
 name = "cosmos-cu130-torch29"
 url = "https://nvidia-cosmos.github.io/cosmos-dependencies/v1.2.0/cu130_torch29/simple"

before running uv sync --extra=cu130 in cosmos-predict2.5. For testing the setup, run this code (e.g. via putting it into a run.sh and running sh run.sh):

. .venv/bin/activate &&
python examples/inference.py \
  -i assets/base/robot_pouring.json \
  -o outputs/base_video2world \
  --inference-type=video2world

Do not use uv run ... (or if you do, then uv run --no-sync ...), as (by default) it will trigger a uv sync without --extra=cu130, which will lead to a non-cuda version of torch to be installed.