Add Ascend NPU Support for Trinity-RFT

[huajinguo12]

Add Native Ascend NPU Support to Trinity-RFT

Summary

This PR adds native Huawei Ascend NPU (910B and above) support to Trinity-RFT as a first-class backend, alongside the existing CUDA backend. With this change, trinity run --config <yaml> launches end-to-end (Ray + vLLM rollout + verl FSDP trainer) on an NPU-only cluster without forking the codebase.

All changes are fully backward-compatible — GPU users see no behavior change. The NPU path is gated behind a device abstraction layer (trinity/utils/device.py) that auto-detects the runtime device and returns the correct Ray resource key, collective backend, and torch device name.

Validated on Trinity-RFT 0.5.2 with the software stack listed below.

Validated Software Stack

All components below were installed in a single conda environment (trinity) and verified end-to-end on Ascend 910B3.

Component	Version	Notes
Trinity-RFT	v0.5.2	Base version under test
CANN Toolkit	8.5.1	/home/cann/cann851/ascend-toolkit/set_env.sh
torch	2.9.0+cpu	CPU-only wheel; NPU kernels provided by torch_npu
torch_npu	2.9.0	Ascend PyTorch backend extension
vLLM	0.18.0	Rollout engine
vllm-ascend	0.18.0rc1	Ascend plugin for vLLM
verl	0.7.1	FSDP training backend
Hardware	Ascend 910B3 (64 GB HBM) × 8	1-node, 8 NPUs

NPU is detected automatically when torch_npu is importable. To force a device explicitly, set TRINITY_DEVICE=npu|cuda|cpu.

NPU Test Results

Full test suite executed on Ascend 910B3 (CANN 8.5.1, 8 NPUs, conda env trinity). Summary: 253 passed / 8 failed / 6 skipped out of 267 cases. All failures are environment/version issues, not code bugs introduced by this PR.

Test file	Total	Pass	Skip	Failed cases & brief reason
tests/algorithm/advantage_fn_test.py	8	8	0	—
tests/algorithm/kl_fn_test.py	11	11	0	—
tests/algorithm/policy_loss_test.py	8	8	0	—
tests/buffer/experience_pipeline_test.py	2	2	0	—
tests/buffer/experience_storage_test.py	3	3	0	—
tests/buffer/file_test.py	2	2	0	—
tests/buffer/formatter_test.py	6	6	0	—
tests/buffer/queue_test.py	6	6	0	—
tests/buffer/reader_test.py	1	1	0	—
tests/buffer/reward_shaping_mapper_test.py	1	1	0	—
tests/buffer/sample_strategy_test.py	10	10	0	—
tests/buffer/sql_test.py	3	3	0	—
tests/buffer/task_scheduler_test.py	8	8	0	—
tests/buffer/task_storage_test.py	6	6	0	—
tests/cli/launcher_test.py	6	6	0	—
tests/common/config_test.py	9	9	0	—
tests/common/experience_test.py	14	14	0	—
tests/common/sudoku_test.py	11	11	0	—
tests/common/models/utils_test.py	4	4	0	—
tests/common/external_model_test.py	2	2	0	—
tests/common/vllm_test.py	19	10	1	TestModelLen_0/1/2::test_model_len (3), TestModelLenWithoutPromptTruncation::test_model_len (1): max_prompt_tokens/max_response_tokens deprecated in vLLM > 0.16.0, no longer passed to SamplingParams; OpenAI API path doesn't honor these params either. TestLogprobs::test_logprobs_api (1): NPU precision divergence in prompt_logprobs when temperature ≠ 1.0 (rtol=0.5 still fails). TestAPIServerToolCall_0_deepseek_r1::test_api_tool_calls (1): Qwen2.5-0.5B lacks <think> tokens; Qwen3.5 uses qwen3_5 arch unsupported by transformers 4.57.6. TestTinkerAPI::test_tinker_api (1): tinker proprietary package not installed. TestTinkerAsyncAPIServer::test_api_async skipped (no TINKER_API_KEY).
tests/explorer/explorer_test.py	5	5	0	—
tests/explorer/proxy_test.py	1	1	0	—
tests/explorer/scheduler_test.py	18	18	0	—
tests/explorer/step_wise_workflow_test.py	6	6	0	—
tests/explorer/workflow_test.py	19	17	2	—
tests/manager/log_manager_test.py	3	3	0	—
tests/manager/synchronizer_test.py	16	16	0	—
tests/service/data_juicer_test.py	4	4	0	—
tests/utils/eval_utils_test.py	12	12	0	—
tests/utils/log_test.py	3	3	0	—
tests/utils/plugin_test.py	6	6	0	—
tests/utils/registry_test.py	6	6	0	—
tests/utils/swanlab_test.py	1	0	1	—
tests/trainer/trainer_test.py	27	25	2	—
Total	267	253	6	8 failed — all environment/version/backend-scope issues, none caused by this PR**

Failure Category Breakdown

Category	Count	Root cause
vLLM API changes	4	max_prompt_tokens / max_response_tokens deprecated in vLLM > 0.16.0.
NPU precision divergence	1	prompt_logprobs at temperature ≠ 1.0 (rtol=0.5 still fails).
Missing model / deps	3	Qwen2.5-0.5B lacks <think> tokens; Qwen3.5 arch unsupported by transformers 4.57.6; tinker package not installed (includes 1 undocumented vLLM failure).

Total: 8 unfixed — all environment/version/backend-scope issues, none caused by this PR.

Environment Setup

Ascend NPU users set (CUDA users do nothing):

export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7   # equivalent to CUDA_VISIBLE_DEVICES
source /home/cann/cann851/ascend-toolkit/set_env.sh  # CANN 8.5.1
source /home/cann/cann851/nnal/atb/set_env.sh        # libatb.so for vLLM

NPU users disable torch.compile in YAML:

trainer:
  use_torch_compile: false  # inductor has no NPU driver

Checklist

Please check the following items before code is ready to be reviewed.

• Code has passed all tests
• Docstrings have been added/updated in Google Style
• Documentation has been updated
• Code is ready for review

[pan-x-c]

/unittest-all

[huajinguo12]

/unittest-all

[pan-x-c]

/unittest-all

[github-actions]

unittest: Run #1840

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Pending ⏳	Other ❓	Flaky 🍂	Duration ⏱️
0	0	0	0	0	0	0	14.4s

🎉 All tests passed!

Github Test Reporter by CTRF 💚

[pan-x-c]

/unittest-all

[pan-x-c]

The current implementation does not have a get_device_name method. Please check the existing implementation and fix the pre-commit issues. For the specific error details, see: https://github.com/agentscope-ai/Trinity-RFT/actions/runs/30065533393/job/89395620057

[huajinguo12]

The current implementation does not have a get_device_name method. Please check the existing implementation and fix the pre-commit issues. For the specific error details, see: https://github.com/agentscope-ai/Trinity-RFT/actions/runs/30065533393/job/89395620057

Thank you for your review. We have noted the failed unit test cases in the CI workflow.

We are currently investigating the issues, within our internal GPU environment. We also set up external access to NPU resources. If you require access to an NPU environment for further verification or collaboration, please let us know your preferred instant messaging account (e.g., DingTalk) so we can coordinate with you directly.

[huajinguo12]

The current implementation does not have a get_device_name method. Please check the existing implementation and fix the pre-commit issues. For the specific error details, see: https://github.com/agentscope-ai/Trinity-RFT/actions/runs/30065533393/job/89395620057

I have pushed the updated code and please re-trigger the CI pipeline for this PR.

[pan-x-c]

/unittest-all

[huajinguo12]

/unittest-all

A total of 11 test cases failed in the CI gate, but when I ran the same failing tests locally on a GPU environment based on this PR using the following commands:

python -m pytest tests/manager/synchronizer_test.py::TestSynchronizerExit_0::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestSynchronizerExit_1::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestStateDictBasedSynchronizer_0::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestStateDictBasedSynchronizer_1::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestStateDictBasedSynchronizer_2::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestStateDictBasedSynchronizer_3::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestStateDictBasedSynchronizer_4::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestStateDictBasedSynchronizer_5::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestNCCLBasedSynchronizer_0::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestNCCLBasedSynchronizer_1::test_synchronizer -v -s
python -m pytest tests/manager/synchronizer_test.py::TestNCCLBasedSynchronizer_2::test_synchronizer -v -s
python -m pytest tests/trainer/trainer_test.py::TestTrainerCountdown_0_fsdp2::test_trainer -v -s

only 3 test cases failed. Two of them are from synchronizer_test.py and failed due to timeout while waiting for trainer initialization. The other one is tests/trainer/trainer_test.py::TestTrainerCountdown_0_fsdp2::test_trainer, which failed due to an OOM (Out of Memory) error. (which appear to be unrelated to this PR).

I did not see the same errors as in the CI, such as:

(autoscaler +1h33m6s) Error: No available node types can fulfill resource request defaultdict(<class 'float'>, {'cpu': 2.0, 'CPU': 6.0}). Add suitable node types to this cluster to resolve this issue.
(autoscaler +1h33m41s) Error: No available node types can fulfill resource request defaultdict(<class 'float'>, {'cpu': 2.0, 'CPU': 6.0}). Add suitable node types to this cluster to resolve this issue.

It seems that the CI execution environment needs to be checked.

[pan-x-c]

Please refer to docker/README.md to configure the required test environment.

[huajinguo12]

We have set up the test environment following docker/README.md. The tests passed locally without the issues that appeared on CI. Logs are attached below for reference.

log.txt

We would like to verify whether there might be any issues with the CI environment itself.

For easier communication, feel free to reach me directly via phone: 15161904875

[pan-x-c]

/unittest-all

[huajinguo12]

/unittest-all
2026-08-07 04:48:29,926 INFO worker.py:1672 -- Using address auto set in the environment variable RAY_ADDRESS 2026-08-07 04:48:29,928 INFO worker.py:1814 -- Connecting to existing Ray cluster at address: 172.22.0.2:6379... 2026-08-07 04:48:29,932 INFO worker.py:2003 -- Connected to Ray cluster. View the dashboard at http://172.22.0.2:8265/ 2026-08-07 04:48:30,517 INFO worker.py:1672 -- Using address auto set in the environment variable RAY_ADDRESS 2026-08-07 04:48:30,519 INFO worker.py:1814 -- Connecting to existing Ray cluster at address: 172.22.0.2:6379... 2026-08-07 04:48:30,519 INFO worker.py:1828 -- Calling ray.init() again after it has already been called. tests/trainer/trainer_test.py::TestTrainerCountdown_0_fsdp2::test_trainer [START] 04:48:30 - Running: tests/trainer/trainer_test.py::TestTrainerCountdown_0_fsdp2::test_trainer WARNING 08-07 04:48:30 [config_validator.py:613] max_prompt_tokensis set to 3072,max_response_tokensis set to 1024. WARNING 08-07 04:48:30 [config_validator.py:638]enable_prompt_truncationis set to True; the prompt will be truncated tomax_prompt_tokens=3072 tokens if it is too long. INFO 08-07 04:48:30 [config_validator.py:976] buffer.train_batch_sizeis set tobuffer.batch_size*algorithm.repeat_timesINFO 08-07 04:48:30 [config_validator.py:1015]buffer.explorer_input.taskset.repeat_timesis set toalgorithm.repeat_times(=3). INFO 08-07 04:48:30 [config_validator.py:1114] Auto setbuffer.trainer_input.experience_bufferto ExperienceBufferConfig(name='experience_buffer', storage_type='queue', path=None, capacity=10000, max_read_timeout=1800, replay_buffer=ReplayBufferConfig(enable=False, priority_fn='linear_decay', reuse_cooldown_time=None, priority_fn_args={'decay': 2.0}), max_retry_times=3, max_retry_interval=1, split='train', subset_name=None, format=FormatConfig(prompt_type=<PromptType.MESSAGES: 'messages'>, prompt_key='prompt', response_key='response', system_prompt_key=None, system_prompt=None, messages_key='message', tools_key='tools', image_key=None, video_key=None, reply_prefix=None, workflow_key='', reward_fn_key='', chosen_key='chosen', rejected_key='rejected', enable_concatenated_multi_turn=False, chat_template=None, enable_thinking=None), enable_progress_bar=False, data_selector=DataSelectorConfig(selector_type='sequential', seed=42, feature_keys=[], kwargs={}), schema_type=None, index=0, batch_size=0, tokenizer_path=None, total_epochs=1, total_steps=None, ray_namespace=None) WARNING 08-07 04:48:30 [config_validator.py:1134] Auto setbuffer.trainer_input.experience_buffer.pathto /mnt/checkpoints-unittest/Trainer-unittest/trainer-20260807044830/buffer/experience_buffer.jsonl (pid=100797) W0807 04:48:35.814000 100797 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' (Trainer pid=29212, ip=172.22.0.3) Using blocking ray.get inside async actor. This blocks the event loop. Please useawait on object ref with asyncio.gather if you want to yield execution to the event loop instead. (Synchronizer pid=29270, ip=172.22.0.3) W0807 04:48:41.189000 29270 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' [repeated 2x across cluster] (Trainer pid=29212, ip=172.22.0.3) INFO 08-07 04:48:37 [trainer.py:227] Initializing verl Trainer with fsdp2 backend (Trainer pid=29212, ip=172.22.0.3) WARNING 08-07 04:48:37 [config.py:769] Stripping trainer_config override 'actor.megatron' — incompatible with strategy 'fsdp2' (Trainer pid=29212, ip=172.22.0.3) WARNING 08-07 04:48:37 [config.py:769] Stripping trainer_config override 'ref.megatron' — incompatible with strategy 'fsdp2' (Trainer pid=29212, ip=172.22.0.3) WARNING 08-07 04:48:37 [config.py:769] Stripping trainer_config override 'critic.megatron' — incompatible with strategy 'fsdp2' (Explorer pid=100797) INFO 08-07 04:48:37 [explorer.py:97] Finished initializing Explorer. (Trainer pid=29212, ip=172.22.0.3) INFO 08-07 04:48:39 [trainer.py:404] worker group kwargs: {'device_name': 'cpu'} (Trainer pid=29212, ip=172.22.0.3) colocated worker base class <class 'verl.single_controller.base.worker.Worker'> (autoscaler +1h32m47s) Error: No available node types can fulfill resource request defaultdict(<class 'float'>, {'CPU': 6.0, 'cpu': 2.0}). Add suitable node types to this cluster to resolve this issue. (autoscaler +1h33m22s) Error: No available node types can fulfill resource request defaultdict(<class 'float'>, {'CPU': 6.0, 'cpu': 2.0}). Add suitable node types to this cluster to resolve this issue.
It belongs to CI cluster resource/configuration issues. The cluster has been running for an extended period: +1h32m47s indicates the autoscaler has been running for over 1.5 hours. Preceding test cases may not have fully released placement groups / actors, leading to progressive resource exhaustion. ray.shutdown(_exiting_interpreter=True) at trainer_test.py is only called at the end of the test, while RayUnittestBase.tearDownClass uses ray.shutdown , and setUp calls ray.init(ignore_reinit_error=True) again — if a preceding test case exits abnormally without cleanup, resources will remain leaked.

[pan-x-c]

/unittest-module-trainer

[pan-x-c]

Please carefully review the current implementation. The currently failing test is the first trainer test (TestTrainerCountdown), and the failure is caused by a timeout due to insufficient resources. This suggests that there may be a bug in the current Ray cluster resource allocation logic.

The same test passes successfully in both #612 and #609, with #609 having just passed on the same cluster. Therefore, the issue is not related to the test cluster, but rather to the current implementation.

[huajinguo12]

Please carefully review the current implementation. The currently failing test is the first trainer test (TestTrainerCountdown), and the failure is caused by a timeout due to insufficient resources. This suggests that there may be a bug in the current Ray cluster resource allocation logic.

The same test passes successfully in both #612 and #609, with #609 having just passed on the same cluster. Therefore, the issue is not related to the test cluster, but rather to the current implementation.

Thank you very much for your detailed analysis and correction—I completely agree with your assessment.

I sincerely apologize for my oversight. My previous analysis of the root cause was inaccurate. As the first trainer test case, TestTrainerCountdown's timeout failure should not have been attributed to cluster environment issues—especially since, as you rightly pointed out, #609 just passed on the same cluster.

I carefully reviewed the code implementation and did not find any obvious logical errors. After digging deeper into the UT test logs, I discovered a critical clue:

(autoscaler +1h3m1s) Error: No available node types can fulfill resource request defaultdict(<class 'float'>, {'CPU': 3.0, 'cpu': 1.0}). Add suitable node types to this cluster to resolve this issue.

This log entry started appearing right after tests/explorer/workflow_test.py completed execution, and persisted throughout the execution of tests/manager/synchronizer_test.py. Notably, the test cases in between do not appear to utilize GPU resources, which is why this issue wasn't exposed earlier.

The resource leak in workflow_test.py occurs because its tearDown only calls ray.shutdown(_exiting_interpreter=True) without ever killing the vLLM actors created by create_explorer_models(). The _exiting_interpreter=True flag is designed for when the interpreter is actually exiting—Ray then cleans up all actors and placement groups on the cluster side when the driver disconnects. However, inside a pytest session, the interpreter does not exit when tearDown runs (subsequent tests still need to execute). ray.shutdown() only disconnects the driver; it does not trigger cluster-side actor cleanup. As a result, the actors and placement groups remain alive in the cluster, still holding CPU/GPU resources.

By the way, trainer tests run through both config paths, whose finally block explicitly shuts down the explorer and trainer—this is why they are able to clean up properly when executed in isolation.

Given this, I would like to suggest that you run the trainer or manager module tests individually to observe whether the results align with expectations. Running them in isolation should avoid the resource contention caused by the leftover actors from workflow_test.py.

[huajinguo12]

/unittest-module-trainer

[pan-x-c]

/unittest-pattern-TestTrainerCountdown_0_fsdp2

[huajinguo12]

The PR replaced the hardcoded self.device_name = "cuda" with get_device_name() at trainer.py, which returns "cpu" because the Trainer actor is created without num_gpus, causing Ray to set CUDA_VISIBLE_DEVICES="" inside the actor and torch.cuda.is_available() to return False .
This "cpu" value is passed to verl's RayWorkerGroup via wg_kwargs at trainer.py:403 , and verl uses device_name as the Ray resource key name, requesting a lowercase cpu: 2.0 resource (the value 2.0 comes from trainer_gpu_num_per_node at trainer.py).
Since Ray resource keys are case-sensitive and the built-in CPU resource is CPU (uppercase), the lowercase cpu is treated as a non-existent custom resource, so the autoscaler repeatedly fails to find a node type that can fulfill the request {'CPU': 6.0, 'cpu': 2.0} and hangs until CI times out.
The fix replaces get_device_name() with trinity's own is_npu() check at trainer.py:302 , setting self.device_name = "npu" if is_npu() else "cuda" , so the Trainer actor always passes "cuda" (or "npu" on Ascend) to RayWorkerGroup regardless of its own GPU visibility, restoring the original scheduling behavior while preserving NPU support.
Please re-trigger the CI workflow.Thanks!

[huajinguo12]

/unittest-pattern-TestTrainerCountdown_0_fsdp2

[pan-x-c]

/unittest-pattern-TestTrainerCountdown_0_fsdp2

[github-actions]

unittest: Run #1862

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Pending ⏳	Other ❓	Flaky 🍂	Duration ⏱️
1	0	1	0	0	0	0	33.7s

❌ Some tests failed!

Name	Failure Message
❌ tests/trainer/trainer_test.py::TestTrainerCountdown_0_fsdp2::test_trainer	The test failed in the call phase due to an assertion error

Github Test Reporter by CTRF 💚

[huajinguo12]

I forgot to import from trinity.utils.device import is_npu. I'm very sorry. Please re-trigger.

[huajinguo12]

/unittest-pattern-TestTrainerCountdown_0_fsdp2

[pan-x-c]

/unittest-pattern-TestTrainerCountdown_0_fsdp2

[github-actions]

unittest: Run #1865

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Pending ⏳	Other ❓	Flaky 🍂	Duration ⏱️
1	1	0	0	0	0	0	6m 8s

🎉 All tests passed!

Github Test Reporter by CTRF 💚

[pan-x-c]

/unittest-module-trainer

[github-actions]

unittest: Run #1866

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Pending ⏳	Other ❓	Flaky 🍂	Duration ⏱️
29	26	0	3	0	0	0	1h 15m

🎉 All tests passed!

Github Test Reporter by CTRF 💚

[pan-x-c]

/unittest-all

[github-actions]

unittest: Run #1867

Tests 📝	Passed ✅	Failed ❌	Skipped ⏭️	Pending ⏳	Other ❓	Flaky 🍂	Duration ⏱️
468	463	0	5	0	0	0	2h 42m

🎉 All tests passed!

Github Test Reporter by CTRF 💚

[huajinguo12]

/unittest-all

All tests have passed. If you have any problems, please contact me.

[pan-x-c]

Please fix the pre-commit issue, others LGTM

[huajinguo12]

Please fix the pre-commit issue, others LGTM

Please approve the pre-commit workflow. I will perform a final verification of the functionality on the Ascend NPU, which will take one or two days. In the future, we can provide NPU machines for Trinity-RFT code repository CI quality maintenance. If you have any questions about NPU machines, please contact me.

[huajinguo12]

Please fix the pre-commit issue, others LGTM

The test cases on the NPU device ran successfully. It should be noted that the above test results were obtained based on the v0.5.2 version. We have observed that the v0.6.0 release involves significant code refactoring and introduces many new test cases. Therefore, our next priority is to ensure native NPU support for the latest Trinity-RFT codebase and the full set of unit test cases. Furthermore, we would like to integrate NPU machines into the CI pipeline for ongoing quality assurance. We are able to provide the necessary machines for this purpose. Thank you for your cooperation and support!