Repository navigation
[ckpt, rollout, vllm] feat: support delta weight sync on Ascend NPU (HCCL transport) - #331
Merged
Merged
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds delta (sparse) weight sync over HCCL for Ascend NPU to the
delta_shardedcheckpoint engine, so NPU RL runs can skip full-weight resyncs between trainer and
rollout the same way CUDA hosts do.
Highlights
manifests live in a transport-agnostic
DeltaSyncMixinshared by the NCCL and HCCLengines. The CUDA path is refactored but behaviourally unchanged.
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.
backend=delta_shardedauto-selects the HCCL engine on NPU and theNCCL engine on CUDA; the vLLM weight-transfer backend likewise defaults to
twinkle_delta_hcclon NPU andtwinkle_delta_ipcelsewhere.tie_word_embeddings=Truethe HFstate_dict lists one tensor twice (
lm_head.weightand*.embed_tokens.weight). vLLMloads 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=Trueis not supported (the HCCLgroup 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 onCPU (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.