Skip to content

[ckpt, rollout, vllm] feat: support delta weight sync on Ascend NPU (HCCL transport) - #331

Merged
kevssim merged 16 commits into
modelscope:dev-swift-v5from
kevssim:delta-sync
Oct 22, 2026
Merged

kevssim merged 16 commits into
modelscope:dev-swift-v5from
kevssim:delta-sync

Conversation

@kevssim

@kevssim kevssim commented Oct 21, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

Adds delta (sparse) weight sync over HCCL for Ascend NPU to the delta_sharded
checkpoint engine, so NPU RL runs can skip full-weight resyncs between trainer and
rollout the same way CUDA hosts do.

Highlights

  • One pipeline, two transports — the seed/steady state machine, bucketing and flush
    manifests live in a transport-agnostic DeltaSyncMixin shared by the NCCL and HCCL
    engines. The CUDA path is refactored but behaviourally unchanged.
  • Direct broadcast into the vLLM workers — NPU has no device IPC, so trainer rank 0
    broadcasts each flush into a persistent PyHccl group whose members are the vLLM worker
    processes themselves: one zmq manifest subscription plus one or two broadcasts per flush,
    no IPC-handle negotiation.
  • Zero new config — backend=delta_sharded auto-selects the HCCL engine on NPU and the
    NCCL engine on CUDA; the vLLM weight-transfer backend likewise defaults to
    twinkle_delta_hccl on NPU and twinkle_delta_ipc elsewhere.
  • Tied-embedding fix on the full-sync path — with tie_word_embeddings=True the HF
    state_dict lists one tensor twice (lm_head.weight and *.embed_tokens.weight). vLLM
    loads only the name its traversal reaches first and then asserts the canonical one did
    load (AutoWeightsLoader._check_skipped_aliases), so sending both made that check fail.
    The duplicate head is now dropped; one write covers both.

Scope and limitations

vLLM rollout only; BF16 non-quantized dense models; trainer FSDP2 with TP=PP=EP=1; vLLM
internal data_parallel_size=1; no LoRA; rebuild_group=True is not supported (the HCCL
group is persistent). Multiple rollout replicas are supported (replica-major, TP-minor).

Test

Unit and multi-process suites under tests/checkpoint_engine/:

  • test_delta_hccl.py — trainer-side engine: rendezvous, lazy communicator, close.
  • test_delta_hccl_receiver.py — receiver side: wire decode, pump protocol, channel lifecycle.
  • test_delta_hccl_loopback.py + _hccl_loopback.py — multi-process end-to-end protocol on
    CPU (gloo stands in for PyHccl, zmq is real), covering seed/steady rounds, zero-change
    rounds, bucket slicing, reseed, corrupted payloads and shuffled receiver start order.
  • test_delta_npu_orchestration.py — manager rendezvous and the rollout RPC sequence, mocked.
  • test_delta_vllm.py — vLLM integration: patch loaders, tied aliases, TP geometry.

Validated on a 2-node Ascend cluster (2-card FSDP2 trainer + 2-card vLLM rollout, TP=2)
with Qwen3.5-9B and Qwen3.5-4B: delta-synced rollout weights are byte-identical to a
full resync across change ratios from 0% to 3%, plus a real optimizer step.

…ixins

Rename *_cuda_delta_* methods to *_delta_* and CudaDelta{Trainer,Sampler}Mixin to Delta{Trainer,Sampler}Mixin. CudaDeltaReceiver keeps its name since it is CUDA-specific. No behavior change.
Replace torch.hash_tensor (CPU/CUDA kernels only) with a byte-wise XOR fold that runs on any device, including NPU. Sender and receiver must be upgraded together.
…CCL engine

Move the seed/steady state machine, gather batching, bucketing and flush/manifest assembly into delta/pipeline.py. The NCCL engine keeps only transport concerns and implements four hooks: _wire_device, _ensure_wire, _broadcast_flush and _release_after_round. No behavior change for CUDA.
…rainers

Add resolve_delta_engine_cls (GPU -> NCCL engine, NPU -> HCCL engine, imported lazily) and use it in the manager and trainer. The manager no longer rejects non-GPU platforms and fails at construction when no transport exists. The trainer derives its device from the platform and validate_fsdp2 accepts NPU. The NPU transport itself is not implemented yet, so selecting it raises NotImplementedError.
DeltaHCCLCheckpointEngine implements the four DeltaSyncMixin hooks over a PyHccl broadcast group whose other members are the vLLM workers. Rank 0 builds the communicator lazily at the first publish, so a ready group implies every receiver has subscribed to the zmq manifest channel. The group persists across rounds; rebuild_group=True is rejected. The zmq PUB/SUB helpers move to zmq_pubsub.py and are shared with the NCCL engine.
Extract the device-agnostic flush decode into delta/decode.py and share the decode-and-apply path between the IPC and HCCL weight-transfer engines. The new twinkle_delta_hccl engine runs inside each vLLM worker: it connects the zmq manifest SUB before joining the trainer's PyHccl group, then pumps manifests and broadcasts into vLLM checkpoint patches. The worker extension registers the engine, reports the real device type and forwards per-worker group ranks.
# Conflicts:
#	src/twinkle/checkpoint_engine/__init__.py
#	src/twinkle/checkpoint_engine/manager.py
#	src/twinkle/checkpoint_engine/nccl_checkpoint_engine.py
#	src/twinkle/model/transformers/transformers.py
#	src/twinkle/sampler/vllm_sampler/vllm_sampler.py
# Conflicts:
#	src/twinkle/sampler/vllm_sampler/vllm_sampler.py
- checkpoint_engine/delta: protocol/ (wire, decode), trainer/ (mixin, spec,
  export, backends), transport/ (pipeline, gather, nccl_engine, hccl_engine),
  rollout/ (mixin, cuda_receiver, hccl_receiver); CudaDeltaReceiver split out
  of the rollout mixin.
- sampler/vllm_sampler/delta: transfer_engine (was delta_weight_transfer),
  cuda_ipc (was bucketed_weight_transfer), weight_update_utils.
- transformers: restore _checkpoint_tensor so CPU-offloaded DTensors are
  staged on device before full_tensor (dropped in the upstream merge).
# Conflicts:
#	src/twinkle/sampler/vllm_sampler/vllm_sampler.py
A refactor removed _not_encoded but left its @staticmethod behind, which
then decorated _lazy_wrap_model(self) and broke every self._lazy_wrap_model()
call with a missing 'self' argument.
@kevssim
kevssim merged commit cae2a29 into modelscope:dev-swift-v5 Oct 22, 2026
0 of 5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant