diff --git a/.claude/docs/architecture.md b/.claude/docs/architecture.md index f439440bf2..330336bcaa 100644 --- a/.claude/docs/architecture.md +++ b/.claude/docs/architecture.md @@ -11,7 +11,7 @@ skyrl/ # Core library │ └── skyrl_train/ # FSDP/Megatron training backend │ ├── distributed/ # Dispatch, FSDP/Megatron strategies │ ├── inference_servers/ # HTTP inference path (RemoteInferenceClient, vLLM servers, router) -│ ├── weight_sync/ # Weight extraction and transfer +│ ├── weight_sync/ # WeightSources + trainer-side transfer engines │ └── workers/ # FSDP/Megatron workers ├── train/ # Training entrypoints, config, dataset, generators, trainer │ ├── config/ # Hydra YAML configs (ppo_base, megatron, skyrl_gym) diff --git a/.claude/docs/backends/fsdp.md b/.claude/docs/backends/fsdp.md index fdb0645a3d..576de10d97 100644 --- a/.claude/docs/backends/fsdp.md +++ b/.claude/docs/backends/fsdp.md @@ -6,7 +6,7 @@ Default backend (`trainer.strategy=fsdp`). Uses PyTorch FSDP2 for distributed tr - **FSDPConfig** in `skyrl/train/config.py`. - **FSDPStrategy** in `skyrl/backends/skyrl_train/distributed/fsdp_strategy.py`. -- **FSDPWeightExtractor** for extracting weights from sharded parameters (in `skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py`). +- **FsdpWeightSource** (`skyrl/backends/skyrl_train/weight_sync/sources.py`) presents the sharded model to weight sync as vLLM's `WeightSource`; the worker builds it in `FSDPPolicyWorkerBase._build_weight_source`. ## CPU Offload diff --git a/.claude/docs/weight_sync.md b/.claude/docs/weight_sync.md index 85b4baa25b..2429b735ac 100644 --- a/.claude/docs/weight_sync.md +++ b/.claude/docs/weight_sync.md @@ -4,63 +4,155 @@ Training-to-inference weight transfer. Runs after every training step (or on the ## Architecture -Two-sided protocol with sender (training) / receiver (inference): +Weight sync runs on vLLM's **trainer-send** abstraction. Each training worker builds +a `WeightSource` over its live model and hands it to a `TrainerWeightTransferEngine`, +whose `send_weights()` owns the whole round trip (`start_weight_update` → +`update_weights` → `finish_weight_update`, plus barriers and the non-sender collective +replay). SkyRL owns three things and nothing else: the source, the control-plane +client, and choosing an init info. + +**Chunking is vLLM's responsibility.** SkyRL passes a source and stops — it does not +decide, express, or bound how the stream is cut on the wire. The packed NCCL and IPC +producers already chunk out of a fixed reusable buffer and consume the source lazily, +so a source's only obligation is to be a lazy generator that does not retain. ``` skyrl/backends/skyrl_train/weight_sync/ -├── base.py # WeightUpdateRequest, LoraLoadRequest, WeightChunk -├── transfer_strategy.py # WeightSyncInitInfo / Sender / Strategy ABCs (sender-side only; receive is vLLM-native) -├── broadcast_strategy.py # NCCL broadcast (non-colocated) -├── nccl_trainer_send.py # vendored: vLLM 0.26's NCCL trainer-send statics -├── cuda_ipc_strategy.py # CUDA IPC (colocated) -├── delta_strategy.py # Checkpoint-delta sender + strategy (disk / gs:// / s3://) +├── __init__.py # backend selection: get_transfer_strategy / get_vllm_receive_backend +├── base.py # LoraLoadRequest (not a weight transfer -- an adapter path) +├── sources.py # FsdpWeightSource / MegatronWeightSource (vLLM's metadata()+__iter__) +├── trainer_engines.py # build_trainer_engine: init info + client -> trainer_init +├── control_plane.py # SkyrlWeightSyncClient (blocking HTTP) + per-server init rewrites +├── skyrl_engines.py # receive side: skyrl_nccl / skyrl_ipc (+ the drafter-reload proxy) +├── delta_trainer.py # DeltaTrainerWeightTransferEngine (send side) +├── delta_engine.py # DeltaWeightTransferEngine (receive side, in the vLLM worker) ├── delta_checkpoint.py # DeltaCheckpointPublisher, LocalCheckpointStore, manifest + XOR payloads -├── delta_engine.py # DeltaWeightTransferEngine (receive side, runs in the vLLM worker) ├── delta_payload.py # zstd compress/decompress + uint8 tensor <-> bytes helpers -├── weight_extractor.py # Sharded-param -> dense tensor extraction -├── weight_extractor_utils.py └── sharded_rdt/ # the sharded_rdt (NIXL pull) backend; __init__ is import-free - ├── sharded_rdt_strategy.py # sharded_rdt as a WeightTransferStrategy (thin adapter) - ├── rdt_send.py # trainer-side driver: WeightSource impls + RdtWeightSyncSender - ├── rdt_control_plane.py # SyncRdtControlPlaneClient (blocking HTTP to /collective_rpc) + ├── rdt_send.py # its WeightSources + build_rdt_trainer_init_info + ├── sharded_rdt_base.py # GroupedWeightSource + layerwise_groups (the two RDT-only channels) ├── rdt_vllm_register.py # registers the sharded_rdt engine into vLLM's factory ├── rdt_libfabric_shim.py # LIBFABRIC provider shim for NIXL - ├── sharded_rdt_base.py # vendored: trainer-side ABCs (WeightSource, layerwise_groups) ├── sharded_rdt_trainer.py # vendored: trainer engine + the _RDTProducerServer sidecar ├── sharded_rdt_engine.py # vendored: consumer engine (runs in the vLLM worker) ├── sharded_rdt_common.py # vendored: RdtRouter, op-chain allowlist, buffer sizing └── sharded_rdt_fake.py # vendored: FakeRDTTensor placeholders for the bake ``` -The subpackage's `__init__.py` deliberately imports nothing: `sharded_rdt_engine` and -`sharded_rdt_trainer` import `vllm` at module scope, so re-exporting from it would pull -vllm into every `weight_sync` import and break the CPU CI job that runs without the wheel. -Public names (`ShardedRdtTransferStrategy` and friends) are re-exported from -`weight_sync/__init__.py`, which only reaches the vllm-free `sharded_rdt_strategy`. +Neither `__init__.py` imports anything: `sources`, `trainer_engines`, `delta_trainer`, +`sharded_rdt_engine` and `sharded_rdt_trainer` import `vllm` at module scope, so a +re-export would pull vllm into every `weight_sync` import and break the CPU CI job that +runs without the wheel. Import those modules at their call sites. vLLM worker-extension class (loaded via `--worker-extension-cls`): -- `skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py` — `NewInferenceWorkerWrap`. Three-phase chunked lifecycle. - -The weight sync implementation relies on the native vLLM weight sync APIs - `WeightTransferEngine` abstractions as well as native RPC endpoints for weight updates. - -## Transfer Strategies - -- **Broadcast** (`BroadcastTransferStrategy`): NCCL collective. Used for **non-colocated** setups. Training and inference are on different GPUs; weights cross the wire over a dedicated process group. -- **CUDA IPC** (`CudaIpcTransferStrategy`): Per-chunk packed buffer + one IPC handle per rank. Used for **colocated** setups (`colocate_all=true`). Both sides live on the same GPU; the receiver maps the sender's CUDA allocation directly. -- **Delta** (`DeltaTransferStrategy`): Weights travel as compressed XOR deltas against the base checkpoint, through a shared filesystem or object store instead of the network fabric. Selected with `generator.inference_engine.weight_sync_backend=delta`; intended for **non-colocated** setups where the two sides are not NCCL-reachable (separate clusters, PD-disaggregated serving). Not supported with LoRA (`validate_cfg` rejects it). +- `inference_servers/new_inference_worker_wrap.py` — `NewInferenceWorkerWrap`. Holds + exactly **two** weight-sync things, both limits of *dispatch* rather than of the + engine abstraction: `fetch_weights` (`/collective_rpc` reaches worker methods by name + only, and no native route can invoke an engine method) and the sleep/wake pair + (`EngineCore.sleep` hardcodes the prefix-cache clear). It is also where the + receive-side factory registrations and the model-runner recorder patch are installed, + because vLLM imports it in every worker process before model init. + +### Trainer side, per backend + +| logical backend | trainer engine | receive engine (`WeightTransferConfig.backend`) | +|---|---|---| +| `nccl` | vLLM's `NCCLTrainerWeightTransferEngine` | `skyrl_nccl` | +| `ipc` | vLLM's `IPCTrainerWeightTransferEngine` | `skyrl_ipc` | +| `delta` | `DeltaTrainerWeightTransferEngine` | `delta` | +| `sharded_rdt` | `ShardedRDTTrainerWeightTransferEngine` | `sharded_rdt` | + +The receive side takes new names for NCCL and IPC because SkyRL subclasses vLLM's +engines (to reload the spec-decode drafter) and `register_engine` raises on an +already-registered name. The trainer-side factory has its own registry, so the send side +keeps vLLM's names. + +## Transfer backends + +- **Broadcast** (`nccl`): NCCL collective, packed. Used for **non-colocated** setups. Training and inference are on different GPUs; weights cross the wire over a dedicated process group. +- **CUDA IPC** (`ipc`): packed IPC handles out of one reusable buffer. Used for **colocated** setups (`colocate_all=true`). Both sides live on the same GPU; the receiver maps the sender's CUDA allocation directly and clones out of it. +- **Delta** (`delta`): Weights travel as compressed XOR deltas against the base checkpoint, through a shared filesystem or object store instead of the network fabric. Selected with `generator.inference_engine.weight_sync_backend=delta`; intended for **non-colocated** setups where the two sides are not NCCL-reachable (separate clusters, PD-disaggregated serving). Not supported with LoRA (`validate_cfg` rejects it). - **Sharded RDT** (`sharded_rdt`): the inference workers **pull** the slices they consume from the trainer ranks over NIXL/RDMA, instead of the trainer pushing every tensor to every worker. Selected with `generator.inference_engine.weight_sync_backend=sharded_rdt`; non-colocated only (`placement.colocate_all=false`), Megatron or FSDP, and it forces `distributed_executor_backend=ray` because the workers dial named trainer actors. See - the dedicated section below for the capabilities it declares. - -Strategy choice is decided by the sender (`get_transfer_strategy_cls`). The init info is expanded per server via `for_servers()` / `to_api_payload()` and pushed to the servers through the HTTP control plane (`init_weight_update_communicator` → vLLM's native `/init_weight_transfer_engine`); the receive side is vLLM's native weight-transfer engine, driven by `NewInferenceWorkerWrap`. + the dedicated section below. + +Selection is `get_transfer_strategy(weight_sync_backend, colocate_all)`, called from two +places that must agree: `inference_servers/utils.build_vllm_cli_args` on the driver (via +`get_vllm_receive_backend`, to build the servers' `WeightTransferConfig`, and to force +`distributed_executor_backend=ray` for `sharded_rdt`), and `build_trainer_engine` on each +worker. Both read the same two config values, so the trainer and receive engines cannot +disagree. + +### Weight sources + +`sources.py` implements exactly vLLM's two-channel contract, and the two channels must +agree element for element — `metadata()` declares what iteration will yield, and the +engine sizes the worker's receive buffers and cuts its packed chunk boundaries from it. +`NCCLTrainerWeightTransferEngine._checked_iter` enforces that at runtime and names the +first divergent parameter; the Megatron GPU test +(`gpu_ci/megatron/test_megatron_weight_source.py`) is the regression check that gets +there first. + +- `FsdpWeightSource` — `state_dict()` for metadata (FSDP2 `DTensor.shape` is already the + global shape, so declaring costs no collective); iteration all-gathers via + `materialize_full_tensor`. `weight_prefix` handles syncing a CausalLM backbone into a + vLLM multimodal namespace. +- `MegatronWeightSource` — `bridge.export_hf_weights(conversion_tasks=None)`, a lazy + generator in HF-canonical order that gathers TP/PP/EP internally. `metadata()` must + materialize once to learn shapes, so it runs a dry export and caches. + +**There is no bucketing, and that is deliberate.** Bucketing does not bound memory — it +*accumulates* a whole bucket before handing it on, where the unbucketed export yields one +parameter at a time. What it was for, IPC handle count and Flash-RL fused-loader grouping, +`packed_ipc_producer`'s single reusable buffer subsumes. And one whole-model +`export_hf_weights` call satisfies `_accumulate_grouped_export`'s "every task of a +`group_key` in one call" requirement by construction, where bucketing has to special-case +it (splitting them means expert weights are silently never yielded). + +### Control plane + +`control_plane.SkyrlWeightSyncClient` is a **blocking** HTTP client over vLLM's native +RLHF routes, because `VLLMWeightSyncClient` is a synchronous protocol and the engine runs +off the event loop (`asyncio.to_thread(engine.send_weights)`). Four properties are +load-bearing: `Connection: close` (a full training step elapses between syncs, so any +keep-alive is stale), concurrent fan-out (required for *correctness* on the RDT path — a +serial `update_weights` deadlocks the producer's ref-counted group free), body-aware error +messages, and no timeout / no retry. + +Two per-server init rewrites live there, because each engine builds only **one** +worker-side init dict: + +- `nccl_init_payloads` — cumulative `rank_offset`, advancing one deployment's worth per + deployment and staying put across the DP servers *within* one (vLLM's + `data_parallel_index` already separates those). This is the highest-risk arithmetic in + weight sync: a wrong offset mis-maps ranks and **hangs in the NCCL rendezvous** rather + than erroring, so it cross-checks `world_size` against what the offsets imply. +- `rdt_init_payloads` — the deployment ordinal as `replica_rank` plus `num_replicas`, so + the engine can offset its consumer ids into globally distinct ranges. + +### Capability probes + +Three things the trainer engine cannot do for itself are decided by `getattr` probes on +the engine in `Worker._sync_weights_to_inference_engines`: + +| probe | default | who sets it | +|---|---|---| +| `skyrl_handles_prefix_cache_reset` | False | delta (it resets inside its own pause bracket) | +| `skyrl_force_disable_expandable_segments` | False | sharded_rdt (CUDA-IPC shares on every run, not only under colocation) | +| `skyrl_empty_cache_after_send` | True | sharded_rdt sets False (buffers are reused next step) | + +They are probes and not declared attributes because two of the four engines are vLLM's +own classes and cannot carry SkyRL attributes at all — so the *absence* of a flag is the +common case and must mean the default. `skyrl_set_reset_prefix_cache(bool)` is the same +idea for a per-round value, since `send_weights()` takes no arguments. ## Delta backend -Unlike the other two strategies, delta sync does not push tensors to the receiver at all — it +Unlike the push backends, delta sync does not send tensors to the receiver at all — it publishes bytes to `sync_dir` and the receiver pulls them. **Publish (trainer, rank 0).** `DeltaCheckpointPublisher` keeps a CPU `uint8` snapshot of the @@ -69,48 +161,75 @@ patch, and writes them as safetensors payload files plus a `manifest.json` under `/delta-/`. Unchanged tensors are omitted. Payload files roll over at `max_file_size_in_gb`. -**Fetch (receiver, before pause).** A control-plane operation the other strategies do not have: -`RemoteInferenceClient.fetch_weights` → `/fetch_weights` on every server, driven by -`DeltaWeightTransferSender._apply_receiver_update` *before* generation is paused, so the -download and patch-apply happen off the critical path. `LocalCheckpointStore` maintains a +**Fetch (receiver, before pause).** A control-plane operation the other backends do not have: +`SkyrlWeightSyncClient.fetch_weights` → `/fetch_weights` on every server, driven by +`DeltaTrainerWeightTransferEngine._apply_receiver_update` *before* generation is paused, so +the download and patch-apply happen off the critical path. It is also why the worker +extension still exists (see Architecture): `/collective_rpc` reaches worker methods only. + +`LocalCheckpointStore` maintains a mutable copy of the checkpoint under `local_checkpoint_dir/weights`, replaying every delta from its current version up to the target (so a late-joining engine catches up), and applies each patch by XOR-ing directly into the mmap'd safetensors files. **Reload.** Only then is generation paused and the local checkpoint reloaded into vLLM via `iter_tensors`. Note the delta shrinks the *transfer*, not the reload — the whole checkpoint is -re-read even for an empty delta. +re-read even for an empty delta. The pause/resume bracket and the prefix-cache reset ride the +engine, not the four-method client protocol: this is the only backend that reloads a whole +checkpoint in place, so it is the only one that needs generation stopped. + +The engine owns its layerwise-reload lifecycle like every other one +(`start_weight_update` → `initialize_layerwise_reload`, `finish_weight_update` → +`finalize_layerwise_reload`), which is what runs `process_weights_after_loading`. Its +drafter reload takes a **second** `iter_tensors` pass rather than the push backends' +proxy: materializing a whole checkpoint into a list so the drafter can re-read it would +put the entire model in memory at once. `DeltaWeightSyncConfig.__post_init__` derives `local_checkpoint_dir` and `publish_staging_dir` from `sync_dir` when unset, so consuming classes never invent their own defaults. ## Sharded RDT (`sharded_rdt`) -`sharded_rdt` is selected through `get_transfer_strategy_cls` like every other backend -(`ShardedRdtTransferStrategy`, in `weight_sync/sharded_rdt/sharded_rdt_strategy.py`), but underneath it -is vLLM's *trainer-send* model — a `WeightSource`, a `TrainerWeightTransferEngine` and a -`VLLMWeightSyncClient` — where the engine owns the whole round trip and the workers pull. -The strategy is a thin adapter over `RdtWeightSyncSender` (`weight_sync/sharded_rdt/rdt_send.py`), -which is where the real work lives. - -Where the pull model does not fit the push backends' shape, declared as capabilities -so the workers hold no backend conditional: +`sharded_rdt` is selected through `get_transfer_strategy` and built by the same +`build_trainer_engine` call as every other backend; nothing about its trainer side is +special. What *is* special is that it **pulls**, and pulling needs two +channels vLLM's `WeightSource` has no concept of. Those live on +`sharded_rdt_base.GroupedWeightSource` and stay inside this package +(`weight_sync/sources.py` is vLLM's contract verbatim, so a future upstream chunking +change arrives for free): -| flag | why | +| channel | what it really is | |---|---| -| `sender_initializes_receivers = True` | `trainer_init` opens the inference side itself (the bake needs the source metadata and must run under `set_current_vllm_config`), so worker rank 0 must **not** also call `init_weight_update_communicator`. | -| `force_disable_expandable_segments = True` | The sidecar shares gathered tensors over CUDA IPC on every run, not only under colocation. | -| `empty_cache_after_send = False` | Publish buffers are reused by the next training step; scrubbing them back to CUDA costs 0.25-0.53s/rank at 235B. The worker still empties under `colocate_all`. | - -The sender overrides `WeightTransferSender.send` rather than implementing `send_chunks`: -there is no chunk stream to push, and the base `send()` would call -`get_weight_metadata`, which on the Megatron extractor is a whole-model -`export_hf_weights` pass — the exact gather the RDT weight source exists to avoid. -`send_chunks` raises here. - -This is still intermediate: when the pinned vLLM ships trainer-send for NCCL/IPC too, -those backends collapse into this shape and the strategy layer goes away — the adapter is -what disappears, not `rdt_send.py`. +| `held_names()` | **ownership** under PP / EP. Feeds `_resolve_ownership` → `_spawn_server(held)`, and consumers route pulls to producers by name. Not a chunking concern, and not optional: the default (hold everything) is correct at pp=1/ep=1 and wrong above it. | +| `groups()` | the coordination index the producer's free barrier counts (`_inflight` is keyed by group). | +| `iter_groups()` | batching — driving the gather per group instead of per tensor turns ~37k generator resumes into ~95 on a per-expert MoE model (~0.9s/sync at 235B). | + +Three source flavors, chosen by `make_megatron_weight_source` / `make_fsdp_weight_source`: + +- `RdtFsdpWeightSource` — the shared FSDP source re-ordered **group-major**, so + `layerwise_groups` partitions `metadata()` exactly. The push backends do not care about + order (any permutation transfers identically as long as the two channels agree), which + is why only RDT pays for the reorder. +- `MegatronStackedWeightSource` — PP-local and EP-local, gathering experts at stack + granularity. `_pp_local_export_ctx` patches the bridge's PP broadcast so the owning + stage gets its tensor and every other stage gets `None` (which `megatron_to_hf` already + skips). Without it every rank materializes the whole model under PP; this is the actual + OOM fix, measured on Qwen3-32B at tp4/pp2 (70.56 GiB) and tp8/pp2 (73.06 GiB) of 79.18. +- `RdtMegatronWeightSource` — the whole-model fallback, taken on four conditions + (`_demoted`: a gather group spans PP stages, i.e. tied embeddings or MTP; grouped-export + archs; `etp != 1`; dense at `pp == 1`). + +**PP/EP locality cannot be shared with the push backends.** It needs per-rank ownership, +so only a pull backend can consume a PP-local source: NCCL broadcasts from rank 0, IPC +shares whole-tensor handles, and delta publishes a whole checkpoint — all three require +the cross-stage export. Conversely, RDT's whole-model residency is *the point* for a pull +backend (each producer must be able to serve its bound consumer the complete model), which +is why the Qwen3-32B OOM does not transfer to the push backends: they stream and drop. + +Two capability probes (see Architecture): `skyrl_force_disable_expandable_segments = True` +(the sidecar shares gathered tensors over CUDA IPC on every run, and VMM memory makes +export/rebuild 5-10x slower per storage) and `skyrl_empty_cache_after_send = False` +(publish buffers are reused by the next training step). Three kinds of process: the **trainer ranks** (each builds a `WeightSource` and a `ShardedRDTTrainerWeightTransferEngine`), one **producer sidecar** Ray actor per trainer @@ -212,10 +331,45 @@ waits on the same channel, so a sharer that never arrives fails the same way. The `sharded_rdt_*.py` files are vendored from the vLLM PR: github.com/vllm-project/vllm/pull/43375; see each file's header for the removal plan once the pinned vLLM ships the trainer-side ABCs natively. -## Lifecycle (`NewInferenceWorkerWrap`) -1. `start_weight_update(is_checkpoint_format=True)` — initializes layerwise reload (moves layers to meta device, wraps loaders). -2. `update_weights_chunk(update_info)` — called repeatedly. Unpacks the SkyRL packed CUDA-IPC payload, slices the contiguous buffer per param, calls `model.load_weights(weights=...)` under `set_current_vllm_config`. -3. `finish_weight_update()` — runs `finalize_layerwise_reload` (quantization repacking, attention weight postprocessing). +## Receive-side lifecycle + +vLLM's native routes, driven by the trainer engine through `SkyrlWeightSyncClient`: + +1. `/init_weight_transfer_engine` → `engine.init_transfer_engine(...)`. The **one** + lifecycle method `GPUWorker` does not wrap in `set_current_vllm_config`, which is why + the RDT engine opens that context itself (its bake drives `model.load_weights` against + meta params, and `process_weights_after_loading` on MoE models reads + `get_current_vllm_config()` to build kernels). +2. `/start_weight_update` → `engine.start_weight_update()` → `initialize_layerwise_reload` + (moves layers to meta device, wraps loaders). +3. `/update_weights` → `engine.update_weights(...)`, once or many times. The actual + receive; wrapped in `disable_mtp_completeness_check()` because MTP architectures raise + on incomplete layer coverage. +4. `/finish_weight_update` → `engine.finish_weight_update()` → `finalize_layerwise_reload` + (quantization repacking, attention weight postprocessing), then + `model_runner.reset_lora_state()` on the worker. + +`GPUWorker` opens `set_current_vllm_config` around steps 2-4 itself, which is why the +receive path needs no SkyRL wrapper. + +### Spec-decode drafter reload + +vLLM's engines call `self.model.load_weights(...)` directly and there is still no +`load_weights` callback, so `skyrl_engines.SkyrlDrafterReloadMixin` swaps `self.model` for +a `_LoadWeightsProxy` for the duration of `receive_weights` — the drafter +(`model_runner.drafter.model`, a separate module the main load never touches) is then +reloaded from exactly the weights the main model just received. The proxy is only +installed when this process actually *has* a drafter, so a non-MTP deployment runs vLLM's +path verbatim. + +Delta is the exception: it re-streams `iter_tensors` a second time instead, because the +proxy has to materialize the weight list so the drafter can re-read it, and for a whole +checkpoint that is the entire model resident at once. + +An engine has no route to its worker — it is constructed inside `Worker.load_model`, and +the worker-extension class is appended to `Worker.__bases__` *after* `Worker` — so +`patches/vllm/patch_model_runner_registry.py` wraps `GPUModelRunner.load_model` to record +the runner in a process-global weakref. ## KV offload during non-colocated weight sync @@ -233,19 +387,37 @@ Validated in `validate_inference_engine_cfg`. vLLM-version coupled (mirrors `GPU ## Convention: vLLM imports -`vllm` is a Linux-only optional dep. Import it **lazily inside methods**, not at module top. Match the existing pattern in `new_inference_worker_wrap.py`. +`vllm` is a Linux-only optional dep, and half the CPU suite runs without the wheel. Two +rules, not one: + +- A module that **anything** vllm-free imports must import `vllm` **lazily inside + methods**: `weight_sync/__init__.py`, `base.py`, `delta_checkpoint.py`, + `delta_engine.py`, `control_plane.py` and `new_inference_worker_wrap.py` all follow this. +- A module that is *itself* vllm-only may import at module top — `sources.py`, + `trainer_engines.py`, `delta_trainer.py`, `sharded_rdt_base.py`, + `sharded_rdt_{engine,trainer}.py`. The rule then moves up: **nothing vllm-free may + import them at module scope.** The workers import `sources` / `trainer_engines` inside + `_build_weight_source` / `init_weight_sync_state` for exactly this reason, and + `weight_sync/__init__.py` re-exports neither. + +Tests for the second group carry `pytest.importorskip("vllm")` plus +`pytestmark = pytest.mark.vllm`, which is what puts them in the `-m "vllm"` CPU half. ## Tests ```bash -# CPU — chunk packing, transfer strategies, and the sharded_rdt pull plan / -# producer sidecar / weight source / control plane +# CPU — the control plane's per-server init rewrites, the delta publisher, and the +# sharded_rdt pull plan / producer sidecar / grouped source contract uv run --extra dev --extra fsdp pytest tests/backends/skyrl_train/weight_sync/ -v -# GPU — end-to-end weight sync (NCCL + CUDA IPC paths, TP=1 and TP=2) +# GPU — end-to-end weight sync: all four backends through build_trainer_engine uv run --isolated --extra dev --extra fsdp \ pytest tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py -v +# GPU — the Megatron source's two channels must agree (what _checked_iter enforces at runtime) +uv run --isolated --extra dev --extra megatron \ + pytest tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_weight_source.py -v + # GPU — end-to-end delta sync (sparse perturbation, fsdp and megatron) uv run --isolated --extra dev --extra fsdp \ pytest tests/backends/skyrl_train/gpu/gpu_ci/test_delta_weight_sync_e2e.py -m "not megatron" -v @@ -259,43 +431,54 @@ The CPU tests do **not** import `NewInferenceWorkerWrap`. Any change to the work | Change | Run | |--------|-----| -| `WeightChunk` packing / size accounting | `tests/backends/skyrl_train/weight_sync/test_weight_chunk.py` | -| Broadcast or CUDA IPC sender | `test_transfer_strategies.py` (CPU) **and** GPU `test_weight_sync.py` | +| A `WeightSource` (`sources.py` or `rdt_send.py`) | `test_sources.py` + `test_sharded_rdt_source.py` (CPU), GPU `test_weight_sync.py`, and — for Megatron — GPU `test_megatron_weight_source.py` | +| `control_plane.py`, especially the init rewrites | `test_control_plane.py` (CPU) **and** GPU `test_weight_sync.py` | +| `trainer_engines.py` / a trainer engine | `test_trainer_engines.py` (CPU) **and** GPU `test_weight_sync.py` | +| `skyrl_engines.py` (receive side) | GPU `test_weight_sync.py` only — it runs inside the vLLM worker | | `NewInferenceWorkerWrap` | GPU `test_weight_sync.py` (CPU tests will not catch regressions) | | Delta publish / manifest / payload format | `test_delta_checkpoint.py` **and** GPU `test_delta_weight_sync_e2e.py` | | `LocalCheckpointStore` (fetch, replay, apply, cache keys) | `test_delta_checkpoint.py` | -| `DeltaWeightTransferEngine` | GPU `test_delta_weight_sync_e2e.py` only — it runs inside the vLLM worker | +| `DeltaWeightTransferEngine` (receive side) | GPU `test_delta_weight_sync_e2e.py` only — it runs inside the vLLM worker | | Who pauses / resets the prefix cache | `test_prefix_cache_reset.py` **and** `distributed/test_worker_dispatch.py` | | `DeltaWeightSyncConfig` defaults or validation | `tests/train/test_config.py` | ## vLLM version coupling -`vllm` is pinned in `pyproject.toml` (currently `0.28.0`, in both the `fsdp` and `megatron` extras). Weight-sync code paths are tightly coupled to vLLM internals (`model_runner.load_weights`, `initialize_layerwise_reload`, `SKIP_LOAD_TENSORS`). When bumping the pin, re-verify the GPU weight-sync tests. - -### `packed` is an init-time wire param (0.28.0+) - -Through 0.26.0 the NCCL sender put `packed` on the **per-round** update info. 0.28.0 moved it -to the **init** info: `NCCLWeightTransferInitInfo.packed` is recorded once during -`/init_weight_transfer_engine`, `receive_weights` reads `self.packed`, and -`NCCLWeightTransferUpdateInfo` now rejects the key outright. - -So `BroadcastInitInfo.packed` must be sent by `to_api_payload()` and must agree with the -`packed` handed to `nccl_trainer_send_weights`. Its default is `False` on the vLLM side, so -forgetting it does not raise — the trainer packs while the worker unpacks one-by-one, and the -transfer hangs inside NCCL. Only the GPU weight-sync test catches this. - -### Vendored trainer-send (`weight_sync/nccl_trainer_send.py`) - -0.28.0 deleted `NCCLWeightTransferEngine.trainer_send_weights` and the engine's `trainer_init` -staticmethod, replacing them with a trainer-side engine abstraction -(`NCCLTrainerWeightTransferEngine`, via `WeightTransferTrainerFactory`). SkyRL still drives the -send itself, so `nccl_trainer_send.py` keeps the 0.26.0 behaviour: it wraps -`packed_nccl_broadcast_producer` and re-exports `nccl_common.trainer_init`, both of which -survive in 0.28.0 unchanged. Delete that module when SkyRL's sender layer migrates onto vLLM's -trainer-send engines. - -The `sharded_rdt` backend is unaffected: it vendors its own trainer-side ABCs in -`sharded_rdt_base.py`, which already match 0.28.0's `Generic[TTrainerInitInfo]` shape. +`vllm` is pinned in `pyproject.toml` (currently `0.28.0`, in both the `fsdp` and `megatron` extras). Weight-sync code paths are tightly coupled to vLLM internals. When bumping the pin, re-verify the GPU weight-sync tests and the extension points listed below. + +### `packed` is chosen by SkyRL, agreed by vLLM + +`packed` rides the **init** info, and the trainer engine propagates it to the worker at +`trainer_init`, so the two sides structurally cannot disagree. What SkyRL chooses is +`packed=True` on both push backends: + +- NCCL: packed broadcasts out of a fixed reusable buffer instead of one NCCL call per + parameter. vLLM's unpacked path is only reached by its own tests. +- IPC: `packed=True` is **not** vLLM's default but is required here. The unpacked path + holds a strong ref to a contiguous copy of *every* parameter until past + `finish_weight_update` (so the consumer's IPC views stay valid) — i.e. the whole model + resident on the trainer. Packed streams through one reusable buffer, and the consumer + clones out of it. + +### Extension points relied on + +- `WeightTransferEngineFactory` / `WeightTransferTrainerFactory` — `register_engine` + raises `ValueError` on a duplicate name, and `WeightTransferConfig.backend` is typed + `Literal[...] | str` and validated against the registry at engine creation. +- `TrainerInitInfo` — `backend` is a `ClassVar` the factory dispatches on; `rank` is + keyword-only. +- `VLLMWeightSyncClient` — a **structural** (PEP 544) four-method protocol, so + `SkyrlWeightSyncClient` needs no import or subclassing. Note 0.28 declares + `finish_weight_update(weight_version: str | None = None)`. +- `set_weight_update_target` / `reset_weight_update_target` on `WeightTransferEngine` — + the draft-session hook. SkyRL's proxy swap composes with it by restoring the exact + object it found. +- `GPUModelRunner.load_model` — wrapped by `patch_model_runner_registry`. +- `CuMemAllocator` via `vllm.device_allocator.get_mem_allocator_instance` — the KV-offload + path drives it directly because `EngineCore.sleep` hardcodes + `clear_prefix_cache = level >= 1` and `CuMemBackend.suspend` cannot express "discard + weights, offload kv_cache". Registering a custom `SleepModeBackend` does not help: the + problem is on the dispatch path, not in the suspend mechanism. ### DeepGEMM is unavailable under the torch override @@ -310,8 +493,8 @@ when those wheels publish torch 2.13 builds. ## Gotchas -- After `update_weights_chunk` runs, call `torch.accelerator.synchronize()` before returning so the sender doesn't drop its packed buffer mid-copy on the next barrier. -- Delta: `DeltaWeightTransferEngine` is registered as an **import side effect** of `new_inference_worker_wrap.py`, which is the module vLLM loads via `--worker-extension-cls`. Registering anywhere else (e.g. while building CLI args in the driver) is a no-op — it has to happen in the process that owns the engine. +- Every receive-side engine is registered as an **import side effect** of `new_inference_worker_wrap.py`, the module vLLM loads via `--worker-extension-cls`. Registering anywhere else (e.g. while building CLI args in the driver) is a no-op for the workers — it has to happen in the process that owns the engine. The driver registers too, separately, because it validates `WeightTransferConfig.backend` against the registry. +- `send_weights()` must be called on **every** trainer rank, and every rank must drain the source: iterating it is what drives the gather collectives, so a rank that skipped it deadlocks its peers. Only rank 0 touches the wire (the engines resolve `is_sender` from `init_info.rank`, never from a global process group, which is ambiguous once FSDP/TP/PP/EP groups exist). - Delta: the receive-side `delta checkpoint fetch:` / `receive reload-only:` log lines are emitted inside the nested vLLM worker process and do **not** reach the driver log, even with `SKYRL_DUMP_INFRA_LOG_TO_STDOUT=1`. Find them with `grep -rhE "delta checkpoint (fetch|receive)" /tmp/ray/session_latest/logs/`. Filter by mtime — that directory accumulates lines from earlier runs. - Delta: `_safe_path_name` appends a digest of the full value because sibling delta URIs differ only in their trailing `delta-`; a plain length cap collapses every version onto one cache directory. Don't "simplify" it back to truncation. - Delta: `s3://` needs the `s5cmd` CLI, which the `aws` extra installs into the run's venv (`--extra aws`). `gs://` needs the `gcloud` CLI as a *system* binary — the `gcp` extra only provides a Python library and will not satisfy it. diff --git a/docs/content/docs/getting-started/inference_architecture.mdx b/docs/content/docs/getting-started/inference_architecture.mdx index 828683862c..5235ad7303 100644 --- a/docs/content/docs/getting-started/inference_architecture.mdx +++ b/docs/content/docs/getting-started/inference_architecture.mdx @@ -102,29 +102,31 @@ SkyRL uses the native weight syncing APIs in vLLM, with the following four-stage - `POST /update_weights` — Updates all or a subset of the weights. SkyRL uses chunked weight transfer for efficiency. - `POST /finish_weight_update` — Finishes the current weight update. - -For colocated training, SkyRL currently uses chunked transfers with CUDA IPC handles and currently implements a [custom `Worker` extension](https://github.com/NovaSky-AI/SkyRL/blob/fb87f35dfe9f74f71445da60c79549773b15ba5e/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py) as a transitional implementation, pending [vllm-project/vllm#39212](https://github.com/vllm-project/vllm/issues/39212). We plan to migrate to the native APIs in the next vLLM release. - +The trainer side runs on vLLM's *trainer-send* abstraction: each training worker builds a `WeightSource` over its live model and hands it to a `TrainerWeightTransferEngine`, whose `send_weights()` drives the whole four-stage round trip above. See `skyrl/backends/skyrl_train/weight_sync/trainer_engines.py`. + +Four backends, selected by `generator.inference_engine.weight_sync_backend` plus `trainer.placement.colocate_all`: -SkyRL implements two transfer strategies in `skyrl/backends/skyrl_train/weight_sync/`: +- **`nccl`** — non-colocated. Tensor data is broadcast over NCCL from trainer rank 0 to all inference workers, concurrently with the `/update_weights` HTTP call that ships the metadata (both sides rendezvous inside the same NCCL calls). Combined with `/pause?mode=keep` and `/resume` so in-flight rollouts pause correctly during the sync. +- **`ipc`** — colocated, where the trainer and inference engines share GPUs. Weights are packed into one reusable buffer and exchanged via CUDA IPC handles. Combined with `/sleep` and `/wake_up` for memory management — the inference engine sleeps to free VRAM during training, then wakes for rollouts. +- **`delta`** — weights travel as compressed deltas against the base checkpoint through a shared filesystem or object store, for setups where the two sides are not NCCL-reachable. +- **`sharded_rdt`** — the inference workers *pull* the slices they consume from the trainer ranks over NIXL/RDMA. -- **`BroadcastTransferStrategy`** (`broadcast_strategy.py`) — used for non-colocated training. Tensor data is broadcast over NCCL from trainer rank 0 to all inference workers, concurrently with `/update_weights` HTTP calls that ship the metadata. This is used in combination with `/pause?mode=keep` and `/resume` so that in-flight rollouts are paused correctly during the sync. -- **`CudaIpcTransferStrategy`** (`cuda_ipc_strategy.py`) — used for colocated training where the trainer and inference engines share GPUs. Weights are exchanged via CUDA IPC handles. Combined with `/sleep` and `/wake_up` for memory management — the inference engine sleeps to free VRAM during training, then wakes for rollouts. + +The receive side is a vLLM `WeightTransferEngine` in each inference worker. SkyRL registers its own under `skyrl_nccl` / `skyrl_ipc`, subclassing vLLM's to also reload the speculative-decoding drafter (a separate module the main model's `load_weights` never touches). A small [`Worker` extension](https://github.com/NovaSky-AI/SkyRL/blob/main/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py) remains for the two things an engine cannot reach: the checkpoint-delta `fetch_weights` call, and driving the allocator directly for KV offload. + Endpoint summary: | Endpoint | Plane | Purpose | |----------|-------|---------| | `/init_weight_transfer_engine` | fan-out | One-time communicator setup | -| `/start_weight_update` * | fan-out | Begin a chunked update | -| `/update_weights` * | fan-out | Send a tensor chunk | -| `/finish_weight_update` * | fan-out | Commit the update | +| `/start_weight_update` | fan-out | Begin a chunked update | +| `/update_weights` | fan-out | Send a tensor chunk | +| `/finish_weight_update` | fan-out | Commit the update | | `/pause`, `/resume` | fan-out | Generation control | | `/sleep`, `/wake_up` | fan-out | Colocated memory management | | `/v1/completions`, `/v1/chat/completions`, generate | routed | Generation | -\* Currently we use custom `/collective_rpc` + worker methods that mimick the native APIs because SkyRL makes some vLLM fixes for Qwen 3.6 model loading. We will migrate to the native `/start_weight_update` and `/finish_weight_update` soon. - ## End-to-end Weight Sync Flow **Non-colocated mode (NCCL broadcast):** diff --git a/examples/train/rlm/openrouter_client.py b/examples/train/rlm/openrouter_client.py index 9cf74eed4b..50ec6140b5 100644 --- a/examples/train/rlm/openrouter_client.py +++ b/examples/train/rlm/openrouter_client.py @@ -182,9 +182,6 @@ async def sleep(self, *args: Any, **kwargs: Any) -> None: async def wake_up(self, *args: Any, **kwargs: Any) -> None: pass - async def init_weight_update_communicator(self, *args: Any, **kwargs: Any) -> None: - pass - async def update_named_weights(self, *args: Any, **kwargs: Any) -> None: pass diff --git a/skyrl/backends/skyrl_train/inference_servers/base.py b/skyrl/backends/skyrl_train/inference_servers/base.py index aac9ee6f6d..236e8e8122 100644 --- a/skyrl/backends/skyrl_train/inference_servers/base.py +++ b/skyrl/backends/skyrl_train/inference_servers/base.py @@ -4,10 +4,7 @@ from skyrl.backends.skyrl_train.utils.routed_experts import RoutedExpertIndices if TYPE_CHECKING: - from skyrl.backends.skyrl_train.weight_sync import WeightUpdateRequest - from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( - WeightSyncInitInfo, - ) + from skyrl.backends.skyrl_train.weight_sync import LoraLoadRequest MessageType = Dict[str, str] ConversationType = List[MessageType] @@ -125,19 +122,16 @@ async def sleep(self, *args: Any, **kwargs: Any): raise NotImplementedError @abstractmethod - async def init_weight_update_communicator(self, init_info: "WeightSyncInitInfo"): - """Initialize weight update communicator from init info. + async def update_named_weights(self, request: "LoraLoadRequest | Dict[str, Any]"): + """Load weights the engine can reach itself, rather than transferring them. - Args: - init_info: WeightSyncInitInfo from the sender containing all info needed - to create the appropriate receiver. + The only caller is the LoRA path, which passes a + :class:`LoraLoadRequest` naming an adapter directory on disk. Tensor + transfer goes through the trainer-side engines (see + ``weight_sync/trainer_engines.py``). """ raise NotImplementedError() - @abstractmethod - async def update_named_weights(self, request: "WeightUpdateRequest"): - raise NotImplementedError() - @abstractmethod async def teardown(self): raise NotImplementedError diff --git a/skyrl/backends/skyrl_train/inference_servers/layerwise_reload.py b/skyrl/backends/skyrl_train/inference_servers/layerwise_reload.py deleted file mode 100644 index 1a7406e67e..0000000000 --- a/skyrl/backends/skyrl_train/inference_servers/layerwise_reload.py +++ /dev/null @@ -1,188 +0,0 @@ -"""Shared vLLM layerwise-reload lifecycle for SkyRL's vLLM worker-extension classes. - -Provides `LayerwiseReloadWorkerMixin`, the start/finish bracket that -`new_inference_worker_wrap.NewInferenceWorkerWrap` uses to run vLLM's layerwise -reload once per weight sync rather than once per chunk. -""" - -import inspect -from collections.abc import Callable -from typing import TYPE_CHECKING - -import torch - -if TYPE_CHECKING: - from vllm.config import ModelConfig, VllmConfig - from vllm.v1.worker.gpu_model_runner import GPUModelRunner - - -def get_numel_loaded(weight_loader: Callable, args: inspect.BoundArguments) -> tuple[int, object]: - """ - Determine how many elements would be loaded by a weight loader call. - - Args: - weight_loader: used to load weights - args: bound arguments to weight loader - - Returns: - number of elements loaded by the weight loader, the return value of the - weight loader - """ - # Lazy import: vllm is a Linux-only optional dependency, so this module stays importable on macOS / CI. - from vllm.model_executor.model_loader.reload.meta import CopyCounter - - with CopyCounter() as counter: - return_value = weight_loader(*args.args, **args.kwargs) - - # A weight loader fills a single destination parameter, so the number of - # loaded elements is at most that parameter's size. Some loaders copy into - # the parameter more than once -- e.g. ``composed_weight_loader`` runs an - # in-place post-load transform (``param.copy_(fn(param))``) on top of the - # initial copy -- which would make CopyCounter report twice the parameter - # size. Over-counting inflates the layer's loaded-element total and can - # finalize the layer before every parameter is loaded, silently dropping - # the trailing parameter(s) (e.g. Mamba ``mixer.D``). Cap the count at the - # destination size to keep the per-layer accounting correct. - numel = counter.copied_numel - param = args.arguments.get("param", None) - if isinstance(param, torch.Tensor): - numel = min(numel, param.numel()) - return numel, return_value - - -def patch_numel_loaded(): - # vLLM's layerwise reload binds get_numel_loaded at import time - # (`from .meta import get_numel_loaded`), so its call site at - # layerwise.py uses the `layerwise` module's own binding. Rebind that - # attribute to our patched version to substitute the symbol. - from vllm.model_executor.model_loader.reload import layerwise as _layerwise - from vllm.model_executor.model_loader.reload import meta as _meta - - _layerwise.get_numel_loaded = get_numel_loaded - _meta.get_numel_loaded = get_numel_loaded - - -_PATCHED_LAYERWISE_NUMEL_LOADED = False - - -def _empty_cuda_cache_rocm() -> None: - """Release unused ROCm cached blocks after full-weight sync.""" - is_rocm = torch.version.hip is not None - if not torch.cuda.is_available() or not is_rocm: - return - - device = torch.cuda.current_device() - torch.cuda.synchronize(device) - torch.cuda.empty_cache() - torch.cuda.synchronize(device) - - -class LayerwiseReloadWorkerMixin: - """Bracket a multi-chunk weight sync with one vLLM layerwise-reload init/finalize. - - `skyrl_start_weight_update` initializes the layerwise reload once; each chunk then loads - its weights raw; `skyrl_finish_weight_update` finalizes once over the whole weight set. - A per-chunk `reload_weights` is the wrong approach: it re-finalizes on every call - and restores layers absent from that chunk, corrupting a multi-chunk sync. - """ - - vllm_config: "VllmConfig" - model_runner: "GPUModelRunner" - model_config: "ModelConfig" - device: torch.device - - # NOTE: named with a `skyrl_` prefix to avoid colliding with vLLM's own - # Worker.start_weight_update / finish_weight_update (added in vllm-project/vllm - # #39212, merge e3b65a5, shipped in vLLM 0.22.0+). vLLM injects the - # worker-extension class as a *base* of Worker and asserts the extension - # defines no attribute already present on Worker, so same-named methods abort - # engine init. The skyrl_-prefixed variants keep SkyRL's IPC weight-sync path - # (and the MoE set_current_vllm_config wrapping) intact alongside vLLM's native API. - def skyrl_start_weight_update(self, is_checkpoint_format: bool = True) -> None: - """ - Prepare the model for a new weight update. - - For checkpoint-format weights, initializes the layerwise reload - machinery which moves layers to meta device and wraps weight loaders - to defer processing until all weights for each layer are loaded. - - Must be called before any update_weights_ipc calls. - - Args: - is_checkpoint_format: True if incoming weights are in checkpoint - format (need layerwise processing). False if weights are - already in kernel format (direct copy). - """ - if getattr(self, "_skyrl_weight_update_active", False): - raise RuntimeError( - "skyrl_start_weight_update called while a weight update is " - "already active. Call skyrl_finish_weight_update first." - ) - if getattr(self, "_weight_update_active", False): - raise RuntimeError("vLLM native weight update is already active. Call finish_weight_update first.") - - # Ensure the get_numel_loaded patch is in effect before layerwise - # reload runs. - global _PATCHED_LAYERWISE_NUMEL_LOADED - if not _PATCHED_LAYERWISE_NUMEL_LOADED: - # use patched version, based on https://github.com/vllm-project/vllm/pull/44814 - patch_numel_loaded() - _PATCHED_LAYERWISE_NUMEL_LOADED = True - - if is_checkpoint_format: - # Lazy import: vllm is a Linux-only optional dependency, so this module stays importable on macOS / CI. - from vllm.config import set_current_vllm_config - from vllm.model_executor.model_loader.reload import ( - initialize_layerwise_reload, - ) - - model = self.model_runner.model - with set_current_vllm_config(self.vllm_config), torch.device(self.device): - initialize_layerwise_reload(model) - - self._skyrl_is_checkpoint_format = is_checkpoint_format - self._skyrl_weight_update_active = True - # vLLM's native /update_weights endpoint checks these flags before - # calling the configured WeightTransferEngine. Mirroring them lets - # SkyRL keep its patched layerwise start/finish while using native - # update_weights for transports such as checkpoint-delta. - self._is_checkpoint_format = is_checkpoint_format - self._weight_update_active = True - - def skyrl_finish_weight_update(self) -> None: - """ - Finalize the current weight update. - - For checkpoint-format weights, runs layerwise postprocessing - (quantization repacking, attention weight processing, etc.). - Must be called after all update_weights_ipc calls are done. - """ - if not getattr(self, "_skyrl_weight_update_active", False): - raise RuntimeError("skyrl_start_weight_update must be called before skyrl_finish_weight_update.") - - # The sharded_rdt engine defers its GPU post-processing (scatter/quant/ - # kernel-copy) to background threads during update, so drain it here — - # before finalize, which needs every layer fully loaded + reset. No-op - # for the ipc/nccl engines (they process synchronously per chunk). - engine = getattr(self, "weight_transfer_engine", None) - if engine is not None and getattr(engine, "defers_processing", False): - drain_pending = getattr(engine, "drain_pending", None) - if drain_pending is not None: - drain_pending() - - if self._skyrl_is_checkpoint_format: - # Lazy import: vllm is a Linux-only optional dependency, so this module stays importable on macOS / CI. - from vllm.config import set_current_vllm_config - from vllm.model_executor.model_loader.reload import ( - finalize_layerwise_reload, - ) - - model = self.model_runner.model - with set_current_vllm_config(self.vllm_config), torch.device(self.device): - finalize_layerwise_reload(model, self.model_config) - - self._skyrl_weight_update_active = False - self._skyrl_is_checkpoint_format = True - self._weight_update_active = False - self._is_checkpoint_format = True - _empty_cuda_cache_rocm() diff --git a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py index 9be72ae9ab..b256bd2220 100644 --- a/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py +++ b/skyrl/backends/skyrl_train/inference_servers/new_inference_worker_wrap.py @@ -1,21 +1,27 @@ -""" -vLLM Worker Extension for native weight sync with chunked transfer support. - -This module provides NewInferenceWorkerWrap, a vLLM worker extension that -enables chunked weight updates from training to inference using the -start/update/finish lifecycle: - - skyrl_start_weight_update -> one or more update_weights_ipc -> skyrl_finish_weight_update - -This separates the layerwise reload initialization/finalization from individual -chunk transfers, allowing weights to be sent in bounded-memory chunks rather -than all at once. - -TODO: Once https://github.com/vllm-project/vllm/pull/39212 lands, vLLM will -natively support start_weight_update / update_weights / finish_weight_update -on GPUWorker with dedicated HTTP endpoints. At that point this worker extension -can be removed and SkyRL can call the native endpoints directly instead of -routing through /collective_rpc. +"""vLLM worker extension for the two weight-sync things an engine cannot do. + +Weight transfer itself does not route through here: the receive path is an engine +subclass (``weight_sync/skyrl_engines.py``, ``weight_sync/delta_engine.py``, +``weight_sync/sharded_rdt/sharded_rdt_engine.py``) driven over vLLM's native RLHF +routes, which wrap ``set_current_vllm_config`` themselves and give each engine +its own layerwise-reload lifecycle. + +What remains are two limits of *dispatch*: + +``fetch_weights`` + ``/collective_rpc`` dispatches to worker methods by name and refuses + callables (``entrypoints/serve/dev/rpc/api_router.py``), and no native route + can invoke an arbitrary engine method. SkyRL's ``/fetch_weights`` route + (``vllm_server_actor``) collective-RPCs into this method. It is called + *before* ``pause_generation`` so the checkpoint-delta download overlaps live + generation. + +sleep / wake + ``EngineCore.sleep`` hardcodes ``clear_prefix_cache = level >= 1`` + (``v1/engine/core.py``) with no parameter, and ``CuMemBackend.suspend`` maps + level to tags as ``("weights",)`` / ``()`` with no way to express "discard + weights, offload kv_cache". A custom ``SleepModeBackend`` does not help: the + problem is on the dispatch path, not in the suspend mechanism. Usage: Pass as --worker-extension-cls to vLLM: @@ -24,12 +30,36 @@ skyrl.backends.skyrl_train.inference_servers.new_inference_worker_wrap.NewInferenceWorkerWrap """ +from typing import TYPE_CHECKING + import torch -from skyrl.backends.skyrl_train.inference_servers.layerwise_reload import ( - LayerwiseReloadWorkerMixin, - _empty_cuda_cache_rocm, -) +if TYPE_CHECKING: + from vllm.config import ModelConfig, VllmConfig + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +# Everything below must run inside EVERY vLLM worker process: Worker.load_model +# builds the weight-transfer engine through the factory, and the model-runner +# recorder must be installed before load_model runs. vLLM loads this module +# before model init, which is what guarantees it. Each is guarded because this +# module is also imported from processes without the optional deps. +try: + from skyrl.backends.skyrl_train.patches.vllm.patch_model_runner_registry import ( + apply_model_runner_registry_patch, + ) + + apply_model_runner_registry_patch() +except ModuleNotFoundError: + pass + +try: + from skyrl.backends.skyrl_train.weight_sync.skyrl_engines import ( + register_skyrl_engines, + ) + + register_skyrl_engines() +except ModuleNotFoundError: + pass try: from skyrl.backends.skyrl_train.weight_sync.delta_engine import ( @@ -40,12 +70,6 @@ except ModuleNotFoundError: pass -# Registering the sharded_rdt engine into vLLM's WeightTransferEngineFactory must -# happen inside every worker process (GPUWorker.load_model builds the engine via -# the factory). Importing here — the worker-extension module vLLM loads before -# model init — guarantees it runs on each worker. Guarded like the delta engine -# above: this module is also imported from processes without the RDT dependencies -# (e.g. a trainer process), and a missing optional dep must not break them. try: from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import ( rdt_vllm_register, # noqa: F401 @@ -58,39 +82,17 @@ VLLM_NEW_INFERENCE_WORKER_EXTENSION_CLS = f"{__name__}.NewInferenceWorkerWrap" -class _LoadWeightsProxy: - """Wraps a model, overriding only ``load_weights``. +class NewInferenceWorkerWrap: + """The weight-sync methods that must live on the worker, not an engine. - vLLM's weight transfer engines call ``self.model.load_weights(...)`` - internally (as of 0.26 there is no injectable callback). Handing them this - proxy via ``set_weight_update_target`` lets SkyRL interpose its own loader - while every other attribute access falls through to the real model. + Attributes come from the host GPUWorker: vLLM appends this class to + ``Worker.__bases__``. """ - def __init__(self, model, load_weights): - self._model = model - self.load_weights = load_weights - - def __getattr__(self, name): - # Only reached for attributes not set on the proxy itself. - return getattr(self._model, name) - - -class NewInferenceWorkerWrap(LayerwiseReloadWorkerMixin): - """ - vLLM worker extension for chunked weight sync (new inference path). - - Provides a three-phase weight update protocol via collective_rpc: - 1. skyrl_start_weight_update: Prepare model for receiving weights - 2. update_weights_ipc: Receive and load one chunk of weights - 3. skyrl_finish_weight_update: Finalize the model after all chunks - - Attributes accessed from the host GPUWorker (via mixin inheritance): - self.weight_transfer_engine - self.model_runner - self.model_config - self.device - """ + vllm_config: "VllmConfig" + model_runner: "GPUModelRunner" + model_config: "ModelConfig" + device: torch.device def fetch_weights(self, target_version: int, sync_dir: str | None = None, uri: str | None = None): """Fetch/apply a checkpoint delta before the paused reload phase.""" @@ -103,157 +105,15 @@ def fetch_weights(self, target_version: int, sync_dir: str | None = None, uri: s raise RuntimeError(f"{type(self.weight_transfer_engine).__name__} does not support fetch_weights") return fetch(target_version=target_version, sync_dir=sync_dir, uri=uri) - def update_weights_ipc(self, update_info: dict) -> None: - """ - Receive and load a single chunk of weights. - - SkyRL packs each chunk's tensors into a single contiguous CUDA buffer and sends - one IPC handle per rank plus per-param `sizes` metadata. We rebuild - the packed tensor here, slice it per param, and hand the list to - model.load_weights (checkpoint format) or copy per-param directly - (kernel format). - - Args: - update_info: Dict with keys: - - names: list[str] - - dtype_names: list[str] - - shapes: list[list[int]] - - sizes: list[int] (element count per param; used for slicing) - - ipc_handles_pickled: b64(pickle({gpu_uuid: (func, args)})) - """ - if not getattr(self, "_skyrl_weight_update_active", False): - raise RuntimeError("skyrl_start_weight_update must be called before update_weights_ipc.") - - if self.weight_transfer_engine is None: - raise RuntimeError( - "Weight transfer not configured. " "Please set weight_transfer_config to enable weight transfer." - ) - - # --- unpack SkyRL packed CUDA IPC format --- - import base64 - import pickle - - names = update_info["names"] - shapes = update_info["shapes"] - sizes = update_info["sizes"] - pickled = update_info["ipc_handles_pickled"] - handles = pickle.loads(base64.b64decode(pickled)) - - device_index = torch.cuda.current_device() - physical_gpu_id = str(torch.cuda.get_device_properties(device_index).uuid) - if physical_gpu_id not in handles: - raise ValueError(f"IPC handle not found for GPU UUID {physical_gpu_id}. " f"Available: {list(handles)}") - func, args = handles[physical_gpu_id] - # Remap device index to the LOCAL current-device. - list_args = list(args) - list_args[6] = device_index - packed_tensor = func(*list_args) - - weights: list[tuple[str, torch.Tensor]] = [] - offset = 0 - for name, shape, size in zip(names, shapes, sizes): - weights.append((name, packed_tensor[offset : offset + size].view(*shape))) - offset += size - - # process_weights_after_loading reads get_current_vllm_config() (e.g. - # flashinfer_cutlass_moe needs the compilation config to build kernels), - # and vllm only sets that context around init_device / load_model. - from vllm.config import set_current_vllm_config - - model = self.model_runner.model - with set_current_vllm_config(self.vllm_config), torch.device(self.device): - if self._skyrl_is_checkpoint_format: - model.load_weights(weights=weights) - # vLLM's load only updates the main model; the spec-decode (MTP/Eagle) - # drafter is a separate module and must be reloaded from the same - # checkpoint-format weights (see spec_decode_utils). - from skyrl.backends.skyrl_train.inference_servers.spec_decode_utils import ( - _reload_spec_decode_drafter, - ) - - _reload_spec_decode_drafter(self.model_runner, weights) - else: - for name, weight in weights: - param = model.get_parameter(name) - param.copy_(weight) - - # Ensure consumption of packed_tensor finishes before we return (and - # before the sender drops its reference on the next barrier). - torch.accelerator.synchronize() - - def update_weights_nccl(self, update_info: dict) -> None: - """ - Receive a batched weight update via vLLM's NCCL weight transfer engine. - - Alternative to update_weights_ipc for the broadcast (non-IPC) sender: - the trainer initiates an NCCL broadcast via the vendored - ``nccl_trainer_send_weights``, and each inference worker calls - weight_transfer_engine.receive_weights here. - - ``update_info`` carries only names/dtype_names/shapes. Since vLLM 0.28.0 - whether the transfer is packed is fixed at init (from the - ``/init_weight_transfer_engine`` payload), not per round. - - Routed through this skyrl wrap (rather than vLLM's native - /update_weights endpoint) so the load is wrapped with - set_current_vllm_config — process_weights_after_loading on MoE - models can otherwise instantiate kernels (e.g. FlashInfer CUTLASS) - whose __init__ reads get_current_vllm_config(). - - TODO: remove once the upstream vLLM patch lands (vllm-project/vllm - weight-sync-fix), then route via the native /update_weights endpoint. - https://github.com/vllm-project/vllm/pull/42577 - """ - if not getattr(self, "_skyrl_weight_update_active", False): - raise RuntimeError("skyrl_start_weight_update must be called before update_weights_nccl.") - - if self.weight_transfer_engine is None: - raise RuntimeError( - "Weight transfer not configured. Please set weight_transfer_config to enable weight transfer." - ) - - from vllm.config import set_current_vllm_config - - from skyrl.backends.skyrl_train.inference_servers.spec_decode_utils import ( - _reload_spec_decode_drafter, - ) - - engine = self.weight_transfer_engine - typed_update_info = engine.parse_update_info(update_info) - model = self.model_runner.model - - def _load_weights(weights): - weights = list(weights) - loaded = model.load_weights(weights=weights) - _reload_spec_decode_drafter(self.model_runner, weights) - return loaded - - # vLLM 0.26 dropped the `load_weights` callback parameter from - # WeightTransferEngine.receive_weights; the engines now call - # `self.model.load_weights` directly. Retarget the engine at a proxy - # whose load_weights is ours (so the spec-decode drafter still gets - # reloaded), using vLLM's own set/reset_weight_update_target hooks. - engine.set_weight_update_target( - _LoadWeightsProxy(model, _load_weights), - self.model_config, - ) - try: - with set_current_vllm_config(self.vllm_config), torch.device(self.device): - engine.receive_weights(typed_update_info) - finally: - engine.reset_weight_update_target() - - torch.accelerator.synchronize() - _empty_cuda_cache_rocm() - # Suspend / resume for non-colocated weight sync. # - # Drive the per-worker CuMemAllocator directly instead of GPUWorker.sleep/ - # wake_up (only reachable via EngineCore.sleep, which force-clears the prefix - # cache and preempts running requests at level >= 1). Touching the allocator - # alone lets the caller hold a KEEP pause across the sync and resume frozen - # requests with their KV restored to the same virtual addresses -- no abort, - # no prefill recompute. Mirrors GPUWorker.sleep/wake_up; re-verify on vLLM bumps. + # Drives the per-worker CuMemAllocator directly instead of GPUWorker.sleep/ + # wake_up, which is only reachable via EngineCore.sleep and force-clears the + # prefix cache and preempts running requests at level >= 1. Touching the + # allocator alone lets the caller hold a KEEP pause across the sync and + # resume frozen requests with their KV at the same virtual addresses -- no + # abort, no prefill recompute. Mirrors GPUWorker.sleep/wake_up; re-verify on + # vLLM bumps. def skyrl_sleep_for_weight_sync(self, offload_kv: bool = True) -> None: """Free GPU memory for weight sync by sleeping the allocator. @@ -262,13 +122,19 @@ def skyrl_sleep_for_weight_sync(self, offload_kv: bool = True) -> None: every parameter on wake. ``offload_kv`` controls whether the KV cache is offloaded to CPU (preserved for frozen in-flight requests) or discarded. Model buffers live in the weights pool but are not sent by the broadcast (e.g. - non-persistent rotary ``inv_freq``), so save them here and restore on wake -- - as GPUWorker.sleep(level=2) does. + non-persistent rotary ``inv_freq``), so save them here and restore on + wake -- as GPUWorker.sleep(level=2) does. + + The drafter's buffers are saved too: the weight sync reloads its + parameters, but nothing else restores its buffers. """ from vllm.device_allocator import get_mem_allocator_instance - model = self.model_runner.model - self._skyrl_saved_buffers = {name: buf.cpu().clone() for name, buf in model.named_buffers()} + self._skyrl_saved_buffers = {name: buf.cpu().clone() for name, buf in self.model_runner.model.named_buffers()} + draft = self._skyrl_draft_model() + self._skyrl_saved_draft_buffers = ( + {name: buf.cpu().clone() for name, buf in draft.named_buffers()} if draft is not None else {} + ) get_mem_allocator_instance().sleep(offload_tags=("kv_cache",) if offload_kv else ()) def skyrl_wake_for_weight_sync(self, tags: list) -> None: @@ -284,71 +150,28 @@ def skyrl_wake_for_weight_sync(self, tags: list) -> None: torch.cuda.empty_cache() get_mem_allocator_instance().wake_up(tags) - # Restore model buffers (not covered by the broadcast) once weights remap. - saved = getattr(self, "_skyrl_saved_buffers", None) - if saved and (tags is None or "weights" in tags): - model = self.model_runner.model - for name, buf in model.named_buffers(): - if name in saved: - buf.data.copy_(saved[name].data) - self._skyrl_saved_buffers = {} + # Restore buffers (not covered by the broadcast) once weights remap. + if tags is None or "weights" in tags: + self._skyrl_restore_buffers(self.model_runner.model, "_skyrl_saved_buffers") + draft = self._skyrl_draft_model() + if draft is not None: + self._skyrl_restore_buffers(draft, "_skyrl_saved_draft_buffers") # Re-init fp8 KV scales after the KV pool remaps (no-op without fp8 KV cache). if tags is None or "kv_cache" in tags: post_wake = getattr(self.model_runner, "post_kv_cache_wake_up", None) if post_wake is not None: post_wake() - def init_weight_transfer_engine_rdt(self, init_info: dict) -> None: - """ - Initialize + bake the sharded_rdt weight-transfer engine. - - GPUWorker.load_model already constructed the engine via the factory (the - sharded_rdt backend is registered in rdt_vllm_register); here we run its - one-time bake. Routed through this skyrl wrap (rather than vLLM's native - /init_weight_transfer_engine endpoint) so the bake runs under - set_current_vllm_config + torch.device(self.device): the bake drives - model.load_weights against meta params, and process_weights_after_loading - on MoE models reads get_current_vllm_config() to build kernels. - - Args: - init_info: asdict(ShardedRDTWeightTransferInitInfo) — trainer actor - name/namespace, produce method name, M:N + ring knobs, and the - group-major names/dtype_names/shapes/group_lens the bake plans over. - """ - if self.weight_transfer_engine is None: - raise RuntimeError( - "Weight transfer not configured. Set weight_transfer_config with " - "backend='sharded_rdt' to enable the RDT weight-transfer engine." - ) - - from vllm.config import set_current_vllm_config - - typed_init_info = self.weight_transfer_engine.parse_init_info(init_info) - with set_current_vllm_config(self.vllm_config), torch.device(self.device): - self.weight_transfer_engine.init_transfer_engine(typed_init_info) - - def update_weights_rdt(self, update_info: dict) -> None: - """ - Pull this worker's consumed slices via the sharded_rdt engine. - - Called once per sync (the engine pre-built its static whole-model plan at - init, so update_info is empty). The engine pulls every slice over NIXL, - pipelined across its receive-buffer ring, and DEFERS the GPU - post-processing (materialize/scatter/quant/kernel-copy) to background - threads — so, unlike the ipc/nccl paths, we do NOT synchronize here. - skyrl_finish_weight_update drains the deferred work before finalize. - """ - if not getattr(self, "_skyrl_weight_update_active", False): - raise RuntimeError("skyrl_start_weight_update must be called before update_weights_rdt.") - - if self.weight_transfer_engine is None: - raise RuntimeError( - "Weight transfer not configured. Set weight_transfer_config with " - "backend='sharded_rdt' to enable the RDT weight-transfer engine." - ) - - from vllm.config import set_current_vllm_config - - typed_update_info = self.weight_transfer_engine.parse_update_info(update_info) - with set_current_vllm_config(self.vllm_config), torch.device(self.device): - self.weight_transfer_engine.receive_weights(typed_update_info) + def _skyrl_draft_model(self): + """The spec-decode drafter module, or None when there is no drafter.""" + get_draft_model = getattr(self.model_runner, "get_draft_model", None) + return get_draft_model() if callable(get_draft_model) else None + + def _skyrl_restore_buffers(self, module, attr: str) -> None: + saved = getattr(self, attr, None) + if not saved: + return + for name, buf in module.named_buffers(): + if name in saved: + buf.data.copy_(saved[name].data) + setattr(self, attr, {}) diff --git a/skyrl/backends/skyrl_train/inference_servers/rdt_control_protocol.py b/skyrl/backends/skyrl_train/inference_servers/rdt_control_protocol.py deleted file mode 100644 index cfeba81438..0000000000 --- a/skyrl/backends/skyrl_train/inference_servers/rdt_control_protocol.py +++ /dev/null @@ -1,63 +0,0 @@ -"""Wire protocol for the sharded-RDT control plane (dependency-light). - -The RDT control plane is a serial ``init -> start -> update -> finish`` handshake -routed through the SkyRL worker extension's ``/collective_rpc`` (so the bake runs -under ``set_current_vllm_config``), not vLLM's native weight-transfer endpoints. - -Two clients speak it: the async ``RemoteInferenceClient`` (over aiohttp) and the -trainer-side ``SyncRdtControlPlaneClient`` (over blocking HTTP; see -``weight_sync/sharded_rdt/rdt_control_plane.py``). This module is the single source of truth -for the method names and the per-server ``replica_rank`` fan-out so the two can -never drift. It is intentionally stdlib-only — no ray / torch / vllm — so either -client can import it without pulling backend deps. -""" - -from typing import Any, Dict, List, Sequence, Tuple - -COLLECTIVE_RPC_ENDPOINT = "/collective_rpc" - -# Worker-extension methods each /collective_rpc call dispatches to. -RDT_INIT_METHOD = "init_weight_transfer_engine_rdt" -RDT_START_METHOD = "skyrl_start_weight_update" -RDT_UPDATE_METHOD = "update_weights_rdt" -RDT_FINISH_METHOD = "skyrl_finish_weight_update" - - -def build_rdt_init_payloads( - init_info: Dict[str, Any], - server_urls: Sequence[str], - data_parallel_size: int, -) -> List[Tuple[str, Dict[str, Any]]]: - """Per-server ``/collective_rpc`` payloads for the RDT engine bake + init. - - Each independent inference *deployment* has its own self-contained parallel - config, so the vLLM engine can't tell deployments apart — every deployment's - internal worker index restarts at 0 and would collide under the M:N block - assignment. So we stamp each server with its deployment ordinal - (``server_index // data_parallel_size``) as ``replica_rank`` and the - deployment count as ``num_replicas``; the engine offsets its consumers into a - globally distinct range from those two fields. Every other field is shared. - - The replica ordinal divides by ``data_parallel_size`` because the DP servers - of one deployment share a parallel config (vLLM's ``data_parallel_index`` - already separates them), so they must share ONE ``replica_rank`` or a DP - deployment would double-count. - """ - dp = max(1, data_parallel_size) - num_replicas = max(1, len(server_urls) // dp) - return [ - ( - url, - { - "method": RDT_INIT_METHOD, - "kwargs": { - "init_info": { - **init_info, - "replica_rank": i // dp, - "num_replicas": num_replicas, - }, - }, - }, - ) - for i, url in enumerate(server_urls) - ] diff --git a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py index 7bfbd42a64..4923d1f8bd 100644 --- a/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py +++ b/skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py @@ -51,7 +51,6 @@ from dataclasses import dataclass, field from enum import Enum from typing import ( - TYPE_CHECKING, Any, Dict, List, @@ -97,12 +96,6 @@ "stop_tokens": "stop_token_ids", } -if TYPE_CHECKING: - from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( - WeightSyncInitInfo, - ) - - logger = logging.getLogger(__name__) @@ -1102,35 +1095,6 @@ async def reset_prefix_cache( # Weight Sync (control plane - fan-out) # --------------------------- - async def init_weight_update_communicator( - self, - init_info: "WeightSyncInitInfo", - ) -> Dict[str, Any]: - """ - Initialize weight sync via vLLM native /init_weight_transfer_engine. - - Fetches per-server world sizes, expands init_info into per-server - payloads (with correct NCCL rank offsets), and fans out to all servers. - - Args: - init_info: A WeightSyncInitInfo (e.g. BroadcastInitInfo) that supports - for_servers() and to_api_payload(). - - Returns: - Dict mapping server_url to response. - """ - _, world_size_per_server = await self.get_world_size() - num_servers = len(self.server_urls) - server_infos = init_info.for_servers(world_size_per_server, num_servers, dp_size=self.data_parallel_size) - payloads = [{"init_info": x.to_api_payload()} for x in server_infos] - results = await asyncio.gather( - *[ - self._call_server(url, "/init_weight_transfer_engine", payload) - for url, payload in zip(self.server_urls, payloads) - ] - ) - return {url: resp for url, resp in results} - async def update_named_weights( self, update_info: Dict[str, Any], @@ -1168,106 +1132,12 @@ async def fetch_weights( kwargs["uri"] = uri return await self._call_all_servers("/fetch_weights", kwargs) - # TODO: Once https://github.com/vllm-project/vllm/pull/39212 lands, switch - # these three methods from /collective_rpc to the native vLLM endpoints - # (/start_weight_update, /update_weights, /finish_weight_update) and remove - # the NewInferenceWorkerWrap worker extension. - - async def start_weight_update( - self, - is_checkpoint_format: bool = True, - ) -> Dict[str, Any]: - """ - Start a new chunked weight update via /collective_rpc. - - Calls the NewInferenceWorkerWrap.skyrl_start_weight_update method on all - workers. For checkpoint-format weights this initializes layerwise - reload. Must be called before any update_weights_ipc calls. - - Args: - is_checkpoint_format: True if weights are in checkpoint format - (need layerwise processing), False for kernel format. - - Returns: - Dict mapping server_url to response. - """ - return await self._call_all_servers( - "/collective_rpc", - { - "method": "skyrl_start_weight_update", - "kwargs": {"is_checkpoint_format": is_checkpoint_format}, - }, - ) - - async def update_weights_ipc( - self, - update_info: Dict[str, Any], - ) -> Dict[str, Any]: - """ - Send a single weight chunk via /collective_rpc. - - Calls NewInferenceWorkerWrap.update_weights_ipc on all workers. - Can be called multiple times between skyrl_start_weight_update and - skyrl_finish_weight_update. - - Args: - update_info: Dict with backend-specific update info (names, - dtype_names, shapes, ipc_handles_pickled or packed flag). - - Returns: - Dict mapping server_url to response. - """ - return await self._call_all_servers( - "/collective_rpc", - { - "method": "update_weights_ipc", - "kwargs": {"update_info": update_info}, - }, - ) - - async def update_weights_nccl( - self, - update_info: Dict[str, Any], - ) -> Dict[str, Any]: - """ - Send batched weight update via /collective_rpc to the broadcast receiver. - - Calls NewInferenceWorkerWrap.update_weights_nccl on all workers, - which routes weight_transfer_engine.receive_weights through the - set_current_vllm_config wrap. Used by the broadcast (NCCL) sender as - a temporary substitute for vLLM's native /update_weights endpoint - until the upstream patch (vllm-project/vllm weight-sync-fix) lands. - - Args: - update_info: Dict with backend-specific update info (names, - dtype_names, shapes, packed flag, etc.) — same shape vLLM's - native /update_weights expects. - - Returns: - Dict mapping server_url to response. - """ - return await self._call_all_servers( - "/collective_rpc", - { - "method": "update_weights_nccl", - "kwargs": {"update_info": update_info}, - }, - ) - - async def finish_weight_update(self) -> Dict[str, Any]: - """ - Finish the current chunked weight update via /collective_rpc. - - Calls NewInferenceWorkerWrap.skyrl_finish_weight_update on all workers. - For checkpoint-format weights, runs layerwise postprocessing. - - Returns: - Dict mapping server_url to response. - """ - return await self._call_all_servers( - "/collective_rpc", - {"method": "skyrl_finish_weight_update"}, - ) + # The weight-sync lifecycle (/start_weight_update, /update_weights, + # /finish_weight_update) is driven by the trainer-side engines through the + # blocking SkyrlWeightSyncClient (weight_sync/control_plane.py), which they + # need because the protocol is synchronous and they run off the event loop. + # What is left here is what the driver drives: pause/resume, prefix-cache + # reset, /fetch_weights, LoRA, and /get_world_size at init. async def load_lora_adapter( self, diff --git a/skyrl/backends/skyrl_train/inference_servers/utils.py b/skyrl/backends/skyrl_train/inference_servers/utils.py index 8abc5e0578..360083b15a 100644 --- a/skyrl/backends/skyrl_train/inference_servers/utils.py +++ b/skyrl/backends/skyrl_train/inference_servers/utils.py @@ -11,14 +11,18 @@ SKYRL_LORA_ADAPTER_NAME, ) -# Importing the weight_sync package registers the sharded_rdt engine into vLLM's -# factory and relaxes WeightTransferConfig.backend so backend="sharded_rdt" -# validates below. This must run on the driver before the WeightTransferConfig -# is constructed; the import above already triggers it, this is just explicit. -from skyrl.backends.skyrl_train.weight_sync import get_transfer_strategy +# The receive-side engines must be registered in vLLM's factory before the +# WeightTransferConfig below is built: `backend` is validated against the +# registry. Needed on the driver (here) and in every worker process (the +# worker-extension module imported above). +from skyrl.backends.skyrl_train.weight_sync import ( + get_transfer_strategy, + get_vllm_receive_backend, +) from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import ( rdt_vllm_register, # noqa: F401,E402 ) +from skyrl.backends.skyrl_train.weight_sync.skyrl_engines import register_skyrl_engines from skyrl.train.config import ( InferenceEngineConfig, SkyRLTrainConfig, @@ -72,6 +76,8 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace: from vllm.platforms import current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser + register_skyrl_engines() + # This function may run a GPU-less Ray head # node, where ``current_platform`` resolves to ``UnspecifiedPlatform`` with # ``device_type == ""``. vLLM's ``add_cli_args`` walks ``VllmConfig`` defaults @@ -110,7 +116,7 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace: enable_sleep_mode=cfg.trainer.placement.colocate_all or ie_cfg.offload_kv_for_weight_sync, enable_return_routed_experts=ie_cfg.enable_return_routed_experts, weight_transfer_config=WeightTransferConfig( - backend=get_transfer_strategy(ie_cfg.weight_sync_backend, cfg.trainer.placement.colocate_all), + backend=get_vllm_receive_backend(ie_cfg.weight_sync_backend, cfg.trainer.placement.colocate_all), ), worker_extension_cls=VLLM_NEW_INFERENCE_WORKER_EXTENSION_CLS, # NOTE (sumanthrh): We set generation config to be vLLM so that the generation behaviour of the server is same as using the vLLM Engine APIs directly diff --git a/skyrl/backends/skyrl_train/patches/vllm/patch_model_runner_registry.py b/skyrl/backends/skyrl_train/patches/vllm/patch_model_runner_registry.py new file mode 100644 index 0000000000..804e16892a --- /dev/null +++ b/skyrl/backends/skyrl_train/patches/vllm/patch_model_runner_registry.py @@ -0,0 +1,59 @@ +"""Record this process's ``GPUModelRunner`` so a weight-transfer engine can reach it. + +A ``WeightTransferEngine`` is constructed with ``(config, vllm_config, device, +model)`` -- the main model only. SkyRL's receive path also needs the +speculative-decoding drafter (``model_runner.drafter.model``), a separate module +that ``load_weights`` on the main model never touches, so an MTP model would keep +drafting with pre-sync weights (see ``inference_servers/spec_decode_utils``). + +There is no route from an engine to its worker: the engine is created inside +``Worker.load_model``, and the worker-extension class is appended to +``Worker.__bases__`` *after* ``Worker`` (``v1/worker/worker_base.py``), so it can +neither override ``load_model`` nor hand the engine anything. So wrap +``GPUModelRunner.load_model`` and record the runner in a process-global weakref: +one runner per worker process, recorded before any weight sync can run, read +lazily at sync time. + +REMOVAL: delete this once vLLM gives ``WeightTransferEngine`` a post-load hook +(or reaches the drafter itself). Nothing else depends on it -- the engines treat a +missing runner as "no drafter". +""" + +import weakref +from typing import Any, Optional + +_PATCHED = False +_CURRENT_MODEL_RUNNER: Optional["weakref.ReferenceType[Any]"] = None + + +def apply_model_runner_registry_patch() -> None: + """Install the recorder on ``GPUModelRunner.load_model`` (idempotent).""" + global _PATCHED + if _PATCHED: + return + + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + + original = GPUModelRunner.load_model + + def load_model(self, *args, **kwargs): + result = original(self, *args, **kwargs) + global _CURRENT_MODEL_RUNNER + _CURRENT_MODEL_RUNNER = weakref.ref(self) + return result + + load_model.__wrapped__ = original + GPUModelRunner.load_model = load_model + _PATCHED = True + + +def current_model_runner() -> Optional[Any]: + """This process's ``GPUModelRunner``, or None if it was never recorded. + + None means the patch did not run -- not a vLLM worker process, or a vLLM + version whose runner no longer goes through ``GPUModelRunner.load_model``. + Callers must treat it as "no drafter available", not as an error. + """ + if _CURRENT_MODEL_RUNNER is None: + return None + return _CURRENT_MODEL_RUNNER() diff --git a/skyrl/backends/skyrl_train/weight_sync/__init__.py b/skyrl/backends/skyrl_train/weight_sync/__init__.py index d28c29f65e..c76eace4f0 100644 --- a/skyrl/backends/skyrl_train/weight_sync/__init__.py +++ b/skyrl/backends/skyrl_train/weight_sync/__init__.py @@ -1,66 +1,45 @@ -"""Weight synchronization abstractions for distributed RL training.""" +"""Weight synchronization for distributed RL training. + +SkyRL drives weight sync through vLLM's trainer-send abstraction: each training +worker builds a ``WeightSource`` over its live model (``sources.py``) and hands +it to a ``TrainerWeightTransferEngine`` (``trainer_engines.py``) whose +``send_weights()`` owns the whole round trip. Four backends: + +=============== ========================================================== +``nccl`` broadcast from trainer rank 0; non-colocated +``ipc`` CUDA IPC handles; colocated (``placement.colocate_all``) +``delta`` compressed checkpoint deltas via disk / ``gs://`` / ``s3://`` +``sharded_rdt`` the inference workers pull slices over NIXL +=============== ========================================================== + +This package imports no vLLM at module scope: half the CPU suite runs without the +(Linux-only, optional) wheel. The vLLM-facing modules -- ``sources``, +``trainer_engines``, ``skyrl_engines``, ``delta_trainer`` -- are imported at their +call sites instead. +""" + +from .base import LoraLoadRequest + +#: Logical backend name -> the name its *receive*-side engine is registered +#: under in vLLM's ``WeightTransferEngineFactory``. +#: +#: NCCL and IPC take new names because SkyRL subclasses vLLM's engines and +#: ``register_engine`` refuses an already-registered name. The trainer-side +#: factory has a separate registry, so the send side keeps vLLM's own names. +_VLLM_RECEIVE_BACKENDS = { + "nccl": "skyrl_nccl", + "ipc": "skyrl_ipc", + "delta": "delta", + "sharded_rdt": "sharded_rdt", +} -from typing import Type -from .base import LoraLoadRequest, WeightChunk, WeightUpdateRequest -from .broadcast_strategy import ( - BroadcastInitInfo, - BroadcastTransferStrategy, - BroadcastWeightTransferSender, - BroadcastWeightUpdateRequest, -) -from .cuda_ipc_strategy import ( - CudaIpcInitInfo, - CudaIpcTransferStrategy, - CudaIpcWeightTransferSender, - CudaIpcWeightUpdateRequest, -) -from .delta_strategy import ( - DeltaInitInfo, - DeltaTransferStrategy, - DeltaWeightTransferSender, -) -from .sharded_rdt.sharded_rdt_strategy import ( - ShardedRdtInitInfo, - ShardedRdtTransferStrategy, - ShardedRdtWeightTransferSender, -) -from .transfer_strategy import ( - WeightSyncInitInfo, - WeightTransferSender, - WeightTransferStrategy, -) -from .weight_extractor import WeightExtractor - - -def get_transfer_strategy_cls(weight_sync_backend: str, colocate_all: bool) -> Type[WeightTransferStrategy]: - """Get the appropriate transfer strategy class based on config. - - Uses CUDA IPC when: - - weight_sync_backend is "nccl" - - colocate_all is True (training and inference on same nodes) - - Otherwise uses broadcast. - - Args: - weight_sync_backend: The weight sync backend ("nccl" or other). - colocate_all: Whether training and inference are colocated on same nodes. +def get_transfer_strategy(weight_sync_backend: str, colocate_all: bool) -> str: + """Resolve the logical weight-sync backend from config + placement. - Returns: - The strategy class (CudaIpcTransferStrategy or BroadcastTransferStrategy). + A configured ``nccl`` means CUDA IPC when training and inference share GPUs, + and broadcast otherwise. """ - strategy = get_transfer_strategy(weight_sync_backend, colocate_all) - if strategy == "sharded_rdt": - return ShardedRdtTransferStrategy - if strategy == "delta": - return DeltaTransferStrategy - if strategy == "ipc": - return CudaIpcTransferStrategy - return BroadcastTransferStrategy - - -def get_transfer_strategy(weight_sync_backend: str, colocate_all: bool) -> str: - """Get the appropriate transfer strategy string based on config.""" if weight_sync_backend in ("sharded_rdt", "rdt"): return "sharded_rdt" if weight_sync_backend == "delta": @@ -70,27 +49,13 @@ def get_transfer_strategy(weight_sync_backend: str, colocate_all: bool) -> str: return "nccl" +def get_vllm_receive_backend(weight_sync_backend: str, colocate_all: bool) -> str: + """The ``WeightTransferConfig.backend`` the inference servers must be built with.""" + return _VLLM_RECEIVE_BACKENDS[get_transfer_strategy(weight_sync_backend, colocate_all)] + + __all__ = [ - "WeightChunk", - "WeightExtractor", - "WeightUpdateRequest", "LoraLoadRequest", - "BroadcastWeightUpdateRequest", - "CudaIpcWeightUpdateRequest", - "WeightTransferStrategy", - "WeightTransferSender", - "WeightSyncInitInfo", - "BroadcastInitInfo", - "CudaIpcInitInfo", - "BroadcastTransferStrategy", - "BroadcastWeightTransferSender", - "CudaIpcTransferStrategy", - "CudaIpcWeightTransferSender", - "DeltaInitInfo", - "DeltaTransferStrategy", - "DeltaWeightTransferSender", - "ShardedRdtInitInfo", - "ShardedRdtTransferStrategy", - "ShardedRdtWeightTransferSender", - "get_transfer_strategy_cls", + "get_transfer_strategy", + "get_vllm_receive_backend", ] diff --git a/skyrl/backends/skyrl_train/weight_sync/base.py b/skyrl/backends/skyrl_train/weight_sync/base.py index fd8753c24b..d8bd993bbc 100644 --- a/skyrl/backends/skyrl_train/weight_sync/base.py +++ b/skyrl/backends/skyrl_train/weight_sync/base.py @@ -1,106 +1,29 @@ """Base data structures for weight synchronization.""" -from dataclasses import asdict, dataclass, field -from functools import cached_property -from typing import Any, Dict, List - -import torch +from dataclasses import asdict, dataclass +from typing import Any, Dict @dataclass -class WeightUpdateRequest: - """Base class for weight update requests. - - Each transfer strategy has its own request type with strategy-specific fields. - """ +class LoraLoadRequest: + """Request to load LoRA weights from disk. - names: List[str] - dtypes: List[str] - shapes: List[List[int]] + Not a weight *transfer*: it tells the inference engine to load an adapter + from a path rather than moving any tensor. - def __post_init__(self): - lengths = [len(self.names), len(self.dtypes), len(self.shapes)] - if len(set(lengths)) != 1: - raise ValueError( - f"names, dtypes, shapes must have the same length. " - f"Got names={len(self.names)}, dtypes={len(self.dtypes)}, shapes={len(self.shapes)}" - ) + ``lora_name`` is the name vLLM registers the adapter under, and what callers + later pass as ``model=`` when sampling. Empty string preserves the + legacy single-tenant behavior where the engine generates a numeric name. + """ - def __len__(self) -> int: - return len(self.names) + lora_path: str = "" + lora_name: str = "" def to_json_dict(self) -> Dict[str, Any]: """Serialize the request to JSON.""" return asdict(self) @classmethod - def from_json_dict(cls, data: Dict[str, Any]) -> "WeightUpdateRequest": + def from_json_dict(cls, data: Dict[str, Any]) -> "LoraLoadRequest": """Deserialize the request from JSON.""" return cls(**data) - - -@dataclass -class LoraLoadRequest(WeightUpdateRequest): - """Request to load LoRA weights from disk. - - This is a special request type used for loading LoRA adapters - from disk rather than transferring weights over network in training. Unlike other - WeightUpdateRequest subclasses, this doesn't transfer weights - it tells - the inference engine to load LoRA from a path. - - ``lora_name`` is the name vLLM should register the adapter under and is - what callers later pass as ``model=`` when sampling. Empty - string preserves the legacy single-tenant behavior where the engine - generates a numeric name itself. - """ - - names: List[str] = field(default_factory=list) - dtypes: List[str] = field(default_factory=list) - shapes: List[List[int]] = field(default_factory=list) - lora_path: str = "" - lora_name: str = "" - - -@dataclass -class WeightChunk: - """Represents one or more model parameters to be transferred. - - A WeightChunk can contain multiple parameters grouped together for efficient - transfer (e.g., Q/K/V projections for fused-weight loaders). - - Attributes: - names: List of parameter names (e.g., ["model.layer.0.weight"]) - dtypes: List of dtype strings (e.g., ["torch.bfloat16"]) - shapes: List of tensor shapes (e.g., [[4096, 4096]]) - tensors: List of actual tensor data (populated during extraction) - total_numel: Total number of elements (cached property, auto-calculated) - total_size_bytes: Total memory footprint (cached property, auto-calculated) - """ - - names: List[str] - dtypes: List[str] - shapes: List[List[int]] - tensors: List[torch.Tensor] - - def __post_init__(self): - """Validate that all input lists have the same length.""" - lengths = [len(self.names), len(self.dtypes), len(self.shapes), len(self.tensors)] - if len(set(lengths)) != 1: - raise ValueError( - f"All lists must have the same length. Got names={len(self.names)}, " - f"dtypes={len(self.dtypes)}, shapes={len(self.shapes)}, tensors={len(self.tensors)}" - ) - - def __len__(self) -> int: - """Return the number of parameters in this chunk.""" - return len(self.names) - - @cached_property - def total_numel(self) -> int: - """Calculate total number of elements across all tensors.""" - return sum(t.numel() for t in self.tensors) - - @cached_property - def total_size_bytes(self) -> int: - """Calculate total memory footprint in bytes.""" - return sum(t.numel() * t.element_size() for t in self.tensors) diff --git a/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py b/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py deleted file mode 100644 index 7aa9c455c1..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/broadcast_strategy.py +++ /dev/null @@ -1,292 +0,0 @@ -"""Broadcast-based weight transfer strategy using torch.distributed. - -This module implements the broadcast transfer strategy for synchronizing model weights -from training workers to inference engines using NCCL/Gloo broadcast operations. -""" - -import asyncio -import socket -from dataclasses import dataclass, replace -from typing import TYPE_CHECKING, Any, Dict, Iterable, Iterator, List, Optional, Tuple - -if TYPE_CHECKING: - from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( - RemoteInferenceClient, - ) - from skyrl.train.config.config import InferenceEngineConfig - -import ray -import torch - -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk, WeightUpdateRequest -from skyrl.backends.skyrl_train.weight_sync.nccl_trainer_send import ( - nccl_trainer_init, - nccl_trainer_send_weights, -) -from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( - WeightSyncInitInfo, - WeightTransferSender, - WeightTransferStrategy, -) - - -@dataclass -class BroadcastInitInfo(WeightSyncInitInfo): - """Initialization info for broadcast-based weight transfer.""" - - master_addr: str - master_port: int - rank_offset: int - world_size: int - packed: bool = True - """Whether the transfer is packed. As of vLLM 0.28.0 this is an init-time - wire param: the worker records it from ``NCCLWeightTransferInitInfo`` during - ``/init_weight_transfer_engine`` and ``receive_weights`` reads it from there, - so it can no longer ride the per-round update info. It must match the - ``packed`` the sender passes to ``nccl_trainer_send_weights``, or the two - sides split the stream differently and the broadcast hangs in NCCL.""" - - def for_servers(self, world_size_per_server: int, num_servers: int, dp_size: int = 1) -> List["BroadcastInitInfo"]: - """Return one BroadcastInitInfo per server with rank_offset for each. - - Used when calling init_weight_update_communicator on the new inference path: - expand the single init_info into a list (one per server), then pass - [x.to_api_payload() for x in server_infos] to the client. - - server_urls are ordered as [engine0_dp0, engine0_dp1, ..., engine1_dp0, ...]. - All DP servers within one deployment share the same rank_offset because - vLLM's init_transfer_engine already accounts for dp_rank internally. - The offset only advances at deployment (num_engines) boundaries. - - Args: - world_size_per_server: Number of workers per server (same for all servers). - num_servers: Total number of servers (num_engines * dp_size). - dp_size: Data parallel size. Servers are grouped into deployments - of dp_size servers each. - - Returns: - List of BroadcastInitInfo, one per server, with cumulative rank_offset. - """ - result: List[BroadcastInitInfo] = [] - rank_offset = self.rank_offset - for i in range(num_servers): - result.append(replace(self, rank_offset=rank_offset)) - # Advance rank_offset only at deployment boundaries (every dp_size servers) - if (i + 1) % dp_size == 0: - rank_offset += world_size_per_server - return result - - def to_api_payload(self) -> Dict[str, Any]: - """Return JSON-serializable payload for the /init_weight_transfer_engine endpoint.""" - return { - "master_address": self.master_addr, - "master_port": self.master_port, - "rank_offset": self.rank_offset, - "world_size": self.world_size, - "packed": self.packed, - } - - -@dataclass -class BroadcastWeightUpdateRequest(WeightUpdateRequest): - """Request for broadcast-based weight transfer. - - When sizes is provided, tensors are packed into a single contiguous buffer - and broadcast as one NCCL operation per chunk. The receiver uses sizes to unpack. - When sizes is None, falls back to per-tensor broadcast (backward compatible). - """ - - sizes: Optional[List[int]] = None - - -class BroadcastWeightTransferSender(WeightTransferSender): - """Sends weights via torch.distributed.broadcast or vLLM NCCL (new inference path). - - When using new inference, uses the vendored ``nccl_trainer_send_weights`` - (see ``nccl_trainer_send.py``) with batched update_weights. Otherwise uses - per-chunk HTTP + torch.distributed.broadcast. - """ - - def __init__( - self, - init_info: BroadcastInitInfo, - model_update_group: Optional[Any], - inference_client: "RemoteInferenceClient", - ) -> None: - """Initialize the broadcast sender. - - Args: - init_info: BroadcastInitInfo containing all config-derived args. - model_update_group: Communication group for weight transfer. Either a - torch.distributed.ProcessGroup (legacy) or a vLLM NCCL - communicator (new path). None on non-rank-0 workers. - inference_client: Client for coordinating with inference engines. - """ - self._init_info = init_info - self._model_update_group = model_update_group - self._inference_client = inference_client - - async def send_chunks( - self, - chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, - **kwargs, - ) -> None: - """Send chunks via broadcast or vLLM native NCCL. - - Args: - chunks: Iterable of WeightChunk objects to send. - weight_metadata: Pre-computed metadata dict with "names", "dtype_names", - "shapes". Avoids materializing all chunks to collect metadata. - """ - await self._send_chunks_vllm_native(chunks, weight_metadata) - - async def _send_chunks_vllm_native( - self, - chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, - ) -> None: - """Batched path: one update_weights call + nccl_trainer_send_weights. - - All ranks must evaluate the chunks iterator (extract_weights uses - collective all-gather internally). Only rank 0 sends the gathered - tensors to vLLM via the NCCL weight transfer engine. - """ - if weight_metadata is None: - raise ValueError( - "weight_metadata is required for vLLM native path. " - "Call weight_extractor.get_weight_metadata() and pass it to send_chunks." - ) - - def weight_iterator() -> Iterator[Tuple[str, torch.Tensor]]: - for chunk in chunks: - yield from zip(chunk.names, chunk.tensors) - - # Route via the skyrl wrap (start_weight_update + update_weights_nccl - # + finish_weight_update) rather than vLLM's native /update_weights so - # the receive is wrapped with set_current_vllm_config. Matches how - # CUDA IPC already routes through skyrl's wrap. - # TODO: switch back to update_named_weights once the upstream vLLM - # patch lands (vllm-project/vllm weight-sync-fix). - # https://github.com/vllm-project/vllm/pull/42577 - if torch.distributed.get_rank() == 0: - await self._inference_client.start_weight_update(is_checkpoint_format=True) - - # vLLM 0.28.0 dropped `packed` (and the buffer geometry) from - # NCCLWeightTransferUpdateInfo -- it is agreed once at init instead, - # via BroadcastInitInfo.packed. Sending it here is now a TypeError. - update_info = dict(weight_metadata) - update_task = asyncio.create_task(self._inference_client.update_weights_nccl(update_info)) - - # Run in thread so the HTTP update_task can progress concurrently - await asyncio.to_thread( - nccl_trainer_send_weights, - weight_iterator(), - self._model_update_group, - packed=self._init_info.packed, - ) - await update_task - - await self._inference_client.finish_weight_update() - else: - # Non-rank-0 still needs to participate in the all-gather - for _ in weight_iterator(): - pass - - torch.distributed.barrier() - - def teardown(self) -> None: - """Destroy the process group used for weight transfer.""" - if self._model_update_group is not None and isinstance( - self._model_update_group, torch.distributed.ProcessGroup - ): - torch.distributed.destroy_process_group(self._model_update_group) - self._model_update_group = None - - -class BroadcastTransferStrategy(WeightTransferStrategy): - """Factory for broadcast-based weight transfer. - - This strategy uses NCCL/Gloo broadcast operations to transfer weights from - training workers to inference engines. - - All methods are static - no instance state needed. - """ - - @staticmethod - def create_init_info( - ie_cfg: "InferenceEngineConfig", - inference_world_size: int, - base_model_path: Optional[str] = None, - ) -> BroadcastInitInfo: - """Create init info with all config-derived args. - - Args: - ie_cfg: InferenceEngineConfig containing inference engine settings. - inference_world_size: Total number of inference workers (from client.get_world_size()). - - Returns: - BroadcastInitInfo containing all args needed for sender/receiver creation. - """ - # Use world_size reported by the inference servers (+1 for trainer rank 0). - world_size = inference_world_size + 1 - - master_addr = ray._private.services.get_node_ip_address() - with socket.socket() as sock: - sock.bind(("", 0)) - master_port = sock.getsockname()[1] - - return BroadcastInitInfo( - master_addr=master_addr, - master_port=master_port, - rank_offset=1, - world_size=world_size, - override_existing_receiver=not ie_cfg.run_engines_locally, - ) - - @staticmethod - def create_sender( - init_info: BroadcastInitInfo, - inference_client: "RemoteInferenceClient", - weight_extractor: Optional[Any] = None, - ) -> BroadcastWeightTransferSender: - """Create a broadcast sender. - - On rank 0, joins the weight-transfer group via ``nccl_trainer_init`` - (vLLM's ``nccl_common.trainer_init``). Other ranks hold no communicator. - - Args: - init_info: BroadcastInitInfo from create_init_info. - inference_client: Client for coordinating with inference engines. - """ - rank = torch.distributed.get_rank() - model_update_group = None - - if rank == 0: - model_update_group = nccl_trainer_init( - dict( - master_address=init_info.master_addr, - master_port=init_info.master_port, - world_size=init_info.world_size, - ) - ) - - return BroadcastWeightTransferSender( - init_info=init_info, - model_update_group=model_update_group, - inference_client=inference_client, - ) - - @staticmethod - def get_vllm_transfer_engine() -> type: - """Return the vLLM weight-transfer engine class for this strategy (NCCL). - - Reference for the receive side: the inference servers drive this engine - natively. Currently unused on the sender side (we route through the - SkyRL ``/collective_rpc`` wrap), kept as the canonical mapping. - """ - from vllm.distributed.weight_transfer.nccl_engine import ( - NCCLWeightTransferEngine, - ) - - return NCCLWeightTransferEngine diff --git a/skyrl/backends/skyrl_train/weight_sync/control_plane.py b/skyrl/backends/skyrl_train/weight_sync/control_plane.py new file mode 100644 index 0000000000..5668c01e81 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/control_plane.py @@ -0,0 +1,287 @@ +"""Blocking control-plane client for the trainer-side weight transfer engines. + +vLLM's ``VLLMWeightSyncClient`` is a synchronous four-method protocol +(``init_weight_transfer_engine`` / ``start_weight_update`` / ``update_weights`` / +``finish_weight_update``), and the engine that drives it runs off the worker's +event loop. So this client uses blocking HTTP against the same routes rather +than SkyRL's async ``RemoteInferenceClient``, keeping the whole engine +sync-to-sync with no event-loop involvement. + +Four properties are load-bearing: + +* **Per-call fresh connections** (``Connection: close``). A full training step + elapses between syncs, so a pooled keep-alive connection is stale by the next + call and races the server's ``timeout_keep_alive`` (uvicorn 5s) into + ECONNRESET. +* **Concurrent fan-out** across servers. Required for correctness on the RDT + path, not just speed: consumers pull in lockstep and a producer frees a served + group only once every bound consumer has pulled, so a serial + ``update_weights`` stalls the producer's gather loop and deadlocks. +* **Body-aware error messages** (:func:`_error_message`) — surface the response + body's error detail, not the bare HTTP reason phrase. +* **No timeout** (an RDT bake + NIXL pull are long) and **no retry** (retrying a + half-done stateful call would be wrong). + +``requests`` rather than a new dependency: Ray provides it, and Ray is a hard +requirement of every path that reaches here. The import is local so nothing else +pays for it. +""" + +import logging +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple + +logger = logging.getLogger(__name__) + +# vLLM dev-mode RLHF routes (entrypoints/serve/dev/rlhf/api_router.py), plus +# /fetch_weights and /reset_prefix_cache, which SkyRL adds in vllm_server_actor. +INIT_ENGINE_ENDPOINT = "/init_weight_transfer_engine" +START_UPDATE_ENDPOINT = "/start_weight_update" +UPDATE_WEIGHTS_ENDPOINT = "/update_weights" +FINISH_UPDATE_ENDPOINT = "/finish_weight_update" +FETCH_WEIGHTS_ENDPOINT = "/fetch_weights" +RESET_PREFIX_CACHE_ENDPOINT = "/reset_prefix_cache" +PAUSE_ENDPOINT = "/pause" +RESUME_ENDPOINT = "/resume" + + +class SkyrlWeightSyncClient: + """``VLLMWeightSyncClient`` over the inference servers' native HTTP routes. + + Args: + server_urls: every inference server, ordered + ``[engine0_dp0, engine0_dp1, ..., engine1_dp0, ...]`` — the order the + per-server init rewrite (``init_payload_fn``) depends on. + data_parallel_size: DP replicas per deployment. Servers are grouped into + deployments of this many. + init_payload_fn: rewrites the engine's single worker-side init dict into + one payload per server, for the backends that need per-server fields: + NCCL a cumulative ``rank_offset``, sharded RDT a deployment ordinal. + See :func:`nccl_init_payloads` / :func:`rdt_init_payloads`. ``None`` + sends the same dict to every server. + """ + + def __init__( + self, + server_urls: Sequence[str], + data_parallel_size: int = 1, + *, + init_payload_fn: Optional[Callable[[Dict[str, Any], Sequence[str], int], List[Dict[str, Any]]]] = None, + ) -> None: + import requests # local: Ray (a hard dep of every path reaching here) provides it. + + self._urls = list(server_urls) + self._dp = max(1, int(data_parallel_size)) + self._init_payload_fn = init_payload_fn + if not self._urls: + raise ValueError("SkyrlWeightSyncClient requires at least one server_url.") + + self._session = requests.Session() + # See module docstring: fresh connection per call. + self._session.headers["Connection"] = "close" + # One worker per server so a fan-out call issues every POST concurrently. + self._pool = ThreadPoolExecutor(max_workers=len(self._urls), thread_name_prefix="weight-sync-ctrl") + + # ---- VLLMWeightSyncClient protocol (the four methods vLLM's engines call) ---- + + def init_weight_transfer_engine(self, init_info: Dict[str, Any]) -> None: + if self._init_payload_fn is None: + per_server = [dict(init_info) for _ in self._urls] + else: + per_server = self._init_payload_fn(init_info, self._urls, self._dp) + if len(per_server) != len(self._urls): + raise ValueError( + f"init_payload_fn returned {len(per_server)} payloads for {len(self._urls)} servers; " + "it must return exactly one per server, in server order." + ) + self._fanout([(url, INIT_ENGINE_ENDPOINT, {"init_info": info}) for url, info in zip(self._urls, per_server)]) + + def start_weight_update(self) -> None: + self._fanout_uniform(START_UPDATE_ENDPOINT, None) + + def update_weights(self, update_info: Dict[str, Any]) -> None: + self._fanout_uniform(UPDATE_WEIGHTS_ENDPOINT, {"update_info": _json_safe(update_info)}) + + def finish_weight_update(self, weight_version: Optional[str] = None) -> None: + # Omit the key entirely when unset so the route's default applies. + body = {"weight_version": weight_version} if weight_version is not None else None + self._fanout_uniform(FINISH_UPDATE_ENDPOINT, body) + + # ---- extras: the checkpoint-delta lifecycle ---- + # + # Not the transport contract, and here anyway, because they have to be driven + # from inside `send_weights`: `fetch_weights` needs the `target_version` the + # publish produces mid-send, and must run before the pause. `send_weights` is + # sync and runs off the event loop, so `RemoteInferenceClient` -- where these + # routes otherwise live -- is reachable only via `run_coroutine_threadsafe`, + # the loop coupling this client exists to avoid (see module docstring). + # Splitting them into a second sync object would duplicate the session and + # pool to narrow one type. + # + # The cost: an engine calling these needs a SkyRL client, not any object + # satisfying VLLMWeightSyncClient. Only the delta engine does. + + def fetch_weights(self, target_version: int, sync_dir: Optional[str] = None, uri: Optional[str] = None) -> None: + body: Dict[str, Any] = {"target_version": int(target_version)} + if sync_dir is not None: + body["sync_dir"] = sync_dir + if uri is not None: + body["uri"] = uri + self._fanout_uniform(FETCH_WEIGHTS_ENDPOINT, body) + + def reset_prefix_cache(self, reset_running_requests: bool = True) -> None: + self._fanout_uniform(RESET_PREFIX_CACHE_ENDPOINT, {"reset_running_requests": reset_running_requests}) + + def pause_generation(self, clear_cache: bool = False) -> None: + # /pause takes query params, not a body (mirrors RemoteInferenceClient.pause). + self._fanout( + [(url, f"{PAUSE_ENDPOINT}?mode=keep&clear_cache={str(clear_cache).lower()}", None) for url in self._urls] + ) + + def resume_generation(self) -> None: + self._fanout_uniform(RESUME_ENDPOINT, None) + + def close(self) -> None: + """Release the HTTP session + fan-out pool. Idempotent.""" + self._pool.shutdown(wait=True) + self._session.close() + + # ---- internals ---- + + def _fanout_uniform(self, endpoint: str, body: Optional[Dict[str, Any]]) -> None: + self._fanout([(url, endpoint, body) for url in self._urls]) + + def _fanout(self, calls: Sequence[Tuple[str, str, Optional[Dict[str, Any]]]]) -> None: + """POST to every server concurrently; raise the first failure after all + return. + + Every future is drained before raising, so a failure on one server never + leaves POSTs in flight against the others.""" + futures = [self._pool.submit(self._post, url, endpoint, body) for url, endpoint, body in calls] + first_exc: Optional[BaseException] = None + for fut in futures: + try: + fut.result() + except Exception as exc: # noqa: BLE001 + first_exc = first_exc or exc + if first_exc is not None: + raise first_exc + + def _post(self, url: str, endpoint: str, body: Optional[Dict[str, Any]]) -> None: + # No timeout, no retry -- see module docstring. + resp = self._session.post(f"{url}{endpoint}", json=body, timeout=None) + if resp.status_code >= 400: + raise RuntimeError(_error_message(url, endpoint, resp)) + + +def _json_safe(update_info: Dict[str, Any]) -> Dict[str, Any]: + """Make an update_info dict JSON-serializable for HTTP transport. + + ``IPCTrainerWeightTransferEngine`` emits raw CUDA IPC handles -- tuples + containing storage types and handle bytes -- because a Ray transport carries + them natively. Over HTTP they must be pickled into ``ipc_handles_pickled``, + which the worker auto-deserializes when + ``VLLM_ALLOW_INSECURE_SERIALIZATION=1`` (set by ``vllm_server_actor``). + + Other backends carry only JSON-native metadata and pass through unchanged. + """ + handles = update_info.get("ipc_handles") + if handles is None: + return update_info + + import base64 + import pickle + + out = {k: v for k, v in update_info.items() if k != "ipc_handles"} + out["ipc_handles_pickled"] = base64.b64encode(pickle.dumps(handles)).decode("utf-8") + return out + + +def nccl_init_payloads( + init_info: Dict[str, Any], + server_urls: Sequence[str], + data_parallel_size: int, +) -> List[Dict[str, Any]]: + """Per-server NCCL init payloads with cumulative ``rank_offset``. + + ``worker_init_process_group`` derives a worker's rank as + ``dp_rank * world_size_per_dp + rank_within_dp + rank_offset``, and each + deployment's indices restart at 0, so every deployment's offset must advance + past the previous one's workers or their ranges collide. The engine builds + one init info with ``rank_offset=1``, which is only correct for a single + deployment. + + DP servers within one deployment share an offset -- vLLM's + ``data_parallel_index`` already separates them, so advancing per DP server + would double-count. The offset advances only at deployment boundaries. + + A wrong offset mis-maps ranks and hangs in the NCCL rendezvous rather than + failing, so ``world_size`` is cross-checked against what the offsets imply. + """ + dp = max(1, int(data_parallel_size)) + num_servers = len(server_urls) + if num_servers % dp != 0: + raise ValueError( + f"Number of servers ({num_servers}) must be divisible by data_parallel_size ({dp})." + ) + base_offset = int(init_info.get("rank_offset", 1)) + world_size = int(init_info["world_size"]) + num_deployments = max(1, num_servers // dp) + # world_size counts the trainer sender (rank 0) plus every inference worker, + # so the per-deployment worker count follows from it. + workers_total = world_size - base_offset + if workers_total <= 0 or workers_total % num_deployments != 0: + raise ValueError( + f"NCCL weight sync world_size={world_size} with rank_offset={base_offset} implies " + f"{workers_total} inference workers, which does not divide across {num_deployments} " + f"deployment(s) ({num_servers} servers / dp={dp})." + ) + world_size_per_deployment = workers_total // num_deployments + + payloads: List[Dict[str, Any]] = [] + offset = base_offset + for i in range(num_servers): + payloads.append({**init_info, "rank_offset": offset}) + if (i + 1) % dp == 0: + offset += world_size_per_deployment + return payloads + + +def rdt_init_payloads( + init_info: Dict[str, Any], + server_urls: Sequence[str], + data_parallel_size: int, +) -> List[Dict[str, Any]]: + """Per-server sharded-RDT init payloads, stamped with the deployment ordinal. + + Each deployment has its own self-contained parallel config, so its internal + worker index restarts at 0 and would collide under the M:N block assignment. + Stamping each server with its ordinal (``server_index // data_parallel_size``) + as ``replica_rank``, plus the deployment count as ``num_replicas``, lets the + engine offset its consumers into a globally distinct range. + + The ordinal divides by ``data_parallel_size`` because the DP servers of one + deployment share a parallel config, so they must share one ``replica_rank``. + """ + dp = max(1, int(data_parallel_size)) + num_replicas = max(1, len(server_urls) // dp) + return [{**init_info, "replica_rank": i // dp, "num_replicas": num_replicas} for i in range(len(server_urls))] + + +def _error_message(url: str, endpoint: str, resp: Any) -> str: + """Surface the response body's error detail (``{"error": {"message": ...}}`` + or FastAPI's ``{"detail": ...}``) rather than the bare HTTP reason phrase.""" + detail = resp.reason + try: + body = resp.json() + if isinstance(body, dict): + err = body.get("error") + if isinstance(err, dict): + detail = err.get("message", detail) + elif isinstance(err, str): + detail = err + elif body.get("detail") is not None: + detail = body["detail"] + except Exception: # noqa: BLE001 + detail = (resp.text or resp.reason)[:1000] + return f"Weight-sync control-plane call {endpoint} to {url} failed [{resp.status_code}]: {detail}" diff --git a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py b/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py deleted file mode 100644 index 18642cca92..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/cuda_ipc_strategy.py +++ /dev/null @@ -1,316 +0,0 @@ -"""CUDA IPC-based weight transfer strategy. - -This module implements the CUDA IPC transfer strategy for synchronizing model weights -from training workers to inference engines using CUDA IPC handles. -""" - -import base64 -import copy -import pickle -from dataclasses import asdict, dataclass -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - Iterable, - List, - Optional, - Tuple, -) - -if TYPE_CHECKING: - from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( - RemoteInferenceClient, - ) - from skyrl.train.config import InferenceEngineConfig - -import torch -from torch.multiprocessing.reductions import reduce_tensor - -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk, WeightUpdateRequest -from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( - WeightSyncInitInfo, - WeightTransferSender, - WeightTransferStrategy, -) -from skyrl.train.utils.utils import str_to_torch_dtype - -# IPC handle type: (rebuild_func, args) returned by reduce_tensor -IpcHandle = Tuple[Callable[..., torch.Tensor], Tuple[Any, ...]] - - -@dataclass -class CudaIpcInitInfo(WeightSyncInitInfo): - """Initialization info for CUDA IPC-based weight transfer.""" - - model_dtype_str: str - - def for_servers(self, world_size_per_server: int, num_servers: int, dp_size: int = 1) -> List["CudaIpcInitInfo"]: - """IPC init is a no-op, so return identical copies for each server.""" - return [copy.deepcopy(self) for _ in range(num_servers)] - - def to_api_payload(self) -> Dict[str, Any]: - """IPC needs no initialization parameters.""" - return {} - - -_IPC_REQUEST_END_MARKER = b"__END_OF_REQUEST__" - - -@dataclass -class CudaIpcWeightUpdateRequest(WeightUpdateRequest): - """Request for CUDA IPC-based weight transfer. - - Contains IPC handles for direct GPU memory access. Tensors are packed into - a contiguous buffer to reduce the number of IPC handles. - """ - - sizes: List[int] # Size in elements per parameter (for unpacking) - ipc_handles: Dict[str, IpcHandle] # Physical GPU UUID -> IPC handle for the packed buffer - - def serialize(self) -> bytes: - """Serialize the request to bytes.""" - import base64 - import pickle - - request_data = pickle.dumps(self) - request_data_encoded = base64.b64encode(request_data) - data_with_marker = request_data_encoded + _IPC_REQUEST_END_MARKER - - # Pad for 4-byte alignment - data_size = len(data_with_marker) - padded_size = ((data_size + 3) // 4) * 4 - result = bytearray(data_with_marker) - result.extend(b"\x00" * (padded_size - data_size)) - return bytes(result) - - @classmethod - def deserialize(cls, data: bytes) -> "CudaIpcWeightUpdateRequest": - """Deserialize the request from bytes.""" - import base64 - import pickle - - end_index = data.find(_IPC_REQUEST_END_MARKER) - if end_index == -1: - raise ValueError("End marker not found in serialized data") - request_data = data[:end_index] - try: - request_data_decoded = base64.b64decode(request_data) - return pickle.loads(request_data_decoded) - except Exception as e: - raise ValueError("Failed to deserialize request") from e - - def to_json_dict(self) -> Dict[str, Any]: - """Serialize the request to JSON.""" - data = asdict(self) - # serialize the ipc handle - import base64 - import pickle - - data["ipc_handles"] = base64.b64encode(pickle.dumps(self.ipc_handles)).decode("utf-8") - return data - - @classmethod - def from_json_dict(cls, data: Dict[str, Any]) -> "CudaIpcWeightUpdateRequest": - """Deserialize the request from JSON.""" - import base64 - import pickle - - data = data.copy() - data["ipc_handles"] = pickle.loads(base64.b64decode(data["ipc_handles"])) - return cls(**data) - - -class CudaIpcWeightTransferSender(WeightTransferSender): - """Sends weights via CUDA IPC handles. - - Creates IPC handles for tensors, gathers them across ranks, and sends - the handle metadata to inference engines. When using the new inference - path, sends handles via vLLM's native /update_weights endpoint. - """ - - def __init__( - self, - init_info: CudaIpcInitInfo, - inference_client: "RemoteInferenceClient", - ) -> None: - """Initialize the CUDA IPC sender. - - Args: - init_info: CudaIpcInitInfo containing config-derived args. - inference_client: Client for coordinating with inference engines. - """ - self._init_info = init_info - self._inference_client = inference_client - - async def send_chunks( - self, - chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, - **kwargs, - ) -> None: - """Send chunks via CUDA IPC. - - Args: - chunks: Iterable of WeightChunk objects to send. - weight_metadata: Unused for IPC (metadata is derived from chunks - directly to avoid ordering mismatches). Kept for interface - compatibility with the base class. - """ - await self._send_chunks_vllm_native(chunks, weight_metadata) - - async def _send_chunks_vllm_native( - self, - chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, - ) -> None: - """Send weights chunk-by-chunk via vLLM native IPC (new inference path). - - Uses the start/update/finish lifecycle to enable chunked transfers. - Per chunk, all tensors are packed into a single contiguous CUDA buffer - (one dtype per chunk, guaranteed by the weight extractor) and one IPC - handle is created for the packed buffer per rank. - - All ranks iterate chunks (weight extraction may use collective ops). - Per chunk, each rank packs + creates one IPC handle, handles are - all_gather_object'd into a single {gpu_uuid: handle} dict, and rank 0 - sends the dict (plus per-param `sizes` metadata) via - update_weights_ipc. The receiver rebuilds the packed tensor, slices - it per param, and loads into vLLM. - - TODO: Once https://github.com/vllm-project/vllm/pull/39212 lands, - replace update_weights_ipc with the native /update_weights endpoint - and start/finish with /start_weight_update and /finish_weight_update. - """ - rank = torch.distributed.get_rank() - world_size = torch.distributed.get_world_size() - device = torch.cuda.current_device() - gpu_uuid = str(torch.cuda.get_device_properties(device).uuid) - dtype = str_to_torch_dtype(self._init_info.model_dtype_str) - dtype_name = self._init_info.model_dtype_str.split(".")[-1] - - if rank == 0: - await self._inference_client.start_weight_update(is_checkpoint_format=True) - torch.distributed.barrier() - - for chunk in chunks: - # --- pack all tensors in this chunk into one contiguous buffer --- - # Chunk tensors share a single dtype by construction (see - # weight_extractor_utils.py), so offsets in element units are safe. - names: List[str] = [] - dtype_names: List[str] = [] - shapes: List[List[int]] = [] - sizes: List[int] = [] - - total_numel = sum(t.numel() for t in chunk.tensors) - packed_tensor = torch.empty( - total_numel, - device=device, - dtype=dtype, - requires_grad=False, - ) - - offset = 0 - for name, tensor, shape in zip(chunk.names, chunk.tensors, chunk.shapes): - size = tensor.numel() - packed_tensor[offset : offset + size].copy_(tensor.detach().reshape(-1)) - offset += size - names.append(name) - dtype_names.append(dtype_name) - shapes.append(list(shape) if not isinstance(shape, list) else shape) - sizes.append(size) - - # --- one IPC handle per rank for the packed buffer --- - ipc_handle: IpcHandle = reduce_tensor(packed_tensor) - local_handle_dict: Dict[str, IpcHandle] = {gpu_uuid: ipc_handle} - gathered: List[Optional[Dict[str, IpcHandle]]] = [None] * world_size - torch.distributed.all_gather_object(gathered, local_handle_dict) - - torch.distributed.barrier() - torch.cuda.synchronize() - - if rank == 0: - merged_handles: Dict[str, IpcHandle] = {} - for d in gathered: - if d is not None: - merged_handles.update(d) - - pickled = base64.b64encode(pickle.dumps(merged_handles)).decode("utf-8") - chunk_update_info: Dict[str, Any] = { - "names": names, - "dtype_names": dtype_names, - "shapes": shapes, - "sizes": sizes, - "ipc_handles_pickled": pickled, - } - await self._inference_client.update_weights_ipc(chunk_update_info) - - # Keep packed_tensor alive past the barrier so the receiver's - # rebuilt view has valid backing storage while it copies into - # the model. Post-barrier drops the local ref safely. - torch.distributed.barrier() - torch.cuda.ipc_collect() - torch.cuda.synchronize() - - if rank == 0: - await self._inference_client.finish_weight_update() - torch.distributed.barrier() - - def teardown(self) -> None: - """No-op for CUDA IPC sender (no custom process group to clean up).""" - pass - - -class CudaIpcTransferStrategy(WeightTransferStrategy): - """Factory for CUDA IPC-based weight transfer. - - This strategy uses CUDA IPC handles to share GPU memory between training - workers and inference engines on the same machine. - - All methods are static - no instance state needed. - """ - - @staticmethod - def create_init_info( - ie_cfg: "InferenceEngineConfig", - inference_world_size: Optional[int] = None, - base_model_path: Optional[str] = None, - ) -> CudaIpcInitInfo: - """Create init info with all config-derived args.""" - return CudaIpcInitInfo( - model_dtype_str=ie_cfg.model_dtype, - override_existing_receiver=not ie_cfg.run_engines_locally, - ) - - @staticmethod - def create_sender( - init_info: CudaIpcInitInfo, - inference_client: "RemoteInferenceClient", - weight_extractor: Optional[Any] = None, - ) -> CudaIpcWeightTransferSender: - """Create a CUDA IPC sender. - - Args: - init_info: CudaIpcInitInfo containing config-derived args. - inference_client: Client for coordinating with inference engines. - - Returns: - A configured CudaIpcWeightTransferSender instance. - """ - return CudaIpcWeightTransferSender( - init_info=init_info, - inference_client=inference_client, - ) - - @staticmethod - def get_vllm_transfer_engine() -> type: - """Return the vLLM weight-transfer engine class for this strategy (CUDA IPC). - - Reference for the receive side: the inference servers drive this engine - natively. Currently unused on the sender side (we route through the - SkyRL ``/collective_rpc`` wrap), kept as the canonical mapping. - """ - from vllm.distributed.weight_transfer.ipc_engine import IPCWeightTransferEngine - - return IPCWeightTransferEngine diff --git a/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py b/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py index 41c3a55a64..12d895a2f7 100644 --- a/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py +++ b/skyrl/backends/skyrl_train/weight_sync/delta_checkpoint.py @@ -25,7 +25,6 @@ from safetensors import safe_open from safetensors.torch import save, save_file -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk from skyrl.backends.skyrl_train.weight_sync.delta_payload import ( bytes_to_uint8_tensor, compress_bytes, @@ -1105,9 +1104,20 @@ def _staging_pool_for(self, max_tensor_bytes: int, num_workers: int) -> _PinnedS def create_delta_files( self, - chunks: Iterable[WeightChunk], + source: Iterable[Tuple[str, torch.Tensor]], ) -> DeltaPublishResult: - """Iterate over weight chunks and create delta files locally""" + """Drain a ``WeightSource`` and create this rank's delta files locally. + + Takes the ``(name, tensor)`` stream directly rather than a chunk stream: + the chunk boundary was never load-bearing here — no batching, no flush + decision, no size accounting (that is ``payload_file_bytes``, tracked per + tensor). All it ever did was get flattened. + + **Every rank must call this**, source or not: iterating the source is + what drives its gather collectives (FSDP ``full_tensor()``, a Megatron + export), so a rank that skipped it would deadlock its peers. Non-source + ranks drain and drop, and return an empty result. + """ base_version = self.version target_version = base_version + 1 publish_uri = _join_uri(self.sync_dir, _version_name(target_version)) @@ -1184,47 +1194,43 @@ def drain_one() -> None: max_tensor_bytes = max((loc.nbytes for loc in locations.values()), default=0) staging_pool = self._staging_pool_for(max_tensor_bytes, num_workers) executor = self._publish_executor_for(num_workers) if is_source_rank else None - chunk_iter = iter(chunks) - while True: - try: - chunk = next(chunk_iter) - except StopIteration: - break + for name, tensor in source: if not is_source_rank: + # Drain and drop: the iteration itself is the collective. + del tensor continue - for name, tensor in zip(chunk.names, chunk.tensors): - try: - loc = self._base_location(name) - base = self._snapshot_bytes(name) - except KeyError: - skipped_names.add(name) - stats["skipped_tensors"] += 1 - continue - - t = time.perf_counter() - staged, dtype_name, shape = self._stage_tensor_for_publish( - tensor, - target_dtype=_safetensors_dtype_to_torch(loc.dtype), - target_shape=loc.shape, - staging_pool=staging_pool, - ) - stats["cpu_stage_s"] += time.perf_counter() - t - if executor is None: - raise RuntimeError("Delta checkpoint source rank is missing a publish executor") - pending.append( - executor.submit( - self._process_tensor_delta, - name, - staged, - base, - dtype_name, - shape, - staging_pool, - self.checksum_algorithm, - ) + try: + loc = self._base_location(name) + base = self._snapshot_bytes(name) + except KeyError: + skipped_names.add(name) + stats["skipped_tensors"] += 1 + continue + + t = time.perf_counter() + staged, dtype_name, shape = self._stage_tensor_for_publish( + tensor, + target_dtype=_safetensors_dtype_to_torch(loc.dtype), + target_shape=loc.shape, + staging_pool=staging_pool, + ) + stats["cpu_stage_s"] += time.perf_counter() - t + if executor is None: + raise RuntimeError("Delta checkpoint source rank is missing a publish executor") + pending.append( + executor.submit( + self._process_tensor_delta, + name, + staged, + base, + dtype_name, + shape, + staging_pool, + self.checksum_algorithm, ) - if len(pending) >= max_inflight: - drain_one() + ) + if len(pending) >= max_inflight: + drain_one() while pending: drain_one() flush_payload_file() diff --git a/skyrl/backends/skyrl_train/weight_sync/delta_engine.py b/skyrl/backends/skyrl_train/weight_sync/delta_engine.py index 017deb3dce..09e46d70ec 100644 --- a/skyrl/backends/skyrl_train/weight_sync/delta_engine.py +++ b/skyrl/backends/skyrl_train/weight_sync/delta_engine.py @@ -5,13 +5,14 @@ import logging import time from dataclasses import dataclass -from typing import Any, Iterator +from typing import Any import torch from skyrl.backends.skyrl_train.weight_sync.delta_checkpoint import ( LocalCheckpointStore, ) +from skyrl.backends.skyrl_train.weight_sync.skyrl_engines import empty_cuda_cache_rocm try: from vllm.logger import init_logger @@ -36,10 +37,6 @@ class DeltaTransferUpdateInfo: sync_dir: str | None = None uri: str | None = None version: int | None = None - # vLLM 0.23's native /update_weights path checks this attribute before it - # dispatches into the custom transfer engine. Delta checkpoint sync always - # reloads dense prepared checkpoint tensors. - update_kind: str = "dense" @property def resolved_target_version(self) -> int: @@ -65,7 +62,11 @@ def register_delta_weight_transfer_engine() -> None: class DeltaWeightTransferEngine: - """Receive compressed checkpoint deltas and load updated weights into vLLM.""" + """Receive compressed checkpoint deltas and load updated weights into vLLM. + + Duck-typed against vLLM's ``WeightTransferEngine`` rather than subclassing + it, so this module stays importable without vLLM installed. + """ init_info_cls = DeltaTransferInitInfo update_info_cls = DeltaTransferUpdateInfo @@ -75,10 +76,8 @@ class DeltaWeightTransferEngine: supports_draft_weight_update = False def __init__(self, config: Any, vllm_config: Any, device: Any, model: torch.nn.Module) -> None: - # Signature mirrors vLLM's WeightTransferEngine base (0.26+): - # WeightTransferEngineFactory.create_engine calls - # engine_cls(config, vllm_config, device, model). Duck-typed rather than - # subclassed so this module stays importable without vLLM installed. + # Signature fixed by WeightTransferEngineFactory.create_engine, which + # calls engine_cls(config, vllm_config, device, model). self.config = config self.vllm_config = vllm_config self.parallel_config = getattr(vllm_config, "parallel_config", None) @@ -103,20 +102,19 @@ def reset_weight_update_target(self) -> None: self.model_config = self._default_model_config def start_weight_update(self) -> None: - """No-op: SkyRL drives the layerwise-reload lifecycle from the worker. + """Initialize layerwise reloading for the incoming checkpoint weights.""" + from vllm.model_executor.model_loader.reload import initialize_layerwise_reload - vLLM's ``Worker.start_weight_update`` delegates here, but SkyRL's delta - flow calls ``NewInferenceWorkerWrap.skyrl_start_weight_update`` instead - (see DeltaWeightTransferSender._apply_receiver_update), which is what - initializes layerwise reload. Doing it again here would double-initialize. - """ + with torch.device(self.device): + initialize_layerwise_reload(self.model) def finish_weight_update(self) -> None: - """No-op counterpart to :meth:`start_weight_update`. + """Finalize layerwise reloading after the checkpoint has been loaded.""" + from vllm.model_executor.model_loader.reload import finalize_layerwise_reload - SkyRL finalizes layerwise reload via - ``NewInferenceWorkerWrap.skyrl_finish_weight_update``. - """ + with torch.device(self.device): + finalize_layerwise_reload(self.model, self.model_config) + empty_cuda_cache_rocm() def update_weights(self, update_info: dict[str, Any]) -> None: """Load one update, as vLLM's native ``/update_weights`` endpoint expects.""" @@ -176,12 +174,27 @@ def receive_weights(self, update_info: DeltaTransferUpdateInfo) -> None: prepare_s = time.perf_counter() - t0 load_s = 0.0 t1 = time.perf_counter() - self.model.load_weights( - self._store.iter_tensors( + + def tensors(): + return self._store.iter_tensors( load_format=self._checkpoint_load_format, multi_thread_safetensors_max_workers=self._multi_thread_safetensors_max_workers, ) + + # MTP architectures raise on incomplete layer coverage. + from vllm.model_executor.model_loader.mtp_validation import ( + disable_mtp_completeness_check, ) + + with torch.device(self.device), disable_mtp_completeness_check(): + self.model.load_weights(tensors()) + # The spec-decode drafter is a separate module the main load never + # touches. Unlike the push backends this does NOT go through + # skyrl_drafter_reload's proxy: that materializes the weight list so + # the drafter can re-read it, which for a whole checkpoint would be + # the entire model resident at once. The store re-streams instead. + self._reload_drafter(tensors) + load_s = time.perf_counter() - t1 total_s = time.perf_counter() - t0 message = ( @@ -197,11 +210,25 @@ def receive_weights(self, update_info: DeltaTransferUpdateInfo) -> None: logger.info(message) print(message, flush=True) + def _reload_drafter(self, tensors) -> None: + """Reload the spec-decode drafter from a fresh pass over the checkpoint. + + No-op, and no second pass, when this process has no loadable proposer. + """ + from skyrl.backends.skyrl_train.patches.vllm.patch_model_runner_registry import ( + current_model_runner, + ) + + model_runner = current_model_runner() + drafter = getattr(model_runner, "drafter", None) if model_runner is not None else None + if drafter is None or getattr(drafter, "model", None) is None: + return + + from skyrl.backends.skyrl_train.inference_servers.spec_decode_utils import ( + _reload_spec_decode_drafter, + ) + + _reload_spec_decode_drafter(model_runner, tensors()) + def shutdown(self): self._store = None - - @staticmethod - def trainer_send_weights( - _iterator: Iterator[tuple[str, torch.Tensor]], _trainer_args: dict[str, Any] | Any - ) -> None: - raise NotImplementedError("Delta weight sync publishes through SkyRL's DeltaWeightTransferSender") diff --git a/skyrl/backends/skyrl_train/weight_sync/delta_strategy.py b/skyrl/backends/skyrl_train/weight_sync/delta_strategy.py deleted file mode 100644 index 4d4aae37ad..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/delta_strategy.py +++ /dev/null @@ -1,190 +0,0 @@ -"""Disk/cloud checkpoint-delta transfer strategy.""" - -from __future__ import annotations - -import asyncio -import copy -import logging -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional - -import torch - -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk -from skyrl.backends.skyrl_train.weight_sync.delta_checkpoint import ( - SUPPORTED_CHECKPOINT_LOAD_FORMATS, - DeltaCheckpointPublisher, - DeltaPublishResult, -) -from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( - WeightSyncInitInfo, - WeightTransferSender, - WeightTransferStrategy, -) - -if TYPE_CHECKING: - from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( - RemoteInferenceClient, - ) - from skyrl.train.config import InferenceEngineConfig - - -logger = logging.getLogger(__name__) - - -@dataclass -class DeltaInitInfo(WeightSyncInitInfo): - base_model_path: str - sync_dir: str - local_checkpoint_dir: str - publish_staging_dir: str - max_file_size_in_gb: float = 1.0 - cloud_download_workers: int = 4 - publish_num_workers: Optional[int] = None - checkpoint_load_format: str = "vllm_multi_thread_safetensors" - multi_thread_safetensors_max_workers: int = 8 - - def for_servers(self, world_size_per_server: int, num_servers: int, dp_size: int = 1) -> List["DeltaInitInfo"]: - return [copy.deepcopy(self) for _ in range(num_servers)] - - def to_api_payload(self) -> Dict[str, Any]: - return { - "base_model_path": self.base_model_path, - "local_checkpoint_dir": self.local_checkpoint_dir, - "cloud_download_workers": self.cloud_download_workers, - "checkpoint_load_format": self.checkpoint_load_format, - "multi_thread_safetensors_max_workers": self.multi_thread_safetensors_max_workers, - } - - -class DeltaWeightTransferSender(WeightTransferSender): - - handles_prefix_cache_reset = True - """Indicates whether the transfer strategy handles resetting prefix cache - for the inference engines internally.""" - - def __init__(self, init_info: DeltaInitInfo, inference_client: "RemoteInferenceClient") -> None: - self._init_info = init_info - self._inference_client = inference_client - self._publisher: Optional[DeltaCheckpointPublisher] = None - - async def _apply_receiver_update(self, update_info: Dict[str, Any], rank: int, reset_prefix_cache: bool) -> None: - - target_version = int(update_info.get("target_version", update_info.get("version"))) - await self._inference_client.fetch_weights( - target_version=target_version, - sync_dir=update_info.get("sync_dir", self._init_info.sync_dir), - uri=update_info.get("uri"), - ) - await self._inference_client.pause_generation() - try: - if reset_prefix_cache: - await self._inference_client.reset_prefix_cache(reset_running_requests=True) - await self._inference_client.start_weight_update(is_checkpoint_format=True) - await self._inference_client.update_named_weights(update_info) - await self._inference_client.finish_weight_update() - finally: - await self._inference_client.resume_generation() - - async def send_chunks( - self, - chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, - reset_prefix_cache: bool = False, - ) -> None: - if torch.distributed.is_available() and torch.distributed.is_initialized(): - rank = torch.distributed.get_rank() - else: - rank = 0 - - if self._publisher is None: - self._publisher = DeltaCheckpointPublisher( - base_model_path=self._init_info.base_model_path, - sync_dir=self._init_info.sync_dir, - publish_staging_dir=self._init_info.publish_staging_dir, - max_file_size_in_gb=self._init_info.max_file_size_in_gb, - publish_num_workers=self._init_info.publish_num_workers, - ) - - # All ranks call publish to drain the chunk stream and drive the - # extractor's collectives. Only rank 0 publishes the deltas and - # runs finalization. - local_result = await asyncio.to_thread(self._publisher.create_delta_files, chunks) - if not isinstance(local_result, DeltaPublishResult): - raise TypeError(f"Expected DeltaPublishResult from sharded publisher, got {type(local_result)}") - - is_distributed = torch.distributed.is_available() and torch.distributed.is_initialized() - world_size = torch.distributed.get_world_size() if is_distributed else 1 - - if is_distributed: - gathered_results: list[Optional[DeltaPublishResult]] = [None] * world_size - torch.distributed.all_gather_object(gathered_results, local_result) - else: - gathered_results = [local_result] - - update_info = None - if rank == 0: - source_results = [result for result in gathered_results if result is not None and result.rank == 0] - update_info = self._publisher.publish(source_results) - - if is_distributed: - update_info_box = [update_info] - torch.distributed.broadcast_object_list(update_info_box, src=0) - update_info = update_info_box[0] - - if rank == 0 and update_info is not None: - await self._apply_receiver_update(update_info, rank, reset_prefix_cache) - - def teardown(self) -> None: - self._publisher = None - - -class DeltaTransferStrategy(WeightTransferStrategy): - @staticmethod - def create_init_info( - ie_cfg: "InferenceEngineConfig", - inference_world_size: Optional[int] = None, - base_model_path: Optional[str] = None, - ) -> DeltaInitInfo: - if base_model_path is None: - raise ValueError("Delta weight sync requires base_model_path") - delta_cfg = ie_cfg.delta_weight_sync - # Ensure valid `delta_cfg` - if delta_cfg is None or not delta_cfg.sync_dir: - raise ValueError("Delta weight sync requires generator.inference_engine.delta_weight_sync.sync_dir") - - # local_checkpoint_dir and publish_staging_dir are resolved by - # DeltaWeightSyncConfig.__post_init__, so they are already concrete here. - if delta_cfg.checkpoint_load_format not in SUPPORTED_CHECKPOINT_LOAD_FORMATS: - raise ValueError( - "Delta checkpoint_load_format must be one of " - f"{sorted(SUPPORTED_CHECKPOINT_LOAD_FORMATS)}, got {delta_cfg.checkpoint_load_format!r}" - ) - return DeltaInitInfo( - base_model_path=base_model_path, - sync_dir=delta_cfg.sync_dir, - local_checkpoint_dir=delta_cfg.local_checkpoint_dir, - publish_staging_dir=delta_cfg.publish_staging_dir, - max_file_size_in_gb=delta_cfg.max_file_size_in_gb, - cloud_download_workers=delta_cfg.cloud_download_workers, - publish_num_workers=delta_cfg.publish_num_workers, - checkpoint_load_format=delta_cfg.checkpoint_load_format, - multi_thread_safetensors_max_workers=delta_cfg.multi_thread_safetensors_max_workers, - override_existing_receiver=not ie_cfg.run_engines_locally, - ) - - @staticmethod - def create_sender( - init_info: DeltaInitInfo, - inference_client: "RemoteInferenceClient", - weight_extractor: Optional[Any] = None, - ) -> DeltaWeightTransferSender: - return DeltaWeightTransferSender(init_info=init_info, inference_client=inference_client) - - @staticmethod - def get_vllm_transfer_engine() -> type: - from skyrl.backends.skyrl_train.weight_sync.delta_engine import ( - DeltaWeightTransferEngine, - ) - - return DeltaWeightTransferEngine diff --git a/skyrl/backends/skyrl_train/weight_sync/delta_trainer.py b/skyrl/backends/skyrl_train/weight_sync/delta_trainer.py new file mode 100644 index 0000000000..acc805ee2d --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/delta_trainer.py @@ -0,0 +1,198 @@ +"""Trainer-side engine for the checkpoint-delta weight-sync backend. + +The trainer publishes a compressed delta against the base checkpoint to disk or +object storage; each inference worker fetches and reloads it. Nothing crosses the +network fabric, but the shape is vLLM's trainer-send: a ``WeightSource`` in, a +four-method ``VLLMWeightSyncClient`` out, ``send_weights()`` owning the round +trip. + +Two parts of that round trip ride the engine rather than the client protocol, and +cannot be hoisted to the worker: ``fetch_weights`` needs the ``target_version`` +the publish produces mid-send, and must run **before** the pause so the download +overlaps live generation. The pause / resume bracket and the prefix-cache reset +then wrap the reload -- this is the only backend that reloads a whole checkpoint +in place, so the only one that needs generation stopped. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar, Optional + +import torch +from vllm.distributed.weight_transfer.base import ( + TrainerInitInfo, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightSource, +) + +from skyrl.backends.skyrl_train.weight_sync.delta_checkpoint import ( + DeltaCheckpointPublisher, + DeltaPublishResult, +) + +if TYPE_CHECKING: + from typing_extensions import Self + +logger = logging.getLogger(__name__) + +DELTA_BACKEND = "delta" + + +@dataclass +class DeltaTrainerInitInfo(TrainerInitInfo): + """Trainer-side init info for checkpoint-delta weight sync. + + ``base_model_path``, ``local_checkpoint_dir``, the load format and its worker + counts are propagated to the worker at ``trainer_init``; the rest are + publisher-side only. + """ + + backend: ClassVar[str] = DELTA_BACKEND + + base_model_path: str + sync_dir: str + local_checkpoint_dir: str + publish_staging_dir: str + max_file_size_in_gb: float = 1.0 + cloud_download_workers: int = 4 + publish_num_workers: Optional[int] = None + checkpoint_load_format: str = "vllm_multi_thread_safetensors" + multi_thread_safetensors_max_workers: int = 8 + + +class DeltaTrainerWeightTransferEngine(TrainerWeightTransferEngine[DeltaTrainerInitInfo]): + """Publish a checkpoint delta and drive the inference-side reload.""" + + init_info_cls = DeltaTrainerInitInfo + + # Resets the prefix cache itself, inside the pause bracket, so the worker + # must not also fire a concurrent reset (see workers/worker.py). + skyrl_handles_prefix_cache_reset = True + + def __init__( + self, + *, + client: VLLMWeightSyncClient, + source: WeightSource, + is_sender: bool = True, + init_info: DeltaTrainerInitInfo, + ) -> None: + super().__init__(client=client, source=source, is_sender=is_sender) + self._init_info = init_info + self._publisher: Optional[DeltaCheckpointPublisher] = None + self._reset_prefix_cache = False + + @classmethod + def trainer_init( + cls, + init_info: DeltaTrainerInitInfo, + *, + client: VLLMWeightSyncClient, + source: Optional[WeightSource] = None, + ) -> "Self": + if source is None: + raise ValueError("Delta trainer weight transfer requires a WeightSource.") + engine = cls( + client=client, + source=source, + is_sender=init_info.is_sender, + init_info=init_info, + ) + if engine.is_sender: + # No data-plane rendezvous; this builds each worker's local + # checkpoint store and records the load format. + engine.client.init_weight_transfer_engine( + { + "base_model_path": init_info.base_model_path, + "local_checkpoint_dir": init_info.local_checkpoint_dir, + "cloud_download_workers": init_info.cloud_download_workers, + "checkpoint_load_format": init_info.checkpoint_load_format, + "multi_thread_safetensors_max_workers": init_info.multi_thread_safetensors_max_workers, + } + ) + return engine + + def skyrl_set_reset_prefix_cache(self, reset: bool) -> None: + """Whether this round should reset the prefix cache inside the pause.""" + self._reset_prefix_cache = bool(reset) + + def send_weights(self) -> None: + """Publish this round's delta and reload it on every inference worker. + + Called on **every** trainer rank: draining the source drives its gather + collectives, so a rank that skipped it would hang its peers. Only rank 0 + stages and uploads payloads, and only rank 0 drives the inference side. + """ + assert self.source is not None # guaranteed by trainer_init + distributed = torch.distributed.is_available() and torch.distributed.is_initialized() + rank = torch.distributed.get_rank() if distributed else 0 + + if self._publisher is None: + self._publisher = DeltaCheckpointPublisher( + base_model_path=self._init_info.base_model_path, + sync_dir=self._init_info.sync_dir, + publish_staging_dir=self._init_info.publish_staging_dir, + max_file_size_in_gb=self._init_info.max_file_size_in_gb, + publish_num_workers=self._init_info.publish_num_workers, + ) + + local_result = self._publisher.create_delta_files(self.source) + if not isinstance(local_result, DeltaPublishResult): + raise TypeError(f"Expected DeltaPublishResult from the publisher, got {type(local_result)}") + + if distributed: + world_size = torch.distributed.get_world_size() + gathered: list[Optional[DeltaPublishResult]] = [None] * world_size + torch.distributed.all_gather_object(gathered, local_result) + else: + gathered = [local_result] + + update_info = None + if rank == 0: + source_results = [r for r in gathered if r is not None and r.rank == 0] + update_info = self._publisher.publish(source_results) + + if distributed: + box = [update_info] + torch.distributed.broadcast_object_list(box, src=0) + update_info = box[0] + + if self.is_sender and update_info is not None: + self._apply_receiver_update(update_info) + + def _apply_receiver_update(self, update_info: dict) -> None: + target_version = int(update_info.get("target_version", update_info.get("version"))) + # Before the pause, so the download overlaps live generation. + self.client.fetch_weights( + target_version=target_version, + sync_dir=update_info.get("sync_dir", self._init_info.sync_dir), + uri=update_info.get("uri"), + ) + self.client.pause_generation() + try: + if self._reset_prefix_cache: + self.client.reset_prefix_cache(reset_running_requests=True) + self.client.start_weight_update() + self.client.update_weights(update_info) + self.client.finish_weight_update() + finally: + self.client.resume_generation() + + def shutdown(self) -> None: + self._publisher = None + + +def register_delta_trainer_engine() -> None: + """Register ``delta`` in ``WeightTransferTrainerFactory`` (idempotent).""" + from vllm.distributed.weight_transfer.factory import WeightTransferTrainerFactory + + if DELTA_BACKEND in WeightTransferTrainerFactory._registry: + return + WeightTransferTrainerFactory.register_engine( + DELTA_BACKEND, + "skyrl.backends.skyrl_train.weight_sync.delta_trainer", + "DeltaTrainerWeightTransferEngine", + ) diff --git a/skyrl/backends/skyrl_train/weight_sync/nccl_trainer_send.py b/skyrl/backends/skyrl_train/weight_sync/nccl_trainer_send.py deleted file mode 100644 index 0f2998db8a..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/nccl_trainer_send.py +++ /dev/null @@ -1,110 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Trainer-side NCCL send, vendored from vLLM 0.26.0. - -REMOVAL: ``NCCLWeightTransferEngine.trainer_send_weights`` and the engine's -``trainer_init`` staticmethod were deleted in vLLM 0.28.0, which replaced them -with a trainer-side engine abstraction (``NCCLTrainerWeightTransferEngine``, -dispatched through ``WeightTransferTrainerFactory``). SkyRL still drives the send -itself -- the sender owns the ``/collective_rpc`` round trip in -``broadcast_strategy.py`` -- so this module keeps 0.26.0's behaviour available -under the new pin. Delete it when SkyRL's sender layer moves onto vLLM's -trainer-send engines; that migration also retires ``WeightTransferStrategy``, -``WeightTransferSender`` and ``NewInferenceWorkerWrap``. - -``packed_nccl_broadcast_producer`` (the real work) and ``nccl_common.trainer_init`` -(the rendezvous) both survive in 0.28.0 unchanged, so this is a thin re-wrap -rather than a copy of the transfer itself. -""" - -from collections.abc import Callable, Iterator -from typing import TYPE_CHECKING, Any, Optional, Tuple - -import torch - -if TYPE_CHECKING: - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator - -__all__ = ["nccl_trainer_send_weights", "nccl_trainer_init"] - - -def nccl_trainer_init(init_info: Any) -> "PyNcclCommunicator": - """Open the trainer-side (rank 0) endpoint of the weight-transfer group. - - Was ``NCCLWeightTransferEngine.trainer_init`` in 0.26.0, which was itself a - ``staticmethod`` alias of this helper. 0.28.0 dropped the alias but kept the - helper, so this just forwards. - - Args: - init_info: object or dict carrying ``master_address``, ``master_port`` - and ``world_size``. - """ - # Lazy import: vllm is a Linux-only optional dependency (see - # .claude/docs/weight_sync.md), so this module stays importable without it. - from vllm.distributed.weight_transfer.nccl_common import trainer_init - - return trainer_init(init_info) - - -def nccl_trainer_send_weights( - iterator: Iterator[Tuple[str, torch.Tensor]], - group: "PyNcclCommunicator", - *, - src: int = 0, - packed: bool = True, - post_iter_func: Optional[Callable[[Tuple[str, torch.Tensor]], torch.Tensor]] = None, - stream: Optional[torch.cuda.Stream] = None, - packed_buffer_size_bytes: Optional[int] = None, - packed_num_buffers: Optional[int] = None, -) -> None: - """Broadcast dense weights from the trainer to the vLLM workers. - - Vendored from ``NCCLWeightTransferEngine.trainer_send_weights`` (vLLM 0.26.0), - with the ``trainer_args`` dict/dataclass flattened into keyword arguments -- - the dataclass (``NCCLTrainerSendWeightsArgs``) was deleted alongside the method. - - ``packed`` must match what the receiving worker was told at init. As of - vLLM 0.28.0 the worker reads it from ``NCCLWeightTransferInitInfo.packed`` - (set once during ``/init_weight_transfer_engine``) rather than from the - per-round update info, so the two sides can no longer disagree per round -- - see ``BroadcastInitInfo.to_api_payload``. - - Args: - iterator: ``(name, tensor)`` pairs to send. - group: the ``PyNcclCommunicator`` from ``nccl_trainer_init``. - src: source rank; the trainer is rank 0. - packed: batch tensors into packed buffers instead of one broadcast each. - post_iter_func: maps each pair to the tensor to send. Defaults to taking - the tensor. - stream: CUDA stream for the unpacked path. Defaults to the current - stream. Ignored when ``packed`` (the producer makes its own). - packed_buffer_size_bytes: packed buffer size. Defaults to vLLM's - ``DEFAULT_PACKED_BUFFER_SIZE_BYTES``, which is also what the worker's - init info defaults to -- leave unset so the two cannot drift. - packed_num_buffers: packed buffer count, likewise defaulting to vLLM's - ``DEFAULT_PACKED_NUM_BUFFERS``. - """ - from vllm.distributed.weight_transfer.packed_tensor import ( - DEFAULT_PACKED_BUFFER_SIZE_BYTES, - DEFAULT_PACKED_NUM_BUFFERS, - packed_nccl_broadcast_producer, - ) - - if post_iter_func is None: - post_iter_func = lambda item: item[1] # noqa: E731 - - if packed: - packed_nccl_broadcast_producer( - iterator=iterator, - group=group, - src=src, - post_iter_func=post_iter_func, - buffer_size_bytes=( - DEFAULT_PACKED_BUFFER_SIZE_BYTES if packed_buffer_size_bytes is None else packed_buffer_size_bytes - ), - num_buffers=(DEFAULT_PACKED_NUM_BUFFERS if packed_num_buffers is None else packed_num_buffers), - ) - else: - send_stream = stream or torch.cuda.current_stream() - for item in iterator: - group.broadcast(post_iter_func(item), src=src, stream=send_stream) diff --git a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/__init__.py b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/__init__.py index d970026993..9cca447e42 100644 --- a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/__init__.py +++ b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/__init__.py @@ -1,10 +1,12 @@ """The ``sharded_rdt`` weight-transfer backend: RDMA/NIXL pull instead of push. -``sharded_rdt_{base,common,engine,fake,trainer}.py`` are vendored from the -``vllm-rdt-weight-sync`` fork and are deleted once SkyRL's pinned vLLM carries the -trainer-side transfer API natively. The rest is SkyRL glue: ``sharded_rdt_strategy`` -(the ``WeightTransferStrategy`` adapter), ``rdt_send`` (trainer-side driver and weight -sources), ``rdt_control_plane``, ``rdt_vllm_register``, ``rdt_libfabric_shim``. +``sharded_rdt_{common,engine,fake,trainer}.py`` are vendored from the +``vllm-rdt-weight-sync`` fork and are deleted once SkyRL's pinned vLLM carries +this backend natively. ``sharded_rdt_base`` is what remains of the vendored +trainer-side ABCs now that vLLM 0.28 ships them: the two channels a *pull* +backend needs and vLLM has no concept of (per-rank ownership, and a group index). +The rest is SkyRL glue: ``rdt_send`` (weight sources + the trainer init info), +``rdt_vllm_register``, ``rdt_libfabric_shim``. This ``__init__`` imports nothing: ``sharded_rdt_engine`` and ``sharded_rdt_trainer`` import ``vllm`` at module scope, so a re-export here would pull vllm into every diff --git a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_control_plane.py b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_control_plane.py deleted file mode 100644 index 00f91da7b5..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_control_plane.py +++ /dev/null @@ -1,152 +0,0 @@ -"""Synchronous control-plane client for the sharded-RDT trainer engine. - -The vendored trainer engine (``sharded_rdt_trainer.py``) drives the inference -side through the *synchronous* ``VLLMWeightSyncClient`` protocol — a serial -``init -> start -> update -> finish`` handshake issued once per weight sync. -SkyRL's ``RemoteInferenceClient`` exposes those same routes as **coroutines** on -the worker's event loop, so the previous glue bounced every call back onto that -loop with ``run_coroutine_threadsafe`` — coupling the (thread-run) engine to a -specific loop instance and requiring a documented "never call from the loop -thread or you deadlock" invariant. - -None of these four calls needs the async client's connection pooling or its -generation concurrency, so this client talks to the same ``/collective_rpc`` -endpoints over **blocking HTTP**. The engine then runs entirely sync-to-sync in -its worker thread with zero event-loop involvement, and the deadlock class -disappears structurally rather than by convention. - -What we deliberately keep from ``RemoteInferenceClient`` so sidestepping it costs -nothing (see ``_post`` / ``_fanout``): - -* **Per-call fresh connections** (``Connection: close``). A full training step - elapses between syncs, so any pooled keep-alive connection is stale by the next - call; reusing it races the server's ``timeout_keep_alive`` (uvicorn 5s) and - yields ECONNRESET. The async client dodges this with ``keepalive_timeout=2``; - for once-per-step control calls it is simpler and strictly safer to not keep - connections at all. -* **Concurrent fan-out** across servers. This is REQUIRED, not merely faster: - the consumers pull over NIXL in lockstep and the producer only frees a served - group once every consumer bound to it has pulled, so issuing ``update_weights`` - server-by-server would stall the producer's gather loop and deadlock. Mirrors - ``RemoteInferenceClient._call_all_servers``. -* **Body-aware error messages** (``_error_message``), matching the client's - ``raise_for_status`` — surface the response body's error detail, not just the - HTTP reason phrase. -* **No timeout** (bake + NIXL pull are long) and **no retry** (retrying a - half-done bake/pull would be wrong; ``Connection: close`` already removes the - only transient race the async control path guarded against). - -``requests`` is used rather than adding a new dependency: the sharded_rdt path -already hard-requires Ray, and Ray depends on ``requests``, so it is always -importable wherever this client runs. The import is local so non-RDT / non-Ray -code paths never pay for it. -""" - -import logging -from concurrent.futures import ThreadPoolExecutor -from typing import Any, Dict, List, Optional, Sequence, Tuple - -from skyrl.backends.skyrl_train.inference_servers.rdt_control_protocol import ( - COLLECTIVE_RPC_ENDPOINT, - RDT_FINISH_METHOD, - RDT_START_METHOD, - RDT_UPDATE_METHOD, - build_rdt_init_payloads, -) - -logger = logging.getLogger(__name__) - - -class SyncRdtControlPlaneClient: - """Blocking ``VLLMWeightSyncClient`` over the inference servers' ``/collective_rpc``. - - Implements the four synchronous control-plane calls the vendored trainer - engine makes (``sharded_rdt_base.VLLMWeightSyncClient``). Constructed by - ``RdtWeightSyncSender`` from the ``RemoteInferenceClient``'s ``server_urls`` - and ``data_parallel_size`` — the async session and event loop are NOT used. - Safe to call from any thread; the engine calls it from the worker thread - ``RdtWeightSyncSender`` runs ``send_weights`` on. - """ - - def __init__(self, server_urls: Sequence[str], data_parallel_size: int) -> None: - import requests # local: Ray (an RDT hard-dep) provides it; keep it off non-RDT paths. - - self._urls = list(server_urls) - self._dp = int(data_parallel_size) - if not self._urls: - raise ValueError("SyncRdtControlPlaneClient requires at least one server_url.") - - self._session = requests.Session() - # Per-call fresh connections — see module docstring (avoids the stale - # keep-alive / ECONNRESET race across the long idle gap between syncs). - self._session.headers["Connection"] = "close" - # One worker per server so a fan-out call issues every POST concurrently. - self._pool = ThreadPoolExecutor(max_workers=len(self._urls), thread_name_prefix="rdt-ctrl") - - # ---- VLLMWeightSyncClient protocol ---- - - def init_weight_transfer_engine(self, init_info: Dict[str, Any]) -> None: - self._fanout(build_rdt_init_payloads(init_info, self._urls, self._dp)) - - def start_weight_update(self) -> None: - self._fanout_uniform(RDT_START_METHOD, {"is_checkpoint_format": True}) - - def update_weights(self, update_info: Dict[str, Any]) -> None: - self._fanout_uniform(RDT_UPDATE_METHOD, {"update_info": update_info}) - - def finish_weight_update(self) -> None: - self._fanout_uniform(RDT_FINISH_METHOD, None) - - def close(self) -> None: - """Release the HTTP session + fan-out pool. Idempotent.""" - self._pool.shutdown(wait=True) - self._session.close() - - # ---- internals ---- - - def _fanout_uniform(self, method: str, kwargs: Optional[Dict[str, Any]]) -> None: - payload: Dict[str, Any] = {"method": method} - if kwargs is not None: - payload["kwargs"] = kwargs - self._fanout([(url, payload) for url in self._urls]) - - def _fanout(self, url_payloads: List[Tuple[str, Dict[str, Any]]]) -> None: - """POST to every server concurrently; raise the first failure after all - return. Concurrency is required for correctness, not speed (see module - docstring): serial ``update_weights`` deadlocks the producer's - ref-counted group free. We drain ALL futures before raising so a failure - on one server never leaves POSTs in flight against the others.""" - futures = [self._pool.submit(self._post, url, payload) for url, payload in url_payloads] - first_exc: Optional[BaseException] = None - for fut in futures: - try: - fut.result() - except Exception as exc: # noqa: BLE001 - first_exc = first_exc or exc - if first_exc is not None: - raise first_exc - - def _post(self, url: str, payload: Dict[str, Any]) -> None: - # No timeout (bake + NIXL pull are long) and no retry (retrying a - # half-done stateful call would be wrong; Connection: close already - # removes the stale-keepalive race the async path guarded against). - resp = self._session.post(f"{url}{COLLECTIVE_RPC_ENDPOINT}", json=payload, timeout=None) - if resp.status_code >= 400: - raise RuntimeError(_error_message(url, payload, resp)) - - -def _error_message(url: str, payload: Dict[str, Any], resp: Any) -> str: - """Mirror ``RemoteInferenceClient.raise_for_status``: surface the response - body's error detail (``{"error": {"message": ...}}``) rather than the bare - HTTP reason phrase, which is usually unhelpful.""" - method = payload.get("method") - detail = resp.reason - try: - body = resp.json() - if isinstance(body, dict): - err = body.get("error") - if isinstance(err, dict): - detail = err.get("message", detail) - except Exception: # noqa: BLE001 - detail = (resp.text or resp.reason)[:1000] - return f"RDT control-plane call {method!r} to {url} failed [{resp.status_code}]: {detail}" diff --git a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_send.py b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_send.py index 91806f0766..5f2fe5c7b8 100644 --- a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_send.py +++ b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_send.py @@ -1,43 +1,36 @@ -"""Trainer-side driver for the sharded-RDT (NIXL pull) weight-sync backend. - -RDT matches vLLM's trainer-send model — a ``WeightSource`` + a -``TrainerWeightTransferEngine`` + a ``VLLMWeightSyncClient``, where the engine owns -the round trip and the inference workers pull. ``ShardedRdtTransferStrategy`` -(``sharded_rdt_strategy.py``) presents :class:`RdtWeightSyncSender` through SkyRL's -``WeightTransferSender`` interface, overriding ``send`` so no chunk stream or -metadata is ever materialized. - -Three ``WeightSource`` flavors, selected by extractor flavor in -``make_weight_source``: ``_FsdpWeightSource`` (all-gather via ``full_tensor()``), -``MegatronStackedWeightSource`` (PP-local + EP-local, stack-granularity gathers) and -``MegatronWeightSource`` (the whole-model Megatron-Bridge ``export_hf_weights`` -fallback). +"""Sharded-RDT (NIXL pull) weight sources and trainer init info. + +Sources here are :class:`GroupedWeightSource`s, not the plain +``weight_sync/sources.py`` ones: a pull backend needs per-rank ownership +(``held_names()``) and a group index its free barrier counts. See +``sharded_rdt_base``. + +Three flavors, chosen in :func:`make_megatron_weight_source` / +:func:`make_fsdp_weight_source`: :class:`RdtFsdpWeightSource` (all-gather via +``full_tensor()``), :class:`MegatronStackedWeightSource` (PP-local + EP-local, +stack-granularity gathers) and :class:`RdtMegatronWeightSource` (the whole-model +Megatron-Bridge ``export_hf_weights`` fallback). """ -import asyncio import contextlib import logging import os from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Iterator, List, Optional +from typing import Any, Iterator, List, Optional import torch from loguru import logger as _loguru -from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_control_plane import ( - SyncRdtControlPlaneClient, -) from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_base import ( + GroupedWeightSource, ParamMeta, - WeightSource, layerwise_groups, + materialize_full_tensor, +) +from skyrl.backends.skyrl_train.weight_sync.sources import ( + FsdpWeightSource, + MegatronWeightSource, ) -from skyrl.train.utils.utils import str_to_torch_dtype - -if TYPE_CHECKING: - from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( - RemoteInferenceClient, - ) logger = logging.getLogger(__name__) @@ -177,112 +170,60 @@ def _local_obj(self, obj, cache_key=None): ) = saved -class _FsdpWeightSource(WeightSource): - """``WeightSource`` over the FSDP policy model for the sidecar trainer engine. +class RdtFsdpWeightSource(FsdpWeightSource, GroupedWeightSource): + """The shared FSDP ``WeightSource``, re-ordered group-major. - Yields ``(name, full tensor)`` pairs in group-contiguous (pre / per-layer / - post) order, cast to the inference dtype, using the worker's - ``FSDPWeightExtractor`` so the names (incl. ``weight_prefix``) and the - all-gather match exactly what the consumer engine baked its plan over. - - ``metadata()`` reads state_dict shapes only (no gather); iteration all-gathers - each parameter (``full_tensor()``) and is therefore a collective every trainer - rank must run in lockstep — which the vendored engine's ``send_weights`` - guarantees (all ranks iterate the source). + Same state_dict and ``full_tensor()`` gather as + ``weight_sync.sources.FsdpWeightSource``; only the order differs. RDT needs + each ``model.layers..*`` block contiguous so ``layerwise_groups`` + partitions ``metadata()`` exactly -- that partition is the group index the + consumers' pull plans and the producer's free barrier are keyed on. The push + backends transfer any permutation identically, so they use the plain source. """ - def __init__(self, weight_extractor: Any, dtype: torch.dtype) -> None: - self._extractor = weight_extractor - self._dtype = dtype - meta = weight_extractor.get_weight_metadata(dtype) - names = list(meta["names"]) - shapes = [list(s) for s in meta["shapes"]] - # Reorder into group-major order so layerwise_groups(names) partitions the - # list exactly and each model.layers..* block is contiguous (the order - # the trainer engine validates + the gather loop drives). + def __init__(self, model: Any, dtype: torch.dtype, weight_prefix: str = "") -> None: + super().__init__(model, dtype, weight_prefix) + sd = model.state_dict() + prefix = weight_prefix or "" + names = [f"{prefix}{key}" for key in sd.keys()] + shapes = {f"{prefix}{key}": tuple(param.shape) for key, param in sd.items()} idx = {n: i for i, n in enumerate(names)} order = [idx[n] for g in layerwise_groups(names) for n in g] self._names = [names[i] for i in order] - self._shapes = [shapes[i] for i in order] + self._shapes = [shapes[names[i]] for i in order] def metadata(self) -> List[ParamMeta]: return [ParamMeta(name, self._dtype, tuple(shape)) for name, shape in zip(self._names, self._shapes)] def __iter__(self) -> Iterator[tuple]: - # A worker thread does not inherit the main thread's current CUDA device; - # set it so the gather collectives + casts land on this rank's GPU. - if torch.cuda.is_available(): - torch.cuda.set_device(torch.cuda.current_device()) - prefix = getattr(self._extractor, "weight_prefix", "") or "" - # Use the SAME model handle the metadata came from: get_weight_metadata - # reads ``weight_extractor.model.state_dict()`` (the inner HF model), - # whereas the worker's ``self.model`` may be a wrapper prefixed otherwise. - sd = self._extractor.model.state_dict() + # The caller selects this rank's CUDA device before iterating; a worker + # thread does not inherit it (see Worker._weight_sync_thread). + device = torch.cuda.current_device() if torch.cuda.is_available() else None + prefix = self.weight_prefix + # The same handle metadata() was built from, so names and gather match + # what the consumer engine baked its plan over. + sd = self.model.state_dict() for name in self._names: raw = name[len(prefix) :] if prefix and name.startswith(prefix) else name - full = self._extractor._gather_tensor(sd[raw]).to(self._dtype).detach().contiguous() - yield name, full - - -class MegatronWeightSource(WeightSource): - """``WeightSource`` over the Megatron policy model for the sidecar trainer engine. - - Wraps ``MegatronWeightExtractor`` and streams ``(HF name, full tensor)`` pairs - via the Megatron-Bridge (``bridge.export_hf_weights``), which gathers each - parameter across TP / PP / EP internally and — per its contract — "All ranks - get full tensors". That is exactly the FSDP ``full_tensor()`` semantics the RDT - producer model needs: every trainer rank materializes the whole model (one - parameter at a time, so peak memory is bounded), so each producer can serve - its bound consumer the complete model. - - We call the bridge **directly** with ``conversion_tasks=None`` rather than the - extractor's ``extract_weights`` because the extractor is built with - ``enable_bucketing=True``, and bucketing hoists the grouped MoE-expert tasks - into dedicated leading buckets — which would break the group-major - (pre / per-decoder-layer / post) contiguity the trainer engine's - ``trainer_init`` validates. The non-bucketed export yields HF-canonical order, - which ``layerwise_groups`` partitions cleanly. ``metadata()`` and iteration run - the same export, so their order always agrees. - - Megatron only: this reads ``weight_extractor.bridge`` / ``.actor_module``. - ``rdt_send`` selects it (vs ``_FsdpWeightSource``) by the presence of - ``bridge`` on the extractor. - """ + param = sd[raw] + if device is not None: + param = param.to(device, non_blocking=True) + yield name, materialize_full_tensor(param).to(self._dtype).detach().contiguous() - def __init__(self, weight_extractor: Any, dtype: torch.dtype) -> None: - self._bridge = weight_extractor.bridge - self._module = weight_extractor.actor_module - self._dtype = dtype - self._meta: Optional[List[ParamMeta]] = None - def _export(self) -> Iterator[tuple]: - # conversion_tasks=None -> full model in HF-canonical (group-contiguous) - # order; the bridge gathers TP/PP/EP and yields full tensors on every rank. - return self._bridge.export_hf_weights(self._module, show_progress=False, conversion_tasks=None) +class RdtMegatronWeightSource(MegatronWeightSource, GroupedWeightSource): + """The shared Megatron ``WeightSource``, with RDT's group channels. - def metadata(self) -> List[ParamMeta]: - # One collective dry-run export to learn names/shapes; tensors are - # discarded. Runs once (trainer_init caches the result), so the extra - # gather is a one-time init cost, not per-sync. - if self._meta is None: - meta: List[ParamMeta] = [] - for name, tensor in self._export(): - meta.append(ParamMeta(name, self._dtype, tuple(tensor.shape))) - del tensor - self._meta = meta - return self._meta + Identical stream to ``weight_sync.sources.MegatronWeightSource``: the + whole-model export, in HF-canonical (already group-contiguous) order, with + every rank getting full tensors. For a pull backend that whole-model + residency is the point -- each producer must be able to serve its bound + consumer the complete model -- so the inherited "hold everything" defaults + are correct and this subclass adds nothing else. - def __iter__(self) -> Iterator[tuple]: - # A worker thread does not inherit the main thread's current CUDA device; - # set it so the bridge's gather collectives + casts land on this rank's GPU. - if torch.cuda.is_available(): - torch.cuda.set_device(torch.cuda.current_device()) - device = torch.cuda.current_device() - else: - device = None - for name, tensor in self._export(): - full = tensor.to(device=device, dtype=self._dtype).detach().contiguous() - yield name, full + :class:`MegatronStackedWeightSource` is the narrower alternative; + ``make_megatron_weight_source`` falls back here when it cannot serve a layout. + """ @dataclass(frozen=True) @@ -314,7 +255,7 @@ def E(self) -> int: return self.n_local * self.ep_size -class MegatronStackedWeightSource(WeightSource): +class MegatronStackedWeightSource(GroupedWeightSource): """Megatron ``WeightSource`` with STACKED expert gathers (per-tensor-overhead fix). The plain ``MegatronWeightSource`` streams every HF tensor through the @@ -365,9 +306,9 @@ class MegatronStackedWeightSource(WeightSource): _EXPERT_PRED = ".experts.linear_fc" # model_bridge.py uses the same predicate - def __init__(self, weight_extractor: Any, dtype: torch.dtype) -> None: - self._bridge = weight_extractor.bridge - self._module = weight_extractor.actor_module + def __init__(self, bridge: Any, module: Any, dtype: torch.dtype) -> None: + self._bridge = bridge + self._module = module self._dtype = dtype self._meta: Optional[List[ParamMeta]] = None self._verified = False @@ -391,7 +332,7 @@ def __init__(self, weight_extractor: Any, dtype: torch.dtype) -> None: # exports only its own parameters; EP-local (ep>1) — # only this coordinate's experts are materialized, foreign experts # yield None. Both are declared through held_names(). The escape hatch to naive whole-model - # extraction is the plain MegatronWeightSource + # extraction is the plain RdtMegatronWeightSource # (SKYRL_RDT_STACKED_EXPERTS=0), which make_weight_source also # delegates to automatically when a gather group spans pipeline stages # (tied embeddings, MTP) — see held_names / make_weight_source. @@ -413,7 +354,7 @@ def __init__(self, weight_extractor: Any, dtype: torch.dtype) -> None: self._ep_local = self._ep_size > 1 # Set by held_names when a gather group spans pipeline stages: this # source cannot serve that layout, and make_weight_source delegates to - # the plain MegatronWeightSource instead. Iteration refuses to run. + # the plain RdtMegatronWeightSource instead. Iteration refuses to run. self._demoted = False # Per-layer (F, H) recorded as layers are walked, so metadata() can # synthesize the shapes of foreign experts (their tensors are None). @@ -801,7 +742,7 @@ def _iter_groups_impl(self, collect_meta: bool) -> Iterator[tuple]: raise RuntimeError( "[stacked-source] this source was demoted (a gather group spans " "pipeline stages) and cannot iterate; make_weight_source " - "delegates to the plain MegatronWeightSource automatically." + "delegates to the plain RdtMegatronWeightSource automatically." ) if torch.cuda.is_available(): torch.cuda.set_device(torch.cuda.current_device()) @@ -1055,7 +996,7 @@ def _owned_group_indices(self) -> Optional[List[int]]: logger.warning( "[stacked-source] gather groups %s are produced by more than one pipeline " "stage (tied embeddings / MTP), which this source cannot serve per-stage. " - "Demoting: make_weight_source delegates to the plain MegatronWeightSource " + "Demoting: make_weight_source delegates to the plain RdtMegatronWeightSource " "(naive whole-model extraction, correct but slower).", shared[:4], ) @@ -1195,260 +1136,202 @@ def _verify_against_bridge(self) -> None: logger.info("[stacked-verify] layer %d: expert tensors match bridge export", layer) -def make_weight_source(weight_extractor: Any, dtype: torch.dtype) -> WeightSource: - """Pick the RDT ``WeightSource`` for the trainer's weight extractor. +def make_fsdp_weight_source(model: Any, dtype: torch.dtype, weight_prefix: str = "") -> GroupedWeightSource: + """The RDT ``WeightSource`` for an FSDP policy model.""" + return RdtFsdpWeightSource(model, dtype, weight_prefix) - The Megatron extractor exposes a Megatron-Bridge (``bridge`` / - ``export_hf_weights``); the FSDP extractor exposes an all-gatherable - ``state_dict`` (``model`` / ``_gather_tensor``). Selection is by the presence - of ``bridge`` so neither backend module has to be imported here. - For Megatron, prefer ``MegatronStackedWeightSource`` (stack-granularity - expert gathers; ~20x fewer collectives on per-expert MoE archs) unless the - arch uses grouped-export mappings (fused HF expert names — different - contract) or ``SKYRL_RDT_STACKED_EXPERTS=0``. +def make_megatron_weight_source(bridge: Any, module: Any, dtype: torch.dtype) -> GroupedWeightSource: + """Pick the RDT ``WeightSource`` for a Megatron policy model. + + Prefer :class:`MegatronStackedWeightSource` (stack-granularity expert + gathers; ~20x fewer collectives on per-expert MoE archs, and PP-local export + at pp>1) unless the arch uses grouped-export mappings (fused HF expert names + — a different contract) or ``SKYRL_RDT_STACKED_EXPERTS=0``. Everything else + falls back to the whole-model :class:`RdtMegatronWeightSource`. """ - if hasattr(weight_extractor, "bridge"): - if os.environ.get("SKYRL_RDT_STACKED_EXPERTS", "1") != "0": + if os.environ.get("SKYRL_RDT_STACKED_EXPERTS", "1") != "0": + try: + tasks = bridge.get_conversion_tasks(module) + expert_tasks = [t for t in tasks if MegatronStackedWeightSource._EXPERT_PRED in t.param_name] + grouped = any(getattr(t.mapping, "is_grouped_export", False) for t in expert_tasks) + wrapped = any(".to_wrap." in t.global_param_name for t in expert_tasks) + # ETP>1 means no rank holds a WHOLE expert (shapes would be read + # off shards; EP-local stamping has no truthful answer), so probe + # for EVERY config — Megatron defaults etp to tp when unset, not + # to 1. Unprobeable (no mpu) means no ETP — except under LoRA, + # where the stack merge is at stake and the conservative answer + # is to fall back (the pre-existing behaviour). try: - tasks = weight_extractor.bridge.get_conversion_tasks(weight_extractor.actor_module) - expert_tasks = [t for t in tasks if MegatronStackedWeightSource._EXPERT_PRED in t.param_name] - grouped = any(getattr(t.mapping, "is_grouped_export", False) for t in expert_tasks) - wrapped = any(".to_wrap." in t.global_param_name for t in expert_tasks) - # ETP>1 means no rank holds a WHOLE expert (shapes would be read - # off shards; EP-local stamping has no truthful answer), so probe - # for EVERY config — Megatron defaults etp to tp when unset, not - # to 1. Unprobeable (no mpu) means no ETP — except under LoRA, - # where the stack merge is at stake and the conservative answer - # is to fall back (the pre-existing behaviour). + from megatron.core import parallel_state + + etp = torch.distributed.get_world_size(parallel_state.get_expert_tensor_parallel_group()) + except Exception: # noqa: BLE001 + etp = 1 if not wrapped else 0 + etp_ok = etp == 1 + # A DENSE model has no expert tasks, but the source's two grains + # engage INDEPENDENTLY: PP-local ("this stage exports only its own + # layers") needs no experts at all. At pp>1 that is the difference + # between fitting and not -- measured on Qwen3-32B, where the plain + # source's whole-model export OOMed at both tp4/pp2 (70.56 GiB) and + # tp8/pp2 (73.06 GiB) of 79.18. Halving per-rank params did not help, + # which is what identifies the export as model-sized rather than + # shard-sized. + # + # At pp==1 the stacked source degenerates to the plain + # filtered==full export, so there is nothing to win and the simpler + # path is kept. The `_demoted` check below is unchanged and still + # catches layouts a stage cannot serve alone (tied embeddings, MTP). + pp_local_gain = False + if not expert_tasks: try: from megatron.core import parallel_state - etp = torch.distributed.get_world_size(parallel_state.get_expert_tensor_parallel_group()) + pp_local_gain = parallel_state.get_pipeline_model_parallel_world_size() > 1 except Exception: # noqa: BLE001 - etp = 1 if not wrapped else 0 - etp_ok = etp == 1 - # A DENSE model has no expert tasks, but the source's two grains - # engage INDEPENDENTLY: PP-local ("this stage exports only its own - # layers") needs no experts at all. At pp>1 that is the difference - # between fitting and not -- measured on Qwen3-32B, where the plain - # source's whole-model export OOMed at both tp4/pp2 (70.56 GiB) and - # tp8/pp2 (73.06 GiB) of 79.18. Halving per-rank params did not help, - # which is what identifies the export as model-sized rather than - # shard-sized. - # - # At pp==1 the stacked source degenerates to the plain - # filtered==full export, so there is nothing to win and the simpler - # path is kept. The `_demoted` check below is unchanged and still - # catches layouts a stage cannot serve alone (tied embeddings, MTP). - pp_local_gain = False - if not expert_tasks: - try: - from megatron.core import parallel_state - - pp_local_gain = parallel_state.get_pipeline_model_parallel_world_size() > 1 - except Exception: # noqa: BLE001 - pp_local_gain = False - if (expert_tasks or pp_local_gain) and not grouped and etp_ok: - if wrapped: - logger.info("[rdt] stacked expert source with LoRA stack-merge (etp=1)") - src = MegatronStackedWeightSource(weight_extractor, dtype) - # pp>1: run the shared-group discovery NOW (one metadata - # walk, rank-identical on every rank) so a tied-embeddings/ - # MTP layout delegates to the plain source here instead of - # failing mid-init. At pp==1 this returns without work. - src.held_names() - if src._demoted: - logger.warning( - "[rdt] a gather group spans pipeline stages; delegating to the " - "plain MegatronWeightSource (naive whole-model extraction)" - ) - return MegatronWeightSource(weight_extractor, dtype) - return src - if not etp_ok and expert_tasks: - logger.info( - "[rdt] expert_tensor_parallel_size != 1: no rank holds a whole " - "expert; using plain MegatronWeightSource (the bridge gathers " - "ETP shards internally — correct, slower)" - ) - elif grouped: - logger.info("[rdt] grouped-export arch; using plain MegatronWeightSource") - else: - logger.info( - "[rdt] dense model at pp==1: PP-local serving would narrow nothing " - "and there are no experts to stack; using plain MegatronWeightSource" + pp_local_gain = False + if (expert_tasks or pp_local_gain) and not grouped and etp_ok: + if wrapped: + logger.info("[rdt] stacked expert source with LoRA stack-merge (etp=1)") + src = MegatronStackedWeightSource(bridge, module, dtype) + # pp>1: run the shared-group discovery NOW (one metadata + # walk, rank-identical on every rank) so a tied-embeddings/ + # MTP layout delegates to the plain source here instead of + # failing mid-init. At pp==1 this returns without work. + src.held_names() + if src._demoted: + logger.warning( + "[rdt] a gather group spans pipeline stages; delegating to the " + "plain RdtMegatronWeightSource (naive whole-model extraction)" ) - except Exception: # noqa: BLE001 - logger.warning("[rdt] stacked-source probe failed; using plain MegatronWeightSource", exc_info=True) - return MegatronWeightSource(weight_extractor, dtype) - return _FsdpWeightSource(weight_extractor, dtype) - - -class RdtWeightSyncSender: - """Drives sharded-RDT weight sync through the vendored vLLM trainer-send engine. - ``ShardedRdtWeightTransferSender`` wraps this to satisfy SkyRL's - ``WeightTransferSender`` interface. - - Built once per worker in ``init_weight_sync_state``; ``send()`` is called from - ``broadcast_to_inference_engines`` each RL step. The heavy ``trainer_init`` - (spawn the per-rank ``_RDTProducerServer`` sidecar + rank-0 bake) and every - ``send_weights`` block on Ray + ``torch.distributed`` collectives, so both run - off the worker's event loop via ``asyncio.to_thread``. The engine drives the - inference side through the synchronous ``SyncRdtControlPlaneClient`` (blocking - HTTP to the servers' ``/collective_rpc``), so the whole engine runs - sync-to-sync in that worker thread with no event-loop involvement — see - ``rdt_control_plane`` for why sidestepping the async ``RemoteInferenceClient`` - here is safe. + return RdtMegatronWeightSource(bridge, module, dtype) + return src + if not etp_ok and expert_tasks: + logger.info( + "[rdt] expert_tensor_parallel_size != 1: no rank holds a whole " + "expert; using plain RdtMegatronWeightSource (the bridge gathers " + "ETP shards internally — correct, slower)" + ) + elif grouped: + logger.info("[rdt] grouped-export arch; using plain RdtMegatronWeightSource") + else: + logger.info( + "[rdt] dense model at pp==1: PP-local serving would narrow nothing " + "and there are no experts to stack; using plain RdtMegatronWeightSource" + ) + except Exception: # noqa: BLE001 + logger.warning("[rdt] stacked-source probe failed; using plain RdtMegatronWeightSource", exc_info=True) + return RdtMegatronWeightSource(bridge, module, dtype) + + +def build_rdt_trainer_init_info( + rank: int, + inference_world_size: int, + server_urls: List[str], + data_parallel_size: int, +): + """Build ``ShardedRDTTrainerInitInfo`` for this rank. + + Every knob is resolved on the **trainer**: the producer sidecar is a Ray + actor that inherits the raylet's environment, so a launch-time ``SKYRL_*`` + override never reaches it. + + Args: + rank: this trainer process's rank. Rank 0 is the sender. + inference_world_size: total inference workers across the fleet — the + consumer count the producers' free barrier counts against. + server_urls: every inference server, ordered + ``[engine0_dp0, ..., engine1_dp0, ...]``. Only its length and the DP + size are used, to derive the deployment count. + data_parallel_size: DP replicas per deployment. """ + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_trainer import ( + ShardedRDTTrainerInitInfo, + ) - def __init__( - self, - inference_client: "RemoteInferenceClient", - model_dtype: str, - inference_world_size: int, - trainer_actor_namespace: Optional[str], - ) -> None: - if not inference_world_size: - raise ValueError( - f"sharded_rdt requires the inference world size (consumer count); got {inference_world_size!r}." - ) - self._model_dtype = model_dtype - self._world_size = int(inference_world_size) - self._namespace = trainer_actor_namespace - # Snapshot only what the sync control plane needs (server URLs + DP size); - # the async client, its aiohttp session, and the event loop are NOT used. - self._server_urls = list(inference_client.server_urls) - self._data_parallel_size = int(inference_client.data_parallel_size) - # Deployment count, derived the same way the control plane derives each - # server's `replica_rank` (see `build_rdt_init_payloads`): the DP servers of - # one deployment share a parallel config, so they share one ordinal. - # [RDT-SHARE-SLOTS] reads it to size a deployment's consumer-id block. - self._num_replicas = max(1, len(self._server_urls) // max(1, self._data_parallel_size)) - self._engine: Any = None - self._control_plane: Optional[SyncRdtControlPlaneClient] = None - - def initialize(self, weight_extractor: Any) -> None: - """Eagerly rendezvous (spawn sidecar servers + rank-0 bake) at - ``init_weight_sync_state`` time, BEFORE any weight-gather collective can - be in flight. Deferring this to the first ``send()`` deadlocks: rank 0 - blocks in the inference-side init RPC while the other ranks spin in the - gather NCCL collectives waiting for it, and the producer servers (which - share those ranks' GPUs) then can't finish NIXL agent creation — - libfabric's fi_getinfo CUDA probe blocks behind the spinning kernels. - Called with every rank inside ``init_weight_sync_state`` (followed by a - barrier), ranks that finish early sit idle, so the window is empty. - Blocking is fine here: ``init_weight_sync_state`` already blocks on a - ``torch.distributed.barrier()``.""" - if weight_extractor is None: - raise RuntimeError( - "sharded_rdt weight sync requires the worker's weight_extractor " "(built in init_weight_sync_state)." - ) - if self._engine is None: - self._engine = self._trainer_init_blocking(weight_extractor) - - async def send(self, weight_extractor: Any) -> None: - """Sync weights once; every rank must call it (the gather is a - collective). ``initialize()`` should already have run; the lazy fallback - here only covers callers that skipped it (and reintroduces the - first-send deadlock risk described there — do not rely on it).""" - if weight_extractor is None: - raise RuntimeError( - "sharded_rdt weight sync requires the worker's weight_extractor " "(built in init_weight_sync_state)." - ) - if self._engine is None: - self._engine = await asyncio.to_thread(self._trainer_init_blocking, weight_extractor) - await asyncio.to_thread(self._engine.send_weights) - - def _trainer_init_blocking(self, weight_extractor: Any) -> Any: - """Build the WeightSource + control-plane client + trainer init info and - rendezvous. Runs in a worker thread (blocks on the Ray spawn, an - all-gather collective, and — on rank 0 — the inference-side bake).""" - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_trainer import ( - ShardedRDTTrainerInitInfo, - ShardedRDTTrainerWeightTransferEngine, + if not inference_world_size: + raise ValueError( + f"sharded_rdt requires the inference world size (consumer count); got {inference_world_size!r}." ) - # Constructed on every rank (the engine holds a client on all ranks), but - # only the sender rank actually issues control-plane calls. - self._control_plane = SyncRdtControlPlaneClient(self._server_urls, self._data_parallel_size) - dtype = str_to_torch_dtype(self._model_dtype) - source = make_weight_source(weight_extractor, dtype) - # Positive configuration tripwire. make_weight_source's own logs use the - # vLLM logger, which does NOT forward to the driver log from a Megatron - # rank actor (only loguru does), and the STACKED_EXPERTS=0 short-circuit - # logs nothing at all -- so an ablation could silently run the wrong - # source. Log the CHOSEN CLASS (not the env var) plus both knobs. - _loguru.info( - "[rdt-config] source={} lookahead_env={} stacked_experts_env={}", - type(source).__name__, - os.environ.get("SKYRL_RDT_LOOKAHEAD", ""), - os.environ.get("SKYRL_RDT_STACKED_EXPERTS", ""), - ) - rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0 + def _knob(env: str, default): + v = os.environ.get(env) + return v if v is not None else default + + # Deployment count, derived the same way the control plane derives each + # server's ``replica_rank`` (see ``control_plane.rdt_init_payloads``): the DP + # servers of one deployment share a parallel config, so they share one + # ordinal. [RDT-SHARE-SLOTS] reads it to size a deployment's consumer-id block. + dp = max(1, int(data_parallel_size)) + num_replicas = max(1, len(server_urls) // dp) + + # [RDT-SHARE-SLOTS] Consumers per deployment, which is what groups the + # workers that can share one serve slot on each producer. 0 turns sharing + # off, and one deployment makes it a no-op anyway (the width equals the + # consumer count, so every group is a singleton). + share_slots = os.environ.get("SKYRL_RDT_SHARE_SLOTS", "1") not in ("0", "false", "False") + workers_per_replica = int(inference_world_size) // num_replicas if share_slots else 0 + _loguru.info( + "[rdt-config] num_consumers={} num_replicas={} workers_per_replica={} (slot sharing {})", + inference_world_size, + num_replicas, + workers_per_replica, + "on" if share_slots and num_replicas > 1 else "off", + ) - # Pipeline-depth knobs are env-driven (SKYRL_* env vars are forwarded - # into every Ray worker by prepare_runtime_environment). The env - # override exists because the sync is a credit-limited latency - # pipeline: its steady-state period is (publish->serve->pull->free - # RTT) / credits, so deepening K/lookahead hides per-link latency - # without touching any link. - def _knob(env: str, default): - v = os.environ.get(env) - return v if v is not None else default - - # [RDT-SHARE-SLOTS] Consumers per deployment, which is what groups the - # workers that can share one serve slot on each producer. Resolved here - # because the sidecar is a Ray actor that never sees this process's env; - # 0 turns sharing off, and one deployment makes it a no-op anyway (the - # width equals the consumer count, so every group is a singleton). - share_slots = os.environ.get("SKYRL_RDT_SHARE_SLOTS", "1") not in ("0", "false", "False") - workers_per_replica = self._world_size // max(1, self._num_replicas) if share_slots else 0 - _loguru.info( - "[rdt-config] num_consumers={} num_replicas={} workers_per_replica={} (slot sharing {})", - self._world_size, - self._num_replicas, - workers_per_replica, - "on" if share_slots and self._num_replicas > 1 else "off", - ) + return ShardedRDTTrainerInitInfo( + rank=rank, + num_consumers=int(inference_world_size), + workers_per_replica=workers_per_replica, + trainer_actor_namespace=ray_namespace(), + num_rdt_buffers=int(_knob("SKYRL_RDT_NUM_BUFFERS", _DEFAULT_NUM_RDT_BUFFERS)), + buffer_presize_gb=float(_knob("SKYRL_RDT_BUFFER_PRESIZE_GB", _DEFAULT_BUFFER_PRESIZE_GB)), + gather_lookahead=int(_knob("SKYRL_RDT_LOOKAHEAD", _DEFAULT_GATHER_LOOKAHEAD)), + stall_timeout_s=float(_knob("SKYRL_RDT_STALL_TIMEOUT_S", _DEFAULT_STALL_TIMEOUT_S)), + ) - init_info = ShardedRDTTrainerInitInfo( - rank=rank, - num_consumers=self._world_size, - workers_per_replica=workers_per_replica, - trainer_actor_namespace=self._namespace, - num_rdt_buffers=int(_knob("SKYRL_RDT_NUM_BUFFERS", _DEFAULT_NUM_RDT_BUFFERS)), - buffer_presize_gb=float(_knob("SKYRL_RDT_BUFFER_PRESIZE_GB", _DEFAULT_BUFFER_PRESIZE_GB)), - gather_lookahead=int(_knob("SKYRL_RDT_LOOKAHEAD", _DEFAULT_GATHER_LOOKAHEAD)), - # Resolved here, not in the sidecar: the sidecar is a Ray actor that - # inherits the raylet's environment, so a launch-time SKYRL_* override - # never reaches it. - stall_timeout_s=float(_knob("SKYRL_RDT_STALL_TIMEOUT_S", _DEFAULT_STALL_TIMEOUT_S)), - ) - engine = ShardedRDTTrainerWeightTransferEngine.trainer_init( - init_info, - client=self._control_plane, - source=source, - ) - # The producer sidecar already freezes its object graph "so gen-2 GC - # never stops the world mid-serve". The trainer needs it more: it - # rebuilds the whole conversion-task graph every sync on top of a 235B - # heap, and a gen-2 pass costs up to 1.2s on a rank -- which the sync's - # PP all_gather_object then propagates to that rank's partner, so the - # slowest pair sets the sync. Freezing leaves gen-0/1 untouched and makes - # gen-2 ~40x cheaper by skipping the static graph. - if os.environ.get("SKYRL_RDT_GC_FREEZE", "1") not in ("0", "false", "False"): - import gc - - gc.collect() - gc.freeze() - return engine - - def teardown(self) -> None: - engine = self._engine - control_plane = self._control_plane - self._engine = None - self._control_plane = None - if engine is not None: - engine.shutdown() - if control_plane is not None: - control_plane.close() + +def ray_namespace() -> Optional[str]: + """This process's Ray namespace, or None outside a Ray runtime. + + The producer sidecars are named actors, so the consumers need the namespace + they were created in to resolve them. + """ + try: + import ray + + return ray.get_runtime_context().namespace or None + except Exception: # noqa: BLE001 - no Ray runtime: the caller falls back to the default namespace + return None + + +def freeze_trainer_heap() -> None: + """``gc.freeze()`` the trainer's static object graph after the rendezvous. + + The trainer rebuilds the whole conversion-task graph every sync on top of a + 235B heap, and a gen-2 pass costs up to 1.2s on a rank -- which the sync's PP + ``all_gather_object`` propagates to that rank's partner, so the slowest pair + sets the sync. Freezing leaves gen-0/1 untouched and makes gen-2 ~40x cheaper. + """ + if os.environ.get("SKYRL_RDT_GC_FREEZE", "1") in ("0", "false", "False"): + return + import gc + + gc.collect() + gc.freeze() + + +def log_source_choice(source: Any) -> None: + """Log the chosen source class and both knobs. + + ``make_megatron_weight_source``'s own logs use the vLLM logger, which does + not forward to the driver log from a Megatron rank actor (only loguru does), + and the ``STACKED_EXPERTS=0`` short-circuit logs nothing -- so without this + an ablation could silently run the wrong source. + """ + _loguru.info( + "[rdt-config] source={} lookahead_env={} stacked_experts_env={}", + type(source).__name__, + os.environ.get("SKYRL_RDT_LOOKAHEAD", ""), + os.environ.get("SKYRL_RDT_STACKED_EXPERTS", ""), + ) diff --git a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_base.py b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_base.py index 9511c965b9..78ef7c6f2f 100644 --- a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_base.py +++ b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_base.py @@ -1,58 +1,53 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Trainer-side weight-transfer base ABCs, vendored from the vLLM RDT fork. - -These classes live in ``vllm/distributed/weight_transfer/base.py`` in the -``vllm-rdt-weight-sync`` fork. The pinned ``vllm==0.26.0`` wheel DOES ship -trainer-side ABCs at that path (``WeightSource``, ``TrainerWeightTransferEngine``, -``TrainerInitInfo``, ``VLLMWeightSyncClient``, ``ParamMeta``, -``materialize_full_tensor``) — but a NARROWER version of them: it has no -``layerwise_groups``, and its ``WeightSource`` declares only ``metadata`` / -``__iter__``, without the ownership and group extensions -(``held_names`` / ``groups`` / ``iter_groups``) the sharded-RDT engine is built on. -So the fork's versions are copied here VERBATIM rather than imported. - -REMOVAL: once SkyRL's pinned vLLM carries the *fork's* versions of these — the -group and ownership channels included — delete this module and repoint -``sharded_rdt_trainer.py``'s import back to -``vllm.distributed.weight_transfer.base``. - -Two consequences of that split, both load-bearing: - -* ``layerwise_groups``, ``WeightSource.groups()`` and ``WeightSource.iter_groups()`` - live HERE (as in the fork's ``base.py``), because they define what a group index - means for a ``WeightSource``. ``sharded_rdt_common`` re-exports - ``layerwise_groups`` for callers that already import it from there. -* ``WeightTransferEngine.defers_processing`` / ``drain_pending()`` are declared on - the fork's ABC, but the *worker*-side ABC comes from the pinned wheel (verified - on 0.26.0), which has neither. So ``sharded_rdt_engine`` declares both on the concrete - engine class, and ``inference_servers/layerwise_reload.py`` probes them with - ``getattr`` — which works against either ABC. +"""What sharded RDT needs on top of vLLM's trainer-side weight-transfer ABCs. + +Those ABCs come from the wheel. This module holds the two extra ``WeightSource`` +channels a *pull* backend needs and vLLM has no concept of, plus +``layerwise_groups``, which makes a group index mean the same thing on every rank +and every consumer: + +* ``held_names()`` -- per-rank **ownership** under pipeline / expert parallelism. + A consumer routes each parameter to a rank that holds it. Not a chunking + concern, and not optional: the default (hold everything) is correct at + pp=1/ep=1 and wrong above it. +* ``groups()`` / ``iter_groups()`` -- the coordination index the producer's free + barrier counts (``_inflight`` is keyed by group), plus batching: gathering per + group rather than per tensor turns ~37k generator resumes into ~95 on a + per-expert MoE model. + +They stay here rather than in ``weight_sync/sources.py``, which is vLLM's +contract verbatim, so a future upstream chunking change costs nothing. If vLLM +grows ownership or group channels, :class:`GroupedWeightSource` collapses into +them. """ -from abc import ABC, abstractmethod +from abc import abstractmethod from collections.abc import Collection, Iterator -from dataclasses import dataclass, field -from typing import ( - Any, - ClassVar, - Generic, - Protocol, - TypeVar, - runtime_checkable, -) import torch -from typing_extensions import Self - -TTrainerInitInfo = TypeVar("TTrainerInitInfo", bound="TrainerInitInfo") +from vllm.distributed.weight_transfer.base import ( + ParamMeta, + TrainerInitInfo, + TrainerWeightTransferEngine, + VLLMWeightSyncClient, + WeightSource, + materialize_full_tensor, +) -# A trainer supplies its parameters as a `WeightSource` (defined below): a -# re-iterable stream of materialized `(name, tensor)` pairs plus a `metadata()` -# channel. The built-in `ModuleSource` uses `materialize_full_tensor`. +__all__ = [ + "GroupedWeightSource", + "ParamMeta", + "TrainerInitInfo", + "TrainerWeightTransferEngine", + "VLLMWeightSyncClient", + "WeightSource", + "layerwise_groups", + "materialize_full_tensor", +] -def _stack_key(name: str) -> tuple[str, int] | None: +def _stack_key(name: str) -> "tuple[str, int] | None": """``(prefix, index)`` of the OUTERMOST integer segment, or None if there is none. @@ -70,10 +65,10 @@ def layerwise_groups(names: list[str]) -> list[list[str]]: """Partition flat parameter names into one group per decoder layer, keyed on the outermost index segment of each name. - This defines what a *group index* means for `WeightSource.groups` and - `WeightSource.iter_groups`: index *g* names the same group on every trainer - rank and every consumer, because it is derived from one rank's `metadata()` - order. + This defines what a *group index* means for `GroupedWeightSource.groups` and + `GroupedWeightSource.iter_groups`: index *g* names the same group on every + trainer rank and every consumer, because it is derived from one rank's + `metadata()` order. Keying on the index rather than a literal prefix needs no per-architecture naming table: ``model.layers.0.``, ``model.language_model.layers.0.``, @@ -122,49 +117,16 @@ def layerwise_groups(names: list[str]) -> list[list[str]]: return groups -def materialize_full_tensor(tensor: torch.Tensor) -> torch.Tensor: - """Return a full, locally-materialized tensor ready to send. - - FSDP shards (DTensors) expose `full_tensor()`, a collective all-gather; - regular tensors do not and are returned unchanged. Trainer engines call - this at send time so the (potentially expensive) gather happens exactly - once — reading `.shape`/`.dtype` for metadata does not trigger it. - """ - full_tensor = getattr(tensor, "full_tensor", None) - return full_tensor() if callable(full_tensor) else tensor - - -@dataclass(frozen=True) -class ParamMeta: - """Name / wire dtype / full (HF) shape for one output parameter.""" +class GroupedWeightSource(WeightSource): + """A ``WeightSource`` with the ownership and group channels sharded RDT pulls over. - name: str - dtype: torch.dtype - shape: tuple[int, ...] + Adds two channels to vLLM's ``metadata()`` + ``__iter__`` contract: - -class WeightSource(ABC): - """A re-iterable source of the trainer's weights, handed to a trainer engine. - - Two channels: - - * `metadata()` — `(name, wire dtype, full shape)` for every parameter, - *without* transferring. Cheap when shapes are known locally (FSDP - `DTensor` global shape); may be expensive on first call for backends that - must materialize to learn shapes (e.g. a Megatron-Bridge export), in which - case it should cache. - * iteration — yields fully-materialized `(name, tensor)` pairs, one at a - time. Materializing is typically a collective (FSDP `full_tensor()`, a - Megatron export), so the ranks that share a parameter must iterate it in - the same order in lockstep, or they deadlock. * `held_names()` — which parameters this rank holds, for producers that are split so each rank holds only part of the model. Defaults to all. * `iter_groups()` — the same stream batched per gather group (see `layerwise_groups`). Defaults to batching `__iter__`; override to materialize a whole group in one step. - - `iter(source)` must yield a *fresh* pass each round. Backends with custom - producer logic (Megatron export, RDT plans, MoE re-fusing) subclass this. """ @abstractmethod @@ -189,16 +151,15 @@ def held_names(self) -> "Collection[str] | None": * `metadata()` must still describe the WHOLE model on every rank. The group partition, the iteration checks and the consumers' pull plans are - all built from one rank's metadata, so a rank that reported only its own - share would leave the rest of the model silently un-transferred. The - sharded-RDT engine cross-checks this across ranks at init. + all built from one rank's metadata, so a rank reporting only its own + share would leave the rest silently un-transferred. The engine + cross-checks this across ranks at init. * Every name must be held by at least one rank, or it can never be served. The engine raises at init naming the first orphan. * Iteration must cover exactly `groups()` in metadata order, yielding a - real tensor for each held name and `None` for the rest. A group's - gather is a collective among the ranks that hold part of it, so the - name must still appear (to keep the order check aligned) while the - data is absent. + real tensor for each held name and `None` for the rest -- a group's + gather is a collective among the ranks holding part of it, so the name + must appear to keep the order check aligned even when the data does not. Returns: The held parameter names, or None to hold every one. @@ -227,11 +188,8 @@ def iter_groups(self) -> Iterator[tuple[list[str], list[torch.Tensor]]]: materialize it with a collective, so a rank that iterates out of order deadlocks its peers rather than returning wrong data. - Override when a backend can produce a whole group at once. Materializing - is usually a collective, and driving it per group instead of per tensor - turns ~37k generator resumes into ~95 on a per-expert MoE model (worth - ~0.9s per sync there). An override must yield the same batches in the same - order as this default. + Override when a backend can produce a whole group at once; it must yield + the same batches in the same order as this default. """ it = iter(self) for group in self.groups(): @@ -247,164 +205,3 @@ def iter_groups(self) -> Iterator[tuple[list[str], list[torch.Tensor]]]: names.append(name) tensors.append(tensor) yield names, tensors - - -class ModuleSource(WeightSource): - """`WeightSource` over `module.named_parameters()` — the common case. - - Handles both plain dense modules and FSDP-sharded ones with no special - casing: iteration all-gathers each `DTensor` via `full_tensor()` (a - collective) and passes regular tensors through. `metadata()` reads the - *global* `.shape` / `.dtype`, so it never triggers a gather. - """ - - def __init__(self, module: torch.nn.Module) -> None: - self._module = module - - def metadata(self) -> list[ParamMeta]: - return [ParamMeta(name, p.dtype, tuple(p.shape)) for name, p in self._module.named_parameters()] - - def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: - for name, param in self._module.named_parameters(): - yield name, materialize_full_tensor(param) - - -@dataclass -class TrainerInitInfo: - """Base trainer-side init info: which trainer rank drives the transfer. - - `rank` is this trainer process's rank, provided **explicitly** by the - caller — the engine does not read it from a global process group, which is - ambiguous once several groups (FSDP / TP / PP / EP) exist. Rank 0 is always - the sender: only it opens the endpoint and drives the inference-side RPCs, - while every rank still runs the trainer-side collectives. Backend subclasses - add their own (positional) fields; `rank` is keyword-only so that ordering - never conflicts. - - Every concrete subclass sets a class-level `backend` string (the same key it - registers under in `WeightTransferTrainerFactory`). The factory reads it to - dispatch, so callers pass only the init info/ It is a `ClassVar` - (a fixed per-backend constant), so it is not an ``__init__`` field. - """ - - backend: ClassVar[str] - - rank: int = field(kw_only=True) - - def __init_subclass__(cls, **kwargs: Any) -> None: - super().__init_subclass__(**kwargs) - if not getattr(cls, "backend", None): - raise TypeError( - f"{cls.__name__} must set a class-level `backend` string " - "(the WeightTransferTrainerFactory registry key)." - ) - - @property - def is_sender(self) -> bool: - return self.rank == 0 - - -@runtime_checkable -class VLLMWeightSyncClient(Protocol): - """Trainer-side stub for the inference engine's weight-sync control plane. - - Mirrors the weight-sync methods that the inference engine exposes - (`EngineClient` / the HTTP RLHF routes / Ray actors). A - `TrainerWeightTransferEngine` drives the full handshake through this - protocol so trainer code never has to know the transport. - - All methods are synchronous and accept plain dicts (matching what the - inference side already accepts). Concurrency that some backends need - (e.g. NCCL must run `update_weights` concurrently with the trainer-side - broadcast) is the engine's responsibility, not the client's, so the - protocol stays a flat four-method surface that any wrapper can implement. - - The protocol is structural (PEP 544), so user implementations need only - define these four methods — no import or subclassing required. - """ - - def init_weight_transfer_engine(self, init_info: dict[str, Any]) -> None: ... - - def start_weight_update(self) -> None: ... - - def update_weights(self, update_info: dict[str, Any]) -> None: ... - - def finish_weight_update(self) -> None: ... - - -class TrainerWeightTransferEngine(ABC, Generic[TTrainerInitInfo]): - """Trainer-side weight transfer engine. - - Symmetric to `WeightTransferEngine` but lives in the training process. - Constructed via the `trainer_init` factory classmethod; carries any - backend-specific state (NCCL communicators, IPC device info, transfer - plans) on `self`. Full-resync backends (NCCL, IPC) take a `WeightSource` at - `trainer_init` and replay it each round via the no-argument - `send_weights()`. Backends that push per-round deltas instead (e.g. sparse - patches) leave `source` as `None` and take their payload as a `send_weights` - argument. - - Unlike the worker engine, the trainer side does not take a - `WeightTransferConfig`: the backend is selected from the init info's - `backend` `ClassVar` (so callers pass only the init info), and the static - wire params (packed, buffer sizes) ride the backend-specific - `TrainerInitInfo`, which the sender also propagates to the worker at the init - handshake. - - Multi-rank trainers: `trainer_init` and `send_weights` are - called on *every* trainer rank. Rank 0 is the sender, resolved once at - `trainer_init` into `is_sender`. Non-sender ranks still run every - collective (iterating the source, metadata export, IPC handle all-gather) so - the group stays aligned, but each engine explicitly guards the control-plane - RPCs and the transmit on `self.is_sender`, so only the sender touches the - client. - - Subclasses should define: - init_info_cls: Type of backend-specific trainer init info - """ - - # Subclasses should override this class attribute - init_info_cls: type[TTrainerInitInfo] - - def __init__( - self, - *, - client: "VLLMWeightSyncClient", - source: "WeightSource | None" = None, - is_sender: bool = True, - ) -> None: - self.is_sender = is_sender - # The real client is held on every rank; each engine only *calls* it when - # `is_sender`, so non-sender ranks never touch the wire. - self.client = client - self.source = source - - @classmethod - @abstractmethod - def trainer_init( - cls, - init_info: TTrainerInitInfo, - *, - client: "VLLMWeightSyncClient", - source: "WeightSource | None" = None, - ) -> Self: - """Rendezvous with the inference side and return a ready instance. - - Called on every trainer rank. The sender drives the full handshake via - `client` (build the worker-side init info, call - `client.init_weight_transfer_engine`, open the trainer-side endpoint); - non-sender ranks skip the rendezvous and the RPC. - """ - raise NotImplementedError - - @abstractmethod - def send_weights(self) -> None: - """Push weights to inference workers and drive the full update round - trip: `start_weight_update`, `update_weights` (run concurrently with the - trainer-side broadcast when the backend requires it), then - `finish_weight_update`. Called on every trainer rank. - """ - raise NotImplementedError - - def shutdown(self) -> None: - """Tear down communicators / process groups. Default no-op.""" diff --git a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_engine.py b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_engine.py index 893f3109ed..9ab2b03a4f 100644 --- a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_engine.py +++ b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_engine.py @@ -51,8 +51,9 @@ VENDORED from the ``vllm-rdt-weight-sync`` fork's ``vllm/distributed/weight_transfer/sharded_rdt_engine.py``. Keep the two in sync; the only intended differences are the import paths, the LIBFABRIC shim below, and -the no-op ``trainer_send_weights`` the pinned wheel's ABC still requires (the -fork's ABC dropped it). +the ``set_current_vllm_config`` + ``torch.device`` brackets this copy opens itself +(``Worker.init_weight_transfer_engine`` is the one lifecycle method vLLM does not +wrap). REMOVAL: delete this module once SkyRL's pinned vLLM registers the ``sharded_rdt`` engine itself, and repoint the factory registration at @@ -60,7 +61,6 @@ """ import time -from collections.abc import Iterator from dataclasses import dataclass, field from math import prod from typing import TYPE_CHECKING, Any, cast @@ -303,9 +303,9 @@ class ShardedRDTWeightTransferEngine( init_info_cls = ShardedRDTWeightTransferInitInfo update_info_cls = ShardedRDTWeightTransferUpdateInfo # receive_weights pulls synchronously but defers GPU post-processing to a - # background thread so it overlaps the next chunk's pull, so ``update_weights`` - # skips the base's device sync and ``finish_weight_update`` drains the - # deferred work before finalize. + # background thread so it overlaps the next chunk's pull. So ``update_weights`` + # skips the base's device sync, and ``finish_weight_update`` drains the + # deferred work before finalize. Declarative: nothing reads the flag. defers_processing = True # The baked replay plan is a function of one concrete model's parameter # layout, so a separate draft model cannot reuse it. @@ -401,6 +401,20 @@ def __init__( self._proc_error: BaseException | None = None def init_transfer_engine(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: + """Resolve producers and run the one-time bake. + + Opens ``set_current_vllm_config`` + ``torch.device`` itself: + ``Worker.init_weight_transfer_engine`` is the one lifecycle method vLLM + does not wrap, and the bake drives ``model.load_weights`` against meta + params, so ``process_weights_after_loading`` on MoE models reads + ``get_current_vllm_config()`` to build its kernels. + """ + from vllm.config import set_current_vllm_config + + with set_current_vllm_config(self.vllm_config), torch.device(self.device): + self._init_transfer_engine(init_info) + + def _init_transfer_engine(self, init_info: ShardedRDTWeightTransferInitInfo) -> None: """Configure the ring, bind the producers, bake the replay plan, and pre-register every NIXL buffer -- in that order, because each step depends on the previous one. @@ -658,7 +672,8 @@ def start_weight_update(self) -> None: initialize_layerwise_reload, ) - initialize_layerwise_reload(self.model) + with torch.device(self.device): + initialize_layerwise_reload(self.model) def finish_weight_update(self) -> None: """Drain the deferred pull/process pipeline (so every layer is fully @@ -668,7 +683,8 @@ def finish_weight_update(self) -> None: ) self.drain_pending() - finalize_layerwise_reload(self.model, self.model_config) + with torch.device(self.device): + finalize_layerwise_reload(self.model, self.model_config) def update_weights(self, update_info: dict[str, Any]) -> None: """Receive one update. Unlike the base, does NOT issue a per-update @@ -676,7 +692,8 @@ def update_weights(self, update_info: dict[str, Any]) -> None: sync here would block on them and serialize the pull/process pipeline. Completion is guaranteed by ``drain_pending`` in ``finish_weight_update``.""" - self.receive_weights(self.parse_update_info(update_info)) + with torch.device(self.device): + self.receive_weights(self.parse_update_info(update_info)) def receive_weights( self, @@ -1442,22 +1459,6 @@ def shutdown(self) -> None: # process lifetime; freeing the tensors just drops our strong refs). self._dest_buffers = [{} for _ in range(self._ring_depth)] - @staticmethod - def trainer_send_weights( - iterator: Iterator[tuple[str, torch.Tensor]], - trainer_args: dict[str, Any] | Any, - ) -> None: - """No-op for the pull-based sharded RDT backend. - - Workers initiate the transfer themselves via the trainer's - ``@ray.method(tensor_transport="nixl")`` batched accessor. - - SkyRL-only: the pinned vLLM's worker-side ``WeightTransferEngine`` - declares this abstract, so the class cannot be instantiated without it. - The fork's ABC dropped it, which is why the fork's copy has no stub. - """ - del iterator, trainer_args - def _plan_digest(keys_per_chunk: list) -> str: """Digest of the chunks one consumer pulls from one producer, in pull order. diff --git a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_strategy.py b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_strategy.py deleted file mode 100644 index 9a84e1be9a..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_strategy.py +++ /dev/null @@ -1,229 +0,0 @@ -"""``sharded_rdt`` as a :class:`WeightTransferStrategy`. - -A thin adapter over the vendored vLLM trainer-send stack — a ``WeightSource`` feeding -a ``ShardedRDTTrainerWeightTransferEngine`` that the inference workers pull from, -driven by :class:`RdtWeightSyncSender` (``rdt_send.py``). Presenting it through the -shared interface keeps the workers to one sender attribute and one send call. - -Two places RDT does not fit the push backends' shape, declared as capabilities so the -workers hold no backend conditional: - -* ``trainer_init`` opens the inference side itself, through the blocking - ``SyncRdtControlPlaneClient``, since the bake needs the source metadata and must run - under ``set_current_vllm_config``. ``sender_initializes_receivers`` is True, and - worker rank 0 must not also call ``init_weight_update_communicator``. -* The consumers pull, so there is no chunk stream: this sender overrides - :meth:`WeightTransferSender.send` and never materializes chunks or metadata. - ``get_weight_metadata`` on the Megatron extractor is a whole-model - ``export_hf_weights`` pass, the gather the RDT weight source avoids. - -REMOVAL: when SkyRL's pinned vLLM ships trainer-send for NCCL/IPC too, those backends -collapse into this shape and the ``WeightTransferStrategy`` layer goes away. This -adapter is what disappears then; ``rdt_send.py`` stays. -""" - -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Dict, Iterable, Optional - -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk -from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( - WeightSyncInitInfo, - WeightTransferSender, - WeightTransferStrategy, -) - -if TYPE_CHECKING: - import torch - - from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( - RemoteInferenceClient, - ) - from skyrl.train.config import InferenceEngineConfig - - -@dataclass -class ShardedRdtInitInfo(WeightSyncInitInfo): - """Config-derived args for the sharded-RDT sender. - - Deliberately does NOT implement ``for_servers`` / ``to_api_payload``: those - exist for the driver-side ``init_weight_update_communicator`` fan-out, which - this backend does not use (see ``sender_initializes_receivers``). The - inference-side payload is built by the trainer engine instead, from the - weight source's metadata, and shipped over ``/collective_rpc``. - """ - - model_dtype: str - """Inference dtype, as a ``torch.dtype`` string. The weight source casts to it.""" - - inference_world_size: int - """Total inference workers across the fleet — the consumer count the - producers' free barrier counts against.""" - - trainer_actor_namespace: Optional[str] - """Ray namespace the per-rank producer sidecars are registered in, so the - consumers can resolve them by name.""" - - -class ShardedRdtWeightTransferSender(WeightTransferSender): - """Presents :class:`RdtWeightSyncSender` as a ``WeightTransferSender``.""" - - # The producer sidecar shares every gathered group with this rank over CUDA - # IPC on every run, not only under colocation, and expandable-segment (VMM) - # memory makes that export/rebuild 5-10x slower per storage. - force_disable_expandable_segments = True - - # Publish buffers freed during a sync stay in this process's allocator cache and - # are reused by the next training step; returning them to CUDA costs 0.25-0.53s - # per rank at 235B and buys nothing. The worker still empties under colocate_all, - # where an inference engine wants the physical memory. - empty_cache_after_send = False - - def __init__(self, sender: Any) -> None: - self._sender = sender - - async def send( - self, - weight_extractor: Any, - dtype: "torch.dtype", - **kwargs, - ) -> None: - """Run one pull-based sync. Every rank must call it: the gather is a - collective. - - ``dtype`` is ignored — the weight source was built with the inference - dtype at init and does the cast itself. The push backends' kwargs - (``reset_prefix_cache``) are ignored too: this sender does not handle the - prefix cache, so the worker resets it (``handles_prefix_cache_reset`` - stays False). - """ - del dtype, kwargs - await self._sender.send(weight_extractor) - - async def send_chunks( - self, - chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, - **kwargs, - ) -> None: - """Not applicable: the consumers pull, so there is no chunk stream to - push. :meth:`send` is this backend's entry point.""" - raise NotImplementedError( - "sharded_rdt does not push a chunk stream: the inference workers pull " - "the slices they consume. Call send(weight_extractor, dtype) instead." - ) - - def teardown(self) -> None: - self._sender.teardown() - - -class ShardedRdtTransferStrategy(WeightTransferStrategy): - """Factory for the sharded-RDT (NIXL pull) sender.""" - - # trainer_init drives the inference-side init through the engine's own - # control-plane client; the worker must not also push init_info. - sender_initializes_receivers = True - - @staticmethod - def create_init_info( - ie_cfg: "InferenceEngineConfig", - inference_world_size: Optional[int] = None, - base_model_path: Optional[str] = None, - ) -> ShardedRdtInitInfo: - """Collect the config-derived args. Runs on every training rank. - - Raises: - ValueError: the inference world size is missing. It is the consumer - count the whole ownership arithmetic is sized from, and a wrong - value silently mis-maps consumers onto slices, so there is no - safe default. - """ - del base_model_path # weights come off the live model, never a checkpoint - if not inference_world_size: - raise ValueError( - f"sharded_rdt requires the inference world size (consumer count); got {inference_world_size!r}." - ) - return ShardedRdtInitInfo( - override_existing_receiver=not ie_cfg.run_engines_locally, - model_dtype=ie_cfg.model_dtype, - inference_world_size=int(inference_world_size), - trainer_actor_namespace=_ray_namespace(), - ) - - @staticmethod - def create_sender( - init_info: WeightSyncInitInfo, - inference_client: "RemoteInferenceClient", - weight_extractor: Any = None, - ) -> ShardedRdtWeightTransferSender: - """Build the sender AND rendezvous, on every rank. - - The rendezvous (spawn each rank's producer sidecar, then the sender rank's - bake) happens here rather than on the first send because deferring it - deadlocks: rank 0 would block in the inference-side init RPC while the - other ranks spin in gather collectives, and the sidecars — which share - those ranks' GPUs — could then never finish NIXL agent creation, because - libfabric's CUDA probe blocks behind the spinning kernels. Every rank is - inside ``init_weight_sync_state`` here, so that window is empty. - - Blocking is expected: the worker calls this off the event loop, and - ``init_weight_sync_state`` already ends in a barrier. - - Raises: - ValueError: ``init_info`` is not a :class:`ShardedRdtInitInfo`. - RuntimeError: no ``weight_extractor``. The rendezvous needs the model - to build its weight source, so this is a wiring error, not a - config one. - """ - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_vllm_register - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - RdtWeightSyncSender, - ) - - if not isinstance(init_info, ShardedRdtInitInfo): - raise ValueError(f"sharded_rdt requires a ShardedRdtInitInfo, got {type(init_info).__name__}.") - if weight_extractor is None: - raise RuntimeError( - "sharded_rdt weight sync requires the worker's weight_extractor, which " - "must be built before init_weight_sync_state runs." - ) - - # Registers the consumer engine in vLLM's factory under "sharded_rdt". - # Needed on the driver and every vLLM worker; harmless (idempotent) here. - rdt_vllm_register.ensure_registered() - - sender = RdtWeightSyncSender( - inference_client, - init_info.model_dtype, - init_info.inference_world_size, - init_info.trainer_actor_namespace, - ) - sender.initialize(weight_extractor) - return ShardedRdtWeightTransferSender(sender) - - @staticmethod - def get_vllm_transfer_engine() -> type: - """The receive-side engine, as registered in vLLM's factory. - - Unlike the push strategies' mapping this one is live: the consumers really - do construct this class, via ``WeightTransferEngineFactory`` under the - name ``rdt_vllm_register`` registers it as. - """ - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_engine import ( - ShardedRDTWeightTransferEngine, - ) - - return ShardedRDTWeightTransferEngine - - -def _ray_namespace() -> Optional[str]: - """This process's Ray namespace, or None outside a Ray runtime. - - The producer sidecars are named actors, so the consumers need the namespace - they were created in to resolve them. - """ - try: - import ray - - return ray.get_runtime_context().namespace or None - except Exception: # noqa: BLE001 - no Ray runtime: the caller falls back to the default namespace - return None diff --git a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_trainer.py b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_trainer.py index f2191b8d51..0a5bf86ad9 100644 --- a/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_trainer.py +++ b/skyrl/backends/skyrl_train/weight_sync/sharded_rdt/sharded_rdt_trainer.py @@ -24,10 +24,12 @@ consumer to rejoin, both of which are deliberately omitted here. The consumer set is the provisioned fleet, fixed at ``trainer_init``. Do not "restore" them when syncing; the producer-side stall watchdog IS kept, because the slot-sharing -rendezvous relies on it as its failure mode. +rendezvous relies on it as its failure mode. This copy also carries the +``skyrl_*`` capability flags the workers probe; they are additive, so the fork can +be re-synced over them. REMOVAL: delete this module once SkyRL's pinned vLLM ships the trainer-side -engine, and repoint ``weight_sync/sharded_rdt/rdt_send.py`` at +engine, and repoint ``weight_sync/trainer_engines.py``'s registration at ``vllm.distributed.weight_transfer.sharded_rdt_trainer``. """ @@ -799,6 +801,18 @@ class ShardedRDTTrainerWeightTransferEngine(TrainerWeightTransferEngine[ShardedR init_info_cls = ShardedRDTTrainerInitInfo + # Worker capability probes (see workers/worker.py). + # + # The producer sidecar shares every gathered group with this rank over CUDA + # IPC on every run, not only under colocation, and expandable-segment (VMM) + # memory makes that export/rebuild 5-10x slower per storage (measured: + # publish rebuild 7.2s/rank/sync at 30B, the dominant weight-sync cost). + skyrl_force_disable_expandable_segments = True + # Publish buffers stay in this process's allocator cache and are reused by + # the next training step; returning them to CUDA costs 0.25-0.53s per rank at + # 235B and buys nothing. + skyrl_empty_cache_after_send = False + def __init__( self, *, diff --git a/skyrl/backends/skyrl_train/weight_sync/skyrl_engines.py b/skyrl/backends/skyrl_train/weight_sync/skyrl_engines.py new file mode 100644 index 0000000000..c8538b8963 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/skyrl_engines.py @@ -0,0 +1,195 @@ +"""SkyRL's receive-side weight-transfer engines (the inference-worker half). + +vLLM's NCCL and IPC engines, subclassed to add one thing: reloading the +speculative-decoding drafter. The drafter (``model_runner.drafter.model``) is a +separate module that the main model's ``load_weights`` never touches, and vLLM's +engines call ``self.model.load_weights(...)`` directly with no callback, so the +only injection point is the ``self.model`` handle they read. + +Registered under ``skyrl_nccl`` / ``skyrl_ipc`` rather than shadowing vLLM's +``nccl`` / ``ipc``: ``register_engine`` raises on a duplicate name, and +``WeightTransferConfig.backend`` is typed ``Literal[...] | str`` and validated +against the registry, so a new name is all that is needed. +""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from typing import Any, Iterator + +import torch + +logger = logging.getLogger(__name__) + +SKYRL_NCCL_BACKEND = "skyrl_nccl" +SKYRL_IPC_BACKEND = "skyrl_ipc" + +_REGISTERED = False + + +def empty_cuda_cache_rocm() -> None: + """Release unused ROCm cached blocks after a full-weight sync. + + ROCm's allocator does not return the reload's transient blocks on its own. + """ + if torch.version.hip is None or not torch.cuda.is_available(): + return + device = torch.cuda.current_device() + torch.cuda.synchronize(device) + torch.cuda.empty_cache() + torch.cuda.synchronize(device) + + +class _LoadWeightsProxy: + """Wraps a model, overriding only ``load_weights``. + + Every other attribute access falls through to the real model. + """ + + def __init__(self, model: Any, load_weights: Any) -> None: + self._model = model + self.load_weights = load_weights + + def __getattr__(self, name: str) -> Any: + # Only reached for attributes not set on the proxy itself. + return getattr(self._model, name) + + +class SkyrlDrafterReloadMixin: + """Reload the spec-decode drafter from the same weights the main model got. + + Wrap the engine's own load in :meth:`skyrl_drafter_reload`. Costs nothing + when this process has no drafter: the proxy is not installed at all. + """ + + @contextmanager + def skyrl_drafter_reload(self) -> Iterator[None]: + """Install the drafter-reloading proxy over ``self.model``, if needed.""" + from skyrl.backends.skyrl_train.patches.vllm.patch_model_runner_registry import ( + current_model_runner, + ) + + model_runner = current_model_runner() + drafter = getattr(model_runner, "drafter", None) if model_runner is not None else None + if drafter is None or getattr(drafter, "model", None) is None: + # No speculative decoding, or a proposer with no loadable model + # (ngram): nothing to interpose. + yield + return + + from skyrl.backends.skyrl_train.inference_servers.spec_decode_utils import ( + _reload_spec_decode_drafter, + ) + + model = self.model + + def load_weights(weights: Any, **kwargs: Any) -> Any: + # The engines hand us a one-shot generator, and the drafter needs its + # own pass to filter for the names it consumes. + weight_list = list(weights) + loaded = model.load_weights(weights=weight_list, **kwargs) + _reload_spec_decode_drafter(model_runner, weight_list) + return loaded + + self.model = _LoadWeightsProxy(model, load_weights) + try: + yield + finally: + # Restore the exact object we found, so this composes with the + # worker's set_weight_update_target / reset_weight_update_target. + self.model = model + + +# Each engine brackets its lifecycle in `torch.device(self.device)`. vLLM's own +# path does not (it passes `device=` where it matters), but SkyRL's loaders have +# always run under it, so keep it as the default device weight loading sees. + + +def _build_skyrl_nccl_engine() -> type: + from vllm.distributed.weight_transfer.nccl_engine import NCCLWeightTransferEngine + + class SkyrlNCCLWeightTransferEngine(SkyrlDrafterReloadMixin, NCCLWeightTransferEngine): + """vLLM's dense NCCL receive engine plus the drafter reload.""" + + def start_weight_update(self) -> None: + with torch.device(self.device): + super().start_weight_update() + + def receive_weights(self, update_info: Any) -> None: + with torch.device(self.device), self.skyrl_drafter_reload(): + super().receive_weights(update_info) + + def finish_weight_update(self) -> None: + with torch.device(self.device): + super().finish_weight_update() + empty_cuda_cache_rocm() + + return SkyrlNCCLWeightTransferEngine + + +def _build_skyrl_ipc_engine() -> type: + from vllm.distributed.weight_transfer.ipc_engine import IPCWeightTransferEngine + + class SkyrlIPCWeightTransferEngine(SkyrlDrafterReloadMixin, IPCWeightTransferEngine): + """vLLM's CUDA IPC receive engine plus the drafter reload.""" + + def start_weight_update(self) -> None: + with torch.device(self.device): + super().start_weight_update() + + def receive_weights(self, update_info: Any) -> None: + with torch.device(self.device), self.skyrl_drafter_reload(): + super().receive_weights(update_info) + + def finish_weight_update(self) -> None: + with torch.device(self.device): + super().finish_weight_update() + empty_cuda_cache_rocm() + + return SkyrlIPCWeightTransferEngine + + +# Built lazily and cached: the classes subclass vLLM's engines, and this module +# must stay importable without the wheel. +_ENGINE_CACHE: dict[str, type] = {} + + +def get_skyrl_nccl_engine() -> type: + if SKYRL_NCCL_BACKEND not in _ENGINE_CACHE: + _ENGINE_CACHE[SKYRL_NCCL_BACKEND] = _build_skyrl_nccl_engine() + return _ENGINE_CACHE[SKYRL_NCCL_BACKEND] + + +def get_skyrl_ipc_engine() -> type: + if SKYRL_IPC_BACKEND not in _ENGINE_CACHE: + _ENGINE_CACHE[SKYRL_IPC_BACKEND] = _build_skyrl_ipc_engine() + return _ENGINE_CACHE[SKYRL_IPC_BACKEND] + + +def register_skyrl_engines() -> None: + """Register ``skyrl_nccl`` / ``skyrl_ipc`` in vLLM's factory (idempotent). + + Must run in every vLLM worker process (``Worker.load_model`` builds the + engine through the factory) and on the driver (which validates + ``WeightTransferConfig.backend`` against the registry). + """ + global _REGISTERED + if _REGISTERED: + return + try: + from vllm.distributed.weight_transfer.factory import WeightTransferEngineFactory + except ImportError: + # vLLM is a Linux-only optional dependency; nothing to register. + logger.debug("vLLM not importable; skipping SkyRL weight-transfer engine registration.") + return + + for name, loader in ( + (SKYRL_NCCL_BACKEND, get_skyrl_nccl_engine), + (SKYRL_IPC_BACKEND, get_skyrl_ipc_engine), + ): + if name in WeightTransferEngineFactory._registry: + continue + # Direct-class registration, so resolve now; vLLM is importable by here. + WeightTransferEngineFactory.register_engine(name, loader()) + _REGISTERED = True diff --git a/skyrl/backends/skyrl_train/weight_sync/sources.py b/skyrl/backends/skyrl_train/weight_sync/sources.py new file mode 100644 index 0000000000..33886f8b23 --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/sources.py @@ -0,0 +1,120 @@ +"""``WeightSource`` implementations over SkyRL's training backends. + +A ``WeightSource`` is vLLM's trainer-side weight contract +(``vllm.distributed.weight_transfer.base``): ``metadata()`` and ``__iter__``, +which must agree element for element. Chunking is vLLM's: the packed producers +cut bounded-memory chunks out of a fixed buffer and consume the source lazily, +so a source's only obligation is to be a lazy generator that does not retain. + +Sharded RDT needs two more channels (per-rank ownership under PP/EP, and a group +index) that vLLM has no concept of; they live in +``sharded_rdt/sharded_rdt_base.GroupedWeightSource``. + +Imports vLLM at module scope, so import this lazily from anything that must work +without the wheel. +""" + +from typing import Any, Iterator, List, Optional, Tuple + +import torch +from vllm.distributed.weight_transfer.base import ( + ParamMeta, + WeightSource, + materialize_full_tensor, +) + +__all__ = [ + "FsdpWeightSource", + "MegatronWeightSource", + "ParamMeta", + "WeightSource", + "materialize_full_tensor", +] + + +class FsdpWeightSource(WeightSource): + """``WeightSource`` over an FSDP2-sharded HF model. + + ``metadata()`` reads ``state_dict()`` shapes only: an FSDP2 ``DTensor``'s + ``.shape`` is already the global shape, so declaring the stream costs no + collective. Iteration all-gathers each parameter (``full_tensor()``), which + IS a collective, so every trainer rank must iterate the same source in the + same order. + + Args: + model: the inner HF module (``self.model.model`` on the worker), whose + ``state_dict()`` keys are the names vLLM expects. + dtype: inference dtype. Both channels use it. + weight_prefix: prepended to every name (``"language_model."`` when + syncing a CausalLM backbone into a vLLM multimodal namespace). + """ + + def __init__(self, model: torch.nn.Module, dtype: torch.dtype, weight_prefix: str = "") -> None: + self._model = model + self._dtype = dtype + self._prefix = weight_prefix or "" + + @property + def model(self) -> torch.nn.Module: + return self._model + + @property + def weight_prefix(self) -> str: + return self._prefix + + def metadata(self) -> List[ParamMeta]: + sd = self._model.state_dict() + return [ParamMeta(f"{self._prefix}{key}", self._dtype, tuple(param.shape)) for key, param in sd.items()] + + def __iter__(self) -> Iterator[Tuple[str, torch.Tensor]]: + # The caller selects this rank's CUDA device before iterating; a worker + # thread does not inherit it (see Worker._weight_sync_thread). + device = torch.cuda.current_device() if torch.cuda.is_available() else None + for key, param in self._model.state_dict().items(): + if device is not None: + param = param.to(device, non_blocking=True) + full = materialize_full_tensor(param).to(self._dtype).detach().contiguous() + yield f"{self._prefix}{key}", full + + +class MegatronWeightSource(WeightSource): + """``WeightSource`` over a Megatron model, via Megatron-Bridge. + + ``export_hf_weights(conversion_tasks=None)`` streams the whole model in + HF-canonical order, gathering each parameter across TP / PP / EP internally + and giving all ranks full tensors. It is a generator, so its collectives run + as the consumer pulls and a lazy consumer never holds more than the tensors + it is working on. + + Exporting in ONE call is required: the bridge's ``_accumulate_grouped_export`` + needs every task sharing a ``group_key`` present in the same call, or the + expert weights are silently never yielded. + + There is no shape-only export, so ``metadata()`` runs a dry export and + caches. Engines call it every round; the cost is one-time. + """ + + def __init__(self, bridge: Any, module: Any, dtype: torch.dtype) -> None: + self._bridge = bridge + self._module = module + self._dtype = dtype + self._meta: Optional[List[ParamMeta]] = None + + def _export(self) -> Iterator[Tuple[str, torch.Tensor]]: + return self._bridge.export_hf_weights(self._module, show_progress=False, conversion_tasks=None) + + def metadata(self) -> List[ParamMeta]: + if self._meta is None: + meta: List[ParamMeta] = [] + for name, tensor in self._export(): + meta.append(ParamMeta(name, self._dtype, tuple(tensor.shape))) + del tensor + self._meta = meta + return self._meta + + def __iter__(self) -> Iterator[Tuple[str, torch.Tensor]]: + # See FsdpWeightSource.__iter__ on the device. + device = torch.cuda.current_device() if torch.cuda.is_available() else None + for name, tensor in self._export(): + full = tensor.to(device=device, dtype=self._dtype, non_blocking=True).detach().contiguous() + yield name, full diff --git a/skyrl/backends/skyrl_train/weight_sync/trainer_engines.py b/skyrl/backends/skyrl_train/weight_sync/trainer_engines.py new file mode 100644 index 0000000000..4af9002a1a --- /dev/null +++ b/skyrl/backends/skyrl_train/weight_sync/trainer_engines.py @@ -0,0 +1,310 @@ +"""Build the trainer-side weight-transfer engine for a backend. + +:func:`build_trainer_engine` is the whole trainer-side weight-sync surface SkyRL +owns: pick the init info, build the control-plane client, and hand both to +``WeightTransferTrainerFactory.trainer_init``, which rendezvouses and returns an +engine whose ``send_weights()`` owns the round trip. + +=============== ================================================== +``nccl`` vLLM's ``NCCLTrainerWeightTransferEngine`` +``ipc`` vLLM's ``IPCTrainerWeightTransferEngine`` +``delta`` ``weight_sync/delta_trainer.py`` +``sharded_rdt`` ``weight_sync/sharded_rdt/sharded_rdt_trainer.py`` +=============== ================================================== + +The trainer- and worker-side factories keep separate registries, so the trainer +engines use vLLM's ``nccl`` / ``ipc`` keys even though the receive side registers +under ``skyrl_nccl`` / ``skyrl_ipc`` (see ``skyrl_engines.py``). +""" + +from __future__ import annotations + +import logging +import math +import socket +from typing import TYPE_CHECKING, Any, Callable, Optional + +from vllm.distributed.weight_transfer.packed_tensor import ( + DEFAULT_PACKED_BUFFER_SIZE_BYTES, +) + +from skyrl.backends.skyrl_train.weight_sync.control_plane import ( + SkyrlWeightSyncClient, + nccl_init_payloads, + rdt_init_payloads, +) + +if TYPE_CHECKING: + import torch + + from skyrl.train.config.config import InferenceEngineConfig + +logger = logging.getLogger(__name__) + + +def build_trainer_engine( + *, + ie_cfg: "InferenceEngineConfig", + colocate_all: bool, + rank: int, + inference_world_size: int, + source_factory: Callable[["torch.dtype", str], Any], + server_urls: list, + data_parallel_size: int, + base_model_path: Optional[str] = None, +) -> Any: + """Resolve the backend, build this rank's source, rendezvous, and return the engine. + + Called on **every** trainer rank, off the event loop: rank 0 drives the + inference-side handshake while the others return without touching the wire. + + The backend is resolved here, from the same two config values the driver uses + to configure the inference servers (``get_transfer_strategy``), so the two + sides cannot pick different engines. + + Args: + ie_cfg: inference engine config. Supplies the backend and the inference + dtype; ``delta`` also reads its ``delta_weight_sync`` block. + colocate_all: ``trainer.placement.colocate_all``. + rank: this trainer process's rank. Rank 0 is the sender. + inference_world_size: total inference workers, from + ``client.get_world_size()``. + source_factory: ``(dtype, backend) -> WeightSource``. A callback because + the source reads the live model, which only the caller has, and it + cannot be built until the backend is known -- sharded RDT needs an + ownership-aware subclass. + server_urls: every inference server, in deployment-major order. + data_parallel_size: DP replicas per deployment. + base_model_path: policy model path. Required by ``delta``, which + publishes against that checkpoint. + """ + from vllm.distributed.weight_transfer.factory import WeightTransferTrainerFactory + + from skyrl.backends.skyrl_train.weight_sync import get_transfer_strategy + from skyrl.train.utils.utils import str_to_torch_dtype + + backend = get_transfer_strategy(ie_cfg.weight_sync_backend, colocate_all) + source = source_factory(str_to_torch_dtype(ie_cfg.model_dtype), backend) + + init_info, init_payload_fn = _build_init_info( + backend=backend, + ie_cfg=ie_cfg, + rank=rank, + inference_world_size=inference_world_size, + source=source, + server_urls=server_urls, + data_parallel_size=data_parallel_size, + base_model_path=base_model_path, + ) + client = SkyrlWeightSyncClient( + server_urls, + data_parallel_size, + init_payload_fn=init_payload_fn, + ) + if backend == "sharded_rdt": + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_send + + rdt_send.log_source_choice(source) + + engine = WeightTransferTrainerFactory.trainer_init(init_info, client=client, source=source) + + if backend == "sharded_rdt": + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_send + + rdt_send.freeze_trainer_heap() + return engine + + +def _packed_buffer_size_bytes(source: Any) -> int: + """Packed-buffer size that fits the model's largest single parameter. + + The packed producers stream through a fixed reusable buffer, and a parameter + too large for one raises on the IPC path and over-allocates on NCCL. vLLM's + 1 GiB default is smaller than a large-vocab embedding matrix (Qwen3-235B's is + 151936 x 4096 in bf16 = 1.24 GiB), so size it from the source. + + ``metadata()`` is a collective on a Megatron source, so this runs on every + rank -- it is called before the sender split, which keeps them in lockstep -- + and the source caches it. + """ + meta = source.metadata() + if not meta: + return DEFAULT_PACKED_BUFFER_SIZE_BYTES + largest = max(math.prod(m.shape) * m.dtype.itemsize for m in meta) + return max(DEFAULT_PACKED_BUFFER_SIZE_BYTES, largest) + + +def _build_init_info( + *, + backend: str, + ie_cfg: "InferenceEngineConfig", + rank: int, + inference_world_size: int, + source: Any, + server_urls: list, + data_parallel_size: int, + base_model_path: Optional[str], +): + """Return ``(init_info, init_payload_fn)`` for an already-resolved backend. + + ``init_payload_fn`` expands the engine's single worker-side init dict to one + payload per server; only NCCL and sharded RDT need it (see ``control_plane``). + """ + _register_skyrl_trainer_engines() + + if backend == "nccl": + import ray + from vllm.distributed.weight_transfer.nccl_engine import NCCLTrainerInitInfo + + # Only rank 0 opens the endpoint, so only its address/port reaches a + # worker; the other ranks build and discard theirs. + master_address = ray._private.services.get_node_ip_address() + with socket.socket() as sock: + sock.bind(("", 0)) + master_port = sock.getsockname()[1] + return ( + NCCLTrainerInitInfo( + master_address=master_address, + master_port=master_port, + # Every inference worker plus the single trainer sender (rank 0). + world_size=inference_world_size + 1, + # Broadcast out of a fixed reusable buffer instead of one NCCL + # call per parameter. The engine propagates this to the worker at + # the handshake, so the two sides cannot disagree. + packed=True, + packed_buffer_size_bytes=_packed_buffer_size_bytes(source), + rank=rank, + ), + nccl_init_payloads, + ) + + if backend == "ipc": + from vllm.distributed.weight_transfer.ipc_engine import IPCTrainerInitInfo + + return ( + # packed=True overrides the vLLM default: the unpacked path holds a + # strong ref to a contiguous copy of EVERY parameter until past + # `finish_weight_update` (so the consumer's IPC views stay valid), + # i.e. the whole model resident on the trainer. Packed streams + # through one reusable buffer. + IPCTrainerInitInfo( + packed=True, + packed_buffer_size_bytes=_packed_buffer_size_bytes(source), + rank=rank, + ), + None, + ) + + if backend == "delta": + from skyrl.backends.skyrl_train.weight_sync.delta_checkpoint import ( + SUPPORTED_CHECKPOINT_LOAD_FORMATS, + ) + from skyrl.backends.skyrl_train.weight_sync.delta_trainer import ( + DeltaTrainerInitInfo, + ) + + if base_model_path is None: + raise ValueError("Delta weight sync requires base_model_path") + delta_cfg = ie_cfg.delta_weight_sync + if delta_cfg is None or not delta_cfg.sync_dir: + raise ValueError("Delta weight sync requires generator.inference_engine.delta_weight_sync.sync_dir") + if delta_cfg.checkpoint_load_format not in SUPPORTED_CHECKPOINT_LOAD_FORMATS: + raise ValueError( + "Delta checkpoint_load_format must be one of " + f"{sorted(SUPPORTED_CHECKPOINT_LOAD_FORMATS)}, got {delta_cfg.checkpoint_load_format!r}" + ) + # local_checkpoint_dir and publish_staging_dir are resolved by + # DeltaWeightSyncConfig.__post_init__, so they are already concrete here. + return ( + DeltaTrainerInitInfo( + base_model_path=base_model_path, + sync_dir=delta_cfg.sync_dir, + local_checkpoint_dir=delta_cfg.local_checkpoint_dir, + publish_staging_dir=delta_cfg.publish_staging_dir, + max_file_size_in_gb=delta_cfg.max_file_size_in_gb, + cloud_download_workers=delta_cfg.cloud_download_workers, + publish_num_workers=delta_cfg.publish_num_workers, + checkpoint_load_format=delta_cfg.checkpoint_load_format, + multi_thread_safetensors_max_workers=delta_cfg.multi_thread_safetensors_max_workers, + rank=rank, + ), + None, + ) + + if backend == "sharded_rdt": + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_send + + return ( + rdt_send.build_rdt_trainer_init_info( + rank=rank, + inference_world_size=inference_world_size, + server_urls=list(server_urls), + data_parallel_size=data_parallel_size, + ), + rdt_init_payloads, + ) + + raise ValueError(f"Unknown weight sync backend {backend!r}.") + + +_TRAINER_ENGINES_REGISTERED = False + + +def _register_skyrl_trainer_engines() -> None: + """Register SkyRL's trainer engines (``delta``, ``sharded_rdt``) once.""" + global _TRAINER_ENGINES_REGISTERED + if _TRAINER_ENGINES_REGISTERED: + return + from vllm.distributed.weight_transfer.factory import WeightTransferTrainerFactory + + if "delta" not in WeightTransferTrainerFactory._registry: + WeightTransferTrainerFactory.register_engine( + "delta", + "skyrl.backends.skyrl_train.weight_sync.delta_trainer", + "DeltaTrainerWeightTransferEngine", + ) + if "sharded_rdt" not in WeightTransferTrainerFactory._registry: + WeightTransferTrainerFactory.register_engine( + "sharded_rdt", + "skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_trainer", + "ShardedRDTTrainerWeightTransferEngine", + ) + _TRAINER_ENGINES_REGISTERED = True + + +def engine_capability(engine: Any, name: str, default: Any) -> Any: + """Read a SkyRL capability flag off a trainer engine. + + A ``getattr`` probe rather than a declared attribute: two of the four engines + are vLLM's own classes and cannot carry SkyRL attributes, so an engine that + declares nothing must get the default. + + Flags: ``skyrl_handles_prefix_cache_reset``, + ``skyrl_force_disable_expandable_segments``, + ``skyrl_empty_cache_after_send``. + """ + return getattr(engine, f"skyrl_{name}", default) + + +def maybe_set_reset_prefix_cache(engine: Any, reset: bool) -> None: + """Tell an engine whether to reset the prefix cache this round, if it cares. + + ``send_weights()`` takes no arguments, so a per-round flag has to ride the + engine. Only delta implements the setter. + """ + setter = getattr(engine, "skyrl_set_reset_prefix_cache", None) + if setter is not None: + setter(reset) + + +def teardown_engine(engine: Any) -> None: + """Shut an engine down and close its control-plane client.""" + if engine is None: + return + try: + engine.shutdown() + finally: + client = getattr(engine, "client", None) + close = getattr(client, "close", None) + if close is not None: + close() diff --git a/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py b/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py deleted file mode 100644 index 3dbe2de01d..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/transfer_strategy.py +++ /dev/null @@ -1,195 +0,0 @@ -"""Weight transfer strategy abstractions for distributed RL training. - -This module defines the abstract interfaces for transferring model weights -from training workers to inference engines. The strategy pattern allows different -transfer mechanisms (broadcast, CUDA IPC) to be used interchangeably. -""" - -from abc import ABC, abstractmethod -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, ClassVar, Dict, Iterable, Optional - -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk - -if TYPE_CHECKING: - import torch - -if TYPE_CHECKING: - from skyrl.backends.skyrl_train.inference_servers.remote_inference_client import ( - RemoteInferenceClient, - ) - from skyrl.train.config import InferenceEngineConfig - - -@dataclass -class WeightSyncInitInfo(ABC): - """Base class for weight sync initialization info.""" - - override_existing_receiver: bool - """Whether to override an existing weight receiver. If False and a receiver exists, init is skipped.""" - - -class WeightTransferSender(ABC): - """Strategy-specific component that sends WeightChunk data to inference actors. - - Implementations handle the transfer primitive (broadcast, CUDA IPC) and coordinate - with inference actors. - """ - - handles_prefix_cache_reset: bool = False - """Indicates whether the transfer strategy handles resetting prefix cache - for the inference engines internally.""" - - force_disable_expandable_segments: ClassVar[bool] = False - """Disable expandable_segments around the send even when NOT colocated. - - The push backends only need it under ``colocate_all`` (CUDA IPC calls - cudaIpcGetMemHandle, which VMM addresses break). A backend that shares GPU - memory on every run regardless of colocation sets this True.""" - - empty_cache_after_send: ClassVar[bool] = True - """Whether the worker should ``torch.cuda.empty_cache()`` after the send. - - False for backends whose send buffers are reused by the next step, where - scrubbing them back to CUDA is pure cost. A colocated inference engine needs - the physical memory regardless, so the worker still empties under - ``colocate_all``.""" - - async def send( - self, - weight_extractor: Any, - dtype: "torch.dtype", - **kwargs, - ) -> None: - """Send this rank's weights. Called on every training rank. - - The default materializes the extractor's chunk stream plus its metadata - and hands both to :meth:`send_chunks` — the push backends' contract. - Backends that do not consume a chunk stream override this instead, which - is what keeps ``get_weight_metadata`` (a whole-model gather on the - Megatron extractor) off their critical path entirely. - - Args: - weight_extractor: The worker's extractor, already built. - dtype: Inference dtype to convert to. - **kwargs: Forwarded to :meth:`send_chunks`. - """ - await self.send_chunks( - weight_extractor.extract_weights(dtype), - weight_metadata=weight_extractor.get_weight_metadata(dtype), - **kwargs, - ) - - @abstractmethod - async def send_chunks( - self, - chunks: Iterable[WeightChunk], - weight_metadata: Optional[Dict[str, list]] = None, - **kwargs, - ) -> None: - """Send chunks using this transfer strategy. - - This method must be called on all training ranks. Implementations may have - different behavior for different ranks. - - Args: - chunks: Iterable of WeightChunk objects to send. - weight_metadata: Optional pre-computed metadata (names, dtype_names, shapes). - When provided, allows the sender to avoid materializing all chunks - to collect metadata upfront. - """ - ... - - @abstractmethod - def teardown(self) -> None: - """Clean up resources used by the sender (e.g., destroy process groups).""" - ... - - -# NOTE (sumanthrh): WeightTransferStrategy is assymetric - only dictates sender send APIs -# because we rely on the native vLLM WeightTransferEngine for the receive logic. -# For CUDA IPC, we use a custom send implementation and for NCCL, we rely on -# `nccl_trainer_send.py` -- vLLM 0.26's NCCLWeightTransferEngine send statics, -# vendored after 0.28 replaced them with a trainer-side engine abstraction. -class WeightTransferStrategy(ABC): - """Stateless factory for creating init info and senders. - - Each strategy implementation provides static methods to create: - - init_info: Contains all config-derived args - - sender: Uses init_info + inference_client - - Usage on sender side: - init_info = Strategy.create_init_info(ie_cfg, inference_world_size) - sender = Strategy.create_sender(init_info, inference_client) - - The receiver side lives inside the inference servers (vLLM's native weight - transfer engine), driven via the inference client's HTTP control plane. - """ - - sender_initializes_receivers: ClassVar[bool] = False - """The sender drives the inference-side init itself, so the worker must NOT - also call ``init_weight_update_communicator``. - - False for the push backends: worker rank 0 pushes ``init_info`` to the - servers, concurrently with ``create_sender`` (broadcast needs both sides in - the same process group at once). True for a backend whose own engine owns the - handshake.""" - - @staticmethod - @abstractmethod - def create_init_info( - ie_cfg: "InferenceEngineConfig", - inference_world_size: Optional[int] = None, - base_model_path: Optional[str] = None, - ) -> WeightSyncInitInfo: - """Create init info with all config-derived args. - - Args: - ie_cfg: Inference engine configuration. - inference_world_size: Total number of inference workers (from - ``client.get_world_size()``). Required by strategies that use it - (broadcast); strategies that don't (CUDA IPC) ignore it. - base_model_path: Policy model path. - - Returns: - WeightSyncInitInfo containing all args needed for sender/receiver creation. - """ - ... - - @staticmethod - @abstractmethod - def get_vllm_transfer_engine() -> type: - """Return the vLLM weight-transfer engine class for this strategy. - - Broadcast -> ``NCCLWeightTransferEngine``; CUDA IPC -> - ``IPCWeightTransferEngine``. This is the receive-side engine the - inference servers drive natively. Currently unused on the sender side - (we route through the SkyRL ``/collective_rpc`` wrap); kept as the - canonical strategy->engine mapping. - """ - ... - - @staticmethod - @abstractmethod - def create_sender( - init_info: WeightSyncInitInfo, - inference_client: "RemoteInferenceClient", - weight_extractor: Optional[Any] = None, - ) -> WeightTransferSender: - """Create a sender for the training worker side. - - This method must be called on all training ranks. Implementations may - have different initialization logic for different ranks (e.g., only rank 0 - joins a process group for broadcast, while all ranks participate for IPC). - - Args: - init_info: WeightSyncInitInfo containing config-derived args. - inference_client: Client for coordinating with inference engines. - weight_extractor: The worker's extractor. Only backends that - rendezvous at init rather than on the first send need it - (sharded_rdt); the others ignore it. - - Returns: - A configured WeightTransferSender instance. - """ - ... diff --git a/skyrl/backends/skyrl_train/weight_sync/weight_extractor.py b/skyrl/backends/skyrl_train/weight_sync/weight_extractor.py deleted file mode 100644 index 055a7903ce..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/weight_extractor.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Weight extractor interface for extracting weights from training backends.""" - -from abc import ABC, abstractmethod -from typing import Dict, Iterator, List - -import torch - -from .base import WeightChunk - - -class WeightExtractor(ABC): - """Extracts weights from training backend models. - - Subclasses implement backend-specific logic to extract model weights, - handle sharding, and prepare them for transfer to inference engines. - """ - - @abstractmethod - def extract_weights(self, dtype: torch.dtype) -> Iterator[WeightChunk]: - """Extract weights from the model as WeightChunk objects. - - Implementations should: - - Gather sharded weights into full tensors - - Convert tensors to the specified dtype for inference - - Ensure tensors are contiguous in memory - - Optionally group related parameters (e.g., QKV for efficiency) - - Args: - dtype: Target dtype for inference (e.g., torch.bfloat16, torch.float16) - - Yields: - WeightChunk objects containing model parameters ready for transfer - """ - ... - - @abstractmethod - def get_weight_metadata(self, dtype: torch.dtype) -> Dict[str, List]: - """Return weight metadata without materializing tensors. - - Args: - dtype: Target dtype for inference (used for dtype name). - - Returns: - Dict with keys "names", "dtype_names", "shapes". - """ - ... diff --git a/skyrl/backends/skyrl_train/weight_sync/weight_extractor_utils.py b/skyrl/backends/skyrl_train/weight_sync/weight_extractor_utils.py deleted file mode 100644 index 745f91bd81..0000000000 --- a/skyrl/backends/skyrl_train/weight_sync/weight_extractor_utils.py +++ /dev/null @@ -1,108 +0,0 @@ -"""Utility functions for weight extraction.""" - -from collections import defaultdict -from typing import Any, Callable, Dict, Iterator, List - -import torch - -from skyrl.backends.skyrl_train.weight_sync import WeightChunk - - -def yield_module_grouped_chunks( - params: Dict[str, Any], - dtype: torch.dtype, - gather_tensor_fn: Callable[[Any], torch.Tensor], - get_shape_fn: Callable[[str, Any, torch.Tensor], List[int]], - batch_size_threshold_gb: float = 0.0, -) -> Iterator[WeightChunk]: - """Yield WeightChunk objects grouped by module. - - This helper function eliminates duplication between different weight extractors - that need to group parameters by module (e.g., for fused QKV loaders). - - Groups parameters by their parent module by removing the last two components - from the parameter name. For example: - "model.layers.0.self_attn.q_proj.weight" -> "model.layers.0.self_attn" - - Args: - params: Dictionary mapping parameter names to parameter objects - dtype: Target dtype for inference - gather_tensor_fn: Backend-specific function to gather sharded tensors into full tensors - get_shape_fn: Function to extract shape from param_name, param, and prepared tensor - batch_size_threshold_gb: If > 0, batch complete modules together until threshold is reached - - Yields: - WeightChunk objects containing all parameters for each module (or batched modules if threshold set) - """ - # Group parameters by module for integrations that load fused QKV weights. - # NOTE (sumanthrh): We sync weights module by module. Ex: weights for self attn together, weights for mlp together - # We allocate new storage for each param. Since q, k and v layer weights are fused internally by vllm, - # we need to pass the weights for all of these together. - # Overall, this doesn't hurt perf even in the general case - module_to_params: Dict[str, List[str]] = defaultdict(list) - for param_name in params.keys(): - # Extract module name (e.g., "model.layers.0.self_attn" from "model.layers.0.self_attn.q_proj.weight") - # TODO (sumanthrh): When would this fail? Works for many AutoModelForCausalLM models for now - module_name = ".".join(param_name.split(".")[:-2]) - module_to_params[module_name].append(param_name) - - # Accumulate complete modules until threshold reached - batch_tensors = [] - batch_names = [] - batch_shapes = [] - batch_dtypes = [] - current_size = 0 - threshold_bytes = batch_size_threshold_gb * 1024**3 - - for module_name, param_names in module_to_params.items(): - module_tensors = [] - module_names = [] - module_shapes = [] - module_dtypes = [] - module_size = 0 - - # Prepare all tensors for this module - # TODO: Allow gather_tensor_fn to accept a list of params for batched gathering - # to improve efficiency for sharded backends that support multi-parameter collects. - for param_name in param_names: - param = params[param_name] - tensor = gather_tensor_fn(param) - tensor = tensor.to(dtype).detach().contiguous() - shape = get_shape_fn(param_name, param, tensor) - module_tensors.append(tensor) - module_names.append(param_name) - module_shapes.append(shape) - module_dtypes.append(str(dtype)) - module_size += tensor.nbytes - - # Check if adding this module would exceed threshold - if current_size > 0 and current_size + module_size > threshold_bytes: - # Yield current batch before adding this module - yield WeightChunk( - names=batch_names, - dtypes=batch_dtypes, - shapes=batch_shapes, - tensors=batch_tensors, - ) - # Start new batch - batch_tensors = [] - batch_names = [] - batch_shapes = [] - batch_dtypes = [] - current_size = 0 - - # Add module to current batch - batch_tensors.extend(module_tensors) - batch_names.extend(module_names) - batch_shapes.extend(module_shapes) - batch_dtypes.extend(module_dtypes) - current_size += module_size - - # Yield final batch if non-empty - if batch_tensors: - yield WeightChunk( - names=batch_names, - dtypes=batch_dtypes, - shapes=batch_shapes, - tensors=batch_tensors, - ) diff --git a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py index a74118d805..f325aa0122 100644 --- a/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py +++ b/skyrl/backends/skyrl_train/workers/fsdp/fsdp_worker.py @@ -1,18 +1,12 @@ import io import os -from typing import TYPE_CHECKING, Optional +from typing import Optional import ray import torch import torch.distributed from transformers import AutoConfig -try: - # for torch 2.5+ - from torch.distributed.tensor import DTensor -except ImportError: - from torch.distributed._tensor import DTensor - from skyrl.backends.skyrl_train.distributed.dispatch import WorkerOutput from skyrl.backends.skyrl_train.distributed.fsdp_strategy import FSDPStrategy from skyrl.backends.skyrl_train.distributed.fsdp_utils import ( @@ -25,14 +19,7 @@ TrainingInputBatch, ) from skyrl.backends.skyrl_train.utils.profiler import build_profiler_from_policy_cfg -from skyrl.backends.skyrl_train.weight_sync import ( - LoraLoadRequest, - WeightChunk, - WeightExtractor, -) -from skyrl.backends.skyrl_train.weight_sync.weight_extractor_utils import ( - yield_module_grouped_chunks, -) +from skyrl.backends.skyrl_train.weight_sync import LoraLoadRequest from skyrl.backends.skyrl_train.workers.model_wrapper import ( HFModelWrapper, get_llm_for_sequence_regression, @@ -42,96 +29,6 @@ PolicyWorkerBase, RefWorkerBase, ) -from skyrl.train.utils.utils import str_to_torch_dtype - -if TYPE_CHECKING: - from skyrl.train.config.config import InferenceEngineConfig - - -class FSDPWeightExtractor(WeightExtractor): - """Extracts weights from FSDP-sharded models. - - Args: - model: FSDP model to extract weights from - enable_bucketing: If True, group parameters by module (e.g., for fused QKV loaders) - batch_size_threshold_gb: If > 0, batch complete modules together until threshold is reached - weight_prefix: Prefix to prepend to all weight names (e.g., ``"language_model."`` - when syncing a CausalLM backbone to a vLLM instance which always uses the namespace of the - multimodal model, even if vision encoder weights are not initialized). - """ - - def __init__( - self, - model: torch.nn.Module, - enable_bucketing: bool = False, - batch_size_threshold_gb: float = 0.0, - weight_prefix: str = "", - ): - self.model = model - self.enable_bucketing = enable_bucketing - self.batch_size_threshold_gb = batch_size_threshold_gb - self.weight_prefix = weight_prefix - - def extract_weights(self, dtype: torch.dtype): - """Extract weights from FSDP model. - - Args: - dtype: Target dtype for inference - - Yields: - WeightChunk objects (one per parameter, or grouped by module) - """ - # FSDP2 state_dict returns DTensors directly; no state_dict_type configuration needed. - params = self.model.state_dict() - - if self.weight_prefix: - params = {f"{self.weight_prefix}{k}": v for k, v in params.items()} - - if not self.enable_bucketing: - # Simple path: yield one chunk per parameter - for name, param in params.items(): - tensor = self._gather_tensor(param).to(dtype).detach().contiguous() - yield WeightChunk( - names=[name], - dtypes=[str(dtype)], - shapes=[list(tensor.shape)], - tensors=[tensor], - ) - else: - # NOTE (sumanthrh): By default, when we bucket parameters with FSDP, we group by module - # This was done to support the older FlashRL integration where vLLM required - # Q,K, V tensors for the same layer to be sent at the same time - # Grouping by module is also beneficial because of layerwise reloading in vLLM - # We will be able to complete the reload for a full layer in one /update_weights - # call and clear the layerwise buffer held for that layer. - for chunk in yield_module_grouped_chunks( - params=params, - dtype=dtype, - gather_tensor_fn=self._gather_tensor, - get_shape_fn=lambda name, param, tensor: list(tensor.shape), - batch_size_threshold_gb=self.batch_size_threshold_gb, - ): - yield chunk - - def get_weight_metadata(self, dtype: torch.dtype) -> dict: - """Return weight metadata without materializing full tensors. - - Reads state_dict() shapes; sharded DTensors are not gathered. - """ - names = [] - dtype_names = [] - shapes = [] - dtype_name = str(dtype).split(".")[-1] - for name, param in self.model.state_dict().items(): - names.append(f"{self.weight_prefix}{name}" if self.weight_prefix else name) - dtype_names.append(dtype_name) - shapes.append(list(param.shape)) - return {"names": names, "dtype_names": dtype_names, "shapes": shapes} - - def _gather_tensor(self, param: torch.Tensor) -> torch.Tensor: - """Gather sharded tensor into full tensor.""" - device = torch.cuda.current_device() - return param.to(device, non_blocking=True).full_tensor() if isinstance(param, DTensor) else param class FSDPPolicyWorkerBase(PolicyWorkerBase): @@ -205,42 +102,27 @@ def init_model(self, model_path, num_training_steps: int = None): # Created only on profiled ranks. self.profiler = build_profiler_from_policy_cfg(self.cfg) - async def init_weight_sync_state(self, inference_engine_client, inference_engine_cfg: "InferenceEngineConfig"): - # Initialize the weight extractor BEFORE super(): a strategy that - # rendezvouses at init (sharded_rdt) is handed this extractor by - # create_sender. It only depends on - # the already-built model, not on super(). - # TODO(haochen): Module grouping for fused-weight loaders is only enabled for CUDA IPC. - # transfer strategy, we can enable it for other strategies as well. - from skyrl.backends.skyrl_train.weight_sync import ( - CudaIpcTransferStrategy, - get_transfer_strategy_cls, - ) + def _build_weight_source(self, dtype: "torch.dtype", backend: str): + """``WeightSource`` over the FSDP-sharded policy model. - # TODO (sumanthrh): bucketing can be enabled for other strategies as well - # Historically, this was only enabled for CUDA IPC in order to support - # Flash-RL - enable_bucketing = ( - get_transfer_strategy_cls( - weight_sync_backend=inference_engine_cfg.weight_sync_backend, - colocate_all=self.cfg.placement.colocate_all, - ) - is CudaIpcTransferStrategy - ) + ``self.model.model`` is the inner HF module, whose ``state_dict()`` keys + are the names vLLM expects; the outer wrapper's are prefixed differently. + ``weight_prefix`` covers the one case where they still differ: syncing a + CausalLM backbone into a vLLM multimodal namespace. + """ weight_prefix = "language_model." if self._is_multimodal_lm_only else "" - self.weight_extractor = FSDPWeightExtractor( - self.model.model, - enable_bucketing=enable_bucketing, - batch_size_threshold_gb=( - inference_engine_cfg.weight_transfer_threshold_cuda_ipc_GB if enable_bucketing else 0.0 - ), - weight_prefix=weight_prefix, - ) + if backend == "sharded_rdt": + # RDT pulls, so it needs the ownership + group channels its own + # source subclass adds (see sharded_rdt/sharded_rdt_base.py). + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( + make_fsdp_weight_source, + ) + + return make_fsdp_weight_source(self.model.model, dtype, weight_prefix) + + from skyrl.backends.skyrl_train.weight_sync.sources import FsdpWeightSource - # super picks the strategy and creates the sender (for sharded_rdt that - # includes the eager rendezvous + bake, which is why the extractor is - # built first). - await super().init_weight_sync_state(inference_engine_client, inference_engine_cfg) + return FsdpWeightSource(self.model.model, dtype, weight_prefix) async def _save_lora_adapters_and_sync( self, @@ -295,21 +177,6 @@ async def broadcast_to_inference_engines( ): if inference_engine_client is None: inference_engine_client = self._weight_sync_inference_client - use_prefix_cache = inference_engine_cfg.enable_prefix_caching - generator_dtype = str_to_torch_dtype(inference_engine_cfg.model_dtype) - cache_reset_task = None - sender_handles_prefix_cache_reset = self._weight_transfer_sender.handles_prefix_cache_reset - # Clear prefix cache for synchronous training or for async training if `clear_kv_cache_on_weight_sync` is set - reset_prefix_cache: bool = use_prefix_cache and ( - not self.cfg.fully_async.enabled or self.cfg.fully_async.clear_kv_cache_on_weight_sync - ) - send_chunks_kwargs = {"reset_prefix_cache": reset_prefix_cache} - - if reset_prefix_cache and torch.distributed.get_rank() == 0 and not sender_handles_prefix_cache_reset: - # clear prefix cache - cache_reset_task = inference_engine_client.reset_prefix_cache(reset_running_requests=True) - - torch.cuda.empty_cache() # Check if this is a LoRA model peft_model = getattr(self.model.model, "_fsdp_wrapped_module", self.model.model) @@ -322,30 +189,20 @@ async def broadcast_to_inference_engines( # name. _resolve_lora_sync_target (shared with Megatron, defined on # PolicyWorkerBase) basename-guards against a malformed model_id # escaping lora_sync_path even though api.py already validates IDs. + cache_reset_task = self._reset_prefix_cache_task(inference_engine_client, inference_engine_cfg) + torch.cuda.empty_cache() lora_name, lora_sync_path = self._resolve_lora_sync_target(model_id) await self._save_lora_adapters_and_sync( peft_model, lora_sync_path, inference_engine_client, lora_name=lora_name ) - else: - # Send with the sender created at init time. Disable expandable_segments - # around it: under colocate_all the CUDA-IPC path calls - # cudaIpcGetMemHandle, which is incompatible with the VMM addresses - # expandable segments uses, and some senders (sharded_rdt) share GPU - # memory on every run and ask for the toggle unconditionally. - with self._expandable_segments_disabled_for_sync( - force=self._weight_transfer_sender.force_disable_expandable_segments - ): - await self._weight_transfer_sender.send( - self.weight_extractor, - generator_dtype, - **send_chunks_kwargs, - ) - - if cache_reset_task is not None: - await cache_reset_task - if self._weight_transfer_sender.empty_cache_after_send or self.cfg.placement.colocate_all: - torch.cuda.empty_cache() - torch.distributed.barrier() + if cache_reset_task is not None: + await cache_reset_task + if self.cfg.placement.colocate_all: + torch.cuda.empty_cache() + torch.distributed.barrier() + return + + await self._sync_weights_to_inference_engines(inference_engine_client, inference_engine_cfg) def _set_pad_token_id(self, pad_token_id): # NOTE (sumanthrh): self.model -> HFModelWrapper; self.model.model -> AutoModelForCausalLM diff --git a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py index d513f7fa5e..edb5d2da8b 100644 --- a/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py +++ b/skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py @@ -26,7 +26,6 @@ from skyrl.backends.skyrl_train.distributed.megatron.megatron_utils import ( _clear_mtp_hybrid_pattern, _convert_moe_experts_lora_to_vllm, - broadcast_object_across_pp_ranks, freeze_moe_router, gdn_in_proj_lora_is_safe, get_model_config, @@ -53,11 +52,7 @@ ) from skyrl.backends.skyrl_train.utils.profiler import build_profiler_from_policy_cfg from skyrl.backends.skyrl_train.utils.replay_utils import make_replay_padding_indices -from skyrl.backends.skyrl_train.weight_sync import ( - LoraLoadRequest, - WeightChunk, - WeightExtractor, -) +from skyrl.backends.skyrl_train.weight_sync import LoraLoadRequest from skyrl.backends.skyrl_train.workers.megatron.adapter_store import ( AdapterStore, LoraSignature, @@ -81,7 +76,7 @@ ) from skyrl.env_vars import SKYRL_WORKER_NCCL_TIMEOUT_IN_S from skyrl.train.config.config import MegatronDDPConfig, get_config_as_dict -from skyrl.train.utils.utils import str_to_torch_dtype, update_model_config +from skyrl.train.utils.utils import update_model_config from skyrl.utils.tok import get_tokenizer if TYPE_CHECKING: @@ -96,253 +91,6 @@ ) -class MegatronWeightExtractor(WeightExtractor): - """Extracts weights from Megatron model-parallel models. - - Uses Megatron's bridge to export weights in HuggingFace format. - - Args: - bridge: Megatron AutoBridge instance for weight conversion - actor_module: The actor module to extract weights from - enable_bucketing: If True, group parameters into size-based buckets for packing - bucket_size_threshold_GB: Size threshold in GB for bucketing (only used if enable_bucketing=True) - training_dtype: Training dtype for size calculation (only used if enable_bucketing=True) - """ - - def __init__( - self, - bridge, - actor_module, - enable_bucketing: bool = False, - bucket_size_threshold_GB: float = 1.0, - training_dtype: torch.dtype = torch.bfloat16, - ): - self.bridge = bridge - self.actor_module = actor_module - self.enable_bucketing = enable_bucketing - self.bucket_size_threshold_GB = bucket_size_threshold_GB - self.training_dtype = training_dtype - - # Defer bucket init to first extract_weights call. - # At __init__ time the model may be CPU-offloaded (colocate_all), - # so param.numel()==0 and bucketing collapses to a single bucket. - # By the time extract_weights runs, the dispatch has already - # called prepare_for_weight_sync → _ensure_on_gpu. - self.bucket_index_groups = None - self._buckets_initialized = False - - def _init_param_buckets(self): - """Compute bucket boundaries (index groups) from parameter sizes. - - Only the bucket *structure* (which task indices go in which bucket) is - persisted. The actual ``WeightConversionTask`` objects are rebuilt on - every ``extract_weights`` call so that mapping objects start with clean - PP-collective caches, avoiding stale cached state across offload/reload - and training cycles. - - Tasks that participate in grouped export (e.g., fused MoE expert - weights) are collected first and placed into dedicated buckets so that - all tasks sharing the same ``group_key`` end up in a single - ``export_hf_weights`` call. The bridge's - ``_accumulate_grouped_export`` requires every task for a group to be - present in one call; splitting them across buckets causes expert - weights to never be yielded. - """ - weight_conversion_tasks = self.bridge.get_conversion_tasks(self.actor_module) - - def calculate_size_in_bytes(param, tp_size, ep_size): - if param is None: - size_in_bytes = None - else: - prec_to_bytes = { - torch.bfloat16: 2, - torch.float32: 4, - } - scale = prec_to_bytes[self.training_dtype] / prec_to_bytes[param.dtype] - size_in_bytes = param.element_size() * param.numel() * tp_size * ep_size * scale - # allow_missing: a task may correspond to no parameter on any PP rank - # (see the layout note below), in which case there is no size to agree on. - return broadcast_object_across_pp_ranks(size_in_bytes, allow_missing=True) - - sizes = [ - calculate_size_in_bytes( - task.param_weight, - task.mapping.tp_size, - task.mapping.ep_size if task.mapping.is_expert else 1, - ) - for task in weight_conversion_tasks - ] - - # ---- Separate grouped-export tasks from regular tasks ---- - # Grouped-export tasks (is_grouped_export=True, e.g. FusedGatedExpertMapping / - # FusedExpertMapping for MoE expert weights) must ALL be present in a single - # export_hf_weights call for the bridge's _accumulate_grouped_export to produce - # the fused tensor. Collect them by group_key and give each group its own bucket. - grouped_task_indices: dict[str, list[int]] = {} # group_key -> list of task indices - regular_task_indices: list[int] = [] - - for idx, task in enumerate(weight_conversion_tasks): - # Skip tasks that own no parameter on any PP rank. megatron-bridge can - # register mappings for BOTH MoE expert layouts -- grouped-GEMM - # (`mlp.experts.linear_fc1`) and SequentialMLP - # (`mlp.experts.local_experts.*.linear_fc1`) -- so a model built with one - # layout still gets conversion tasks for the other. Those have no weights - # to export, and including them would break bucket-size accounting. - if sizes[idx] is None: - continue - if getattr(task.mapping, "is_grouped_export", False): - gk = getattr(task.mapping, "group_key", None) - grouped_task_indices.setdefault(gk, []).append(idx) - else: - regular_task_indices.append(idx) - - self.bucket_index_groups: list[list[int]] = [] - - # Pack grouped-export tasks into buckets by size, keeping each - # group_key's tasks together (they must not be split across calls). - curr_size = 0 - threshold = self.bucket_size_threshold_GB * 1024**3 - for gk, indices in grouped_task_indices.items(): - group_size = sum(sizes[idx] for idx in indices if sizes[idx] is not None) - if not self.bucket_index_groups or curr_size + group_size > threshold: - self.bucket_index_groups.append([]) - curr_size = 0 - self.bucket_index_groups[-1].extend(indices) - curr_size += group_size - - # Bucket regular (non-grouped) tasks by size as before. - if regular_task_indices: - self.bucket_index_groups.append([]) - curr_size = 0 - for idx in regular_task_indices: - size = sizes[idx] - if curr_size + size > threshold: - self.bucket_index_groups.append([]) - curr_size = 0 - self.bucket_index_groups[-1].append(idx) - curr_size += size - - def get_weight_metadata(self, dtype: torch.dtype) -> dict: - """Return weight metadata without keeping tensors in memory. - - On first call, runs export_hf_weights to discover HF names and shapes - (tensors are discarded immediately). Result is cached for subsequent calls. - TODO (aaron): find a better way to get all metadata without materializing tensors. - """ - if hasattr(self, "_weight_metadata_cache"): - return self._weight_metadata_cache - - self._ensure_buckets_initialized() - names = [] - dtype_names = [] - shapes = [] - dtype_name = str(dtype).split(".")[-1] - # Collect parameter metadata in the same order - # as provided by `.extract_weights`. - if not self.enable_bucketing: - for name, tensor in self.bridge.export_hf_weights( - self.actor_module, - show_progress=False, - conversion_tasks=None, - ): - names.append(name) - dtype_names.append(dtype_name) - shapes.append(list(tensor.shape)) - del tensor - else: - # Build fresh tasks each sync so mapping objects have clean - # PP-collective caches; reuse the pre-computed bucket structure. - fresh_tasks = self.bridge.get_conversion_tasks(self.actor_module) - for index_group in self.bucket_index_groups: - bucket_tasks = [fresh_tasks[i] for i in index_group] - for name, tensor in self.bridge.export_hf_weights( - self.actor_module, - show_progress=False, - conversion_tasks=bucket_tasks, - ): - names.append(name) - shapes.append(list(tensor.shape)) - dtype_names.append(dtype_name) - del tensor - - self._weight_metadata_cache = {"names": names, "dtype_names": dtype_names, "shapes": shapes} - return self._weight_metadata_cache - - def _ensure_buckets_initialized(self): - """Lazily initialize param buckets on first use (model must be on GPU).""" - if self._buckets_initialized: - return - if self.enable_bucketing: - self._init_param_buckets() - self._buckets_initialized = True - - def extract_weights(self, dtype: torch.dtype): - """Extract weights from Megatron model. - - Args: - dtype: Target dtype for inference - - Yields: - WeightChunk objects (one per parameter, or one per bucket if bucketing enabled) - """ - self._ensure_buckets_initialized() - device = torch.cuda.current_device() - - if not self.enable_bucketing: - # No bucketing: yield one chunk per parameter - hf_params_generator = self.bridge.export_hf_weights( - self.actor_module, - show_progress=False, - conversion_tasks=None, - ) - - for name, tensor in hf_params_generator: - tensor = tensor.to(device=device, dtype=dtype, non_blocking=True) - - yield WeightChunk( - names=[name], - dtypes=[str(dtype)], - shapes=[list(tensor.shape)], - tensors=[tensor], - ) - else: - # Build fresh tasks each sync so mapping objects have clean - # PP-collective caches; reuse the pre-computed bucket structure. - fresh_tasks = self.bridge.get_conversion_tasks(self.actor_module) - - for index_group in self.bucket_index_groups: - bucket_tasks = [fresh_tasks[i] for i in index_group] - hf_params_generator = self.bridge.export_hf_weights( - self.actor_module, - show_progress=False, - conversion_tasks=bucket_tasks, - ) - - # Collect all parameters in this bucket into one chunk - names = [] - dtypes_list = [] - shapes = [] - tensors = [] - - for name, tensor in hf_params_generator: - # Move to device and convert dtype - tensor = tensor.to(device=device, dtype=dtype, non_blocking=True) - - names.append(name) - dtypes_list.append(str(dtype)) - shapes.append(list(tensor.shape)) - tensors.append(tensor) - - # Yield one chunk containing all parameters in this bucket - if tensors: - yield WeightChunk( - names=names, - dtypes=dtypes_list, - shapes=shapes, - tensors=tensors, - ) - - class MegatronWorker: def _maybe_setup_fake_int4_qat(self): """Wire up INT4-served training and return the BF16 bridge-weights path. @@ -1438,23 +1186,20 @@ def set_lr(self, learning_rate: float) -> None: for param_group in self.optimizer.param_groups: param_group["lr"] = learning_rate - async def init_weight_sync_state(self, inference_engine_client, inference_engine_cfg: "InferenceEngineConfig"): - # Initialize the weight extractor BEFORE super(): a strategy that - # rendezvouses at init (sharded_rdt) is handed this extractor by - # create_sender. It only depends on - # the already-built bridge/actor_module, not on super(). - self.weight_extractor = MegatronWeightExtractor( - bridge=self.bridge, - actor_module=self.actor_module, - enable_bucketing=True, - bucket_size_threshold_GB=inference_engine_cfg.weight_transfer_threshold_cuda_ipc_GB, - training_dtype=torch.bfloat16 if self.cfg.bf16 else torch.float32, - ) + def _build_weight_source(self, dtype: "torch.dtype", backend: str): + """``WeightSource`` over the Megatron policy model, via Megatron-Bridge.""" + if backend == "sharded_rdt": + # RDT pulls, so it needs the ownership + group channels its own + # source subclasses add, and can serve PP/EP-local exports. + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( + make_megatron_weight_source, + ) - # super picks the strategy and creates the sender (for sharded_rdt that - # includes the eager rendezvous + bake, which is why the extractor is - # built first). - await super().init_weight_sync_state(inference_engine_client, inference_engine_cfg) + return make_megatron_weight_source(self.bridge, self.actor_module, dtype) + + from skyrl.backends.skyrl_train.weight_sync.sources import MegatronWeightSource + + return MegatronWeightSource(self.bridge, self.actor_module, dtype) async def _save_lora_adapters_and_sync( self, lora_sync_path, inference_engine_client, lora_name: str = SKYRL_LORA_ADAPTER_NAME @@ -1525,53 +1270,24 @@ async def broadcast_to_inference_engines( ): if inference_engine_client is None: inference_engine_client = self._weight_sync_inference_client - use_prefix_cache = inference_engine_cfg.enable_prefix_caching - generator_dtype = str_to_torch_dtype(inference_engine_cfg.model_dtype) - cache_reset_task = None - sender_handles_prefix_cache_reset = self._weight_transfer_sender.handles_prefix_cache_reset - # Clear prefix cache for synchronous training or for async training if `clear_kv_cache_on_weight_sync` is set - reset_prefix_cache: bool = use_prefix_cache and ( - not self.cfg.fully_async.enabled or self.cfg.fully_async.clear_kv_cache_on_weight_sync - ) - send_chunks_kwargs = {"reset_prefix_cache": reset_prefix_cache} - - if reset_prefix_cache and torch.distributed.get_rank() == 0 and not sender_handles_prefix_cache_reset: - # clear prefix cache - cache_reset_task = inference_engine_client.reset_prefix_cache(reset_running_requests=True) - - torch.cuda.empty_cache() if self._is_lora and not self.cfg.policy.megatron_config.lora_config.merge_lora: # AdapterStore.swap_to has already made `model_id` the live adapter # before we get here; sync that adapter to vLLM under its own name # so sample(model=) routes correctly. Single-tenant # (model_id=None) keeps the legacy shared path + name. + cache_reset_task = self._reset_prefix_cache_task(inference_engine_client, inference_engine_cfg) + torch.cuda.empty_cache() lora_name, lora_sync_path = self._resolve_lora_sync_target(model_id) await self._save_lora_adapters_and_sync(lora_sync_path, inference_engine_client, lora_name=lora_name) - else: - # Send with the sender created at init time. Disable expandable_segments - # around it: under colocate_all the CUDA-IPC path calls - # cudaIpcGetMemHandle, which is incompatible with the VMM addresses - # expandable segments uses, and some senders (sharded_rdt) share GPU - # memory on every run and ask for the toggle unconditionally. - with self._expandable_segments_disabled_for_sync( - force=self._weight_transfer_sender.force_disable_expandable_segments - ): - await self._weight_transfer_sender.send( - self.weight_extractor, - generator_dtype, - **send_chunks_kwargs, - ) + if cache_reset_task is not None: + await cache_reset_task + if self.cfg.placement.colocate_all: + torch.cuda.empty_cache() + torch.distributed.barrier() + return - if cache_reset_task is not None: - await cache_reset_task - # A sender whose send buffers are reused next step (sharded_rdt) declares - # empty_cache_after_send=False: scrubbing them back to CUDA costs 0.25-0.53s - # per rank at 235B and buys nothing. Under colocation the physical memory is - # wanted by an inference engine, so empty regardless. - if self._weight_transfer_sender.empty_cache_after_send or self.cfg.placement.colocate_all: - torch.cuda.empty_cache() - torch.distributed.barrier() + await self._sync_weights_to_inference_engines(inference_engine_client, inference_engine_cfg) def _set_pad_token_id(self, pad_token_id): # this already gets set in the init_model method diff --git a/skyrl/backends/skyrl_train/workers/worker.py b/skyrl/backends/skyrl_train/workers/worker.py index 5d3275ebb8..8c693221cb 100644 --- a/skyrl/backends/skyrl_train/workers/worker.py +++ b/skyrl/backends/skyrl_train/workers/worker.py @@ -479,15 +479,21 @@ async def init_weight_sync_state( inference_engine_client: "RemoteInferenceClient", inference_engine_cfg: "InferenceEngineConfig", ): - """Initialize state for weight syncing with Inference Engine Client + """Build this rank's weight source and rendezvous its trainer engine. - Creates init info and sender, then sends init info to inference engines - so they can create receivers. + After this, ``broadcast_to_inference_engines`` is one + ``engine.send_weights()`` call per sync — the engine owns the whole round + trip (start / update / finish, the barriers, and the non-sender + collective replay). .. note:: - This function should be called on all the ranks in the worker group simultaneously. + This function must be called on all the ranks in the worker group + simultaneously: rank 0 blocks driving the inference-side handshake + while the others reach the collectives it needs. """ - from skyrl.backends.skyrl_train.weight_sync import get_transfer_strategy_cls + from skyrl.backends.skyrl_train.weight_sync.trainer_engines import ( + build_trainer_engine, + ) assert inference_engine_client is not None # Cache the client so per-sync broadcast_to_inference_engines calls can @@ -500,47 +506,28 @@ async def init_weight_sync_state( # Fetch the total inference world size from the servers. inference_world_size, _ = await inference_engine_client.get_world_size() - # Determine transfer strategy based on inference engine config and placement - self._transfer_strategy_cls = get_transfer_strategy_cls( - weight_sync_backend=inference_engine_cfg.weight_sync_backend, + # Off the event loop, on every rank: rank 0 drives the inference-side + # handshake with blocking HTTP while the others return immediately. + # + # Rendezvous happens HERE, not on the first send. sharded_rdt requires + # that: deferring it deadlocks, because rank 0 would block in the + # inference-side init RPC while the other ranks spin in gather + # collectives, and the producer sidecars -- sharing those ranks' GPUs -- + # could then never finish NIXL agent creation, since libfabric's CUDA + # probe blocks behind the spinning kernels. Every rank is inside this + # method here, so that window is empty. + self._weight_sync_engine = await self._weight_sync_thread( + build_trainer_engine, + ie_cfg=inference_engine_cfg, colocate_all=self.cfg.placement.colocate_all, - ) - - # Create init info on all ranks (it's deterministic from cfg or fetched world_size) - init_info = self._transfer_strategy_cls.create_init_info( - inference_engine_cfg, + rank=torch.distributed.get_rank(), inference_world_size=inference_world_size, + source_factory=self._build_weight_source, + server_urls=list(inference_engine_client.server_urls), + data_parallel_size=int(inference_engine_client.data_parallel_size), base_model_path=self.cfg.policy.model.path, ) - # Create sender on all ranks - # Strategy implementations may have different logic for different ranks - # The extractor is passed to every strategy; only those that rendezvous at - # init rather than on the first send use it (sharded_rdt). Both workers build - # it before calling super(), so it is available here. - tasks = [ - asyncio.to_thread( - self._transfer_strategy_cls.create_sender, - init_info=init_info, - inference_client=inference_engine_client, - weight_extractor=getattr(self, "weight_extractor", None), - ), - ] - - # Only rank 0 initializes receivers on inference engines - # NOTE: For broadcast strategy, sender and receiver init must run concurrently - # because both need to join the same process group to avoid deadlock - # NOTE: strategies whose sender drives the inference-side handshake itself - # (sharded_rdt) must NOT be initialized from here as well. - # TODO (sumanthrh): `sender_initializes_receivers` is currently used as a workaround for - # supporting RDT. We can probably move the inference-side init - # as a sender method in all the classes to unify this. RDT doesn't call `init_weight_update_communicator` itself - if torch.distributed.get_rank() == 0 and not self._transfer_strategy_cls.sender_initializes_receivers: - tasks.append(inference_engine_client.init_weight_update_communicator(init_info)) - - results = await asyncio.gather(*tasks) - self._weight_transfer_sender = results[0] # sender is always first task - # # Register signal handlers for termination only on rank 0 # NOTE (sumanthrh): This doesn't work yet, and is thus commented out. # The better way is to just have this specified in __del__, but there is @@ -551,6 +538,107 @@ async def init_weight_sync_state( torch.distributed.barrier() + def _build_weight_source(self, dtype: "torch.dtype", backend: str) -> Any: + """Build this worker's ``WeightSource`` over the live model. + + Implemented per training backend, and passed to ``build_trainer_engine`` + as its ``source_factory``. ``backend`` only picks between the plain + source and sharded RDT's ownership-aware subclass. + """ + raise NotImplementedError() + + def _weight_sync_thread(self, fn, *args, **kwargs): + """Run ``fn`` off the event loop with **this rank's** CUDA device selected. + + Every weight-sync entry point is fully synchronous -- it blocks on + collectives and, on rank 0, on blocking HTTP -- so it must not run on the + worker's event loop. But the current CUDA device is thread-local and + defaults to 0, and under + ``RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES`` this rank's device is not + 0, so the NCCL communicator, the IPC handles and the gather collectives + would all land on the wrong GPU. Capture it on the main thread, where it + is correct, and select it inside the thread. + """ + device = torch.cuda.current_device() if torch.cuda.is_available() else None + + def run(): + if device is not None: + torch.cuda.set_device(device) + return fn(*args, **kwargs) + + return asyncio.to_thread(run) + + def _should_reset_prefix_cache(self, inference_engine_cfg) -> bool: + """Whether this sync should clear the inference engines' prefix cache. + + Always for synchronous training; for fully-async only when + ``clear_kv_cache_on_weight_sync`` is set (otherwise in-flight requests + keep generating against their cached prefixes across the sync). + """ + return inference_engine_cfg.enable_prefix_caching and ( + not self.cfg.fully_async.enabled or self.cfg.fully_async.clear_kv_cache_on_weight_sync + ) + + def _reset_prefix_cache_task(self, inference_engine_client, inference_engine_cfg): + """Fire the prefix-cache reset from rank 0, or return None. + + Returned rather than awaited so the caller can overlap it with the send. + """ + if not self._should_reset_prefix_cache(inference_engine_cfg): + return None + if torch.distributed.get_rank() != 0: + return None + return inference_engine_client.reset_prefix_cache(reset_running_requests=True) + + async def _sync_weights_to_inference_engines( + self, + inference_engine_client, + inference_engine_cfg, + ) -> None: + """Run one weight sync: the engine's round trip, plus the worker-side + memory bracket around it. + + Three things the engine cannot do for itself, each decided by a + capability probe on it (see ``trainer_engines.engine_capability``): + + * the prefix-cache reset, unless the engine handles it internally + (checkpoint-delta does, inside its own pause bracket); + * the ``expandable_segments`` toggle, a trainer-process allocator setting; + * ``empty_cache`` afterwards. + """ + from skyrl.backends.skyrl_train.weight_sync.trainer_engines import ( + engine_capability, + maybe_set_reset_prefix_cache, + ) + + engine = self._weight_sync_engine + maybe_set_reset_prefix_cache(engine, self._should_reset_prefix_cache(inference_engine_cfg)) + + cache_reset_task = None + if not engine_capability(engine, "handles_prefix_cache_reset", False): + cache_reset_task = self._reset_prefix_cache_task(inference_engine_client, inference_engine_cfg) + + torch.cuda.empty_cache() + + # Under colocate_all the CUDA-IPC path calls cudaIpcGetMemHandle, which + # is incompatible with the VMM addresses expandable segments use; some + # engines (sharded_rdt) share GPU memory on every run and ask for the + # toggle unconditionally. + with self._expandable_segments_disabled_for_sync( + force=engine_capability(engine, "force_disable_expandable_segments", False) + ): + await self._weight_sync_thread(engine.send_weights) + + if cache_reset_task is not None: + await cache_reset_task + # An engine whose send buffers are reused next step declares + # empty_cache_after_send=False: scrubbing them back to CUDA costs + # 0.25-0.53s per rank at 235B and buys nothing. Under colocation an + # inference engine wants the physical memory, so empty regardless. + if engine_capability(engine, "empty_cache_after_send", True) or self.cfg.placement.colocate_all: + torch.cuda.empty_cache() + torch.distributed.barrier() + def forward(self, *args, **kwargs) -> WorkerOutput: """Run forward pass on the input batch. diff --git a/tests/backends/skyrl_train/distributed/test_worker_dispatch.py b/tests/backends/skyrl_train/distributed/test_worker_dispatch.py index 6d862d973d..c8e7ae0612 100644 --- a/tests/backends/skyrl_train/distributed/test_worker_dispatch.py +++ b/tests/backends/skyrl_train/distributed/test_worker_dispatch.py @@ -60,7 +60,7 @@ async def test_non_colocated_calls_pause_and_resume(self): async def test_non_colocated_delta_does_not_pause(self): """Delta sync owns pause/resume itself. - ``DeltaWeightTransferSender._apply_receiver_update`` fetches before pausing and + ``DeltaTrainerWeightTransferEngine._apply_receiver_update`` fetches before pausing and pauses only around the final reload, so the dispatcher must not pause as well -- doing so would hold generation down across the whole publish+upload+fetch window instead of just the reload. diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py b/tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py index 9c54895a34..1b99e0a37d 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py @@ -1,27 +1,29 @@ """ GPU CI tests for weight synchronization from trainer to inference server. -1. Non-colocated (NCCL broadcast), TP=2: - - Trainer on GPUs 0-1, server (TP=2) on GPUs 2-3 (4 GPUs total) - - Uses NCCL broadcast for weight sync via HTTP router - -2. Colocated (CUDA IPC), TP=1: - - Trainer and server share GPU 0 (2 GPUs total, 1 shared) - - Uses CUDA IPC handles for zero-copy weight transfer - -3. Legacy `WorkerWrap.load_weights` MoE reload, TP=1: - - Server on GPU 0 (1 GPU total, no separate trainer process) - - The NCCL or CUDA-IPC receiver is stubbed with safetensors-from-disk - to skip trainer-side sender setup +Each case drives the production trainer-side path: ``build_trainer_engine`` picks +the init info and builds the client, then ``engine.send_weights()`` owns the round +trip over vLLM's native RLHF routes. The trainer actor stands in for +FSDP/Megatron only -- it holds a plain HF model on one GPU, which +``FsdpWeightSource`` handles unchanged. + +1. Non-colocated NCCL broadcast, TP=2, plus a 1P1D PD variant. Covers + ``NCCLTrainerWeightTransferEngine`` against ``skyrl_nccl``, and the per-server + ``rank_offset`` rewrite in ``nccl_init_payloads``. +2. Colocated CUDA IPC, TP=1. Covers ``IPCTrainerWeightTransferEngine`` (packed) + against ``skyrl_ipc``, and the ``nccl`` + ``colocate_all`` -> ``ipc`` + resolution. +3. Non-colocated sharded RDT (NIXL pull), TP=1. Covers + ``ShardedRDTTrainerWeightTransferEngine``, the ownership-aware source, and the + ``replica_rank`` rewrite in ``rdt_init_payloads``. Run: uv run --extra dev --extra fsdp pytest tests/backends/skyrl_train/gpu/gpu_ci/inference_servers/test_weight_sync.py -v -s """ import asyncio -import base64 import os -import pickle +from types import SimpleNamespace import httpx import pytest @@ -31,96 +33,172 @@ from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy from transformers import AutoModelForCausalLM -from skyrl.backends.skyrl_train.inference_servers.common import ( - get_node_ip, - get_open_port, -) -from skyrl.backends.skyrl_train.weight_sync import ( - BroadcastInitInfo, - CudaIpcInitInfo, -) from skyrl.train.config import SkyRLTrainConfig from tests.backends.skyrl_train.gpu.utils import InferenceEngineState MODEL = os.environ.get("SKYRL_RDT_TEST_MODEL", "Qwen/Qwen2.5-0.5B-Instruct") +PROMPT = { + "model": MODEL, + "prompt": "What is the capital of France?", + "max_tokens": 32, + "temperature": 0.0, +} -@ray.remote -class Trainer: - """ - Simple trainer emulator that holds the real model weights. - This is a simplified version of the trainer side for testing weight sync - via NCCL broadcast in non-colocated scenarios. +class WeightSyncTrainerBase: + """Single-GPU stand-in for a training worker, driving the real engine path. + + No ``torch.distributed`` group, which is the rank-0-only case: the engines + resolve ``is_sender`` from ``init_info.rank``, and the IPC handle all-gather + and delta result gather both no-op without a group. + + Not ``@ray.remote`` itself -- each backend needs different Ray resources, so + the decorators are applied per backend below. """ - def __init__(self, model_name: str, device: str = "cuda"): - self.device = torch.device(device) - self.model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype=torch.bfloat16, - ).to(self.device) - self.pg = None - self.model_name = model_name + def __init__( + self, model_name, weight_sync_backend, colocate_all, server_urls, data_parallel_size, inference_world_size + ): + self._model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to("cuda") + # The two config values the backend is resolved from, so this exercises + # the same resolution the driver uses to configure the servers. + self._ie_cfg = SimpleNamespace(weight_sync_backend=weight_sync_backend, model_dtype="bfloat16") + self._colocate_all = colocate_all + self._server_urls = list(server_urls) + self._data_parallel_size = int(data_parallel_size) + self._inference_world_size = int(inference_world_size) + self._model_name = model_name + self._engine = None def ready(self): - """Check if the trainer is ready.""" return True - def init_weight_sync(self, master_address: str, master_port: int, world_size: int, group_name: str): - """Initialize the weight sync process group as rank 0 (trainer).""" - from skyrl.backends.skyrl_train.weight_sync.nccl_trainer_send import ( - nccl_trainer_init, + def _build_source(self, dtype, backend): + """The ``source_factory`` build_trainer_engine calls once the backend is known.""" + if backend == "sharded_rdt": + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( + make_fsdp_weight_source, + ) + + return make_fsdp_weight_source(self._model, dtype) + + from skyrl.backends.skyrl_train.weight_sync.sources import FsdpWeightSource + + return FsdpWeightSource(self._model, dtype) + + def sync_once(self): + """Rendezvous on the first call, then run one full weight sync.""" + from skyrl.backends.skyrl_train.weight_sync.trainer_engines import ( + build_trainer_engine, ) - self.pg = nccl_trainer_init( - dict( - master_address=master_address, - master_port=master_port, - world_size=world_size, + if self._engine is None: + self._engine = build_trainer_engine( + ie_cfg=self._ie_cfg, + colocate_all=self._colocate_all, + rank=0, + inference_world_size=self._inference_world_size, + source_factory=self._build_source, + server_urls=self._server_urls, + data_parallel_size=self._data_parallel_size, + base_model_path=self._model_name, + ) + self._engine.send_weights() + + def shutdown(self): + if self._engine is not None: + from skyrl.backends.skyrl_train.weight_sync.trainer_engines import ( + teardown_engine, ) - ) - return True - def get_weight_info(self) -> dict: - """ - Get weight metadata (names, dtypes, shapes) without doing NCCL. + teardown_engine(self._engine) + self._engine = None - Returns: - dict with names, dtypes, shapes for the weight update request. - """ - names = [] - dtypes = [] - shapes = [] - for name, param in self.model.named_parameters(): - names.append(name) - dtypes.append(str(param.dtype).split(".")[-1]) # e.g. "bfloat16" - shapes.append(list(param.shape)) +# A whole GPU for NCCL and RDT; RDT additionally needs it so its engine can pin +# the producer sidecar it spawns to this GPU for CUDA IPC. +NcclTrainer = ray.remote(num_gpus=1)(WeightSyncTrainerBase) +RdtTrainer = ray.remote(num_gpus=1, max_concurrency=4)(WeightSyncTrainerBase) +# IPC shares a GPU with the colocated server, so a fraction. +IpcTrainer = ray.remote(WeightSyncTrainerBase) - return {"names": names, "dtypes": dtypes, "shapes": shapes} - def broadcast_weights(self): - """ - Broadcast all model weights to inference workers via NCCL. +async def _completion(http_client, router_url): + resp = await http_client.post(f"{router_url}/v1/completions", json=PROMPT) + assert resp.status_code == 200 + return resp.json()["choices"][0]["text"] - This is a blocking operation - server must call receive concurrently. - """ - from skyrl.backends.skyrl_train.weight_sync.nccl_trainer_send import ( - nccl_trainer_send_weights, - ) - params = list(self.model.named_parameters()) - print( - f"[Trainer.broadcast_weights] Starting send of {len(params)} params, pg={self.pg}, pg.rank={self.pg.rank}, pg.world_size={self.pg.world_size}" +async def _assert_sync_replaces_dummy_weights(env, timeout_s: float = 120.0): + """Dummy weights -> sync -> real weights. + + ``load_format="dummy"`` starts the server with garbage, so "Paris" appearing + proves the transfer landed *and* that ``process_weights_after_loading`` ran + (an unfinalized layerwise reload produces garbage too). + """ + router_url = env["router_url"] + trainer = env["trainer"] + + async with httpx.AsyncClient(timeout=httpx.Timeout(timeout_s)) as http_client: + text_before = await _completion(http_client, router_url) + print(f"[step 1] dummy weights output: {text_before!r}") + assert "Paris" not in text_before, "Dummy weights unexpectedly produced the correct answer" + + print("[step 2] trainer.sync_once() -- rendezvous + one full send_weights()") + await asyncio.to_thread(lambda: ray.get(trainer.sync_once.remote())) + + text_after = await _completion(http_client, router_url) + print(f"[step 3] synced weights output: {text_after!r}") + assert "Paris" in text_after, f"Weight sync failed - expected 'Paris' but got: {text_after!r}" + + +async def _make_env(cfg, create_kwargs, trainer_cls, weight_sync_backend, *, colocate_with_engine=False): + """Bring up the servers plus a trainer actor, and yield the pair.""" + async with InferenceEngineState.create(cfg, **create_kwargs) as engines: + client = engines.client + inference_world_size, _ = await client.get_world_size() + options = {} + if colocate_with_engine: + options = dict( + num_gpus=0.2, + num_cpus=0.2, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=engines.pg, + placement_group_bundle_index=0, + ), + ) + trainer = trainer_cls.options(**options).remote( + MODEL, + weight_sync_backend, + # From create_kwargs, not cfg: InferenceEngineState.create deep-copies + # cfg before applying the override, so the outer cfg still has the + # default. This must be the value the servers were built with. + create_kwargs["colocate_all"], + client.server_urls, + client.data_parallel_size, + inference_world_size, ) - try: - nccl_trainer_send_weights(iter(params), self.pg, packed=True) - torch.cuda.synchronize() - print("[Trainer.broadcast_weights] Send complete") - except Exception as e: - print(f"[Trainer.broadcast_weights] ERROR: {e}") - raise + ray.get(trainer.ready.remote()) + + yield { + "engines": engines, + "trainer": trainer, + "client": client, + "router_url": client.proxy_url, + } + + ray.get(trainer.shutdown.remote()) + await client.teardown() + ray.kill(trainer) + # cleanup manually in colocated case + if engines.pg: + ray.util.remove_placement_group(engines.pg) + + +# ----------------------------------------------------------------- +# Non-colocated NCCL broadcast +# ----------------------------------------------------------------- @pytest_asyncio.fixture( @@ -140,14 +218,15 @@ async def weight_update_env(class_scoped_ray_init_fixture, request): - no_pd: TP=2 server on its own GPUs, trainer on separate GPU(s) (4 GPUs). - pd_1P1D_non_colocated: 1P1D (2 engines, TP=1), trainer on separate GPU (3 GPUs). Exercises non-colocated PD path in create_inference_servers with separate - prefill/decode placement groups. + prefill/decode placement groups, and -- being two deployments -- the + per-deployment ``rank_offset`` advance, which a single deployment cannot + distinguish from a constant. """ pd_cfg = request.param - enable_pd = pd_cfg["enable_pd"] cfg = SkyRLTrainConfig() cfg.trainer.policy.model.path = MODEL - if enable_pd: + if pd_cfg["enable_pd"]: num_prefill = pd_cfg["num_prefill"] num_decode = pd_cfg["num_decode"] create_kwargs = dict( @@ -174,220 +253,27 @@ async def weight_update_env(class_scoped_ray_init_fixture, request): engine_init_kwargs={"load_format": "dummy"}, ) - async with InferenceEngineState.create(cfg, **create_kwargs) as engines: - trainer = Trainer.options(num_gpus=1.0).remote(MODEL) - ray.get(trainer.ready.remote()) - - yield { - "engines": engines, - "trainer": trainer, - "client": engines.client, - "router_url": engines.client.proxy_url, - } - - await engines.client.teardown() - ray.kill(trainer) - # cleanup manually in colocated case - if engines.pg: - ray.util.remove_placement_group(engines.pg) + async for env in _make_env(cfg, create_kwargs, NcclTrainer, "nccl"): + yield env @pytest.mark.asyncio(loop_scope="class") class TestWeightUpdateFlow: - """Tests for weight synchronization from trainer to inference server (non-colocated).""" + """Weight sync via NCCL broadcast (non-colocated).""" async def test_update_weights_flow(self, weight_update_env): - """ - Full E2E weight sync test (non-colocated, NCCL broadcast): - 1. Query with dummy weights → gibberish - 2. Init weight transfer (both sides concurrently via client) - 3. Broadcast weights from trainer (concurrent with server receive) - 4. Finalize weight update - 5. Query again → correct output - """ - router_url = weight_update_env["router_url"] - trainer = weight_update_env["trainer"] - client = weight_update_env["client"] - - print("\n[TEST] Running non-colocated weight sync test") - - async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as http_client: - # ===== Step 1: Verify dummy weights produce gibberish ===== - payload = { - "model": MODEL, - "prompt": "What is the capital of France?", - "max_tokens": 32, - "temperature": 0.0, - } - - resp = await http_client.post(f"{router_url}/v1/completions", json=payload) - assert resp.status_code == 200 - - text_before = resp.json()["choices"][0]["text"] - print(f"[Step 1] Dummy weights output: {text_before!r}") - - # Dummy weights should NOT produce coherent output about Paris - assert "Paris" not in text_before, "Dummy weights unexpectedly produced correct answer" - - # ===== Step 2: Init weight transfer (both sides concurrently) ===== - master_address = get_node_ip() - master_port = get_open_port() - - # Query all servers for world_size via client (fans out to all backends) - inference_world_size, _ = await client.get_world_size() - world_size = 1 + inference_world_size # 1 trainer + all inference workers - group_name = f"weight_sync_test_{master_port}" - - print(f"[Step 2] Init weight transfer: master={master_address}:{master_port}, world_size={world_size}") - - init_info = BroadcastInitInfo( - master_addr=master_address, - master_port=master_port, - rank_offset=1, - world_size=world_size, - override_existing_receiver=True, - ) - - # Both sides must init concurrently (NCCL blocks until all ranks join) - # Start trainer init (returns immediately, runs in Ray actor) - trainer_init_ref = trainer.init_weight_sync.remote(master_address, master_port, world_size, group_name) - - # Await server init via client (fans out to all backends) - result = await client.init_weight_update_communicator(init_info) - for server_url, resp in result.items(): - assert resp["status"] == 200, f"Server {server_url} init failed: {resp}" - - # Trainer should be done now (NCCL group formed) - ray.get(trainer_init_ref) - print("[Step 2] Both sides init complete") - - # ===== Step 3: Broadcast weights (concurrent send/receive) ===== - print("[Step 3] Broadcasting weights from trainer to server...") - - # Get weight metadata first (no NCCL yet) - weight_info = ray.get(trainer.get_weight_info.remote()) - print(f"[Step 3] Weight info: {len(weight_info['names'])} parameters") - - # Start trainer broadcast (returns immediately, runs in Ray actor) - print("[Step 3] Launching trainer broadcast_weights.remote()...") - trainer_broadcast_ref = trainer.broadcast_weights.remote() - - # Await server receive via client (fans out to all backends) - dtype_names = [(d.split(".")[-1] if "." in d else d) for d in weight_info["dtypes"]] - # No "packed" here: since vLLM 0.28.0 it is an init-time wire param - # (BroadcastInitInfo.packed -> NCCLWeightTransferInitInfo.packed), and - # NCCLWeightTransferUpdateInfo rejects it. - update_info = { - "names": weight_info["names"], - "dtype_names": dtype_names, - "shapes": weight_info["shapes"], - } - print( - f"[Step 3] Calling update_weights_nccl with {len(update_info['names'])} names, " - f"packed={init_info.packed} (from init)" - ) - # Use SkyRL's chunked weight-sync API (skyrl_start_weight_update -> - # update_weights_nccl -> skyrl_finish_weight_update) rather than vLLM's - # native /update_weights endpoint, which in vLLM 0.22.0+ requires - # vLLM's own native start_weight_update to be called first. - # skyrl_start_weight_update is local (layerwise-reload init), so it is - # safe to call while the trainer is blocked on the NCCL send; the - # actual receive happens in update_weights_nccl. - await client.start_weight_update() - result = await client.update_weights_nccl(update_info) - print(f"[Step 3] update_weights_nccl returned: {list(result.keys())}") - for server_url, resp in result.items(): - assert resp["status"] == 200, f"Server {server_url} update weights failed: {resp}" - await client.finish_weight_update() - - # Trainer should be done now (NCCL broadcast complete) - ray.get(trainer_broadcast_ref) - print("[Step 3] Weight sync complete") - - # ===== Step 4: Query again - should produce correct output ===== - resp = await http_client.post(f"{router_url}/v1/completions", json=payload) - assert resp.status_code == 200 - - text_after = resp.json()["choices"][0]["text"] - print(f"[Step 5] Real weights output: {text_after!r}") - - assert "Paris" in text_after, f"Weight sync failed - expected 'Paris' but got: {text_after!r}" - - print("[SUCCESS] Non-colocated weight sync test passed!") + """``send_weights()`` runs the inference-side ``update_weights`` + concurrently with the trainer-side broadcast, so a mis-sized receive + buffer or a mismatched ``packed`` hangs rather than failing. Treat a + timeout here as that bug, not as flake.""" + await _assert_sync_replaces_dummy_weights(weight_update_env) # ----------------------------------------------------------------- -# Colocated CUDA IPC Weight Sync Test +# Colocated CUDA IPC # ----------------------------------------------------------------- -@ray.remote -class IpcTrainer: - """ - Trainer emulator that creates CUDA IPC handles for weight transfer. - - Unlike the NCCL Trainer, this does not create a process group. - Instead it creates per-tensor IPC handles that the colocated - inference server opens to read weights directly from GPU memory. - """ - - def __init__(self, model_name: str, device: str = "cuda"): - self.device = torch.device(device) - self.model = AutoModelForCausalLM.from_pretrained( - model_name, - torch_dtype=torch.bfloat16, - ).to(self.device) - self._tensor_refs: list = [] - - def ready(self): - return True - - def create_ipc_update_info(self) -> dict: - """Create a single packed CUDA-IPC buffer for all model parameters. - - Matches SkyRL's ``update_weights_ipc`` contract (the packed format - produced by ``CudaIpcTransferStrategy``): all parameters are copied into - one contiguous CUDA buffer, a single IPC handle is created for that - buffer, and per-parameter ``sizes`` let the receiver slice it back out. - This differs from vLLM's native ``/update_weights`` (one handle per - parameter), which we no longer use. - """ - from torch.multiprocessing.reductions import reduce_tensor - - gpu_uuid = str(torch.cuda.get_device_properties(torch.cuda.current_device()).uuid) - - params = list(self.model.named_parameters()) - # The model is loaded in a single dtype (bfloat16), so element offsets - # into one packed buffer are well-defined across all parameters. - dtype = params[0][1].dtype - total_numel = sum(p.numel() for _, p in params) - packed_tensor = torch.empty(total_numel, device=self.device, dtype=dtype) - - names, dtype_names, shapes, sizes = [], [], [], [] - offset = 0 - for name, param in params: - size = param.numel() - packed_tensor[offset : offset + size].copy_(param.detach().reshape(-1)) - offset += size - names.append(name) - dtype_names.append(str(param.dtype).split(".")[-1]) - shapes.append(list(param.shape)) - sizes.append(size) - - # Keep the packed buffer alive so the IPC handle stays valid on the receiver. - self._tensor_refs = [packed_tensor] - - ipc_handle = reduce_tensor(packed_tensor) - pickled = base64.b64encode(pickle.dumps({gpu_uuid: ipc_handle})).decode("utf-8") - return { - "names": names, - "dtype_names": dtype_names, - "shapes": shapes, - "sizes": sizes, - "ipc_handles_pickled": pickled, - } - - @pytest_asyncio.fixture(scope="class") async def ipc_weight_update_env(class_scoped_ray_init_fixture): """Create environment for colocated IPC weight update testing.""" @@ -401,212 +287,40 @@ async def ipc_weight_update_env(class_scoped_ray_init_fixture): engine_init_kwargs={"load_format": "dummy"}, ) - async with InferenceEngineState.create(cfg, **create_kwargs) as engines: - # Trainer on same PG bundle as server (colocated) with fractional GPU - trainer = IpcTrainer.options( - num_gpus=0.2, - num_cpus=0.2, - scheduling_strategy=PlacementGroupSchedulingStrategy( - placement_group=engines.pg, - placement_group_bundle_index=0, - ), - ).remote(MODEL) - ray.get(trainer.ready.remote()) - - yield { - "engines": engines, - "trainer": trainer, - "client": engines.client, - "router_url": engines.client.proxy_url, - } - - await engines.client.teardown() - ray.kill(trainer) - # cleanup manually in colocated case - if engines.pg: - ray.util.remove_placement_group(engines.pg) + # weight_sync_backend="nccl" + colocate_all resolves to ipc, exactly as the + # driver resolves it for the servers. + async for env in _make_env(cfg, create_kwargs, IpcTrainer, "nccl", colocate_with_engine=True): + yield env @pytest.mark.asyncio(loop_scope="class") class TestColocatedIpcWeightUpdateFlow: - """Tests for weight synchronization via CUDA IPC (colocated, TP=1).""" + """Weight sync via CUDA IPC (colocated, TP=1).""" async def test_update_weights_ipc(self, ipc_weight_update_env): - """ - Full E2E weight sync test (colocated, CUDA IPC): - 1. Query with dummy weights → gibberish - 2. Init IPC weight transfer engine (no-op for IPC) - 3. Create IPC handles from trainer weights and send to server - 4. Query again → correct output - """ - router_url = ipc_weight_update_env["router_url"] - trainer = ipc_weight_update_env["trainer"] - client = ipc_weight_update_env["client"] - - print("\n[TEST] Running colocated IPC weight sync test") - - async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as http_client: - # ===== Step 1: Verify dummy weights produce gibberish ===== - payload = { - "model": MODEL, - "prompt": "What is the capital of France?", - "max_tokens": 32, - "temperature": 0.0, - } - - resp = await http_client.post(f"{router_url}/v1/completions", json=payload) - assert resp.status_code == 200 - - text_before = resp.json()["choices"][0]["text"] - print(f"[Step 1] Dummy weights output: {text_before!r}") - assert "Paris" not in text_before, "Dummy weights unexpectedly produced correct answer" - - # ===== Step 2: Init IPC engine (no-op but verifies endpoint) ===== - init_info = CudaIpcInitInfo( - model_dtype_str="bfloat16", - override_existing_receiver=True, - ) - result = await client.init_weight_update_communicator(init_info) - for server_url, resp_data in result.items(): - assert resp_data["status"] == 200, f"Server {server_url} IPC init failed: {resp_data}" - print("[Step 2] IPC engine init complete (no-op)") - - # ===== Step 3: Create IPC handles and send to server ===== - print("[Step 3] Creating IPC handles from trainer weights...") - update_info = ray.get(trainer.create_ipc_update_info.remote()) - print(f"[Step 3] Created handles for {len(update_info['names'])} parameters") - - # Use SkyRL's chunked weight-sync API (skyrl_start_weight_update -> - # update_weights_ipc -> skyrl_finish_weight_update) rather than vLLM's - # native /update_weights endpoint. - await client.start_weight_update() - result = await client.update_weights_ipc(update_info) - for server_url, resp_data in result.items(): - assert resp_data["status"] == 200, f"Server {server_url} IPC update failed: {resp_data}" - await client.finish_weight_update() - print("[Step 3] IPC weight update complete") - - # ===== Step 4: Query again — should produce correct output ===== - resp = await http_client.post(f"{router_url}/v1/completions", json=payload) - assert resp.status_code == 200 - - text_after = resp.json()["choices"][0]["text"] - print(f"[Step 4] Real weights output: {text_after!r}") - assert "Paris" in text_after, f"IPC weight sync failed - expected 'Paris' but got: {text_after!r}" - - print("[SUCCESS] Colocated IPC weight sync test passed!") + """Packed IPC: the producer streams through one reusable buffer and the + consumer clones out of it. A wrong refcount contract surfaces as garbage + weights -- the buffer reused under a reader -- which "Paris" catches.""" + await _assert_sync_replaces_dummy_weights(ipc_weight_update_env) # ----------------------------------------------------------------- -# Sharded RDT (NIXL pull) Weight Sync Test +# Sharded RDT (NIXL pull) # ----------------------------------------------------------------- -# The trainer side is the VENDORED vLLM RDT sidecar engine -# (ShardedRDTTrainerWeightTransferEngine): trainer_init spawns a per-rank -# _RDTProducerServer actor (the NIXL serve surface) and shares gathered weights -# into it over CUDA IPC; send_weights drives the concurrent start/update/finish -# handshake. This test drives that engine from a single-GPU (no-FSDP) Ray actor -# via the SkyRL adapter (SyncRdtControlPlaneClient + _FsdpWeightSource) — the -# same code path FSDPPolicyWorkerBase uses in production, minus FSDP sharding. - - -class _ShimExtractor: - """Minimal weight-extractor for a single-GPU (no-FSDP) trainer: params are - already whole, so gather is identity and there is no name prefix. Matches the - surface ``_FsdpWeightSource`` reads (model / weight_prefix / _gather_tensor / - get_weight_metadata).""" - - weight_prefix = "" - - def __init__(self, model): - self.model = model - - def _gather_tensor(self, param): - return param - - def get_weight_metadata(self, dtype): - names, dtype_names, shapes = [], [], [] - dtype_name = str(dtype).split(".")[-1] - for name, param in self.model.state_dict().items(): - names.append(name) - dtype_names.append(dtype_name) - shapes.append(list(param.shape)) - return {"names": names, "dtype_names": dtype_names, "shapes": shapes} - - -@ray.remote(num_gpus=1, max_concurrency=4) -class RdtTrainerActor: - """Single-GPU (no-FSDP) trainer that drives the vendored sidecar RDT engine. - - Holds the model on its own GPU and drives the engine's synchronous - control-plane calls through ``SyncRdtControlPlaneClient`` (blocking HTTP to - the servers' ``/collective_rpc``). No event loop and no async client are - involved: the July control-plane rework replaced the old - ``_SyncInferenceClient`` loop bridge with direct blocking HTTP. - ``num_gpus=1`` (a real GPU assignment) is what lets the trainer engine pin - its spawned ``_RDTProducerServer`` to this GPU for CUDA IPC. - """ - - def __init__(self, model_name, server_urls, data_parallel_size, namespace, num_consumers): - self._model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16).to("cuda") - self._extractor = _ShimExtractor(self._model) - self._namespace = namespace - self._num_consumers = num_consumers - self._server_urls = list(server_urls) - self._data_parallel_size = data_parallel_size - self._engine = None - - def ready(self): - return True - - def sync_once(self): - """Rendezvous (first call: spawn server + bake the inference side) and run - one full weight sync through the vendored engine.""" - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_control_plane import ( - SyncRdtControlPlaneClient, - ) - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - _FsdpWeightSource, - ) - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_trainer import ( - ShardedRDTTrainerInitInfo, - ShardedRDTTrainerWeightTransferEngine, - ) - - sync_client = SyncRdtControlPlaneClient(self._server_urls, self._data_parallel_size) - source = _FsdpWeightSource(self._extractor, torch.bfloat16) - if self._engine is None: - init_info = ShardedRDTTrainerInitInfo( - rank=0, - num_consumers=self._num_consumers, - trainer_actor_namespace=self._namespace, - ) - self._engine = ShardedRDTTrainerWeightTransferEngine.trainer_init( - init_info, - client=sync_client, - source=source, - ) - self._engine.send_weights() - - def shutdown(self): - if self._engine is not None: - self._engine.shutdown() - self._engine = None - @pytest_asyncio.fixture(scope="class") async def rdt_weight_update_env(class_scoped_ray_init_fixture): """Non-colocated sharded_rdt (NIXL pull) environment, TP=1. - The trainer actor (1 GPU) drives the vendored sidecar engine, which spawns - its own producer server; the vLLM server (TP=1, - distributed_executor_backend=ray) runs on another GPU. 2 GPUs + the sidecar - (shares the trainer's GPU). + The trainer actor (1 GPU) drives the engine, which spawns its own producer + sidecar on that GPU; the vLLM server (TP=1, + distributed_executor_backend=ray) runs on another GPU. 2 GPUs + the sidecar. """ cfg = SkyRLTrainConfig() cfg.trainer.policy.model.path = MODEL - # Select the sharded_rdt weight-sync backend (build_vllm_cli_args reads this - # and sets WeightTransferConfig(backend="sharded_rdt") + executor=ray). + # Selects the sharded_rdt backend: build_vllm_cli_args reads this and sets + # WeightTransferConfig(backend="sharded_rdt") + executor=ray. cfg.generator.inference_engine.weight_sync_backend = "sharded_rdt" create_kwargs = dict( @@ -617,30 +331,8 @@ async def rdt_weight_update_env(class_scoped_ray_init_fixture): engine_init_kwargs={"load_format": "dummy"}, ) - async with InferenceEngineState.create(cfg, **create_kwargs) as engines: - client = engines.client - namespace = ray.get_runtime_context().namespace or None - trainer = RdtTrainerActor.remote( - MODEL, - client.server_urls, - client.data_parallel_size, - namespace, - 1, # num_consumers (TP=1, single engine) - ) - ray.get(trainer.ready.remote()) - - yield { - "engines": engines, - "trainer": trainer, - "client": client, - "router_url": client.proxy_url, - } - - ray.get(trainer.shutdown.remote()) - await engines.client.teardown() - ray.kill(trainer) - if engines.pg: - ray.util.remove_placement_group(engines.pg) + async for env in _make_env(cfg, create_kwargs, RdtTrainer, "sharded_rdt"): + yield env @pytest.mark.asyncio(loop_scope="class") @@ -648,49 +340,7 @@ class TestShardedRdtWeightUpdateFlow: """Weight sync via the sharded_rdt (NIXL pull) backend (non-colocated, TP=1).""" async def test_update_weights_rdt(self, rdt_weight_update_env): - """ - Full E2E weight sync test (non-colocated, sharded RDT / NIXL pull) via - the VENDORED sidecar trainer engine: - 1. Query with dummy weights -> gibberish. - 2. trainer.sync_once(): the vendored engine spawns its producer server, - bakes the plan on the inference side (init_weight_transfer_engine_rdt), - then drives start_weight_update -> concurrent gather/publish + - update_weights_rdt (workers pull their slices over NIXL) -> - finish_weight_update. - 3. Query again -> correct output. - """ - router_url = rdt_weight_update_env["router_url"] - trainer = rdt_weight_update_env["trainer"] - - print("\n[TEST] Running sharded_rdt (NIXL pull) weight sync test (sidecar)") - - async with httpx.AsyncClient(timeout=httpx.Timeout(180.0)) as http_client: - payload = { - "model": MODEL, - "prompt": "What is the capital of France?", - "max_tokens": 32, - "temperature": 0.0, - } - - # ===== Step 1: dummy weights -> gibberish ===== - resp = await http_client.post(f"{router_url}/v1/completions", json=payload) - assert resp.status_code == 200 - text_before = resp.json()["choices"][0]["text"] - print(f"[Step 1] Dummy weights output: {text_before!r}") - assert "Paris" not in text_before, "Dummy weights unexpectedly produced correct answer" - - # ===== Step 2: drive the vendored sidecar engine end-to-end ===== - # sync_once rendezvouses (spawn server + bake) on the first call and - # runs the full concurrent gather/pull handshake. - print("[Step 2] trainer.sync_once() — bake + NIXL pull weight sync") - await asyncio.to_thread(lambda: ray.get(trainer.sync_once.remote())) - print("[Step 2] Weight sync complete") - - # ===== Step 3: real weights -> correct output ===== - resp = await http_client.post(f"{router_url}/v1/completions", json=payload) - assert resp.status_code == 200 - text_after = resp.json()["choices"][0]["text"] - print(f"[Step 3] Real weights output: {text_after!r}") - assert "Paris" in text_after, f"RDT weight sync failed - expected 'Paris' but got: {text_after!r}" - - print("[SUCCESS] sharded_rdt (sidecar) weight sync test passed!") + """The first ``sync_once`` spawns the producer sidecar and bakes the + replay plan on the inference side, inside the engine's + ``init_transfer_engine``. Then the workers pull their slices over NIXL.""" + await _assert_sync_replaces_dummy_weights(rdt_weight_update_env, timeout_s=180.0) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_extractor_consistency.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_weight_source.py similarity index 56% rename from tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_extractor_consistency.py rename to tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_weight_source.py index cd66279d64..adb151482d 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_extractor_consistency.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_weight_source.py @@ -1,74 +1,67 @@ -"""Iteration-order consistency test for ``MegatronWeightExtractor``. - -When ``enable_bucketing=True``, ``extract_weights`` (the producer that streams -parameter tensors to the inference engine) and ``get_weight_metadata`` (the -consumer's ``update_info`` source) must yield the same parameters, in the -same order, with the same count. Any divergence between the two methods -breaks downstream weight-sync consumers that rely on positional alignment -between the two streams. - -This test loads a Megatron ref worker for two parametrizations -- a -multimodal MoE model that exercises the bucketed grouped-export path and a -small dense model that serves as a non-MoE sanity check -- builds a fresh -``MegatronWeightExtractor`` per rank with bucketing enabled, calls -``get_weight_metadata`` and ``extract_weights`` on it, and asserts that the -two iteration sequences are identical (same length and same per-position -``name``). +"""Channel-agreement test for the Megatron ``WeightSource``. + +``metadata()`` declares what iteration will yield, and the trainer engine sizes +the worker's receive buffers -- and, in packed mode, cuts its chunk boundaries -- +from it. A source that disagrees splits the stream differently on each side. +``NCCLTrainerWeightTransferEngine._checked_iter`` catches that at runtime and +names the first divergent parameter; this test gets there first. + +Two parametrizations: a multimodal MoE model that exercises the grouped expert +export, and a small dense model as a non-MoE sanity check. Run with:: - uv run --isolated --extra megatron --extra dev pytest -s -vvv tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_extractor_consistency.py + uv run --isolated --extra megatron --extra dev pytest -s -vvv tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_megatron_weight_source.py """ import pytest import ray -import torch +from skyrl.backends.skyrl_train.weight_sync.sources import MegatronWeightSource from skyrl.backends.skyrl_train.workers.megatron import ( megatron_worker as _megatron_worker_mod, ) from skyrl.backends.skyrl_train.workers.megatron.megatron_worker import ( MegatronRefWorkerBase, - MegatronWeightExtractor, ) from skyrl.train.config import SkyRLTrainConfig -from skyrl.train.utils.utils import str_to_torch_dtype, validate_cfg +from skyrl.train.utils.utils import validate_cfg from tests.backends.skyrl_train.gpu.utils import init_worker_with_type class _ProbeMegatronRefWorker(MegatronRefWorkerBase): - """Test-only ``MegatronRefWorkerBase`` subclass that exposes a probe of - ``MegatronWeightExtractor`` iteration sequences. + """Exposes a probe of the Megatron ``WeightSource``'s two channels. - The probe is added on the test side (rather than on the production - ``MegatronRefWorkerBase``) so production code stays free of test-only - instrumentation. + Test-side rather than on the production worker, so production code stays free + of test-only instrumentation. """ - def probe_extractor_iteration_sequences(self, dtype_str: str) -> dict: - """Return per-rank ``get_weight_metadata`` and ``extract_weights`` - name sequences captured from a fresh ``MegatronWeightExtractor``.""" - dtype = str_to_torch_dtype(dtype_str) - - extractor = MegatronWeightExtractor( - bridge=self.bridge, - actor_module=self.actor_module, - enable_bucketing=True, - bucket_size_threshold_GB=1.0, - training_dtype=torch.bfloat16, - ) + def probe_source_channel_agreement(self, dtype_str: str) -> dict: + """Return this rank's ``metadata()`` and iteration name sequences. + + ``metadata()`` caches its dry export and iteration runs a second one, so + this also covers the two exports agreeing. + """ + from skyrl.train.utils.utils import str_to_torch_dtype + + source = MegatronWeightSource(self.bridge, self.actor_module, str_to_torch_dtype(dtype_str)) - metadata = extractor.get_weight_metadata(dtype) - meta_names = list(metadata["names"]) + meta = source.metadata() + meta_names = [m.name for m in meta] + meta_shapes = [list(m.shape) for m in meta] - extract_names: list[str] = [] - for chunk in extractor.extract_weights(dtype): - extract_names.extend(chunk.names) - del chunk + iter_names: list[str] = [] + iter_shapes: list[list[int]] = [] + for name, tensor in source: + iter_names.append(name) + iter_shapes.append(list(tensor.shape)) + del tensor return { "meta_names": meta_names, - "extract_names": extract_names, + "meta_shapes": meta_shapes, + "iter_names": iter_names, + "iter_shapes": iter_shapes, } @@ -118,9 +111,9 @@ def _make_ref_cfg(model_name: str) -> SkyRLTrainConfig: pytest.param("Qwen/Qwen2.5-1.5B-Instruct", id="qwen2_5_1_5b_dense"), ], ) -def test_megatron_extractor_iteration_order_consistency(ray_init_fixture, model_name): - """Per rank, assert ``get_weight_metadata`` and ``extract_weights`` - yield the same parameter names in the same order with the same count.""" +def test_megatron_source_channel_agreement(ray_init_fixture, model_name): + """Per rank, assert ``metadata()`` and iteration yield the same parameter + names and shapes, in the same order, with the same count.""" cfg = _make_ref_cfg(model_name) # Monkey-patch the production ``RefWorker`` symbol so @@ -138,30 +131,39 @@ def test_megatron_extractor_iteration_order_consistency(ray_init_fixture, model_ num_gpus_per_node=4, cfg=cfg, ) - results = ray.get(ref.async_run_ray_method("pass_through", "probe_extractor_iteration_sequences", "bfloat16")) + results = ray.get(ref.async_run_ray_method("pass_through", "probe_source_channel_agreement", "bfloat16")) assert results, "expected at least one Megatron ref rank" for rank_idx, result in enumerate(results): meta_names = result["meta_names"] - extract_names = result["extract_names"] + iter_names = result["iter_names"] assert len(meta_names) > 0, f"[rank {rank_idx}] empty iteration sequence" - assert len(meta_names) == len(extract_names), ( + assert len(meta_names) == len(iter_names), ( f"[rank {rank_idx}] count divergence: " - f"get_weight_metadata yielded {len(meta_names)} params, " - f"extract_weights yielded {len(extract_names)}" + f"metadata() declared {len(meta_names)} params, " + f"iteration yielded {len(iter_names)}" ) # First-divergence index for a useful failure message. first_diff = next( - (i for i, (a, b) in enumerate(zip(meta_names, extract_names)) if a != b), + ( + i + for i, (a, b) in enumerate( + zip( + zip(meta_names, result["meta_shapes"]), + zip(iter_names, result["iter_shapes"]), + ) + ) + if a != b + ), None, ) assert first_diff is None, ( - f"[rank {rank_idx}] order divergence at index {first_diff}: " - f"metadata={meta_names[first_diff]!r}, " - f"extract={extract_names[first_diff]!r}" + f"[rank {rank_idx}] channel divergence at index {first_diff}: " + f"metadata()={meta_names[first_diff]!r} {result['meta_shapes'][first_diff]}, " + f"iteration={iter_names[first_diff]!r} {result['iter_shapes'][first_diff]}" ) print( - f"[rank {rank_idx}] iteration sequences match: N={len(meta_names)} params, " + f"[rank {rank_idx}] channels agree: N={len(meta_names)} params, " f"first={meta_names[0]!r}, last={meta_names[-1]!r}," ) finally: diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py b/tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py index d10da48eb2..9b8c6a86a9 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py @@ -1,12 +1,16 @@ -"""Tests for prefix cache reset behaviour in `PolicyWorker.broadcast_to_inference_engines` +"""Tests for prefix cache reset behaviour in ``PolicyWorker.broadcast_to_inference_engines``. -The worker's send path goes through ``WeightTransferSender.send``, whose default -implementation materializes the chunk stream and calls ``send_chunks``; senders that -pull instead (sharded_rdt) override ``send``. When -``WeightTransferSender.handles_prefix_cache_reset`` is ``True``, -the sender will resets the cache itself as part of its own pause/update sequence -(``DeltaWeightTransferSender`` does this inside ``_apply_receiver_update``), and the worker must skip -this. +The worker calls the trainer engine's ``send_weights()`` inside a memory bracket +whose three decisions come from ``getattr`` capability probes on the engine +(``skyrl_handles_prefix_cache_reset``, +``skyrl_force_disable_expandable_segments``, ``skyrl_empty_cache_after_send``). +Two of the four engines are vLLM's own classes and cannot declare SkyRL +attributes, so the *absence* of a flag is the common case and must mean the +default. + +An engine declaring ``skyrl_handles_prefix_cache_reset`` resets the cache itself, +at the right point in its own pause/update sequence, and the worker must skip its +own concurrent reset. uv run --extra dev --extra fsdp -- pytest -s tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py -m "not megatron" uv run --extra dev --extra megatron -- pytest -s tests/backends/skyrl_train/gpu/gpu_ci/test_prefix_cache_reset.py -m "megatron" @@ -23,32 +27,49 @@ from skyrl.backends.skyrl_train.workers.worker import PolicyWorkerBase +class _FakeEngine: + """A trainer engine with only what the worker's send bracket touches. + + A plain class, not a Mock: a Mock answers every attribute, so it could never + exercise the probe defaults. + """ + + def __init__(self, **flags) -> None: + self.sends = 0 + self.reset_prefix_cache_told = None + for name, value in flags.items(): + setattr(self, f"skyrl_{name}", value) + + def send_weights(self) -> None: + self.sends += 1 + + +class _DeltaLikeEngine(_FakeEngine): + """Also implements the per-round setter, as the delta engine does.""" + + def skyrl_set_reset_prefix_cache(self, reset: bool) -> None: + self.reset_prefix_cache_told = reset + + # TODO (sumanthrh): Ideally we avoid all this mocking with an easier way to construct dummy policy workers -def _make_worker(worker_cls, handles_prefix_cache_reset: bool): +def _make_worker(worker_cls, engine): """Build a policy worker with just enough state for broadcast_to_inference_engines.""" worker = worker_cls.__new__(worker_cls) worker._is_lora = False worker.cfg = SimpleNamespace( fully_async=SimpleNamespace(enabled=False, clear_kv_cache_on_weight_sync=False), + placement=SimpleNamespace(colocate_all=False), policy=SimpleNamespace(megatron_config=SimpleNamespace(lora_config=SimpleNamespace(merge_lora=False))), ) - worker._weight_transfer_sender = AsyncMock() - # AsyncMock would make this attribute a coroutine; the production code reads it as a - # plain bool, so set it explicitly. - worker._weight_transfer_sender.handles_prefix_cache_reset = handles_prefix_cache_reset - # Same reason: the worker reads these as plain bools to decide the - # expandable-segments toggle and the post-send empty_cache. - worker._weight_transfer_sender.force_disable_expandable_segments = False - worker._weight_transfer_sender.empty_cache_after_send = True - worker._weight_transfer_sender._inference_client = None - worker.weight_extractor = MagicMock() + worker._weight_sync_engine = engine + worker._weight_sync_inference_client = None # FSDP dereferences self.model.model before the LoRA branch, so it must exist even # though _is_lora is False and the resulting peft_model goes unused. worker.model = SimpleNamespace(model=SimpleNamespace()) @contextmanager def _noop_ctx(*args, **kwargs): - # The worker calls this with force=. + # The worker calls this with force=. yield worker._expandable_segments_disabled_for_sync = _noop_ctx @@ -61,8 +82,8 @@ def _patch_collectives(monkeypatch): monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) -def _ie_cfg(): - return SimpleNamespace(enable_prefix_caching=True, model_dtype="bfloat16") +def _ie_cfg(enable_prefix_caching: bool = True): + return SimpleNamespace(enable_prefix_caching=enable_prefix_caching, model_dtype="bfloat16") def get_worker_cls(strategy: str) -> Type[PolicyWorkerBase]: @@ -82,50 +103,110 @@ def get_worker_cls(strategy: str) -> Type[PolicyWorkerBase]: raise ValueError(f"Invalid worker cls: {strategy}") -@pytest.mark.asyncio -@pytest.mark.parametrize("strategy", ["fsdp", pytest.param("megatron", marks=pytest.mark.megatron)]) -async def test_worker_skips_prefix_cache_reset_when_sender_handles_it(strategy, monkeypatch): - worker_cls = get_worker_cls(strategy) +STRATEGIES = ["fsdp", pytest.param("megatron", marks=pytest.mark.megatron)] + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", STRATEGIES) +async def test_worker_skips_prefix_cache_reset_when_engine_handles_it(strategy, monkeypatch): _patch_collectives(monkeypatch) - worker = _make_worker(worker_cls, handles_prefix_cache_reset=True) + engine = _DeltaLikeEngine(handles_prefix_cache_reset=True) + worker = _make_worker(get_worker_cls(strategy), engine) client = AsyncMock() await worker.broadcast_to_inference_engines(client, _ie_cfg()) client.reset_prefix_cache.assert_not_awaited() - # The sync still happens, and the sender is still told the cache needs resetting so - # it can do it at the right point in its own sequence. - worker._weight_transfer_sender.send.assert_awaited_once() - assert worker._weight_transfer_sender.send.await_args.kwargs["reset_prefix_cache"] is True + # The sync still happens, and the engine is still told the cache needs + # resetting so it can do it at the right point in its own sequence. + assert engine.sends == 1 + assert engine.reset_prefix_cache_told is True @pytest.mark.asyncio -@pytest.mark.parametrize("strategy", ["fsdp", pytest.param("megatron", marks=pytest.mark.megatron)]) -async def test_worker_resets_prefix_cache_when_sender_does_not(strategy, monkeypatch): - worker_cls = get_worker_cls(strategy) - +@pytest.mark.parametrize("strategy", STRATEGIES) +async def test_worker_resets_prefix_cache_when_engine_does_not(strategy, monkeypatch): _patch_collectives(monkeypatch) - worker = _make_worker(worker_cls, handles_prefix_cache_reset=False) + # No flags at all -- the shape of vLLM's own NCCL / IPC trainer engines. + engine = _FakeEngine() + worker = _make_worker(get_worker_cls(strategy), engine) client = AsyncMock() await worker.broadcast_to_inference_engines(client, _ie_cfg()) client.reset_prefix_cache.assert_awaited_once_with(reset_running_requests=True) - worker._weight_transfer_sender.send.assert_awaited_once() + assert engine.sends == 1 @pytest.mark.asyncio -@pytest.mark.parametrize("strategy", ["fsdp", pytest.param("megatron", marks=pytest.mark.megatron)]) +@pytest.mark.parametrize("strategy", STRATEGIES) async def test_no_reset_when_prefix_caching_disabled(strategy, monkeypatch): - worker_cls = get_worker_cls(strategy) - _patch_collectives(monkeypatch) - worker = _make_worker(worker_cls, handles_prefix_cache_reset=False) + engine = _DeltaLikeEngine() + worker = _make_worker(get_worker_cls(strategy), engine) client = AsyncMock() - ie_cfg = SimpleNamespace(enable_prefix_caching=False, model_dtype="bfloat16") - await worker.broadcast_to_inference_engines(client, ie_cfg) + await worker.broadcast_to_inference_engines(client, _ie_cfg(enable_prefix_caching=False)) client.reset_prefix_cache.assert_not_awaited() - assert worker._weight_transfer_sender.send.await_args.kwargs["reset_prefix_cache"] is False + assert engine.reset_prefix_cache_told is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", STRATEGIES) +async def test_expandable_segments_force_comes_from_the_engine(strategy, monkeypatch): + """sharded_rdt asks for the toggle unconditionally; everything else leaves it + to ``colocate_all``. Default False for an engine that declares nothing.""" + _patch_collectives(monkeypatch) + seen = [] + + for engine in (_FakeEngine(), _FakeEngine(force_disable_expandable_segments=True)): + worker = _make_worker(get_worker_cls(strategy), engine) + + @contextmanager + def _record(force=False): + seen.append(force) + yield + + worker._expandable_segments_disabled_for_sync = _record + await worker.broadcast_to_inference_engines(AsyncMock(), _ie_cfg(enable_prefix_caching=False)) + + assert seen == [False, True] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", STRATEGIES) +async def test_empty_cache_after_send_defaults_on_and_can_be_declined(strategy, monkeypatch): + """The default must be True -- that is what vLLM's own engines get.""" + _patch_collectives(monkeypatch) + for engine, expected in ((_FakeEngine(), 2), (_FakeEngine(empty_cache_after_send=False), 1)): + calls = [] + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: calls.append(1)) + worker = _make_worker(get_worker_cls(strategy), engine) + await worker.broadcast_to_inference_engines(AsyncMock(), _ie_cfg(enable_prefix_caching=False)) + # One before the send always; one after only if the engine allows it. + assert len(calls) == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("strategy", STRATEGIES) +async def test_lora_sync_skips_the_engine_entirely(strategy, monkeypatch): + """A LoRA adapter sync writes safetensors and calls the LoRA route; there is + no tensor transfer, so it must not drive the engine.""" + _patch_collectives(monkeypatch) + engine = _FakeEngine() + worker = _make_worker(get_worker_cls(strategy), engine) + worker._is_lora = True + worker.model = SimpleNamespace(model=SimpleNamespace(peft_config={"default": {}})) + worker._resolve_lora_sync_target = lambda model_id: ("adapter", "/tmp/lora") + saved = MagicMock() + worker._save_lora_adapters_and_sync = AsyncMock(side_effect=saved) + + client = AsyncMock() + await worker.broadcast_to_inference_engines(client, _ie_cfg()) + + assert engine.sends == 0 + worker._save_lora_adapters_and_sync.assert_awaited_once() + # A LoRA sync still invalidates the prefix cache: the adapter changes what + # the cached prefixes decode to. + client.reset_prefix_cache.assert_awaited_once_with(reset_running_requests=True) diff --git a/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py b/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py index 75267428b6..171fdfa811 100644 --- a/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py +++ b/tests/backends/skyrl_train/inference_servers/test_remote_inference_client.py @@ -650,24 +650,10 @@ async def test_reset_prefix_cache(self, client): class TestWeightSync: """Test weight sync methods.""" - @pytest.mark.asyncio - async def test_init_weight_update_communicator(self, client): - """Test init_weight_update_communicator expands init_info via to_api_payload and fans out.""" - api_payload = {"master_address": "127.0.0.1", "master_port": 29500, "rank_offset": 1, "world_size": 5} - - class MockInitInfo: - """Lightweight mock satisfying the for_servers / to_api_payload protocol.""" - - def for_servers(self, world_size_per_server, num_servers, dp_size=1): - return [self] * num_servers - - def to_api_payload(self): - return dict(api_payload) - - result = await client.init_weight_update_communicator(MockInitInfo()) - assert set(result) == set(client.server_urls) - for response in result.values(): - assert response["body"]["body"] == {"init_info": api_payload} + # The init handshake and the start/update/finish lifecycle moved off this + # async client onto the trainer-side engines' blocking SkyrlWeightSyncClient + # (weight_sync/control_plane.py, covered by test_control_plane.py). What is + # left here is what the driver still drives. @pytest.mark.asyncio async def test_update_named_weights(self, client): diff --git a/tests/backends/skyrl_train/weight_sync/test_control_plane.py b/tests/backends/skyrl_train/weight_sync/test_control_plane.py new file mode 100644 index 0000000000..56e835d25f --- /dev/null +++ b/tests/backends/skyrl_train/weight_sync/test_control_plane.py @@ -0,0 +1,224 @@ +"""Tests for the trainer-side weight-sync control plane. + +The per-server init payload rewrites are what matter: both fail +silently-then-fatally when wrong. A bad NCCL ``rank_offset`` mis-maps ranks and +hangs in the rendezvous rather than erroring; a bad RDT ``replica_rank`` collides +two deployments' consumer id blocks. +""" + +import base64 +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List + +import pytest + +from skyrl.backends.skyrl_train.weight_sync.control_plane import ( + FINISH_UPDATE_ENDPOINT, + INIT_ENGINE_ENDPOINT, + START_UPDATE_ENDPOINT, + UPDATE_WEIGHTS_ENDPOINT, + SkyrlWeightSyncClient, + nccl_init_payloads, + rdt_init_payloads, +) + + +class _NotJsonable: + """Stands in for the storage type / handle bytes inside reduce_tensor's args.""" + + def __eq__(self, other): + return isinstance(other, _NotJsonable) + + +class _Resp: + def __init__(self, status_code: int = 200, body: Any = None, reason: str = "OK") -> None: + self.status_code = status_code + self.reason = reason + self._body = body if body is not None else {"status": "ok"} + self.text = str(self._body) + + def json(self): + return self._body + + +class _FakeSession: + """Records every POST; optionally fails for one URL.""" + + def __init__(self, fail_url: str = None, fail_body: Any = None) -> None: + self.headers: Dict[str, str] = {} + self.calls: List[tuple] = [] + self.closed = False + self._fail_url = fail_url + self._fail_body = fail_body + + def post(self, url, json=None, timeout=None): + self.calls.append((url, json)) + if self._fail_url is not None and url.startswith(self._fail_url): + return _Resp(500, self._fail_body, reason="Internal Server Error") + return _Resp() + + def close(self): + self.closed = True + + +@pytest.fixture +def make_client(monkeypatch): + """Build a client whose HTTP session is a recording fake.""" + + def _make(urls, dp=1, init_payload_fn=None, fail_url=None, fail_body=None): + client = SkyrlWeightSyncClient.__new__(SkyrlWeightSyncClient) + client._urls = list(urls) + client._dp = max(1, dp) + client._init_payload_fn = init_payload_fn + client._session = _FakeSession(fail_url=fail_url, fail_body=fail_body) + client._pool = ThreadPoolExecutor(max_workers=len(urls)) + return client + + return _make + + +class TestNcclInitPayloads: + """``rank_offset`` advances one deployment's worth per deployment, and stays + put across the DP servers within a deployment.""" + + def test_single_deployment_keeps_the_engines_offset(self): + init = {"master_address": "h", "master_port": 1, "rank_offset": 1, "world_size": 5} + payloads = nccl_init_payloads(init, ["u0"], 1) + assert [p["rank_offset"] for p in payloads] == [1] + # Everything else rides through untouched. + assert payloads[0]["master_port"] == 1 and payloads[0]["world_size"] == 5 + + def test_offset_advances_per_deployment(self): + # 2 deployments x 4 workers each + the trainer sender = world_size 9. + init = {"master_address": "h", "master_port": 1, "rank_offset": 1, "world_size": 9} + payloads = nccl_init_payloads(init, ["u0", "u1"], 1) + assert [p["rank_offset"] for p in payloads] == [1, 5] + + def test_dp_servers_of_one_deployment_share_an_offset(self): + # 2 deployments x dp=2 servers = 4 servers; 4 workers per deployment. + init = {"master_address": "h", "master_port": 1, "rank_offset": 1, "world_size": 9} + payloads = nccl_init_payloads(init, ["u0", "u1", "u2", "u3"], 2) + assert [p["rank_offset"] for p in payloads] == [1, 1, 5, 5] + + def test_world_size_that_does_not_divide_is_rejected(self): + # 3 deployments cannot share 7 workers; this is the shape that would + # otherwise mis-map ranks and hang in the rendezvous. + init = {"master_address": "h", "master_port": 1, "rank_offset": 1, "world_size": 8} + with pytest.raises(ValueError, match="does not divide"): + nccl_init_payloads(init, ["u0", "u1", "u2"], 1) + + def test_no_workers_is_rejected(self): + init = {"master_address": "h", "master_port": 1, "rank_offset": 1, "world_size": 1} + with pytest.raises(ValueError, match="inference workers"): + nccl_init_payloads(init, ["u0"], 1) + + +class TestRdtInitPayloads: + def test_stamps_deployment_ordinal_and_count(self): + payloads = rdt_init_payloads({"num_consumers": 8}, ["u0", "u1", "u2"], 1) + assert [p["replica_rank"] for p in payloads] == [0, 1, 2] + assert {p["num_replicas"] for p in payloads} == {3} + assert {p["num_consumers"] for p in payloads} == {8} + + def test_dp_servers_share_a_replica_rank(self): + payloads = rdt_init_payloads({}, ["u0", "u1", "u2", "u3"], 2) + assert [p["replica_rank"] for p in payloads] == [0, 0, 1, 1] + assert {p["num_replicas"] for p in payloads} == {2} + + +class TestClientFanout: + def test_init_uses_the_rewrite_and_hits_every_server(self, make_client): + client = make_client(["http://a", "http://b"], dp=1, init_payload_fn=rdt_init_payloads) + client.init_weight_transfer_engine({"num_consumers": 4}) + calls = dict(client._session.calls) + assert set(calls) == {f"http://a{INIT_ENGINE_ENDPOINT}", f"http://b{INIT_ENGINE_ENDPOINT}"} + assert calls[f"http://a{INIT_ENGINE_ENDPOINT}"]["init_info"]["replica_rank"] == 0 + assert calls[f"http://b{INIT_ENGINE_ENDPOINT}"]["init_info"]["replica_rank"] == 1 + + def test_init_without_a_rewrite_sends_the_same_dict_everywhere(self, make_client): + client = make_client(["http://a", "http://b"]) + client.init_weight_transfer_engine({"packed": True}) + assert [body["init_info"] for _, body in client._session.calls] == [{"packed": True}] * 2 + + def test_a_rewrite_returning_the_wrong_count_is_rejected(self, make_client): + client = make_client(["http://a", "http://b"], init_payload_fn=lambda info, urls, dp: [dict(info)]) + with pytest.raises(ValueError, match="one per server"): + client.init_weight_transfer_engine({}) + + def test_lifecycle_endpoints(self, make_client): + client = make_client(["http://a"]) + client.start_weight_update() + client.update_weights({"names": ["w"]}) + client.finish_weight_update() + assert [url for url, _ in client._session.calls] == [ + f"http://a{START_UPDATE_ENDPOINT}", + f"http://a{UPDATE_WEIGHTS_ENDPOINT}", + f"http://a{FINISH_UPDATE_ENDPOINT}", + ] + bodies = [body for _, body in client._session.calls] + assert bodies[0] is None + assert bodies[1] == {"update_info": {"names": ["w"]}} + # weight_version is omitted entirely when unset, so the route's default applies. + assert bodies[2] is None + + def test_ipc_handles_are_pickled_for_http(self, make_client): + """``IPCTrainerWeightTransferEngine`` emits raw ipc_handles, which are not + JSON-serializable; pickling them is this transport's job.""" + import json + import pickle + + # A non-JSON-native payload standing in for reduce_tensor's args, which + # carry storage *types* and raw handle bytes. Module-level so it pickles. + handles = {"GPU-uuid": ("rebuild", _NotJsonable(), 7)} + with pytest.raises(TypeError): + json.dumps(handles) + client = make_client(["http://a"]) + client.update_weights({"names": ["w"], "shapes": [[2, 2]], "ipc_handles": handles}) + + body = client._session.calls[0][1]["update_info"] + assert "ipc_handles" not in body + assert pickle.loads(base64.b64decode(body["ipc_handles_pickled"])) == handles + # Everything else rides through untouched. + assert body["names"] == ["w"] and body["shapes"] == [[2, 2]] + + def test_non_ipc_update_info_passes_through_unchanged(self, make_client): + client = make_client(["http://a"]) + update_info = {"names": ["w"], "dtype_names": ["bfloat16"], "shapes": [[2, 2]]} + client.update_weights(update_info) + assert client._session.calls[0][1] == {"update_info": update_info} + + def test_finish_carries_a_weight_version_when_given(self, make_client): + client = make_client(["http://a"]) + client.finish_weight_update("v7") + assert client._session.calls[0][1] == {"weight_version": "v7"} + + def test_every_server_is_posted_even_when_one_fails(self, make_client): + """A failure must not leave POSTs in flight against the other servers: + every future is drained before the first exception is raised.""" + client = make_client(["http://a", "http://b", "http://c"], fail_url="http://b") + with pytest.raises(RuntimeError): + client.start_weight_update() + assert len(client._session.calls) == 3 + + def test_error_surfaces_the_response_body_detail(self, make_client): + client = make_client( + ["http://a"], + fail_url="http://a", + fail_body={"error": {"message": "engine not initialized"}}, + ) + with pytest.raises(RuntimeError, match="engine not initialized"): + client.start_weight_update() + + def test_error_falls_back_to_fastapi_detail(self, make_client): + client = make_client(["http://a"], fail_url="http://a", fail_body={"detail": "missing init_info"}) + with pytest.raises(RuntimeError, match="missing init_info"): + client.start_weight_update() + + def test_close_is_idempotent(self, make_client): + client = make_client(["http://a"]) + client.close() + assert client._session.closed + + def test_no_servers_is_rejected(self): + with pytest.raises(ValueError, match="at least one server_url"): + SkyrlWeightSyncClient([], 1) diff --git a/tests/backends/skyrl_train/weight_sync/test_delta_checkpoint.py b/tests/backends/skyrl_train/weight_sync/test_delta_checkpoint.py index da64bf9b5f..e54b7eead4 100644 --- a/tests/backends/skyrl_train/weight_sync/test_delta_checkpoint.py +++ b/tests/backends/skyrl_train/weight_sync/test_delta_checkpoint.py @@ -8,7 +8,6 @@ from safetensors import safe_open from safetensors.torch import save_file -from skyrl.backends.skyrl_train.weight_sync.base import WeightChunk from skyrl.backends.skyrl_train.weight_sync.delta_checkpoint import ( _MANIFEST_NAME, _MAX_SAFE_PATH_NAME_LEN, @@ -26,13 +25,9 @@ ) -def _chunk_from_tensors(tensors): - return WeightChunk( - names=list(tensors.keys()), - dtypes=[str(t.dtype) for t in tensors.values()], - shapes=[list(t.shape) for t in tensors.values()], - tensors=list(tensors.values()), - ) +def _source(tensors): + """A minimal ``WeightSource``-shaped stream: ``create_delta_files`` only iterates.""" + return list(tensors.items()) def _write_checkpoint(path, tensors): @@ -76,7 +71,7 @@ def test_delta_checkpoint_publish_fetch_and_reload_roundtrip(tmp_path): sync_dir=str(tmp_path / "sync"), publish_staging_dir=str(tmp_path / "staging_dir"), ) - result = publisher.create_delta_files([_chunk_from_tensors(updated_tensors)]) + result = publisher.create_delta_files(_source(updated_tensors)) update_info = publisher.publish(result) store = LocalCheckpointStore(base_model_path=str(base_dir), local_checkpoint_dir=str(receiver_dir)) @@ -107,7 +102,7 @@ def test_delta_checkpoint_payload_stores_xor_patch(tmp_path): publish_staging_dir=str(tmp_path / "staging_dir"), publish_num_workers=1, ) - results = publisher.create_delta_files([_chunk_from_tensors({"a.weight": updated})]) + results = publisher.create_delta_files(_source({"a.weight": updated})) publisher.publish(results) with (sync_dir / "delta-00000001" / _MANIFEST_NAME).open(encoding="utf-8") as f: @@ -148,7 +143,7 @@ def test_delta_checkpoint_vllm_multi_thread_safetensors_iterator_roundtrip(tmp_p sync_dir=str(tmp_path / "sync"), publish_staging_dir=str(tmp_path / "staging_dir"), ) - result = publisher.create_delta_files([_chunk_from_tensors(updated_tensors)]) + result = publisher.create_delta_files(_source(updated_tensors)) update_info = publisher.publish(result) store = LocalCheckpointStore(base_model_path=str(base_dir), local_checkpoint_dir=str(receiver_dir)) @@ -173,7 +168,7 @@ def test_delta_checkpoint_publisher_converts_to_base_checkpoint_dtype(tmp_path): sync_dir=str(tmp_path / "sync"), publish_staging_dir=str(tmp_path / "staging_dir"), ) - update_info = publisher.create_delta_files([_chunk_from_tensors(runtime_updated)]) + update_info = publisher.create_delta_files(_source(runtime_updated)) update_info = publisher.publish(update_info) with open(tmp_path / "sync" / "delta-00000001" / "manifest.json", encoding="utf-8") as f: @@ -197,7 +192,7 @@ def test_delta_checkpoint_non_source_rank_drains_without_publishing(tmp_path, mo name: tensor + torch.tensor(idx + 1, dtype=torch.bfloat16) for idx, (name, tensor) in enumerate(base_tensors.items()) } - chunks = [_chunk_from_tensors({name: tensor}) for name, tensor in updated_tensors.items()] + stream = _source(updated_tensors) base_dir = tmp_path / "base" sync_dir = tmp_path / "sync" _write_checkpoint(base_dir, base_tensors) @@ -208,10 +203,10 @@ def test_delta_checkpoint_non_source_rank_drains_without_publishing(tmp_path, mo publish_staging_dir=str(tmp_path / "staging_dir"), ) - # Simulate a non-source rank (rank != 0): it drains the chunk stream but must + # Simulate a non-source rank (rank != 0): it drains the weight stream but must # not compute or upload any deltas. monkeypatch.setattr(publisher, "_current_rank", lambda: 1) - result = publisher.create_delta_files(chunks) + result = publisher.create_delta_files(stream) assert isinstance(result, DeltaPublishResult) assert result.records == [] assert result.payload_files == [] @@ -232,13 +227,13 @@ def test_delta_checkpoint_replays_multiple_versions_for_late_join(tmp_path): sync_dir=str(tmp_path / "sync"), publish_staging_dir=str(tmp_path / "staging_dir"), ) - v1_result = publisher.create_delta_files([_chunk_from_tensors(v1_tensors)]) + v1_result = publisher.create_delta_files(_source(v1_tensors)) publisher.publish(v1_result) first_snapshot = publisher.snapshot["a.weight"] first_snapshot_id = id(first_snapshot) assert first_snapshot.tobytes() == v1_tensors["a.weight"].contiguous().view(torch.uint8).numpy().tobytes() - v2_result = publisher.create_delta_files([_chunk_from_tensors(v2_tensors)]) + v2_result = publisher.create_delta_files(_source(v2_tensors)) update_info = publisher.publish(v2_result) assert id(publisher.snapshot["a.weight"]) == first_snapshot_id assert ( @@ -273,7 +268,7 @@ def test_delta_checkpoint_splits_payload_files_by_size(tmp_path): publish_staging_dir=str(tmp_path / "staging_dir"), max_file_size_in_gb=1e-9, ) - update_info = publisher.create_delta_files([_chunk_from_tensors(updated_tensors)]) + update_info = publisher.create_delta_files(_source(updated_tensors)) update_info = publisher.publish(update_info) with open(tmp_path / "sync" / "delta-00000001" / "manifest.json", encoding="utf-8") as f: @@ -297,12 +292,8 @@ def test_delta_checkpoint_skips_missing_lm_head_when_checkpoint_ties_embeddings( ) update_info = publisher.create_delta_files( [ - WeightChunk( - names=["model.embed_tokens.weight", "lm_head.weight"], - dtypes=["torch.bfloat16", "torch.bfloat16"], - shapes=[list(updated_embed.shape), list(updated_lm_head.shape)], - tensors=[updated_embed, updated_lm_head], - ) + ("model.embed_tokens.weight", updated_embed), + ("lm_head.weight", updated_lm_head), ] ) update_info = publisher.publish(update_info) @@ -331,7 +322,7 @@ def test_local_checkpoint_store_fetch_is_single_writer_with_concurrent_ray_actor publisher = DeltaCheckpointPublisher( base_model_path=str(base_dir), sync_dir=str(sync_dir), publish_staging_dir=str(tmp_path / "staging_dir") ) - update_info = publisher.create_delta_files([_chunk_from_tensors(updated_tensors)]) + update_info = publisher.create_delta_files(_source(updated_tensors)) update_info = publisher.publish(update_info) class FetchActor: @@ -405,7 +396,7 @@ def test_delta_checkpoint_unchanged_publish_advances_version(tmp_path): sync_dir=str(tmp_path / "sync"), publish_staging_dir=str(tmp_path / "staging_dir"), ) - update_info = publisher.create_delta_files([_chunk_from_tensors({"a.weight": base_tensors["a.weight"].clone()})]) + update_info = publisher.create_delta_files(_source({"a.weight": base_tensors["a.weight"].clone()})) update_info = publisher.publish(update_info) # An unchanged publish is not skipped: it still advances the version and @@ -483,7 +474,7 @@ def fake_run(cmd, stdout=None, stderr=None, text=None): sync_dir="gs://bucket/sync", publish_staging_dir=str(staging_dir), ) - update_info = publisher.create_delta_files([_chunk_from_tensors(updated_tensors)]) + update_info = publisher.create_delta_files(_source(updated_tensors)) update_info = publisher.publish(update_info) assert update_info["uri"] == "gs://bucket/sync/delta-00000001" @@ -542,7 +533,7 @@ def fake_run(cmd, stdout=None, stderr=None, text=None): sync_dir="s3://bucket/sync", publish_staging_dir=str(staging_dir), ) - update_info = publisher.create_delta_files([_chunk_from_tensors(updated_tensors)]) + update_info = publisher.create_delta_files(_source(updated_tensors)) update_info = publisher.publish(update_info) assert update_info["uri"] == "s3://bucket/sync/delta-00000001" @@ -570,7 +561,7 @@ def test_delta_checkpoint_checksum_failure_marks_write_in_progress(tmp_path): sync_dir=str(tmp_path / "sync"), publish_staging_dir=str(tmp_path / "staging_dir"), ) - update_info = publisher.create_delta_files([_chunk_from_tensors(updated_tensors)]) + update_info = publisher.create_delta_files(_source(updated_tensors)) update_info = publisher.publish(update_info) manifest_path = tmp_path / "sync" / "delta-00000001" / "manifest.json" manifest = json.loads(manifest_path.read_text(encoding="utf-8")) diff --git a/tests/backends/skyrl_train/weight_sync/test_rdt_control_plane.py b/tests/backends/skyrl_train/weight_sync/test_rdt_control_plane.py deleted file mode 100644 index 89f1278e93..0000000000 --- a/tests/backends/skyrl_train/weight_sync/test_rdt_control_plane.py +++ /dev/null @@ -1,192 +0,0 @@ -"""Unit tests for the synchronous sharded-RDT control-plane client. - -Covers the wire payloads (per-server replica_rank fan-out + uniform calls), the -body-aware error surfacing, and — most importantly — that the per-call fan-out is -*concurrent*, which is a correctness requirement (serial update_weights would -deadlock the producer's ref-counted group free), not just a perf choice. -""" - -import threading - -import pytest - -from skyrl.backends.skyrl_train.inference_servers.rdt_control_protocol import ( - COLLECTIVE_RPC_ENDPOINT, - RDT_FINISH_METHOD, - RDT_INIT_METHOD, - RDT_START_METHOD, - RDT_UPDATE_METHOD, -) -from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_control_plane import ( - SyncRdtControlPlaneClient, -) - - -class _FakeResponse: - def __init__(self, status_code=200, body=None, text="", reason="OK"): - self.status_code = status_code - self._body = body - self.text = text - self.reason = reason - - def json(self): - if self._body is None: - raise ValueError("no json") - return self._body - - -class _RecordingSession: - """Stand-in for requests.Session that records every POST. - - ``barrier`` (optional) is entered inside each POST so a test can assert the - calls actually run concurrently: if the client issued them serially the - barrier would never fill and time out. - """ - - def __init__(self, barrier=None, status_for=None): - self.headers = {} - self.calls = [] # (url, payload) - self.closed = False - self._lock = threading.Lock() - self._barrier = barrier - self._status_for = status_for or {} - - def post(self, url, json=None, timeout=None): - # The client POSTs to "{base}/collective_rpc"; record/key on the base url. - base = url[: -len(COLLECTIVE_RPC_ENDPOINT)] if url.endswith(COLLECTIVE_RPC_ENDPOINT) else url - if self._barrier is not None: - self._barrier.wait(timeout=5) - with self._lock: - self.calls.append((base, json)) - status = self._status_for.get(base, 200) - if status >= 400: - return _FakeResponse( - status_code=status, body={"error": {"message": "boom on " + base}}, reason="Server Error" - ) - return _FakeResponse() - - def close(self): - self.closed = True - - -def _client(urls, data_parallel_size=1, session=None): - c = SyncRdtControlPlaneClient(urls, data_parallel_size) - if session is not None: - c._session = session # swap the real requests.Session for the fake - return c - - -def _init_infos(session): - """Extract per-url init_info dicts from recorded init calls.""" - out = {} - for url, payload in session.calls: - assert payload["method"] == RDT_INIT_METHOD - out[url] = payload["kwargs"]["init_info"] - return out - - -def test_connection_close_header_set(): - c = _client(["http://a"]) - try: - assert c._session.headers.get("Connection") == "close" - finally: - c.close() - - -def test_init_fans_out_per_server_replica_rank(): - urls = ["http://a", "http://b", "http://c"] - sess = _RecordingSession() - c = _client(urls, data_parallel_size=1, session=sess) - try: - c.init_weight_transfer_engine({"num_consumers": 3, "names": ["w"]}) - finally: - c.close() - infos = _init_infos(sess) - assert set(infos) == set(urls) - # dp=1 -> one replica per server, distinct ranks 0..N-1, num_replicas=N. - assert sorted(i["replica_rank"] for i in infos.values()) == [0, 1, 2] - assert all(i["num_replicas"] == 3 for i in infos.values()) - # Shared fields untouched on every server. - assert all(i["num_consumers"] == 3 and i["names"] == ["w"] for i in infos.values()) - - -def test_init_replica_rank_is_per_deployment_under_dp(): - # 4 servers, dp=2 -> two deployments; the dp servers of a deployment share - # one replica_rank (server_index // dp), so ranks are [0, 0, 1, 1]. - urls = ["http://a", "http://b", "http://c", "http://d"] - sess = _RecordingSession() - c = _client(urls, data_parallel_size=2, session=sess) - try: - c.init_weight_transfer_engine({"k": "v"}) - finally: - c.close() - infos = _init_infos(sess) - assert [infos[u]["replica_rank"] for u in urls] == [0, 0, 1, 1] - assert all(i["num_replicas"] == 2 for i in infos.values()) - - -@pytest.mark.parametrize( - "call, method, expect_kwargs", - [ - (lambda c: c.start_weight_update(), RDT_START_METHOD, {"is_checkpoint_format": True}), - (lambda c: c.update_weights({"names": ["x"]}), RDT_UPDATE_METHOD, {"update_info": {"names": ["x"]}}), - (lambda c: c.finish_weight_update(), RDT_FINISH_METHOD, None), - ], -) -def test_uniform_calls_hit_every_server(call, method, expect_kwargs): - urls = ["http://a", "http://b"] - sess = _RecordingSession() - c = _client(urls, session=sess) - try: - call(c) - finally: - c.close() - assert {u for u, _ in sess.calls} == set(urls) - for _, payload in sess.calls: - assert payload["method"] == method - if expect_kwargs is None: - assert "kwargs" not in payload - else: - assert payload["kwargs"] == expect_kwargs - - -def test_fanout_is_concurrent(): - # If the client issued the POSTs serially, the barrier for N servers would - # never fill and each wait() would time out -> BrokenBarrierError. Reaching - # the barrier from all workers proves concurrent issue. - urls = ["http://a", "http://b", "http://c", "http://d"] - barrier = threading.Barrier(len(urls)) - sess = _RecordingSession(barrier=barrier) - c = _client(urls, session=sess) - try: - c.update_weights({"names": []}) # must not raise (barrier released) - finally: - c.close() - assert len(sess.calls) == len(urls) - - -def test_error_surfaces_body_message_and_drains_all(): - urls = ["http://a", "http://b", "http://c"] - sess = _RecordingSession(status_for={"http://b": 500}) - c = _client(urls, session=sess) - try: - with pytest.raises(RuntimeError) as ei: - c.update_weights({"names": []}) - finally: - c.close() - # Body error detail is surfaced, not the bare reason phrase. - assert "boom on http://b" in str(ei.value) - # All servers were still called (failure drains the whole fan-out). - assert {u for u, _ in sess.calls} == set(urls) - - -def test_close_releases_session(): - sess = _RecordingSession() - c = _client(["http://a"], session=sess) - c.close() - assert sess.closed is True - - -def test_empty_server_urls_rejected(): - with pytest.raises(ValueError): - SyncRdtControlPlaneClient([], 1) diff --git a/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_source.py b/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_source.py index c6283e17b3..5a938316d2 100644 --- a/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_source.py +++ b/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_source.py @@ -1,33 +1,38 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""`WeightSource` group-contract tests for the vendored sharded-RDT base. +"""`GroupedWeightSource` group-contract tests for the sharded-RDT base. `groups()` / `iter_groups()` define what a *group index* means -- the unit -`owned_groups()` names and that the RDT trainer gathers, publishes and frees by. +`held_names()` narrows and that the RDT trainer gathers, publishes and frees by. A source whose batches disagree with the trainer's own group partition does not return wrong data, it deadlocks the ranks sharing a gather collective, so the contract is pinned here. -Vendored alongside `sharded_rdt_base.py` from the vLLM RDT fork -(`tests/distributed/test_weight_transfer.py`). Needs no vLLM: torch only. +Adapted from the vLLM RDT fork (`tests/distributed/test_weight_transfer.py`). +Marked `vllm` because `sharded_rdt_base` now imports the base ABCs from the +wheel rather than vendoring them. """ import pytest import torch -from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_base import ( +pytest.importorskip("vllm", reason="sharded_rdt_base imports vLLM's trainer-side ABCs") + +pytestmark = pytest.mark.vllm + +from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_base import ( # noqa: E402 + GroupedWeightSource, ParamMeta, - WeightSource, layerwise_groups, ) -class TestWeightSourceGroupContract: - """`groups()` / `iter_groups()` on the WeightSource ABC. Group indices are +class TestGroupedWeightSourceContract: + """`groups()` / `iter_groups()` on the GroupedWeightSource ABC. Group indices are what backends gather and free by, and `held_names()` is what narrows them, so the default must agree with `layerwise_groups` over `metadata()`.""" - class _Source(WeightSource): + class _Source(GroupedWeightSource): """Minimal source over an ordered (name, tensor) list, optionally owning only some groups (in which case it iterates only those, per contract).""" @@ -131,10 +136,153 @@ class TestHeldNamesDefault: def test_the_default_is_none(self): names = ["embed.w", "model.layers.0.a"] - src = TestWeightSourceGroupContract._Source(names) + src = TestGroupedWeightSourceContract._Source(names) assert src.held_names() is None +class TestRdtFsdpWeightSource: + """RDT's FSDP source is the shared one re-ordered **group-major**. + + That reorder is not cosmetic: ``layerwise_groups`` must partition + ``metadata()`` exactly, because that partition IS the group index the + consumers' pull plans and the producer's free barrier are keyed on. A + ``state_dict()`` whose layers interleave with the pre/post block would + otherwise produce groups that are not contiguous runs of the metadata. + """ + + @staticmethod + def _model(names): + model = torch.nn.Module() + for name in names: + # register_parameter rejects dots, so nest a holder module per name. + holder = model + parts = name.split(".") + for part in parts[:-1]: + child = getattr(holder, part, None) + if child is None: + child = torch.nn.Module() + holder.add_module(part, child) + holder = child + holder.register_parameter(parts[-1], torch.nn.Parameter(torch.ones(2, 2))) + return model + + def _source(self, names): + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( + make_fsdp_weight_source, + ) + + return make_fsdp_weight_source(self._model(names), torch.bfloat16) + + # state_dict order with the layers out of numeric order, which is what the + # reorder actually has to fix. + _INTERLEAVED = [ + "model.embed_tokens.weight", + "model.layers.1.a", + "model.layers.0.a", + "lm_head.weight", + ] + + def test_reorders_into_group_major_order(self): + src = self._source(self._INTERLEAVED) + assert [m.name for m in src.metadata()] == [ + "model.embed_tokens.weight", + "model.layers.0.a", + "model.layers.1.a", + "lm_head.weight", + ] + + def test_groups_partition_metadata_exactly(self): + src = self._source(self._INTERLEAVED) + flat = [n for g in src.groups() for n in g] + assert flat == [m.name for m in src.metadata()] + assert src.groups() == [ + ["model.embed_tokens.weight"], + ["model.layers.0.a"], + ["model.layers.1.a"], + ["lm_head.weight"], + ] + + def test_unindexed_names_split_by_position_not_by_role(self): + """``layerwise_groups`` puts un-indexed names before the first indexed + one in ``pre`` and the rest in ``post`` — it does not know what an + embedding is. So a state_dict that emitted a layer before the embedding + would file the embedding under ``post``. Real ``state_dict()`` order is + definition order, so this does not arise; it is pinned because the + partition is what the consumers' pull plans are keyed on, and getting it + silently different per rank is the failure mode that matters.""" + src = self._source(["model.layers.0.a", "model.embed_tokens.weight"]) + assert src.groups() == [["model.layers.0.a"], ["model.embed_tokens.weight"]] + + def test_channels_agree_after_the_reorder(self): + src = self._source(self._INTERLEAVED) + meta = src.metadata() + pairs = list(src) + assert [m.name for m in meta] == [n for n, _ in pairs] + assert [m.shape for m in meta] == [tuple(t.shape) for _, t in pairs] + assert [m.dtype for m in meta] == [t.dtype for _, t in pairs] + + def test_iter_groups_batches_per_layer(self): + src = self._source(["model.layers.0.a", "model.layers.0.b", "model.layers.1.a"]) + assert [names for names, _ in src.iter_groups()] == [ + ["model.layers.0.a", "model.layers.0.b"], + ["model.layers.1.a"], + ] + + def test_holds_everything_by_default(self): + """FSDP replicates the whole model on every rank after the gather, so + there is no ownership to declare — unlike the PP/EP-local Megatron + source.""" + assert self._source(["model.layers.0.a"]).held_names() is None + + +class TestRdtMegatronWeightSource: + """The whole-model Megatron fallback: the shared source plus RDT's inherited + group channels, and nothing else. + + Worth a test of its own because it is a diamond — ``MegatronWeightSource`` + and ``GroupedWeightSource`` both derive from vLLM's ``WeightSource``, and + ``GroupedWeightSource`` re-declares ``metadata`` / ``__iter__`` abstract. If + the MRO put the abstract declarations first the class would not instantiate + at all. + """ + + class _Bridge: + def __init__(self, names): + self._names = names + + def export_hf_weights(self, module, show_progress=False, conversion_tasks=None): + assert conversion_tasks is None, "the fallback exports the whole model in one call" + return ((n, torch.ones(2, 2)) for n in self._names) + + def _source(self, names): + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( + RdtMegatronWeightSource, + ) + + return RdtMegatronWeightSource(self._Bridge(names), object(), torch.bfloat16) + + def test_instantiates_and_streams(self): + src = self._source(["model.embed_tokens.weight", "model.layers.0.a"]) + assert [m.name for m in src.metadata()] == ["model.embed_tokens.weight", "model.layers.0.a"] + assert [n for n, _ in src] == ["model.embed_tokens.weight", "model.layers.0.a"] + + def test_groups_partition_the_bridges_canonical_order(self): + """The bridge yields HF-canonical order, which is already + group-contiguous, so no reorder is needed here.""" + src = self._source(["model.embed_tokens.weight", "model.layers.0.a", "model.layers.1.a", "lm_head.weight"]) + assert src.groups() == [ + ["model.embed_tokens.weight"], + ["model.layers.0.a"], + ["model.layers.1.a"], + ["lm_head.weight"], + ] + + def test_holds_everything(self): + """Whole-model residency is the point for a pull backend: each producer + must be able to serve its bound consumer the complete model.""" + assert self._source(["model.layers.0.a"]).held_names() is None + + class TestExpertNameResolution: """Expert HF names come from the bridge's mapping registry, never an assumed layout: architectures differ (Kimi K2.5-VL nests its decoder stack under diff --git a/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_strategy.py b/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_strategy.py deleted file mode 100644 index a1b0f4421a..0000000000 --- a/tests/backends/skyrl_train/weight_sync/test_sharded_rdt_strategy.py +++ /dev/null @@ -1,247 +0,0 @@ -"""The sharded-RDT ``WeightTransferStrategy`` adapter: capability flags, init info, -and that ``send`` never materializes metadata or chunks. - -The rendezvous itself (Ray actors, NIXL, CUDA IPC) is covered by the producer/plan -suites and the GPU tests. -""" - -from dataclasses import dataclass -from typing import Optional - -import pytest - -from skyrl.backends.skyrl_train.weight_sync import ( - ShardedRdtInitInfo, - ShardedRdtTransferStrategy, - ShardedRdtWeightTransferSender, - get_transfer_strategy_cls, -) -from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import ( - WeightSyncInitInfo, - WeightTransferSender, -) - - -@dataclass -class _IeCfg: - """Just the fields the strategy reads.""" - - model_dtype: str = "torch.bfloat16" - run_engines_locally: bool = False - - -class _Client: - server_urls = ["http://a:1", "http://b:2"] - data_parallel_size = 1 - - -class _Extractor: - """Records whether the metadata/chunk channels were touched at all.""" - - def __init__(self) -> None: - self.metadata_calls = 0 - self.extract_calls = 0 - - def get_weight_metadata(self, dtype): - self.metadata_calls += 1 - return {"names": [], "dtype_names": [], "shapes": []} - - def extract_weights(self, dtype): - self.extract_calls += 1 - return iter(()) - - -class _FakeRdtSender: - """Stands in for RdtWeightSyncSender: records the calls it receives.""" - - def __init__(self) -> None: - self.sent = [] - self.torn_down = 0 - - async def send(self, weight_extractor): - self.sent.append(weight_extractor) - - def teardown(self): - self.torn_down += 1 - - -class TestInitInfo: - def test_carries_the_config_derived_args(self, monkeypatch): - monkeypatch.setattr( - "skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_strategy._ray_namespace", - lambda: "ns", - ) - info = ShardedRdtTransferStrategy.create_init_info(_IeCfg(), inference_world_size=16) - assert isinstance(info, ShardedRdtInitInfo) - assert info.model_dtype == "torch.bfloat16" - assert info.inference_world_size == 16 - assert info.trainer_actor_namespace == "ns" - - def test_a_missing_world_size_is_refused(self): - """The consumer count sizes every ownership decision; defaulting it would - mis-map consumers onto slices with nothing downstream to notice.""" - for bad in (None, 0): - with pytest.raises(ValueError, match="inference world size"): - ShardedRdtTransferStrategy.create_init_info(_IeCfg(), inference_world_size=bad) - - def test_base_model_path_is_ignored(self): - """RDT reads the live model, never a checkpoint, so the delta backend's - argument must not become a hidden requirement here.""" - info = ShardedRdtTransferStrategy.create_init_info( - _IeCfg(), inference_world_size=2, base_model_path="/does/not/matter" - ) - assert info.inference_world_size == 2 - - def test_override_existing_receiver_follows_run_engines_locally(self): - assert ShardedRdtTransferStrategy.create_init_info( - _IeCfg(run_engines_locally=False), inference_world_size=1 - ).override_existing_receiver - assert not ShardedRdtTransferStrategy.create_init_info( - _IeCfg(run_engines_locally=True), inference_world_size=1 - ).override_existing_receiver - - -class TestCapabilityFlags: - def test_the_sender_owns_the_inference_side_handshake(self): - """trainer_init opens the inference side itself, so worker rank 0 must not - also push init_info — doing both would init the engines twice.""" - assert ShardedRdtTransferStrategy.sender_initializes_receivers is True - assert get_transfer_strategy_cls("nccl", False).sender_initializes_receivers is False - - def test_every_strategy_accepts_the_weight_extractor(self): - """The worker passes `weight_extractor` to every strategy uniformly, so all - of them must accept the keyword even though only sharded_rdt uses it.""" - import inspect - - for backend, colocate in (("nccl", False), ("nccl", True), ("sharded_rdt", False)): - sig = inspect.signature(get_transfer_strategy_cls(backend, colocate).create_sender) - assert "weight_extractor" in sig.parameters, f"{backend}/{colocate} rejects weight_extractor" - - def test_expandable_segments_are_forced_off(self): - """The sidecar shares gathered tensors over CUDA IPC on every run, not only - under colocation, and VMM-backed memory makes that 5-10x slower.""" - assert ShardedRdtWeightTransferSender.force_disable_expandable_segments is True - assert WeightTransferSender.force_disable_expandable_segments is False - - def test_the_post_send_empty_cache_is_skipped(self): - """Publish buffers are reused by the next training step, so returning them to - CUDA costs 0.25-0.53s per rank at 235B and buys nothing.""" - assert ShardedRdtWeightTransferSender.empty_cache_after_send is False - assert WeightTransferSender.empty_cache_after_send is True - - def test_the_worker_still_resets_the_prefix_cache(self): - """Unlike delta, this sender does not touch the prefix cache, so the worker - must keep doing it.""" - assert ShardedRdtWeightTransferSender.handles_prefix_cache_reset is False - - -class TestSend: - @pytest.mark.asyncio - async def test_send_delegates_to_the_rdt_sender(self): - inner = _FakeRdtSender() - sender = ShardedRdtWeightTransferSender(inner) - extractor = _Extractor() - - await sender.send(extractor, "torch.bfloat16") - - assert inner.sent == [extractor] - - @pytest.mark.asyncio - async def test_send_never_materializes_metadata_or_chunks(self): - """get_weight_metadata on the Megatron extractor is a whole-model - export_hf_weights pass -- ~20s at 235B and model-sized memory, which is what - the RDT weight source avoids. The push backends' default send() calls it, so - this override must not.""" - extractor = _Extractor() - await ShardedRdtWeightTransferSender(_FakeRdtSender()).send(extractor, "torch.bfloat16") - - assert extractor.metadata_calls == 0 - assert extractor.extract_calls == 0 - - @pytest.mark.asyncio - async def test_the_push_backends_kwargs_are_accepted_and_ignored(self): - """The workers pass one kwarg set to every sender; an unexpected keyword - here would break the shared call site.""" - inner = _FakeRdtSender() - await ShardedRdtWeightTransferSender(inner).send(_Extractor(), "torch.bfloat16", reset_prefix_cache=True) - assert len(inner.sent) == 1 - - @pytest.mark.asyncio - async def test_send_chunks_refuses_with_an_explanation(self): - """There is no chunk stream to push; the error has to say what to call.""" - with pytest.raises(NotImplementedError, match="pull"): - await ShardedRdtWeightTransferSender(_FakeRdtSender()).send_chunks(iter(())) - - def test_teardown_forwards(self): - inner = _FakeRdtSender() - ShardedRdtWeightTransferSender(inner).teardown() - assert inner.torn_down == 1 - - -class TestCreateSender: - @staticmethod - def _info(namespace: Optional[str] = None) -> ShardedRdtInitInfo: - return ShardedRdtInitInfo( - override_existing_receiver=True, - model_dtype="torch.bfloat16", - inference_world_size=2, - trainer_actor_namespace=namespace, - ) - - def test_a_missing_extractor_is_a_loud_error(self): - """It would otherwise surface deep inside the rendezvous, or worse, as a - deferred init on the first send (which deadlocks).""" - with pytest.raises(RuntimeError, match="weight_extractor"): - ShardedRdtTransferStrategy.create_sender(self._info(), _Client(), weight_extractor=None) - - def test_a_foreign_init_info_is_rejected(self): - @dataclass - class _Other(WeightSyncInitInfo): - pass - - with pytest.raises(ValueError, match="ShardedRdtInitInfo"): - ShardedRdtTransferStrategy.create_sender( - _Other(override_existing_receiver=True), _Client(), weight_extractor=_Extractor() - ) - - def test_it_rendezvouses_eagerly_and_wraps_the_sender(self, monkeypatch): - """create_sender must initialize() before returning: every rank is inside - init_weight_sync_state here, which is the only window where rank 0 can wait - on the inference side without the others spinning in gather collectives.""" - built = {} - - class _Recorder: - def __init__(self, client, model_dtype, world_size, namespace): - built.update( - client=client, model_dtype=model_dtype, world_size=world_size, namespace=namespace, inited=None - ) - - def initialize(self, weight_extractor): - built["inited"] = weight_extractor - - monkeypatch.setattr( - "skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send.RdtWeightSyncSender", - _Recorder, - ) - extractor = _Extractor() - sender = ShardedRdtTransferStrategy.create_sender( - self._info(namespace="ns"), _Client(), weight_extractor=extractor - ) - - assert isinstance(sender, ShardedRdtWeightTransferSender) - assert built["model_dtype"] == "torch.bfloat16" - assert built["world_size"] == 2 - assert built["namespace"] == "ns" - assert built["inited"] is extractor - - -@pytest.mark.vllm -class TestVllmEngineMapping: - def test_it_returns_the_registered_consumer_engine(self): - """The consumers construct this class through vLLM's factory.""" - pytest.importorskip("vllm") - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_engine import ( - ShardedRDTWeightTransferEngine, - ) - - assert ShardedRdtTransferStrategy.get_vllm_transfer_engine() is ShardedRDTWeightTransferEngine diff --git a/tests/backends/skyrl_train/weight_sync/test_sources.py b/tests/backends/skyrl_train/weight_sync/test_sources.py new file mode 100644 index 0000000000..2d55f7415b --- /dev/null +++ b/tests/backends/skyrl_train/weight_sync/test_sources.py @@ -0,0 +1,179 @@ +"""Tests for the shared ``WeightSource`` implementations. + +Two properties matter: + +* **channel agreement.** ``metadata()`` declares what iteration will yield, and + the trainer engine sizes the worker's receive buffers -- and in packed mode + cuts its chunk boundaries -- from it. A source that disagrees splits the stream + differently on each side, which hangs in NCCL rather than erroring. +* **laziness.** The packed producers bound their memory only because they consume + the source lazily; a source that materialized eagerly would silently hold the + whole model. +""" + +import pytest +import torch + +pytest.importorskip("vllm", reason="sources.py implements vLLM's WeightSource contract") + +pytestmark = pytest.mark.vllm + +from skyrl.backends.skyrl_train.weight_sync import sources as sources_mod # noqa: E402 +from skyrl.backends.skyrl_train.weight_sync.sources import ( # noqa: E402 + FsdpWeightSource, + MegatronWeightSource, +) + + +def _model(): + model = torch.nn.Module() + model.register_parameter("w", torch.nn.Parameter(torch.ones(2, 3, dtype=torch.float32))) + model.register_parameter("b", torch.nn.Parameter(torch.zeros(3, dtype=torch.float32))) + return model + + +def _assert_channels_agree(source): + """The contract vLLM validates at runtime in ``_checked_iter``.""" + meta = source.metadata() + pairs = list(source) + assert len(meta) == len(pairs) + for m, (name, tensor) in zip(meta, pairs): + assert m.name == name + assert m.dtype == tensor.dtype + assert m.shape == tuple(tensor.shape) + return meta, pairs + + +class TestFsdpWeightSource: + def test_channels_agree(self): + meta, pairs = _assert_channels_agree(FsdpWeightSource(_model(), torch.bfloat16)) + assert [m.name for m in meta] == ["w", "b"] + assert [m.shape for m in meta] == [(2, 3), (3,)] + + def test_casts_to_the_inference_dtype(self): + source = FsdpWeightSource(_model(), torch.bfloat16) + # Declared as well as yielded: the worker allocates from metadata(). + assert {m.dtype for m in source.metadata()} == {torch.bfloat16} + assert {t.dtype for _, t in source} == {torch.bfloat16} + + def test_yields_contiguous_tensors(self): + # NCCL sends `numel` elements straight from data_ptr(), so a + # non-contiguous view would ship whatever follows its base pointer. + assert all(t.is_contiguous() for _, t in FsdpWeightSource(_model(), torch.bfloat16)) + + def test_weight_prefix_applies_to_both_channels(self): + source = FsdpWeightSource(_model(), torch.bfloat16, weight_prefix="language_model.") + meta, pairs = _assert_channels_agree(source) + assert [m.name for m in meta] == ["language_model.w", "language_model.b"] + # The prefix is a *wire* name; the state_dict is still keyed without it. + assert set(source.model.state_dict()) == {"w", "b"} + + def test_is_re_iterable(self): + source = FsdpWeightSource(_model(), torch.bfloat16) + assert [n for n, _ in source] == [n for n, _ in source] + + def test_metadata_does_not_gather(self, monkeypatch): + """FSDP2 ``DTensor.shape`` is already the global shape, so declaring the + stream must not run the gather collective. + + Asserted at the seam rather than with a fake DTensor: ``state_dict()`` + returns detached tensors, so a hand-attached ``full_tensor`` would not + survive it.""" + monkeypatch.setattr( + sources_mod, + "materialize_full_tensor", + lambda t: pytest.fail("metadata() must not materialize"), + ) + meta = FsdpWeightSource(_model(), torch.bfloat16).metadata() + assert [m.shape for m in meta] == [(2, 3), (3,)] + + def test_iteration_gathers_every_parameter(self, monkeypatch): + """The collective must run for every parameter on every rank: under + pipeline parallelism a rank may not own one, but iterating still drives + the collective its peers are waiting in.""" + gathered = [] + + def _spy(tensor): + gathered.append(tuple(tensor.shape)) + return tensor + + monkeypatch.setattr(sources_mod, "materialize_full_tensor", _spy) + list(FsdpWeightSource(_model(), torch.bfloat16)) + assert gathered == [(2, 3), (3,)] + + +class _FakeBridge: + """Stands in for Megatron-Bridge. Records calls so laziness is observable.""" + + def __init__(self, tensors, *, expect_module=None): + self._tensors = tensors + self.export_calls = 0 + self.yielded = 0 + self.tasks_args = [] + self._expect_module = expect_module + + def export_hf_weights(self, module, show_progress=False, conversion_tasks=None): + assert self._expect_module is None or module is self._expect_module + self.export_calls += 1 + self.tasks_args.append(conversion_tasks) + + def gen(): + for name, tensor in self._tensors: + self.yielded += 1 + yield name, tensor + + return gen() + + +class TestMegatronWeightSource: + def _tensors(self): + return [ + ("model.embed_tokens.weight", torch.ones(4, 2, dtype=torch.float32)), + ("model.layers.0.self_attn.q_proj.weight", torch.ones(2, 2, dtype=torch.float32)), + ] + + def test_channels_agree(self): + bridge = _FakeBridge(self._tensors()) + meta, pairs = _assert_channels_agree(MegatronWeightSource(bridge, object(), torch.bfloat16)) + assert [m.name for m in meta] == [n for n, _ in self._tensors()] + + def test_exports_the_whole_model_in_one_call(self): + """`_accumulate_grouped_export` needs every task of a `group_key` in one + call, or expert weights are silently never yielded.""" + bridge = _FakeBridge(self._tensors()) + list(MegatronWeightSource(bridge, object(), torch.bfloat16)) + assert bridge.tasks_args == [None] + + def test_metadata_caches_its_dry_export(self): + bridge = _FakeBridge(self._tensors()) + source = MegatronWeightSource(bridge, object(), torch.bfloat16) + first = source.metadata() + assert bridge.export_calls == 1 + # Engines call metadata() every round; it must not re-run the export. + assert source.metadata() is first + assert bridge.export_calls == 1 + + def test_iteration_is_lazy(self): + """Nothing may be pulled from the bridge until the consumer asks -- the + packed producers consume lazily into a fixed buffer, so streaming alone + bounds peak memory.""" + bridge = _FakeBridge(self._tensors()) + source = MegatronWeightSource(bridge, object(), torch.bfloat16) + it = iter(source) + assert bridge.yielded == 0 + next(it) + assert bridge.yielded == 1 + next(it) + assert bridge.yielded == 2 + + def test_is_re_iterable_and_re_exports(self): + bridge = _FakeBridge(self._tensors()) + source = MegatronWeightSource(bridge, object(), torch.bfloat16) + assert [n for n, _ in source] == [n for n, _ in source] + assert bridge.export_calls == 2 + + def test_reads_the_module_it_was_given(self): + module = object() + bridge = _FakeBridge(self._tensors(), expect_module=module) + list(MegatronWeightSource(bridge, module, torch.bfloat16)) + assert bridge.export_calls == 1 diff --git a/tests/backends/skyrl_train/weight_sync/test_trainer_engines.py b/tests/backends/skyrl_train/weight_sync/test_trainer_engines.py new file mode 100644 index 0000000000..e075eb654e --- /dev/null +++ b/tests/backends/skyrl_train/weight_sync/test_trainer_engines.py @@ -0,0 +1,368 @@ +"""Tests for ``weight_sync/trainer_engines.py``. + +What is worth pinning are the choices it encodes that are silent when wrong: + +* ``world_size = inference_world_size + 1`` (the trainer sender is a rank in the + NCCL group). Off by one and the rendezvous never completes. +* ``packed=True`` on IPC, which overrides vLLM's default. Unpacked IPC holds a + strong ref to a contiguous copy of every parameter until past + ``finish_weight_update``, i.e. the whole model resident on the trainer. +* which backends need a per-server init rewrite. +* the capability probes defaulting correctly for an engine that declares nothing. +""" + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("vllm", reason="trainer_engines builds vLLM trainer init infos") + +pytestmark = pytest.mark.vllm + +import torch # noqa: E402 +from vllm.distributed.weight_transfer.base import ParamMeta # noqa: E402 + +from skyrl.backends.skyrl_train.weight_sync.control_plane import ( # noqa: E402 + nccl_init_payloads, + rdt_init_payloads, +) +from skyrl.backends.skyrl_train.weight_sync.trainer_engines import ( # noqa: E402 + _build_init_info, + _packed_buffer_size_bytes, + engine_capability, + maybe_set_reset_prefix_cache, + teardown_engine, +) + +_1GiB = 1024**3 + + +class _Source: + """Only ``metadata()`` is read here; the engines are not constructed.""" + + def __init__(self, shapes=((4, 4),), dtype=None): + self.metadata_calls = 0 + self._meta = [ParamMeta(f"w{i}", dtype or torch.bfloat16, s) for i, s in enumerate(shapes)] + + def metadata(self): + self.metadata_calls += 1 + return self._meta + + def __iter__(self): + return iter(()) + + +def _init_info(backend, *, inference_world_size=4, ie_cfg=None, base_model_path=None, rank=0, source=None): + return _build_init_info( + backend=backend, + ie_cfg=ie_cfg if ie_cfg is not None else SimpleNamespace(), + rank=rank, + inference_world_size=inference_world_size, + source=source if source is not None else _Source(), + server_urls=["http://a"], + data_parallel_size=1, + base_model_path=base_model_path, + ) + + +class TestPackedBufferSize: + """A parameter too large for the buffer raises on the IPC path, and vLLM's + 1 GiB default is smaller than a large-vocab embedding matrix.""" + + def test_small_models_keep_vllms_default(self): + assert _packed_buffer_size_bytes(_Source(shapes=((4, 4),))) == _1GiB + + def test_grows_to_fit_the_largest_single_parameter(self): + # 151936 x 4096 bf16 = Qwen3-235B's embedding: 1.24 GiB, over the default. + source = _Source(shapes=((151936, 4096), (4096, 4096))) + assert _packed_buffer_size_bytes(source) == 151936 * 4096 * 2 + + def test_sizes_from_the_largest_not_the_total(self): + # Four 0.5 GiB parameters total 2 GiB but each fits the default; the + # buffer bounds one chunk, not the model. + half_gib_rows = (1024**3) // 2 // 2 // 1024 + source = _Source(shapes=((half_gib_rows, 1024),) * 4) + assert _packed_buffer_size_bytes(source) == _1GiB + + def test_an_empty_source_falls_back_to_the_default(self): + assert _packed_buffer_size_bytes(_Source(shapes=())) == _1GiB + + +class TestNcclInitInfo: + def test_world_size_counts_the_trainer_sender(self): + info, _ = _init_info("nccl", inference_world_size=4) + assert info.world_size == 5 + + def test_packed_is_on(self): + info, _ = _init_info("nccl") + assert info.packed is True + + def test_buffer_is_sized_from_the_source(self): + source = _Source(shapes=((151936, 4096),)) + info, _ = _init_info("nccl", source=source) + assert info.packed_buffer_size_bytes == 151936 * 4096 * 2 + # metadata() is a collective on a Megatron source, so a rank that + # skipped it would hang its peers. + assert source.metadata_calls == 1 + + def test_backend_is_vllms_own_key(self): + """Separate registries, so only the receive side takes a new name.""" + info, _ = _init_info("nccl") + assert info.backend == "nccl" + + def test_rank_decides_the_sender(self): + assert _init_info("nccl", rank=0)[0].is_sender is True + assert _init_info("nccl", rank=3)[0].is_sender is False + + def test_uses_the_per_server_rank_offset_rewrite(self): + _, payload_fn = _init_info("nccl") + assert payload_fn is nccl_init_payloads + + def test_picks_a_free_port(self): + first, _ = _init_info("nccl") + assert first.master_port > 0 + assert first.master_address + + +class TestIpcInitInfo: + def test_packed_is_forced_on(self): + info, payload_fn = _init_info("ipc") + assert info.packed is True + assert payload_fn is None + + def test_buffer_is_sized_from_the_source(self): + source = _Source(shapes=((151936, 4096),)) + info, _ = _init_info("ipc", source=source) + assert info.packed_buffer_size_bytes == 151936 * 4096 * 2 + + def test_backend_is_vllms_own_key(self): + assert _init_info("ipc")[0].backend == "ipc" + + +class TestShardedRdtInitInfo: + def test_uses_the_replica_rank_rewrite(self): + _, payload_fn = _init_info("sharded_rdt") + assert payload_fn is rdt_init_payloads + + def test_carries_the_consumer_count(self): + info, _ = _init_info("sharded_rdt", inference_world_size=8) + assert info.backend == "sharded_rdt" + assert info.num_consumers == 8 + + def test_a_missing_consumer_count_is_rejected(self): + """The ownership arithmetic is sized from it, and a wrong value silently + mis-maps consumers onto slices, so there is no safe default.""" + with pytest.raises(ValueError, match="inference world size"): + _init_info("sharded_rdt", inference_world_size=0) + + +def _delta_cfg(**overrides): + delta = SimpleNamespace( + sync_dir="/tmp/sync", + local_checkpoint_dir="/tmp/local", + publish_staging_dir="/tmp/staging", + max_file_size_in_gb=1.0, + cloud_download_workers=4, + publish_num_workers=None, + checkpoint_load_format="vllm_multi_thread_safetensors", + multi_thread_safetensors_max_workers=8, + ) + for key, value in overrides.items(): + setattr(delta, key, value) + return SimpleNamespace(delta_weight_sync=delta) + + +class TestDeltaInitInfo: + def test_does_not_touch_the_source(self): + """No wire buffer to size, and a Megatron ``metadata()`` is a whole-model + dry export.""" + source = _Source() + _init_info("delta", ie_cfg=_delta_cfg(), base_model_path="/m", source=source) + assert source.metadata_calls == 0 + + def test_carries_the_publisher_and_worker_settings(self): + info, payload_fn = _init_info("delta", ie_cfg=_delta_cfg(), base_model_path="/models/base") + assert info.backend == "delta" + assert info.base_model_path == "/models/base" + assert info.sync_dir == "/tmp/sync" + # Identical for every server, so no rewrite. + assert payload_fn is None + + def test_requires_a_base_model_path(self): + with pytest.raises(ValueError, match="base_model_path"): + _init_info("delta", ie_cfg=_delta_cfg()) + + def test_requires_a_sync_dir(self): + with pytest.raises(ValueError, match="sync_dir"): + _init_info("delta", ie_cfg=_delta_cfg(sync_dir=""), base_model_path="/models/base") + + def test_rejects_an_unsupported_load_format(self): + with pytest.raises(ValueError, match="checkpoint_load_format"): + _init_info("delta", ie_cfg=_delta_cfg(checkpoint_load_format="nope"), base_model_path="/m") + + +def test_unknown_backend_is_rejected(): + with pytest.raises(ValueError, match="Unknown weight sync backend"): + _init_info("telepathy") + + +def test_skyrl_trainer_engines_are_registered(): + """``delta`` and ``sharded_rdt`` are ours; vLLM registers the rest.""" + from vllm.distributed.weight_transfer.factory import WeightTransferTrainerFactory + + _init_info("ipc") # any call performs the registration + for name in ("nccl", "ipc", "delta", "sharded_rdt"): + assert name in WeightTransferTrainerFactory._registry + + +class TestBuildTrainerEngineResolvesTheBackend: + """``build_trainer_engine`` resolves the backend from the same two config + values the driver uses to configure the servers, and only then builds the + source. A mismatch here means the trainer and receive engines disagree, + which the driver has no way to catch.""" + + def _build(self, monkeypatch, weight_sync_backend, colocate_all): + from vllm.distributed.weight_transfer.factory import ( + WeightTransferTrainerFactory, + ) + + from skyrl.backends.skyrl_train.weight_sync.trainer_engines import ( + build_trainer_engine, + ) + + seen = {} + + def _fake_trainer_init(init_info, *, client, source=None): + seen["init_info"] = init_info + seen["source"] = source + return object() + + monkeypatch.setattr(WeightTransferTrainerFactory, "trainer_init", _fake_trainer_init) + + def source_factory(dtype, backend): + seen["factory_args"] = (dtype, backend) + return _Source() + + build_trainer_engine( + ie_cfg=SimpleNamespace(weight_sync_backend=weight_sync_backend, model_dtype="bfloat16"), + colocate_all=colocate_all, + rank=0, + inference_world_size=4, + source_factory=source_factory, + server_urls=["http://a"], + data_parallel_size=1, + base_model_path=None, + ) + return seen + + @pytest.mark.parametrize( + "weight_sync_backend,colocate_all,expected", + [ + ("nccl", False, "nccl"), + # The one resolution with no config field of its own. + ("nccl", True, "ipc"), + ("sharded_rdt", False, "sharded_rdt"), + ("rdt", False, "sharded_rdt"), + ], + ) + def test_resolution_reaches_both_the_factory_and_the_source( + self, monkeypatch, weight_sync_backend, colocate_all, expected + ): + seen = self._build(monkeypatch, weight_sync_backend, colocate_all) + assert seen["init_info"].backend == expected + # The source factory is told the SAME backend, so sharded RDT gets its + # ownership-aware subclass and nothing else does. + assert seen["factory_args"][1] == expected + + def test_source_factory_gets_the_inference_dtype(self, monkeypatch): + seen = self._build(monkeypatch, "nccl", False) + assert seen["factory_args"][0] is torch.bfloat16 + + def test_the_built_source_is_handed_to_the_engine(self, monkeypatch): + seen = self._build(monkeypatch, "nccl", False) + assert seen["source"] is not None + + +class _Bare: + """An engine that declares nothing — the shape of vLLM's own engines.""" + + +class TestCapabilityProbes: + def test_defaults_for_an_engine_that_declares_nothing(self): + engine = _Bare() + assert engine_capability(engine, "handles_prefix_cache_reset", False) is False + assert engine_capability(engine, "force_disable_expandable_segments", False) is False + assert engine_capability(engine, "empty_cache_after_send", True) is True + + def test_reads_a_declared_flag(self): + engine = _Bare() + engine.skyrl_empty_cache_after_send = False + assert engine_capability(engine, "empty_cache_after_send", True) is False + + def test_delta_declares_it_resets_the_prefix_cache(self): + from skyrl.backends.skyrl_train.weight_sync.delta_trainer import ( + DeltaTrainerWeightTransferEngine, + ) + + assert engine_capability(DeltaTrainerWeightTransferEngine, "handles_prefix_cache_reset", False) is True + + def test_rdt_declares_its_two_memory_flags(self): + from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_trainer import ( + ShardedRDTTrainerWeightTransferEngine as E, + ) + + assert engine_capability(E, "force_disable_expandable_segments", False) is True + assert engine_capability(E, "empty_cache_after_send", True) is False + + def test_set_reset_prefix_cache_is_optional(self): + # No setter must be a no-op, not an AttributeError. + maybe_set_reset_prefix_cache(_Bare(), True) + + class _WithSetter: + told = None + + def skyrl_set_reset_prefix_cache(self, reset): + self.told = reset + + engine = _WithSetter() + maybe_set_reset_prefix_cache(engine, True) + assert engine.told is True + + +class TestTeardown: + def test_shuts_down_the_engine_and_closes_its_client(self): + calls = [] + + class _Client: + def close(self): + calls.append("close") + + class _Engine: + client = _Client() + + def shutdown(self): + calls.append("shutdown") + + teardown_engine(_Engine()) + assert calls == ["shutdown", "close"] + + def test_closes_the_client_even_if_shutdown_raises(self): + """A half-torn-down engine must not leak the session and fan-out pool.""" + calls = [] + + class _Client: + def close(self): + calls.append("close") + + class _Engine: + client = _Client() + + def shutdown(self): + raise RuntimeError("boom") + + with pytest.raises(RuntimeError, match="boom"): + teardown_engine(_Engine()) + assert calls == ["close"] + + def test_none_is_a_no_op(self): + teardown_engine(None) diff --git a/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py b/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py deleted file mode 100644 index f73b810a05..0000000000 --- a/tests/backends/skyrl_train/weight_sync/test_transfer_strategies.py +++ /dev/null @@ -1,969 +0,0 @@ -import pytest - -from skyrl.backends.skyrl_train.weight_sync import ( - BroadcastInitInfo, - BroadcastTransferStrategy, - BroadcastWeightUpdateRequest, - CudaIpcInitInfo, - CudaIpcTransferStrategy, - CudaIpcWeightUpdateRequest, - DeltaInitInfo, - DeltaTransferStrategy, - LoraLoadRequest, - ShardedRdtTransferStrategy, - get_transfer_strategy, - get_transfer_strategy_cls, -) -from skyrl.train.config import InferenceEngineConfig -from skyrl.train.config.config import DeltaWeightSyncConfig - - -class TestGetTransferStrategyCls: - """Tests for get_transfer_strategy_cls function.""" - - @pytest.mark.parametrize( - "backend,colocate_all,expected_strategy", - [ - ("nccl", True, CudaIpcTransferStrategy), - ("nccl", False, BroadcastTransferStrategy), - ("gloo", True, BroadcastTransferStrategy), - ("gloo", False, BroadcastTransferStrategy), - ("delta", True, DeltaTransferStrategy), - ("delta", False, DeltaTransferStrategy), - ("sharded_rdt", False, ShardedRdtTransferStrategy), - # colocate_all is rejected elsewhere (build_vllm_cli_args); selection - # must still never hand sharded_rdt a push strategy. - ("sharded_rdt", True, ShardedRdtTransferStrategy), - ], - ) - def test_returns_correct_strategy(self, backend, colocate_all, expected_strategy): - """Should return correct strategy based on backend and colocate_all.""" - assert get_transfer_strategy_cls(backend, colocate_all) is expected_strategy - - @pytest.mark.parametrize( - "backend,colocate_all,expected", - [ - ("nccl", True, "ipc"), - ("nccl", False, "nccl"), - ("sharded_rdt", True, "sharded_rdt"), - ("sharded_rdt", False, "sharded_rdt"), - ], - ) - def test_backend_string(self, backend, colocate_all, expected): - """get_transfer_strategy maps to the vLLM WeightTransferConfig.backend string.""" - assert get_transfer_strategy(backend, colocate_all) == expected - - -class TestRdtSend: - """The sharded_rdt trainer-send glue in ``weight_sync/sharded_rdt/rdt_send.py`` - (the ``WeightSource`` implementations driven by ``RdtWeightSyncSender``), covered - without the vendored engine, whose ``trainer_init`` needs Ray + GPU.""" - - def test_weight_source_reorders_group_major(self): - """The FSDP WeightSource reorders metadata into group-major order (pre / - per-layer / post) so the vendored trainer's group-contiguity check passes.""" - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - _FsdpWeightSource, - ) - - class _FakeExtractor: - weight_prefix = "" - - def get_weight_metadata(self, dtype): - # Layer 1 before layer 0 => must reorder to pre / layer-0 / layer-1 / post. - return { - "names": [ - "model.embed_tokens.weight", - "model.layers.1.mlp.gate_proj.weight", - "model.layers.0.mlp.gate_proj.weight", - "lm_head.weight", - ], - "dtype_names": ["bfloat16", "bfloat16", "bfloat16", "bfloat16"], - "shapes": [[4, 8], [1, 8], [0, 8], [8, 4]], - } - - source = _FsdpWeightSource(_FakeExtractor(), torch.bfloat16) - meta = source.metadata() - assert [m.name for m in meta] == [ - "model.embed_tokens.weight", - "model.layers.0.mlp.gate_proj.weight", - "model.layers.1.mlp.gate_proj.weight", - "lm_head.weight", - ] - # shapes travel with their names through the reorder; dtype is the wire dtype. - assert [list(m.shape) for m in meta] == [[4, 8], [0, 8], [1, 8], [8, 4]] - assert all(m.dtype is torch.bfloat16 for m in meta) - - def test_megatron_weight_source_streams_bridge_export(self): - """MegatronWeightSource wraps the extractor's Megatron-Bridge: metadata() - and iteration both run the non-bucketed export (so their order agrees), and - iteration casts each full tensor to the wire dtype.""" - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - MegatronWeightSource, - ) - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_base import ( - layerwise_groups, - ) - - # HF-canonical (group-contiguous) order, fp32 source tensors. - items = [ - ("model.embed_tokens.weight", torch.ones(4, 8, dtype=torch.float32)), - ("model.layers.0.mlp.gate_proj.weight", torch.ones(2, 8, dtype=torch.float32)), - ("model.layers.1.mlp.gate_proj.weight", torch.ones(2, 8, dtype=torch.float32)), - ("lm_head.weight", torch.ones(8, 4, dtype=torch.float32)), - ] - export_calls = [] - - class _FakeBridge: - def export_hf_weights(self, module, show_progress=False, conversion_tasks=None): - # RDT must use the NON-bucketed export (conversion_tasks=None) so - # MoE-expert grouping doesn't break group-major contiguity. - assert conversion_tasks is None - export_calls.append(module) - for name, tensor in items: - yield name, tensor - - class _FakeMegatronExtractor: - bridge = _FakeBridge() - actor_module = object() - - source = MegatronWeightSource(_FakeMegatronExtractor(), torch.bfloat16) - - meta = source.metadata() - names = [m.name for m in meta] - assert names == [n for n, _ in items] # order preserved (no reorder) - assert [list(m.shape) for m in meta] == [[4, 8], [2, 8], [2, 8], [8, 4]] - assert all(m.dtype is torch.bfloat16 for m in meta) - # Order is already group-contiguous -> layerwise_groups partitions it exactly - # (this is what the trainer engine's trainer_init validates). - assert [n for g in layerwise_groups(names) for n in g] == names - - yielded = list(source) - assert [n for n, _ in yielded] == names - assert all(t.dtype is torch.bfloat16 and t.is_contiguous() for _, t in yielded) - # metadata() cached: iteration ran a fresh export, so the bridge was - # exported twice total (one dry-run for metadata, one for the stream). - assert len(export_calls) == 2 - - def test_make_weight_source_selects_by_extractor_flavor(self): - """make_weight_source picks Megatron (has .bridge) vs FSDP (has .model).""" - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - MegatronWeightSource, - _FsdpWeightSource, - make_weight_source, - ) - - class _FakeBridge: - def export_hf_weights(self, module, show_progress=False, conversion_tasks=None): - return iter(()) - - class _FakeMegatronExtractor: - bridge = _FakeBridge() - actor_module = object() - - class _FakeFsdpExtractor: - weight_prefix = "" - model = object() - - def get_weight_metadata(self, dtype): - return {"names": [], "dtype_names": [], "shapes": []} - - assert isinstance(make_weight_source(_FakeMegatronExtractor(), torch.bfloat16), MegatronWeightSource) - assert isinstance(make_weight_source(_FakeFsdpExtractor(), torch.bfloat16), _FsdpWeightSource) - - -class TestRdtReplicaConsumerMapping: - """The per-replica consumer identity the engine computes from the injected - replica_rank/num_replicas must give every worker in a multi-engine fleet a - DISTINCT global id and a correct 1:1 producer binding (the fix for the - multi-engine deadlock). This mirrors the engine's arithmetic over the shared - M:N helpers, so it runs without a GPU/vLLM.""" - - @staticmethod - def _consumer_id(replica_rank, num_replicas, num_consumers, local_index): - # Mirrors ShardedRDTWeightTransferEngine.init_transfer_engine. - workers_per_replica = num_consumers // max(1, num_replicas) - return replica_rank * workers_per_replica + local_index - - def test_two_dense_engines_bind_distinct_producers(self): - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_common import ( - RdtRouter, - assign_producer_indices, - ) - - # 2 independent TP=1 engines: each engine's local index is 0, so the - # replica_rank offset is what separates them into consumer ids 0 and 1. - num_consumers, num_producers, num_replicas = 2, 2, 2 - cids = [self._consumer_id(r, num_replicas, num_consumers, 0) for r in range(2)] - assert cids == [0, 1] - # Each consumer binds its own producer; each producer serves exactly one. - assert assign_producer_indices(num_producers, num_consumers, cids[0]) == [0] - assert assign_producer_indices(num_producers, num_consumers, cids[1]) == [1] - names = ["g0.w", "g1.w", "g2.w"] - router = RdtRouter(num_producers, num_consumers, None, None, names, [1, 1, 1]) - assert router.producer_for(cids[0], "g0.w") == 0 - assert router.producer_for(cids[1], "g0.w") == 1 - - def test_single_replica_offset_is_zero(self): - # num_replicas=1 (default / single deployment) => offset 0, id == local index. - assert self._consumer_id(0, 1, 4, 3) == 3 - - def test_multi_engine_multi_worker_ids_are_contiguous(self): - # 2 engines x TP=2 = 4 consumers; ids must cover 0..3 with no collision. - num_consumers, num_replicas = 4, 2 - ids = [self._consumer_id(r, num_replicas, num_consumers, local) for r in range(2) for local in range(2)] - assert sorted(ids) == [0, 1, 2, 3] - - -class TestPpLocalOwnership: - """``MegatronStackedWeightSource`` ownership detection (the PP grain of - SKYRL_RDT_SHARD_AWARE). - - In PP-local mode a stage exports only its own parameters, so the source has - to (a) rebuild WHOLE-model metadata from what the stages exchange — the RDT - contract requires every rank to describe the whole model — and (b) notice when - one gather group is produced by two stages, which cannot be served per-stage. - Both are exercised against the assembly directly (the walk itself needs - Megatron + GPUs).""" - - @staticmethod - def _source(pp_size, my_pp, gathered): - """A source in PP-local mode with the stages' exchange stubbed out. - - ``metadata()`` is pre-populated the way the real one does it (the walk's - local result handed to ``_assemble_pp_metadata``), so the assembly and - ``held_names`` can be exercised without Megatron or a GPU.""" - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - MegatronStackedWeightSource, - ) - - src = MegatronStackedWeightSource.__new__(MegatronStackedWeightSource) - src._dtype = torch.bfloat16 - src._pp_local = True - src._ep_local = True # demotion must flip BOTH shard-aware grains - src._ep_size = 1 - src._my_ep_rank = 0 - src._demoted = False - src._verified = True - src._group_stages = [] - src._owned_group_idx = [] - src._pp_geometry = lambda: (pp_size, my_pp) # type: ignore[method-assign] - src._exchange_pp_names = lambda mine: gathered # type: ignore[method-assign] - src._meta = src._assemble_pp_metadata([]) - return src - - def test_metadata_is_the_whole_model_group_major_on_every_stage(self): - """Stage 1 walks only its own layer, but metadata() must come back as the - whole model in group-major order — identical on both stages, since the - engine cross-checks a digest of it and bakes the consumers' plan from one - rank's copy.""" - stage0 = [("model.embed_tokens.weight", [8, 4]), ("model.layers.0.w", [4, 4])] - stage1 = [("model.layers.1.w", [4, 4]), ("model.norm.weight", [4])] - expected = [ - "model.embed_tokens.weight", - "model.layers.0.w", - "model.layers.1.w", - "model.norm.weight", - ] - for my_pp in (0, 1): - meta = self._source(2, my_pp, [stage0, stage1]).metadata() - assert [m.name for m in meta] == expected - assert [tuple(m.shape) for m in meta] == [(8, 4), (4, 4), (4, 4), (4,)] - # ... and each stage claims exactly the names it produced. - assert self._source(2, 0, [stage0, stage1]).held_names() == [ - "model.embed_tokens.weight", - "model.layers.0.w", - ] - assert self._source(2, 1, [stage0, stage1]).held_names() == [ - "model.layers.1.w", - "model.norm.weight", - ] - - def test_metadata_is_group_contiguous_even_when_a_stage_holds_both_ends(self): - """The assembled order must satisfy the engine's group-contiguity check — - ``flat(layerwise_groups(names)) == names`` — for whatever the stages - produce, and must be the same list on every stage. - - Here stage 0 holds the output block too (a tied-embedding layout), so it - yields two non-layer names before any layer exists. ``layerwise_groups`` - splits pre/post by POSITION, so those land in one leading group rather - than a pre and a post block. That is still a valid partition — ownership - follows it, and both sides derive it from the same list — which is why the - invariant to hold is contiguity, not a canonical pre/layers/post shape.""" - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_base import ( - layerwise_groups, - ) - - stage0 = [("model.embed_tokens.weight", [8, 4]), ("model.lm_head.weight", [8, 4])] - stage1 = [("model.layers.0.w", [4, 4])] - per_stage = [] - for my_pp in (0, 1): - src = self._source(2, my_pp, [stage0, stage1]) - names = [m.name for m in src.metadata()] - assert [n for g in layerwise_groups(names) for n in g] == names - per_stage.append(names) - assert per_stage[0] == per_stage[1] - # Stage 0 produced both names of the leading group; stage 1 the layer. - src = self._source(2, 1, [stage0, stage1]) - assert src.held_names() == ["model.layers.0.w"] - assert src._group_stages == [{0}, {1}] - - def test_a_group_produced_by_two_stages_disables_pp_local(self): - """Tied embeddings / MTP put one group's names on two stages. Serving that - per-stage would publish half a group, so the source must fall back to - gather-to-all instead of silently truncating it.""" - stage0 = [("model.embed_tokens.weight", [8, 4]), ("model.layers.0.w", [4, 4])] - # Stage 1 also produces a post-block name -> the post group spans stages. - stage1 = [("model.layers.1.w", [4, 4]), ("model.norm.weight", [4])] - stage0 = stage0 + [("model.lm_head.weight", [8, 4])] - src = self._source(2, 0, [stage0, stage1]) - assert src.held_names() is None - assert src._demoted is True - assert src._pp_local is False - # BOTH shard-aware grains demote together: a stamped name a rank no - # longer serves per-stage would misroute pulls. - assert src._ep_local is False - # Re-asking is still None: a demoted source holds everything. - assert src.held_names() is None - # Metadata is still the whole model, so the digest check still passes. - assert [m.name for m in src.metadata()] == [ - "model.embed_tokens.weight", - "model.layers.0.w", - "model.layers.1.w", - "model.lm_head.weight", - "model.norm.weight", - ] - - def test_a_tied_name_on_two_stages_is_not_duplicated(self): - """A weight both stages hold (Megatron keeps a copy of a tied embedding on - the last stage) must appear ONCE in metadata — a duplicate name would give - the consumers two plan entries for one tensor — and mark its group shared.""" - tied = ("model.embed_tokens.weight", [8, 4]) - src = self._source(2, 0, [[tied, ("model.layers.0.w", [4, 4])], [tied]]) - assert [m.name for m in src.metadata()] == ["model.embed_tokens.weight", "model.layers.0.w"] - assert src._group_stages[0] == {0, 1} - assert src.held_names() is None - - def test_walk_is_reordered_into_partition_order(self): - """The bridge streams a stage's tasks in ITS order, which need not match the - partition: at 235B the last stage exports the output block BEFORE its layers, - while layerwise_groups places that block last. The gather loop walks the held - groups ascending and raises on anything else, so the walk must be reordered.""" - stage0 = [("model.embed_tokens.weight", [8, 4]), ("model.layers.0.w", [4, 4])] - stage1 = [("model.norm.weight", [4]), ("model.layers.1.w", [4, 4])] - src = self._source(2, 1, [stage0, stage1]) - # layer 1 and the post block, in partition order - assert src.held_names() == ["model.layers.1.w", "model.norm.weight"] - - # Stage 1's walk emits the post block first, as the real bridge does. - walk = iter([(["model.norm.weight"], ["N"]), (["model.layers.1.w"], ["L1"])]) - assert list(src._walk_in_group_order(walk)) == [ - (["model.layers.1.w"], ["L1"]), - (["model.norm.weight"], ["N"]), - ] - - def test_walk_reorder_refuses_to_hold_layer_stacks(self): - """Deferring a group pins its gathered tensors (~4.6 GiB for a 235B layer), - so an unexpected permutation must raise rather than quietly inflate trainer - memory.""" - gathered = [[(f"model.layers.{i}.w", [4, 4]) for i in range(5)], []] - src = self._source(2, 0, gathered) - assert src.held_names() == [f"model.layers.{i}.w" for i in range(5)] - # Every group arrives in reverse: nothing can be released. - walk = iter([([f"model.layers.{i}.w"], [i]) for i in (4, 3, 2, 1, 0)]) - with pytest.raises(RuntimeError, match="groups ahead of the partition order"): - list(src._walk_in_group_order(walk)) - - def test_walk_reorder_rejects_a_name_outside_the_partition(self): - stage0 = [("model.layers.0.w", [4, 4])] - src = self._source(1, 0, [stage0]) - walk = iter([(["model.layers.9.w"], ["X"])]) - with pytest.raises(RuntimeError, match="not in the assembled partition"): - list(src._walk_in_group_order(walk)) - - def test_a_demoted_source_refuses_to_iterate(self): - """A demoted source cannot serve (its gather paths are gone); iterating - it is a wiring bug — make_weight_source should have delegated to the - plain MegatronWeightSource.""" - src = self._source(2, 0, [[("a.weight", [2])], [("a.weight", [2])]]) - assert src.held_names() is None # tied name on both stages -> demoted - assert src._demoted is True - with pytest.raises(RuntimeError, match="demoted"): - list(src.iter_groups()) - - -class TestExpertOwnership: - """The EP grain of shard-aware serving: `held_names()` + the - walk's real-vs-None emission must agree (the stamps are what the consumers - route by, the Nones are what the trainer drops before publishing).""" - - @staticmethod - def _stub(ep_size=2, my_ep_rank=1, pp_local=False): - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - MegatronStackedWeightSource, - ) - - src = MegatronStackedWeightSource.__new__(MegatronStackedWeightSource) - src._dtype = torch.bfloat16 - src._pp_local = pp_local - src._ep_local = True - src._ep_size = ep_size - src._my_ep_rank = my_ep_rank - src._demoted = False - # These tests are about the EP ownership grain, not about name resolution: - # the source has no bridge, so seed the per-layer expert-name cache with the - # Qwen3-MoE layout the assertions below use. (Production always resolves - # these through the bridge's mapping registry and raises if it cannot — - # there is no synthesized fallback to lean on here.) - src._expert_names = { - layer: [ - name - for e in range(ep_size * 2) - for name in ( - f"model.layers.{layer}.mlp.experts.{e}.gate_proj.weight", - f"model.layers.{layer}.mlp.experts.{e}.up_proj.weight", - f"model.layers.{layer}.mlp.experts.{e}.down_proj.weight", - ) - ] - for layer in range(8) - } - src._expert_name_source = "test stub" - src._layer_geom = {} - src._phase = {} - src._phase_prefix = "" - return src - - @staticmethod - def _meta_of(names): - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_base import ( - ParamMeta, - ) - - return [ParamMeta(n, torch.bfloat16, (2, 2)) for n in names] - - def test_stamps_are_expert_index_over_n_local_and_minus_one_elsewhere(self): - src = self._stub(ep_size=2, my_ep_rank=1) - src._meta = self._meta_of( - ["embed.weight"] - + [f"model.layers.0.mlp.experts.{e}.gate_proj.weight" for e in range(4)] - + ["model.layers.0.input_layernorm.weight", "lm_head.weight"] - ) - owners = src._name_owner() - assert src._my_ep_rank == 1 - assert owners == [-1, 0, 0, 1, 1, -1, -1] - - def test_an_expert_count_not_divisible_by_ep_size_raises(self): - src = self._stub(ep_size=2) - src._meta = self._meta_of([f"model.layers.0.mlp.experts.{e}.up_proj.weight" for e in range(3)]) - with pytest.raises(RuntimeError, match="not divisible"): - src._name_owner() - - def test_ep_local_off_returns_none(self): - src = self._stub() - src._ep_local = False - src._pp_local = False - assert src.held_names() is None - - def test_the_walk_materializes_exactly_the_stamped_experts(self): - """The truthfulness clause of the ABC contract: within an owned layer, a - name stamped my_ep_rank yields a real view of the LOCAL stack; every - other expert's entries are None. Zero collectives on this path.""" - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - _ExpertLayer, - ) - - class _Task: - def __init__(self, t): - self.param_weight = t - - F, H, n_local = 3, 2, 2 - src = self._stub(ep_size=2, my_ep_rank=1) - fc1 = [_Task(torch.full((2 * F, H), float(10 + i))) for i in range(n_local)] - fc2 = [_Task(torch.full((H, F), float(20 + i))) for i in range(n_local)] - lay = _ExpertLayer(layer=0, fc1=fc1, fc2=fc2, owned=True, n_local=n_local, F=F, H=H, ep_size=2) - - names, tensors = [], [] - src._extend_layer_experts(lay, None, names, tensors) - - E = 4 - assert len(names) == len(tensors) == 3 * E - by_name = dict(zip(names, tensors)) - # Foreign coordinate (0): experts 0..1 are None. - for e in (0, 1): - for proj in ("gate_proj", "up_proj", "down_proj"): - assert by_name[f"model.layers.0.mlp.experts.{e}.{proj}.weight"] is None - # Own coordinate (1): experts 2..3 are real views with today's shapes. - for i, e in enumerate((2, 3)): - gate = by_name[f"model.layers.0.mlp.experts.{e}.gate_proj.weight"] - up = by_name[f"model.layers.0.mlp.experts.{e}.up_proj.weight"] - down = by_name[f"model.layers.0.mlp.experts.{e}.down_proj.weight"] - assert gate.shape == (F, H) and up.shape == (F, H) and down.shape == (H, F) - assert torch.equal(gate, torch.full((F, H), float(10 + i), dtype=torch.bfloat16)) - assert torch.equal(down, torch.full((H, F), float(20 + i), dtype=torch.bfloat16)) - - def test_foreign_expert_shapes_are_synthesized_from_local_geometry(self): - """metadata() cannot read .shape off a None; the walk records each - layer's (F, H) before the expert names are emitted.""" - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - _ExpertLayer, - ) - - class _Task: - def __init__(self, t): - self.param_weight = t - - F, H = 3, 2 - src = self._stub(ep_size=2, my_ep_rank=0) - lay = _ExpertLayer( - layer=7, - fc1=[_Task(torch.zeros((2 * F, H)))], - fc2=[_Task(torch.zeros((H, F)))], - owned=True, - n_local=1, - F=F, - H=H, - ep_size=2, - ) - src._extend_layer_experts(lay, None, [], []) - assert src._foreign_expert_shape("model.layers.7.mlp.experts.1.gate_proj.weight") == (F, H) - assert src._foreign_expert_shape("model.layers.7.mlp.experts.1.down_proj.weight") == (H, F) - - -class TestHeldNamesComposition: - """``held_names`` is the misroute guard's source of truth: exactly the names - this rank publishes — its stage's groups, narrowed to the replicated names - plus its own coordinate's experts. The trainer copies it verbatim into the - ``served_names`` it hands the sidecar.""" - - @staticmethod - def _source(ep_size, my_ep_rank, names, *, pp_local=False, owned=None): - import torch - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.rdt_send import ( - MegatronStackedWeightSource, - ) - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_base import ( - ParamMeta, - ) - - src = MegatronStackedWeightSource.__new__(MegatronStackedWeightSource) - src._dtype = torch.bfloat16 - src._pp_local = pp_local - src._ep_local = ep_size > 1 - src._ep_size = ep_size - src._my_ep_rank = my_ep_rank - src._demoted = False - src._expert_names = {} - src._layer_geom = {} - src._phase = {} - src._phase_prefix = "" - src._meta = [ParamMeta(n, torch.bfloat16, (2, 2)) for n in names] - if owned is not None: - src._owned_group_idx = owned - src._group_stages = [] - return src - - def test_ep_local_holds_replicated_names_plus_its_own_experts(self): - names = [ - "model.layers.0.input_layernorm.weight", - "model.layers.0.mlp.experts.0.gate_proj.weight", - "model.layers.0.mlp.experts.1.gate_proj.weight", - "model.norm.weight", - ] - src = self._source(2, 1, names) - assert src.held_names() == [ - "model.layers.0.input_layernorm.weight", - "model.layers.0.mlp.experts.1.gate_proj.weight", - "model.norm.weight", - ] - - def test_the_other_coordinate_holds_the_complement(self): - names = [ - "model.layers.0.mlp.experts.0.gate_proj.weight", - "model.layers.0.mlp.experts.1.gate_proj.weight", - ] - held0 = self._source(2, 0, names).held_names() - held1 = self._source(2, 1, names).held_names() - assert held0 == [names[0]] and held1 == [names[1]] - assert sorted(held0 + held1) == sorted(names) - - def test_neither_grain_holds_everything(self): - names = ["a", "model.layers.0.w", "b"] - assert self._source(1, 0, names).held_names() is None - - -@pytest.mark.vllm -class TestStampedYieldValidation: - """The gather loop checks stamps against yields per group — the ABC's - truthfulness invariant, enforced where both sit side by side. Without it a - stamps/yield mismatch is a 300s stall-watchdog death instead of an - immediate, named error.""" - - @staticmethod - def _engine(held): - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_trainer import ( - ShardedRDTTrainerWeightTransferEngine, - ) - - e = ShardedRDTTrainerWeightTransferEngine.__new__(ShardedRDTTrainerWeightTransferEngine) - e._held_names = held - return e - - def test_matching_yields_pass(self): - import torch - - e = self._engine({"norm", "e1"}) - e._validate_held_yields(0, ["norm", "e0", "e1"], [torch.zeros(1), None, torch.zeros(1)]) - - def test_a_held_name_yielding_none_raises(self): - """The dangerous direction: served_names advertises the name, consumers - route pulls here, and the cache wait would never complete.""" - import torch - - e = self._engine({"norm", "e0"}) - with pytest.raises(RuntimeError, match="disagrees with the yielded tensors"): - e._validate_held_yields(0, ["norm", "e0"], [torch.zeros(1), None]) - - def test_a_foreign_name_yielding_a_tensor_raises(self): - import torch - - e = self._engine({"norm"}) - with pytest.raises(RuntimeError, match="disagrees with the yielded tensors"): - e._validate_held_yields(0, ["norm", "e0"], [torch.zeros(1), torch.zeros(1)]) - - def test_unstamped_sources_are_never_checked(self): - e = self._engine(None) - e._validate_held_yields(0, ["anything"], [None]) - - -class TestQkvIndexDeviceCtx: - """``_qkv_index_device_ctx`` keeps the QKV split's index tensors on the weight's - device instead of the host, which is worth ~0.65s/sync of ne_bridge at 235B (a - CPU index tensor against a CUDA weight forces an H2D copy + stream sync per - gather). It deliberately copies NO upstream logic — it only changes where - ``torch.arange`` allocates — so these cover the wrapping, the injection, and - the restore. A meta device stands in for CUDA so this needs no GPU.""" - - @staticmethod - def _fake_modules(monkeypatch): - """Stub split fns on the real modules, so the context wraps something we can - observe. Records the device each torch.arange call landed on.""" - import torch - from megatron.bridge.models.conversion import param_mapping as pm - - seen = [] - - def _split(config, qkv, *a, **kw): - seen.append(torch.arange(4).device.type) - return ("q", "k", "v") - - for name in ("split_qkv_weights", "split_qkv_biases", "split_qkv_weights_scale"): - monkeypatch.setattr(pm, name, _split, raising=False) - return pm, seen - - def test_index_tensors_follow_the_weight_device(self, monkeypatch): - import torch - - pytest.importorskip("megatron.bridge.models.conversion.param_mapping") - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_send - - pm, seen = self._fake_modules(monkeypatch) - weight = torch.empty(2, 2, device="meta") - with rdt_send._qkv_index_device_ctx(): - pm.split_qkv_weights(None, weight) - assert seen == ["meta"], "arange should have been redirected to the weight's device" - - def test_cpu_weights_are_left_alone(self, monkeypatch): - """The redirect must not fire for a host weight — there is nothing to fix and - forcing a device would be a behaviour change.""" - import torch - - pytest.importorskip("megatron.bridge.models.conversion.param_mapping") - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_send - - pm, seen = self._fake_modules(monkeypatch) - with rdt_send._qkv_index_device_ctx(): - pm.split_qkv_weights(None, torch.empty(2, 2)) - assert seen == ["cpu"] - - def test_originals_and_torch_arange_are_restored(self, monkeypatch): - """torch.arange is patched process-wide for the duration of ONE call, so a - leak would silently put every later index tensor on a device.""" - import torch - - pytest.importorskip("megatron.bridge.models.conversion.param_mapping") - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_send - - pm, _seen = self._fake_modules(monkeypatch) - before, real_arange = pm.split_qkv_weights, torch.arange - with rdt_send._qkv_index_device_ctx(): - assert pm.split_qkv_weights is not before # wrapped - pm.split_qkv_weights(None, torch.empty(2, 2, device="meta")) - assert torch.arange is real_arange, "arange must be restored after each call" - assert pm.split_qkv_weights is before - assert torch.arange is real_arange - assert torch.arange(3).device.type == "cpu" - - -@pytest.mark.vllm -class TestShardedRdtVllmRegistration: - """The factory registration (requires the vLLM wheel).""" - - def test_engine_registered(self): - pytest.importorskip("vllm") - # Importing the weight_sync package's register module runs ensure_registered(). - from vllm.config import WeightTransferConfig - from vllm.distributed.weight_transfer import WeightTransferEngineFactory - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_vllm_register - - rdt_vllm_register.ensure_registered() - assert "sharded_rdt" in WeightTransferEngineFactory._registry - # vLLM 0.23.0 already accepts arbitrary backend strings (Literal | str); - # no runtime relaxation needed, and the built-ins still validate. - assert WeightTransferConfig(backend="sharded_rdt").backend == "sharded_rdt" - assert WeightTransferConfig(backend="nccl").backend == "nccl" - assert WeightTransferConfig(backend="ipc").backend == "ipc" - - def test_the_registered_module_path_actually_resolves(self): - """`register_engine` stores the engine's module path as a string and imports - it only when a worker constructs the backend, so a stale path fails on a real - inference worker rather than here. Force the loader to resolve it.""" - pytest.importorskip("vllm") - from vllm.distributed.weight_transfer import WeightTransferEngineFactory - - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt import rdt_vllm_register - from skyrl.backends.skyrl_train.weight_sync.sharded_rdt.sharded_rdt_engine import ( - ShardedRDTWeightTransferEngine, - ) - - rdt_vllm_register.ensure_registered() - loader = WeightTransferEngineFactory._registry["sharded_rdt"] - assert loader() is ShardedRDTWeightTransferEngine - - -class TestCreateInitInfo: - """Tests for create_init_info static methods.""" - - def _make_ie_cfg( - self, - weight_sync_backend: str = "nccl", - model_dtype: str = "torch.bfloat16", - num_engines: int = 1, - tensor_parallel_size: int = 1, - pipeline_parallel_size: int = 1, - data_parallel_size: int = 1, - run_engines_locally: bool = True, - ): - """Create an InferenceEngineConfig for create_init_info.""" - return InferenceEngineConfig( - weight_sync_backend=weight_sync_backend, - model_dtype=model_dtype, - num_engines=num_engines, - tensor_parallel_size=tensor_parallel_size, - pipeline_parallel_size=pipeline_parallel_size, - data_parallel_size=data_parallel_size, - run_engines_locally=run_engines_locally, - ) - - def test_cuda_ipc_create_init_info(self): - """CudaIpcTransferStrategy.create_init_info should create CudaIpcInitInfo with model_dtype_str.""" - ie_cfg = self._make_ie_cfg(model_dtype="torch.float32") - init_info = CudaIpcTransferStrategy.create_init_info(ie_cfg) - - assert isinstance(init_info, CudaIpcInitInfo) - assert init_info.model_dtype_str == "torch.float32" - - def test_broadcast_create_init_info(self, monkeypatch): - """BroadcastTransferStrategy.create_init_info should create BroadcastInitInfo with correct fields.""" - # Mock ray to avoid actual network operations - import skyrl.backends.skyrl_train.weight_sync.broadcast_strategy as broadcast_module - - monkeypatch.setattr(broadcast_module.ray._private.services, "get_node_ip_address", lambda: "192.168.1.1") - - ie_cfg = self._make_ie_cfg( - weight_sync_backend="gloo", - model_dtype="torch.bfloat16", - num_engines=2, - tensor_parallel_size=2, - pipeline_parallel_size=1, - data_parallel_size=1, - run_engines_locally=False, - ) - init_info = BroadcastTransferStrategy.create_init_info(ie_cfg, inference_world_size=4) - - assert isinstance(init_info, BroadcastInitInfo) - assert init_info.master_addr == "192.168.1.1" - assert isinstance(init_info.master_port, int) - assert init_info.rank_offset == 1 - # world_size = inference_world_size + 1 = 4 + 1 = 5 - assert init_info.world_size == 5 - assert init_info.override_existing_receiver is True - - def test_broadcast_create_init_info_override_existing_receiver_disabled_for_local_engines(self, monkeypatch): - """BroadcastTransferStrategy.create_init_info should set override_existing_receiver=False for local engines.""" - import skyrl.backends.skyrl_train.weight_sync.broadcast_strategy as broadcast_module - - monkeypatch.setattr(broadcast_module.ray._private.services, "get_node_ip_address", lambda: "192.168.1.1") - - ie_cfg = self._make_ie_cfg(run_engines_locally=True) - init_info = BroadcastTransferStrategy.create_init_info(ie_cfg, inference_world_size=1) - - assert init_info.override_existing_receiver is False - - def test_delta_create_init_info(self): - ie_cfg = self._make_ie_cfg(weight_sync_backend="delta", run_engines_locally=False) - # delta_weight_sync defaults to None, so a delta run must supply the whole sub-config. - ie_cfg.delta_weight_sync = DeltaWeightSyncConfig( - sync_dir="gs://bucket/prefix", - local_checkpoint_dir="/tmp/receiver", - max_file_size_in_gb=2, - publish_num_workers=3, - checkpoint_load_format="vllm_multi_thread_safetensors", - multi_thread_safetensors_max_workers=4, - ) - - init_info = DeltaTransferStrategy.create_init_info( - ie_cfg, - base_model_path="Qwen/Qwen2.5-1.5B-Instruct", - ) - - assert isinstance(init_info, DeltaInitInfo) - assert init_info.sync_dir == "gs://bucket/prefix" - assert init_info.base_model_path == "Qwen/Qwen2.5-1.5B-Instruct" - assert init_info.local_checkpoint_dir == "/tmp/receiver" - assert init_info.checkpoint_load_format == "vllm_multi_thread_safetensors" - assert init_info.multi_thread_safetensors_max_workers == 4 - assert init_info.max_file_size_in_gb == 2 - assert init_info.publish_num_workers == 3 - assert init_info.override_existing_receiver is True - - def test_delta_create_init_info_requires_sync_dir(self): - ie_cfg = self._make_ie_cfg(weight_sync_backend="delta") - - with pytest.raises(ValueError, match="sync_dir"): - DeltaTransferStrategy.create_init_info(ie_cfg, base_model_path="model") - - -class TestBroadcastWeightUpdateRequest: - """Tests for BroadcastWeightUpdateRequest.""" - - def test_len(self): - """__len__ should return number of weights.""" - request = BroadcastWeightUpdateRequest( - names=["layer1.weight", "layer2.weight"], - dtypes=["bfloat16", "bfloat16"], - shapes=[[4096, 4096], [1024]], - ) - assert len(request) == 2 - - def test_mismatched_lengths_raises(self): - """Mismatched lengths should raise ValueError.""" - with pytest.raises(ValueError, match="must have the same length"): - BroadcastWeightUpdateRequest( - names=["layer1.weight", "layer2.weight"], - dtypes=["bfloat16"], - shapes=[[4096, 4096]], - ) - - -class TestCudaIpcWeightUpdateRequest: - """Tests for CudaIpcWeightUpdateRequest.""" - - def test_serialize_roundtrip(self): - """Serialization/deserialization roundtrip preserves data.""" - request = CudaIpcWeightUpdateRequest( - names=["model.layer.weight"], - dtypes=["bfloat16"], - shapes=[[4096, 4096]], - sizes=[4096 * 4096], - ipc_handles={"gpu-uuid": "test_handle"}, - ) - - data = request.serialize() - result = CudaIpcWeightUpdateRequest.deserialize(data) - - assert result.names == request.names - assert result.dtypes == request.dtypes - assert result.shapes == request.shapes - assert result.sizes == request.sizes - assert result.ipc_handles == request.ipc_handles - - def test_serialize_roundtrip_multiple_weights(self): - """Roundtrip with multiple weights.""" - request = CudaIpcWeightUpdateRequest( - names=["layer1.weight", "layer2.weight", "layer3.bias"], - dtypes=["bfloat16", "bfloat16", "bfloat16"], - shapes=[[4096, 4096], [4096, 1024], [1024]], - sizes=[4096 * 4096, 4096 * 1024, 1024], - ipc_handles={"gpu-0": "handle1"}, - ) - - data = request.serialize() - result = CudaIpcWeightUpdateRequest.deserialize(data) - - assert result.names == request.names - assert result.dtypes == request.dtypes - assert result.shapes == request.shapes - assert result.sizes == request.sizes - assert result.ipc_handles == request.ipc_handles - - def test_deserialize_missing_end_marker(self): - """Missing end marker raises ValueError.""" - - invalid_data = b"some_invalid_data" - - with pytest.raises(ValueError, match="End marker not found"): - CudaIpcWeightUpdateRequest.deserialize(invalid_data) - - def test_deserialize_invalid_data(self): - """Invalid base64/pickle data raises ValueError.""" - from skyrl.backends.skyrl_train.weight_sync.cuda_ipc_strategy import ( - _IPC_REQUEST_END_MARKER, - ) - - invalid_data = b"not_valid_base64!!!" + _IPC_REQUEST_END_MARKER - - with pytest.raises(ValueError, match="Failed to deserialize"): - CudaIpcWeightUpdateRequest.deserialize(invalid_data) - - def test_serialize_aligned_to_4_bytes(self): - """Serialized data is 4-byte aligned.""" - request = CudaIpcWeightUpdateRequest( - names=["test"], - dtypes=["bfloat16"], - shapes=[[10]], - sizes=[10], - ipc_handles={}, - ) - data = request.serialize() - - assert len(data) % 4 == 0 - - -class TestLoraLoadRequest: - """Tests for LoraLoadRequest.""" - - def test_lora_path(self): - """lora_path should be stored correctly with empty defaults for base fields.""" - request = LoraLoadRequest(lora_path="/path/to/lora") - assert request.lora_path == "/path/to/lora" - assert request.names == [] - assert request.dtypes == [] - assert request.shapes == [] diff --git a/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py b/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py deleted file mode 100644 index a66c4f5159..0000000000 --- a/tests/backends/skyrl_train/weight_sync/test_weight_chunk.py +++ /dev/null @@ -1,97 +0,0 @@ -import pytest -import torch - -from skyrl.backends.skyrl_train.weight_sync import WeightChunk - - -class TestWeightChunk: - """Tests for WeightChunk dataclass.""" - - def test_initialization_basic(self): - """Test basic initialization with valid data.""" - names = ["layer1.weight", "layer1.bias"] - dtypes = ["torch.float32", "torch.float32"] - shapes = [[4, 3], [4]] - tensors = [torch.randn(4, 3), torch.randn(4)] - - chunk = WeightChunk(names=names, dtypes=dtypes, shapes=shapes, tensors=tensors) - - assert chunk.names == names - assert chunk.dtypes == dtypes - assert chunk.shapes == shapes - assert len(chunk.tensors) == 2 - - def test_validation_length_mismatch(self): - """Test that validation catches length mismatches.""" - with pytest.raises(ValueError, match="All lists must have the same length"): - WeightChunk( - names=["layer1.weight", "layer1.bias"], - dtypes=["torch.float32"], # Wrong length - shapes=[[4, 3], [4]], - tensors=[torch.randn(4, 3), torch.randn(4)], - ) - - with pytest.raises(ValueError, match="All lists must have the same length"): - WeightChunk( - names=["layer1.weight"], - dtypes=["torch.float32"], - shapes=[[4, 3], [4]], # Wrong length - tensors=[torch.randn(4, 3)], - ) - - def test_len(self): - """Test __len__ returns number of parameters.""" - chunk = WeightChunk( - names=["layer1.weight", "layer1.bias", "layer2.weight"], - dtypes=["torch.float32"] * 3, - shapes=[[4, 3], [4], [3, 2]], - tensors=[torch.randn(4, 3), torch.randn(4), torch.randn(3, 2)], - ) - - assert len(chunk) == 3 - - def test_total_numel(self): - """Test total_numel cached property.""" - tensors = [ - torch.randn(4, 3), # 12 elements - torch.randn(4), # 4 elements - torch.randn(3, 2), # 6 elements - ] - chunk = WeightChunk( - names=["layer1.weight", "layer1.bias", "layer2.weight"], - dtypes=["torch.float32"] * 3, - shapes=[[4, 3], [4], [3, 2]], - tensors=tensors, - ) - - assert chunk.total_numel == 12 + 4 + 6 - - def test_total_size_bytes(self): - """Test total_size_bytes cached property.""" - tensors = [ - torch.randn(4, 3, dtype=torch.float32), # 12 * 4 = 48 bytes - torch.randn(4, dtype=torch.float32), # 4 * 4 = 16 bytes - ] - chunk = WeightChunk( - names=["layer1.weight", "layer1.bias"], - dtypes=["torch.float32"] * 2, - shapes=[[4, 3], [4]], - tensors=tensors, - ) - - assert chunk.total_size_bytes == 48 + 16 - - def test_total_size_bytes_mixed_dtypes(self): - """Test total_size_bytes with mixed dtypes.""" - tensors = [ - torch.randn(10, dtype=torch.float32), # 10 * 4 = 40 bytes - torch.randn(10, dtype=torch.bfloat16), # 10 * 2 = 20 bytes - ] - chunk = WeightChunk( - names=["layer1.weight", "layer1.bias"], - dtypes=["torch.float32", "torch.bfloat16"], - shapes=[[10], [10]], - tensors=tensors, - ) - - assert chunk.total_size_bytes == 40 + 20 diff --git a/tests/backends/skyrl_train/weight_sync/test_weight_extractor_utils.py b/tests/backends/skyrl_train/weight_sync/test_weight_extractor_utils.py deleted file mode 100644 index aa64497c36..0000000000 --- a/tests/backends/skyrl_train/weight_sync/test_weight_extractor_utils.py +++ /dev/null @@ -1,231 +0,0 @@ -import pytest -import torch - -from skyrl.backends.skyrl_train.weight_sync import WeightChunk -from skyrl.backends.skyrl_train.weight_sync.weight_extractor_utils import ( - yield_module_grouped_chunks, -) - - -class TestModuleGrouping: - """Tests for yield_module_grouped_chunks utility function.""" - - def test_basic_module_grouping(self): - """Test that parameters are grouped by module correctly.""" - params = { - "model.layers.0.self_attn.q_proj.weight": torch.randn(768, 768), - "model.layers.0.self_attn.k_proj.weight": torch.randn(768, 768), - "model.layers.0.self_attn.v_proj.weight": torch.randn(768, 768), - "model.layers.0.mlp.fc1.weight": torch.randn(3072, 768), - "model.layers.0.mlp.fc2.weight": torch.randn(768, 3072), - } - - def gather_tensor(param): - return param - - def get_shape(name, param, tensor): - return list(tensor.shape) - - chunks = list( - yield_module_grouped_chunks( - params=params, - dtype=torch.float32, - gather_tensor_fn=gather_tensor, - get_shape_fn=get_shape, - batch_size_threshold_gb=0.0, - ) - ) - - # Should have 2 modules: self_attn and mlp - assert len(chunks) == 2 - - # Find self_attn and mlp chunks - self_attn_chunks = [c for c in chunks if "self_attn" in c.names[0]] - mlp_chunks = [c for c in chunks if "mlp" in c.names[0]] - - assert len(self_attn_chunks) == 1 - assert len(mlp_chunks) == 1 - - # Check self_attn has 3 params (q, k, v) - self_attn_chunk = self_attn_chunks[0] - assert len(self_attn_chunk) == 3 - assert "q_proj" in self_attn_chunk.names[0] - assert "k_proj" in self_attn_chunk.names[1] - assert "v_proj" in self_attn_chunk.names[2] - - # Check mlp has 2 params (fc1, fc2) - mlp_chunk = mlp_chunks[0] - assert len(mlp_chunk) == 2 - assert "fc1" in mlp_chunk.names[0] - assert "fc2" in mlp_chunk.names[1] - - @pytest.mark.parametrize( - "params_config,threshold_gb,expected_chunks,expected_params_per_chunk,description", - [ - # No batching: each module in separate chunk - ( - {"model.layers.{}.attn.weight": [(i, 100) for i in range(4)]}, - 0.0, - 4, - [1, 1, 1, 1], - "no batching", - ), - # Small threshold (~1000 bytes): batch 2 modules at a time - # Each param is 100 elements * 4 bytes = 400 bytes - # First batch: 400 + 400 = 800 < 1000, Second batch: 400 + 400 = 800 < 1000 - ( - {"model.layers.{}.attn.weight": [(i, 100) for i in range(4)]}, - 0.000001, - 2, - [2, 2], - "small threshold batches 2 modules", - ), - # Large threshold (1 GB): batch all modules together - ( - {"model.layers.{}.attn.weight": [(i, 100) for i in range(4)]}, - 1.0, - 1, - [4], - "large threshold batches all", - ), - # Module boundaries: 3 params (1200 bytes) exceeds threshold (600 bytes) by 2x - ( - {"model.layer.attn.proj{}.weight": [(i, 100) for i in range(3)]}, - 0.0000006, - 1, - [3], - "module boundary: 3 params exceed threshold", - ), - # Module boundaries: 2 params (800 bytes) exceeds threshold (400 bytes) by 2x - ( - {"model.layer.attn.proj{}.weight": [(i, 100) for i in range(2)]}, - 0.0000004, - 1, - [2], - "module boundary: 2 params exceed threshold", - ), - # Module boundaries: 5 params (2000 bytes) exceeds threshold (500 bytes) by 4x - ( - {"model.layer.attn.proj{}.weight": [(i, 100) for i in range(5)]}, - 0.0000005, - 1, - [5], - "module boundary: 5 params exceed threshold", - ), - ], - ) - def test_batching_with_different_thresholds( - self, params_config, threshold_gb, expected_chunks, expected_params_per_chunk, description - ): - """Test batching behavior with various threshold values and module structures.""" - # Build params dict from config - # params_config is {"template": [(idx, size), ...]} - params = {} - for template, indices_and_sizes in params_config.items(): - for idx, size in indices_and_sizes: - param_name = template.format(idx) - params[param_name] = torch.randn(size, dtype=torch.float32) - - chunks = list( - yield_module_grouped_chunks( - params=params, - dtype=torch.float32, - gather_tensor_fn=lambda p: p, - get_shape_fn=lambda n, p, t: list(t.shape), - batch_size_threshold_gb=threshold_gb, - ) - ) - - assert len(chunks) == expected_chunks - for i, expected_params in enumerate(expected_params_per_chunk): - assert len(chunks[i]) == expected_params - - def test_gather_tensor_callback(self): - """Test that gather_tensor_fn callback is called correctly.""" - original = torch.randn(10, 10, dtype=torch.float32) - params = {"model.layer.weight": original} - - def gather_tensor(param): - # Simulate gathering by adding 1.0 - return param + 1.0 - - chunks = list( - yield_module_grouped_chunks( - params=params, - dtype=torch.bfloat16, - gather_tensor_fn=gather_tensor, - get_shape_fn=lambda n, p, t: list(t.shape), - ) - ) - - assert len(chunks) == 1 - tensor = chunks[0].tensors[0] - # Should be bfloat16 (dtype cast in utils) and have +1.0 applied (from gather) - assert tensor.dtype == torch.bfloat16 - expected = (original + 1.0).to(torch.bfloat16) - assert torch.allclose(tensor, expected) - - def test_get_shape_callback(self): - """Test that get_shape_fn callback is called correctly.""" - params = {"model.layer.weight": torch.randn(10, 20)} - - def get_shape(name, param, tensor): - # Return custom shape - return [999, 888] - - chunks = list( - yield_module_grouped_chunks( - params=params, - dtype=torch.float32, - gather_tensor_fn=lambda p: p, - get_shape_fn=get_shape, - ) - ) - - assert len(chunks) == 1 - assert chunks[0].shapes[0] == [999, 888] - - def test_empty_params(self): - """Test with empty params dict.""" - params = {} - - chunks = list( - yield_module_grouped_chunks( - params=params, - dtype=torch.float32, - gather_tensor_fn=lambda p: p, - get_shape_fn=lambda n, p, t: list(t.shape), - ) - ) - - assert len(chunks) == 0 - - def test_chunk_properties(self): - """Test that returned chunks have correct WeightChunk properties.""" - params = { - "model.layer.attn.weight": torch.randn(10, 10), - "model.layer.attn.bias": torch.randn(10), - } - - chunks = list( - yield_module_grouped_chunks( - params=params, - dtype=torch.float32, - gather_tensor_fn=lambda p: p, - get_shape_fn=lambda n, p, t: list(t.shape), - ) - ) - - assert len(chunks) == 1 - chunk = chunks[0] - - # Check all required fields are present - assert isinstance(chunk, WeightChunk) - assert len(chunk.names) == 2 - assert len(chunk.dtypes) == 2 - assert len(chunk.shapes) == 2 - assert len(chunk.tensors) == 2 - - # Check total_numel - expected_numel = 10 * 10 + 10 - assert chunk.total_numel == expected_numel