Home
# Add Ascend NPU Support While Preserving GPU Compatibility
## Summary
This proposal document extends Trinity-RFT to support Huawei Ascend NPUs (910B and above) as a first-class backend, alongside the existing CUDA backend. All changes are fully backward-compatible — GPU users see no behavior change.
Trinity-RFT is a post-training framework built on verl, Ray, and vLLM — all originally GPU-centric. Ascend NPU is now a first-class PyTorch backend via torch_npu, with HCCL mirroring NCCL's API and vllm-ascend providing vLLM support. Growing demand for RLHF / GRPO post-training on Ascend 910B/C NPUs motivates This proposal document. The goal: trinity run --config <yaml> should launch end-to-end (Ray + vLLM rollout + verl FSDP trainer) on an NPU-only cluster without forking the codebase, while keeping the GPU path untouched.
Trinity用例执行结果.xlsx
This proposal document documents the migration points validated on a 1-node × 8-NPU (910B3, 64 GB HBM) server.
| Area | File | Change |
|---|---|---|
| Device abstraction layer | trinity/utils/device.py |
get_device_type() / is_npu() / `is_cu |
| npu-tensorboard.zip | ||
da()/is_cpu()for detection;get_ray_resource_key()→"NPU"/"GPU"; get_collective_backend()→"hccl"/"nccl". TRINITY_DEVICEenv var forcescuda/npu/cpu` and skips auto-detection. |
||
| Ray resource detection | trinity/common/config_validator.py |
get_ray_resource_key() reads "NPU" or "GPU" from Ray cluster Resources |
| Placement group | trinity/common/models/__init__.py |
Bundle key {resource_key: 1, "CPU": 8}; STRICT_PACK strategy kept unconditionally (benefits both NPU HCCL and GPU NVLink) |
| Sync method enum | trinity/common/constants.py |
"online" alias dynamically resolved to "hccl" (NPU) or "nccl" (GPU) via metaclass override |
| Sync method default | trinity/common/config.py |
SynchronizerConfig.sync_method uses default_factory to pick HCCL/NCCL by device |
| vLLM weight-sync backend | trinity/common/models/vllm_model.py |
backend: Optional[str] = None → get_collective_backend() if unset |
| vLLM worker backend | trinity/common/models/vllm_worker.py |
Same as above |
| Checkpoint converter | trinity/manager/checkpoint_converter.py |
init_process_group(get_collective_backend()) (verl API, manager layer) |
| FSDP Actor init | trinity/trainer/verl/fsdp_workers.py |
getattr(torch, get_device_name()).set_device(local_rank) + .to(get_device_name()); backend=f"cpu:gloo,{get_device_name()}:{get_collective_backend()}"
|
| FSDP weight-sync | trinity/trainer/verl/fsdp_workers.py |
backend=get_collective_backend(), device_id=torch.device(f"{get_device_name()}:{get_device_id()}"), getattr(torch, get_device_name()).synchronize()
|
| FSDP Critic init | trinity/trainer/verl/fsdp_workers.py |
backend=get_collective_backend() |
| Distributed guard | trinity/utils/distributed.py |
Relaxed backend == "nccl" assertion so HCCL callers pass through |
| Trainer device | trinity/trainer/verl/verl_config.py |
device: str = "npu" (was "cuda"); auto_set_device() flips Ray resource request from GPU to NPU |
torch.compile |
trinity/trainer/verl/verl_trainer.py |
TrainerConfig.use_torch_compile exposed as top-level field; NPU users set false (inductor has no NPU driver) |
Every modification is gated by the device abstraction layer, so on CUDA:
- No manual
torch_npuimport in launcher — NPU init handled by environment, launcher stays device-agnostic -
get_ray_resource_key()returns"GPU"— Ray resource detection unchanged -
get_collective_backend()returns"nccl"— allinit_process_groupcalls unchanged -
SyncMethod("online")resolves toNCCL— GPU users settingsync_method: onlineare unaffected -
SynchronizerConfig.sync_methoddefaults toNCCL - FSDP
getattr(torch, "cuda").set_device()+.to("cuda")— redundant but idempotent (verl base class already sets device) -
TrainerConfig.deviceremains configurable; GPU users can set"cuda"explicitly (auto-detection also handles it) -
use_torch_compiledefaults toTrue— GPU path verbatim
Full test suite executed on Ascend 910B3 (CANN 8.5.1, 8 NPUs, conda env trinity). Summary: 225 passed / 11 failed / 4 skipped out of 240 cases. All failures are environment/version issues, not code bugs introduced by This proposal document.
| 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 | 1 | 0 |
test_normal_conversation_data, test_no_assistant_messages, test_first_message_is_assistant: tokenizer version drift — snapshot tests hardcode exact token counts, transformers upgrade shifted tokenization results (pre-existing, not PR-introduced) |
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 | — |
| Total | 240 | 225 | 4 | 11 failed — all environment/version/test-isolation issues, none caused by This proposal document |
Failure category breakdown: tokenizer version drift (3, pre-existing), vLLM API changes (4), NPU precision (1), missing model/deps (3, includes 1 undocumented vllm failure). Total 11 unfixed.
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 vLLMNPU users disable torch.compile in YAML:
trainer:
use_torch_compile: false # inductor has no NPU driver-
Megatron training backend: Currently the NPU backend only supports the FSDP training backend. Adapting the Megatron training backend for NPU will be the focus of the next phase of work.