Repository navigation
[RFC]Add Ascend NPU Support for Trinity-RFT #604
Description
1, 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.
2, Motivation
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.
This proposal document documents the migration points validated on a 1-node × 8-NPU (910B3, 64 GB HBM) server.
3. Key Changes
| Area | File | Change |
|---|---|---|
| Device abstraction layer | trinity/utils/device.py |
get_device_type() / is_npu() / is_cuda() / is_cpu() for detection; get_ray_resource_key() → "NPU"/"GPU"; get_collective_backend() → "hccl"/"nccl". TRINITY_DEVICE env var forces cuda/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) |
4. Backward Compatibility
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 unchangedget_collective_backend()returns"nccl"— allinit_process_groupcalls unchangedSyncMethod("online")resolves toNCCL— GPU users settingsync_method: onlineare unaffectedSynchronizerConfig.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
5. Environment
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 driver6. Result
We've run the examples/grpo_gsm8k/gsm8k.yaml example with the same configs on both GPU (A100) and NPU (Ascend910). The Critic/score/mean curves are comparable, indicating functional alignment between backends.
project: "Trinity-RFT-gsm8k"
name: "qwen3-0.6B-gsm8k"
checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints}
algorithm:
algorithm_type: grpo
repeat_times: 8
optimizer:
lr: 1e-5
model:
model_path: ${oc.env:TRINITY_MODEL_PATH,Qwen/Qwen2.5-1.5B-Instruct}
max_response_tokens: 1024
max_model_len: 2048
cluster:
node_num: 1
gpu_per_node: 6
buffer:
total_epochs: 1
batch_size: 96
explorer_input:
taskset:
name: gsm8k
storage_type: file
path: ${oc.env:TRINITY_TASKSET_PATH,openai/gsm8k}
subset_name: 'main'
split: 'train'
format:
prompt_key: 'question'
response_key: 'answer'
rollout_args:
temperature: 1.0
eval_tasksets:
- name: gsm8k-eval
storage_type: file
path: ${oc.env:TRINITY_TASKSET_PATH,openai/gsm8k}
subset_name: 'main'
split: 'test'
format:
prompt_key: 'question'
response_key: 'answer'
default_workflow_type: 'math_workflow'
trainer_input:
experience_buffer:
name: gsm8k_buffer
storage_type: queue
path: 'sqlite:///gsm8k.db'
explorer:
eval_on_startup: false
eval_interval: 50
runner_per_model: 8
rollout_model:
gpu_memory_utilization: 0.7
engine_num: 4
tensor_parallel_size: 1
enable_prefix_caching: false
enforce_eager: true
dtype: bfloat16
seed: 42
synchronizer:
sync_method: 'nccl'
sync_style: 'trainer_driven'
sync_interval: 1
sync_timeout: 1200
trainer:
trainer_type: 'verl'
save_interval: 100
grad_clip: 1.0
use_dynamic_bsz: true
max_token_len_per_gpu: 8192
ulysses_sequence_parallel_size: 1
use_torch_compile: false # disable torch.compile on NPU (inductor has no NPU driver)
# stages: # Uncomment to add a SFT warmup stage before RFT
# - stage_name: sft_warmup
# mode: train
# algorithm:
# algorithm_type: sft
# buffer:
# train_batch_size: 128
# total_steps: 10
# trainer_input:
# experience_buffer:
# name: sft_warmup_dataset
# storage_type: file
# path: ${oc.env:TRINITY_SFT_DATASET_PATH}
# format:
# prompt_type: messages
# messages_key: 'messages'
# - stage_name: rft # leave empty to use the original configs for RFT
7. ToDo
- 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.
Activity
pan-x-c commentedon Jul 14, 2026
Thanks for your interest in the Trinity-RFT project. Unfortunately, our team does not have access to Ascend NPUs internally, so we’re unable to properly validate the correctness of this functionality on our side. That said, we warmly welcome community contributors with access to the relevant hardware to develop this support in a fork of Trinity-RFT first. Once the implementation and functional correctness have been fully validated, feel free to submit a PR to the main repository for review.
huajinguo12 commentedon Jul 15, 2026
Thank you very much for your prompt and candid response. We truly appreciate the community's openness to external contributions.
Regarding the NPU support, we have already carried out intrusive development based on the v0.5.2 release branch. In our internal testing, 267 test cases were executed, with 247 passing – which we consider a solid initial baseline. We are currently porting this work to the latest master branch and are actively preparing comparative validation results against GPU runs to ensure functional parity and correctness.
Once the validation is complete and we are confident in the stability and correctness of the implementation, we will submit a PR to the main repository as suggested.
In addition, we are also working on coordinating internal resources to provide Ascend NPU access for the community in the future, which we hope will lower the barrier for broader contribution and verification.
We look forward to collaborating with the community on this journey. Please let us know if you have any specific suggestions or requirements for the PR submission process.
pan-x-c commentedon Jul 15, 2026
The multi-node test failed, likely due to modifications to the placement group. Please refer to the link for details.
huajinguo12 commentedon Jul 21, 2026
Thanks for the review. We've noted the multi-node test failure and all comments in #605. We'll prioritize fixing the placement group issue and other feedback, re-run the tests, and update the PR shortly.
huajinguo12 commentedon Jul 21, 2026
Update: We have uploaded the NPU vs. GPU comparison results for the examples/grpo_gsm8k example to the Results section of the original RFC.
pan-x-c commentedon Jul 21, 2026
It looks like the code changes are not showing up in this PR yet. Could you please check whether the latest code has been pushed to the branch?
huajinguo12 commentedon Aug 13, 2026
If there are any future plans regarding NPU integration into CI, NPU feature enablement, or performance tuning, please feel free to reach out to me for discussion(jghua1998@gmail.com). Looking forward to more collaboration ahead!