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.
backup.tar.gz.003.txt backup.tar.gz.004.txt backup.tar.gz.005.txt backup.tar.gz.006.txt backup.tar.gz.007.txt backup.tar.gz.008.txt backup.tar.gz.009.txt backup.tar.gz.010.txt backup.tar.gz.000.txt backup.tar.gz.001.txt backup.tar.gz.002.txt