[RFC]Add Ascend NPU Support for Trinity-RFT

[huajinguo12]

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_npu import 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" — all init_process_group calls unchanged
• SyncMethod("online") resolves to NCCL — GPU users setting sync_method: online are unaffected
• SynchronizerConfig.sync_method defaults to NCCL
• FSDP getattr(torch, "cuda").set_device() + .to("cuda") — redundant but idempotent (verl base class already sets device)
• TrainerConfig.device remains configurable; GPU users can set "cuda" explicitly (auto-detection also handles it)
• use_torch_compile defaults to True — 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 vLLM

NPU users disable torch.compile in YAML:

trainer:
  use_torch_compile: false  # inductor has no NPU driver

6. 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.

[pan-x-c]

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]

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.

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]

The multi-node test failed, likely due to modifications to the placement group. Please refer to the link for details.