[WS1][CI] GPU test installs incompatible Torch 2.4.1 and Transformers 5.x

[hihaluemen]

Summary

The GPU CI pins PyTorch to torch==2.4.1, but later installs the latest unconstrained transformers release.

In the affected run, pip installed transformers==5.14.1. This version attempts to import DTensor from torch.distributed.tensor, which is incompatible with the pinned PyTorch 2.4.1 environment.

As a result, the GPU test suite fails during pytest collection before the operator tests are executed.

Affected run

• Workflow run: https://github.com/RL-Align/RL-Kernel/actions/runs/29518774291
• H100 job: https://github.com/RL-Align/RL-Kernel/actions/runs/29518774291/job/87690353599
• Observed while testing PR: [WS1][kernels] Batch-invariant logprob (Native, Triton) #199
• Tested commit: fb22b981380fbdfb186e76b3c130a368536cb7b8

Observed environment

The job installed:

torch 2.4.1+cu124
transformers 5.14.1

The relevant CI commands are in ci/run_gpu_ci.sh:

TORCH_SPEC="${TORCH_SPEC:-torch==2.4.1}"
"$PY" -m pip install --no-cache-dir "$TORCH_SPEC" \
    --index-url "$TORCH_INDEX_URL"

"$PY" -m pip install --no-build-isolation --no-deps -e .
"$PY" -m pip install --no-cache-dir \
    numpy tabulate accelerate transformers pytest

The editable project installation uses --no-deps, and the later manual dependency installation does not constrain the Transformers version.

Error

Pytest fails while collecting tests/test_stateless_hf_integration.py:

transformers/distributed/sharding_utils.py:28: in <module>
    from torch.distributed.tensor import DTensor

ImportError: cannot import name 'DTensor'
from 'torch.distributed.tensor'

The traceback is printed twice because the H100 job runs pytest with two distributed workers. Both workers report the same collection error.

Evidence that the CUDA extension is working

Before pytest collection, the GPU smoke test succeeds:

[smoke] torch: 2.4.1+cu124 (cuda 12.4)
[smoke] capability: sm_90
[smoke] _C file: /workspace/repo/rl_engine/_C.cpython-311-x86_64-linux-gnu.so
[smoke] OK: rl_engine._C built and fused_logp ran on sm_90.

Pytest then reports:

collected 705 items / 1 error

Therefore, this is a dependency-resolution failure rather than a CUDA extension or kernel failure.

Root cause

ci/run_gpu_ci.sh fixes the Torch version to ensure a deterministic extension ABI, but installs an unconstrained Transformers version.

The project metadata also currently allows this combination:

"torch>=2.4.1",
"transformers",

This lets pip install a dependency set that succeeds during installation but fails when Hugging Face model classes are imported.

The relevant CI installation logic was introduced by commit:

c6b33a2 CI: pin torch and make the extension build order deterministic

Pinning Torch is useful and should likely remain. The missing part is a matching Transformers compatibility constraint.

Expected behavior

GPU CI should install a Torch/Transformers combination supported by the repository, and tests/test_stateless_hf_integration.py should be collected and executed successfully.

Possible solutions

One lower-risk option is to keep the existing Torch 2.4.1/CUDA 12.4 toolchain and constrain Transformers to a validated compatible version range.

The constraint should be kept consistent in:

• ci/run_gpu_ci.sh
• pyproject.toml
• setup.py

For example, transformers<5 could be evaluated, but the exact supported range should be verified against Torch 2.4.1 and the Hugging Face integration tests.

The alternative is upgrading Torch and the CI image, but that would require broader validation of the CUDA extension ABI, Triton, CUDA wheels, and SM86/SM90 kernels.

CI fix validation note

GPU CI uses pull_request_target and obtains the CI orchestrator from the PR's base SHA. Therefore, changes to ci/run_gpu_ci.sh in a PR will not automatically be used by that PR's own GPU CI run.

The change may need maintainer-controlled manual validation. After the fix is merged into main, a new workflow run should be triggered to confirm that the issue is resolved.