From b88fc27d1aa7188f0d97fb95aa69f3c27a9bf8fc Mon Sep 17 00:00:00 2001 From: Nan Date: Fri, 22 May 2026 05:59:06 +0000 Subject: [PATCH 1/2] =?UTF-8?q?feat:=20delta=20weight=20sync=20=E2=80=94?= =?UTF-8?q?=20disk=20+=20NCCL=20transports?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds delta weight sync: ship only changed positions + values instead of full parameters. Two transports (disk for cross-DC disaggregation, NCCL for intra-DC), three encodings (indices, deltas, deltas_zstd), lossless selective overwrite via NaN sentinel. Examples, docs, and performance comparison to follow in a separate PR. Co-Authored-By: Claude Opus 4.6 --- docker/patch/latest/sglang.patch | 569 +++++++++++- slime/backends/megatron_utils/actor.py | 10 +- slime/backends/megatron_utils/data.py | 3 +- slime/backends/megatron_utils/sglang.py | 4 + .../update_weight_from_distributed.py | 178 ++-- .../update_weight_from_distributed_delta.py | 856 ++++++++++++++++++ .../update_weight_from_tensor.py | 9 + slime/backends/sglang_utils/sglang_engine.py | 46 +- slime/utils/arguments.py | 74 ++ slime/utils/train_metric_utils.py | 8 +- 10 files changed, 1621 insertions(+), 136 deletions(-) create mode 100644 slime/backends/megatron_utils/update_weight/update_weight_from_distributed_delta.py diff --git a/docker/patch/latest/sglang.patch b/docker/patch/latest/sglang.patch index 4a13e2f9b4..622b5e630e 100644 --- a/docker/patch/latest/sglang.patch +++ b/docker/patch/latest/sglang.patch @@ -1,5 +1,5 @@ diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py -index 691f06411d..671ac81c48 100644 +index 691f064..671ac81 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -294,6 +294,7 @@ class ModelConfig: @@ -11,7 +11,7 @@ index 691f06411d..671ac81c48 100644 ]: self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py -index f7d4092d85..3aae51c849 100644 +index f7d4092..3aae51c 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -17,6 +17,7 @@ class KVArgs: @@ -23,7 +23,7 @@ index f7d4092d85..3aae51c849 100644 aux_data_lens: List[int] aux_item_lens: List[int] diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py -index f54c882cc2..03832002f0 100644 +index f54c882..0383200 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -21,6 +21,7 @@ Life cycle of a request in the decode server @@ -209,7 +209,7 @@ index f54c882cc2..03832002f0 100644 if not hasattr(self, "polling_count"): diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py -index 64d97f5c69..4ef08446aa 100644 +index 64d97f5..4ef0844 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -31,6 +31,7 @@ from sglang.srt.disaggregation.mooncake.utils import ( @@ -341,7 +341,7 @@ index 64d97f5c69..4ef08446aa 100644 # Only the last chunk we need to send the aux data ret = self.send_aux( diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py -index 8eadf81954..c180ce79f3 100644 +index 8eadf81..c180ce7 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -20,6 +20,8 @@ Life cycle of a request in the prefill server @@ -484,7 +484,7 @@ index 8eadf81954..c180ce79f3 100644 release_kv_cache(req, self.tree_cache) # unlock the tree req.finished_reason = FINISH_LENGTH(length=0) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py -index d7956a6048..0ced278713 100644 +index d7956a6..0ced278 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py @@ -28,6 +28,17 @@ if TYPE_CHECKING: @@ -691,7 +691,7 @@ index d7956a6048..0ced278713 100644 ######################### diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py -index d864e4abaa..3a000a80f2 100644 +index d864e4a..3a000a8 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -69,6 +69,7 @@ from sglang.srt.managers.io_struct import ( @@ -728,7 +728,7 @@ index d864e4abaa..3a000a80f2 100644 """Get weights by parameter name.""" obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py -index 6978e0c062..80dc159e8f 100644 +index 6978e0c..80dc159 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -127,6 +127,7 @@ from sglang.srt.managers.io_struct import ( @@ -800,7 +800,7 @@ index 6978e0c062..80dc159e8f 100644 @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Request): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py -index dfc5507de0..be9501b05a 100644 +index dfc5507..be9501b 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -242,6 +242,7 @@ class Envs: @@ -812,7 +812,7 @@ index dfc5507de0..be9501b05a 100644 # Extra slots in req_to_token_pool for decode workers (only effective when # max_num_reqs > 32). Increases pool capacity so more KV cache transfers diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py -index 02ef4e2440..fd5a43cce8 100644 +index 02ef4e2..fd5a43c 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1,6 +1,7 @@ @@ -890,7 +890,7 @@ index 02ef4e2440..fd5a43cce8 100644 if enable_dual_stream: current_stream = torch.cuda.current_stream() diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py -index 72483f4ea6..2e1148d189 100644 +index 72483f4..2e1148d 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -702,6 +702,7 @@ class FusedMoE(torch.nn.Module): @@ -910,7 +910,7 @@ index 72483f4ea6..2e1148d189 100644 ) diff --git a/python/sglang/srt/layers/moe/routed_experts_capturer.py b/python/sglang/srt/layers/moe/routed_experts_capturer.py -index 00bd687555..12d5577af2 100644 +index 00bd687..12d5577 100644 --- a/python/sglang/srt/layers/moe/routed_experts_capturer.py +++ b/python/sglang/srt/layers/moe/routed_experts_capturer.py @@ -8,10 +8,15 @@ import torch @@ -971,7 +971,7 @@ index 00bd687555..12d5577af2 100644 def get_routed_experts( diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py -index a13c53af4d..1d80d06b13 100644 +index a13c53a..1d80d06 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -500,7 +500,7 @@ class CompressedTensorsConfig(QuantizationConfig): @@ -994,7 +994,7 @@ index a13c53af4d..1d80d06b13 100644 self, layer: torch.nn.Module, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py -index 7a8fb65421..f1c85899cd 100644 +index 7a8fb65..f1c8589 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -17,7 +17,10 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( @@ -1106,10 +1106,70 @@ index 7a8fb65421..f1c85899cd 100644 is_k_full=self.is_k_full, routed_scaling_factor=self.moe_runner_config.routed_scaling_factor, diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index bd97965345..e6a147c1b4 100644 +index bd97965..af97f11 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -1449,6 +1449,18 @@ class ResumeMemoryOccupationReqOutput(BaseReq): +@@ -1224,6 +1224,8 @@ class ContinueGenerationReqInput(BaseReq): + class UpdateWeightFromDiskReqInput(BaseReq): + # The model path with the new weights + model_path: str ++ # Required iff ``load_format == "delta"``: basenames under ``model_path`` to apply. ++ files: Optional[List[str]] = None + # The format to load the weights + load_format: Optional[str] = None + # Whether to abort all requests before updating weights +@@ -1254,6 +1256,41 @@ class UpdateWeightFromDiskReqOutput(BaseReq): + num_paused_requests: Optional[int] = 0 + + ++class DeltaEncoding(str, Enum): ++ """Position encoding for delta weight updates.""" ++ ++ # int32 absolute nonzero offsets. ++ INDICES = "indices" ++ # uint16 gap-deltas between consecutive sorted positions; uint32 per-param fallback. ++ DELTAS = "deltas" ++ # ``deltas`` wrapped in zstd L1. ++ DELTAS_ZSTD = "deltas_zstd" ++ ++ ++@dataclass ++class DeltaParam: ++ """Per-param slice into the shared (positions, values) bucket. ``pos_*`` index ++ into the uint8 byte blob; ``val_*`` index into the param-dtype value tensor.""" ++ ++ name: str ++ dtype: str ++ shape: List[int] ++ pos_start: int ++ pos_end: int ++ pos_width: int # 2 or 4 ++ val_start: int ++ val_end: int ++ ++ ++@dataclass ++class DeltaSpec: ++ """Decoding manifest for one delta bucket. ``checksum`` is verified on apply.""" ++ ++ encoding: DeltaEncoding ++ params: List[DeltaParam] ++ checksum: int = 0 ++ ++ + @dataclass + class UpdateWeightsFromDistributedReqInput(BaseReq): + names: List[str] +@@ -1269,6 +1306,8 @@ class UpdateWeightsFromDistributedReqInput(BaseReq): + weight_version: Optional[str] = None + # Optional format specification for loading + load_format: Optional[str] = None ++ # JSON-encoded DeltaSpec; required iff load_format == "delta". ++ delta: Optional[str] = None + + + @dataclass +@@ -1449,6 +1488,18 @@ class ResumeMemoryOccupationReqOutput(BaseReq): pass @@ -1128,7 +1188,7 @@ index bd97965345..e6a147c1b4 100644 @dataclass class CheckWeightsReqInput(BaseReq): action: str -@@ -1753,6 +1765,8 @@ class GetLoadReqOutput(BaseReq): +@@ -1753,6 +1804,8 @@ class GetLoadReqOutput(BaseReq): num_waiting_reqs: int num_tokens: int ts_tic: float @@ -1138,7 +1198,7 @@ index bd97965345..e6a147c1b4 100644 @dataclass diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py -index e0a1669fb3..fbbb6bb12b 100644 +index e0a1669..fbbb6bb 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py @@ -496,6 +496,35 @@ def monkey_patch_uvicorn_multiprocessing(timeout: float = 10): @@ -1178,7 +1238,7 @@ index e0a1669fb3..fbbb6bb12b 100644 class SenderWrapper: def __init__(self, port_args: PortArgs, send_to_scheduler: zmq.Socket): diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py -index 0b26be6c6d..2ea1042cf9 100644 +index 0b26be6..2ea1042 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1972,7 +1972,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): @@ -1194,7 +1254,7 @@ index 0b26be6c6d..2ea1042cf9 100644 break diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index 67af2d0de9..122ddb3874 100644 +index 67af2d0..122ddb3 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -120,6 +120,7 @@ from sglang.srt.managers.io_struct import ( @@ -1214,7 +1274,7 @@ index 67af2d0de9..122ddb3874 100644 (ReleaseMemoryOccupationReqInput, self.release_memory_occupation), (ResumeMemoryOccupationReqInput, self.resume_memory_occupation), diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py -index 496cd96656..cf2d43015a 100644 +index 496cd96..cf2d430 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -1154,7 +1154,7 @@ class SchedulerOutputProcessorMixin: @@ -1227,7 +1287,7 @@ index 496cd96656..cf2d43015a 100644 BatchTokenIDOutput( rids=rids, diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py -index c02ed7997d..61733c4127 100644 +index c02ed79..61733c4 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/python/sglang/srt/managers/scheduler_profiler_mixin.py @@ -349,7 +349,7 @@ class SchedulerProfilerMixin: @@ -1240,7 +1300,7 @@ index c02ed7997d..61733c4127 100644 if self.profile_in_progress: # force trace flush diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py -index abcda67946..a53848b79d 100644 +index abcda67..a53848b 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py @@ -12,6 +12,7 @@ from sglang.srt.constants import ( @@ -1305,7 +1365,7 @@ index abcda67946..a53848b79d 100644 def check_weights(self: Scheduler, recv_req: CheckWeightsReqInput): diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py -index 544c609401..841658c30e 100644 +index 544c609..841658c 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -59,6 +59,8 @@ from sglang.srt.managers.io_struct import ( @@ -1357,7 +1417,7 @@ index 544c609401..841658c30e 100644 self: TokenizerManager, obj: InitWeightsSendGroupForRemoteInstanceReqInput, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 81424329a0..2c132be63d 100644 +index 8142432..5378268 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1383,7 +1383,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): @@ -1378,6 +1438,24 @@ index 81424329a0..2c132be63d 100644 self.is_pause_cond.notify_all() async def update_weights_from_disk( +@@ -1446,7 +1446,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): + self.model_update_result = asyncio.Future() + if self.server_args.dp_size == 1: + result = await self.model_update_result +- if result.success: ++ if result.success and obj.load_format != "delta": + self._update_model_path_info(obj.model_path, obj.load_format) + return result.success, result.message, result.num_paused_requests + else: # self.server_args.dp_size > 1 +@@ -1454,7 +1454,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): + result = await self.model_update_result + + all_success = all([r.success for r in result]) +- if all_success is True: ++ if all_success is True and obj.load_format != "delta": + self._update_model_path_info(obj.model_path, obj.load_format) + all_message = [r.message for r in result] + all_message = " | ".join(all_message) @@ -1965,25 +1965,23 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): priority = getattr(state.obj, "priority", None) if priority is not None: @@ -1412,7 +1490,7 @@ index 81424329a0..2c132be63d 100644 if state.finished: retraction_count = ( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py -index 7f63610da8..fb56de1583 100644 +index 7f63610..8f5dd6e 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -29,6 +29,7 @@ from sglang.srt.managers.io_struct import ( @@ -1423,7 +1501,23 @@ index 7f63610da8..fb56de1583 100644 SendWeightsToRemoteInstanceReqInput, UnloadLoRAAdapterReqInput, UpdateWeightFromDiskReqInput, -@@ -170,6 +171,11 @@ class BaseTpWorker(ABC): +@@ -96,6 +97,7 @@ class BaseTpWorker(ABC): + success, message = self.model_runner.update_weights_from_disk( + recv_req.model_path, + recv_req.load_format, ++ files=recv_req.files, + recapture_cuda_graph=recv_req.recapture_cuda_graph, + ) + return success, message +@@ -151,6 +153,7 @@ class BaseTpWorker(ABC): + recv_req.shapes, + recv_req.group_name, + recv_req.load_format, ++ recv_req.delta, + ) + return success, message + +@@ -170,6 +173,11 @@ class BaseTpWorker(ABC): success, message = self.model_runner.update_weights_from_ipc(recv_req) return success, message @@ -1436,7 +1530,7 @@ index 7f63610da8..fb56de1583 100644 parameter = self.model_runner.get_weights_by_name( recv_req.name, recv_req.truncate_size diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py -index 3c1e97daab..e5128e5ee2 100644 +index 3c1e97d..e5128e5 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -755,9 +755,8 @@ class HiRadixCache(RadixCache): @@ -1476,7 +1570,7 @@ index 3c1e97daab..e5128e5ee2 100644 self._inc_hit_count(new_node, chunked) total_prefix_length += prefix_len diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py -index e4c158cda9..cf7333235f 100644 +index e4c158c..cf73332 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1854,9 +1854,12 @@ class NSATokenToKVPool(MLATokenToKVPool): @@ -1559,7 +1653,7 @@ index e4c158cda9..cf7333235f 100644 kv_size_bytes = super().get_kv_size_bytes() for index_k_cache in self.index_k_with_scale_buffer: diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py -index 7d16160372..70fbdc702f 100644 +index 7d16160..70fbdc7 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -512,7 +512,17 @@ class RadixCache(BasePrefixCache): @@ -1594,10 +1688,40 @@ index 7d16160372..70fbdc702f 100644 return DecLockRefResult(delta=delta) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index a59742b943..a7347c15b8 100644 +index a59742b..d20b8d0 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py -@@ -406,7 +406,12 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -15,18 +15,20 @@ + + from __future__ import annotations + ++import contextlib + import datetime + import gc + import inspect + import json + import logging ++import math + import os + import socket + import threading + import time + from collections import defaultdict + from dataclasses import dataclass +-from typing import Callable, List, Optional, Tuple, Union ++from typing import Callable, Dict, List, Optional, Tuple, Union + + import torch + import torch.distributed as dist +@@ -118,6 +120,7 @@ from sglang.srt.layers.sampler import create_sampler + from sglang.srt.layers.torchao_utils import apply_torchao_config_to_model + from sglang.srt.lora.lora_manager import LoRAManager + from sglang.srt.lora.lora_registry import LoRARef ++from sglang.srt.managers.io_struct import DeltaEncoding, DeltaParam, DeltaSpec + from sglang.srt.managers.schedule_batch import sanity_check_mm_pad_shift_value + from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +@@ -406,7 +409,12 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.forward_stream = torch.get_device_module(self.device).Stream() # CPU offload @@ -1611,7 +1735,7 @@ index a59742b943..a7347c15b8 100644 self._weight_checker = WeightChecker(model_runner=self) -@@ -646,7 +651,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -646,7 +654,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) # Init routed experts capturer @@ -1621,7 +1745,209 @@ index a59742b943..a7347c15b8 100644 if self.device == "cuda" or self.device == "musa": self.init_cublas() -@@ -2767,11 +2773,19 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -1339,8 +1348,16 @@ class ModelRunner(ModelRunnerKVCacheMixin): + load_format: str, + weight_name_filter: Optional[Callable[[str], bool]] = None, + recapture_cuda_graph: bool = False, ++ files: Optional[List[str]] = None, + ) -> tuple[bool, str]: +- """Update engine weights in-place from the disk.""" ++ """Update weights in-place from disk. For ``load_format="delta"``, read + ++ apply each basename in ``files`` under ``model_path``; otherwise reload the ++ HF checkpoint at ``model_path``.""" ++ if load_format == "delta": ++ if not files: ++ return False, "load_format='delta' requires non-empty ``files``" ++ return self._apply_delta([os.path.join(model_path, f) for f in files]) ++ + logger.info( + f"Update engine weights online from disk begin. " + f"avail mem={get_available_gpu_memory(self.device, self.gpu_id, empty_cache=False):.2f} GB" +@@ -1563,6 +1580,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): + shapes, + group_name, + load_format: Optional[str] = None, ++ delta: Optional[str] = None, + ): + """ + Update specific parameter in the model weights online +@@ -1583,6 +1601,18 @@ class ModelRunner(ModelRunnerKVCacheMixin): + return self._update_bucketed_weights_from_distributed( + names, dtypes, shapes, group_name + ) ++ if load_format == "delta": ++ if delta is None: ++ return False, "load_format='delta' requires a DeltaSpec in the request" ++ spec_dict = json.loads(delta) ++ spec = DeltaSpec( ++ encoding=DeltaEncoding(spec_dict["encoding"]), ++ params=[DeltaParam(**p) for p in spec_dict["params"]], ++ checksum=int(spec_dict["checksum"]), ++ ) ++ return self._apply_delta_from_distributed( ++ names, dtypes, shapes, group_name, spec ++ ) + try: + weights = [] + handles = [] +@@ -1646,6 +1676,156 @@ class ModelRunner(ModelRunnerKVCacheMixin): + logger.error(error_msg) + return False, error_msg + ++ def _decode_delta_one_param( ++ self, ++ encoding: DeltaEncoding, ++ positions: torch.Tensor, ++ values: torch.Tensor, ++ p: DeltaParam, ++ ) -> torch.Tensor: ++ """Decode one param's (positions, values) into a full-shape NaN-masked tensor. ++ NaN at unchanged positions triggers the patched-copy on apply.""" ++ numel = math.prod(p.shape) ++ param_dtype = p.dtype if isinstance(p.dtype, torch.dtype) else getattr(torch, p.dtype) ++ flat = torch.full((numel,), float("nan"), dtype=param_dtype, device=self.device) ++ val_slice = values[p.val_start : p.val_end] ++ if val_slice.numel() == 0: ++ return flat.view(tuple(p.shape)) ++ ++ pos_bytes = positions[p.pos_start : p.pos_end] ++ if encoding is DeltaEncoding.INDICES: ++ width = 4 # int32 absolute indices ++ elif encoding in (DeltaEncoding.DELTAS, DeltaEncoding.DELTAS_ZSTD): ++ width = p.pos_width # uint16 or uint32 gap-deltas ++ else: ++ raise ValueError(f"unsupported delta encoding: {encoding!r}") ++ ++ n_elems = pos_bytes.numel() // width ++ b = pos_bytes.view(n_elems, width).to(torch.int64) ++ if width == 2: ++ unpacked = b[:, 0] | (b[:, 1] << 8) ++ else: # 4 ++ unpacked = b[:, 0] | (b[:, 1] << 8) | (b[:, 2] << 16) | (b[:, 3] << 24) ++ ++ if encoding is DeltaEncoding.INDICES: ++ idx = unpacked ++ else: ++ # Sender encodes ``delta[k] = idx[k] - idx[k-1] - 1`` with idx[-1] := -1; ++ # receiver inverts with ``idx = cumsum(delta + 1) - 1``. ++ idx = (unpacked + 1).cumsum(dim=0) - 1 ++ # Sender may concat values across params of mixed dtypes (bf16 weights ++ # + fp32 norms in one bucket); torch.cat promotes to the widest dtype, ++ # so re-cast each slice back to the param's own dtype. The promoted ++ # round-trip is exact (bf16 ⊂ fp32), no precision loss. ++ flat.index_copy_(0, idx, val_slice.to(param_dtype)) ++ return flat.view(tuple(p.shape)) ++ ++ def _apply_delta_payload( ++ self, ++ encoding: DeltaEncoding, ++ params: List[DeltaParam], ++ positions: torch.Tensor, ++ values: torch.Tensor, ++ expected_checksum: int, ++ ) -> None: ++ """Verify checksum, decode each param, apply via the patched-copy context. ++ ``load_weights`` is called per ``update_weight_delta_chunk_bytes`` budget.""" ++ actual_checksum = _delta_checksum(positions, values) ++ if actual_checksum != expected_checksum: ++ raise RuntimeError( ++ f"delta checksum mismatch: expected={expected_checksum} got={actual_checksum}; " ++ "indicates corruption between sender encode and receiver apply" ++ ) ++ chunk_byte_cap = self.server_args.update_weight_delta_chunk_bytes ++ with _delta_apply_context(self.model): ++ chunk: List[Tuple[str, torch.Tensor]] = [] ++ chunk_bytes = 0 ++ for p in params: ++ t = self._decode_delta_one_param(encoding, positions, values, p) ++ tensor_bytes = t.numel() * t.element_size() ++ if chunk_bytes + tensor_bytes > chunk_byte_cap and chunk: ++ self.model.load_weights(chunk) ++ chunk = [] ++ chunk_bytes = 0 ++ chunk.append((p.name, t)) ++ chunk_bytes += tensor_bytes ++ if chunk: ++ self.model.load_weights(chunk) ++ ++ def _decode_and_apply_blob(self, blob: bytes) -> None: ++ """Decode + apply one decompressed safetensors blob from the delta sender.""" ++ from safetensors.torch import load as st_load ++ ++ # st_load only returns tensors, so parse the header for metadata. ++ hdr_len = int.from_bytes(blob[:8], "little") ++ meta = json.loads(blob[8:8 + hdr_len]).get("__metadata__", {}) ++ encoding = DeltaEncoding(meta["encoding"]) ++ params = [DeltaParam(**p) for p in json.loads(meta["params"])] ++ expected_checksum = int(meta["checksum"]) ++ ++ tensors = st_load(blob) ++ positions = tensors["__positions__"].to(self.device, non_blocking=True) ++ values = tensors["__values__"].to(self.device, non_blocking=True) ++ self._apply_delta_payload(encoding, params, positions, values, expected_checksum) ++ ++ def _apply_delta_from_distributed( ++ self, ++ names: List[str], ++ dtypes: List[str], ++ shapes: List[List[int]], ++ group_name: str, ++ delta: DeltaSpec, ++ ) -> tuple[bool, str]: ++ """NCCL receive: broadcast (positions, values) from sender, then apply.""" ++ try: ++ recv: Dict[str, torch.Tensor] = {} ++ handles = [] ++ for name, dtype, shape in zip(names, dtypes, shapes): ++ target_dtype = dtype if isinstance(dtype, torch.dtype) else getattr(torch, dtype) ++ t = torch.empty(shape, dtype=target_dtype, device=self.device) ++ handles.append( ++ torch.distributed.broadcast( ++ t, src=0, group=self._model_update_group[group_name], async_op=True ++ ) ++ ) ++ recv[name] = t ++ for h in handles: ++ h.wait() ++ ++ self._apply_delta_payload( ++ delta.encoding, delta.params, recv["__positions__"], recv["__values__"], delta.checksum ++ ) ++ return True, "ok" ++ except Exception as e: ++ error_msg = f"Failed to apply delta from distributed: {e}." ++ logger.error(error_msg) ++ return False, error_msg ++ ++ def _apply_delta(self, paths: List[str]) -> tuple[bool, str]: ++ """Read + decompress delta safetensors files in parallel, decode + apply each.""" ++ import concurrent.futures ++ ++ n_files = len(paths) ++ workers = min(n_files, self.server_args.update_weight_delta_read_workers) ++ ++ def _read_and_decompress(path: str) -> bytes: ++ with open(path, "rb") as fh: ++ return _maybe_zstd_decompress(fh.read()) ++ ++ try: ++ # Cap peak memory at workers × file_size by applying each batch before ++ # prefetching the next. ++ for i in range(0, n_files, workers): ++ with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool: ++ batch = list(pool.map(_read_and_decompress, paths[i : i + workers])) ++ for blob in batch: ++ self._decode_and_apply_blob(blob) ++ return True, f"Applied {n_files} delta file(s)" ++ except Exception as e: ++ error_msg = f"Failed to apply delta update from disk: {e}." ++ logger.error(error_msg) ++ return False, error_msg ++ + def update_weights_from_tensor( + self, + named_tensors: List[Tuple[str, Union[torch.Tensor, "LocalSerializedTensor"]]], +@@ -2767,11 +2947,19 @@ class ModelRunner(ModelRunnerKVCacheMixin): output.expert_distribution_metrics = recorder_outputs.get("metrics") # Copy cached routing experts' buffers back to CPU cache @@ -1646,7 +1972,7 @@ index a59742b943..a7347c15b8 100644 if self.eplb_manager is not None: self.eplb_manager.on_forward_pass_end() -@@ -3021,6 +3035,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): +@@ -3021,6 +3209,161 @@ class ModelRunner(ModelRunnerKVCacheMixin): device=self.device, ) @@ -1685,12 +2011,131 @@ index a59742b943..a7347c15b8 100644 + quant_method.process_weights_after_loading(module) + + return True, "Success" ++ ++ ++def _param_storage_index(model): ++ """Build ``find_parent(dst)``: looks up the param/buffer owning ``dst``'s storage, ++ or None. Used by ``_delta_apply_context`` to scope its patched copy_/fill_.""" ++ import bisect ++ ++ starts: List[int] = [] ++ ends: List[int] = [] ++ owners: List[torch.Tensor] = [] ++ seen: set = set() ++ for tensors in (model.named_parameters(), model.named_buffers()): ++ for _, t in tensors: ++ if t.is_meta: ++ continue ++ try: ++ ptr = t.data_ptr() ++ except RuntimeError: ++ continue ++ if ptr == 0 or ptr in seen: ++ continue ++ seen.add(ptr) ++ sz = t.numel() * t.element_size() ++ starts.append(ptr) ++ ends.append(ptr + sz) ++ owners.append(t) ++ order = sorted(range(len(starts)), key=lambda i: starts[i]) ++ starts = [starts[i] for i in order] ++ ends = [ends[i] for i in order] ++ owners = [owners[i] for i in order] ++ ++ def find_parent(dst): ++ try: ++ ptr = dst.data_ptr() ++ except RuntimeError: ++ return None ++ idx = bisect.bisect_right(starts, ptr) - 1 ++ if 0 <= idx < len(starts) and starts[idx] <= ptr < ends[idx]: ++ return owners[idx] ++ return None ++ ++ return find_parent ++ ++ ++@contextlib.contextmanager ++def _delta_apply_context(model): ++ """Patch ``copy_`` / ``fill_`` so writes into ``model``'s param storage skip ++ positions whose source is NaN. Non-param writes go through unmodified. ++ ``post_load_weights`` runs in the original env so derived tensors (fp8 scales, ++ MoE biases, w_kc/w_vc) overwrite as usual.""" ++ is_param_target = _param_storage_index(model) ++ original_copy_ = torch.Tensor.copy_ ++ original_fill_ = torch.Tensor.fill_ ++ ++ def patched_copy_(self, src, *args, **kwargs): ++ if is_param_target(self) is not None: ++ src_aligned = ( ++ src.to(device=self.device, dtype=self.dtype) ++ if src.dtype != self.dtype ++ else src ++ ) ++ mask = ~torch.isnan(src_aligned) ++ self[mask] = src_aligned[mask] ++ return self ++ return original_copy_(self, src, *args, **kwargs) ++ ++ def patched_fill_(self, value): ++ if is_param_target(self) is not None: ++ # NaN scalar means "don't change the param" (per-element analog of ++ # patched_copy_). Non-NaN scalars write through. ++ try: ++ if math.isnan(value): ++ return self ++ except TypeError: ++ pass ++ return original_fill_(self, value) ++ return original_fill_(self, value) ++ ++ original_post_load = getattr(model, "post_load_weights", None) ++ if original_post_load is not None: ++ def wrapped_post_load(*args, **kwargs): ++ current_copy = torch.Tensor.copy_ ++ current_fill = torch.Tensor.fill_ ++ torch.Tensor.copy_ = original_copy_ ++ torch.Tensor.fill_ = original_fill_ ++ try: ++ return original_post_load(*args, **kwargs) ++ finally: ++ torch.Tensor.copy_ = current_copy ++ torch.Tensor.fill_ = current_fill ++ ++ model.post_load_weights = wrapped_post_load ++ ++ torch.Tensor.copy_ = patched_copy_ ++ torch.Tensor.fill_ = patched_fill_ ++ try: ++ yield ++ finally: ++ torch.Tensor.copy_ = original_copy_ ++ torch.Tensor.fill_ = original_fill_ ++ if original_post_load is not None: ++ model.post_load_weights = original_post_load ++ ++ ++def _delta_checksum(positions: torch.Tensor, values: torch.Tensor) -> int: ++ """Wire-corruption check, must match the sender's computation.""" ++ p = int(torch.hash_tensor(positions).item()) if positions.numel() else 0 ++ v = int(torch.hash_tensor(values).item()) if values.numel() else 0 ++ return p ^ (v << 1) ++ ++ ++def _maybe_zstd_decompress(blob: bytes) -> bytes: ++ """Decompress if zstd-framed (sender uses zstd when encoding=deltas_zstd).""" ++ # Zstandard frame magic: 0xFD2FB528 little-endian (RFC 8478 §3.1.1). ++ if blob.startswith(b"\x28\xb5\x2f\xfd"): ++ import zstandard ++ ++ return zstandard.ZstdDecompressor().decompress(blob) ++ return blob + def _model_load_weights_direct(model, named_tensors: List[Tuple[str, torch.Tensor]]): params_dict = dict(model.named_parameters()) diff --git a/python/sglang/srt/models/glm4v_moe.py b/python/sglang/srt/models/glm4v_moe.py -index 2f0074924d..1f991932c6 100644 +index 2f00749..1f99193 100644 --- a/python/sglang/srt/models/glm4v_moe.py +++ b/python/sglang/srt/models/glm4v_moe.py @@ -52,11 +52,31 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): @@ -1795,7 +2240,7 @@ index 2f0074924d..1f991932c6 100644 continue diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py -index 912891b6a7..fd67a7b580 100644 +index 912891b..fd67a7b 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -325,7 +325,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module): @@ -1808,7 +2253,7 @@ index 912891b6a7..fd67a7b580 100644 if ( diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py -index 7746b24459..57b65fe06f 100644 +index 7746b24..57b65fe 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py @@ -1005,14 +1005,19 @@ class Qwen3LLMModel(Qwen3Model): @@ -1836,7 +2281,7 @@ index 7746b24459..57b65fe06f 100644 positions, hidden_states, diff --git a/python/sglang/srt/multimodal/processors/glm4v.py b/python/sglang/srt/multimodal/processors/glm4v.py -index a44f14b6ca..6d6c65ea49 100644 +index a44f14b..6d6c65e 100644 --- a/python/sglang/srt/multimodal/processors/glm4v.py +++ b/python/sglang/srt/multimodal/processors/glm4v.py @@ -1,7 +1,13 @@ @@ -1904,7 +2349,7 @@ index a44f14b6ca..6d6c65ea49 100644 image_grid_thw = None video_grid_thw = None diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py -index 3f102567d0..6fb3899021 100644 +index 3f10256..6fb3899 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/python/sglang/srt/multimodal/processors/qwen_vl.py @@ -499,7 +499,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor): @@ -1917,7 +2362,7 @@ index 3f102567d0..6fb3899021 100644 image_data=image_data, video_data=request_obj.video_data, diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py -index 8caf21c320..51d1edc584 100644 +index 8caf21c..51d1edc 100644 --- a/python/sglang/srt/observability/req_time_stats.py +++ b/python/sglang/srt/observability/req_time_stats.py @@ -21,7 +21,10 @@ import uuid @@ -2157,7 +2602,7 @@ index 8caf21c320..51d1edc584 100644 def format_duration(self, duration: float) -> str: diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py -index ff5695ce2e..588379a85d 100644 +index ff5695c..588379a 100644 --- a/python/sglang/srt/observability/scheduler_metrics_mixin.py +++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py @@ -883,12 +883,42 @@ class SchedulerMetricsMixin: @@ -2204,7 +2649,7 @@ index ff5695ce2e..588379a85d 100644 def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index d91ced805f..4c8774bb64 100644 +index d91ced8..26f8ddf 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -670,6 +670,7 @@ class ServerArgs: @@ -2215,7 +2660,16 @@ index d91ced805f..4c8774bb64 100644 enable_fused_qk_norm_rope: bool = False enable_precise_embedding_interpolation: bool = False enable_fused_moe_sum_all_reduce: bool = False -@@ -5659,6 +5660,12 @@ class ServerArgs: +@@ -711,6 +712,8 @@ class ServerArgs: + # For model weight update and weight loading + custom_weight_loader: Optional[List[str]] = None + weight_loader_disable_mmap: bool = False ++ update_weight_delta_chunk_bytes: int = 512 * 1024 * 1024 ++ update_weight_delta_read_workers: int = 4 + remote_instance_weight_loader_seed_instance_ip: Optional[str] = None + remote_instance_weight_loader_seed_instance_service_port: Optional[int] = None + remote_instance_weight_loader_send_weights_group_ports: Optional[List[int]] = None +@@ -5659,6 +5662,12 @@ class ServerArgs: help="Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: 'round-robin-split'(default), 'in-seq-split' " "'round-robin-split' distributes tokens across ranks based on token_idx %% cp_size. It supports multi-batch prefill, fused MoE, and FP8 KV cache.", ) @@ -2228,8 +2682,27 @@ index d91ced805f..4c8774bb64 100644 parser.add_argument( "--enable-prefill-context-parallel", action="store_true", +@@ -5830,6 +5839,18 @@ class ServerArgs: + action="store_true", + help="Disable mmap while loading weight using safetensors.", + ) ++ parser.add_argument( ++ "--update-weight-delta-chunk-bytes", ++ type=int, ++ default=ServerArgs.update_weight_delta_chunk_bytes, ++ help="Byte cap per load_weights call when applying a delta update.", ++ ) ++ parser.add_argument( ++ "--update-weight-delta-read-workers", ++ type=int, ++ default=ServerArgs.update_weight_delta_read_workers, ++ help="Max parallel I/O threads for reading delta files from disk.", ++ ) + parser.add_argument( + "--remote-instance-weight-loader-seed-instance-ip", + type=str, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py -index 40e859b2d6..2604ae037c 100644 +index 40e859b..2604ae0 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -377,6 +377,10 @@ class EAGLEDraftCudaGraphRunner: @@ -2259,7 +2732,7 @@ index 40e859b2d6..2604ae037c 100644 buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py -index dbb91f555e..a04caefc34 100644 +index dbb91f5..a04caef 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -776,6 +776,10 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): @@ -2302,7 +2775,7 @@ index dbb91f555e..a04caefc34 100644 @dataclass diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py -index b0be70d751..44a78d684e 100644 +index b0be70d..44a78d6 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -2157,6 +2157,7 @@ class SafeUnpickler(pickle.Unpickler): @@ -2314,7 +2787,7 @@ index b0be70d751..44a78d684e 100644 DENY_CLASSES = { diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py -index 3be16446e0..1b2371c839 100644 +index 3be1644..1b2371c 100644 --- a/python/sglang/srt/utils/weight_checker.py +++ b/python/sglang/srt/utils/weight_checker.py @@ -69,6 +69,9 @@ def _check_tensors( diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 84deb1b7cf..f480dbbe08 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -34,6 +34,7 @@ from .model import forward_only, initialize_model_and_optimizer, save, train from .update_weight.common import named_params_and_buffers from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed +from .update_weight.update_weight_from_distributed_delta import UpdateWeightFromDistributedDelta from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor logging.getLogger("megatron").setLevel(logging.WARNING) @@ -135,7 +136,12 @@ def init( hf_vocab = getattr(self.hf_config, "vocab_size", None) self.args.vocab_size = hf_vocab if hf_vocab is not None else self.tokenizer.vocab_size - update_weight_cls = UpdateWeightFromTensor if self.args.colocate else UpdateWeightFromDistributed + if self.args.colocate: + update_weight_cls = UpdateWeightFromTensor + elif self.args.update_weight_mode == "delta": + update_weight_cls = UpdateWeightFromDistributedDelta + else: + update_weight_cls = UpdateWeightFromDistributed self.weight_updater = update_weight_cls( self.args, self.model, @@ -544,7 +550,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data logger.info(f"Updating ref model at rollout_id {rollout_id}") self.weights_backuper.backup("ref") - log_perf_data(rollout_id, self.args) + log_perf_data(rollout_id, self.args, extra_metrics=self.weight_updater.pop_metrics()) @timer def save_model(self, rollout_id: int, force_sync: bool = False) -> None: diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py index dee5e4705a..e5dd874e14 100644 --- a/slime/backends/megatron_utils/data.py +++ b/slime/backends/megatron_utils/data.py @@ -503,7 +503,7 @@ def log_passrate(rollout_id: int, args: Namespace, rollout_data: RolloutBatch) - gather_log_data("passrate", args, rollout_id, log_dict) -def log_perf_data(rollout_id: int, args: Namespace) -> None: +def log_perf_data(rollout_id: int, args: Namespace, extra_metrics: dict | None = None) -> None: train_metric_utils.log_perf_data_raw( rollout_id=rollout_id, args=args, @@ -515,6 +515,7 @@ def log_perf_data(rollout_id: int, args: Namespace) -> None: compute_total_fwd_flops=lambda seq_lens: calculate_fwd_flops(seqlens=seq_lens, args=args) / dist.get_world_size() / 1e12, + extra_metrics=extra_metrics, ) diff --git a/slime/backends/megatron_utils/sglang.py b/slime/backends/megatron_utils/sglang.py index 97c82a31cd..c42c9b45fe 100644 --- a/slime/backends/megatron_utils/sglang.py +++ b/slime/backends/megatron_utils/sglang.py @@ -13,6 +13,7 @@ from sglang.srt.patch_torch import monkey_patch_torch_reductions +from sglang.srt.managers.io_struct import DeltaEncoding, DeltaParam, DeltaSpec from sglang.srt.utils import MultiprocessingSerializer @@ -28,4 +29,7 @@ "monkey_patch_torch_reductions", "MultiprocessingSerializer", "FlattenedTensorBucket", + "DeltaEncoding", + "DeltaParam", + "DeltaSpec", ] diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py index 822b801776..1aa4db8bcf 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py @@ -1,7 +1,7 @@ import socket import time from argparse import Namespace -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence import ray import torch @@ -14,6 +14,7 @@ from slime.utils.distributed_utils import get_gloo_group, init_process_group from ..megatron_to_hf import convert_to_hf +from ..sglang import DeltaSpec from .common import all_gather_param, named_params_and_buffers @@ -21,6 +22,7 @@ class UpdateWeightFromDistributed: """ Update distributed engines via NCCL. Each PP rank: group "slime-pp_{pp_rank}", only DP=TP=0 broadcasts. Non-expert (TP) and expert (EP) params separate. + Subclasses override ``_send_weights`` / ``_on_chunk`` to inject per-mode behaviour. """ def __init__( @@ -41,6 +43,14 @@ def __init__( self.quantization_config = quantization_config self.weight_version = 0 self._model_update_groups = None + self.update_weight_metrics: dict[str, float] = {} + + def pop_metrics(self) -> dict[str, float]: + """ + Return and clear ``update_weight_metrics``. Drained by the actor onto the rollout/step log. + """ + out, self.update_weight_metrics = self.update_weight_metrics, {} + return out def connect_rollout_engines( self, @@ -89,7 +99,7 @@ def disconnect_rollout_engines(self) -> None: @torch.no_grad() def update_weights(self) -> None: """ - Pause → flush → non-expert (TP) → expert (EP) → continue. Progress on PP source. + Pause → flush → _send_weights → continue. Progress on PP source. """ self.weight_version += 1 @@ -106,36 +116,9 @@ def update_weights(self) -> None: ) dist.barrier(group=get_gloo_group()) - buffer_size = 0 - converted_named_tensors = [] - # non expert params pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None + self._send_weights(pbar) - for name, param in named_params_and_buffers(self.args, self.model): - if ".experts." in name: - continue - buffer_size = self._update_weight_from_distributed( - name, param, converted_named_tensors, buffer_size, pbar=pbar - ) - - if converted_named_tensors: - self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar) - - dist.barrier(group=get_gloo_group()) - - buffer_size = 0 - named_tensors = [] - for name, param in named_params_and_buffers(self.args, self.model): - if ".experts." not in name: - continue - buffer_size = self._update_expert_weight_from_distributed( - name, param, named_tensors, buffer_size, pbar=pbar - ) - - if named_tensors: - self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar) - - dist.barrier(group=get_gloo_group()) if dist.get_rank() == 0: # int4/fp4 post_process if self.quantization_config and self.quantization_config["quant_method"] in ["compressed-tensors"]: @@ -147,59 +130,84 @@ def update_weights(self) -> None: ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) - def _update_weight_from_distributed( - self, - name: str, - param: torch.nn.Parameter, - converted_named_tensors: list[tuple[str, torch.Tensor]], - buffer_size: int, - pbar: tqdm | None = None, - ) -> int | None: + def _send_weights(self, pbar: tqdm | None) -> None: """ - Non-expert: gather TP → rm pad → HF → buffer (flush if full). All gather, PP source buffers. - Returns updated bytes on source, None on non-source. + Non-expert (TP) pass → barrier → expert (EP) pass → barrier. Each iterator + yields broadcast-ready chunks (bucketing happens internally); subclasses + override ``_on_chunk`` to inject per-chunk behaviour. """ - param = all_gather_param(name, param) - if not self._is_pp_src_rank: - return + for chunk_iter in (self._iter_non_expert_chunks(), self._iter_expert_chunks()): + for hf_chunk in chunk_iter: + self._on_chunk(hf_chunk) + self._update_bucket_weights_from_distributed(hf_chunk, pbar=pbar) + dist.barrier(group=get_gloo_group()) - param_size = param.numel() * param.element_size() - if buffer_size + param_size > self.args.update_weight_buffer_size: - self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar) - buffer_size = 0 - converted_named_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config) - buffer_size += param_size - return buffer_size + def _on_chunk(self, hf_chunk: list[tuple[str, torch.Tensor]]) -> None: + """ + Hook for each HF chunk in ``_send_weights`` before its broadcast. No-op by default. + """ - def _update_expert_weight_from_distributed( + def _iter_non_expert_chunks(self) -> Iterator[list[tuple[str, torch.Tensor]]]: + """ + Yield broadcast-sized HF chunks of non-expert params: TP all-gather + + HF convert per param, then bucket up to ``--update-weight-buffer-size``. + Empty on non-PP-src ranks (they still join all_gather_param). + """ + buffer_size = 0 + buffer: list[tuple[str, torch.Tensor]] = [] + for name, param in named_params_and_buffers(self.args, self.model): + if ".experts." in name: + continue + param = all_gather_param(name, param) + if not self._is_pp_src_rank: + continue + hf_chunk = convert_to_hf(self.args, self.model_name, name, param, self.quantization_config) + chunk_bytes = sum(t.numel() * t.element_size() for _, t in hf_chunk) + if buffer and buffer_size + chunk_bytes > self.args.update_weight_buffer_size: + yield buffer + buffer = [] + buffer_size = 0 + buffer.extend(hf_chunk) + buffer_size += chunk_bytes + if buffer: + yield buffer + + def _iter_expert_chunks( self, - name: str, - param: torch.nn.Parameter, - named_tensors: list[tuple[str, torch.Tensor]], - buffer_size: int, - pbar: tqdm | None = None, - ) -> int: + params: Iterator[tuple[str, torch.Tensor]] | None = None, + ) -> Iterator[list[tuple[str, torch.Tensor]]]: """ - Expert: gather TP → rm pad → buffer. EP gather + HF deferred. Threshold × EP size. + Yield one HF chunk per EP-weighted batch of expert params: TP gather + + buffer until threshold, then EP gather + HF convert. ``params`` lets + callers restrict the iter to a subset (used by delta-sync sub-passes); + defaults to all expert params on this rank. """ - param = all_gather_param(name, param) - - param_size = param.numel() * param.element_size() - if ( - buffer_size + param_size - ) * mpu.get_expert_model_parallel_world_size() > self.args.update_weight_buffer_size: - self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar) - buffer_size = 0 - - named_tensors.append((name, param)) - buffer_size += param_size - return buffer_size - - def _update_expert_bucket_weights_from_distributed( - self, named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None - ) -> None: + if params is None: + params = ((n, p) for n, p in named_params_and_buffers(self.args, self.model) if ".experts." in n) + buffer_size = 0 + batch: list[tuple[str, torch.Tensor]] = [] + for name, param in params: + param = all_gather_param(name, param) + param_size = param.numel() * param.element_size() + if ( + buffer_size + param_size + ) * mpu.get_expert_model_parallel_world_size() > self.args.update_weight_buffer_size: + hf_chunk = self._ep_gather_and_convert(batch) + if hf_chunk: + yield hf_chunk + batch = [] + buffer_size = 0 + batch.append((name, param)) + buffer_size += param_size + if batch: + hf_chunk = self._ep_gather_and_convert(batch) + if hf_chunk: + yield hf_chunk + + def _ep_gather_and_convert(self, named_tensors: list[tuple[str, torch.Tensor]]) -> list[tuple[str, torch.Tensor]]: """ - Gather EP → HF → broadcast. Clears buffer. + EP all-gather a buffered batch + HF convert on PP source. Returns HF tensors on + PP source, [] elsewhere. Clears ``named_tensors``. """ names = [name for name, _ in named_tensors] all_names = [None] * mpu.get_expert_model_parallel_world_size() @@ -224,20 +232,25 @@ def _update_expert_bucket_weights_from_distributed( named_tensors.clear() if not self._is_pp_src_rank: - return + return [] all_gathered_params = sum(all_gathered_params, []) converted_hf_tensors = [] for name, param in all_gathered_params: converted_hf_tensors += convert_to_hf(self.args, self.model_name, name, param, self.quantization_config) - - self._update_bucket_weights_from_distributed(converted_hf_tensors, pbar) + return converted_hf_tensors def _update_bucket_weights_from_distributed( - self, converted_named_tensors: list[tuple[str, torch.Tensor]], pbar: tqdm | None = None + self, + converted_named_tensors: list[tuple[str, torch.Tensor]], + pbar: tqdm | None = None, + load_format: str | None = None, + delta: DeltaSpec | None = None, ) -> None: """ Lock → broadcast → clear → unlock → pbar++. Lock prevents NCCL deadlock. + Delta sync passes ``load_format="delta"`` + a ``DeltaSpec`` describing the + per-param decoding of the (__positions__, __values__) bucket tensors. """ # lock the rollout engines to prevent dead lock on broadcast. while not ray.get(self.rollout_engine_lock.acquire.remote()): @@ -249,6 +262,8 @@ def _update_bucket_weights_from_distributed( self.weight_version, self.rollout_engines, converted_named_tensors, + load_format=load_format, + delta=delta, ) ray.get(refs) @@ -321,9 +336,12 @@ def update_weights_from_distributed( weight_version: int, rollout_engines: Sequence[ActorHandle], converted_named_tensors: Sequence[tuple[str, torch.Tensor]], + load_format: str | None = None, + delta: DeltaSpec | None = None, ) -> list[ObjectRef]: """ Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines). + Delta sync passes ``load_format="delta"`` + ``delta`` (DeltaSpec). """ refs = [ engine.update_weights_from_distributed.remote( @@ -332,6 +350,8 @@ def update_weights_from_distributed( shapes=[param.shape for _, param in converted_named_tensors], group_name=group_name, weight_version=str(weight_version), + load_format=load_format, + delta=delta, ) for engine in rollout_engines ] diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed_delta.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed_delta.py new file mode 100644 index 0000000000..65a7a5a793 --- /dev/null +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed_delta.py @@ -0,0 +1,856 @@ +""" +Delta weight sync. + +For each sync, the sender bytewise-diffs the current weights against a +pinned-CPU snapshot of the last broadcast, packs the changed positions +and values, and ships them via one of two transports: + + - "nccl": each bucket flush goes out via NCCL broadcast (low-latency, + high-bandwidth, intra-datacenter). + - "disk": each bucket flush is written to a versioned shared-FS directory + as one safetensors file; one HTTP push per sync wakes the rollout + engines to read+apply (cross-datacenter, bandwidth-limited). + +Both transports share one wire layout (``__positions__`` uint8 byte blob + +``__values__`` param-dtype tensor + per-param decoding manifest) and one +receiver-side decoder. Three encodings differ only in how positions are +packed: + + indices : int32 absolute positions + deltas : uint16 gap-deltas (uint32 fallback per param) + deltas_zstd : ``deltas`` with the safetensors blob wrapped in zstd L1 + +The receiver overwrites changed positions with the trainer's exact bytes +(no arithmetic), so the apply is lossless and there is no drift to fight +with periodic re-syncs. The first ``update_weights`` call seeds the +snapshot without contacting the rollout engines — they're assumed to have +loaded the same HF checkpoint at init. +""" + +import itertools +import json +import logging +import os +import shutil +import threading +from argparse import Namespace +from collections.abc import Callable, Iterator, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import asdict, dataclass, field, replace +from queue import Queue + +import numpy as np +import ray +import torch +import torch.distributed as dist +from megatron.core import mpu +from ray.actor import ActorHandle +from safetensors.torch import save as st_save_bytes +from tqdm import tqdm + +from slime.utils.distributed_utils import get_gloo_group +from slime.utils.timer import Timer, timer + +from ..sglang import DeltaEncoding, DeltaParam, DeltaSpec +from .update_weight_from_distributed import UpdateWeightFromDistributed + + +logger = logging.getLogger(__name__) + + +# ---------- compute + encode ----------------------------------------------- + + +@dataclass +class ParamDiff: + """ + One per-param compute output. ``values`` is a reference to the full-shape + current tensor (no copy); ``mask`` is a same-shape bool marking the + positions whose bytes differ from the snapshot. + """ + + name: str + values: torch.Tensor + mask: torch.Tensor + + +@dataclass +class EncodedChunk: + """ + One HF chunk after position+value encoding, before bucket merging. + + ``pos_bytes`` and ``val_tensor`` are the chunk-local concatenations across + all params; per-param byte/element offsets live on ``params``. + """ + + pos_bytes: bytes + val_tensor: torch.Tensor + params: list[DeltaParam] + nnz: int + + @classmethod + def empty(cls) -> "EncodedChunk": + return cls(pos_bytes=b"", val_tensor=torch.empty(0, dtype=torch.bfloat16), params=[], nnz=0) + + +def _checksum(positions: torch.Tensor, values: torch.Tensor) -> int: + """ + Wire-corruption check via ``torch.hash_tensor`` (XOR-reduce over uint64 bitcast). + Sender computes pre-flush, receiver computes post-recv; mismatch indicates + corruption between encode and apply. One reduction + one ``.item()`` sync per arg. + """ + p = int(torch.hash_tensor(positions).item()) if positions.numel() else 0 + v = int(torch.hash_tensor(values).item()) if values.numel() else 0 + return p ^ (v << 1) + + +def _bytewise_diff_mask(current: torch.Tensor, snapshot: torch.Tensor) -> torch.Tensor: + """ + Per-element bool mask: True where current and snapshot bytes differ. Dtype-agnostic via view-as-integer. + """ + es = current.element_size() + int_dtype = {1: torch.uint8, 2: torch.int16, 4: torch.int32, 8: torch.int64}.get(es) + if int_dtype is None: + raise ValueError(f"unsupported element size {es}") + return current.view(int_dtype) != snapshot.view(int_dtype) + + +def _sparse_boundaries( + diffs: list[ParamDiff], +) -> tuple[torch.Tensor, list[int], torch.Tensor, list[int]]: + """ + One concat → one nonzero → one searchsorted → one ``tolist()``: collapses + per-param host syncs to one per chunk. Returns ``(big_val, bounds, big_idx, cum)``. + """ + device = diffs[0].values.device + sizes = [d.values.numel() for d in diffs] + cum = list(itertools.accumulate(sizes)) + cum_t = torch.tensor(cum, dtype=torch.int64, device=device) + + big_values = torch.cat([d.values.contiguous().view(-1) for d in diffs], dim=0) + big_mask = torch.cat([d.mask.contiguous().view(-1) for d in diffs], dim=0) + big_idx = big_mask.nonzero(as_tuple=False).view(-1) + big_val = big_values[big_idx] + bounds = torch.searchsorted(big_idx, cum_t).tolist() + return big_val, bounds, big_idx, cum + + +def encode_indices(diffs: list[ParamDiff]) -> EncodedChunk: + """ + int32 absolute positions, per-param. Position blob is uint8 bytes; pos_width=4 for all params. + """ + if not diffs: + return EncodedChunk.empty() + big_val, bounds, big_idx, cum = _sparse_boundaries(diffs) + pos_pieces: list[torch.Tensor] = [] + val_pieces: list[torch.Tensor] = [] + params: list[DeltaParam] = [] + pos_byte_off = val_off = 0 + prev_b = 0 + prev_param_start = 0 + for i, d in enumerate(diffs): + b = bounds[i] + nnz = b - prev_b + if nnz > 0: + local_idx = (big_idx[prev_b:b] - prev_param_start).to(torch.int32) + pos_pieces.append(local_idx) + val_pieces.append(big_val[prev_b:b]) + params.append( + DeltaParam( + name=d.name, + dtype=str(d.values.dtype).replace("torch.", ""), + shape=list(d.values.shape), + pos_start=pos_byte_off, + pos_end=pos_byte_off + nnz * 4, + pos_width=4, + val_start=val_off, + val_end=val_off + nnz, + ) + ) + pos_byte_off += nnz * 4 + val_off += nnz + prev_b = b + prev_param_start = cum[i] + if not params: + return EncodedChunk.empty() + positions = torch.cat(pos_pieces, dim=0) + values = torch.cat(val_pieces, dim=0) + return EncodedChunk( + pos_bytes=positions.cpu().numpy().tobytes(), + val_tensor=values, + params=params, + nnz=val_off, + ) + + +def encode_deltas(diffs: list[ParamDiff]) -> EncodedChunk: + """ + Gap-encode sorted positions: store ``idx[k] - idx[k-1] - 1`` with idx[-1] := -1 + so the first delta equals the first index. Per-param downcast to uint16 if the max + gap fits, otherwise uint32. At ~2% Bernoulli density on bf16 weights, max gap ≈ 300 + — uint16 fits; the fallback covers pathological inputs without correctness risk. + Receiver inverts: ``idx = cumsum(delta + 1) - 1``. + """ + if not diffs: + return EncodedChunk.empty() + big_val, bounds, big_idx, cum = _sparse_boundaries(diffs) + + kept: list[tuple[ParamDiff, int]] = [] # (diff, nnz) for non-empty params + per_param_deltas: list[torch.Tensor] = [] + val_pieces: list[torch.Tensor] = [] + prev_b = 0 + prev_param_start = 0 + for i, d in enumerate(diffs): + b = bounds[i] + nnz = b - prev_b + if nnz > 0: + local_idx = big_idx[prev_b:b] - prev_param_start # int64, sorted + prev = torch.cat( + [ + torch.tensor([-1], dtype=local_idx.dtype, device=local_idx.device), + local_idx[:-1], + ] + ) + per_param_deltas.append(local_idx - prev - 1) + val_pieces.append(big_val[prev_b:b]) + kept.append((d, nnz)) + prev_b = b + prev_param_start = cum[i] + + if not kept: + return EncodedChunk.empty() + + # One CPU sync for per-param width selection. + max_per_param = torch.stack([d.max() for d in per_param_deltas]).cpu().tolist() + pos_byte_pieces: list[bytes] = [] + pos_byte_off = val_off = 0 + params: list[DeltaParam] = [] + for (d, nnz), deltas, max_d in zip(kept, per_param_deltas, max_per_param, strict=True): + width = 2 if int(max_d) <= 65535 else 4 + np_dtype = np.uint16 if width == 2 else np.uint32 + b_chunk = deltas.cpu().numpy().astype(np_dtype, copy=False).tobytes() + pos_byte_pieces.append(b_chunk) + params.append( + DeltaParam( + name=d.name, + dtype=str(d.values.dtype).replace("torch.", ""), + shape=list(d.values.shape), + pos_start=pos_byte_off, + pos_end=pos_byte_off + len(b_chunk), + pos_width=width, + val_start=val_off, + val_end=val_off + nnz, + ) + ) + pos_byte_off += len(b_chunk) + val_off += nnz + + values = torch.cat(val_pieces, dim=0) + return EncodedChunk( + pos_bytes=b"".join(pos_byte_pieces), + val_tensor=values, + params=params, + nnz=val_off, + ) + + +# ---------- snapshot state ------------------------------------------------- + + +class DeltaState: + """ + Pinned-CPU snapshot of every HF tensor we've broadcast, plus the H2D/D2H + side streams that pipeline next-chunk snapshot transfer behind the current + chunk's compute. + """ + + def __init__(self) -> None: + self.snapshot: dict[str, torch.Tensor] = {} + self.d2h_stream: torch.cuda.Stream | None = None + self.h2d_stream: torch.cuda.Stream | None = None + self.snapshot_dirty = False + + def prefetch_snapshot( + self, named_tensors: list[tuple[str, torch.Tensor]] + ) -> tuple[list[torch.Tensor], torch.cuda.Event]: + """ + Start an async H2D copy of the snapshot tensors for ``named_tensors`` on a side stream. + """ + if self.h2d_stream is None: + self.h2d_stream = torch.cuda.Stream() + prev_gpu: list[torch.Tensor] = [] + with torch.cuda.stream(self.h2d_stream): + for name, tensor in named_tensors: + if name not in self.snapshot: + raise KeyError(f"missing snapshot for {name!r}; first update_weights call seeds the snapshot") + prev_gpu.append(self.snapshot[name].to(device=tensor.device, non_blocking=True)) + event = self.h2d_stream.record_event() + return prev_gpu, event + + def compute_diffs( + self, + named_tensors: list[tuple[str, torch.Tensor]], + prefetched: tuple[list[torch.Tensor], torch.cuda.Event], + ) -> list[ParamDiff]: + """ + Wait for the prefetched H2D copy, then per-param bytewise diff against the snapshot. + """ + prev_gpu, event = prefetched + event.wait() + return [ + ParamDiff(name=name, values=current, mask=_bytewise_diff_mask(current, prev)) + for (name, current), prev in zip(named_tensors, prev_gpu, strict=True) + ] + + def update_snapshot_async(self, named_tensors: list[tuple[str, torch.Tensor]]) -> None: + """ + Enqueue a D2H copy of ``named_tensors`` into the pinned-CPU snapshot on a + side stream. Non-blocking; call ``flush_snapshot`` before the next sync. + """ + if self.d2h_stream is None: + self.d2h_stream = torch.cuda.Stream() + event = torch.cuda.current_stream().record_event() + with torch.cuda.stream(self.d2h_stream): + self.d2h_stream.wait_event(event) + for name, tensor in named_tensors: + if name not in self.snapshot: + self.snapshot[name] = torch.empty_like(tensor, device=torch.device("cpu"), pin_memory=True) + self.snapshot[name].copy_(tensor.detach(), non_blocking=True) + self.snapshot_dirty = True + + def flush_snapshot(self) -> None: + """ + Block until all enqueued D2H snapshot copies have landed. + """ + if self.snapshot_dirty: + if self.d2h_stream is not None: + self.d2h_stream.synchronize() + else: + torch.cuda.synchronize() + self.snapshot_dirty = False + + +# ---------- bucket --------------------------------------------------------- + + +@dataclass +class DeltaBucket: + """ + Accumulates encoded chunks for one flush. Per-param offsets are rebased + into the bucket's growing position blob + value tensor on ``add``. + """ + + pos_pieces: list[bytes] = field(default_factory=list) + val_pieces: list[torch.Tensor] = field(default_factory=list) + params: list[DeltaParam] = field(default_factory=list) + pos_total: int = 0 + val_total: int = 0 + byte_size: int = 0 + + @property + def has_updates(self) -> bool: + return bool(self.pos_pieces) + + def should_flush_before_add(self, chunk: EncodedChunk, byte_limit: int) -> bool: + """True iff adding ``chunk`` would push the bucket past ``byte_limit``.""" + chunk_bytes = len(chunk.pos_bytes) + chunk.val_tensor.numel() * chunk.val_tensor.element_size() + return self.has_updates and self.byte_size + chunk_bytes > byte_limit + + def add(self, chunk: EncodedChunk) -> None: + """Append ``chunk``, rebasing each param's byte/element offsets into the bucket.""" + for p in chunk.params: + self.params.append( + replace( + p, + pos_start=p.pos_start + self.pos_total, + pos_end=p.pos_end + self.pos_total, + val_start=p.val_start + self.val_total, + val_end=p.val_end + self.val_total, + ) + ) + self.pos_pieces.append(chunk.pos_bytes) + self.val_pieces.append(chunk.val_tensor) + self.pos_total += len(chunk.pos_bytes) + self.val_total += chunk.val_tensor.numel() + self.byte_size += len(chunk.pos_bytes) + chunk.val_tensor.numel() * chunk.val_tensor.element_size() + + def merged_positions_cpu(self) -> torch.Tensor: + """One CPU uint8 tensor with the bucket's positions blob.""" + merged = b"".join(self.pos_pieces) + if not merged: + return torch.empty(0, dtype=torch.uint8) + return torch.from_numpy(np.frombuffer(merged, dtype=np.uint8).copy()) + + def merged_values(self) -> torch.Tensor: + """One GPU tensor with the bucket's values, concatenated across chunks.""" + if not self.val_pieces: + return torch.empty(0, dtype=torch.bfloat16) + return torch.cat(self.val_pieces, dim=0) + + def clear(self) -> None: + """Reset to empty so the bucket can be reused for the next flush.""" + self.pos_pieces.clear() + self.val_pieces.clear() + self.params.clear() + self.pos_total = 0 + self.val_total = 0 + self.byte_size = 0 + + +# ---------- async safetensors writer (disk transport only) ----------------- + + +class AsyncSafetensorsWriter: + """ + Background thread that drains a queue of file writes. Producers do GPU→CPU + on the default stream and enqueue; the writer does the slow disk I/O + (and optional zstd compress) off the critical path. End-of-sync ``drain()`` + blocks until all enqueued writes have landed. + """ + + def __init__(self, compress_with_zstd: bool, zstd_level: int = 1) -> None: + self._queue: Queue = Queue() + self._error: BaseException | None = None + self._compress_with_zstd = compress_with_zstd + self._zstd_level = zstd_level + if compress_with_zstd: + # Lazy import — non-disk users don't pay the dep. + import zstandard + + self._zstd = zstandard + self._lock = threading.Lock() + self.bytes_pre_compress = 0 + self.bytes_post_compress = 0 + self._thread = threading.Thread(target=self._run, name="delta-disk-writer", daemon=True) + self._thread.start() + + def enqueue( + self, + path: str, + tensors: dict[str, torch.Tensor], + metadata: dict[str, str], + ) -> None: + """Hand a (path, tensors, metadata) tuple to the writer thread.""" + if self._error is not None: + raise RuntimeError(f"writer thread already failed: {self._error!r}") + self._queue.put((path, tensors, metadata)) + + def drain(self) -> None: + """Block until every queued write has landed; re-raise any writer-thread error.""" + self._queue.join() + if self._error is not None: + raise RuntimeError(f"writer thread failed: {self._error!r}") from self._error + + def reset_counters(self) -> None: + """Zero the byte counters at the start of a sync.""" + with self._lock: + self.bytes_pre_compress = 0 + self.bytes_post_compress = 0 + + def _run(self) -> None: + """Writer-thread loop: safetensors-encode → (optional zstd) → atomic replace.""" + cctx = self._zstd.ZstdCompressor(level=self._zstd_level, threads=-1) if self._compress_with_zstd else None + while True: + path, tensors, metadata = self._queue.get() + try: + if self._error is None: + blob = st_save_bytes(tensors, metadata=metadata) + pre = len(blob) + if cctx is not None: + blob = cctx.compress(blob) + post = len(blob) + tmp = path + ".tmp" + with open(tmp, "wb") as f: + f.write(blob) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, path) + with self._lock: + self.bytes_pre_compress += pre + self.bytes_post_compress += post + except BaseException as e: # noqa: BLE001 + self._error = e + finally: + self._queue.task_done() + + +# ---------- main class ----------------------------------------------------- + + +class UpdateWeightFromDistributedDelta(UpdateWeightFromDistributed): + """ + Selective delta sync. ``--update-weight-transport`` picks the per-flush carrier: + "nccl" broadcasts each bucket; "disk" writes each bucket as a safetensors file under + ``--update-weight-delta-dir`` and pushes once at end-of-sync. + """ + + def __init__( + self, + args: Namespace, + model: Sequence[torch.nn.Module], + weights_getter: Callable[[], Mapping[str, torch.Tensor]], + *, + model_name: str, + quantization_config: dict[str, int | str | list[str]] | None, + ) -> None: + super().__init__( + args, + model, + weights_getter, + model_name=model_name, + quantization_config=quantization_config, + ) + self.transport = args.update_weight_transport + self.encoding = DeltaEncoding(args.update_weight_encoding) + self.delta_state = DeltaState() + self._snapshot_seeded = False + # DELTAS_ZSTD shares the gap encoder; zstd is applied at file-write time. + self._encode = encode_indices if self.encoding is DeltaEncoding.INDICES else encode_deltas + + self.writer: AsyncSafetensorsWriter | None = None + self.delta_dir: str | None = None + self._pre_push_hook: Callable | None = None + # Disk transport: each pass boundary publishes its accumulated files + # (the only globally-synced flush points, since ``_publish_batch`` + # contains collectives). ``_pre_push_hook`` may return a Future, in + # which case the receiver RPC is deferred behind it via + # ``_rpc_executor`` so the main encode thread continues immediately. + # ``_pending_publishes`` holds the resulting Future[list[ObjectRef]] + # on rank 0; ``_finalize_sync`` awaits them at end of sync. + self._pending_files: list[str] = [] + self._pending_publishes: list = [] + self._rpc_executor: ThreadPoolExecutor | None = None + if self.transport == "disk": + self.delta_dir = args.update_weight_delta_dir + os.makedirs(self.delta_dir, exist_ok=True) + self.writer = AsyncSafetensorsWriter( + compress_with_zstd=(self.encoding == DeltaEncoding.DELTAS_ZSTD), + ) + self._rpc_executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="delta-publish-rpc") + if getattr(args, "custom_delta_pre_push_path", None): + from slime.utils.misc import load_function + + self._pre_push_hook = load_function(args.custom_delta_pre_push_path) + + def connect_rollout_engines( + self, + rollout_engines: Sequence[ActorHandle], + rollout_engine_lock: ActorHandle, + engine_gpu_counts: Sequence[int] | None = None, + engine_gpu_offsets: Sequence[int] | None = None, + ) -> None: + """ + NCCL transport: delegate to parent (group creation). Disk transport: just + record the engines + PP-src flag (no NCCL group needed). + """ + if self.transport == "nccl": + super().connect_rollout_engines( + rollout_engines, + rollout_engine_lock, + engine_gpu_counts=engine_gpu_counts, + engine_gpu_offsets=engine_gpu_offsets, + ) + return + self.rollout_engines = rollout_engines + self.rollout_engine_lock = rollout_engine_lock + self._engine_gpu_counts = engine_gpu_counts + self._is_pp_src_rank = ( + mpu.get_data_parallel_rank(with_context_parallel=True) == 0 and mpu.get_tensor_model_parallel_rank() == 0 + ) + pp_rank = mpu.get_pipeline_model_parallel_rank() + self._group_name = f"slime-pp_{pp_rank}" + + def disconnect_rollout_engines(self) -> None: + if self.transport == "nccl": + super().disconnect_rollout_engines() + + @torch.no_grad() + def update_weights(self) -> None: + """ + First call: seed the CPU snapshot from current model state, no engine RPCs. + Subsequent calls: pause → diff/encode → finalize → resume. ``delta_encode`` + covers the sender's per-param TP/EP gather + diff + sparse encode + per-publish + commit/RPC handoff; ``delta_finalize`` covers the tail wait for the last + batch's receiver-apply. Their sum is the sync latency the user observes. + """ + if not self._snapshot_seeded: + self._seed_snapshot() + self._snapshot_seeded = True + return + + self.weight_version += 1 + if self.transport == "disk": + self._version_dir = os.path.join(self.delta_dir, f"weight_v{self.weight_version:06d}") + if self._is_pp_src_rank: + os.makedirs(self._version_dir, exist_ok=True) + + if dist.get_rank() == 0: + ray.get([engine.pause_generation.remote() for engine in self.rollout_engines]) + ray.get([engine.flush_cache.remote() for engine in self.rollout_engines]) + dist.barrier(group=get_gloo_group()) + + self.density_nnz = self.density_numel = self.wire_bytes = self._flush_idx = 0 + self._pending_files.clear() + self._pending_publishes.clear() + if self.writer is not None: + self.writer.reset_counters() + pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None + + with timer("delta_encode"): + self._send_weights(pbar) + if self.writer is not None: + self.writer.drain() + self.delta_state.flush_snapshot() + dist.barrier(group=get_gloo_group()) + + with timer("delta_finalize"): + self._finalize_sync() + + self._record_metrics() + + def _seed_snapshot(self) -> None: + """ + Populate the snapshot from current model state (TP/EP gather + HF + convert on PP-src ranks, D2H pinned copy). Cost is one full pass over + params — ~50s blocking on 355B at init. + """ + for chunk_iter in (self._iter_non_expert_chunks(), self._iter_expert_chunks()): + for hf_chunk in chunk_iter: + if hf_chunk: + self.delta_state.update_snapshot_async(hf_chunk) + dist.barrier(group=get_gloo_group()) + self.delta_state.flush_snapshot() + + def _send_weights(self, pbar: tqdm | None) -> None: + """ + Non-expert pass then expert pass, each followed by a barrier + (disk-only) + publish. The expert pass is split into ``_EXPERT_SUBPASSES`` sub-passes so + receiver apply for an earlier batch overlaps with later expert encoding, + instead of bottlenecking at end-of-sync. Megatron splits MoE layers + uniformly across PP ranks, so a per-rank slice of the expert param list + keeps the publish count identical on every rank (no barrier desync). + """ + from .common import named_params_and_buffers + + bucket = DeltaBucket() + self._pipeline_pass(self._iter_non_expert_chunks(), bucket, pbar) + self._flush_and_publish(bucket, pbar) + + expert_params = [(n, p) for n, p in named_params_and_buffers(self.args, self.model) if ".experts." in n] + n = len(expert_params) + for i in range(self._EXPERT_SUBPASSES): + lo = i * n // self._EXPERT_SUBPASSES + hi = (i + 1) * n // self._EXPERT_SUBPASSES + self._pipeline_pass(self._iter_expert_chunks(iter(expert_params[lo:hi])), bucket, pbar) + self._flush_and_publish(bucket, pbar) + + _EXPERT_SUBPASSES = 4 + + def _flush_and_publish(self, bucket: DeltaBucket, pbar: tqdm | None) -> None: + """ + End-of-sub-pass: drain the in-flight bucket, barrier all PP ranks, then + (disk-only) fire one publish RPC for everything since the last call. + """ + if bucket.has_updates: + self._flush_bucket(bucket, pbar) + dist.barrier(group=get_gloo_group()) + if self.transport == "disk": + self._publish_batch() + + def _pipeline_pass( + self, + chunk_iter: Iterator[list[tuple[str, torch.Tensor]]], + bucket: DeltaBucket, + pbar: tqdm | None, + ) -> None: + """ + 1-step H2D snapshot prefetch lookahead: chunk N+1's snapshot transfer + overlaps chunk N's compute+encode on the default stream. + """ + pending_chunk: list[tuple[str, torch.Tensor]] | None = None + pending_prefetch: tuple[list[torch.Tensor], torch.cuda.Event] | None = None + for hf_chunk in chunk_iter: + if not hf_chunk: + continue + next_prefetch = self.delta_state.prefetch_snapshot(hf_chunk) + if pending_prefetch is not None: + self._enqueue_chunk(pending_chunk, pending_prefetch, bucket, pbar) + pending_chunk, pending_prefetch = hf_chunk, next_prefetch + if pending_prefetch is not None: + self._enqueue_chunk(pending_chunk, pending_prefetch, bucket, pbar) + + def _enqueue_chunk( + self, + hf_chunk: list[tuple[str, torch.Tensor]], + prefetched: tuple[list[torch.Tensor], torch.cuda.Event], + bucket: DeltaBucket, + pbar: tqdm | None, + ) -> None: + """ + compute diffs → snapshot new prev → encode → bucket.add (flushing if full). + """ + diffs = self.delta_state.compute_diffs(hf_chunk, prefetched=prefetched) + self.delta_state.update_snapshot_async(hf_chunk) + chunk = self._encode(diffs) + self.density_numel += sum(d.values.numel() for d in diffs) + self.density_nnz += chunk.nnz + self.wire_bytes += len(chunk.pos_bytes) + chunk.val_tensor.numel() * chunk.val_tensor.element_size() + if not chunk.params: + return + if bucket.should_flush_before_add(chunk, self.args.update_weight_buffer_size): + self._flush_bucket(bucket, pbar) + bucket.add(chunk) + + def _flush_bucket(self, bucket: DeltaBucket, pbar: tqdm | None) -> None: + """ + NCCL: broadcast (__positions__, __values__) with a DeltaSpec. + Disk: enqueue one safetensors file with the same payload + metadata. + Both paths embed a checksum the receiver verifies before apply. + """ + if not bucket.has_updates: + return + positions_cpu = bucket.merged_positions_cpu() + values_gpu = bucket.merged_values() + params = list(bucket.params) + bucket.clear() + + # GPU-resident checksum: positions go to the device the values already live on + # (NCCL needs the same move anyway; disk gets it for free at the reduction). + positions_gpu = positions_cpu.to(values_gpu.device, non_blocking=True) + checksum = _checksum(positions_gpu, values_gpu) + + if self.transport == "nccl": + spec = DeltaSpec(encoding=self.encoding, params=params, checksum=checksum) + self._update_bucket_weights_from_distributed( + [("__positions__", positions_gpu), ("__values__", values_gpu)], + pbar=pbar, + load_format="delta", + delta=spec, + ) + else: # disk + tensors = {"__positions__": positions_cpu, "__values__": values_gpu.cpu()} + metadata = { + "encoding": self.encoding.value, + "params": json.dumps([asdict(p) for p in params]), + "current_version": str(self.weight_version), + "checksum": str(checksum), + } + filename = f"rank{dist.get_rank():04d}_flush{self._flush_idx:06d}.safetensors" + path = os.path.join(self._version_dir, filename) + self.writer.enqueue(path, tensors, metadata) + self._pending_files.append(filename) + if pbar is not None: + pbar.update(1) + self._flush_idx += 1 + + def _publish_batch(self) -> None: + """ + Drain pending fsyncs, invoke the pre-push hook (may return a Future for an + async durability step on shared FS), then defer rank 0's + ``update_weights_from_disk`` RPC behind that Future via ``_rpc_executor``. + Each deferred dispatch lands in ``_pending_publishes`` as a + Future[list[ObjectRef]]; ``_finalize_sync`` awaits both layers. Safe to call + with empty ``_pending_files``: the all_gather still synchronizes and rank 0 + skips the dispatch when no rank produced files. + """ + self.writer.drain() + dist.barrier(group=get_gloo_group()) + + commit_future = None + if self._pre_push_hook is not None: + commit_future = self._pre_push_hook(self.args, self._version_dir, list(self.rollout_engines)) + dist.barrier(group=get_gloo_group()) + + # Collect every rank's batch filenames at rank 0; payload is ~KB, gather is cheap. + all_files: list[list[str]] = [None] * dist.get_world_size() # type: ignore[list-item] + dist.all_gather_object(all_files, list(self._pending_files), group=get_gloo_group()) + flat = [f for sub in all_files for f in sub] + self._pending_files.clear() + + if dist.get_rank() == 0 and flat: + version_dir = self._version_dir + engines = list(self.rollout_engines) + weight_version = str(self.weight_version) + + def _fire_when_committed() -> list: + if commit_future is not None: + commit_future.result() + return [ + engine.update_weights_from_disk.remote( + model_path=version_dir, + files=flat, + load_format="delta", + weight_version=weight_version, + ) + for engine in engines + ] + + self._pending_publishes.append(self._rpc_executor.submit(_fire_when_committed)) + + def _finalize_sync(self) -> None: + """ + Per-transport end-of-sync. NCCL: each flush already broadcasted; just resume. + Disk: publish the trailing files, wait for all streamed applies to land, then + cleanup + resume. + """ + if self.transport == "nccl": + if dist.get_rank() == 0: + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + dist.barrier(group=get_gloo_group()) + return + + if self._pending_files: + self._publish_batch() + if dist.get_rank() == 0: + # Each entry is a Future returning a list of ObjectRefs. Awaiting the + # Futures unblocks the (commit-then-RPC) chain; ray.get waits for the + # receivers' apply to finish. + object_refs = [ref for fut in self._pending_publishes for ref in fut.result()] + ray.get(object_refs) + self._pending_publishes.clear() + if not self.args.update_weight_delta_keep_files: + shutil.rmtree(self._version_dir, ignore_errors=True) + ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) + dist.barrier(group=get_gloo_group()) + + def _record_metrics(self) -> None: + """ + Allreduce density/byte counters across PP-src ranks; stash on + ``update_weight_metrics`` for the actor to drain into the next step log. + Wall-clock timings come from the slime ``Timer`` (``delta_encode`` / + ``delta_finalize`` blocks above + the outer ``update_weights`` decorator). + """ + pre_bytes = self.writer.bytes_pre_compress if self.writer is not None else 0 + post_bytes = self.writer.bytes_post_compress if self.writer is not None else 0 + counts = torch.tensor( + [self.density_nnz, self.density_numel, self.wire_bytes, pre_bytes, post_bytes], + dtype=torch.int64, + device=torch.cuda.current_device(), + ) + dist.all_reduce(counts) + nnz, numel, wire_bytes, pre_bytes, post_bytes = counts.tolist() + + density = nnz / max(numel, 1) + compression_ratio = (pre_bytes / post_bytes) if post_bytes > 0 else 1.0 + + m = self.update_weight_metrics + m["perf/update_weights_density"] = density + m["perf/update_weights_wire_bytes"] = wire_bytes + m["perf/update_weights_flushes_per_rank"] = float(self._flush_idx) + if self.transport == "disk": + m["perf/update_weights_disk_bytes_pre_compress"] = pre_bytes + m["perf/update_weights_disk_bytes_post_compress"] = post_bytes + m["perf/update_weights_compression_ratio"] = compression_ratio + + if dist.get_rank() == 0: + t = Timer().log_dict() + logger.info( + "[delta sync v=%s] transport=%s enc=%s density=%.3f%% " "encode=%.2fs finalize=%.2fs flushes/rank=%d", + self.weight_version, + self.transport, + self.encoding.value, + 100.0 * density, + t.get("delta_encode", 0.0), + t.get("delta_finalize", 0.0), + self._flush_idx, + ) diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py index dbba6aeb57..724e05355b 100644 --- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -48,6 +48,7 @@ def __init__( self.model_name = model_name self.quantization_config = quantization_config self.weight_version = 0 + self.update_weight_metrics: dict[str, float] = {} self._hf_weight_iterator = HfWeightIteratorBase.create( args=args, model=model, model_name=model_name, quantization_config=quantization_config @@ -134,6 +135,14 @@ def connect_rollout_engines( if start <= dist.get_rank() < end: self._ipc_engine = engine + def pop_metrics(self) -> dict[str, float]: + """ + Return and clear ``update_weight_metrics``. Empty under colocate today; + kept symmetric with UpdateWeightFromDistributed so the actor can drain unconditionally. + """ + out, self.update_weight_metrics = self.update_weight_metrics, {} + return out + @torch.no_grad() def update_weights(self) -> None: """ diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py index c28e13d5f9..0564915e41 100644 --- a/slime/backends/sglang_utils/sglang_engine.py +++ b/slime/backends/sglang_utils/sglang_engine.py @@ -367,15 +367,28 @@ def resume_memory_occupation(self, tags: list[str] = None): def check_weights(self, action: str): return self._make_request("weights_checker", {"action": action}) - def update_weights_from_disk(self, model_path: str, load_format: str | None = None): + def update_weights_from_disk( + self, + model_path: str, + load_format: str | None = None, + weight_version: str | None = None, + files: list[str] | None = None, + ): """Reload weights from *model_path* without restarting the engine. - Used for non-updatable (frozen) models that overlap with megatron: - after offload, weights are restored from disk instead of CPU cache. + Standard HF reload: ``model_path`` is the checkpoint directory. + Delta (``load_format="delta"``): ``model_path`` is the parent of the + per-sync version subdir and ``files`` is the basenames within it to read + + apply. Each delta call is independent — sender owns batching, sync + boundaries, cleanup. """ - payload = {"model_path": model_path} + payload: dict = {"model_path": model_path} if load_format is not None: payload["load_format"] = load_format + if weight_version is not None: + payload["weight_version"] = weight_version + if files is not None: + payload["files"] = files return self._make_request("update_weights_from_disk", payload) def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend): @@ -404,7 +417,15 @@ def destroy_weights_update_group(self, group_name): pass def update_weights_from_distributed( - self, names, dtypes, shapes, group_name, flush_cache=False, weight_version: str | None = None + self, + names, + dtypes, + shapes, + group_name, + flush_cache=False, + weight_version: str | None = None, + load_format: str | None = None, + delta=None, ): payload = { "names": names, @@ -415,6 +436,21 @@ def update_weights_from_distributed( } if weight_version is not None: payload["weight_version"] = weight_version + if load_format is not None: + payload["load_format"] = load_format + if delta is not None: + # DeltaSpec → JSON string. Receiver reconstructs via DeltaEncoding(...) + + # DeltaParam(**p); avoids depending on FastAPI's nested-dataclass coercion. + import json + from dataclasses import asdict + + payload["delta"] = json.dumps( + { + "encoding": delta.encoding.value, + "params": [asdict(p) for p in delta.params], + "checksum": delta.checksum, + } + ) return self._make_request( "update_weights_from_distributed", payload, diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 778e1aaec4..d9acb2a8af 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -136,6 +136,67 @@ def add_train_arguments(parser): default="raw", help="The method to convert megatron weights to hugging face weights for SGLang.", ) + # Delta weight sync. + parser.add_argument( + "--update-weight-mode", + choices=["full", "delta"], + default="full", + help=( + "Weight sync strategy. 'full' (default) broadcasts every parameter " + "every sync. 'delta' detects byte-level changes against a pinned-CPU " + "snapshot of the previous broadcast and ships only the changed positions + values." + ), + ) + parser.add_argument( + "--update-weight-transport", + choices=["nccl", "disk"], + default="nccl", + help=( + "Per-flush carrier for --update-weight-mode=delta. 'nccl' broadcasts each " + "bucket; 'disk' writes each bucket as a safetensors file under " + "--update-weight-delta-dir and pushes once at end-of-sync." + ), + ) + parser.add_argument( + "--update-weight-encoding", + choices=["indices", "deltas", "deltas_zstd"], + default="indices", + help=( + "Position encoding for partial flushes. 'indices': int32 absolute " + "positions (largest, lowest compute). 'deltas': uint16 gap-deltas " + "with uint32 fallback (smaller). 'deltas_zstd': 'deltas' with the " + "safetensors blob wrapped in zstd L1 (smallest, heaviest compute — " + "best for shared-FS bandwidth ≤ ~300 MB/s)." + ), + ) + parser.add_argument( + "--update-weight-delta-dir", + type=str, + default=None, + help=( + "Filesystem directory for per-sync delta safetensors. Writable by the " + "trainer, readable by every rollout engine. Required when " + "--update-weight-transport=disk. One subdirectory per sync " + "(``weight_v{N:06d}``), removed after every engine has acknowledged." + ), + ) + parser.add_argument( + "--update-weight-delta-keep-files", + action="store_true", + default=False, + help="Skip post-apply cleanup of per-sync version directories. Useful for debugging.", + ) + parser.add_argument( + "--custom-delta-pre-push-path", + type=str, + default=None, + help=( + "Path to a custom function called by --update-weight-transport=disk after each " + "trainer rank's files are durably on local disk, before rank 0 fires the engine " + "RPCs. Signature: ``def hook(args, version_dir: str, rollout_engines) -> None``. " + "Called from every trainer rank; the hook gates itself." + ), + ) parser.add_argument( "--custom-model-provider-path", type=str, @@ -1846,3 +1907,16 @@ def slime_validate_args(args): if args.only_train_params_name_list and args.freeze_params_name_list: raise ValueError("You can only specify ONE of: --only-train-params-name-list, or --freeze-params-name-list.") + + if args.update_weight_mode == "delta": + if args.colocate: + raise ValueError( + "--update-weight-mode=delta is not supported with --colocate. Colocate transfers " + "weights via CUDA IPC (only a handle crosses processes), so the delta bookkeeping " + "(snapshot + diff + sparse encode) is pure overhead." + ) + if args.update_weight_transport == "disk" and not args.update_weight_delta_dir: + raise ValueError( + "--update-weight-transport=disk requires --update-weight-delta-dir to point at " + "a filesystem shared between the trainer and the rollout engines." + ) diff --git a/slime/utils/train_metric_utils.py b/slime/utils/train_metric_utils.py index 9bec049d15..8ecefcfbf3 100644 --- a/slime/utils/train_metric_utils.py +++ b/slime/utils/train_metric_utils.py @@ -11,7 +11,11 @@ def log_perf_data_raw( - rollout_id: int, args: Namespace, is_primary_rank: bool, compute_total_fwd_flops: Callable + rollout_id: int, + args: Namespace, + is_primary_rank: bool, + compute_total_fwd_flops: Callable, + extra_metrics: dict | None = None, ) -> None: timer_instance = Timer() log_dict_raw = deepcopy(timer_instance.log_dict()) @@ -21,6 +25,8 @@ def log_perf_data_raw( return log_dict = {f"perf/{key}_time": val for key, val in log_dict_raw.items()} + if extra_metrics: + log_dict.update(extra_metrics) if ("perf/actor_train_time" in log_dict) and (compute_total_fwd_flops is not None): total_fwd_flops = compute_total_fwd_flops(seq_lens=timer_instance.seq_lens) From 4e90292ba2a701aca5913d7a5513e3b32cdb3761 Mon Sep 17 00:00:00 2001 From: Nan Date: Mon, 25 May 2026 18:16:58 +0000 Subject: [PATCH 2/2] add example --- docs/en/advanced/delta-weight-sync.md | 82 ++++++++ docs/en/index.rst | 1 + docs/zh/advanced/delta-weight-sync.md | 80 ++++++++ docs/zh/index.rst | 1 + examples/README.md | 1 + examples/delta_weight_sync/README.md | 67 ++++++ .../run-glm4.7-355B-A32B-delta.sh | 192 ++++++++++++++++++ 7 files changed, 424 insertions(+) create mode 100644 docs/en/advanced/delta-weight-sync.md create mode 100644 docs/zh/advanced/delta-weight-sync.md create mode 100644 examples/delta_weight_sync/README.md create mode 100755 examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh diff --git a/docs/en/advanced/delta-weight-sync.md b/docs/en/advanced/delta-weight-sync.md new file mode 100644 index 0000000000..6a9e5c18c5 --- /dev/null +++ b/docs/en/advanced/delta-weight-sync.md @@ -0,0 +1,82 @@ +# Delta Weight Sync + +- [Why](#why) +- [Quick Start](#quick-start) +- [How It Works](#how-it-works) +- [Encoding Choice](#encoding-choice) +- [Why Not Colocated](#why-not-colocated) + +## Why + +Slime's default sync broadcasts every parameter every step. The cost scales linearly with model size and dominates the sync phase, even though only a few percent of weights change between consecutive RL steps. Delta sync keeps a pinned-CPU snapshot of the last broadcast and ships only the positions whose bytes differ. + +The motivating use case is **training/inference disaggregation** — running the trainer and the rollout engines in *different datacenters* over a shared filesystem with bandwidth on the order of 100s of MB/s, where a full broadcast is infeasible but a sparse delta (~3% density, ~5 GB for a 355B model) is. The same delta machinery also runs over NCCL inside a single datacenter, where it serves as the validation baseline that proves the wire encoding and apply logic are correct. + +Prior art: selective overwrite is inspired by [arXiv:2509.19128](https://arxiv.org/abs/2509.19128); the cross-DC disaggregation motivation is from [Fireworks AI — Frontier RL Is Cheaper Than You Think](https://fireworks.ai/blog/frontier-rl-is-cheaper-than-you-think). + +## Quick Start + +Disk transport (training/inference disaggregation — the main use case): + +```bash +--update-weight-mode delta +--update-weight-transport disk +--update-weight-encoding deltas_zstd # best for ≤ 300 MB/s shared FS +--update-weight-delta-dir /shared/fs/delta-updates +``` + +NCCL transport (intra-datacenter validation baseline): + +```bash +--update-weight-mode delta +--update-weight-transport nccl +--update-weight-encoding indices # lowest compute, no compression +``` + +Receiver-side tuning (applies to both transports): + +```bash +--sglang-update-weight-delta-chunk-bytes $((2 * 1024 * 1024 * 1024)) # byte cap per load_weights call +--sglang-update-weight-delta-read-workers 4 # parallel I/O threads (disk only) +``` + +See [examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh](../../../examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh) for a complete launcher. + +## How It Works + +Both transports share one sender pipeline, one wire layout, and one receiver-side decoder; only the per-flush carrier differs. + +**Sender (per sync, PP-source rank only):** + +1. **Diff** the current weights against the pinned-CPU snapshot via bytewise compare (`current.view(int_dtype) != snapshot.view(int_dtype)`) — lossless, dtype-agnostic, no arithmetic. +2. **Encode** changed (position, value) pairs into a packed `__positions__` byte blob + `__values__` tensor + per-param decoding manifest. The encoding (`indices`, `deltas`, `deltas_zstd`) governs only how positions are packed; values are sent verbatim in the param's dtype. +3. **Bucket** per-chunk encodes up to `--update-weight-buffer-size` bytes, then flush: + - NCCL: broadcast `(__positions__, __values__)` to the rollout engines with a `DeltaSpec` (encoding + per-param manifest) carried in the Ray RPC. + - Disk: write one safetensors file per flush under `weight_v{N:06d}/`. Async background thread does the I/O + optional zstd compression off the critical path. +4. **Snapshot the just-sent values** via a D2H copy on a side stream so it overlaps with the next chunk's encode. + +**End-of-sync (disk only):** write a `DONE` marker, then rank 0 fires one HTTP push per engine and removes the directory after every engine acknowledges. + +**Receiver:** + +For both transports, the receiver ends up calling the same `_apply_delta_payload(encoding, params, positions, values)` helper. It decodes each param's slice into a full-shape tensor with NaN at unchanged positions, then routes it through `model.load_weights(...)` under a `_delta_apply_context` that patches `Tensor.copy_` / `Tensor.fill_` to perform NaN-masked overwrite. Auxiliary writes (scratch buffers, fp8 scales, MoE biases via `post_load_weights`) keep their normal semantics. + +Selective overwrite has no arithmetic — the receiver writes the trainer's exact bytes at changed positions — so it's lossless by construction and there's no notion of drift to fight with periodic base re-syncs. + +## Encoding Choice + +`--update-weight-encoding` picks how positions are packed. All three share the same on-wire layout (`__positions__` uint8 blob + `__values__` tensor + per-param manifest); decoder dispatches on the metadata. + +| value | positions | when to pick | +|---|---|---| +| `indices` | int32 absolute positions (4 bytes / nnz) | NCCL or fast intra-cluster FS (≥ ~600 MB/s) | +| `deltas` | uint16 gap-deltas with uint32 fallback (~2 bytes / nnz at 2% density) | medium FS bandwidth (~300-500 MB/s) | +| `deltas_zstd` | `deltas` wrapped in zstd L1 on disk | cross-DC / cross-region shared FS (≤ ~300 MB/s) | + +**Why gap-encoded positions are smaller**: positions come out of `mask.nonzero()` already sorted ascending. At density `p`, the expected gap between consecutive nonzero positions is `1/p`, and `P(gap > 65535) ≈ exp(-p · 65535)`. At p = 2% that's effectively zero, so uint16 fits with a uint32 per-param fallback for pathological inputs. Half the position bytes of `indices`, lossless. + +**Break-even with `indices`** at our density (~2%): `deltas` halves the positions blob (which dominates the wire); `zstd` shaves another ~35-40% on top by compressing the gap byte stream, at the cost of ~250ms/file compress + ~150ms/file decompress. The crossover with `indices` is where compress/decompress compute exceeds the bandwidth savings — empirically around 500 MB/s for `deltas` and 300 MB/s for `deltas_zstd`. + +## Why Not Colocated + +Colocated weight sync uses CUDA IPC: only a memory handle (~64 B) crosses processes. Delta encoding's "bytes saved on the wire" benefit is zero, while the bookkeeping (snapshot + diff + sparse encode) is pure overhead. Slime rejects `--update-weight-mode delta --colocate` at argparse time. diff --git a/docs/en/index.rst b/docs/en/index.rst index d5ca098e06..d07321f4cb 100644 --- a/docs/en/index.rst +++ b/docs/en/index.rst @@ -47,6 +47,7 @@ slime is the RL-framework behind GLM-4.7, GLM-4.6 and GLM-4.5. Apart from models advanced/reproducibility.md advanced/fault-tolerance.md advanced/pd-disaggregation.md + advanced/delta-weight-sync.md advanced/sglang-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md diff --git a/docs/zh/advanced/delta-weight-sync.md b/docs/zh/advanced/delta-weight-sync.md new file mode 100644 index 0000000000..256f55c969 --- /dev/null +++ b/docs/zh/advanced/delta-weight-sync.md @@ -0,0 +1,80 @@ +# Delta 权重同步 + +- [背景](#背景) +- [快速开始](#快速开始) +- [工作原理](#工作原理) +- [编码选择](#编码选择) +- [为何不支持 colocated](#为何不支持-colocated) + +## 背景 + +slime 默认的权重同步会在每一步广播全部参数，开销随模型规模线性增长，即使每步真正变化的权重只有几个百分点。Delta 同步在内存中保留上一次同步后的参数快照（pinned CPU），只发送字节发生变化的位置。 + +最主要的应用场景是 **训练 / 推理跨数据中心解耦** —— 训练器和推理引擎运行在不同数据中心，通过共享文件系统通信（带宽通常在百 MB/s 级别）。在这种环境下，全量广播不可行，而 ~3% 密度的稀疏 delta（355B 模型约 5 GB）是可行的。同一套 delta 机制在数据中心内部跑 NCCL，作为验证基线，确认 wire 编码和 apply 逻辑正确。 + +参考资料：选择性覆写借鉴自 [arXiv:2509.19128](https://arxiv.org/abs/2509.19128)，跨数据中心的动机来自 [Fireworks AI — Frontier RL Is Cheaper Than You Think](https://fireworks.ai/blog/frontier-rl-is-cheaper-than-you-think)。 + +## 快速开始 + +磁盘传输（跨数据中心训推解耦，主要场景）： + +```bash +--update-weight-mode delta +--update-weight-transport disk +--update-weight-encoding deltas_zstd # ≤ 300 MB/s 共享 FS 推荐 +--update-weight-delta-dir /shared/fs/delta-updates +``` + +NCCL 传输（数据中心内部验证基线）： + +```bash +--update-weight-mode delta +--update-weight-transport nccl +--update-weight-encoding indices # 计算最少，无压缩 +``` + +接收端调优（两种传输都适用）： + +```bash +--sglang-update-weight-delta-chunk-bytes $((2 * 1024 * 1024 * 1024)) # 每次 load_weights 字节上限 +--sglang-update-weight-delta-read-workers 4 # 并行 I/O 线程数（仅磁盘传输） +``` + +完整启动脚本见 [examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh](../../../examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh)。 + +## 工作原理 + +两种传输共用同一条发送管线、同一种 wire 布局以及同一套接收端解码器；只有每个 bucket 的承载层不同。 + +**发送端（每次同步，仅 PP 源 rank）：** + +1. **求差**：通过逐字节比较 `current.view(int_dtype) != snapshot.view(int_dtype)` 检测变化。无算术、无损、与 dtype 无关。 +2. **编码**：将变化的 (位置, 值) 对打包成 `__positions__` 字节块 + `__values__` 张量 + per-param 解码 manifest。编码方式（`indices` / `deltas` / `deltas_zstd`）只影响位置如何打包，值始终按参数本身的 dtype 原样发送。 +3. **打包并发送**：每个 chunk 编码后累积至 `--update-weight-buffer-size` 字节再 flush： + - NCCL：广播 `(__positions__, __values__)`，Ray RPC 同时携带 `DeltaSpec`（编码 + per-param manifest）。 + - 磁盘：每个 flush 写一个 safetensors 文件到 `weight_v{N:06d}/` 目录，后台线程负责 I/O 和可选的 zstd 压缩，不阻塞关键路径。 +4. **更新快照**：刚发送的值在 side stream 上 D2H 拷贝，与下一个 chunk 的编码重叠。 + +**同步结束（仅磁盘）：** 写 `DONE` 标记，rank 0 对每个引擎触发一次 HTTP push，所有引擎确认后清理目录。 + +**接收端：** 两种传输最终都进入同一个 `_apply_delta_payload(encoding, params, positions, values)` 帮助函数。它把每个参数的切片解码成全形状张量，未变化位置填 NaN，然后通过 `model.load_weights(...)` 应用；过程中 `_delta_apply_context` 替换 `Tensor.copy_` / `Tensor.fill_`，对参数存储执行 NaN 掩码覆写。辅助写入（scratch buffer、fp8 scale、MoE bias 等通过 `post_load_weights` 写入的派生张量）保留正常语义。 + +选择性覆写没有任何算术运算 —— 接收端在变化位置直接写入训练端的精确字节 —— 因此天然无损，也不存在数值漂移问题，无需周期性 base 同步。 + +## 编码选择 + +`--update-weight-encoding` 决定位置如何打包。三种编码共用同一种 wire 布局（`__positions__` uint8 块 + `__values__` 张量 + per-param manifest），解码端根据 metadata 分派。 + +| 取值 | 位置编码 | 推荐场景 | +|---|---|---| +| `indices` | int32 绝对位置（4 字节 / nnz） | NCCL 或高速集群内 FS（≥ ~600 MB/s） | +| `deltas` | uint16 增量（异常时 uint32 兜底，2% 密度下约 2 字节 / nnz） | 中等带宽 FS（~300-500 MB/s） | +| `deltas_zstd` | `deltas` 文件再用 zstd L1 压缩 | 跨数据中心 / 跨区共享 FS（≤ ~300 MB/s） | + +**为何 gap 编码更省**：`mask.nonzero()` 返回的位置已经升序排列。密度 `p` 时连续非零位置的期望间隔为 `1/p`，且 `P(gap > 65535) ≈ exp(-p · 65535)`，p = 2% 时这个概率实际上为零，所以 uint16 完全够用，uint32 仅作 per-param 兜底。位置开销比 `indices` 减半，且无损。 + +**`deltas_zstd` 的额外收益**：在 gap 字节流上做 zstd L1 还能再减少 ~35-40%，代价是每文件约 250ms 压缩 + 150ms 解压。当共享 FS 带宽 ≤ 300 MB/s 时，带宽节省超过额外计算开销。 + +## 为何不支持 colocated + +Colocated 同步通过 CUDA IPC：进程间传递的只是一个内存句柄（~64 B）。Delta 编码的"wire 节省"在此为零，而其簿记开销（快照 + 求差 + 稀疏编码）反而是纯损失。slime 在参数校验阶段拒绝 `--update-weight-mode delta --colocate`。 diff --git a/docs/zh/index.rst b/docs/zh/index.rst index 36d3d79eb2..6b1d6e76c7 100644 --- a/docs/zh/index.rst +++ b/docs/zh/index.rst @@ -47,6 +47,7 @@ slime 是 GLM-4.7、GLM-4.6、GLM-4.5 背后的 RL 训练框架。除此之外 advanced/reproducibility.md advanced/fault-tolerance.md advanced/pd-disaggregation.md + advanced/delta-weight-sync.md advanced/sglang-config.md advanced/megatron-config.md advanced/arch-support-beyond-megatron.md diff --git a/examples/README.md b/examples/README.md index a1e3bd251e..128b1562d4 100644 --- a/examples/README.md +++ b/examples/README.md @@ -11,6 +11,7 @@ These examples provide concrete examples to leverage slime in your own RL workfl - **[low_precision](./low_precision)**: Examples of FP8 training and inference for improved throughput and stability. - **[multi_agent](./multi_agent)**: Example of running multi-agent RL with `slime`. - **[on_policy_distillation](./on_policy_distillation)**: Example implementation for on-policy distillation, extending the reinforcement learning pipeline to support teacher–student distillation directly within on-policy training. +- **[delta_weight_sync](./delta_weight_sync)**: Non-colocated weight sync that ships only changed positions + values over disk (training/inference disaggregation) or NCCL. - **[reproducibility](./reproducibility)**: Guides on achieving bitwise experiment reproduction using deterministic modes. - **[retool](./retool)**: Demonstrates the retool functionality for tool-enabled language model generation. - **[search-r1](./search-r1)**: A minimal reproduction of Search-R1, featuring multi-turn conversation and tool-calling. diff --git a/examples/delta_weight_sync/README.md b/examples/delta_weight_sync/README.md new file mode 100644 index 0000000000..3782b20bb4 --- /dev/null +++ b/examples/delta_weight_sync/README.md @@ -0,0 +1,67 @@ +# Delta Weight Sync + +Non-colocated weight sync that ships only changed positions + values instead of every parameter. Two transports over one wire format and one receiver-side decoder: + +- **Disk** (the point) — write per-flush safetensors to a shared filesystem; one HTTP push per sync. Designed for **training/inference disaggregation** across datacenters where bandwidth between trainer and rollout is on the order of 100s of MB/s. +- **NCCL** (the baseline) — broadcast each per-flush bucket directly. Used intra-datacenter to validate that the wire encoding and apply logic are correct, separate from any shared-FS variable. + +Both modes are lossless by construction (selective overwrite via NaN sentinel; no arithmetic). + +## Files + +- `run-glm4.7-355B-A32B-delta.sh`: 16-node (8 actor + 8 rollout) GLM-4.7-355B-A32B launcher. Disk transport active by default; NCCL block commented below it. + +## Usage + +```bash +bash examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh +``` + +**Disk (default):** + +```bash +DELTA_ARGS=( + --update-weight-mode delta + --update-weight-transport disk + --update-weight-encoding deltas_zstd + --update-weight-delta-dir /shared/fs/delta-updates +) +``` + +**NCCL (baseline):** + +```bash +DELTA_ARGS=( + --update-weight-mode delta + --update-weight-transport nccl + --update-weight-encoding indices +) +``` + +Receiver-side byte cap (both transports): + +```bash +--sglang-update-weight-delta-chunk-bytes $((2 * 1024 * 1024 * 1024)) +``` + +See [docs/en/advanced/delta-weight-sync.md](../../docs/en/advanced/delta-weight-sync.md) for the wire protocol, encoding choice, and design. + +## Results + +W&B traces comparing delta sync against the full-sync baseline on GLM-4.7-355B-A32B / DAPO-Math-17k. + +![Raw reward](./raw_reward.png) + +![Train/rollout logprob abs diff](./train_rollout_logprob_abs_diff.png) + +![Update weights time](./update_weights_time.png) + +> **Note on the small curve-to-curve gap.** RL training is inherently non-deterministic (cuBLAS reductions, FlashAttention split-K, NCCL all-reduce ordering, dynamic-batch token assignment). Two identically-configured *full*-sync runs would diverge the same way. Delta sync's selective overwrite is bit-exact with full sync per step (no arithmetic, no drift); the trajectory matches, the bits don't. + +![Update weights density](./update_weights_density.png) + +*Per-sync change density (`perf/update_weights_density`) — fraction of weight positions that moved between consecutive syncs. Sync 0 is omitted: it's the snapshot-seeding pass with density = 1.0, which would compress the y-axis.* + +## Why these encoding defaults + +Per-sync change density during RL fine-tuning at conservative LRs sits around **2-3%** ([arXiv:2602.03839](https://arxiv.org/pdf/2602.03839) reports ~1% on a related setup; we measured ~2-3% on this run). Below the 3.125% break-even point, gap-encoded positions are smaller than absolute indices — the disk default `deltas_zstd` adds zstd L1 on top to squeeze the gap byte stream further (~35-40%), which is the right tradeoff when shared-FS bandwidth is ≤ 300 MB/s. Intra-datacenter NCCL has no bandwidth pressure, so `indices` (lowest compute, biggest payload) is the cleaner default there. diff --git a/examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh b/examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh new file mode 100755 index 0000000000..9cf2dc9246 --- /dev/null +++ b/examples/delta_weight_sync/run-glm4.7-355B-A32B-delta.sh @@ -0,0 +1,192 @@ +#!/bin/bash + +# Non-colocated GLM-4.7-355B-A32B with delta weight sync. +# 8 actor nodes (TP=8, PP=4, EP=16) + 64 rollout GPUs (8 H100 nodes worth), 16 nodes total. +# Disk transport is active by default; the NCCL block below it is commented out. + +pkill -9 sglang +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python + +set -ex + +export PYTHONBUFFERED=16 +unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +source "/root/slime/scripts/models/glm4.5-355B-A32B.sh" + +CKPT_ARGS=( + --hf-checkpoint /root/GLM-4.7-355B-A32B + --ref-load /root/GLM-4.7-355B-A32B_torch_dist/ +) + +ROLLOUT_ARGS=( + --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl + --input-key prompt + --label-key label + --apply-chat-template + --rollout-shuffle + --rm-type deepscaler + --num-rollout 3000 + --rollout-batch-size 64 + --n-samples-per-prompt 8 + --rollout-max-response-len 8192 + --rollout-temperature 1 + + --num-steps-per-rollout 4 + --balance-data + --rollout-stop-token-ids 151329 151336 151338 +) + +EVAL_ARGS=( + --eval-interval 20 + --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl + --n-samples-per-eval-prompt 8 + --eval-max-response-len 8192 + --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 8 + --sequence-parallel + --pipeline-model-parallel-size 4 + --context-parallel-size 2 + --expert-model-parallel-size 16 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) + +GRPO_ARGS=( + --advantage-estimator gspo + --kl-loss-coef 0.00 + --kl-loss-type low_var_kl + --kl-coef 0.00 + --entropy-coef 0.00 + --eps-clip 1e-4 + --eps-clip-high 2e-4 + --use-tis +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + + --optimizer-cpu-offload + --overlap-cpu-optimizer-d2h-h2d + --use-precision-aware-optimizer +) + +WANDB_ARGS=( + # --use-wandb + # --wandb-project slime-delta + # --wandb-group glm4.7-355B-delta +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 32 + --sglang-mem-fraction-static 0.7 + --sglang-enable-dp-attention + --sglang-dp-size 4 + --sglang-ep-size 32 + --sglang-enable-dp-lm-head + --sglang-moe-dense-tp-size 1 + + # Receiver batches up to this many bytes per model.load_weights call. Bigger + # amortizes per-call cost (name resolution, MoE expert remap) but raises peak HBM. + --sglang-update-weight-delta-chunk-bytes $((2 * 1024 * 1024 * 1024)) + + # Max parallel I/O threads for reading delta files from disk (disk transport only). + --sglang-update-weight-delta-read-workers 4 + + # mtp + --sglang-speculative-algorithm EAGLE + --sglang-speculative-num-steps 3 + --sglang-speculative-eagle-topk 1 + --sglang-speculative-num-draft-tokens 4 +) + +# Delta weight sync. Pick one of the two blocks below. + +# ── Disk (default) — for training/inference disaggregation across datacenters ──── +# `deltas_zstd` is the right pick when shared-FS bandwidth is ≤ ~300 MB/s. +DELTA_ARGS=( + --update-weight-mode delta + --update-weight-transport disk + --update-weight-encoding deltas_zstd + --update-weight-delta-dir /shared/fs/delta-updates +) + +# ── NCCL (baseline) — intra-datacenter, no shared FS ──────────────────────────── +# DELTA_ARGS=( +# --update-weight-mode delta +# --update-weight-transport nccl +# --update-weight-encoding indices +# ) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --moe-token-dispatcher-type flex + --moe-enable-deepep + --update-weight-buffer-size $((2 * 1024 * 1024 * 1024)) +) + +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + +RUNTIME_ENV_JSON=$(cat <