Skip to content

build(data_collection): make TORCH_CUDA_ARCH_LIST a --build-arg (unblocks Blackwell / sm_120) - #191

Open
InfraestruturaRobotec wants to merge 1 commit into
AgibotTech:mainfrom
InfraestruturaRobotec:chore/torch-cuda-arch-build-arg
Open

InfraestruturaRobotec wants to merge 1 commit into
AgibotTech:mainfrom
InfraestruturaRobotec:chore/torch-cuda-arch-build-arg

Conversation

@InfraestruturaRobotec

Copy link
Copy Markdown

What

Promote TORCH_CUDA_ARCH_LIST in source/data_collection/dockerfile from a
hard-coded ENV to an ARG with the same default, so it can be overridden at
build time.

ARG TORCH_CUDA_ARCH_LIST="8.9"
ENV TORCH_CUDA_ARCH_LIST=${TORCH_CUDA_ARCH_LIST}

No behaviour change for anyone who does not pass the flag — the default is
byte-for-byte what it was.

Why

The pin sits immediately before cuRobo is built, so the resulting image can only
run cuRobo on Ada. On Blackwell (RTX 50-series, sm_120) every kinematics call
fails at the first launch:

CUDA error: no kernel image is available for execution on the device
  curobo/curobolib/kinematics.py:73 -> kinematics_fused_cu.forward(...)

Confirmed with cuobjdump, the shipped extensions contain sm_89 only:

kinematics_fused_cu.cpython-311-x86_64-linux-gnu.so -> sm_89

With this change, a Blackwell user builds with:

docker build --build-arg TORCH_CUDA_ARCH_LIST="8.9;12.0+PTX" \
  -f source/data_collection/dockerfile source/data_collection

and all five cuRobo extensions carry sm_120.

Tested on

RTX 5090 (sm_120), driver 580.126.09, torch 2.7.0+cu128, Isaac Sim 5.1,
cuRobo v0.7.6. After rebuilding with the override, MotionGen.plan_single
succeeds and a full geniesim autocollect data-collection run completes with
zero CUDA errors.

Note for 50-series users beyond this PR

This change is necessary but not sufficient on sm_120. Two further things are
needed, and neither belongs in this PR:

  1. --no-deps when reinstalling cuRobo, otherwise pip re-resolves its
    dependency tree and replaces Isaac Sim's bundled torch 2.7.0+cu128 with a
    CUDA-13 stack, breaking the simulator.
  2. -maxrregcount=64 in cuRobo's own nvcc args — the L-BFGS kernel then hits
    too many resources requested for launch because sm_120 codegen exceeds
    the 64K registers-per-block limit at threadsPerBlock = v_dim. Filed
    upstream as L-BFGS kernel: 'too many resources requested for launch' on Blackwell (sm_120) — register pressure, not shared memory NVlabs/curobo#716.

Happy to add a short README note pointing 50-series users at those if that would
be useful.

The dockerfile pins TORCH_CUDA_ARCH_LIST=8.9 (Ada / RTX 4090D) immediately
before building cuRobo, so the resulting image cannot run cuRobo on any other
architecture without editing the dockerfile.

On Blackwell (RTX 50-series, sm_120) this surfaces as
  CUDA error: no kernel image is available for execution on the device
from curobo/curobolib/kinematics.py on the first kinematics call. Confirmed with
cuobjdump --list-elf: the shipped extensions contain sm_89 only.

Promoting the pin to an ARG keeps the existing default byte-for-byte and lets
users override it at build time:

  docker build --build-arg TORCH_CUDA_ARCH_LIST="8.9;12.0+PTX" ...

No behaviour change for anyone who does not pass the flag.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant