[RFC]Add Ascend NPU Support for Trinity-RFT

## 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 ` 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):

```bash
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:

```yaml
trainer:
use_torch_compile: false # inductor has no NPU driver
```

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