From 38b65e56919ec02fd07decac44b59315f278c0a9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=A8=E7=9D=BF?= Date: Sat, 29 Aug 2026 22:34:24 +0800 Subject: [PATCH 01/34] fix(model-source): report remote model metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Include remote model config in start telemetry - Read config.json with the resolved S3 access policy - Attach metadata to a copied start config without mutating training arguments - Continue startup safely when remote metadata is unavailable --- # ✅ Tests ## Cover S3 config loading and telemetry behavior - Verify endpoint, credentials, addressing style, object key, and body cleanup - Cover invalid config payloads, local models, and remote read failures (cherry picked from commit 912b5a8049e1be6159af06e640e8436afd58981f) --- relax/core/controller.py | 18 +++- relax/utils/s3_model_loader.py | 45 ++++++++- .../core/test_controller_s3_model_cleanup.py | 46 ++++++++++ tests/test_s3_model_loader.py | 91 +++++++++++++++++++ 4 files changed, 198 insertions(+), 2 deletions(-) diff --git a/relax/core/controller.py b/relax/core/controller.py index 56cb72b93..5f0405bda 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -1,5 +1,6 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. import concurrent.futures +import copy import os import threading import time @@ -44,6 +45,7 @@ cleanup_s3_model_weights_from_shm, is_s3_uri, prepare_local_model, + read_s3_model_config, remove_stale_s3_model_caches, ) from relax.utils.training.ppo_utils import validate_ppo_config @@ -90,6 +92,20 @@ def _uses_s3_model_prefetch(config: Namespace) -> bool: return source is not None and is_s3_uri(source.uri) +def _build_train_start_config(config: Namespace) -> Namespace: + start_config = copy.copy(config) + if not _uses_s3_model_prefetch(config): + return start_config + try: + start_config._model_source_config = read_s3_model_config(config) + except Exception as exc: + logger.warning( + "Unable to read remote model config for train START telemetry " + f"({type(exc).__name__}); continuing without it." + ) + return start_config + + def _require_positive_timeout(value: float, env_name: str) -> float: if value <= 0: raise ValueError(f"{env_name} must be greater than 0, got {value}") @@ -723,7 +739,7 @@ def register_all_serve(self): roles_to_create.append((role, cls, num_gpus, data_source)) self._maybe_resolve_num_rollout(roles_to_create) - relax_utils.report_train_start(self.config) + relax_utils.report_train_start(_build_train_start_config(self.config)) actor_rollout_pg_roles = _actor_rollout_pg_roles(self.config) self._validate_gpu_resources(roles_to_create, colocate, actor_rollout_pg_roles) diff --git a/relax/utils/s3_model_loader.py b/relax/utils/s3_model_loader.py index 656980846..9790cdd04 100644 --- a/relax/utils/s3_model_loader.py +++ b/relax/utils/s3_model_loader.py @@ -52,6 +52,9 @@ _CACHE_DIR_PATTERN = re.compile(rf"{_MARKER_PREFIX}_[0-9a-f]{{16}}") _CLEANUP_LOCK_TIMEOUT_SECONDS = 300.0 _CLEANUP_LOCK_POLL_INTERVAL_SECONDS = 0.1 +_TELEMETRY_S3_CONNECT_TIMEOUT_SECONDS = 2.0 +_TELEMETRY_S3_READ_TIMEOUT_SECONDS = 3.0 +_TELEMETRY_S3_TOTAL_MAX_ATTEMPTS = 1 def is_s3_uri(uri) -> bool: @@ -138,7 +141,15 @@ def _safe_join(root: str, rel: str) -> str: return dest -def _make_s3_client(*, endpoint, use_placeholder_credentials=False, use_path_style=False): +def _make_s3_client( + *, + endpoint, + use_placeholder_credentials=False, + use_path_style=False, + connect_timeout=None, + read_timeout=None, + total_max_attempts=None, +): import boto3 from botocore.config import Config @@ -147,6 +158,12 @@ def _make_s3_client(*, endpoint, use_placeholder_credentials=False, use_path_sty max_pool_connections=64, proxies={}, # Disable proxies for this client without mutating os.environ. ) + if connect_timeout is not None: + config_kwargs["connect_timeout"] = connect_timeout + if read_timeout is not None: + config_kwargs["read_timeout"] = read_timeout + if total_max_attempts is not None: + config_kwargs["retries"] = {"total_max_attempts": total_max_attempts, "mode": "standard"} if use_path_style: config_kwargs["s3"] = {"addressing_style": "path"} client_kwargs = dict(endpoint_url=endpoint, config=Config(**config_kwargs)) @@ -593,6 +610,32 @@ def _resolve_endpoint(args) -> str | None: return os.environ.get("AWS_ENDPOINT_URL_S3") or os.environ.get("AWS_ENDPOINT_URL") +def read_s3_model_config(args) -> dict: + """Read only ``config.json`` using the resolved model source policy.""" + source = getattr(args, "model_source", None) + if source is None or not is_s3_uri(source.uri): + raise ValueError("an S3 model source is required to read config.json") + + bucket, prefix = _parse_s3_uri(source.uri) + cli = _make_s3_client( + endpoint=_resolve_endpoint(args), + use_placeholder_credentials=source.credential_mode == "placeholder", + use_path_style=source.addressing_style == "path", + connect_timeout=_TELEMETRY_S3_CONNECT_TIMEOUT_SECONDS, + read_timeout=_TELEMETRY_S3_READ_TIMEOUT_SECONDS, + total_max_attempts=_TELEMETRY_S3_TOTAL_MAX_ATTEMPTS, + ) + response = cli.get_object(Bucket=bucket, Key=_normalize_prefix(prefix) + "config.json") + body = response["Body"] + try: + model_config = json.loads(body.read()) + finally: + body.close() + if not isinstance(model_config, Mapping): + raise ValueError(f"S3 model config.json must contain a JSON object, got {type(model_config).__name__}") + return dict(model_config) + + def _resolve_shm_root(args) -> str: """Resolve an existing SHM cache root without falling back to disk.""" root = getattr(args, "s3_model_shm_root", None) or "/dev/shm" diff --git a/tests/core/test_controller_s3_model_cleanup.py b/tests/core/test_controller_s3_model_cleanup.py index f6721e45e..2393dea8f 100644 --- a/tests/core/test_controller_s3_model_cleanup.py +++ b/tests/core/test_controller_s3_model_cleanup.py @@ -20,6 +20,52 @@ ROLES = controller.ROLES +def test_train_start_config_attaches_remote_model_config_without_mutating_args(monkeypatch): + config = SimpleNamespace(model_source=SimpleNamespace(uri="s3://bucket/model/")) + model_config = {"model_type": "qwen3", "hidden_size": 4096} + monkeypatch.setattr(controller, "read_s3_model_config", lambda args: model_config) + + start_config = controller._build_train_start_config(config) + + assert start_config is not config + assert start_config._model_source_config is model_config + assert not hasattr(config, "_model_source_config") + + +def test_train_start_config_fails_open_when_remote_config_is_unavailable(monkeypatch): + config = SimpleNamespace(model_source=SimpleNamespace(uri="s3://bucket/model/")) + warnings = [] + monkeypatch.setattr( + controller, + "read_s3_model_config", + lambda _args: (_ for _ in ()).throw(RuntimeError("not found")), + ) + monkeypatch.setattr(controller.logger, "warning", warnings.append) + + start_config = controller._build_train_start_config(config) + + assert start_config is not config + assert not hasattr(start_config, "_model_source_config") + assert len(warnings) == 1 + assert "RuntimeError" in warnings[0] + assert "not found" not in warnings[0] + assert "s3://" not in warnings[0] + + +def test_train_start_config_does_not_read_local_model(monkeypatch): + config = SimpleNamespace(model_source=SimpleNamespace(uri="/models/qwen3")) + monkeypatch.setattr( + controller, + "read_s3_model_config", + lambda _args: pytest.fail("local model must not use the S3 client"), + ) + + start_config = controller._build_train_start_config(config) + + assert start_config is not config + assert not hasattr(start_config, "_model_source_config") + + class _FakeCleanupTask: def __init__(self): self.node_ids = [] diff --git a/tests/test_s3_model_loader.py b/tests/test_s3_model_loader.py index 510238a7e..ec8e7e9f0 100644 --- a/tests/test_s3_model_loader.py +++ b/tests/test_s3_model_loader.py @@ -205,6 +205,77 @@ def test_pre_parse_cli_model_source_disabled(monkeypatch): assert _pre_parse_cli_model_source() is None +def test_read_s3_model_config_uses_model_source_access_policy(monkeypatch): + class Body: + def __init__(self): + self.closed = False + + def read(self): + return b'{"model_type": "qwen3", "hidden_size": 4096}' + + def close(self): + self.closed = True + + class Client: + def __init__(self): + self.body = Body() + self.request = None + + def get_object(self, **kwargs): + self.request = kwargs + return {"Body": self.body} + + client = Client() + client_options = {} + + def make_client(**kwargs): + client_options.update(kwargs) + return client + + monkeypatch.setattr(m, "_make_s3_client", make_client) + args = SimpleNamespace( + model_source=m.ModelSource( + "s3://bucket/models/qwen3", + "http://s3.example", + credential_mode="placeholder", + addressing_style="path", + ) + ) + + assert m.read_s3_model_config(args) == {"model_type": "qwen3", "hidden_size": 4096} + assert client_options == { + "endpoint": "http://s3.example", + "use_placeholder_credentials": True, + "use_path_style": True, + "connect_timeout": 2.0, + "read_timeout": 3.0, + "total_max_attempts": 1, + } + assert client.request == {"Bucket": "bucket", "Key": "models/qwen3/config.json"} + assert client.body.closed + + +def test_read_s3_model_config_rejects_non_mapping_json(monkeypatch): + class Body: + def read(self): + return b"[]" + + def close(self): + self.closed = True + + body = Body() + monkeypatch.setattr( + m, + "_make_s3_client", + lambda **_kwargs: SimpleNamespace(get_object=lambda **_kwargs: {"Body": body}), + ) + args = SimpleNamespace(model_source=m.ModelSource("s3://bucket/model/")) + + with pytest.raises(ValueError, match="config.json must contain a JSON object"): + m.read_s3_model_config(args) + assert body.closed + + def test_s3_policy_dummy_does_not_affect_genrm(): from relax.backends.sglang.sglang_engine import _compute_genrm_server_args @@ -702,6 +773,26 @@ def test_make_s3_client_preserves_default_credentials_for_generic_s3(monkeypatch assert "aws_access_key_id" not in captured assert "aws_secret_access_key" not in captured + assert captured["config"].retries == {"max_attempts": 10, "mode": "standard"} + + +def test_make_s3_client_applies_bounded_request_policy(monkeypatch): + import boto3 + + captured = {} + monkeypatch.setattr(boto3, "client", lambda *args, **kwargs: captured.update(kwargs) or object()) + + m._make_s3_client( + endpoint="http://s3.example", + connect_timeout=2.0, + read_timeout=3.0, + total_max_attempts=1, + ) + + config = captured["config"] + assert config.connect_timeout == 2.0 + assert config.read_timeout == 3.0 + assert config.retries == {"total_max_attempts": 1, "mode": "standard"} def _fake_s3(objects): From ada460cd038b3623307f217ae36f64ec0a07da70 Mon Sep 17 00:00:00 2001 From: xudonghui Date: Mon, 31 Aug 2026 15:03:06 +0800 Subject: [PATCH 02/34] fix(sglang): use dummy load for async recipe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Avoid legacy SGLang S3 processor failures - initialize the fully-async rollout engines with dummy weights - rely on the mandatory Actor weight sync before the first rollout (cherry picked from commit 67b0e09fb30b77b1675c1e5bc894498d198776d6) --- scripts/training/text/run-qwen35-35B-A3B-16xgpu-async.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/scripts/training/text/run-qwen35-35B-A3B-16xgpu-async.sh b/scripts/training/text/run-qwen35-35B-A3B-16xgpu-async.sh index 2ea32eeae..58f0a37de 100755 --- a/scripts/training/text/run-qwen35-35B-A3B-16xgpu-async.sh +++ b/scripts/training/text/run-qwen35-35B-A3B-16xgpu-async.sh @@ -120,6 +120,7 @@ OPTIMIZER_ARGS=( SGLANG_ARGS=( --rollout-num-gpus-per-engine 2 + --sglang-load-format dummy --sglang-mem-fraction-static 0.75 --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256) ) From 0cd7c4dda838ccc43912dcb54eb0e56e711fdc22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=83=AD=E5=84=92=E8=BE=B0?= Date: Mon, 31 Aug 2026 20:43:54 +0800 Subject: [PATCH 03/34] fix(sglang): repair malformed hunk header in v0.5.15.post1 patch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fixes - `docker/patch/sglang/v0.5.15.post1.patch:401`: the hunk header declared `@@ -1322,6 +1331,20 @@` but the body contains 6 context + 16 added = **22** new lines. `git apply` refuses to parse the file entirely: ```console $ git apply --numstat docker/patch/latest/sglang.patch error: corrupt patch at line 424 ``` Because `--numstat` only parses (it never touches a work tree), this is a defect in the patch file itself, independent of any sglang version. Every image build with `ENABLE_SGLANG_PATCH=1` fails at `docker/Dockerfile:147`. Introduced by `73206573 feat(LoRA): Support LoRA RL MoE`, which hand-edited the `/post_process_weights` hunk without updating its line count, and reached `dev` through merge `f49950c2`. Fixed the count to 22, and shifted the following hunk's new-side start from 1474 to 1476 — the extra 2 lines move every later hunk in that file, and it is the only one. `git apply` tolerates a stale start line (it searches for the context), so the count alone is enough to make the patch apply; correcting the offset keeps the file identical to what `git diff` regenerates. Verified against a real `v0.5.15.post1` worktree: `git apply --check` exits 0, and all 99 hunk headers now match a canonical regeneration byte for byte. Co-Authored-By: Claude (cherry picked from commit 6806055fd36f58f1f45f3e8936e21f1bbbb8e08c) --- docker/Dockerfile | 2 +- docker/patch/latest/sglang.patch | 2 +- docker/patch/sglang/v0.5.15.post1.patch | 2 +- docker/patch/sglang/v0.5.17.patch | 1590 +++++++++++++++++++++ relax/backends/sglang/sglang_engine.py | 2 +- tests/backends/sglang/test_image_patch.py | 28 - 6 files changed, 1594 insertions(+), 32 deletions(-) create mode 100644 docker/patch/sglang/v0.5.17.patch delete mode 100644 tests/backends/sglang/test_image_patch.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 70e98cd9f..312475a81 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,7 +1,7 @@ ARG HTTP_PROXY ARG HTTPS_PROXY ARG NO_PROXY -ARG BASE_IMAGE=mirror.ccs.tencentyun.com/lmsysorg/sglang:v0.5.15.post1-cu129 +ARG BASE_IMAGE=mirror.ccs.tencentyun.com/lmsysorg/sglang:v0.5.17-cu129 ARG TRAIN_IMAGE=train FROM ${BASE_IMAGE} as base diff --git a/docker/patch/latest/sglang.patch b/docker/patch/latest/sglang.patch index 33c1dc275..5cc532326 120000 --- a/docker/patch/latest/sglang.patch +++ b/docker/patch/latest/sglang.patch @@ -1 +1 @@ -../sglang/v0.5.15.post1.patch \ No newline at end of file +../sglang/v0.5.17.patch \ No newline at end of file diff --git a/docker/patch/sglang/v0.5.15.post1.patch b/docker/patch/sglang/v0.5.15.post1.patch index f97f26003..1d66fb6e4 100644 --- a/docker/patch/sglang/v0.5.15.post1.patch +++ b/docker/patch/sglang/v0.5.15.post1.patch @@ -421,7 +421,7 @@ index bd9d7eafa1..f53a32abc0 100644 @app.post("/update_weight_version") @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version( -@@ -1451,6 +1474,19 @@ async def load_lora_adapter_from_tensors( +@@ -1451,6 +1476,19 @@ async def load_lora_adapter_from_tensors( return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) diff --git a/docker/patch/sglang/v0.5.17.patch b/docker/patch/sglang/v0.5.17.patch new file mode 100644 index 000000000..e08a5000a --- /dev/null +++ b/docker/patch/sglang/v0.5.17.patch @@ -0,0 +1,1590 @@ +diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py +index 240e20de..eda3406c 100644 +--- a/python/sglang/srt/disaggregation/decode.py ++++ b/python/sglang/srt/disaggregation/decode.py +@@ -21,6 +21,7 @@ Life cycle of a request in the decode server + from __future__ import annotations + + import logging ++import os + import time + from collections import deque + from dataclasses import dataclass +@@ -680,6 +681,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + def release_memory_occupation(self): + self.queue.clear() + self.retracted_queue.clear() ++ self.pending_reqs.clear() + if hasattr(self.kv_manager, "deregister_buffer_to_engine"): + self.kv_manager.deregister_buffer_to_engine() + +@@ -769,6 +771,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group + ) + ++ bootstrap_timeout = float( ++ os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") ++ ) ++ now = time.perf_counter() ++ + for decode_req, poll in zip(self.queue, polls): + if poll is None: + continue +@@ -776,7 +783,21 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin): + continue + + if poll == KVPoll.Bootstrapping: +- pass ++ entry_time = decode_req.req.time_stats.decode_prealloc_queue_entry_time ++ if entry_time > 0 and now - entry_time > bootstrap_timeout: ++ error_message = ( ++ f"Decode prealloc timeout for request rank={self.tp_rank} " ++ f"{decode_req.req.rid=} {decode_req.req.bootstrap_room=} " ++ f"after {bootstrap_timeout}s" ++ ) ++ logger.error(error_message) ++ prepare_abort( ++ decode_req.req, ++ error_message, ++ status_code=HTTPStatus.GATEWAY_TIMEOUT, ++ ) ++ if self.scheduler.metrics_reporter.enable_metrics: ++ self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() + elif poll == KVPoll.WaitingForInput: + decode_req.waiting_for_input = True + decode_req.req.time_stats.set_bootstrap_done_time() +@@ -1995,6 +2016,11 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + else: + polls = self._poll_with_metadata_gate() + ++ transfer_timeout = float( ++ os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") ++ ) ++ now = time.perf_counter() ++ + transferred_reqs = [] + indices_to_remove = set() + for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): +@@ -2071,7 +2097,17 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + KVPoll.WaitingForInput, + KVPoll.Transferring, + ]: +- pass ++ entry_time = decode_req.req.time_stats.decode_transfer_queue_entry_time ++ if entry_time > 0 and now - entry_time > transfer_timeout: ++ logger.error( ++ "Decode transfer timeout for request rank=%s rid=%s room=%s " ++ "after %ss", ++ self.tp_rank, ++ decode_req.req.rid, ++ decode_req.req.bootstrap_room, ++ transfer_timeout, ++ ) ++ decode_req.kv_receiver.abort() + else: + raise ValueError(f"Unexpected poll case: {poll}") + +@@ -2097,6 +2133,9 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin): + + def release_memory_occupation(self): + """Clean up in-flight transfers before releasing GPU memory.""" ++ for decode_req in self.queue: ++ if decode_req.kv_receiver is not None: ++ decode_req.kv_receiver.abort() + self.queue.clear() + + def resume_memory_occupation(self): +@@ -2320,6 +2359,11 @@ class SchedulerDisaggregationDecodeMixin: + resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs() + self.waiting_queue.extend(resumed_reqs) + if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: ++ transferred_reqs = self.disagg_decode_transfer_queue.pop_transferred() ++ if self.enable_hisparse: ++ for req in transferred_reqs: ++ self.hisparse_coordinator.admit_request_direct(req) ++ self.waiting_queue.extend(transferred_reqs) + # if there are still retracted requests, we do not allocate new requests + return + +diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py +index 1907e0ee..14048ad0 100644 +--- a/python/sglang/srt/disaggregation/mooncake/conn.py ++++ b/python/sglang/srt/disaggregation/mooncake/conn.py +@@ -1050,7 +1050,7 @@ class MooncakeKVManager(CommonKVManager): + for i, dst_aux_ptr in enumerate(dst_aux_ptrs): + length = prefill_aux_item_lens[i] + src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index +- dst_addr = dst_aux_ptrs[i] + length * req.dst_aux_index ++ dst_addr = dst_aux_ptr + length * req.dst_aux_index + transfer_blocks.append((src_addr, dst_addr, length)) + + return self._transfer_data(req.mooncake_session_id, transfer_blocks) +@@ -1714,12 +1714,6 @@ class MooncakeKVManager(CommonKVManager): + if ret != 0: + with self.session_lock: + self.session_failures[req.mooncake_session_id] += 1 +- # Failures should never happen if the session is not dead, if the session fails once, mark it as failed +- if self.session_failures[req.mooncake_session_id] >= 1: +- self.failed_sessions.add(req.mooncake_session_id) +- logger.error( +- f"Session {req.mooncake_session_id} failed." +- ) + self.record_failure( + kv_chunk.room, + f"Failed to send kv chunk of {kv_chunk.room} to " +diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py +index 5642c64e..599b9815 100644 +--- a/python/sglang/srt/disaggregation/prefill.py ++++ b/python/sglang/srt/disaggregation/prefill.py +@@ -21,6 +21,8 @@ from __future__ import annotations + + import hashlib + import logging ++import os ++import time + from array import array + from collections import deque + from http import HTTPStatus +@@ -416,6 +418,11 @@ class PrefillBootstrapQueue: + self.scheduler.attn_tp_cpu_group, + ) + ++ bootstrap_timeout = float( ++ os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") ++ ) ++ now = time.perf_counter() ++ + for i, (req, poll) in enumerate(zip(self.queue, polls)): + if poll is None: + continue +@@ -425,6 +432,27 @@ class PrefillBootstrapQueue: + indices_to_remove.add(i) + failed_reqs.append(req) + elif poll == KVPoll.Bootstrapping: ++ entry_time = req.time_stats.prefill_bootstrap_queue_entry_time ++ if entry_time > 0 and now - entry_time > bootstrap_timeout: ++ error_message = ( ++ f"Prefill bootstrap timed out after {now - entry_time:.1f}s " ++ f"for request rank={self.tp_rank} " ++ f"{req.rid=} {req.bootstrap_room=}" ++ ) ++ logger.error(error_message) ++ prepare_abort( ++ req, ++ error_message, ++ status_code=HTTPStatus.GATEWAY_TIMEOUT, ++ ) ++ self.scheduler.output_streamer.stream_output( ++ [req], req.return_logprob ++ ) ++ indices_to_remove.add(i) ++ failed_reqs.append(req) ++ if self.scheduler.metrics_reporter.enable_metrics: ++ self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() ++ continue + if ( + req.prefill_attempt_count + < self.scheduler.server_args.optimistic_prefill_attempts +@@ -832,6 +860,11 @@ class SchedulerDisaggregationPrefillMixin: + self.attn_tp_cpu_group, + ) + ++ transfer_timeout = float( ++ os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") ++ ) ++ now = time.perf_counter() ++ + undone_reqs: List[Req] = [] + # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue + for req, poll in zip(self.disagg_prefill_inflight_queue, polls): +@@ -866,7 +899,26 @@ class SchedulerDisaggregationPrefillMixin: + + if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: + # todo: set Transferring correctly in backend +- undone_reqs.append(req) ++ entry_time = req.time_stats.prefill_transfer_queue_entry_time ++ if entry_time > 0 and now - entry_time > transfer_timeout: ++ error_message = ( ++ f"Prefill transfer timed out after {now - entry_time:.1f}s " ++ f"(state={poll}) for request rank={self.ps.tp_rank} " ++ f"{req.rid=} {req.bootstrap_room=}" ++ ) ++ logger.error(error_message) ++ release_kv_cache(req, self.tree_cache) ++ prepare_abort( ++ req, ++ error_message, ++ status_code=HTTPStatus.GATEWAY_TIMEOUT, ++ ) ++ req.disagg_kv_sender.clear() ++ done_reqs.append(req) ++ if self.metrics_reporter.enable_metrics: ++ self.metrics_collector.increment_transfer_failed_reqs() ++ else: ++ undone_reqs.append(req) + elif poll == KVPoll.Success: # transfer done + if not isinstance(req.finished_reason, FINISH_ABORT): + req.finished_reason = FINISH_LENGTH(length=0) +diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py +index 6eb0ad1a..d27815f3 100644 +--- a/python/sglang/srt/entrypoints/engine.py ++++ b/python/sglang/srt/entrypoints/engine.py +@@ -73,6 +73,7 @@ from sglang.srt.managers.io_struct import ( + LoadLoRAAdapterReqInput, + MultimodalDataInputFormat, + OpenSessionReqInput, ++ PostProcessWeightsReqInput, + ProfileReq, + ProfileReqType, + ReleaseMemoryOccupationReqInput, +@@ -81,6 +82,7 @@ from sglang.srt.managers.io_struct import ( + RpcReqInput, + RpcReqOutput, + UnloadLoRAAdapterReqInput, ++ UpdateLoRAFromDistributedReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, +@@ -1429,6 +1431,20 @@ class Engine(EngineScoreMixin, EngineBase): + self.tokenizer_manager.update_weights_from_ipc(obj, None) + ) + ++ def post_process_weights( ++ self, ++ restore_weights_before_load: bool = False, ++ post_process_quantization: bool = False, ++ ): ++ """Optional post-processing for updated weights, e.g. quantization packing.""" ++ obj = PostProcessWeightsReqInput( ++ restore_weights_before_load=restore_weights_before_load, ++ post_process_quantization=post_process_quantization, ++ ) ++ return self.loop.run_until_complete( ++ self.tokenizer_manager.post_process_weights(obj, None) ++ ) ++ + def get_weights_by_name(self, name: str, truncate_size: int = 100): + """Get weights by parameter name.""" + obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size) +@@ -1474,6 +1490,36 @@ class Engine(EngineScoreMixin, EngineBase): + self.tokenizer_manager.load_lora_adapter_from_tensors(lora_req, None) + ) + ++ def update_lora_from_distributed( ++ self, ++ lora_name: str, ++ names, ++ dtypes, ++ shapes, ++ config_dict: Dict, ++ group_name: str, ++ pinned: bool = False, ++ ): ++ """Load/refresh a LoRA adapter whose tensors arrive over an NCCL group. ++ ++ In-memory counterpart to ``load_lora_adapter`` (disk) and ++ ``load_lora_adapter_from_tensors`` (serialized blob): the tensors are ++ broadcast ``src=0`` on ``group_name`` (the weight-update group), so ++ only metadata travels through this HTTP call. ++ """ ++ obj = UpdateLoRAFromDistributedReqInput( ++ lora_name=lora_name, ++ config_dict=config_dict, ++ names=names, ++ dtypes=dtypes, ++ shapes=shapes, ++ group_name=group_name, ++ pinned=pinned, ++ ) ++ return self.loop.run_until_complete( ++ self.tokenizer_manager.update_lora_from_distributed(obj, None) ++ ) ++ + def load_lora_adapter(self, lora_name: str, lora_path: str, pinned: bool = False): + """Load a new LoRA adapter without re-launching the engine.""" + +diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py +index ac38b466..ea54bac1 100644 +--- a/python/sglang/srt/entrypoints/http_server.py ++++ b/python/sglang/srt/entrypoints/http_server.py +@@ -132,6 +132,7 @@ from sglang.srt.managers.io_struct import ( + OpenSessionReqInput, + ParseFunctionCallReq, + PauseGenerationReqInput, ++ PostProcessWeightsReqInput, + ProfileReq, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, +@@ -140,6 +141,7 @@ from sglang.srt.managers.io_struct import ( + SetInternalStateReq, + SlowDownReqInput, + UnloadLoRAAdapterReqInput, ++ UpdateLoRAFromDistributedReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, +@@ -762,10 +764,8 @@ async def model_info(): + @app.get("/weight_version") + async def weight_version(): + """Get the current weight version.""" +- raise HTTPException( +- status_code=404, +- detail="Endpoint '/get_weight_version' or '/weight_version' is deprecated. Please use '/model_info' instead.", +- ) ++ result = await model_info() ++ return {"weight_version": result.get("weight_version", None)} + + + @app.get("/get_server_info") +@@ -782,9 +782,18 @@ async def get_server_info(): + async def server_info(): + """Get the server information.""" + # Returns internal states per DP. +- internal_states: List[Dict[Any, Any]] = ( +- await _global_state.tokenizer_manager.get_internal_state() +- ) ++ server_info_timeout = float(os.environ.get("SGLANG_SERVER_INFO_TIMEOUT", "2")) ++ try: ++ internal_states: List[Dict[Any, Any]] = await asyncio.wait_for( ++ _global_state.tokenizer_manager.get_internal_state(), ++ timeout=server_info_timeout, ++ ) ++ except asyncio.TimeoutError: ++ logger.warning( ++ "Timed out getting internal state for /server_info after %.1fs; returning empty internal_states", ++ server_info_timeout, ++ ) ++ internal_states = [] + + server_args = _global_state.tokenizer_manager.server_args + +@@ -1409,6 +1418,22 @@ async def update_weights_from_ipc( + return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) + + ++@app.post("/post_process_weights") ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def post_process_weights( ++ obj: Annotated[PostProcessWeightsReqInput, Body()], request: Request ++): ++ """Optional post-processing for updated weights, e.g. quantization packing.""" ++ success, message = await _global_state.tokenizer_manager.post_process_weights( ++ obj, request ++ ) ++ ++ content = {"success": success, "message": message} ++ return ORJSONResponse( ++ content, status_code=200 if success else HTTPStatus.BAD_REQUEST ++ ) ++ ++ + @app.post("/update_weight_version") + @auth_level(AuthLevel.ADMIN_OPTIONAL) + async def update_weight_version( +@@ -1539,6 +1564,19 @@ async def load_lora_adapter_from_tensors( + return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) + + ++@app.api_route("/update_lora_from_distributed", methods=["POST"]) ++@auth_level(AuthLevel.ADMIN_OPTIONAL) ++async def update_lora_from_distributed( ++ obj: Annotated[UpdateLoRAFromDistributedReqInput, Body()], request: Request ++): ++ """Load/refresh a LoRA adapter whose tensors arrive over an NCCL group.""" ++ result = await _global_state.tokenizer_manager.update_lora_from_distributed( ++ obj, request ++ ) ++ status_code = HTTPStatus.OK if result.success else HTTPStatus.BAD_REQUEST ++ return ORJSONResponse(msgspec_to_builtins(result), status_code=status_code) ++ ++ + @app.api_route("/unload_lora_adapter", methods=["POST"]) + @auth_level(AuthLevel.ADMIN_OPTIONAL) + async def unload_lora_adapter( +diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +index 79136182..ef791247 100644 +--- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py ++++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +@@ -2,6 +2,7 @@ from __future__ import annotations + + import contextlib + import logging ++import os + from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union + + import torch +@@ -30,6 +31,7 @@ from sglang.srt.layers.attention.dsa.utils import ( + is_dsa_enable_prefill_cp, + is_dsa_prefill_cp_in_seq_split, + is_graph_dsa_split_op_surface, ++ match_head_gate_q_scale, + ) + from sglang.srt.layers.layernorm import LayerNorm, RMSNorm + from sglang.srt.layers.utils import MultiPlatformOp +@@ -294,6 +296,17 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + self.k_norm = LayerNorm( + self.head_dim, dtype=torch.bfloat16 if _use_aiter else torch.float32 + ) ++ # NOTE: keep this override here, *after* `use_dsa_indexer_fusion` is computed ++ # above. Hoisting it earlier would let a forced `is_neox_style=False` flip that ++ # flag on and silently enable the fused indexer path. ++ env_neox_style = os.environ.get("INDEXER_ROPE_NEOX_STYLE") ++ if env_neox_style is not None: ++ if env_neox_style not in ("0", "1"): ++ raise ValueError( ++ "INDEXER_ROPE_NEOX_STYLE must be either '0' or '1' when set." ++ ) ++ is_neox_style = env_neox_style == "1" ++ + self.rotary_emb = get_rope_wrapper( + rope_head_dim, + rotary_dim=rope_head_dim, +@@ -365,6 +378,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + ): + weights = self._weights_proj_bf16_in_fp32_out(x) + weights = weights * self.n_heads**-0.5 ++ weights = match_head_gate_q_scale(weights, q_scale) + weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale + return weights + +@@ -372,6 +386,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + def _apply_q_scale_and_softmax_scale( + self, weights: torch.Tensor, q_scale: torch.Tensor + ): ++ weights = match_head_gate_q_scale(weights, q_scale) + return weights.unsqueeze(-1) * q_scale * self.softmax_scale + + @torch.compile(dynamic=True) +@@ -397,6 +412,12 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + return max_kv_len <= self.index_topk + return False + ++ def _maybe_repeat_query_heads(self, query: torch.Tensor) -> torch.Tensor: ++ if query.shape[1] < 32: ++ assert 32 % query.shape[1] == 0 ++ query = query.repeat_interleave(32 // query.shape[1], dim=1) ++ return query ++ + def _get_q_k_bf16( + self, + q_lora: torch.Tensor, +@@ -1604,6 +1625,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + query, key, weights_raw = self._get_q_k_bf16( + q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch + ) ++ query = self._maybe_repeat_query_heads(query) + q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) + with torch.cuda.stream(self.alt_stream): + self._store_index_k_cache( +@@ -1625,6 +1647,7 @@ class Indexer(DSANPUIndexerMixin, MultiPlatformOp): + enable_dual_stream, + forward_batch=forward_batch, + ) ++ query = self._maybe_repeat_query_heads(query) + + if enable_dual_stream: + current_stream = torch.cuda.current_stream() +diff --git a/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py b/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py +index f90e7ba9..407d88d7 100644 +--- a/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py ++++ b/python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py +@@ -14,6 +14,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo + get_tc_piecewise_forward_context, + is_in_tc_piecewise_cuda_graph, + ) ++from sglang.srt.layers.attention.dsa.utils import match_head_gate_q_scale + from sglang.srt.utils import is_cuda + from sglang.srt.utils.custom_op import register_custom_op + +@@ -81,6 +82,7 @@ if _is_cuda: + ) -> torch.Tensor: + out = torch.mm(x, weight.t(), out_dtype=torch.float32) + weights = out * n_heads_inv_sqrt ++ weights = match_head_gate_q_scale(weights, q_scale) + weights = weights.unsqueeze(-1) * q_scale * softmax_scale + return weights + +diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py +index 7df042c3..89054fe6 100644 +--- a/python/sglang/srt/layers/attention/dsa/utils.py ++++ b/python/sglang/srt/layers/attention/dsa/utils.py +@@ -245,6 +245,21 @@ def can_dsa_cp_split(seq_len: int, cp_size: int, use_dsa: bool, forward_batch): + return cur_cp_seq_len != 0 + + ++def match_head_gate_q_scale( ++ weights: torch.Tensor, q_scale: torch.Tensor ++) -> torch.Tensor: ++ """Broadcast head-gate weights up to ``q_scale``'s head count. ++ ++ Models whose indexer has fewer heads than the query projection (e.g. GLM-5) ++ produce ``weights`` with fewer heads than ``q_scale``; the subsequent ++ ``weights.unsqueeze(-1) * q_scale`` would otherwise broadcast incorrectly. ++ """ ++ if weights.shape[1] < q_scale.shape[1]: ++ assert q_scale.shape[1] % weights.shape[1] == 0 ++ weights = weights.repeat_interleave(q_scale.shape[1] // weights.shape[1], dim=1) ++ return weights ++ ++ + from sglang.kernels.ops.attention.dsa.cp_split import ( + dsa_cp_round_robin_split_q_seqs_kernel, + ) +diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +index 3490fc6a..b448595f 100644 +--- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py ++++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +@@ -1021,6 +1021,10 @@ class CompressedTensorsLinearMethod(LinearMethodBase): + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + ++ def restore_weights_before_loading(self, layer: torch.nn.Module) -> None: ++ if hasattr(layer.scheme, "restore_weights_before_loading"): ++ layer.scheme.restore_weights_before_loading(layer) ++ + def create_weights( + self, + layer: torch.nn.Module, +@@ -1075,6 +1079,10 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase): + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + layer.scheme.process_weights_after_loading(layer) + ++ def restore_weights_before_loading(self, layer: torch.nn.Module) -> None: ++ if hasattr(layer.scheme, "restore_weights_before_loading"): ++ layer.scheme.restore_weights_before_loading(layer) ++ + def create_weights( + self, + layer: torch.nn.Module, +diff --git a/python/sglang/srt/lora/layers.py b/python/sglang/srt/lora/layers.py +index 87b2a334..7c70342c 100644 +--- a/python/sglang/srt/lora/layers.py ++++ b/python/sglang/srt/lora/layers.py +@@ -948,6 +948,20 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA): + # initializes FusedMoE with its own moe_runner for base path + super().__init__(base_layer, lora_backend) + ++ # BaseLayerWithLoRA aliases `.weight`/`.bias` so a wrapped layer's base tensors stay ++ # reachable under their un-prefixed name. FusedMoE has neither -- it exposes ++ # `w13_weight`/`w2_weight` (plus quantization scales) -- so without an alias its base ++ # tensors are only reachable as `...experts.base_layer.w13_weight`. Model `load_weights` ++ # implementations resolve the un-prefixed `experts.w13_weight` produced by ++ # `FusedMoE.make_expert_params_mapping`, so reloading base weights into a LoRA-wrapped ++ # model (RL weight sync via /update_weights_from_*) would otherwise fail. Alias every ++ # directly-registered base parameter; this covers quantized MoE (w13_weight_scale, ++ # w13_weight_packed, ...) without a hardcoded name list. `named_parameters()` is called ++ # with remove_duplicate=False on the load path, so both names remain visible. ++ for param_name, param in base_layer.named_parameters(recurse=False): ++ if not hasattr(self, param_name): ++ setattr(self, param_name, param) ++ + lora_backend.is_moe_lora = True + + self.experts_shared_outer_loras: bool = False +diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py +index 928348eb..1c5bf177 100644 +--- a/python/sglang/srt/managers/io_struct.py ++++ b/python/sglang/srt/managers/io_struct.py +@@ -1862,6 +1862,16 @@ class ResumeMemoryOccupationReqOutput(BaseReq, kw_only=True): + pass + + ++class PostProcessWeightsReqInput(BaseReq, kw_only=True): ++ restore_weights_before_load: bool = False ++ post_process_quantization: bool = False ++ ++ ++class PostProcessWeightsReqOutput(BaseReq, kw_only=True): ++ success: bool ++ message: str ++ ++ + class CheckWeightsReqInput(BaseReq, kw_only=True): + action: str = "checksum" + allow_quant_error: bool = False +@@ -2166,6 +2176,45 @@ class LoadLoRAAdapterFromTensorsReqInput(BaseReq, kw_only=True): + ) + + ++class UpdateLoRAFromDistributedReqInput(BaseReq, kw_only=True): ++ """Load/refresh a LoRA adapter whose tensors arrive over an NCCL process ++ group (the RLHF online-update transport), instead of via disk or a ++ serialized blob. ++ ++ Mirrors ``UpdateWeightsFromDistributedReqInput`` (metadata only; the real ++ tensors are broadcast ``src=0`` on ``group_name``) but the payload lands in ++ the ``LoRAManager`` rather than ``model.load_weights``. ``config_dict`` is ++ the HF-PEFT adapter config (same content the disk path writes to ++ ``adapter_config.json``). ++ """ ++ ++ lora_name: str ++ config_dict: Dict[str, Any] ++ names: List[str] ++ dtypes: List[str] ++ shapes: List[List[int]] ++ # Number of tensors per broadcast bucket, in ``names`` order; sums to len(names). ++ # The sender picks the boundaries and the receiver must use the identical ones or the ++ # two desync on the broadcast. ``None`` means one bucket (pre-bucketing senders). ++ bucket_sizes: Optional[List[int]] = None ++ # The NCCL group name to receive the adapter tensors on (reuses the weight-update group). ++ group_name: str = "weight_update_group" ++ # Whether to pin the LoRA adapter in memory. ++ pinned: bool = False ++ added_tokens_config: Optional[Dict[str, Any]] = None ++ # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`. ++ lora_id: Optional[str] = None ++ load_format: Optional[str] = None ++ ++ def to_ref(self) -> LoRARef: ++ return LoRARef( ++ lora_id=self.lora_id, ++ lora_name=self.lora_name, ++ lora_path="__distributed__", ++ pinned=self.pinned, ++ ) ++ ++ + class LoRAUpdateOutput(BaseReq, kw_only=True): + success: bool + error_message: Optional[str] = None +@@ -2174,7 +2223,7 @@ class LoRAUpdateOutput(BaseReq, kw_only=True): + + LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = ( + LoadLoRAAdapterFromTensorsReqOutput +-) = LoRAUpdateOutput ++) = UpdateLoRAFromDistributedReqOutput = LoRAUpdateOutput + + + class BlockReqType(Enum): +diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py +index a45d122f..0334e353 100755 +--- a/python/sglang/srt/managers/schedule_batch.py ++++ b/python/sglang/srt/managers/schedule_batch.py +@@ -1103,6 +1103,7 @@ class Req(ReqDllmMixin): + self.metrics_collector = metrics_collector + if time_stats is not None: + self.time_stats = SchedulerReqTimeStats.new_from_obj(time_stats) ++ self.time_stats.disagg_mode = disagg_mode + else: + self.time_stats = SchedulerReqTimeStats(disagg_mode=disagg_mode) + self.time_stats.set_metrics_collector(metrics_collector) +@@ -2739,11 +2740,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): + + retracted_reqs = [] + first_iter = True ++ num_minimum_reqs = ( ++ 0 if server_args.disaggregation_mode == "decode" else 1 ++ ) + while first_iter or ( + not self.check_decode_mem(selected_indices=sorted_indices) + ): +- if len(sorted_indices) == 1: +- # Always keep at least one request ++ if len(sorted_indices) <= num_minimum_reqs: ++ # Unified mode keeps one request; decode disaggregation may retract all. + break + + first_iter = False +@@ -2754,7 +2758,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): + self.release_req(idx, len(sorted_indices), server_args) + + reqs_to_abort: List[Req] = [] +- if len(sorted_indices) <= 1 and not self.check_decode_mem( ++ if len(sorted_indices) <= num_minimum_reqs and not self.check_decode_mem( + selected_indices=sorted_indices + ): + # Even the last remaining request cannot fit in memory. +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index a9615a14..3fa13809 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -141,6 +141,7 @@ from sglang.srt.managers.io_struct import ( + LoadLoRAAdapterReqOutput, + OpenSessionReqInput, + PauseGenerationReqInput, ++ PostProcessWeightsReqInput, + ProfileReq, + ReleaseMemoryOccupationReqInput, + RemoveExternalCorpusReqInput, +@@ -161,6 +162,8 @@ from sglang.srt.managers.io_struct import ( + TokenizedGenerateReqInput, + UnloadLoRAAdapterReqInput, + UnloadLoRAAdapterReqOutput, ++ UpdateLoRAFromDistributedReqInput, ++ UpdateLoRAFromDistributedReqOutput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, +@@ -1539,6 +1542,10 @@ class Scheduler( + UpdateWeightsFromIPCReqInput, + self.weight_updater.update_weights_from_ipc, + ), ++ ( ++ PostProcessWeightsReqInput, ++ self.weight_updater.post_process_weights, ++ ), + ( + GetWeightsByNameReqInput, + self.weight_updater.get_weights_by_name, +@@ -1571,6 +1578,10 @@ class Scheduler( + LoadLoRAAdapterFromTensorsReqInput, + self.load_lora_adapter_from_tensors, + ), ++ ( ++ UpdateLoRAFromDistributedReqInput, ++ self.update_lora_from_distributed, ++ ), + (UnloadLoRAAdapterReqInput, self.unload_lora_adapter), + (PauseGenerationReqInput, self.pause_generation), + (ContinueGenerationReqInput, self.continue_generation), +@@ -4457,6 +4468,12 @@ class Scheduler( + # The request will still run one decode forward pass. + # Then we reuse all existing code to clean up the KV cache allocation. + logger.debug(f"Abort running request. {req.rid=}") ++ if self.disaggregation_mode == DisaggregationMode.PREFILL and hasattr( ++ req, "disagg_kv_sender" ++ ): ++ sender = getattr(req, "disagg_kv_sender", None) ++ if sender is not None and hasattr(sender, "abort"): ++ sender.abort() + req.to_finish = FINISH_ABORT() + + def _pause_engine(self) -> Tuple[List[Req], int]: +@@ -4669,6 +4686,14 @@ class Scheduler( + result = self.tp_worker.load_lora_adapter_from_tensors(recv_req) + return result + ++ def update_lora_from_distributed( ++ self, recv_req: UpdateLoRAFromDistributedReqInput ++ ) -> UpdateLoRAFromDistributedReqOutput: ++ """In-place loading a new lora adapter whose tensors arrive over NCCL.""" ++ ++ result = self.tp_worker.update_lora_from_distributed(recv_req) ++ return result ++ + def unload_lora_adapter( + self, recv_req: UnloadLoRAAdapterReqInput + ) -> UnloadLoRAAdapterReqOutput: +diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py +index 39eeec2e..12ba5133 100644 +--- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py ++++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py +@@ -397,7 +397,7 @@ class SchedulerProfilerManager: + if self.profiler_prefill_ct > self.profiler_target_prefill_ct: + if self.profile_in_progress: + self._stop_profile(stage=ForwardMode.EXTEND) +- elif batch.forward_mode.is_decode(): ++ elif batch.forward_mode.is_decode() or batch.forward_mode.is_prebuilt(): + if self.profiler_decode_ct == 0: + if self.profile_in_progress: + # force trace flush +diff --git a/python/sglang/srt/managers/scheduler_components/weight_updater.py b/python/sglang/srt/managers/scheduler_components/weight_updater.py +index 653b28c4..d30ce9f9 100644 +--- a/python/sglang/srt/managers/scheduler_components/weight_updater.py ++++ b/python/sglang/srt/managers/scheduler_components/weight_updater.py +@@ -28,6 +28,8 @@ from sglang.srt.managers.io_struct import ( + GetWeightsByNameReqOutput, + InitWeightsUpdateGroupReqInput, + InitWeightsUpdateGroupReqOutput, ++ PostProcessWeightsReqInput, ++ PostProcessWeightsReqOutput, + ReleaseMemoryOccupationReqInput, + ReleaseMemoryOccupationReqOutput, + ResumeMemoryOccupationReqInput, +@@ -183,6 +185,19 @@ class SchedulerWeightUpdaterManager: + parameter = self.tp_worker.get_weights_by_name(recv_req) + return GetWeightsByNameReqOutput(parameter=parameter) + ++ def post_process_weights(self, recv_req: PostProcessWeightsReqInput): ++ success, message = self.tp_worker.post_process_weights(recv_req) ++ if ( ++ success ++ and self.draft_worker is not None ++ and hasattr(self.draft_worker, "post_process_weights") ++ ): ++ success, message = self.draft_worker.post_process_weights(recv_req) ++ if not success: ++ logger.error(message) ++ torch.distributed.barrier(group=self.tp_cpu_group) ++ return PostProcessWeightsReqOutput(success=success, message=message) ++ + def _assert_weight_cache_inactive(self, op: str) -> None: + """Reject freeing/restoring model weights while the CUDA IPC weight + cache is active: the weights are shared with the daemon via CUDA IPC, so +diff --git a/python/sglang/srt/managers/tokenizer_control_mixin.py b/python/sglang/srt/managers/tokenizer_control_mixin.py +index 7cfd98b8..03ac104c 100644 +--- a/python/sglang/srt/managers/tokenizer_control_mixin.py ++++ b/python/sglang/srt/managers/tokenizer_control_mixin.py +@@ -48,6 +48,8 @@ from sglang.srt.managers.io_struct import ( + LoadLoRAAdapterReqOutput, + LoRAUpdateOutput, + OpenSessionReqInput, ++ PostProcessWeightsReqInput, ++ PostProcessWeightsReqOutput, + ProfileReq, + ProfileReqOutput, + ProfileReqType, +@@ -66,6 +68,8 @@ from sglang.srt.managers.io_struct import ( + SlowDownReqOutput, + UnloadLoRAAdapterReqInput, + UnloadLoRAAdapterReqOutput, ++ UpdateLoRAFromDistributedReqInput, ++ UpdateLoRAFromDistributedReqOutput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromDistributedReqOutput, + UpdateWeightsFromIPCReqInput, +@@ -101,6 +105,7 @@ _COMMUNICATOR_SPECS = [ + ("send_weights_to_remote_instance", SendWeightsToRemoteInstanceReqOutput), + ("update_weights_from_tensor", UpdateWeightsFromTensorReqOutput), + ("update_weights_from_ipc", UpdateWeightsFromIPCReqOutput), ++ ("post_process_weights", PostProcessWeightsReqOutput), + ("get_weights_by_name", GetWeightsByNameReqOutput), + ("release_memory_occupation", ReleaseMemoryOccupationReqOutput), + ("resume_memory_occupation", ResumeMemoryOccupationReqOutput), +@@ -738,6 +743,103 @@ class TokenizerControlMixin: + error_message=str(e), + ) + ++ async def update_lora_from_distributed( ++ self: TokenizerManager, ++ obj: UpdateLoRAFromDistributedReqInput, ++ _: Optional[fastapi.Request] = None, ++ ) -> UpdateLoRAFromDistributedReqOutput: ++ self.auto_create_handle_loop() ++ ++ try: ++ if not self.server_args.enable_lora: ++ raise ValueError( ++ "LoRA is not enabled. Please set `--enable-lora` to enable LoRA." ++ ) ++ ++ assert ( ++ self.server_args.dp_size == 1 ++ ), "dp_size must be 1 for dynamic lora loading" ++ logger.info( ++ "Start update Lora adapter from distributed. Lora name=%s", ++ obj.lora_name, ++ ) ++ ++ async with self.lora_update_lock: ++ # Self-contained same-name replacement: drop the prior version of ++ # this adapter first so repeated pushes under the fixed name do not ++ # accumulate (the NCCL caller does not send a separate unload). ++ # ++ # Ask the REGISTRY ("is a version live right now?"), never lora_ref_cache ++ # ("was one ever loaded?"). Upstream's cache is append-only -- unload only ++ # touches the registry -- so a push that dies between the unload above and ++ # the receive below leaves the name in the cache but not in the registry. ++ # Keying the unload off the cache would then make every retry call ++ # unregister() on a name that is gone, which raises before the receive is ++ # ever reached: one transient failure would wedge the engine permanently. ++ if obj.lora_name in self.lora_registry.get_all_adapters(): ++ unload_result = await self._unload_lora_adapter_locked( ++ UnloadLoRAAdapterReqInput(lora_name=obj.lora_name) ++ ) ++ if not unload_result.success: ++ raise ValueError( ++ f"Error while unloading prior LoRA adapter '{obj.lora_name}': " ++ f"{unload_result.error_message}" ++ ) ++ ++ new_adapter = LoRARef( ++ lora_name=obj.lora_name, ++ lora_path="__distributed__", ++ pinned=obj.pinned, ++ ) ++ obj.lora_id = new_adapter.lora_id ++ result = (await self.update_lora_adapter_communicator(obj))[0] ++ ++ if result.success: ++ await self.lora_registry.register(new_adapter) ++ # Deliberately NOT recorded in lora_ref_cache. That cache exists so an ++ # LRU-evicted adapter can be transparently re-read from its lora_path, ++ # but these tensors arrived over NCCL and were never on disk -- ++ # "__distributed__" is a placeholder, not a real path. Caching it would ++ # send the implicit-reload path in TokenizerManager off to load a ++ # directory that does not exist; leaving it out makes an unregistered ++ # adapter surface as "never been loaded", which is the truth. ++ if self.server_args.max_loaded_loras is not None: ++ while ( ++ self.lora_registry.num_registered_loras ++ > self.server_args.max_loaded_loras ++ ): ++ lru_lora_name = await self.lora_registry.lru_lora_name( ++ exclude_pinned=True ++ ) ++ if lru_lora_name is None: ++ raise ValueError( ++ "Didn't find any LoRA adapters when trying to evict LRU LoRA adapter. " ++ f"LoRA registry is: {self.lora_registry._registry}" ++ ) ++ ++ logger.info( ++ f"Unloading least recently used LoRA adapter '{lru_lora_name}' " ++ f"(current number of adapters: {self.lora_registry.num_registered_loras}, " ++ f"max allowed: {self.server_args.max_loaded_loras})" ++ ) ++ ++ unload_result = await self._unload_lora_adapter_locked( ++ UnloadLoRAAdapterReqInput(lora_name=lru_lora_name) ++ ) ++ if not unload_result.success: ++ raise ValueError( ++ f"Error while unloading LRU LoRA adapter '{lru_lora_name}': " ++ f"{unload_result.error_message}" ++ ) ++ del result.loaded_adapters[lru_lora_name] ++ ++ return result ++ except ValueError as e: ++ return UpdateLoRAFromDistributedReqOutput( ++ success=False, ++ error_message=str(e), ++ ) ++ + async def unload_lora_adapter( + self: TokenizerManager, + obj: UnloadLoRAAdapterReqInput, +@@ -797,6 +899,16 @@ class TokenizerControlMixin: + self.auto_create_handle_loop() + await self.resume_memory_occupation_communicator(obj) + ++ async def post_process_weights( ++ self: TokenizerManager, ++ obj: PostProcessWeightsReqInput, ++ request: Optional[fastapi.Request] = None, ++ ) -> Tuple[bool, str]: ++ self.auto_create_handle_loop() ++ async with self.model_update_lock.writer_lock: ++ results = await self.post_process_weights_communicator(obj) ++ return FanOutCommunicator.merge_results(results) ++ + async def check_weights( + self: TokenizerManager, + obj: CheckWeightsReqInput, +diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py +index 80351f33..99e4a1c2 100644 +--- a/python/sglang/srt/managers/tokenizer_manager.py ++++ b/python/sglang/srt/managers/tokenizer_manager.py +@@ -2840,27 +2840,25 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin): + priority = getattr(state.obj, "priority", None) + if priority is not None: + labels["priority"] = str(priority) +- if ( +- not state.ttft_observed +- and self.disaggregation_mode != DisaggregationMode.PREFILL +- ): ++ if not state.ttft_observed: + state.ttft_observed = True + state.last_completion_tokens = completion_tokens +- self.metrics_collector.observe_time_to_first_token( +- labels, +- state.time_stats.get_first_token_latency(), +- stream=getattr(state.obj, "stream", False), +- ) ++ if self.disaggregation_mode != DisaggregationMode.PREFILL: ++ self.metrics_collector.observe_time_to_first_token( ++ labels, ++ state.time_stats.get_first_token_latency(), ++ stream=getattr(state.obj, "stream", False), ++ ) + else: + num_new_tokens = completion_tokens - state.last_completion_tokens +- if num_new_tokens: ++ if num_new_tokens > 0: + self.metrics_collector.observe_inter_token_latency( + labels, + state.time_stats.get_interval(), + num_new_tokens, + ) + state.time_stats.set_last_time() +- state.last_completion_tokens = completion_tokens ++ state.last_completion_tokens = completion_tokens + + if state.finished: + # Get detailed cache breakdown if available +@@ -3731,7 +3729,14 @@ def _get_processor_wrapper(server_args): + + + def _determine_tensor_transport_mode(server_args: ServerArgs) -> TensorTransportMode: +- is_cross_node = server_args.dist_init_addr ++ # NOTE(relax-fix): keying "cross node" off dist_init_addr false-positives for ++ # single-node multi-GPU TP, which MUST also set dist_init_addr to bootstrap the ++ # TP process group (see relax/backends/sglang/sglang_engine.py). The false ++ # positive forces the CPU "default" transport, which makes wrap_shm_features() ++ # a no-op, so multimodal pixel tensors (multi-GB for 5-image doc pages) get ++ # pickled + gloo-broadcast per request (observed 2.2-6.5 GB, 18-33s per ++ # recv_requests). Key off nnodes, which is what actually means multi-node. ++ is_cross_node = server_args.nnodes > 1 + + if is_cross_node: + # Fallback to default CPU transport for multi-node +diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py +index 0fdf5b96..2cda3ba5 100644 +--- a/python/sglang/srt/managers/tp_worker.py ++++ b/python/sglang/srt/managers/tp_worker.py +@@ -30,8 +30,10 @@ from sglang.srt.managers.io_struct import ( + InitWeightsUpdateGroupReqInput, + LoadLoRAAdapterFromTensorsReqInput, + LoadLoRAAdapterReqInput, ++ PostProcessWeightsReqInput, + SendWeightsToRemoteInstanceReqInput, + UnloadLoRAAdapterReqInput, ++ UpdateLoRAFromDistributedReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromIPCReqInput, +@@ -195,6 +197,13 @@ class BaseTpWorker(ABC): + ) + return success, message + ++ def post_process_weights(self, recv_req: PostProcessWeightsReqInput): ++ success, message = self.model_runner.post_process_weights( ++ restore_weights_before_load=recv_req.restore_weights_before_load, ++ post_process_quantization=recv_req.post_process_quantization, ++ ) ++ return success, message ++ + def _deserialize_own_rank(self, serialized_named_tensors): + """Each rank deserializes only its own payload (index ps.tp_rank); + deserializing another rank's copy would break producer-side CUDA-IPC +@@ -285,6 +294,23 @@ class BaseTpWorker(ABC): + ) + return result + ++ def update_lora_from_distributed( ++ self, recv_req: UpdateLoRAFromDistributedReqInput ++ ): ++ # Every TP worker joins the broadcast; the LoRA code slices its own shard ++ # internally (slice_lora_a_weights / slice_lora_b_weights). ++ result = self.model_runner.update_lora_from_distributed( ++ recv_req.to_ref(), ++ recv_req.names, ++ recv_req.dtypes, ++ recv_req.shapes, ++ recv_req.config_dict, ++ recv_req.group_name, ++ recv_req.added_tokens_config, ++ recv_req.bucket_sizes, ++ ) ++ return result ++ + def forward_batch_embedding(self, batch: ScheduleBatch): + forward_batch = ForwardBatch.init_new( + batch, +diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py +index 5ebac56f..e9c2d85d 100644 +--- a/python/sglang/srt/mem_cache/hiradix_cache.py ++++ b/python/sglang/srt/mem_cache/hiradix_cache.py +@@ -1137,9 +1137,7 @@ class HiRadixCache(RadixCache): + self._update_leaf_status(node) + self._update_host_leaf_status(node) + if node.parent is None: +- assert ( +- node is self.root_node +- ), f"This request holds the node from another tree" ++ break + node = node.parent + return DecLockRefResult(delta=delta) + +@@ -1245,6 +1243,7 @@ class HiRadixCache(RadixCache): + self._update_host_leaf_status(node) + # update leaf status for the parent because the node is evicted + self._update_leaf_status(node.parent) ++ self._update_host_leaf_status(node.parent) + return num_evicted + + def _evict_backuped(self, node: TreeNode): +diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py +index a9b2b8b9..27b1c48e 100644 +--- a/python/sglang/srt/mem_cache/memory_pool.py ++++ b/python/sglang/srt/mem_cache/memory_pool.py +@@ -4391,9 +4391,12 @@ class DSATokenToKVPool(MLATokenToKVPool): + def _create_index_buffers(self): + num_pages = (self.index_buf_size + self.page_size + 1) // self.page_size + with ( +- torch.cuda.use_mem_pool(self.custom_mem_pool) +- if self.custom_mem_pool +- else nullcontext() ++ ( ++ torch.cuda.use_mem_pool(self.custom_mem_pool) ++ if self.custom_mem_pool ++ else nullcontext() ++ ), ++ self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE), + ): + self.index_k_with_scale_buffer = [ + torch.zeros( +diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py +index 2916f8c6..32b81c49 100644 +--- a/python/sglang/srt/mem_cache/radix_cache.py ++++ b/python/sglang/srt/mem_cache/radix_cache.py +@@ -491,6 +491,9 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): + return + + token_ids = req.get_fill_ids() ++ kv_committed_len = getattr(req, "kv_committed_len", None) ++ if kv_committed_len is not None and len(token_ids) > kv_committed_len: ++ token_ids = token_ids[:kv_committed_len] + kv_indices = self.req_to_token_pool.req_to_token[ + req.req_pool_idx, : len(token_ids) + ] +@@ -619,9 +622,7 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): + node.lock_ref -= 1 + self._update_leaf_status(node) + if node.parent is None: +- assert ( +- node is self.root_node +- ), "This request holds the node from another tree" ++ break + node = node.parent + return DecLockRefResult(delta=delta) + +diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py +index 8e957f66..1bd46d38 100644 +--- a/python/sglang/srt/model_executor/model_runner.py ++++ b/python/sglang/srt/model_executor/model_runner.py +@@ -376,9 +376,10 @@ class ModelRunner: + self.war_fastpath_read_done_event: Optional[torch.cuda.Event] = None + + # CPU offload +- set_offloader( +- create_offloader_from_server_args(server_args, dp_rank=self.ps.dp_rank) +- ) ++ if not is_draft_worker: ++ set_offloader( ++ create_offloader_from_server_args(server_args, dp_rank=self.ps.dp_rank) ++ ) + + self._weight_checker = WeightChecker(get_model=lambda: self.model, ps=self.ps) + +@@ -1183,10 +1184,130 @@ class ModelRunner: + lora_ref, tensors, config_dict, added_tokens_config + ) + ++ def update_lora_from_distributed( ++ self, ++ lora_ref: LoRARef, ++ names, ++ dtypes, ++ shapes, ++ config_dict, ++ group_name, ++ added_tokens_config=None, ++ bucket_sizes=None, ++ ): ++ """Receive a LoRA adapter's tensors over an NCCL process group and load it. ++ ++ Mirrors ``update_weights_from_distributed`` (each ``(name, dtype, shape)`` ++ is broadcast ``src=0`` on the shared weight-update group), but the ++ gathered tensors are handed to the ``LoRAManager`` instead of ++ ``model.load_weights``. The full (TP-gathered, PP-merged) adapter is ++ broadcast to every worker, which slices its own shard internally. ++ ++ ``bucket_sizes`` is the number of tensors per broadcast bucket, exactly as the ++ sender grouped them. Each bucket is staged on device, received, then moved to ++ host memory and released, so peak device memory is ONE bucket instead of the ++ whole adapter: a MoE adapter runs to several GB, and allocating that on top of ++ the KV cache OOMs mid-broadcast, which wedges every other participant until the ++ NCCL timeout. Host tensors are what ``LoRAAdapter`` keeps anyway ++ (``_process_weight`` calls ``.cpu()``). ``None`` means a single bucket. ++ """ ++ from sglang.srt.managers.io_struct import LoRAUpdateOutput ++ ++ model_update_group = self.weight_updater._model_update_group ++ assert group_name in model_update_group, ( ++ f"Group {group_name} not in {list(model_update_group.keys())}. " ++ "Please call `init_weights_update_group` first." ++ ) ++ if not bucket_sizes: ++ bucket_sizes = [len(names)] ++ assert sum(bucket_sizes) == len(names), ( ++ f"bucket_sizes sum to {sum(bucket_sizes)} but {len(names)} tensors were " ++ "announced; sender and receiver would issue different broadcasts and hang." ++ ) ++ logger.info(f"LoRA adapter loading from distributed starts: {lora_ref}.") ++ try: ++ tensors = {} ++ offset = 0 ++ for count in bucket_sizes: ++ handles = [] ++ staged = [] ++ for name, dtype, shape in zip( ++ names[offset : offset + count], ++ dtypes[offset : offset + count], ++ shapes[offset : offset + count], ++ ): ++ target_dtype = ( ++ dtype ++ if isinstance(dtype, torch.dtype) ++ else getattr(torch, dtype) ++ ) ++ weight = torch.empty(shape, dtype=target_dtype, device=self.device) ++ handles.append( ++ torch.distributed.broadcast( ++ weight, ++ src=0, ++ group=model_update_group[group_name], ++ async_op=True, ++ ) ++ ) ++ staged.append((name, weight)) ++ for handle in handles: ++ handle.wait() ++ handles.clear() ++ # Move off-device and drop the staging buffers before the next bucket is ++ # allocated, so the device high-water mark stays at one bucket. ++ for name, weight in staged: ++ tensors[name] = weight.to("cpu") ++ staged.clear() ++ offset += count ++ except Exception as e: ++ error_msg = f"Failed to receive LoRA adapter over distributed group: {e}." ++ logger.error(error_msg) ++ return LoRAUpdateOutput(success=False, error_message=error_msg) ++ ++ result = self.lora_manager.load_lora_adapter_from_tensors( ++ lora_ref, tensors, config_dict, added_tokens_config ++ ) ++ logger.info(f"LoRA adapter loading from distributed completes: {lora_ref}.") ++ return result ++ + def unload_lora_adapter(self, lora_ref: LoRARef): + """Unload a lora adapter that was previously loaded during initialization or dynamic loading.""" + return self.lora_manager.unload_lora_adapter(lora_ref) + ++ def post_process_weights( ++ self, ++ restore_weights_before_load: bool = False, ++ post_process_quantization: bool = False, ++ ): ++ """Run optional post-loading hooks, such as quantization repacking.""" ++ from sglang.srt.model_loader.loader import device_loading_context ++ ++ if self.device == "cuda": ++ target_device = torch.device("cuda", torch.cuda.current_device()) ++ else: ++ target_device = torch.device(self.device) ++ ++ if restore_weights_before_load: ++ for _, module in self.model.named_modules(): ++ quant_method = getattr(module, "quant_method", None) ++ if quant_method is not None and hasattr( ++ quant_method, "restore_weights_before_loading" ++ ): ++ with device_loading_context(module, target_device): ++ quant_method.restore_weights_before_loading(module) ++ ++ if post_process_quantization: ++ for _, module in self.model.named_modules(): ++ quant_method = getattr(module, "quant_method", None) ++ if quant_method is not None and hasattr( ++ quant_method, "process_weights_after_loading" ++ ): ++ with device_loading_context(module, target_device): ++ quant_method.process_weights_after_loading(module) ++ ++ return True, "Success" ++ + @property + def effective_max_total_num_tokens(self): + """Return the max token pool size considering hybrid swa settings.""" +@@ -1441,6 +1562,12 @@ class ModelRunner: + output.expert_distribution_metrics = recorder_outputs.get("metrics") + + no_copy_to_cpu = not get_schedule().disable_overlap_schedule ++ cuda_graph_num_tokens = None ++ decode_graph_runner = getattr(self, "decode_cuda_graph_runner", None) ++ if getattr(decode_graph_runner, "bs", None): ++ cuda_graph_num_tokens = decode_graph_runner.bs * getattr( ++ decode_graph_runner, "num_tokens_per_bs", 1 ++ ) + if ( + not self.is_draft_worker + and (experts_capturer := get_global_experts_capturer()) is not None +@@ -1448,7 +1575,7 @@ class ModelRunner: + output.routed_experts_output = experts_capturer.on_forward_end( + forward_batch=forward_batch, + can_run_graph=output.can_run_graph, +- cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None), ++ cuda_graph_batch=cuda_graph_num_tokens, + no_copy_to_cpu=no_copy_to_cpu, + ) + +@@ -1456,7 +1583,7 @@ class ModelRunner: + output.indexer_topk_output = indexer_capturer.on_forward_end( + forward_batch=forward_batch, + can_run_graph=output.can_run_graph, +- cuda_graph_batch=getattr(self.decode_cuda_graph_runner, "bs", None), ++ cuda_graph_batch=cuda_graph_num_tokens, + no_copy_to_cpu=no_copy_to_cpu, + ) + +diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py +index 08c47b99..06a222fe 100644 +--- a/python/sglang/srt/models/qwen3_5.py ++++ b/python/sglang/srt/models/qwen3_5.py +@@ -1125,12 +1125,12 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): + forward_batch: ForwardBatch, + ) -> torch.Tensor: + """Full attention forward pass.""" +- if _is_cuda and self.attn_output_gate: ++ if _is_cuda and self.attn_output_gate and positions.ndim == 1: + q, k, v, gate = self.forward_prepare_cuda_fused( + positions=positions, + hidden_states=hidden_states, + ) +- elif (_is_hip or _is_xpu or _is_cpu) and self.attn_output_gate: ++ elif (_is_cuda or _is_hip or _is_xpu or _is_cpu) and self.attn_output_gate: + q, k, v, gate = self.forward_prepare_fused_gate( + positions=positions, + hidden_states=hidden_states, +diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py +index 83697f14..aae8b5c8 100644 +--- a/python/sglang/srt/models/qwen3_vl.py ++++ b/python/sglang/srt/models/qwen3_vl.py +@@ -1171,9 +1171,14 @@ class Qwen3LLMModel(Qwen3Model): + # To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack + # The order matters because addition with different tensors is not associative in practice. + # Deepstack for prev_layer is applied at the start of current layer via post_residual_addition. +- deepstack_embeds = self.get_deepstack_embeds( +- layer_idx - 1, input_deepstack_embeds +- ) ++ deepstack_embeds = None ++ if input_deepstack_embeds is not None: ++ prev_layer_idx = layer_idx - 1 ++ if prev_layer_idx in self.deepstack_embed_to_decoder_layer: ++ sep = self.hidden_size * prev_layer_idx ++ deepstack_embeds = input_deepstack_embeds[ ++ :, sep : sep + self.hidden_size ++ ] + hidden_states, residual = layer( + positions, + hidden_states, +diff --git a/python/sglang/srt/multimodal/processors/glm4v.py b/python/sglang/srt/multimodal/processors/glm4v.py +index db684259..17d2cb69 100644 +--- a/python/sglang/srt/multimodal/processors/glm4v.py ++++ b/python/sglang/srt/multimodal/processors/glm4v.py +@@ -1,7 +1,13 @@ + from typing import List, Union + ++import torch ++ + from sglang.srt.layers.rotary_embedding import MRotaryEmbedding +-from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput ++from sglang.srt.managers.schedule_batch import ( ++ Modality, ++ MultimodalDataItem, ++ MultimodalProcessorOutput, ++) + from sglang.srt.models.glm4v import Glm4vForConditionalGeneration + from sglang.srt.models.glm4v_moe import Glm4vMoeForConditionalGeneration + from sglang.srt.multimodal.processors.base_processor import ( +@@ -46,6 +52,8 @@ class Glm4vImageProcessor(SGLangBaseProcessor): + self.IMAGE_END_TOKEN_ID = hf_config.image_end_token_id + self.VIDEO_START_TOKEN_ID = hf_config.video_start_token_id + self.VIDEO_END_TOKEN_ID = hf_config.video_end_token_id ++ self.IM_START_TOKEN_ID = self.IMAGE_START_TOKEN_ID ++ self.IM_END_TOKEN_ID = self.IMAGE_END_TOKEN_ID + + # Vision config + self.IMAGE_FACTOR = 28 +@@ -60,6 +68,39 @@ class Glm4vImageProcessor(SGLangBaseProcessor): + video_token_id=self.IM_TOKEN_ID, + ).build(_processor) + ++ def get_mm_data(self, prompt, embeddings, img_grid_thw): ++ input_ids, offsets, _ = self.build_input_ids(prompt, img_grid_thw=img_grid_thw) ++ image_embeddings = ( ++ embeddings.get(Modality.IMAGE, embeddings) ++ if isinstance(embeddings, dict) ++ else embeddings ++ ) ++ mm_items = [ ++ MultimodalDataItem( ++ modality=Modality.IMAGE, ++ offsets=offsets, ++ precomputed_embeddings=image_embeddings, ++ ) ++ ] ++ ++ mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index_glm4v( ++ input_ids=torch.tensor(input_ids, dtype=torch.long).unsqueeze(0), ++ hf_config=self.hf_config, ++ image_grid_thw=img_grid_thw, ++ video_grid_thw=None, ++ attention_mask=None, ++ ) ++ ++ return MultimodalProcessorOutput( ++ input_ids=input_ids, ++ mm_items=mm_items, ++ im_start_id=self.IM_START_TOKEN_ID, ++ im_end_id=self.IM_END_TOKEN_ID, ++ im_token_id=self.IM_TOKEN_ID, ++ mrope_positions=mrope_positions.squeeze(1), ++ mrope_position_delta=mrope_position_delta, ++ ) ++ + def compute_mrope_positions(self, input_ids, mm_items): + image_grid_thw = None + video_grid_thw = None +diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py +index a53e9373..1ae2c6f4 100644 +--- a/python/sglang/srt/multimodal/processors/qwen_vl.py ++++ b/python/sglang/srt/multimodal/processors/qwen_vl.py +@@ -723,7 +723,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor): + **kwargs, + ): + entry_time = time.perf_counter() +- base_output = await self.load_mm_data( ++ base_output = await self.legacy_load_mm_data( + prompt=input_text, + image_data=image_data, + video_data=request_obj.video_data, +diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py +index 3bcb6b36..9b971098 100644 +--- a/python/sglang/srt/observability/req_time_stats.py ++++ b/python/sglang/srt/observability/req_time_stats.py +@@ -346,7 +346,7 @@ class ReqTimeStatsBase: + state["trace_ctx"] = TraceNullContext() + + for key in state.keys(): +- if key.endswith("time"): ++ if key.endswith("time") and state[key] > 0.0: + state[key] = convert_time_cross_thread( + state[key], + state["diff_realtime_monotonic"], +@@ -632,9 +632,19 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + + state = { + "has_timing_data": True, ++ "enable_metrics": self.enable_metrics, ++ "disagg_mode": self.disagg_mode, + "wait_queue_entry_time": self.wait_queue_entry_time, + "forward_entry_time": self.forward_entry_time, + "prefill_finished_time": self.prefill_finished_time, ++ "completion_time": self.completion_time, ++ "prefill_bootstrap_queue_entry_time": self.prefill_bootstrap_queue_entry_time, ++ "prefill_transfer_queue_entry_time": self.prefill_transfer_queue_entry_time, ++ "decode_prealloc_queue_entry_time": self.decode_prealloc_queue_entry_time, ++ "decode_transfer_queue_entry_time": self.decode_transfer_queue_entry_time, ++ "bootstrap_done_time": self.bootstrap_done_time, ++ "transfer_speed_gb_s": self.transfer_speed_gb_s, ++ "transfer_total_mb": self.transfer_total_mb, + "diff_realtime_monotonic": global_diff_realtime_monotonic, + } + return state +@@ -1149,6 +1159,13 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + + def convert_to_output_meta_info(self): + meta_data = {} ++ ++ def add_duration(key: str, start: float, end: float): ++ if start > 0.0 and end > 0.0: ++ duration = end - start ++ if duration >= 0.0: ++ meta_data[key] = duration ++ + if self.forward_entry_time > 0.0: + meta_data["forward_entry_time"] = convert_time_to_realtime( + self.forward_entry_time +@@ -1162,6 +1179,66 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + "queue_time": self.get_queueing_time(), + } + ) ++ if self.disagg_mode == DisaggregationMode.PREFILL: ++ add_duration( ++ "pd_prefill_bootstrap_queue_duration", ++ self.prefill_bootstrap_queue_entry_time, ++ self.wait_queue_entry_time, ++ ) ++ add_duration( ++ "pd_prefill_bootstrap_duration", ++ self.prefill_bootstrap_queue_entry_time, ++ self.bootstrap_done_time, ++ ) ++ add_duration( ++ "pd_prefill_alloc_wait_duration", ++ self.bootstrap_done_time, ++ self.wait_queue_entry_time, ++ ) ++ add_duration( ++ "pd_prefill_forward_duration", ++ self.forward_entry_time, ++ self.completion_time, ++ ) ++ add_duration( ++ "pd_prefill_transfer_queue_duration", ++ self.prefill_transfer_queue_entry_time, ++ self.completion_time, ++ ) ++ if self.transfer_speed_gb_s > 0.0: ++ meta_data["pd_transfer_speed_gb_s"] = self.transfer_speed_gb_s ++ if self.transfer_total_mb > 0.0: ++ meta_data["pd_transfer_total_mb"] = self.transfer_total_mb ++ elif self.disagg_mode == DisaggregationMode.DECODE: ++ add_duration( ++ "pd_decode_prealloc_duration", ++ self.decode_prealloc_queue_entry_time, ++ self.decode_transfer_queue_entry_time, ++ ) ++ add_duration( ++ "pd_decode_bootstrap_duration", ++ self.decode_prealloc_queue_entry_time, ++ self.bootstrap_done_time, ++ ) ++ add_duration( ++ "pd_decode_alloc_wait_duration", ++ self.bootstrap_done_time, ++ self.decode_transfer_queue_entry_time, ++ ) ++ add_duration( ++ "pd_decode_transfer_duration", ++ self.decode_transfer_queue_entry_time, ++ self.wait_queue_entry_time, ++ ) ++ add_duration( ++ "pd_decode_forward_duration", ++ self.forward_entry_time, ++ self.completion_time, ++ ) ++ if self.transfer_speed_gb_s > 0.0: ++ meta_data["pd_transfer_speed_gb_s"] = self.transfer_speed_gb_s ++ if self.transfer_total_mb > 0.0: ++ meta_data["pd_transfer_total_mb"] = self.transfer_total_mb + return meta_data + + def format_duration(self, duration: float) -> str: +diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +index 13371f43..1bd91856 100644 +--- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py ++++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +@@ -566,8 +566,10 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): + forward_batch.seq_lens, + forward_batch.out_cache_loc, + forward_batch.positions, +- forward_batch.spec_info.topk_p, +- forward_batch.spec_info.topk_index, ++ forward_batch.spec_info.topk_p.clamp(0.0, 1.0), ++ forward_batch.spec_info.topk_index.clamp( ++ 0, self.model_runner.model_config.vocab_size - 1 ++ ), + forward_batch.req_pool_indices, + ] + if buffers.rids_int is not None and forward_batch.rids_int is not None: +diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py +index 6636d06a..53142e01 100644 +--- a/python/sglang/srt/utils/common.py ++++ b/python/sglang/srt/utils/common.py +@@ -2834,6 +2834,9 @@ class SafeUnpickler(pickle.Unpickler): + "sglang.srt.utils.", + "sglang.srt.disaggregation.", + "sglang.srt.managers.", ++ "slime.", ++ # --- Relax --- ++ "Relax.", + "torch_npu.", + } + diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index 0264ae415..f8e36d861 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -246,7 +246,6 @@ def _resolve_external_model_arch(package_name): def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process: multiprocessing.set_start_method("spawn", force=True) - server_args.host = server_args.host.strip("[]") # Each SGLang patch is controlled by its own env flag and applied # independently (see ``_launch_server_with_patches`` and @@ -589,6 +588,7 @@ def _init_normal(self, server_args_dict, *, apply_policy_load_plan: bool = True) warm_hf_checkpoint_page_cache(server_args_dict.get("model_path")) + server_args_dict = {**server_args_dict, "host": server_args_dict["host"].strip("[]")} self.process = launch_server_process(ServerArgs(**server_args_dict)) bootstrap_port = ( diff --git a/tests/backends/sglang/test_image_patch.py b/tests/backends/sglang/test_image_patch.py deleted file mode 100644 index e41d5a0b3..000000000 --- a/tests/backends/sglang/test_image_patch.py +++ /dev/null @@ -1,28 +0,0 @@ -# Copyright (c) 2026 Relax Authors. All Rights Reserved. - -from pathlib import Path - - -_REPO_ROOT = Path(__file__).resolve().parents[3] -_SGLANG_PATCH = _REPO_ROOT / "docker" / "patch" / "latest" / "sglang.patch" - - -def test_runai_processor_resolves_model_name_uri(): - patch = _SGLANG_PATCH.read_text() - processor_diff = patch.split( - "diff --git a/python/sglang/srt/utils/hf_transformers/processor.py ", - maxsplit=1, - )[1] - processor_diff = processor_diff.split("\ndiff --git ", maxsplit=1)[0] - normalized_diff = processor_diff.replace("\n \n", "\n\n") - - expected_hunk = """@@ -153,6 +153,8 @@ def get_processor( - - revision = kwargs.pop("revision", tokenizer_revision) - tokenizer_name = resolve_runai_obj_uri(tokenizer_name) -+ if model_name is not None: -+ model_name = resolve_runai_obj_uri(model_name) - - if is_mistral_model(tokenizer_name): - config = load_mistral_config(""" - assert expected_hunk in normalized_diff From fdb5cd0a0eaa067aef9d93d4495ee66e03fb32e2 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 31 Aug 2026 20:20:42 +0800 Subject: [PATCH 04/34] upgrade megatron and add deepseek patch (cherry picked from commit 2d5f1cd141a03ae291d1645c0ccb3f6a3fccdced) --- docker/Dockerfile | 20 +- docker/patch/latest/megatron.patch | 2 +- .../patch/megatron/20260728-0e6ac576f.patch | 1773 +++++++++++++++++ 3 files changed, 1792 insertions(+), 3 deletions(-) create mode 100644 docker/patch/megatron/20260728-0e6ac576f.patch diff --git a/docker/Dockerfile b/docker/Dockerfile index 312475a81..1c4a7b792 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -38,9 +38,24 @@ RUN pip install nvidia-cudnn-cu12==9.16.0.29 FROM base as train RUN MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation --no-cache-dir && \ - pip install --no-cache-dir flash-linear-attention==0.4.1 && \ + pip install --no-cache-dir flash-linear-attention==0.4.2 && \ pip install --no-cache-dir tilelang -f https://tile-ai.github.io/whl/nightly/cu128/ +# DeepSeek-V4 DSA deps, both listed under mcore's `no_pypi_wheels` so built from pinned source: +# - fast-hadamard-transform: required by the Lightning Indexer (dsa.rotate_activation asserts on it) +# - FlashMLA: fused DSA sparse attention; without it the CSA backward is quadratic in seq len +# (mcore only allows the fusion on SM90 when dsa_indexer_loss_coeff == 0) +ARG FLASH_MLA_COMMIT=b7643bd54521f563b839b98289b5cd048c062ba2 +RUN MAX_JOBS=64 pip install --no-cache-dir --no-build-isolation --no-deps \ + git+https://github.com/Dao-AILab/fast-hadamard-transform.git@f134af63deb2df17e1171a9ec1ea4a7d8604d5ca && \ + git clone --recurse-submodules https://github.com/deepseek-ai/FlashMLA.git /opt/FlashMLA && \ + cd /opt/FlashMLA && git checkout ${FLASH_MLA_COMMIT} && \ + git submodule update --init --recursive && \ + MAX_JOBS=64 pip install --no-cache-dir --no-build-isolation --no-deps . && \ + cd / && python -c "import flash_mla; assert hasattr(flash_mla, 'flash_mla_sparse_fwd'), \ + 'flash_mla built but flash_mla_sparse_fwd missing'" && \ + rm -rf /opt/FlashMLA + # FA3 (Hopper flash-attention), built from source. This commit's _flash_attn_forward carries # window_size_left/right to match TE 2.14.1 (docs/draft/sglang-0.5.12-upgrade-plan.md §8.4). # The `cp` exposes flash_attn_3.flash_attn_interface, which is what TE imports. BUT this commit's @@ -70,6 +85,7 @@ RUN MAX_JOBS=64 \ RUN pip -v install --no-cache-dir --no-build-isolation "transformer_engine[pytorch]==2.14.1" && \ TMS_CUDA_MAJOR=$(python -c 'import torch; print(torch.version.cuda.split(".")[0])') pip install git+https://github.com/redai-studio/torch_memory_saver.git@afc13785c50119048e2dd8ac497cc9e29ec75bd4 --no-cache-dir --force-reinstall && \ pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation --no-cache-dir && \ + pip install megatron-energon fsspec==2024.3.1 --no-cache-dir&& \ pip install "numpy<2" nvidia-cudnn-cu12==9.16.0.29 --no-cache-dir && \ NVCC_APPEND_FLAGS="--threads 32" \ pip -v install --disable-pip-version-check --no-cache-dir \ @@ -93,7 +109,7 @@ WORKDIR /root ARG PATCH_VERSION=latest ARG ENABLE_SGLANG_PATCH=1 -ARG MEGATRON_BRIDGE_COMMIT=2faedbf6fe3c422835a44b2b360cadcb2a116a54 +ARG MEGATRON_BRIDGE_COMMIT=af17edf52c514c58c79cd291b04a9b42bc5c57f3 ENV MEGATRON_BRIDGE_COMMIT=${MEGATRON_BRIDGE_COMMIT} \ PYTHONPATH=/root/Megatron-LM/ diff --git a/docker/patch/latest/megatron.patch b/docker/patch/latest/megatron.patch index ec9557dc5..c7ffa0d11 120000 --- a/docker/patch/latest/megatron.patch +++ b/docker/patch/latest/megatron.patch @@ -1 +1 @@ -../megatron/20260506-85bced0ae.patch \ No newline at end of file +../megatron/20260728-0e6ac576f.patch \ No newline at end of file diff --git a/docker/patch/megatron/20260728-0e6ac576f.patch b/docker/patch/megatron/20260728-0e6ac576f.patch new file mode 100644 index 000000000..de0823ca1 --- /dev/null +++ b/docker/patch/megatron/20260728-0e6ac576f.patch @@ -0,0 +1,1773 @@ +diff --git a/megatron/bridge/models/gemma/gemma4_bridge.py b/megatron/bridge/models/gemma/gemma4_bridge.py +index a62eb81..c858022 100644 +--- a/megatron/bridge/models/gemma/gemma4_bridge.py ++++ b/megatron/bridge/models/gemma/gemma4_bridge.py +@@ -499,71 +499,81 @@ class Gemma4Bridge(MegatronModelBridge): + """Text-only CausalLM: weights at ``model.*``; override in VL subclass.""" + return "model." + +- def _moe_mapping_registry(self) -> MegatronMappingRegistry: +- """Parameter mappings for the MoE variant.""" ++ def _moe_mapping_registry(self, megatron_prefix: str = "") -> MegatronMappingRegistry: ++ """Parameter mappings for the MoE variant. ++ ++ ``megatron_prefix`` and :meth:`_hf_layer_prefix` mirror ++ :meth:`_dense_mapping_registry`, so a VL subclass can reuse this for a ++ text-only conversion of a ``Gemma4ForConditionalGeneration`` checkpoint, ++ whose language weights sit under ``model.language_model.*``. ++ """ ++ mp = megatron_prefix ++ hp = self._hf_layer_prefix() + param_mappings = { +- "embedding.word_embeddings.weight": "model.embed_tokens.weight", +- "decoder.final_layernorm.weight": "model.norm.weight", +- "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.layers.*.input_layernorm.weight", +- "decoder.layers.*.self_attention.q_layernorm.weight": "model.layers.*.self_attn.q_norm.weight", +- "decoder.layers.*.self_attention.k_layernorm.weight": "model.layers.*.self_attn.k_norm.weight", +- "decoder.layers.*.self_attention.linear_proj.weight": "model.layers.*.self_attn.o_proj.weight", +- "decoder.layers.*.self_attention.linear_proj.post_layernorm.weight": ( +- "model.layers.*.post_attention_layernorm.weight" ++ f"{mp}embedding.word_embeddings.weight": f"{hp}embed_tokens.weight", ++ f"{mp}decoder.final_layernorm.weight": f"{hp}norm.weight", ++ f"{mp}decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": ( ++ f"{hp}layers.*.input_layernorm.weight" ++ ), ++ f"{mp}decoder.layers.*.self_attention.q_layernorm.weight": f"{hp}layers.*.self_attn.q_norm.weight", ++ f"{mp}decoder.layers.*.self_attention.k_layernorm.weight": f"{hp}layers.*.self_attn.k_norm.weight", ++ f"{mp}decoder.layers.*.self_attention.linear_proj.weight": f"{hp}layers.*.self_attn.o_proj.weight", ++ f"{mp}decoder.layers.*.self_attention.linear_proj.post_layernorm.weight": ( ++ f"{hp}layers.*.post_attention_layernorm.weight" + ), +- "decoder.layers.*.pre_mlp_layernorm.weight": "model.layers.*.pre_feedforward_layernorm_2.weight", +- "decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.layers.*.mlp.down_proj.weight", +- "decoder.layers.*.mlp.post_shared_expert_layernorm.weight": ( +- "model.layers.*.post_feedforward_layernorm_1.weight" ++ f"{mp}decoder.layers.*.pre_mlp_layernorm.weight": f"{hp}layers.*.pre_feedforward_layernorm_2.weight", ++ f"{mp}decoder.layers.*.mlp.shared_experts.linear_fc2.weight": f"{hp}layers.*.mlp.down_proj.weight", ++ f"{mp}decoder.layers.*.mlp.post_shared_expert_layernorm.weight": ( ++ f"{hp}layers.*.post_feedforward_layernorm_1.weight" + ), +- "decoder.layers.*.mlp.router.weight": "model.layers.*.router.proj.weight", ++ f"{mp}decoder.layers.*.mlp.router.weight": f"{hp}layers.*.router.proj.weight", + } + + mapping_list = [AutoMapping(megatron_param=m, hf_param=h) for m, h in param_mappings.items()] + mapping_list.extend( + [ + _Gemma4QKVMapping( +- megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", +- q="model.layers.*.self_attn.q_proj.weight", +- k="model.layers.*.self_attn.k_proj.weight", +- v="model.layers.*.self_attn.v_proj.weight", ++ megatron_param=f"{mp}decoder.layers.*.self_attention.linear_qkv.weight", ++ q=f"{hp}layers.*.self_attn.q_proj.weight", ++ k=f"{hp}layers.*.self_attn.k_proj.weight", ++ v=f"{hp}layers.*.self_attn.v_proj.weight", + ), + GatedMLPMapping( +- megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight", +- gate="model.layers.*.mlp.gate_proj.weight", +- up="model.layers.*.mlp.up_proj.weight", ++ megatron_param=f"{mp}decoder.layers.*.mlp.shared_experts.linear_fc1.weight", ++ gate=f"{hp}layers.*.mlp.gate_proj.weight", ++ up=f"{hp}layers.*.mlp.up_proj.weight", + ), + FusedGatedExpertMapping( +- megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", +- hf_param="model.layers.*.experts.gate_up_proj", ++ megatron_param=f"{mp}decoder.layers.*.mlp.experts.linear_fc1.weight*", ++ hf_param=f"{hp}layers.*.experts.gate_up_proj", + ), + FusedExpertMapping( +- megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", +- hf_param="model.layers.*.experts.down_proj", ++ megatron_param=f"{mp}decoder.layers.*.mlp.experts.linear_fc2.weight*", ++ hf_param=f"{hp}layers.*.experts.down_proj", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.layer_scalar", +- hf_param="model.layers.*.layer_scalar", ++ megatron_param=f"{mp}decoder.layers.*.layer_scalar", ++ hf_param=f"{hp}layers.*.layer_scalar", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.mlp.router.per_expert_scale", +- hf_param="model.layers.*.router.per_expert_scale", ++ megatron_param=f"{mp}decoder.layers.*.mlp.router.per_expert_scale", ++ hf_param=f"{hp}layers.*.router.per_expert_scale", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.mlp.router.scale", +- hf_param="model.layers.*.router.scale", ++ megatron_param=f"{mp}decoder.layers.*.mlp.router.scale", ++ hf_param=f"{hp}layers.*.router.scale", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.pre_shared_expert_layernorm.weight", +- hf_param="model.layers.*.pre_feedforward_layernorm.weight", ++ megatron_param=f"{mp}decoder.layers.*.pre_shared_expert_layernorm.weight", ++ hf_param=f"{hp}layers.*.pre_feedforward_layernorm.weight", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.mlp.post_moe_layernorm.weight", +- hf_param="model.layers.*.post_feedforward_layernorm_2.weight", ++ megatron_param=f"{mp}decoder.layers.*.mlp.post_moe_layernorm.weight", ++ hf_param=f"{hp}layers.*.post_feedforward_layernorm_2.weight", + ), + ReplicatedMapping( +- megatron_param="decoder.layers.*.post_ffn_layernorm.weight", +- hf_param="model.layers.*.post_feedforward_layernorm.weight", ++ megatron_param=f"{mp}decoder.layers.*.post_ffn_layernorm.weight", ++ hf_param=f"{hp}layers.*.post_feedforward_layernorm.weight", + ), + ] + ) +diff --git a/megatron/bridge/models/gemma/modeling_gemma4.py b/megatron/bridge/models/gemma/modeling_gemma4.py +index 1c1dda9..b5a0a9b 100644 +--- a/megatron/bridge/models/gemma/modeling_gemma4.py ++++ b/megatron/bridge/models/gemma/modeling_gemma4.py +@@ -456,13 +456,27 @@ class Gemma4DenseSelfAttention(SelfAttention): + self, + hidden_states: Tensor, + key_value_states=None, +- output_gate: bool = False, +- split_qkv: bool = True, ++ **kwargs, + ): ++ """Take core's qkv keyword arguments verbatim and forward them unchanged. ++ ++ Enumerating them instead pins this class to one Megatron-LM lineage: the ++ .dev.commit core (0e6ac576) added `head_wise_gate` to ++ Attention.get_query_key_value_tensors, between `output_gate` and ++ `split_qkv`, and Attention.forward passes it unconditionally by keyword. ++ Gemma4SelfAttention below already takes **kwargs for the same reason. ++ """ ++ split_qkv = kwargs.get("split_qkv", True) ++ output_gate = kwargs.get("output_gate", False) ++ + if self.is_kv_shared_layer: + if not split_qkv or output_gate: +- return super().get_query_key_value_tensors(hidden_states, key_value_states, output_gate, split_qkv) +- query, _k, _v = super().get_query_key_value_tensors(hidden_states, key_value_states, False, True) ++ return super().get_query_key_value_tensors(hidden_states, key_value_states, **kwargs) ++ # Not forwarding **kwargs here is deliberate: this path unpacks a ++ # plain (q, k, v), and any gate flag would make core return more. ++ query, _k, _v = super().get_query_key_value_tensors( ++ hidden_states, key_value_states, output_gate=False, split_qkv=True ++ ) + kv_source = self._kv_source_ref() if self._kv_source_ref is not None else None + if kv_source is not None and kv_source._stored_kv is not None: + key, value = kv_source._stored_kv +@@ -479,7 +493,7 @@ class Gemma4DenseSelfAttention(SelfAttention): + key_value_states, + ) + else: +- result = super().get_query_key_value_tensors(hidden_states, key_value_states, output_gate, split_qkv) ++ result = super().get_query_key_value_tensors(hidden_states, key_value_states, **kwargs) + if not split_qkv: + return result + if output_gate: +diff --git a/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py b/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py +index 918cac7..2bd3011 100644 +--- a/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py ++++ b/megatron/bridge/models/gemma_vl/gemma4_vl_bridge.py +@@ -92,6 +92,15 @@ class Gemma4VLBridge(Gemma4Bridge): + + self._is_dense = False + ++ if self._conversion_mode() == "text": ++ # Same shape as the dense branch above: drop the vision/audio towers and ++ # build the plain MoE GPT provider from ``text_config``. Every field ++ # ``_build_moe_provider`` reads (sliding_window, rope_parameters, head_dim, ++ # global_head_dim, num_global_key_value_heads, attention_k_eq_v, ++ # layer_types, num_experts, top_k_experts, moe_intermediate_size, ++ # intermediate_size, final_logit_softcapping) lives on ``text_config``. ++ return self._build_moe_provider(text_config) ++ + provider_kwargs = self.hf_config_to_provider_kwargs(text_config) + provider = Gemma4VLModelProvider(**provider_kwargs) + +@@ -226,6 +235,10 @@ class Gemma4VLBridge(Gemma4Bridge): + if self._conversion_mode() == "text": + return self._dense_mapping_registry(megatron_prefix="") + return self._dense_vl_mapping_registry() ++ if self._conversion_mode() == "text": ++ # ``_hf_layer_prefix`` still resolves to ``model.language_model.`` here -- ++ # the checkpoint layout does not change, only which submodules we build. ++ return self._moe_mapping_registry(megatron_prefix="") + return self._moe_vl_mapping_registry() + + def _dense_vl_mapping_registry(self) -> MegatronMappingRegistry: +diff --git a/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py b/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +index fa13752..b6c0cb9 100644 +--- a/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py ++++ b/megatron/bridge/models/kimi_vl/kimi_k25_vl_bridge.py +@@ -191,8 +191,15 @@ class KimiK25VLBridge(MegatronModelBridge): + for fqn, tensor in converted_weights_dict.items(): + if self._is_quantized_expert_key(fqn): + base = fqn[:-7] if fqn.endswith(".weight") else fqn +- # Preserve the original scale dtype from the HF checkpoint + orig_scale_key = f"{base}.weight_scale" ++ # When the source HF checkpoint has been pre-cast to BF16 (no ++ # `weight_scale` triplet present), passthrough instead of re- ++ # quantizing — downstream sglang loads BF16 and would reject ++ # the INT4 export names with "not found in params_dict". ++ if orig_scale_key not in hf_state_dict: ++ result[fqn] = tensor ++ continue ++ # Preserve the original scale dtype from the HF checkpoint + scale_dtype = ( + hf_state_dict[orig_scale_key].dtype if orig_scale_key in hf_state_dict else torch.bfloat16 + ) +diff --git a/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py b/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py +index ba632f3..99e8aa9 100644 +--- a/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py ++++ b/megatron/bridge/models/kimi_vl/kimi_k25_vl_provider.py +@@ -60,6 +60,11 @@ class KimiK25VLModelProvider(MLAModelProvider): + pad_token_id: int = 163839 + ignore_index: int = -100 + ++ # Split vision encoder workload across TP ranks (data-parallel over TP). ++ # Each TP rank processes a chunk of images, then all-reduce gathers the ++ # full embedding. Reduces per-GPU peak memory for the vision encoder. ++ vision_dp_when_tp: bool = False ++ + # Freeze options for fine-tuning scenarios + freeze_language_model: bool = False + freeze_vision_model: bool = False +diff --git a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py +index c2bf410..073f570 100644 +--- a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py ++++ b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py +@@ -16,7 +16,9 @@ import logging + from typing import List, Optional + + import torch ++import torch.distributed + from megatron.core import parallel_state ++from megatron.core import parallel_state as mpu + from megatron.core.packed_seq_params import PackedSeqParams + from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region + from megatron.core.transformer.module import MegatronModule +@@ -25,6 +27,7 @@ from transformers.dynamic_module_utils import get_class_from_dynamic_module + from transformers.utils import is_flash_attn_2_available + + from megatron.bridge.models.gpt_provider import GPTModelProvider ++from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.utils import preprocess_packed_seqs + from megatron.bridge.utils.common_utils import hook_hf_module_setattr_for_tp_grad_sync + + +@@ -388,8 +391,79 @@ class KimiK25VLModel(MegatronModule): + + return final_embedding, final_attention_mask, final_labels, position_ids + ++ def _vision_forward_tp_split( ++ self, ++ pixel_values: torch.Tensor, ++ grid_thws: torch.Tensor, ++ ) -> List[torch.Tensor]: ++ """Run vision encoder + projector with workload split across TP ranks. ++ ++ Each TP rank processes a subset of images determined by splitting ++ ``grid_thws``, then the partial feature tensors are all-reduced so ++ every rank holds the complete result. ++ """ ++ tp_rank = mpu.get_tensor_model_parallel_rank() ++ tp_size = mpu.get_tensor_model_parallel_world_size() ++ ++ num_images = grid_thws.shape[0] ++ merge_h, merge_w = self.vision_tower.merge_kernel_size ++ param_dtype = next(self.vision_tower.parameters()).dtype ++ text_hidden = self.projector_config.hidden_size ++ ++ pixel_counts = grid_thws.prod(dim=-1) ++ out_token_counts = (grid_thws[:, 1] // merge_h) * (grid_thws[:, 2] // merge_w) ++ total_out_tokens = out_token_counts.sum().item() ++ ++ chunk_indices = list(range(num_images)) ++ chunks = [chunk_indices[i::tp_size] for i in range(tp_size)] ++ my_indices = chunks[tp_rank] if tp_rank < len(chunks) else [] ++ ++ out_buffer = torch.zeros( ++ (total_out_tokens, text_hidden), ++ device=pixel_values.device, ++ dtype=param_dtype, ++ ) ++ ++ if my_indices: ++ pixel_cumsum = pixel_counts.cumsum(dim=0) ++ pv_parts = [] ++ grid_parts = [] ++ for idx in my_indices: ++ px_start = 0 if idx == 0 else pixel_cumsum[idx - 1].item() ++ px_end = pixel_cumsum[idx].item() ++ pv_parts.append(pixel_values[px_start:px_end]) ++ grid_parts.append(grid_thws[idx : idx + 1]) ++ ++ local_pv = torch.cat(pv_parts, dim=0) ++ local_grid = torch.cat(grid_parts, dim=0) ++ ++ local_vit_out = self.vision_tower(local_pv, local_grid) ++ local_features = self.mm_projector(local_vit_out) ++ ++ out_cumsum = out_token_counts.cumsum(dim=0) ++ for feat_i, img_idx in enumerate(my_indices): ++ out_start = 0 if img_idx == 0 else out_cumsum[img_idx - 1].item() ++ out_end = out_cumsum[img_idx].item() ++ out_buffer[out_start:out_end] = local_features[feat_i].to(param_dtype) ++ ++ tp_group = mpu.get_tensor_model_parallel_group() ++ torch.distributed.all_reduce(out_buffer, group=tp_group) ++ ++ result = [] ++ out_cumsum = out_token_counts.cumsum(dim=0) ++ for i in range(num_images): ++ start = 0 if i == 0 else out_cumsum[i - 1].item() ++ end = out_cumsum[i].item() ++ result.append(out_buffer[start:end]) ++ ++ return result ++ + def _extract_image_features(self, pixel_values, grid_thws): + """Extract and project image features.""" ++ tp_size = mpu.get_tensor_model_parallel_world_size() ++ if getattr(self.config, "vision_dp_when_tp", False) and tp_size > 1: ++ return self._vision_forward_tp_split(pixel_values, grid_thws) ++ + image_features = self.vision_tower(pixel_values, grid_thws) + return self.mm_projector(image_features) + +@@ -430,6 +504,12 @@ class KimiK25VLModel(MegatronModule): + cp_size = parallel_state.get_context_parallel_world_size() + cp_rank = parallel_state.get_context_parallel_rank() if cp_size > 1 else 0 + if self.pre_process: ++ # Save the caller-supplied per-sample attention mask before any rewrite — ++ # _merge_input_ids_with_image_features sets `attention_mask = None` on the ++ # vision path, but the THD repack below needs the original [B, T] mask to ++ # know each sample's valid length. ++ saved_attention_mask = attention_mask ++ + if inputs_embeds is None: + inputs_embeds = self.language_model.embedding( + input_ids=input_ids, position_ids=None +@@ -465,8 +545,30 @@ class KimiK25VLModel(MegatronModule): + # Don't need attention mask for causal attention. + attention_mask = None + +- # Transpose back to (T, B, D) for Megatron language model +- inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() # (B, T, D) -> (T, B, D) ++ # When THD packed_seq_params is provided (VL+CP/SP), repack the raw ++ # padded [B, T_max, D] embedding into compact THD ++ # [sum(padded_seqlens)/cp, 1, D] using the saved per-sample attention ++ # mask. Mirrors the Qwen3VL bridge: without this, downstream MLA ++ # attention sees a tensor whose first dim does not match ++ # cu_seqlens_q_padded (sum=sum(padded_seqlens)), and SP scatter can ++ # hit `T_max % tp_size != 0` since T_max is raw batch-max. ++ # preprocess_packed_seqs also recomputes cu_seqlens with align64. ++ needs_thd_repack = ( ++ packed_seq_params is not None ++ and packed_seq_params.qkv_format == "thd" ++ and saved_attention_mask is not None ++ ) ++ if needs_thd_repack: ++ inputs_embeds, packed_seq_params = preprocess_packed_seqs( ++ inputs_embeds, # [B, T_max, D] ++ saved_attention_mask, ++ pre_process=True, ++ ) ++ # preprocess_packed_seqs returns [1, T_thd, D]; switch to (T_thd, 1, D) ++ inputs_embeds = inputs_embeds.transpose(0, 1).contiguous() ++ else: ++ # Transpose back to (T, B, D) for Megatron language model ++ inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() # (B, T, D) -> (T, B, D) + + if cp_size > 1: + inputs_embeds = _split_on_cp_rank(inputs_embeds, cp_size, cp_rank, seq_dim=0) +diff --git a/megatron/bridge/models/qwen/qwen3_moe_bridge.py b/megatron/bridge/models/qwen/qwen3_moe_bridge.py +index a6476f4..86d83b3 100755 +--- a/megatron/bridge/models/qwen/qwen3_moe_bridge.py ++++ b/megatron/bridge/models/qwen/qwen3_moe_bridge.py +@@ -69,6 +69,40 @@ class Qwen3MoEBridge(MegatronModelBridge): + + return provider + ++ def build_conversion_tasks(self, hf_pretrained, megatron_model): ++ """Inject virtual .weight keys so INT4 checkpoints (weight_packed/weight_scale/ ++ weight_zero_point) pass the hf_keys validation in the base class. ++ ++ When hf_checkpoint points to an INT4 compressed-tensors checkpoint, expert ++ weights are stored as weight_packed/weight_scale/weight_zero_point triplets ++ with no plain .weight key. The base build_conversion_tasks checks that each ++ mapped HF name exists in hf_keys and skips the param if not found, causing ++ all expert weights to be silently dropped. We patch get_all_keys() to return ++ synthetic .weight keys alongside the real packed keys so the check passes. ++ Downstream quantize_params in HfWeightIteratorBridge then converts the BF16 ++ output back to INT4 before sending to the rollout engine. ++ """ ++ if not (hasattr(hf_pretrained, "state") and hasattr(hf_pretrained.state, "source")): ++ return super().build_conversion_tasks(hf_pretrained, megatron_model) ++ ++ original_get_all_keys = hf_pretrained.state.source.get_all_keys ++ ++ def _get_all_keys_with_virtual(): ++ keys = original_get_all_keys() ++ all_keys_set = set(keys) ++ virtual_keys = [ ++ key[:-7] # "...weight_packed" -> "...weight" ++ for key in keys ++ if key.endswith("_packed") and f"{key[:-7]}_scale" in all_keys_set ++ ] ++ return keys + virtual_keys ++ ++ hf_pretrained.state.source.get_all_keys = _get_all_keys_with_virtual ++ try: ++ return super().build_conversion_tasks(hf_pretrained, megatron_model) ++ finally: ++ hf_pretrained.state.source.get_all_keys = original_get_all_keys ++ + def mapping_registry(self) -> MegatronMappingRegistry: + # Return MegatronMappingRegistry containing parameter mappings from Megatron to HF format + # First create simple 1:1 parameter mappings using a dictionary for readability +diff --git a/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py b/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +--- a/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py ++++ b/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/text_model.py +@@ -158,6 +158,7 @@ + *, + inference_params: Optional[BaseInferenceContext] = None, + loss_mask: Optional[Tensor] = None, ++ mtp_kwargs: Optional[dict] = None, + # args for deepstack + visual_pos_masks: Optional[torch.Tensor] = None, + deepstack_visual_embeds: Optional[list[torch.Tensor]] = None, +@@ -259,6 +260,7 @@ + inference_context=inference_context, + output_processor=output_processor, + output_processor_context=output_processor_context, ++ mtp_kwargs=mtp_kwargs, + ) + + if _shadow_embedding: +diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py +index 8d65e29..e4af4f2 100644 +--- a/megatron/core/dist_checkpointing/strategies/torch.py ++++ b/megatron/core/dist_checkpointing/strategies/torch.py +@@ -501,10 +501,12 @@ class MCoreLoadPlanner(DefaultLoadPlanner): + def _validate_global_shapes(self, metadata, sharded_tensors): + for sh_ten in sharded_tensors: + if sh_ten.key not in metadata.state_dict_metadata: +- raise KeyError( +- f"{sh_ten.key} from model not in state dict:" +- f" {sorted(metadata.state_dict_metadata.keys())}" +- ) ++ # raise KeyError( ++ # f"{sh_ten.key} from model not in state dict:" ++ # f" {sorted(metadata.state_dict_metadata.keys())}" ++ # ) ++ print(f"{sh_ten.key} from model not in state dict, will skip") ++ continue + loaded_shape = metadata.state_dict_metadata[sh_ten.key].size + expected_shape = sh_ten.global_shape + if loaded_shape != expected_shape: +@@ -528,7 +530,7 @@ class MCoreLoadPlanner(DefaultLoadPlanner): + tensor_metadata = self.metadata.state_dict_metadata + metadata_with_sizes = [ + (tensor_metadata[key], tensor_metadata[key].size, sharded_tensor) +- for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() ++ for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() if key in tensor_metadata + ] + try: + # Temporarily set sizes to expected shapes +@@ -898,6 +900,7 @@ class TorchDistLoadShardedStrategy: + planner=MCoreLoadPlanner( + shapes_validation_sharded_tensors=flexible_shape_sharded_tensors, + allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors, ++ allow_partial_load=True, + flatten_state_dict=False, + flatten_sharded_tensors=False, + ), +diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py +index 8d797e8..98f6241 100644 +--- a/megatron/core/extensions/transformer_engine.py ++++ b/megatron/core/extensions/transformer_engine.py +@@ -918,6 +918,7 @@ class TELinear(te.pytorch.Linear): + ) + + for param in self.parameters(): ++ setattr(param, "parallel_mode", parallel_mode) + if is_expert: + # Reduce the gradient on the expert_data_parallel group for expert linear layers + setattr(param, "allreduce", not self.expert_parallel) +@@ -1948,6 +1949,61 @@ class TEDotProductAttention(te.pytorch.DotProductAttention): + + + if HAVE_TE and is_te_min_version("1.9.0.dev0"): ++ def ceil_div(x: int, y: int) -> int: ++ return (x + y - 1) // y ++ ++ class _FakeInt4QuantizationSTE(torch.autograd.Function): ++ @staticmethod ++ def forward(ctx, x, group_size): ++ m, n = x.shape ++ block_size_m, block_size_n = 1, group_size ++ ++ ++ m_padded = ceil_div(m, block_size_m) * block_size_m ++ n_padded = ceil_div(n, block_size_n) * block_size_n ++ ++ x_padded = torch.zeros( ++ (m_padded, n_padded), ++ dtype=x.dtype, device=x.device ++ ) ++ x_padded[:m, :n] = x ++ ++ x_view = x_padded.view( ++ m_padded // block_size_m, ++ block_size_m, ++ n_padded // block_size_n, ++ block_size_n ++ ) ++ ++ x_max = x_view.abs().float().amax(dim=(1, 3), keepdim=True) ++ q_max = 7 ++ x_scale = x_max / q_max ++ ++ x_scale = x_scale.clamp(min=1e-5) ++ ++ x_div = x_view / x_scale ++ x_round = torch.round(x_div) ++ ++ x_q_clamped = x_round.clamp(-q_max, q_max) ++ ++ x_dequant_view = x_q_clamped * x_scale ++ ++ x_dequant_full = x_dequant_view.view_as(x_padded) ++ x_out = x_dequant_full[:m, :n].contiguous().to(x.dtype) ++ ++ return x_out ++ ++ @staticmethod ++ def backward(ctx, grad_output): ++ return grad_output, None ++ ++ def fake_int4_quantization_ste(x, group_size): ++ x_out = _FakeInt4QuantizationSTE.apply(x, group_size) ++ ++ if hasattr(x, 'main_grad'): ++ x_out.main_grad = x.main_grad ++ ++ return x_out + + class TEGroupedLinear(te.pytorch.GroupedLinear): + """ +@@ -2179,6 +2235,7 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"): + "amax_history_bwd": torch.cat( + [state["amax_history_bwd"].view(-1, 1) for state in state_list], + dim=1, ++ + ).view(self.fp8_meta["recipe"].amax_history_len, -1), + } + ) +@@ -2304,10 +2361,31 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"): + def _get_weight_tensors(self): + """Get the weight tensors of the module.""" + weight_tensors = super()._get_weight_tensors() +- return maybe_fake_quantize_int4_weight_tensors( ++ ++ # Two independent int4 fake-QAT paths coexist here. Upstream's is ++ # config-driven and ships with the .dev.commit core; Relax's is the ++ # env-var one below, which is load-bearing -- documented in ++ # docs/zh/examples/low-precision-training.md and switched on by ++ # scripts/training/{text,multimodal}/run-{qwen3-30B-A3B,kimi-k2.6}*int4*.sh. ++ # The Relax patch appends its version verbatim, which alongside ++ # upstream's would define this method twice (second silently wins, ++ # disabling upstream's path). Compose them: each is a no-op unless its ++ # own switch is set. Enabling both at once would quantize twice -- a ++ # configuration error, not a case to arbitrate here. ++ weight_tensors = maybe_fake_quantize_int4_weight_tensors( + self.config, self.delay_wgrad_compute, weight_tensors + ) + ++ if os.getenv("OPEN_TRAINING_INT4_FAKE_QAT_FLAG", "0") == "1": ++ group_size = int(os.getenv("OPEN_TRAINING_INT4_GROUP_SIZE", "128")) ++ ++ weight_tensors = [ ++ fake_int4_quantization_ste(w, group_size) ++ for w in weight_tensors ++ ] ++ ++ return weight_tensors ++ + def _encode_extra_state(self, state): + # TE 2.0 changed the format of extra_state to be a byte tensor + if is_te_min_version("2.0.0"): +diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py +index d4cdaa2..31c7c91 100644 +--- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py ++++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py +@@ -513,6 +513,7 @@ def _mla_rope_fwd_kv_split_kernel( + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, ++ k_dim_ceil: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, +@@ -562,21 +563,27 @@ def _mla_rope_fwd_kv_split_kernel( + cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + +- KV_ptr = KV + pid_m * stride_kv_seq + pid_head * BLOCK_H * stride_kv_nheads +- kv_off = tl.arange(0, BLOCK_H)[:, None] * stride_kv_nheads +- mask = kv_off < head_num * stride_kv_nheads +- k_in_off = kv_off + tl.arange(0, k_dim)[None, :] +- v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :] +- k = tl.load(KV_ptr + k_in_off, mask=mask) +- v = tl.load(KV_ptr + v_in_off, mask=mask) ++ KV_ptr = KV + pid_m * stride_kv_seq # + pid_head * BLOCK_H * stride_kv_nheads ++ ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ kj_range = tl.arange(0, k_dim_ceil)[None, :] ++ mask_k = (ki_range < head_num) & (kj_range < k_dim) ++ mask_v = ki_range < head_num ++ k_off = ki_range * stride_kv_nheads + kj_range ++ if v_dim > 0: ++ v_off = ki_range * stride_kv_nheads + k_dim + tl.arange(0, v_dim)[None, :] ++ v = tl.load(KV_ptr + v_off, mask=mask_v) ++ else: ++ v = tl.zeros((BLOCK_H, 1), dtype=KV.dtype.element_ty) ++ k = tl.load(KV_ptr + k_off, mask=mask_k) + +- K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads +- V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads ++ K_ptr = O_KEY + pid_m * stride_k_seq # + pid_head * BLOCK_H * stride_k_nheads ++ V_ptr = O_VALUE + pid_m * stride_v_seq # + pid_head * BLOCK_H * stride_v_nheads + +- k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :] +- v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :] +- tl.store(K_ptr + k_out_off, k, mask=mask) +- tl.store(V_ptr + v_out_off, v, mask=mask) ++ k_out_off = ki_range * stride_k_nheads + kj_range ++ tl.store(K_ptr + k_out_off, k, mask=mask_k) ++ if v_dim > 0: ++ v_out_off = ki_range * stride_v_nheads + tl.arange(0, v_dim)[None, :] ++ tl.store(V_ptr + v_out_off, v, mask=mask_v) + + EMB = K_POS_EMB + pid_m * stride_emb_seq + # x1 = t[..., 0::2], x2 = t[..., 1::2] +@@ -588,24 +595,26 @@ def _mla_rope_fwd_kv_split_kernel( + x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) + ++ x_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ mask_x = x_range < head_num + if REMOVE_INTERLEAVING: + x_1_off = ( +- tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads ++ x_range * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] * 2 + ) + x_2_off = x_1_off + 1 +- tl.store(K_ptr + x_1_off, x_left, mask=mask) +- tl.store(K_ptr + x_2_off, x_right, mask=mask) ++ tl.store(K_ptr + x_1_off, x_left, mask=mask_x) ++ tl.store(K_ptr + x_2_off, x_right, mask=mask_x) + else: + x_left_off = ( +- tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads ++ x_range * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 +- tl.store(K_ptr + x_left_off, x_left, mask=mask) +- tl.store(K_ptr + x_right_off, x_right, mask=mask) ++ tl.store(K_ptr + x_left_off, x_left, mask=mask_x) ++ tl.store(K_ptr + x_right_off, x_right, mask=mask_x) + + + @triton.autotune( +@@ -631,6 +640,7 @@ def _mla_rope_bwd_kv_split_kernel( + SIN, + emb_dim: tl.constexpr, + k_dim: tl.constexpr, ++ k_dim_ceil: tl.constexpr, + v_dim: tl.constexpr, + head_num: tl.constexpr, + batch_size, +@@ -672,27 +682,32 @@ def _mla_rope_bwd_kv_split_kernel( + else: + token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size) + +- dKV_ptr = dKV + pid_m * stride_dkv_seq + pid_head * BLOCK_H * stride_dkv_nheads +- dkv_off = tl.arange(0, BLOCK_H)[:, None] * stride_dkv_nheads +- mask = dkv_off < head_num * stride_dkv_nheads +- dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :] +- dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :] +- +- dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads +- dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads +- dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :] +- dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :] +- dk = tl.load(dK_ptr + dk_in_off, mask=mask) +- dv = tl.load(dV_ptr + dv_in_off, mask=mask) +- tl.store(dKV_ptr + dk_out_off, dk, mask=mask) +- tl.store(dKV_ptr + dv_out_off, dv, mask=mask) ++ dKV_ptr = dKV + pid_m * stride_dkv_seq # + pid_head * BLOCK_H * stride_dkv_nheads ++ ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H ++ kj_range = tl.arange(0, k_dim_ceil)[None, :] ++ mask_k = (ki_range < head_num) & (kj_range < k_dim) ++ mask_v = ki_range < head_num ++ dk_out_off = ki_range * stride_dkv_nheads + kj_range ++ ++ dK_ptr = dK + pid_m * stride_dk_seq # + pid_head * BLOCK_H * stride_dk_nheads ++ dV_ptr = dV + pid_m * stride_dv_seq # + pid_head * BLOCK_H * stride_dv_nheads ++ dk_in_off = ki_range * stride_dk_nheads + kj_range ++ ++ dk = tl.load(dK_ptr + dk_in_off, mask=mask_k) ++ tl.store(dKV_ptr + dk_out_off, dk, mask=mask_k) ++ ++ if v_dim > 0: ++ dv_out_off = ki_range * stride_dkv_nheads + k_dim + tl.arange(0, v_dim)[None, :] ++ dv_in_off = ki_range * stride_dv_nheads + tl.arange(0, v_dim)[None, :] ++ dv = tl.load(dV_ptr + dv_in_off, mask=mask_v) ++ tl.store(dKV_ptr + dv_out_off, dv, mask=mask_v) + + if pid_head == 0: + x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32) + for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)): +- dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads +- x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim ++ dK_ptr = dK + pid_m * stride_dk_seq # + i * BLOCK_H * stride_dk_nheads ++ x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim + i * BLOCK_H * stride_dk_nheads + mask = x_off < head_num * stride_dk_nheads + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 +@@ -779,6 +794,7 @@ class _FusedMLARoPEKVSplit(torch.autograd.Function): + + o_key = kv.new_empty(total_seqlen, nheads, emb_dim + k_dim) + o_value = kv.new_empty(total_seqlen, nheads, v_dim) ++ k_dim_ceil = triton.next_power_of_2(k_dim) + + grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) + _mla_rope_fwd_kv_split_kernel[grid]( +@@ -790,6 +806,7 @@ class _FusedMLARoPEKVSplit(torch.autograd.Function): + sin, + emb_dim, + k_dim, ++ k_dim_ceil, + v_dim, + nheads, + batch_size, +@@ -849,6 +866,7 @@ class _FusedMLARoPEKVSplit(torch.autograd.Function): + + d_kv = dk.new_empty(total_seqlen, nheads, ctx.k_dim + ctx.v_dim) + d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim) ++ k_dim_ceil = triton.next_power_of_2(ctx.k_dim) + + grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) + _mla_rope_bwd_kv_split_kernel[grid]( +@@ -860,6 +878,7 @@ class _FusedMLARoPEKVSplit(torch.autograd.Function): + sin, + ctx.emb_dim, + ctx.k_dim, ++ k_dim_ceil, + ctx.v_dim, + nheads, + batch_size, +diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py +index 4b3e360..c31c80b 100644 +--- a/megatron/core/inference/contexts/dynamic_context.py ++++ b/megatron/core/inference/contexts/dynamic_context.py +@@ -67,7 +67,8 @@ except ImportError: + try: + from torch_memory_saver import torch_memory_saver + +- torch_memory_saver.hook_mode = "torch" ++ # Commented out: breaks SGLang CUDA graph (requires hook_mode="preload") ++ # torch_memory_saver.hook_mode = "torch" + HAVE_TORCH_MEMORY_SAVER = True + except ImportError: + HAVE_TORCH_MEMORY_SAVER = False +diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py +index 5b938cc..b681685 100755 +--- a/megatron/core/models/gpt/gpt_layer_specs.py ++++ b/megatron/core/models/gpt/gpt_layer_specs.py +@@ -195,6 +195,8 @@ def get_gpt_layer_with_transformer_engine_submodules( + mla_down_proj_fusion: bool = False, + dense_grouped_gemm: bool = False, + use_grouped_gemm_for_dense_mlp: bool = False, ++ post_self_attn_layernorm: bool = False, ++ post_mlp_layernorm: bool = False, + ) -> TransformerLayerSubmodules: + """Use these submodules to use lower-level Transformer Engine modules (required for fp8 + training). +@@ -289,9 +291,11 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + ), + self_attn_bda=get_bias_dropout_add, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + sharded_state_dict_keys_map=( + { + "self_attention.linear_q_down_proj.layer_norm_": "input_layernorm.", +@@ -321,10 +325,12 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + self_attn_bda=get_bias_dropout_add, + self_attention_hyper_connection=hc_module, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm(has_residual=True) if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, + mlp_hyper_connection=hc_module, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + ) + else: + qk_norm = backend.layer_norm(for_qk=True) +@@ -346,10 +352,12 @@ def get_gpt_layer_with_transformer_engine_submodules( + ), + self_attn_bda=get_bias_dropout_add, + self_attention_hyper_connection=hc_module, ++ post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp, + pre_mlp_layernorm=backend.layer_norm(has_residual=True) if num_experts else IdentityOp, + mlp=mlp, + mlp_bda=get_bias_dropout_add, + mlp_hyper_connection=hc_module, ++ post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp, + sharded_state_dict_keys_map={ + "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight", + "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias", +diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py +--- a/megatron/core/models/gpt/gpt_model.py ++++ b/megatron/core/models/gpt/gpt_model.py +@@ -34,6 +34,7 @@ + MultiTokenPredictionBlock, + mtp_on_this_rank, + process_mtp_loss, ++ roll_tensor, + ) + from megatron.core.transformer.spec_utils import ModuleSpec + from megatron.core.transformer.transformer_block import TransformerBlock +@@ -529,6 +530,7 @@ + padding_mask: Optional[Tensor] = None, + output_processor: Optional[Callable[..., Tensor]] = None, + output_processor_context: Optional[Any] = None, ++ mtp_kwargs: Optional[dict] = None, + ) -> Tensor: + """Forward function of the GPT Model This function passes the input tensors + through the embedding layer, and then the decoder and finally into the post +@@ -624,6 +626,7 @@ + mhc_multistream=mhc_multistream, + output_processor=output_processor, + output_processor_context=output_processor_context, ++ mtp_kwargs=mtp_kwargs, + ) + + def _postprocess( +@@ -649,6 +652,7 @@ + mhc_multistream=None, + output_processor=None, + output_processor_context=None, ++ mtp_kwargs=None, + ): + """Postprocesses decoder hidden states to generate logits or compute loss. + +@@ -673,7 +677,42 @@ + output_weight = None + if self.share_embeddings_and_output_weights: + output_weight = self.shared_embedding_or_output_weight() +- if mtp_in_postprocess and not (in_inference_mode or is_spec_decode): ++ # Relax always calls forward() with labels=None and computes the main loss ++ # itself, so it hands MTP its labels explicitly via `mtp_kwargs`. Upstream's ++ # `process_mtp_loss(input_ids=...)` fallback is not a substitute: on the ++ # VL+THD+CP unsplit path `input_ids` is the un-CP-split `unsplit_tokens` ++ # (twice the per-chunk length), so the derived labels do not line up with ++ # the CP-split hidden states and vocab_parallel_cross_entropy raises ++ # "shape mismatch: indexing tensors could not be broadcast together". ++ mtp_labels = labels ++ if mtp_kwargs is not None and mtp_kwargs.get("mtp_labels", None) is not None: ++ mtp_labels = mtp_kwargs["mtp_labels"] ++ mtp_labels, _ = roll_tensor( ++ mtp_labels, ++ shifts=-1, ++ dims=-1, ++ cp_group=self.pg_collection.cp, ++ packed_seq_params=packed_seq_params, ++ ) ++ if loss_mask is not None: ++ loss_mask, _ = roll_tensor( ++ loss_mask, ++ shifts=-1, ++ dims=-1, ++ cp_group=self.pg_collection.cp, ++ packed_seq_params=packed_seq_params, ++ ) ++ # `mtp_labels is not None` gates BOTH the MTP block below and the ++ # process_mtp_loss() call further down: Relax's log-prob forward ++ # (model.py forward_step) passes labels=None and no mtp_kwargs, and the ++ # old patch relied on process_mtp_loss returning early in that case. ++ # Upstream replaced that early return with the input_ids fallback, so the ++ # skip has to happen here instead. ++ if ( ++ mtp_in_postprocess ++ and not (in_inference_mode or is_spec_decode) ++ and mtp_labels is not None ++ ): + hidden_states = self.mtp( + input_ids=input_ids, + position_ids=position_ids, +@@ -694,7 +733,7 @@ + if not self.post_process: + return hidden_states + +- if self.config.mtp_num_layers: ++ if self.config.mtp_num_layers and mtp_labels is not None: + assert self.config.mtp_num_layers > 0 + if in_inference_mode or is_spec_decode: + # Cache decoder hidden states for serial MTP computation +@@ -705,7 +744,7 @@ + mtp_cp_group = resolve_cp_group(self.pg_collection.cp, packed_seq_params) + hidden_states = process_mtp_loss( + hidden_states=hidden_states, +- labels=labels, ++ labels=mtp_labels, + loss_mask=loss_mask, + output_layer=self.output_layer, + output_weight=output_weight, +diff --git a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +--- a/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py ++++ b/megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py +@@ -249,6 +249,11 @@ + return cpu_optimizers + + def _get_sub_optimizer_param_groups(self, offload_fraction: float): ++ import warnings ++ warnings.warn( ++ "CPU offload optimizer init can be very slow (potentially minutes) for " ++ "large MoE models due to per-parameter pinned-memory allocation and H2D copies." ++ ) + params = [] + for group in self.param_groups: + params.extend(group["params"]) +@@ -370,8 +375,11 @@ + if not self.param_update_in_fp32: + return + for param, v in self.state.items(): +- fp32_param = self.param_to_fp32_param[param] +- fp32_param.data.copy_(v["master_param"]) ++ # Native FP32 params do not need a separate master parameter and are ++ # intentionally absent from param_to_fp32_param. ++ fp32_param = self.param_to_fp32_param.get(param) ++ if fp32_param is not None: ++ fp32_param.data.copy_(v["master_param"]) + + def update_fp32_param_by_new_param(self): + """ +@@ -379,6 +387,16 @@ + """ + for param, fp32_param in self.param_to_fp32_param.items(): + fp32_param.data.copy_(param) ++ # Params that are already fp32 never get an entry in param_to_fp32_param ++ # (see _get_sub_optimizer_param_groups), so the loop above misses their CPU ++ # snapshot. That snapshot is taken at construction time and copied back over ++ # the model param by the post-step hook, which silently reverts anything ++ # written afterwards -- a checkpoint load, most obviously -- to the values the ++ # model was initialised with. For non-fp32 params the two maps share one ++ # tensor, so this loop only needs to cover the rest. ++ for param, cpu_copy in self.gpu_params_map_cpu_copy.items(): ++ if param not in self.param_to_fp32_param: ++ cpu_copy.data.copy_(param) + + def _register_load_state_dict_hooks(self): + def pre_load_state_dict_hook(self, state_dict): +diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py +index 9baed75..8f68e3d 100644 +--- a/megatron/core/optimizer/distrib_optimizer.py ++++ b/megatron/core/optimizer/distrib_optimizer.py +@@ -471,7 +471,9 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + + # fp32 params. + elif model_param.type() == 'torch.cuda.FloatTensor': +- shard_model_param = model_param.view(-1)[param_range.start : param_range.end] ++ shard_model_param = model_param.detach().view(-1)[ ++ param_range.start : param_range.end ++ ] + model_fp32_params_this_group.append(model_param) + shard_fp32_params_this_group.append(shard_model_param) + tensor_parallel.copy_tensor_model_parallel_attributes( +@@ -888,6 +890,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + # TE FusedAdam will not accumulate step for empty param groups, so we need to + # align the step across param groups. + param_group["step"] = int(step) ++ if "step" in param_group and param_group["step"] is None: ++ del param_group["step"] + + # Grad scaler state. + if self.grad_scaler: +@@ -1971,6 +1975,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer): + # separately via param_groups, not as part of the gradient buffer. + tensors[key] = LocalNonpersistentObject(tensors[key]) + continue ++ if key == 'step': ++ continue + assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), ( + tensors[key].shape, + gbuf_local_start, +diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py +index 7023488..25b07e2 100644 +--- a/megatron/core/parallel_state.py ++++ b/megatron/core/parallel_state.py +@@ -11,6 +11,7 @@ from typing import Callable, List, Optional + + import numpy as np + import torch ++import torch.distributed as dist + + from megatron.core.inference.symmetric_memory import SymmetricMemoryManager + +diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py +index 465e83f..232caef 100644 +--- a/megatron/core/pipeline_parallel/p2p_communication.py ++++ b/megatron/core/pipeline_parallel/p2p_communication.py +@@ -232,34 +232,26 @@ class P2PCommunicator: + group=self.pp_group, + ) + else: +- ops = [] +- if send_prev_shape_tensor is not None: +- send_prev_op = torch.distributed.P2POp( +- torch.distributed.isend, send_prev_shape_tensor, self.prev_rank, self.pp_group +- ) +- ops.append(send_prev_op) +- if recv_prev_shape_tensor is not None: +- recv_prev_op = torch.distributed.P2POp( +- torch.distributed.irecv, recv_prev_shape_tensor, self.prev_rank, self.pp_group +- ) +- ops.append(recv_prev_op) +- if send_next_shape_tensor is not None: +- send_next_op = torch.distributed.P2POp( +- torch.distributed.isend, send_next_shape_tensor, self.next_rank, self.pp_group +- ) +- ops.append(send_next_op) +- if recv_next_shape_tensor is not None: +- recv_next_op = torch.distributed.P2POp( +- torch.distributed.irecv, recv_next_shape_tensor, self.next_rank, self.pp_group +- ) +- ops.append(recv_next_op) +- if len(ops) > 0: +- reqs = torch.distributed.batch_isend_irecv(ops) +- for req in reqs: +- req.wait() ++ # PR #5271 (Megatron-LM): shape exchange MUST use _p2p_ops rather than ++ # batch_isend_irecv. batch_isend_irecv is one tagless NCCL group; when ++ # pp_group.size()==2, prev_rank == next_rank (single physical peer) and ++ # same-peer ops pair FIFO by enqueue order. Both ranks build ops in the ++ # same fixed order but hold opposite prev/next roles → recv_prev_shape ++ # and recv_next_shape get silently crossed. _p2p_ops handles size==2 ++ # via the group.WORLD split + even/odd ordering, correct for size>=4 too. ++ reqs = _p2p_ops( ++ tensor_send_prev=send_prev_shape_tensor, ++ tensor_recv_prev=recv_prev_shape_tensor, ++ tensor_send_next=send_next_shape_tensor, ++ tensor_recv_next=recv_next_shape_tensor, ++ group=self.pp_group, ++ prev_pipeline_rank=self.prev_rank, ++ next_pipeline_rank=self.next_rank, ++ ) ++ for req in reqs.values(): ++ req.wait() + +- # To protect against race condition when using batch_isend_irecv(). +- # should take this out once the bug with batch_isend_irecv is resolved. ++ # keep the CUDA sync as a defensive measure — cheap for 3-int64 tensors. + torch.cuda.synchronize() + + recv_prev_shape = [0, 0, 0] +@@ -371,6 +363,11 @@ class P2PCommunicator: + return [] + + p2p_func = _ring_exchange_wrapper ++ elif self.pp_group.size() == 2: ++ # PR #5271 (Megatron-LM): size==2 has same-peer (prev_rank == next_rank); ++ # batch_isend_irecv cannot pair same-peer bidirectional ops correctly (see ++ # #1450). Force _p2p_ops which handles this via WORLD split + even/odd order. ++ p2p_func = _p2p_ops + elif config.batch_p2p_comm: + assert wait_on_reqs + p2p_func = _batched_p2p_ops +@@ -381,10 +378,12 @@ class P2PCommunicator: + next_rank = self.next_rank + prev_rank = self.prev_rank + +- if config.use_ring_exchange_p2p or config.batch_p2p_comm: +- reqs = [] +- else: ++ # reqs init must match p2p_func return type: _p2p_ops returns dict, ++ # _batched_p2p_ops + _ring_exchange_wrapper return list. ++ if p2p_func is _p2p_ops: + reqs = {} ++ else: ++ reqs = [] + + tensor_recv_prev = None + tensor_recv_next = None +diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py +index 7e28691..3c19562 100644 +--- a/megatron/core/ssm/gated_delta_net.py ++++ b/megatron/core/ssm/gated_delta_net.py +@@ -584,16 +584,17 @@ class GatedDeltaNet(MegatronModule): + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) +- query, key, value, gate, beta, g = self.pre_gated_delta_rule( +- qkvzba, +- batch, +- seq_len_post_headwise, +- cp_size_headwise, +- cp_group_headwise, +- cu_seqlens_q, +- chunkwise_cp_context, +- packed_seq_params=packed_seq_params, +- ) ++ with torch._dynamo.config.patch(disable=True): ++ query, key, value, gate, beta, g = self.pre_gated_delta_rule( ++ qkvzba, ++ batch, ++ seq_len_post_headwise, ++ cp_size_headwise, ++ cp_group_headwise, ++ cu_seqlens_q, ++ chunkwise_cp_context, ++ packed_seq_params=packed_seq_params, ++ ) + nvtx_range_pop(suffix="pre_gated_delta_rule") + + nvtx_range_push(suffix="gated_delta_rule") +@@ -1211,7 +1212,7 @@ def get_parameter_local_cp_headwise( + slices = [slice(None)] * param.dim() + dim_size = param.size(dim=dim) + slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) +- param = param[slices] ++ param = param[tuple(slices)] + return param + + +diff --git a/megatron/core/tensor_parallel/random.py b/megatron/core/tensor_parallel/random.py +index dbecd73..1c52478 100644 +--- a/megatron/core/tensor_parallel/random.py ++++ b/megatron/core/tensor_parallel/random.py +@@ -589,7 +589,30 @@ class CheckpointFunction(torch.autograd.Function): + ) + + # Store everything. +- ctx.save_for_backward(*args) ++ # ++ # save_for_backward accepts tensors only, but not every checkpointed ++ # argument is one: gemma-4 uses dual RoPE, so TransformerBlock passes ++ # rotary_pos_emb as a TUPLE of two tensors, and this raised ++ # TypeError: save_for_backward can only save variables, ++ # but argument 4 is of type tuple ++ # Split the non-tensor arguments out here and rebuild the original ++ # ordering in backward(). ++ # ++ # Deliberately NOT calling _save_args_to_ctx() defined below, even ++ # though it solves the same problem for CheckpointWithoutOutputFunction: ++ # that helper detaches at save time, whereas this Function detaches in ++ # backward via detach_variable(). Splitting inline keeps the all-tensor ++ # path byte-identical to upstream, so no other model changes behaviour. ++ tensor_args = [] ++ non_tensor_entries = [] ++ for index, arg in enumerate(args): ++ if isinstance(arg, torch.Tensor): ++ tensor_args.append(arg) ++ else: ++ non_tensor_entries.append((index, arg)) ++ ctx.save_for_backward(*tensor_args) ++ ctx.non_tensor_entries = tuple(non_tensor_entries) ++ ctx.total_args_count = len(args) + + _unset_checkpointing() + return outputs +@@ -607,7 +630,17 @@ class CheckpointFunction(torch.autograd.Function): + ) + _set_checkpointing() + +- inputs = ctx.saved_tensors ++ # Rebuild the forward arguments in their original order, re-inserting the ++ # non-tensor ones that forward() kept out of save_for_backward. When every ++ # argument was a tensor this is exactly ctx.saved_tensors. ++ # detach_variable() (torch.utils.checkpoint) passes non-tensors through ++ # untouched, so the recompute below needs no further changes. ++ _saved = iter(ctx.saved_tensors) ++ _non_tensors = dict(ctx.non_tensor_entries) ++ inputs = tuple( ++ _non_tensors[i] if i in _non_tensors else next(_saved) ++ for i in range(ctx.total_args_count) ++ ) + if ctx.distribute_saved_activations: + safely_set_viewless_tensor_data( + inputs[0], gather_split_1d_tensor(inputs[0].data).view(ctx.input_0_shape) +@@ -630,7 +663,11 @@ class CheckpointFunction(torch.autograd.Function): + *filter(lambda x: torch.is_tensor(x[0]) and x[0].requires_grad, zip(outputs, args)) + ) + torch.autograd.backward(outputs, args) +- grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs) ++ # One entry per forward argument. autograd requires None -- not the object ++ # itself -- for arguments that are not tensors; returning `inp` here would ++ # raise "expected Variable or None". Unreachable before this commit, ++ # because a non-tensor argument died in save_for_backward first. ++ grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else None for inp in detached_inputs) + + _unset_checkpointing() + return (None, None) + grads +diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +index 3459a19..1a4f130 100644 +--- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py ++++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +@@ -276,4 +276,3 @@ def prepare_cp_compressor_input( +-@torch.compile + def _build_cp_indexer_layout( + cu_seqlens_q: torch.Tensor, + cu_seqlens_compressed: torch.Tensor, +diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py +index c8e197d..c5928f1 100644 +--- a/megatron/core/transformer/moe/moe_utils.py ++++ b/megatron/core/transformer/moe/moe_utils.py +@@ -797,6 +797,9 @@ def topk_routing_with_score_function( + scores, topk, num_groups, group_topk, _compute_topk + ) + ++ from relax.utils.training.routing_replay import get_routing_replay_compute_topk ++ compute_topk = get_routing_replay_compute_topk(compute_topk) ++ + # Precision notes: + # - Logits are converted to fp32 for score functions. + # - All the intermediate calculations are in fp32. +@@ -1356,7 +1359,11 @@ class RouterGatingLinearFunction(torch.autograd.Function): + inp_shape = inp.shape + inp = inp.view(-1, inp_shape[-1]) + +- if te_general_gemm is not None and router_dtype != torch.float64: ++ # TE multiplies in the operand dtype, so BF16 operands silently defeat ++ # --moe-router-dtype fp32. Keep FP32/FP64 routing on torch.mm. ++ if te_general_gemm is not None and router_dtype not in ( ++ torch.float32, torch.float64 ++ ): + output = te_general_gemm(weight, inp, router_dtype, layout="TN", bias=bias) + output = output[0] + elif bias is None: +@@ -1389,7 +1396,9 @@ class RouterGatingLinearFunction(torch.autograd.Function): + inp = inp.view(-1, inp_shape[-1]) + grad_output = grad_output.view(-1, grad_shape[-1]) + +- if te_general_gemm is not None and ctx.router_dtype != torch.float64: ++ if te_general_gemm is not None and ctx.router_dtype not in ( ++ torch.float32, torch.float64 ++ ): + grad_input = te_general_gemm( + weight.to(ctx.router_dtype), grad_output, ctx.router_dtype, layout="NN", grad=True + ) +diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py +index 2796bc6..71e036e 100644 +--- a/megatron/core/transformer/moe/router.py ++++ b/megatron/core/transformer/moe/router.py +@@ -267,6 +267,9 @@ class TopKRouter(Router): + if self.config.moe_enable_routing_replay: + self.router_replay = RouterReplay() + ++ from relax.utils.training.routing_replay import register_routing_replay ++ register_routing_replay(self) ++ + def _maintain_float32_expert_bias(self): + """ + Maintain the expert bias in float32. +diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py +index 43e45b7..8166f94 100644 +--- a/megatron/core/transformer/transformer_config.py ++++ b/megatron/core/transformer/transformer_config.py +@@ -224,6 +224,10 @@ class TransformerConfig(ModelParallelConfig): + """Clamp the output of the linear_fc1 in the activation function. Only used when activation_func + is quick_gelu or weighted SwiGLU (MoE only).""" + ++ activation_func_clamp_shared_expert: bool = True ++ """If False, skip activation_func_clamp_value inside SharedExpertMLP so only routed MoE ++ experts get the clamp.""" ++ + num_moe_experts: Optional[int] = None + """Number of experts to use for MoE layer. When set, it replaces MLP with MoE layer. Set to None + for no MoE.""" +@@ -283,6 +287,9 @@ class TransformerConfig(ModelParallelConfig): + num_query_groups >= tp, and under fp8/fp4 a per-partition + linear_qkv_out_dim aligned to 16/32.""" + ++ post_self_attn_layernorm: bool = False ++ post_mlp_layernorm: bool = False ++ + test_mode: bool = False + """Whether to run real-time tests.""" + +diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py +index 34a456b..8182c7c 100644 +--- a/megatron/core/transformer/transformer_layer.py ++++ b/megatron/core/transformer/transformer_layer.py +@@ -280,6 +280,7 @@ class TransformerLayerSubmodules: + self_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp + self_attention: Union[ModuleSpec, type] = IdentityOp + self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp ++ post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp + + pre_cross_attn_layernorm: LayerNormBuilder = IdentityOp + cross_attention_hyper_connection: Union[ModuleSpec, type] = IdentityOp +@@ -290,6 +291,7 @@ class TransformerLayerSubmodules: + mlp_hyper_connection: Union[ModuleSpec, type] = IdentityOp + mlp: MlpBuilder | type[IdentityOp] = IdentityOp + mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp ++ post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp + + # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method + sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict) +@@ -395,6 +397,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + # [Module 3: BiasDropoutFusion] + self.self_attn_bda = build_module(submodules.self_attn_bda) + ++ self.post_self_attn_layernorm = build_module( ++ submodules.post_self_attn_layernorm, ++ config=self.config, ++ hidden_size=self.config.hidden_size, ++ eps=self.config.layernorm_epsilon, ++ ) ++ + # [Module 4: Post SelfAttention] Optional Layernorm after self-attn + self.pre_cross_attn_layernorm = submodules.pre_cross_attn_layernorm( + config=self.config, +@@ -466,6 +475,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + + self.is_moe_layer = isinstance(self.mlp, MoELayer) + ++ self.post_mlp_layernorm = build_module( ++ submodules.post_mlp_layernorm, ++ config=self.config, ++ hidden_size=self.config.hidden_size, ++ eps=self.config.layernorm_epsilon ++ ) ++ + self.recompute_input_layernorm = False + self.recompute_pre_mlp_layernorm = False + self.recompute_mlp = False +@@ -780,6 +796,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + attention_output_with_bias[0] + ) + ++ attention_output, attention_output_bias = attention_output_with_bias ++ attention_output = self.post_self_attn_layernorm(attention_output) ++ attention_output_with_bias = (attention_output, attention_output_bias) ++ + # TODO: could we move `bias_dropout_add_exec_handler` itself + # inside the module provided in the `bias_dropout_add_spec` module? + nvtx_range_push(suffix="self_attn_bda") +@@ -1036,6 +1056,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer): + ) + mlp_output_with_bias = (mlp_output, mlp_bias) + ++ mlp_output, mlp_output_bias = mlp_output_with_bias ++ mlp_output = self.post_mlp_layernorm(mlp_output) ++ mlp_output_with_bias = (mlp_output, mlp_output_bias) ++ + nvtx_range_pop(suffix="mlp") + return mlp_output_with_bias, residual + +diff --git a/megatron/training/training.py b/megatron/training/training.py +index af94ff6..b309bfa 100644 +--- a/megatron/training/training.py ++++ b/megatron/training/training.py +@@ -221,7 +221,9 @@ except ImportError: + try: + from torch_memory_saver import torch_memory_saver + +- torch_memory_saver.hook_mode = "torch" ++ # NOTE(wuhuan): keep the default hook mode; forcing "torch" triggers ++ # 'torch.AcceleratorError: CUDA error: invalid argument' on weight updates. ++ # torch_memory_saver.hook_mode = "torch" + HAVE_TORCH_MEMORY_SAVER = True + except ImportError: + HAVE_TORCH_MEMORY_SAVER = False +diff --git a/megatron/core/transformer/moe/shared_experts.py b/megatron/core/transformer/moe/shared_experts.py +--- a/megatron/core/transformer/moe/shared_experts.py ++++ b/megatron/core/transformer/moe/shared_experts.py +@@ -123,6 +123,9 @@ class SharedExpertMLP(MLP): + assert config.add_bias_linear == False, "bias is not supported in the shared experts, " + "please set '--disable-bias-linear' instead." + ++ if not config.activation_func_clamp_shared_expert: ++ config.activation_func_clamp_value = None ++ + config.ffn_hidden_size = config.moe_shared_expert_intermediate_size + # TODO(Hepteract): pass pg_collection to MLP after refactoring MLP + super().__init__(config=config, submodules=submodules, tp_group=pg_collection.tp, name=name) +diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py +--- a/megatron/core/fp8_utils.py ++++ b/megatron/core/fp8_utils.py +@@ -339,2 +339,11 @@ def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, + ) ++ ++ recipe_configurator = getattr(quantizer_factory, "configure_custom_recipe", None) ++ if recipe_configurator is not None: ++ if not callable(recipe_configurator): ++ raise ValueError( ++ "fp8 quantizer factory configure_custom_recipe attribute must be callable." ++ ) ++ recipe_configurator(custom_recipe) ++ + return custom_recipe +diff --git a/megatron/core/safe_globals.py b/megatron/core/safe_globals.py +--- a/megatron/core/safe_globals.py ++++ b/megatron/core/safe_globals.py +@@ -102,13 +102,15 @@ class SafeUnpickler(pickle.Unpickler): + ("torch.storage", "_load_from_bytes"), + ("transformer_engine.common.recipe", "DelayedScaling"), + ("transformer_engine.common.recipe", "Float8CurrentScaling"), + ("transformer_engine.common.recipe", "Float8BlockScaling"), ++ ("transformer_engine.common.recipe", "CustomRecipe"), + ("transformer_engine.common.recipe", "MXFP8BlockScaling"), + ("transformer_engine.common.recipe", "NVFP4BlockScaling"), + ("transformer_engine.common.recipe", "Format"), + ("transformer_engine.common.recipe", "_FormatHelper"), + ("transformer_engine.common.recipe", "MMParams"), + ("transformer_engine.common.recipe", "QParams"), ++ ("relax.backends.megatron.fp8_recipes", "_DeepSeekV4SenderAlignedQuantizerFactory"), + ("megatron.core.extensions.transformer_engine", "TEDelayedScaling"), + ("megatron.core.safe_globals", "safe_load_from_bytes"), + ("numpy._core.multiarray", "_reconstruct"), +diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py +--- a/megatron/core/transformer/experimental_attention_variant/csa.py ++++ b/megatron/core/transformer/experimental_attention_variant/csa.py +@@ -1054,2 +1054,4 @@ class Compressor(MegatronModule): +- kv, _ = self.linear_wkv(x) # (total, 1, coff * head_dim) +- score, _ = self.linear_wgate(x) # (total, 1, coff * head_dim) ++ # Run the compressor GEMMs in high precision (BF16) even under FP8 training. ++ with get_fp8_disabled_context(self.config): ++ kv, _ = self.linear_wkv(x) # (total, 1, coff * head_dim) ++ score, _ = self.linear_wgate(x) # (total, 1, coff * head_dim) +diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +--- a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py ++++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py +@@ -261,6 +261,30 @@ + return global_idxs.int() + + ++def _compact_flat_topk_idxs(global_idxs: Tensor) -> Tuple[Tensor, Tensor]: ++ """Pack valid global indices into a per-row prefix. ++ ++ The returned ``topk_length`` selects that prefix in FlashMLA forward and ++ cuDNN DSA backward. Invalid suffix entries remain ``-1`` until forward has ++ consumed them; callers may then replace the ignored suffix with a safe ++ non-negative placeholder before backward. ++ """ ++ if global_idxs.ndim != 2: ++ raise ValueError(f"global_idxs must be 2-D (rows, topk), got {tuple(global_idxs.shape)}") ++ ++ if global_idxs.is_cuda: ++ _ensure_dsa_namespace() ++ res = _DSA.compactify_wrapper(global_idxs) ++ compact_idxs, topk_length = res["indices"], res["topk_length"] ++ else: ++ valid_mask = global_idxs >= 0 ++ sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True) ++ compact_idxs = global_idxs.gather(-1, sorted_indices) ++ topk_length = valid_mask.sum(dim=-1).int() ++ ++ return compact_idxs.int().contiguous(), topk_length.int().contiguous() ++ ++ + def build_flat_topk_idxs( + *idx_groups: Tensor, + batch_size: int, +@@ -309,21 +333,9 @@ + + topk_length_flat = None + if compact: +- if global_idxs.is_cuda: +- # Fast path: single warp-per-row CuTe DSL kernel from cuDNN's DSA +- # namespace. Replaces a stable argsort + gather + sum + permute +- # chain with one global-load + global-store per element. +- _ensure_dsa_namespace() +- res = _DSA.compactify_wrapper(global_idxs) +- global_idxs, topk_length_flat = res["indices"], res["topk_length"] +- else: +- # CPU fallback so the unit tests that exercise this helper without +- # CUDA still work. Production callers always go through the CUDA +- # path above. +- valid_mask = global_idxs >= 0 +- sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True) +- global_idxs = global_idxs.gather(-1, sorted_indices) +- topk_length_flat = valid_mask.sum(dim=-1).int() ++ # The CUDA path is a single warp-per-row CuTe DSL kernel; the helper ++ # retains a stable PyTorch fallback for CPU-only unit tests. ++ global_idxs, topk_length_flat = _compact_flat_topk_idxs(global_idxs) + + return global_idxs, topk_length_flat + +@@ -1136,6 +1148,14 @@ + combined_local = torch.cat([compress_topk_idxs, window_idxs], dim=-1) + global_idxs = local_to_global_flat(combined_local, b) + ++ # When the teacher/indexer loss is disabled, compact the complete ++ # attention set once and share its valid-prefix length between ++ # FlashMLA forward and cuDNN DSA backward. Keep the legacy segmented ++ # layout for loss_coeff > 0 until the full teacher-loss fix is ported. ++ topk_length = None ++ if loss_coeff == 0: ++ global_idxs, topk_length = _compact_flat_topk_idxs(global_idxs) ++ + # ---- 4. FlashMLA forward (flat layout for both SBHD and THD). -------- + if is_thd: + q_flat = query +@@ -1149,9 +1169,13 @@ + global_idxs, + softmax_scale, + attn_sink=attn_sink, +- topk_length=None, +- indexer_topk=indexer_topk, ++ topk_length=topk_length, ++ indexer_topk=0 if topk_length is not None else indexer_topk, + ) ++ if topk_length is not None: ++ # Forward has consumed the -1 sentinels. cuDNN DSA backward still ++ # reads ignored suffix slots, so replace them with a safe address. ++ global_idxs.clamp_min_(0) + + # ---- 4b. Derive padding-row mask for loss exclusion. ----------------- + # When CUDA-graph padding makes cu_seqlens_q cover all total_q rows +@@ -1173,6 +1197,52 @@ + real_len_per_row = real_seg_lens[row_batch_ids].to(torch.int32) + padding_row_mask = pos_in_seg >= real_len_per_row + ++ if topk_length is not None: ++ # cuDNN DSA backward requires at least one tile. The masked ++ # dO/LSE below makes this harmless placeholder gradient-free. ++ topk_length.masked_fill_(padding_row_mask, 1) ++ ++ if loss_coeff == 0: ++ # The old teacher path unconditionally consumes FlashMLA's partial ++ # indexer LSE, which is unavailable in compact mode. Since the ++ # configured coefficient is zero, skip only that inactive work; ++ # loss_coeff > 0 retains the legacy teacher behavior unchanged. ++ indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) ++ precomputed_grad_q_indexer = torch.zeros_like(q_indexer) ++ precomputed_grad_k_indexer = torch.zeros_like(k_indexer) ++ precomputed_grad_weights = torch.zeros_like(weights) ++ ctx.save_for_backward( ++ q_flat, ++ kv_flat, ++ attn_sink, ++ global_idxs, ++ topk_length, ++ out_flat, ++ lse, ++ precomputed_grad_q_indexer, ++ precomputed_grad_k_indexer, ++ precomputed_grad_weights, ++ ) ++ ctx.softmax_scale = softmax_scale ++ ctx.is_thd = is_thd ++ ctx.has_topk_length = True ++ ctx.padding_row_mask = padding_row_mask ++ ctx.np_ = np_ ++ ctx.d = d ++ if is_thd: ++ ctx.total_q = total_q ++ else: ++ ctx.sq = sq ++ ctx.b = b ++ ctx.skv = skv ++ ++ d_v = out_flat.shape[-1] ++ if is_thd: ++ output = out_flat.reshape(total_q, np_ * d_v) ++ else: ++ output = out_flat.reshape(sq, b, np_, d_v).reshape(sq, b, np_ * d_v) ++ return output, indexer_loss ++ + # ---- 5. Derive predict from indexer_scores, compute target. ---------- + # Layout-specific attn tensors (detached — loss is not differentiable + # through them). +@@ -1383,11 +1453,13 @@ + precomputed_grad_weights[padding_row_mask] = 0 + + # ---- 7. Save context (only sparse-attn bwd tensors + indexer grads). - ++ empty_topk_length = torch.empty(0, dtype=torch.int32, device=global_idxs.device) + ctx.save_for_backward( + q_flat, + kv_flat, + attn_sink, + global_idxs, ++ empty_topk_length, + out_flat, + lse, + precomputed_grad_q_indexer, +@@ -1396,6 +1468,8 @@ + ) + ctx.softmax_scale = softmax_scale + ctx.is_thd = is_thd ++ ctx.has_topk_length = False ++ ctx.padding_row_mask = padding_row_mask + ctx.np_ = np_ + ctx.d = d + if is_thd: +@@ -1421,12 +1495,14 @@ + kv_flat, + attn_sink, + global_idxs, ++ saved_topk_length, + out_flat, + lse, + precomputed_grad_q_indexer, + precomputed_grad_k_indexer, + precomputed_grad_weights, + ) = ctx.saved_tensors ++ topk_length = saved_topk_length if ctx.has_topk_length else None + + is_thd = ctx.is_thd + np_, d = ctx.np_, ctx.d +@@ -1439,6 +1515,10 @@ + sq, b, skv = ctx.sq, ctx.b, ctx.skv + dO_flat = grad_output.reshape(sq * b, np_, d_v) + ++ if topk_length is not None and ctx.padding_row_mask is not None: ++ dO_flat = dO_flat.masked_fill(ctx.padding_row_mask[:, None, None], 0) ++ lse = lse.masked_fill(ctx.padding_row_mask[:, None], 0) ++ + attn_bwd = _DSA.sparse_attention_backward_wrapper( + q_flat, + kv_flat, +@@ -1448,7 +1528,7 @@ + attn_sink, + global_idxs, + softmax_scale=ctx.softmax_scale, +- topk_length=None, ++ topk_length=topk_length, + ) + if is_thd: + grad_query = attn_bwd["dq"] +@@ -1534,15 +1614,51 @@ + total_comp = k_indexer.shape[0] + indexer_topk = indexer_topk_idxs.shape[-1] + ++ # The coeff=0 training path does not consume the teacher's segmented ++ # indexer LSE. Compact the complete attention set and reuse the exact ++ # same valid-prefix lengths in forward and backward. Preserve the ++ # legacy non-compact teacher path when loss_coeff > 0. ++ topk_length = None ++ if loss_coeff == 0: ++ topk_idxs, topk_length = _compact_flat_topk_idxs(topk_idxs) ++ + out_flat, lse, lse_indexer = _dsa_fwd_flash_mla( + query, + kv_full, + topk_idxs, + softmax_scale, + attn_sink=attn_sink, +- topk_length=None, +- indexer_topk=indexer_topk, ++ topk_length=topk_length, ++ indexer_topk=0 if topk_length is not None else indexer_topk, + ) ++ if topk_length is not None: ++ topk_idxs.clamp_min_(0) ++ if q_padding_mask is not None: ++ # Avoid cuDNN DSA's zero-tile backward path; dO/LSE for these ++ # rows are masked before the harmless placeholder is used. ++ topk_length.masked_fill_(q_padding_mask, 1) ++ ++ if loss_coeff == 0: ++ indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) ++ saved_grad_q_indexer = torch.zeros_like(q_indexer) ++ saved_grad_k_indexer = torch.zeros_like(k_indexer) ++ saved_grad_weights = torch.zeros_like(weights) ++ ctx.save_for_backward( ++ query, ++ kv_full, ++ attn_sink, ++ topk_idxs, ++ topk_length, ++ out_flat, ++ lse, ++ saved_grad_q_indexer, ++ saved_grad_k_indexer, ++ saved_grad_weights, ++ ) ++ ctx.softmax_scale = softmax_scale ++ ctx.has_topk_length = True ++ ctx.q_padding_mask = q_padding_mask ++ return out_flat.reshape(total_q, np_ * out_flat.shape[-1]), indexer_loss + + bwd_loss_coeff = loss_coeff * total_q / loss_divisor + unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) +@@ -1674,11 +1790,13 @@ + saved_grad_k_indexer = torch.zeros_like(k_indexer) + saved_grad_weights = torch.zeros_like(weights) + ++ empty_topk_length = torch.empty(0, dtype=torch.int32, device=topk_idxs.device) + ctx.save_for_backward( + query, + kv_full, + attn_sink, + topk_idxs, ++ empty_topk_length, + out_flat, + lse, + saved_grad_q_indexer, +@@ -1686,6 +1804,8 @@ + saved_grad_weights, + ) + ctx.softmax_scale = softmax_scale ++ ctx.has_topk_length = False ++ ctx.q_padding_mask = q_padding_mask + + return out_flat.reshape(total_q, np_ * out_flat.shape[-1]), indexer_loss + +@@ -1698,14 +1818,19 @@ + kv_full, + attn_sink, + topk_idxs, ++ saved_topk_length, + out_flat, + lse, + saved_grad_q_indexer, + saved_grad_k_indexer, + saved_grad_weights, + ) = ctx.saved_tensors ++ topk_length = saved_topk_length if ctx.has_topk_length else None + + dO_flat = grad_output.reshape(query.shape[0], query.shape[1], out_flat.shape[-1]) ++ if topk_length is not None and ctx.q_padding_mask is not None: ++ dO_flat = dO_flat.masked_fill(ctx.q_padding_mask[:, None, None], 0) ++ lse = lse.masked_fill(ctx.q_padding_mask[:, None], 0) + attn_bwd = _DSA.sparse_attention_backward_wrapper( + query, + kv_full, +@@ -1715,7 +1840,7 @@ + attn_sink, + topk_idxs, + softmax_scale=ctx.softmax_scale, +- topk_length=None, ++ topk_length=topk_length, + ) + return ( + attn_bwd["dq"], +diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py +--- a/megatron/core/transformer/experimental_attention_variant/csa.py ++++ b/megatron/core/transformer/experimental_attention_variant/csa.py +@@ -2543,7 +2543,8 @@ + self.config, + ) + q_indexer_cp = rotate_activation(q_indexer_cp) +- weights_indexer_cp, _ = indexer.linear_weights_proj(indexer_x) ++ with get_fp8_disabled_context(indexer.config): ++ weights_indexer_cp, _ = indexer.linear_weights_proj(indexer_x) + weights_indexer_cp = weights_indexer_cp.squeeze(1) * (indexer.index_n_heads**-0.5) + + indexer_compressed_local, _ = indexer.compressor._forward_thd( From c113c72a3374b095b6d6908dff4169d9f8d9b792 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=BE=90=E4=B8=9C=E8=BE=89?= Date: Tue, 1 Sep 2026 17:31:06 +0800 Subject: [PATCH 05/34] fix(ray): fence graceful elastic scale-in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Serialize elastic removal with weight updates - Claim DRAINING before waiting for active fully-async or scale-out weight transfers - Reject new fully-async handshakes while graceful removal owns the topology fence - Persist SIGTERM eviction intent across explicit scale-in timeout and cleanup failures ## Batch simultaneous graceful evictions - Collect ready logical engines from one poll and claim the full batch before waiting - Remove each server batch with concurrent Router, DCS, and shutdown phases plus one drain wait - Keep failed engines fenced without blocking successful peers from cleanup ## Coordinate explicit scaling and graceful eviction - Atomically recheck and insert scale requests against lifecycle claims - Let target scale-in adopt matching pending evictions without duplicate cleanup - Retry unmatched URL-target evictions after the explicit request reaches a terminal state - Recover ACTIVE scaled groups while leaving DRAINING and REMOVING groups untouched ## Keep graceful eviction on the live-actor path - Make the SGLang SIGTERM handler publish eviction intent without blocking or performing I/O - Reuse Router, drain, DCS, actor shutdown, and placement-group cleanup for scale-in and eviction - Poll eviction probes per reference and keep elastic ranks monotonic across scale cycles --- # ✅ Tests ## Cover lifecycle and batch coordination - Verify batch claim-before-wait, single drain, concurrent removal phases, and isolated failures - Verify persistent intent, request ordering, target adoption, handshake rollback, and recovery filtering - Verify DCS failure stays out of the hard-kill path and SIGTERM only publishes intent (cherry picked from commit 61e2cddc4f5fee2c8ffaaa8fcca57b28f0d77e6e) --- relax/backends/megatron/actor.py | 23 +- relax/backends/sglang/sglang_engine.py | 58 +- relax/components/rollout.py | 28 +- relax/distributed/ray/rollout.py | 807 ++++++++++++------ .../megatron/test_actor_http_timeout.py | 55 +- .../backends/sglang/test_sigterm_eviction.py | 30 + .../test_rollout_weight_update_handshake.py | 65 +- tests/distributed/ray/conftest.py | 13 +- tests/distributed/ray/test_coordination.py | 322 ++++++- tests/distributed/ray/test_scale_in.py | 363 +++++++- tests/distributed/ray/test_scale_out.py | 11 + tests/distributed/ray/test_utils.py | 26 + tests/distributed/ray/test_weight_sync.py | 23 + 13 files changed, 1479 insertions(+), 345 deletions(-) create mode 100644 tests/backends/sglang/test_sigterm_eviction.py diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 4a7364031..0e93d91fa 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -1835,22 +1835,39 @@ def _check_services_health(self) -> tuple[bool, bool]: # Default: both services healthy → update both rollout_only = False actor_fwd_only = False + process_group = get_gloo_group() # When true_on_policy_mode is enabled, actor_fwd is intentionally absent # (its log_probs are recomputed inline by the train forward). Force # rollout-only weight update and skip the actor_fwd HTTP probe. actor_fwd_absent = getattr(self.args, "true_on_policy_mode", False) - if dist.get_rank() == 0: + if dist.get_rank(process_group) == 0: # Check rollout service try: rollout_serve_url = get_serve_url("rollout") + retry_deadline = None while True: + request_timeout = self.args.rollout_http_timeout + if retry_deadline is not None: + request_timeout = retry_deadline - time.monotonic() + if request_timeout <= 0: + raise requests.exceptions.Timeout("elastic scale-in fence did not clear") response = requests.get( f"{rollout_serve_url}/can_do_update_weight_for_async", - timeout=self.args.rollout_http_timeout, + timeout=request_timeout, ) + if getattr(response, "status_code", 200) == 503: + if retry_deadline is None: + retry_deadline = time.monotonic() + max(float(self.args.rollout_http_timeout), 1.0) + logger.warning("Elastic scale-in is draining; retrying before weight update.") + time.sleep(min(1.0, max(0.0, retry_deadline - time.monotonic()))) + continue response.raise_for_status() + # A successful non-503 response means the scale-in fence + # has cleared. Do not let an earlier draining deadline + # bound the normal readiness polling below. + retry_deadline = None res = response.json() if res: response = requests.get(f"{rollout_serve_url}/recover_rollout_engines") @@ -1890,7 +1907,7 @@ def _check_services_health(self) -> tuple[bool, bool]: dtype=torch.int32, device="cpu", ) - dist.all_reduce(flags, op=dist.ReduceOp.MAX, group=get_gloo_group()) + dist.all_reduce(flags, op=dist.ReduceOp.MAX, group=process_group) rollout_only = bool(flags[0].item()) actor_fwd_only = bool(flags[1].item()) diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py index f8e36d861..ed4dcf950 100644 --- a/relax/backends/sglang/sglang_engine.py +++ b/relax/backends/sglang/sglang_engine.py @@ -353,73 +353,23 @@ def __init__( self.sglang_overrides = sglang_overrides or {} self.num_gpus_per_engine = num_gpus_per_engine self._evicted = threading.Event() - self._is_weight_updating: bool = False self._router_worker_id: str | None = None self._router_unregister_submitted = False if register_sigterm_handler: self._register_sigterm_handler() - def set_weight_updating(self, is_updating: bool) -> None: - """Set whether a weight update is currently in progress. - - Called by RolloutManager before and after each weight sync so that the - SIGTERM handler can wait for the update to finish before unregistering - from the router. - """ - self._is_weight_updating = is_updating - def _register_sigterm_handler(self): """Register SIGTERM handler for platform-initiated pod eviction. - When the platform needs to evict or replace a pod, it sends SIGTERM to the - user process. We catch this signal and perform lightweight cleanup so the - RolloutManager can detect the eviction and treat it as a scale-in event. - - If a weight update is in progress (``_is_weight_updating``), the handler - blocks until the update finishes before unregistering from the router, to - avoid disrupting NCCL communication groups. The k8s PreStop timeout - (default 30s) serves as the hard deadline. + The signal handler only publishes an intent. RolloutManager owns the + weight-update fence and the live-actor removal sequence; doing I/O or + waiting here could block the actor RPC that releases an active update. """ self._original_sigterm_handler = signal.getsignal(signal.SIGTERM) - def _handle_sigterm(signum, frame): - actor_id = "" - try: - actor_id = ray.get_runtime_context().get_actor_id() - except Exception: - pass - logger.warning( - f"[SGLangEngine] Received SIGTERM (rank={self.rank}, actor_id={actor_id}), " - f"marking as evicted for graceful scale-in" - ) + def _handle_sigterm(_signum, _frame): self._evicted.set() - # Wait for any in-progress weight update to finish before cleaning up, - # so we don't disrupt NCCL communication groups during weight sync. - # k8s PreStop hard deadline is typically 30s; leave a safety margin. - weight_update_timeout = 20 - wait_start = time.time() - while self._is_weight_updating: - if time.time() - wait_start > weight_update_timeout: - logger.warning( - f"[SGLangEngine] Weight update did not finish within {weight_update_timeout}s, " - f"proceeding with eviction cleanup (rank={self.rank})" - ) - break - logger.warning( - f"[SGLangEngine] SIGTERM received but weight update in progress " - f"(rank={self.rank}, actor_id={actor_id}), waiting..." - ) - time.sleep(1) - - # Best-effort: unregister from router so new requests are not routed here. - # This is a quick HTTP call; if it fails the router will detect the engine - # as unhealthy anyway. - try: - self.unregister_from_router() - except Exception as e: - logger.warning(f"[SGLangEngine] Failed to unregister from router during SIGTERM handling: {e}") - signal.signal(signal.SIGTERM, _handle_sigterm) def is_evicted(self) -> bool: diff --git a/relax/components/rollout.py b/relax/components/rollout.py index f13983f35..18f84ca7f 100644 --- a/relax/components/rollout.py +++ b/relax/components/rollout.py @@ -591,16 +591,30 @@ async def can_do_update_weight_for_async(self): can_update = await self._async_check_production_for_update_weight(step) if can_update: self._weight_update_ready.clear() + self.status = "paused" try: - self.status = "paused" + prepared = await self.rollout_manager.set_weight_updating.remote(True) + if prepared is False: + raise HTTPException(status_code=503, detail="Elastic scale-in is draining") await self.rollout_manager.health_monitoring_pause.remote() - await self.rollout_manager.set_weight_updating.remote(True) + except Exception: + rollback_succeeded = False + try: + await self.rollout_manager.set_weight_updating.remote(False) + rollback_succeeded = True + except Exception as rollback_error: + self._logger.warning( + "Failed to roll back a partial weight-update handshake: %s", + rollback_error, + ) + if rollback_succeeded: + self.status = "running" + raise finally: - # Always release the handshake gate: even if the remote calls - # above raise (e.g. RayActorError from a dead engine), - # end_update_weight must not block forever. The 500 still - # propagates to the actor, which already has graceful - # degradation (actor_fwd_only). _weight_update_ready only orders + # Always release the handshake gate if a remote call above + # raises, so a later end_update_weight cannot block forever. + # The error still propagates to the actor, which fails closed. + # _weight_update_ready only orders # the can_do <-> end_update_weight handshake; it does not gate # the real weight transfer, so setting it on the failure path # cannot make an engine use wrong weights. diff --git a/relax/distributed/ray/rollout.py b/relax/distributed/ray/rollout.py index e516e94e5..0886413d1 100644 --- a/relax/distributed/ray/rollout.py +++ b/relax/distributed/ray/rollout.py @@ -406,6 +406,13 @@ def to_dict(self) -> dict: } +class EngineGroupLifecycle(str, enum.Enum): + ACTIVE = "ACTIVE" + DRAINING = "DRAINING" + REMOVING = "REMOVING" + REMOVED = "REMOVED" + + @dataclasses.dataclass class EngineGroup: """A group of homogeneous SGLang engines with the same configuration. @@ -429,6 +436,8 @@ class EngineGroup: is_scaled_out: bool = False # True for groups added via scale-out, False for initial groups skip_dcs_registration: bool = False # Skip DCS registration for scaled-out engines skip_router_registration: bool = False # Skip router registration until weight sync completes + lifecycle_status: EngineGroupLifecycle = EngineGroupLifecycle.ACTIVE + eviction_requested: bool = False @property def nodes_per_engine(self): @@ -741,13 +750,21 @@ def nodes_per_engine(self): def recover(self): """Recover dead engines across all active groups, overlapping init.""" - dead_per_group = [[i for i, engine in enumerate(g.all_engines) if engine is None] for g in self.engine_groups] + groups = list(self.engine_groups) + dead_per_group = [[i for i, engine in enumerate(g.all_engines) if engine is None] for g in groups] all_handles = [] port_cursors: dict[int, int] = {} groups_to_remove = [] - for g_idx, g in enumerate(self.engine_groups): + for g_idx, g in enumerate(groups): + if g.is_scaled_out and g.pg is not None and g.lifecycle_status is not EngineGroupLifecycle.ACTIVE: + if any(engine is None for engine in g.all_engines): + logger.warning( + "Skipping recovery for non-active scaled engine group at rank offset %s", + g.rank_offset, + ) + continue if g.pg is None: failed_indices = g.healthcheck_engines() if failed_indices: @@ -764,7 +781,9 @@ def recover(self): all_handles.extend(handles) for g_idx in reversed(groups_to_remove): - self.engine_groups.pop(g_idx) + group = groups.pop(g_idx) + if group in self.engine_groups: + self.engine_groups.remove(group) dead_per_group.pop(g_idx) logger.info(f"Removed dead external engine group {g_idx}") @@ -773,8 +792,8 @@ def recover(self): release_handles = [] new_engines_all = [] - for g, dead_indices in zip(self.engine_groups, dead_per_group, strict=True): - if g.pg is None: + for g, dead_indices in zip(groups, dead_per_group, strict=True): + if g.pg is None or (g.is_scaled_out and g.lifecycle_status is not EngineGroupLifecycle.ACTIVE): continue logger.info(f"Recovered {g.num_new_engines} dead rollout engines (worker_type={g.worker_type})") assert g.num_new_engines == len(dead_indices), "num_new_engines does not match dead_indices length" @@ -894,6 +913,12 @@ def __init__(self, args, pg, data_source=None): self.rollout_id = -1 self._metric_checker = MetricChecker.maybe_create(args) self._tokenizer = None # Lazy-initialized tokenizer for debug data saving + self._engine_lifecycle_lock = threading.RLock() + existing_groups = [group for srv in self.servers.values() for group in srv.engine_groups] + self._next_engine_rank = max( + (group.rank_offset + len(group.all_engines) for group in existing_groups), + default=0, + ) self._health_monitors = [] if not self.args.debug_train_only and self.args.use_fault_tolerance: @@ -924,7 +949,8 @@ def __init__(self, args, pg, data_source=None): # Elastic scale-in tracking self._scale_in_requests: dict[str, ScaleInRequest] = {} - self._is_weight_updating: bool = False + self._training_weight_updating: bool = False + self._scale_out_weight_updating: bool = False # Distributed mutex shared with the Actor process to ensure DCS weight # sync (update_weights_fully_async) and sglang remote instance weight sync # (_sync_weights_from_seed_engine) never run concurrently. @@ -943,6 +969,19 @@ def __init__(self, args, pg, data_source=None): if not self.args.debug_train_only: self._start_eviction_monitor() + def _reserve_engine_ranks(self, count: int, alignment: int = 1) -> int: + """Allocate ranks monotonically so a removed elastic rank is not + reused.""" + with self._engine_lifecycle_lock: + rank_offset = ((self._next_engine_rank + alignment - 1) // alignment) * alignment + self._next_engine_rank = rank_offset + count + return rank_offset + + @property + def _is_weight_updating(self) -> bool: + with self._engine_lifecycle_lock: + return self._training_weight_updating or self._scale_out_weight_updating + def _try_ci_fault_injection(self): """Try to inject fault during generate (when health monitor is running).""" @@ -1436,17 +1475,32 @@ def _collect_in_flight_engine_addrs(self, model_name: str) -> set[str]: def _find_active_scale_request(self) -> Optional[dict]: """Return info about any active (non-terminal) scale-out or scale-in - request. + request or claimed graceful eviction. Returns None if no active request exists, otherwise a dict with ``type``, ``request_id``, and ``status`` of the blocking request. """ - for r in self._scale_out_requests.values(): - if not r.is_terminal(): - return {"type": "scale_out", "request_id": r.request_id, "status": r.status.value} - for r in self._scale_in_requests.values(): - if not r.is_terminal(): - return {"type": "scale_in", "request_id": r.request_id, "status": r.status.value} + with self._engine_lifecycle_lock: + # Explicit requests take precedence because an active scale-in may + # itself have already moved a group into DRAINING. + for r in self._scale_out_requests.values(): + if not r.is_terminal(): + return {"type": "scale_out", "request_id": r.request_id, "status": r.status.value} + for r in self._scale_in_requests.values(): + if not r.is_terminal(): + return {"type": "scale_in", "request_id": r.request_id, "status": r.status.value} + for model_name, srv in self.servers.items(): + for group in srv.engine_groups: + if ( + group.is_scaled_out + and group.eviction_requested + and group.lifecycle_status is EngineGroupLifecycle.DRAINING + ): + return { + "type": "graceful_eviction", + "request_id": f"{model_name}:{group.rank_offset}", + "status": group.lifecycle_status.value, + } return None @ray.method(concurrency_group="scale_coordination") @@ -1589,8 +1643,22 @@ def create_scale_out_request( timeout_secs=timeout_secs or self.args.scale_out_timeout, ) - self._scale_out_requests[request.request_id] = request - self._gc_terminal_requests() + # Recheck and insert under the same lifecycle lock used by graceful + # eviction claims. Validation above may be slow enough for eviction + # to win after the initial fast-path check. + with self._engine_lifecycle_lock: + active = self._find_active_scale_request() + if active is not None: + return { + "request_id": str(uuid.uuid4()), + "status": "CONFLICT", + "message": ( + f"Another {active['type']} request is in progress: " + f"request_id={active['request_id']}, status={active['status']}" + ), + } + self._scale_out_requests[request.request_id] = request + self._gc_terminal_requests() return request.to_dict() @ray.method(concurrency_group="scale_out") @@ -1685,6 +1753,7 @@ async def _scale_out_ray_native(self, request: ScaleOutRequest) -> None: return gpus_per_engine = self.args.rollout_num_gpus_per_engine + actors_per_replica = max(1, gpus_per_engine // self.args.num_gpus_per_node) # Step 2: Create one PG per replica so that replicas with available # resources can proceed immediately without waiting for the others. @@ -1719,9 +1788,7 @@ async def _scale_out_ray_native(self, request: ScaleOutRequest) -> None: # generic "All N replicas failed". failure_reasons: list[ScaleOutFailure] = [] - # Track the running engine offset (may change as replicas succeed) - base_engine_offset = sum(len(g.all_engines) for g in srv.engine_groups) - logger.info(f"[ScaleOut] Current total engines (base offset): {base_engine_offset}") + logger.info(f"[ScaleOut] Next monotonic engine rank: {self._next_engine_rank}") # Phase A: Wait for PGs to become ready, with incremental processing while pending_indices: @@ -1776,8 +1843,13 @@ async def _scale_out_ray_native(self, request: ScaleOutRequest) -> None: # Pre-allocate engine offsets so each coroutine gets a unique rank_offset # without racing on srv.engine_groups mutations. if newly_ready: - current_offset = sum(len(g.all_engines) for g in srv.engine_groups) - replica_offsets = {idx: current_offset + i for i, idx in enumerate(newly_ready)} + replica_offsets = { + idx: self._reserve_engine_ranks( + actors_per_replica, + alignment=actors_per_replica, + ) + for idx in newly_ready + } async def _bring_up_one(idx: int) -> tuple[int, bool, "ScaleOutFailure | None", int]: r = await self._bring_up_single_replica( @@ -2835,18 +2907,43 @@ def _record(failure: ScaleOutFailure) -> None: _record(ScaleOutFailure(ScaleOutFailureCategory.WEIGHT_SYNC_FAILED, "no healthy seed engine")) return False - # Acquire the distributed lock to prevent concurrent DCS weight sync - # (update_weights_fully_async on the Actor side) from overlapping with - # this remote instance weight sync. Both use the seed engine's NCCL stack. - acquired = False - while not acquired: - acquired = await asyncio.to_thread(ray.get, self._weight_sync_lock.acquire.remote()) - if not acquired: - await asyncio.sleep(0.5) - - self._is_weight_updating = True sync_succeeded = False + scale_out_owner_acquired = False + weight_sync_lock_acquired = False try: + # Claim the lifecycle owner before taking the distributed lock. + # Otherwise scale-out can hold that lock while waiting for a + # DRAINING group whose eviction is waiting for a training update + # blocked on the same lock. + owner_deadline = time.monotonic() + timeout + while True: + with self._engine_lifecycle_lock: + has_draining_elastic_group = any( + group.is_scaled_out and group.lifecycle_status is EngineGroupLifecycle.DRAINING + for srv in self.servers.values() + for group in srv.engine_groups + ) + if not has_draining_elastic_group: + self._scale_out_weight_updating = True + scale_out_owner_acquired = True + break + if time.monotonic() >= owner_deadline: + _record( + ScaleOutFailure( + ScaleOutFailureCategory.WEIGHT_SYNC_FAILED, + "timed out waiting for elastic removal", + ) + ) + return False + await asyncio.sleep(0.5) + + # Prevent Actor-side DCS weight sync from overlapping with this + # remote instance sync. Both use the seed engine's NCCL stack. + while not weight_sync_lock_acquired: + weight_sync_lock_acquired = await asyncio.to_thread(ray.get, self._weight_sync_lock.acquire.remote()) + if not weight_sync_lock_acquired: + await asyncio.sleep(0.5) + # Pause generation on all new engines before weight sync # This ensures no pending requests during flush_cache logger.info("[ScaleOut][WeightSync] Pausing generation on new engines...") @@ -3029,11 +3126,13 @@ async def _sync_one(idx, engine, _seed=seed_engine, _addr=master_address, _rem=r logger.warning(f"[ScaleOut][WeightSync] Some continue_generation calls failed: {e}") else: logger.warning("[ScaleOut][WeightSync] Keeping failed new engines paused for rollback") - self._is_weight_updating = False - # Always release the distributed weight-sync lock. This is the SAME - # lock every training weight update acquires; retaining it as an - # isolation mechanism would spin the next training update forever. - ray.get(self._weight_sync_lock.release.remote()) + if scale_out_owner_acquired: + with self._engine_lifecycle_lock: + self._scale_out_weight_updating = False + if weight_sync_lock_acquired: + # This is the SAME lock every training weight update acquires; + # retaining it would spin the next training update forever. + ray.get(self._weight_sync_lock.release.remote()) async def _health_check_engines(self, engines: list, timeout: float = 60.0) -> bool: """Check health of engines. @@ -3392,19 +3491,22 @@ async def cancel_all_scale_out_requests( return result @ray.method(concurrency_group="scale_in") - def set_weight_updating(self, is_updating: bool) -> None: - self._is_weight_updating = is_updating + def set_weight_updating(self, is_updating: bool) -> bool: + """Acquire or release the fully-async topology lease. - # Mirror the flag to every live engine so their SIGTERM handlers can - # check it locally without an extra Ray RPC. - refs = [] - for srv in self.servers.values(): - for group in srv.engine_groups: - for engine in group.all_engines: - if engine is not None: - refs.append(engine.set_weight_updating.remote(is_updating)) - if refs: - ray.get(refs) + A DRAINING group owns the scale-in fence, so a new update must retry. + An update admitted first may finish; scale-in waits for this owner to + release instead of mutating the live NCCL topology underneath it. + """ + with self._engine_lifecycle_lock: + if is_updating and any( + group.is_scaled_out and group.lifecycle_status is EngineGroupLifecycle.DRAINING + for srv in self.servers.values() + for group in srv.engine_groups + ): + return False + self._training_weight_updating = is_updating + return True @ray.method(concurrency_group="scale_in") async def sync_weights_for_scaled_out_engines( @@ -3567,8 +3669,20 @@ def create_scale_in_request( force=force, dry_run=dry_run, ) - self._scale_in_requests[request.request_id] = request - self._gc_terminal_requests() + # Recheck and insert atomically against graceful eviction claims. + with self._engine_lifecycle_lock: + active = self._find_active_scale_request() + if active is not None: + return { + "request_id": str(uuid.uuid4()), + "status": "CONFLICT", + "message": ( + f"Another {active['type']} request is in progress: " + f"request_id={active['request_id']}, status={active['status']}" + ), + } + self._scale_in_requests[request.request_id] = request + self._gc_terminal_requests() return request.to_dict() @ray.method(concurrency_group="scale_in") @@ -3594,68 +3708,85 @@ async def _scale_in(self, request: ScaleInRequest) -> None: request.update_status(ScaleInStatus.FAILED, f"Model '{request.model_name}' not found (no rollout server)") return - # P1-3: Wait for any in-progress weight update to complete before draining. - # Draining engines during a weight update could break NCCL communication groups. - if self._is_weight_updating: - logger.info("[ScaleIn] Weight update in progress, waiting for it to complete...") - wait_start = time.time() - weight_update_timeout = request.timeout_secs - while self._is_weight_updating and (time.time() - wait_start) < weight_update_timeout: - await asyncio.sleep(1) - if self._is_weight_updating: - logger.warning( - f"[ScaleIn] Weight update still in progress after {weight_update_timeout}s, proceeding anyway" - ) - else: - logger.info(f"[ScaleIn] Weight update completed after {time.time() - wait_start:.1f}s, proceeding") - + selected_groups: dict[int, EngineGroup] = {} + completed_eviction_groups: set[int] = set() try: - engine_infos = self._select_engines_for_removal(request, srv) - if not engine_infos: - if request.num_replicas > 0: - request.update_status(ScaleInStatus.COMPLETED) - logger.info( - f"[ScaleIn] No-op: already at or below target replicas (target={request.num_replicas})" - ) + url_candidates = None + if request.engine_urls: + url_candidates = await self._resolve_scale_in_url_candidates(request, srv) + + with self._engine_lifecycle_lock: + # Selection and lifecycle claim are one transaction with the + # SIGTERM handler. Whichever enters first determines whether + # a pending eviction is adopted into this target scale-in. + if request.selected_engines: + logger.warning("[ScaleIn] Request %s already has claimed targets", request.request_id) + return + engine_infos = self._select_engines_for_removal(request, srv, url_candidates=url_candidates) + if not engine_infos: + if request.num_replicas > 0: + request.update_status(ScaleInStatus.COMPLETED) + logger.info( + f"[ScaleIn] No-op: already at or below target replicas (target={request.num_replicas})" + ) + return + request.update_status(ScaleInStatus.FAILED, "No engines selected for removal") return - request.update_status(ScaleInStatus.FAILED, "No engines selected for removal") - return - request.selected_engines = [f"group_{g.rank_offset}_engine_{node0_idx}" for g, node0_idx in engine_infos] - logger.info(f"[ScaleIn] Selected {len(engine_infos)} engines for removal: {request.selected_engines}") + request.selected_engines = [ + f"group_{group.rank_offset}_engine_{node0_idx}" for group, node0_idx in engine_infos + ] + logger.info(f"[ScaleIn] Selected {len(engine_infos)} engines for removal: {request.selected_engines}") - if request.dry_run: - request.update_status(ScaleInStatus.COMPLETED) - logger.info("[ScaleIn] Dry-run complete, no engines removed") - return + if request.dry_run: + request.update_status(ScaleInStatus.COMPLETED) + logger.info("[ScaleIn] Dry-run complete, no engines removed") + return + + selected_groups = {id(group): group for group, _ in engine_infos} + if any( + group.lifecycle_status not in (EngineGroupLifecycle.ACTIVE, EngineGroupLifecycle.DRAINING) + for group in selected_groups.values() + ): + request.update_status(ScaleInStatus.FAILED, "Selected engine group is already being removed") + return + for group in selected_groups.values(): + group.lifecycle_status = EngineGroupLifecycle.DRAINING + + # Publish DRAINING before waiting. Existing transfers may finish, + # but a new fully-async lease cannot enter while removal is pending. + wait_start = time.monotonic() + while self._is_weight_updating: + elapsed = time.monotonic() - wait_start + if elapsed >= request.timeout_secs: + request.update_status( + ScaleInStatus.FAILED, + f"Timed out waiting {request.timeout_secs}s for the weight-update fence", + ) + logger.warning("[ScaleIn] Timed out waiting for weight update; aborting scale-in") + return + await asyncio.sleep(min(1.0, request.timeout_secs - elapsed)) drain_timeout = getattr(self.args, "scale_in_drain_timeout", 30.0) shutdown_timeout = getattr(self.args, "scale_in_shutdown_timeout", 20.0) request.update_status(ScaleInStatus.DRAINING) - unregistered_engine_infos, router_failed = await self._drain_engines( + removed, failed = await self._remove_live_engines( + srv, engine_infos, - timeout=drain_timeout, + drain_timeout=drain_timeout, + shutdown_timeout=shutdown_timeout, force=request.force, ) request.update_status(ScaleInStatus.REMOVING) - removed = [] - failed = list(router_failed) - for group, node0_idx in unregistered_engine_infos: - engine_id = f"group_{group.rank_offset}_engine_{node0_idx}" - try: - await self._remove_engine(group, node0_idx, shutdown_timeout=shutdown_timeout) - removed.append(engine_id) - logger.info(f"[ScaleIn] Removed engine {engine_id}") - except Exception as e: - failed.append(engine_id) - logger.warning(f"[ScaleIn] Failed to remove engine {engine_id}: {e}") - request.removed_engines = removed request.failed_engines = failed - - self._cleanup_engine_groups(srv) + completed_eviction_groups = { + id(group) + for group, node0_idx in engine_infos + if group.eviction_requested and f"group_{group.rank_offset}_engine_{node0_idx}" in removed + } if failed: failure_prefix = "Scale-in partially failed" if removed else "Scale-in failed" @@ -3673,8 +3804,40 @@ async def _scale_in(self, request: ScaleInRequest) -> None: except Exception as e: request.update_status(ScaleInStatus.FAILED, f"Scale-in failed: {e}") logger.exception(f"[ScaleIn] Unhandled error in scale-in for request {request.request_id}") + finally: + with self._engine_lifecycle_lock: + for group in selected_groups.values(): + if id(group) in completed_eviction_groups: + group.eviction_requested = False + if not group.eviction_requested and group.lifecycle_status is EngineGroupLifecycle.DRAINING: + group.lifecycle_status = EngineGroupLifecycle.ACTIVE + + async def _resolve_scale_in_url_candidates(self, request: ScaleInRequest, srv) -> list: + """Resolve URL targets without holding the lifecycle lock.""" + target_urls = {self._normalize_engine_addr(url) for url in request.engine_urls} + with self._engine_lifecycle_lock: + snapshots = [ + (group, node0_idx, engine) + for group in srv.engine_groups + if group.is_scaled_out + and group.lifecycle_status in (EngineGroupLifecycle.ACTIVE, EngineGroupLifecycle.DRAINING) + for node0_idx, engine in enumerate(group.engines) + if engine is not None + ] + + async def _resolve(group, node0_idx, engine): + try: + url = await asyncio.wait_for(engine.get_url.remote(), timeout=5) + if url and self._normalize_engine_addr(url) in target_urls: + return group, node0_idx, engine + except Exception as e: + logger.warning(f"Failed to get URL for engine group_{group.rank_offset}_engine_{node0_idx}: {e}") + return None + + results = await asyncio.gather(*[_resolve(*snapshot) for snapshot in snapshots]) + return [result for result in results if result is not None] - def _select_engines_for_removal(self, request: ScaleInRequest, srv) -> list: + def _select_engines_for_removal(self, request: ScaleInRequest, srv, *, url_candidates: list | None = None) -> list: """Select engines eligible for removal during scale-in. Only engines belonging to groups that were added via scale-out @@ -3682,40 +3845,99 @@ def _select_engines_for_removal(self, request: ScaleInRequest, srv) -> list: touched, regardless of how ``num_replicas`` / ``engine_urls`` are specified. """ - # Collect candidates: only from scale-out groups - engine_infos = [] - for group in srv.engine_groups: - if not group.is_scaled_out: - continue - for node0_idx, engine in enumerate(group.engines): - if engine is not None: - engine_infos.append((group, node0_idx)) + # Collect one coherent snapshot. A DRAINING group can be a SIGTERM + # intent claimed after this request was inserted; prefer it so a + # target-based scale-in does not remove a second elastic engine. + with self._engine_lifecycle_lock: + engine_snapshots = [] + for group in srv.engine_groups: + if not group.is_scaled_out or group.lifecycle_status not in ( + EngineGroupLifecycle.ACTIVE, + EngineGroupLifecycle.DRAINING, + ): + continue + for node0_idx, engine in enumerate(group.engines): + if engine is not None: + engine_snapshots.append((group, node0_idx, engine)) + current_total = sum(1 for group in srv.engine_groups for engine in group.engines if engine is not None) if request.num_replicas > 0: # Count ALL live engines (initial + scaled-out) to decide how many # to remove so the cluster reaches the target size. - current_total = sum(1 for g in srv.engine_groups for e in g.engines if e is not None) num_to_remove = current_total - request.num_replicas if num_to_remove <= 0: return [] - # Remove from the tail (most recently added) first; never exceed - # the number of eligible scale-out engines. - engine_infos = engine_infos[-num_to_remove:] + draining = [item for item in engine_snapshots if item[0].lifecycle_status is EngineGroupLifecycle.DRAINING] + active = [item for item in engine_snapshots if item[0].lifecycle_status is EngineGroupLifecycle.ACTIVE] + # Adopt SIGTERM intents first, then remove the most recently added + # ACTIVE engines only for the remaining target delta. + engine_snapshots = draining[:num_to_remove] + remaining = num_to_remove - len(engine_snapshots) + if remaining > 0: + engine_snapshots.extend(active[-remaining:]) elif request.engine_urls: - # Match by engine URLs (normalize both sides so http:// prefix doesn't matter) - target_urls = {self._normalize_engine_addr(u) for u in request.engine_urls} - matched_infos = [] - for g, idx in engine_infos: - engine = g.engines[idx] - try: - url = ray.get(engine.get_url.remote(), timeout=5) - if url and self._normalize_engine_addr(url) in target_urls: - matched_infos.append((g, idx)) - except Exception as e: - logger.warning(f"Failed to get URL for engine group_{g.rank_offset}_engine_{idx}: {e}") - engine_infos = matched_infos + # URL probes happen outside the lifecycle lock. Revalidate actor + # identity against the current topology before claiming removal. + resolved_by_slot = {(id(group), node0_idx): engine for group, node0_idx, engine in (url_candidates or [])} + engine_snapshots = [ + (group, node0_idx, engine) + for group, node0_idx, engine in engine_snapshots + if resolved_by_slot.get((id(group), node0_idx)) is engine + ] - return engine_infos + return [(group, node0_idx) for group, node0_idx, _ in engine_snapshots] + + async def _remove_live_engines( + self, + srv, + engine_infos: list, + *, + drain_timeout: float, + shutdown_timeout: float, + force: bool, + ) -> tuple[list[str], list[str]]: + """Remove live actors in Router -> drain -> DCS -> shutdown -> PG + order.""" + unregistered_engine_infos, router_failed = await self._drain_engines( + engine_infos, + timeout=drain_timeout, + force=force, + ) + + removal_targets = [ + ( + group, + node0_idx, + f"group_{group.rank_offset}_engine_{node0_idx}", + self._get_live_engine_actors(group, node0_idx), + ) + for group, node0_idx in unregistered_engine_infos + ] + dcs_results = await asyncio.gather( + *[self._unregister_engine_dcs(engine_id, live_actors) for _, _, engine_id, live_actors in removal_targets], + return_exceptions=True, + ) + removed = [] + failed = list(router_failed) + shutdown_targets = [] + for target, result in zip(removal_targets, dcs_results): + _, _, engine_id, _ = target + if isinstance(result, BaseException): + logger.warning(f"[ScaleIn] Failed to unregister DCS for engine {engine_id}: {result}") + shutdown_targets.append(target) + + await asyncio.gather( + *[ + self._shutdown_engine_actors(group, engine_id, live_actors, shutdown_timeout) + for group, _, engine_id, live_actors in shutdown_targets + ] + ) + for _, _, engine_id, _ in shutdown_targets: + removed.append(engine_id) + logger.info(f"[ScaleIn] Removed engine {engine_id}") + + self._cleanup_engine_groups(srv) + return removed, failed async def _drain_engines(self, engine_infos: list, timeout: float, force: bool) -> tuple[list, list[str]]: """Remove all engines from the router, then wait once for the drain @@ -3784,60 +4006,97 @@ async def _mark_one_engine(group, node0_idx): return unregistered_engine_infos, failed_engine_ids - async def _remove_engine(self, group, node0_idx: int, shutdown_timeout: float) -> None: - engine_id = f"group_{group.rank_offset}_engine_{node0_idx}" + @staticmethod + def _get_live_engine_actors(group, node0_idx: int) -> list[tuple[int, object]]: nodes_per_engine = group.nodes_per_engine indices = range(node0_idx * nodes_per_engine, (node0_idx + 1) * nodes_per_engine) + return [ + (i, group.all_engines[i]) + for i in indices + if i < len(group.all_engines) and group.all_engines[i] is not None + ] - for i in indices: - if i >= len(group.all_engines): - continue - engine = group.all_engines[i] - if engine is None: - continue + @staticmethod + async def _unregister_engine_dcs(engine_id: str, live_actors: list[tuple[int, object]]) -> None: + async def _unregister_dcs(engine): + return await asyncio.wait_for(engine.unregister_dcs.remote(), timeout=10) - try: - await asyncio.wait_for(engine.unregister_dcs.remote(), timeout=10) - except Exception as e: - logger.warning(f"[ScaleIn] Failed to unregister DCS for engine {engine_id}[{i}]: {e}") + dcs_results = await asyncio.gather( + *[_unregister_dcs(engine) for _, engine in live_actors], + return_exceptions=True, + ) + dcs_failures = [ + (i, result) for (i, _), result in zip(live_actors, dcs_results) if isinstance(result, BaseException) + ] + if dcs_failures: + failed_indices = ", ".join(str(i) for i, _ in dcs_failures) + raise RuntimeError(f"Failed to unregister DCS for engine {engine_id}[{failed_indices}]") - shutdown_ok = False - try: - await asyncio.wait_for(engine.shutdown.remote(), timeout=shutdown_timeout) - shutdown_ok = True - except Exception as e: - logger.warning(f"[ScaleIn] Failed to shutdown engine {engine_id}[{i}]: {e}") + async def _shutdown_engine_actors( + self, + group, + engine_id: str, + live_actors: list[tuple[int, object]], + shutdown_timeout: float, + ) -> None: + async def _shutdown(engine): + return await asyncio.wait_for(engine.shutdown.remote(), timeout=shutdown_timeout) - if not shutdown_ok: + shutdown_results = await asyncio.gather( + *[_shutdown(engine) for _, engine in live_actors], + return_exceptions=True, + ) + for (i, engine), result in zip(live_actors, shutdown_results): + if isinstance(result, BaseException): + logger.warning(f"[ScaleIn] Failed to shutdown engine {engine_id}[{i}]: {result}") try: ray.kill(engine) except Exception as e: logger.warning(f"[ScaleIn] Failed to kill engine actor {engine_id}[{i}]: {e}") - group.all_engines[i] = None + with self._engine_lifecycle_lock: + for i, _ in live_actors: + group.all_engines[i] = None - def _cleanup_engine_groups(self, srv) -> None: - monitors_to_remove = [] - groups_to_remove = [] + async def _remove_engine(self, group, node0_idx: int, shutdown_timeout: float) -> None: + engine_id = f"group_{group.rank_offset}_engine_{node0_idx}" + live_actors = self._get_live_engine_actors(group, node0_idx) + try: + await self._unregister_engine_dcs(engine_id, live_actors) + except Exception as e: + logger.warning(f"[ScaleIn] Failed to unregister DCS for engine {engine_id}: {e}") + await self._shutdown_engine_actors(group, engine_id, live_actors, shutdown_timeout) - for group in srv.engine_groups: - if all(e is None for e in group.all_engines): - groups_to_remove.append(group) - for monitor in self._health_monitors: - if monitor._engine_group is group: - monitors_to_remove.append(monitor) + def _cleanup_engine_groups(self, srv) -> None: + with self._engine_lifecycle_lock: + groups_to_remove = [ + group + for group in srv.engine_groups + if group.lifecycle_status in (EngineGroupLifecycle.ACTIVE, EngineGroupLifecycle.DRAINING) + and all(engine is None for engine in group.all_engines) + ] + monitors_to_remove = [ + monitor + for monitor in self._health_monitors + if any(monitor._engine_group is group for group in groups_to_remove) + ] + for group in groups_to_remove: + group.lifecycle_status = EngineGroupLifecycle.REMOVING + srv.engine_groups.remove(group) + for monitor in monitors_to_remove: + self._health_monitors.remove(monitor) for monitor in monitors_to_remove: monitor.stop() - self._health_monitors.remove(monitor) for group in groups_to_remove: - srv.engine_groups.remove(group) if group.pg is not None: try: ray.util.remove_placement_group(group.pg[0]) except Exception as e: logger.warning(f"[ScaleIn] Failed to remove placement group: {e}") + with self._engine_lifecycle_lock: + group.lifecycle_status = EngineGroupLifecycle.REMOVED if groups_to_remove: logger.info(f"[ScaleIn] Cleaned up {len(groups_to_remove)} empty engine groups") @@ -3953,120 +4212,160 @@ def _eviction_monitor_loop(self): logger.exception("[Eviction] Unhandled error in eviction monitor loop") def _check_and_handle_evictions(self): - """Check all engines for SIGTERM eviction and handle evicted ones as - scale-in. - - This method polls every live engine via ``is_evicted()`` in parallel. - Evicted engines are removed from the engine group and cleaned up, - similar to the existing scale-in flow but without requiring an external - API call. - """ - # Collect all live engines with their (server_name, group, node0_idx) - engine_refs = [] - engine_info_map = [] - for srv_name, srv in self.servers.items(): - for group in srv.engine_groups: - for node0_idx, engine in enumerate(group.engines): - if engine is None: + """Poll live elastic actors independently for graceful SIGTERM + intent.""" + pending = [] + with self._engine_lifecycle_lock: + for server_name, srv in self.servers.items(): + for group in srv.engine_groups: + if ( + not group.is_scaled_out + or group.pg is None + or group.lifecycle_status not in (EngineGroupLifecycle.ACTIVE, EngineGroupLifecycle.DRAINING) + ): continue - try: - ref = engine.is_evicted.remote() - engine_refs.append(ref) - engine_info_map.append((srv_name, group, node0_idx, engine)) - except Exception: - # Engine actor may already be dead - pass - - if not engine_refs: - return - - # Parallel poll with timeout — an engine that's already dead will - # raise; we treat that as "not evicted" (the health monitor handles dead actors). - try: - results = ray.get(engine_refs, timeout=5) - except ray.exceptions.GetTimeoutError: - logger.warning("[Eviction] Timed out polling engines for eviction status") - return - except Exception as e: - logger.debug(f"[Eviction] Error polling engines: {e}") - return + for actor_idx, engine in enumerate(group.all_engines): + if engine is None: + continue + try: + pending.append( + ( + engine.is_evicted.remote(), + server_name, + group, + actor_idx // group.nodes_per_engine, + ) + ) + except Exception as e: + logger.debug("[Eviction] Failed to submit graceful probe: %s", e) - evicted = [ - (srv_name, group, node0_idx, engine) - for (srv_name, group, node0_idx, engine), is_evict in zip(engine_info_map, results) - if is_evict - ] - if not evicted: + if not pending: return - logger.info( - f"[Eviction] Detected {len(evicted)} evicted engine(s), " - f"processing as scale-in: " - f"{[(srv_name, f'group_{g.rank_offset}_engine_{idx}') for srv_name, g, idx, _ in evicted]}" - ) - - for srv_name, group, node0_idx, engine in evicted: - self._handle_single_eviction(srv_name, group, node0_idx) - - def _handle_single_eviction(self, srv_name: str, group, node0_idx: int): - """Handle a single evicted engine: unregister DCS, kill actor, clean - up. - - This mirrors the scale-in removal path but is triggered by eviction - rather than an API request. The SIGTERM handler in SGLangEngine - already unregistered the engine from the router, so we skip the drain - step. - """ - engine_id = f"group_{group.rank_offset}_engine_{node0_idx}" - logger.info(f"[Eviction] Handling evicted engine: {engine_id}") - - # Mark as intentionally removed in health monitor so it doesn't - # try to recover the engine. - for monitor in self._health_monitors: - if monitor._engine_group is group: - monitor.mark_intentionally_removed(node0_idx) - - nodes_per_engine = group.nodes_per_engine - indices = range(node0_idx * nodes_per_engine, (node0_idx + 1) * nodes_per_engine) - - for i in indices: - if i >= len(group.all_engines): - continue - engine = group.all_engines[i] - if engine is None: + refs = [item[0] for item in pending] + ready, _ = ray.wait(refs, num_returns=len(refs), timeout=5) + ready_set = set(ready) + handled = set() + evicted_engine_infos = [] + for ref, server_name, group, node0_idx in pending: + if ref not in ready_set or (id(group), node0_idx) in handled: continue - - # Best-effort DCS unregister try: - ray.get(engine.unregister_dcs.remote(), timeout=5) + if ray.get(ref): + handled.add((id(group), node0_idx)) + evicted_engine_infos.append((server_name, group, node0_idx)) except Exception as e: - logger.warning(f"[Eviction] Failed to unregister DCS for {engine_id}[{i}]: {e}") + # Hard actor failure belongs to the fault-recovery follow-up, + # not this graceful-eviction path. + logger.debug("[Eviction] Graceful probe failed: %s", e) + + if evicted_engine_infos: + self._handle_evictions(evicted_engine_infos) + + def _handle_evictions(self, eviction_infos: list[tuple[str, EngineGroup, int]]) -> None: + """Fence a ready eviction batch before performing one live removal.""" + claimed = [] + with self._engine_lifecycle_lock: + for srv_name, group, node0_idx in eviction_infos: + srv = self.servers.get(srv_name) + if ( + srv is None + or group not in srv.engine_groups + or group.lifecycle_status not in (EngineGroupLifecycle.ACTIVE, EngineGroupLifecycle.DRAINING) + or not group.is_scaled_out + or group.pg is None + or node0_idx >= len(group.engines) + or group.engines[node0_idx] is None + ): + continue + group.eviction_requested = True + group.lifecycle_status = EngineGroupLifecycle.DRAINING + claimed.append((srv_name, srv, group, node0_idx)) - # Best-effort shutdown — the process may already be terminating - try: - ray.get(engine.shutdown.remote(), timeout=10) - except Exception as e: - logger.debug(f"[Eviction] Engine shutdown failed (expected if pod is terminating): {e}") + if not claimed: + return - # Kill the Ray actor - try: - ray.kill(engine) - except Exception as e: - logger.debug(f"[Eviction] ray.kill failed for {engine_id}[{i}] (may already be dead): {e}") + # Signal-first fence: every ready intent is DRAINING before any + # owner wait or cleanup. An explicit scale-in remains the sole + # removal owner and adopts matching DRAINING engines. + active = self._find_active_scale_request() + if active is not None and active["type"] == "scale_in": + engine_ids = [f"group_{group.rank_offset}_engine_{idx}" for _, _, group, idx in claimed] + logger.info( + "[Eviction] Scale-in request %s will adopt pending evictions %s (%s)", + active["request_id"], + engine_ids, + active["status"], + ) + return + if active is not None and active["type"] not in ("scale_out", "graceful_eviction"): + return - group.all_engines[i] = None + engine_ids = [f"group_{group.rank_offset}_engine_{idx}" for _, _, group, idx in claimed] + logger.info("[Eviction] Gracefully removing engines %s", engine_ids) + try: + wait_timeout = getattr(self.args, "scale_in_drain_timeout", 30.0) + 60.0 + deadline = time.monotonic() + wait_timeout + while self._is_weight_updating: + remaining = deadline - time.monotonic() + if remaining <= 0: + logger.warning("[Eviction] Timed out waiting for the weight-update fence for %s", engine_ids) + return + time.sleep(min(1.0, remaining)) + + batches: dict[str, tuple[RolloutServer, list[tuple[EngineGroup, int]]]] = {} + for srv_name, srv, group, node0_idx in claimed: + if srv_name not in batches: + batches[srv_name] = (srv, []) + batches[srv_name][1].append((group, node0_idx)) + + async def _remove_batches(): + batch_items = list(batches.items()) + results = await asyncio.gather( + *[ + self._remove_live_engines( + srv, + infos, + drain_timeout=getattr(self.args, "scale_in_drain_timeout", 30.0), + shutdown_timeout=getattr(self.args, "scale_in_shutdown_timeout", 20.0), + force=False, + ) + for _, (srv, infos) in batch_items + ], + return_exceptions=True, + ) + return [(srv_name, result) for (srv_name, _), result in zip(batch_items, results)] - logger.info(f"[Eviction] Engine {engine_id} removed from engine group") + batch_results = asyncio.run(_remove_batches()) + except Exception: + logger.exception("[Eviction] Graceful removal failed for %s", engine_ids) + return - # Clean up empty engine groups - srv = self.servers.get(srv_name) - if srv: - self._cleanup_engine_groups(srv) - remaining = sum(1 for g in srv.engine_groups for e in g.engines if e is not None) - logger.info( - f"[Eviction] Server '{srv_name}' now has {remaining} live engine(s) " - f"across {len(srv.engine_groups)} group(s)" - ) + removed_by_server = {} + for srv_name, result in batch_results: + if isinstance(result, BaseException): + logger.warning("[Eviction] Graceful removal batch failed for %s: %s", srv_name, result) + removed_by_server[srv_name] = set() + else: + removed, _ = result + removed_by_server[srv_name] = set(removed) + + claimed_by_group: dict[int, tuple[EngineGroup, list[tuple[str, str]]]] = {} + for srv_name, _, group, node0_idx in claimed: + engine_id = f"group_{group.rank_offset}_engine_{node0_idx}" + claimed_by_group.setdefault(id(group), (group, []))[1].append((srv_name, engine_id)) + + with self._engine_lifecycle_lock: + # Restore a surviving group only when every eviction claimed for + # that group completed. Any timeout/failure remains fail-closed. + for group, group_claims in claimed_by_group.values(): + completed = all( + engine_id in removed_by_server.get(srv_name, set()) for srv_name, engine_id in group_claims + ) + if completed: + group.eviction_requested = False + if completed and group.lifecycle_status is EngineGroupLifecycle.DRAINING: + group.lifecycle_status = EngineGroupLifecycle.ACTIVE def _allocate_rollout_engine_addr_and_ports_external(args, rollout_engines): diff --git a/tests/backends/megatron/test_actor_http_timeout.py b/tests/backends/megatron/test_actor_http_timeout.py index 690c8c7ea..4c2e9110b 100644 --- a/tests/backends/megatron/test_actor_http_timeout.py +++ b/tests/backends/megatron/test_actor_http_timeout.py @@ -13,6 +13,8 @@ timeout, otherwise a legitimate recovery would be interrupted. """ +from unittest.mock import MagicMock + import pytest import requests @@ -29,10 +31,13 @@ class _Resp: - def __init__(self, payload=False): + def __init__(self, payload=False, status_code=200): self._payload = payload + self.status_code = status_code def raise_for_status(self): + if self.status_code >= 400: + raise requests.exceptions.HTTPError(f"status={self.status_code}") return None def json(self): @@ -182,3 +187,51 @@ def _timeout_get(*_a, **_k): rollout_only, actor_fwd_only = shell._check_services_health() assert actor_fwd_only is True assert rollout_only is True + + +def test_check_services_health_retries_scale_in_fence(monkeypatch): + from argparse import Namespace + + calls = [] + attempts = iter([_Resp(status_code=503), _Resp(payload=True)]) + + def _get(url, *args, **kwargs): + calls.append(url) + if url.endswith("/can_do_update_weight_for_async"): + return next(attempts) + return _Resp() + + _patch_health_common(monkeypatch, _get) + shell = _shell() + shell.args = Namespace(true_on_policy_mode=True, hybrid=False, rollout_http_timeout=120.0) + + assert shell._check_services_health() == (True, False) + assert sum(url.endswith("/can_do_update_weight_for_async") for url in calls) == 2 + + +def test_check_services_health_resets_fence_deadline_after_non_503(monkeypatch): + from argparse import Namespace + + calls = [] + attempts = iter( + [ + _Resp(status_code=503), + _Resp(payload=False), + _Resp(payload=False), + _Resp(payload=True), + ] + ) + + def _get(url, *args, **kwargs): + calls.append(url) + if url.endswith("/can_do_update_weight_for_async"): + return next(attempts) + return _Resp() + + _patch_health_common(monkeypatch, _get) + monkeypatch.setattr(actor_mod.time, "monotonic", MagicMock(side_effect=[0.0, 0.0, 1.0, 3.0])) + shell = _shell() + shell.args = Namespace(true_on_policy_mode=True, hybrid=False, rollout_http_timeout=2.0) + + assert shell._check_services_health() == (True, False) + assert sum(url.endswith("/can_do_update_weight_for_async") for url in calls) == 4 diff --git a/tests/backends/sglang/test_sigterm_eviction.py b/tests/backends/sglang/test_sigterm_eviction.py new file mode 100644 index 000000000..b22682939 --- /dev/null +++ b/tests/backends/sglang/test_sigterm_eviction.py @@ -0,0 +1,30 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Signal-safety regression tests for graceful elastic eviction.""" + +import signal +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from relax.backends.sglang.sglang_engine import SGLangEngine + + +def test_sigterm_handler_only_publishes_eviction_intent(): + engine = SimpleNamespace(_evicted=MagicMock()) + captured = {} + + def _capture(_signal_number, handler): + captured["handler"] = handler + + with ( + patch("relax.backends.sglang.sglang_engine.signal.getsignal", return_value=signal.SIG_DFL), + patch("relax.backends.sglang.sglang_engine.signal.signal", side_effect=_capture), + patch("relax.backends.sglang.sglang_engine.time.sleep") as sleep, + patch.object(SGLangEngine, "unregister_from_router") as unregister, + ): + SGLangEngine._register_sigterm_handler(engine) + captured["handler"](signal.SIGTERM, None) + + engine._evicted.set.assert_called_once_with() + sleep.assert_not_called() + unregister.assert_not_called() diff --git a/tests/components/test_rollout_weight_update_handshake.py b/tests/components/test_rollout_weight_update_handshake.py index af3af0f4b..11bc7996e 100644 --- a/tests/components/test_rollout_weight_update_handshake.py +++ b/tests/components/test_rollout_weight_update_handshake.py @@ -19,6 +19,7 @@ import logging import pytest +from fastapi import HTTPException from ray.exceptions import RayActorError from relax.components.rollout import Rollout as RolloutDeployment @@ -44,14 +45,18 @@ async def _ok(*_args, **_kwargs): return None +async def _prepared(*_args, **_kwargs): + return True + + def _raise_dead(*_args, **_kwargs): raise RayActorError() class _ManagerStub: - def __init__(self, set_weight_updating_fn=None): - self.health_monitoring_pause = _RemoteStub(lambda *a, **k: _ok()) - self.set_weight_updating = _RemoteStub(set_weight_updating_fn or (lambda *a, **k: _ok())) + def __init__(self, set_weight_updating_fn=None, health_monitoring_pause_fn=None): + self.health_monitoring_pause = _RemoteStub(health_monitoring_pause_fn or (lambda *a, **k: _ok())) + self.set_weight_updating = _RemoteStub(set_weight_updating_fn or (lambda *a, **k: _prepared())) def _make_rollout(*, can_update: bool, manager: _ManagerStub) -> "Rollout": @@ -99,6 +104,60 @@ def test_can_do_update_weight_dead_engine_releases_gate_then_reraises(): assert shell._weight_update_ready.is_set() +def test_can_do_update_weight_rejected_lease_rolls_back_before_pause(): + calls = [] + + async def _set_weight_updating(value): + calls.append(value) + return not value + + async def _pause(): + calls.append("pause") + + manager = _ManagerStub( + set_weight_updating_fn=_set_weight_updating, + health_monitoring_pause_fn=_pause, + ) + shell = _make_rollout(can_update=True, manager=manager) + + with pytest.raises(HTTPException) as exc_info: + asyncio.run(shell.can_do_update_weight_for_async()) + + assert exc_info.value.status_code == 503 + assert calls == [True, False] + assert shell.status == "running" + assert shell._weight_update_ready.is_set() + + +def test_can_do_update_weight_closes_local_admission_while_lease_is_pending(): + async def _scenario(): + lease_started = asyncio.Event() + finish_lease = asyncio.Event() + + async def _set_weight_updating(value): + if value: + lease_started.set() + await finish_lease.wait() + return False + return True + + shell = _make_rollout( + can_update=True, + manager=_ManagerStub(set_weight_updating_fn=_set_weight_updating), + ) + handshake = asyncio.create_task(shell.can_do_update_weight_for_async()) + await asyncio.wait_for(lease_started.wait(), timeout=1) + + assert shell.status == "paused" + + finish_lease.set() + with pytest.raises(HTTPException): + await handshake + assert shell.status == "running" + + asyncio.run(_scenario()) + + def test_end_update_weight_does_not_block_after_failed_can_do(): async def _scenario(): manager = _ManagerStub(set_weight_updating_fn=_raise_dead) diff --git a/tests/distributed/ray/conftest.py b/tests/distributed/ray/conftest.py index 2fdbba590..bbf18c003 100644 --- a/tests/distributed/ray/conftest.py +++ b/tests/distributed/ray/conftest.py @@ -4,6 +4,7 @@ import os import sys +import threading from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -227,7 +228,17 @@ def create_test_manager(args=None, servers=None): manager.servers = servers if servers is not None else {} manager._scale_out_requests = {} manager._scale_in_requests = {} - manager._is_weight_updating = False + manager._engine_lifecycle_lock = threading.RLock() + manager._training_weight_updating = False + manager._scale_out_weight_updating = False + manager._next_engine_rank = max( + ( + group.rank_offset + len(group.all_engines) + for srv in manager.servers.values() + for group in srv.engine_groups + ), + default=0, + ) # Mock the distributed lock lock = MagicMock() diff --git a/tests/distributed/ray/test_coordination.py b/tests/distributed/ray/test_coordination.py index e8e4cdfc0..bd8c3e6f1 100644 --- a/tests/distributed/ray/test_coordination.py +++ b/tests/distributed/ray/test_coordination.py @@ -3,13 +3,16 @@ """Tests for mutual exclusion, GC of terminal requests, eviction monitoring, and engine info queries.""" -from unittest.mock import MagicMock, patch +import asyncio +import threading +from unittest.mock import AsyncMock, MagicMock, patch import pytest try: from relax.distributed.ray.rollout import ( + EngineGroupLifecycle, ScaleInRequest, ScaleInStatus, ScaleOutRequest, @@ -118,6 +121,35 @@ def test_all_non_terminal_scale_in_detected(self, status_str): ) assert manager._find_active_scale_request() is not None + def test_finds_claimed_graceful_eviction_without_request(self): + group = make_engine_group(is_scaled_out=True) + group.lifecycle_status = EngineGroupLifecycle.DRAINING + group.eviction_requested = True + manager = create_test_manager(servers={"default": make_rollout_server(engine_groups=[group])}) + + result = manager._find_active_scale_request() + + assert result == { + "type": "graceful_eviction", + "request_id": "default:0", + "status": "DRAINING", + } + + def test_explicit_scale_in_precedes_its_draining_group(self): + group = make_engine_group(is_scaled_out=True) + group.lifecycle_status = EngineGroupLifecycle.DRAINING + group.eviction_requested = True + manager = create_test_manager(servers={"default": make_rollout_server(engine_groups=[group])}) + manager._scale_in_requests["scale-in"] = ScaleInRequest( + request_id="scale-in", + status=ScaleInStatus.DRAINING, + ) + + result = manager._find_active_scale_request() + + assert result["type"] == "scale_in" + assert result["request_id"] == "scale-in" + # ===================== _gc_terminal_requests =============================== @@ -201,15 +233,17 @@ def test_detects_evicted_engines(self, patch_ray_get): e_normal = make_mock_engine(evicted=False) e_evicted = make_mock_engine(evicted=True) g = make_engine_group(engines=[e_normal, e_evicted], is_scaled_out=True) + g.pg = (MagicMock(), [], []) srv = make_rollout_server(engine_groups=[g]) manager = create_test_manager(servers={"default": srv}) - with patch.object(manager, "_handle_single_eviction") as mock_handle: + with ( + patch.object(manager, "_handle_evictions") as mock_handle, + patch("relax.distributed.ray.rollout.ray.wait", side_effect=lambda refs, **_kwargs: (refs, [])), + ): manager._check_and_handle_evictions() mock_handle.assert_called_once() - call_args = mock_handle.call_args - assert call_args[0][0] == "default" # srv_name - assert call_args[0][2] == 1 # node0_idx of evicted engine + assert mock_handle.call_args.args[0] == [("default", g, 1)] def test_no_evictions(self, patch_ray_get): e = make_mock_engine(evicted=False) @@ -217,55 +251,313 @@ def test_no_evictions(self, patch_ray_get): srv = make_rollout_server(engine_groups=[g]) manager = create_test_manager(servers={"default": srv}) - with patch.object(manager, "_handle_single_eviction") as mock_handle: + with patch.object(manager, "_handle_evictions") as mock_handle: manager._check_and_handle_evictions() mock_handle.assert_not_called() + def test_pending_probe_does_not_hide_another_eviction(self, patch_ray_get): + pending = make_mock_engine(evicted=False) + evicted = make_mock_engine(evicted=True) + group = make_engine_group(engines=[pending, evicted], is_scaled_out=True) + group.pg = (MagicMock(), [], []) + manager = create_test_manager(servers={"default": make_rollout_server(engine_groups=[group])}) + + def _ready_only(refs, **_kwargs): + return [ref for ref in refs if getattr(ref, "value", False)], [] + + with ( + patch.object(manager, "_handle_evictions") as handle, + patch("relax.distributed.ray.rollout.ray.wait", side_effect=_ready_only), + ): + manager._check_and_handle_evictions() + + handle.assert_called_once_with([("default", group, 1)]) + def test_all_dead_engines_skipped(self, patch_ray_get): g = make_engine_group(engines=[None, None]) srv = make_rollout_server(engine_groups=[g]) manager = create_test_manager(servers={"default": srv}) - with patch.object(manager, "_handle_single_eviction") as mock_handle: + with patch.object(manager, "_handle_evictions") as mock_handle: manager._check_and_handle_evictions() mock_handle.assert_not_called() + def test_collects_ready_evictions_into_one_batch(self, patch_ray_get): + first = make_mock_engine(evicted=True) + second = make_mock_engine(evicted=True) + first_group = make_engine_group(engines=[first], is_scaled_out=True, rank_offset=1) + second_group = make_engine_group(engines=[second], is_scaled_out=True, rank_offset=2) + first_group.pg = (MagicMock(), [], []) + second_group.pg = (MagicMock(), [], []) + manager = create_test_manager( + servers={"default": make_rollout_server(engine_groups=[first_group, second_group])} + ) + + with ( + patch.object(manager, "_handle_evictions") as handle, + patch("relax.distributed.ray.rollout.ray.wait", side_effect=lambda refs, **_kwargs: (refs, [])), + ): + manager._check_and_handle_evictions() + + handle.assert_called_once_with([("default", first_group, 0), ("default", second_group, 0)]) + + +class TestHandleEvictions: + def test_batch_claims_all_groups_before_waiting_once(self, patch_ray_get): + first_group = make_engine_group(engines=[make_mock_engine()], is_scaled_out=True, rank_offset=1) + second_group = make_engine_group(engines=[make_mock_engine()], is_scaled_out=True, rank_offset=2) + first_group.pg = (MagicMock(), [], []) + second_group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[first_group, second_group]) + manager = create_test_manager(servers={"default": server}) + manager._training_weight_updating = True + removed = ["group_1_engine_0", "group_2_engine_0"] + + def release_owner(_seconds): + assert first_group.lifecycle_status is EngineGroupLifecycle.DRAINING + assert second_group.lifecycle_status is EngineGroupLifecycle.DRAINING + manager._training_weight_updating = False + + with ( + patch("relax.distributed.ray.rollout.time.sleep", side_effect=release_owner) as sleep, + patch.object( + manager, "_remove_live_engines", new_callable=AsyncMock, return_value=(removed, []) + ) as remove, + ): + manager._handle_evictions([("default", first_group, 0), ("default", second_group, 0)]) + + sleep.assert_called_once() + remove.assert_awaited_once_with( + server, + [(first_group, 0), (second_group, 0)], + drain_timeout=30.0, + shutdown_timeout=30.0, + force=False, + ) + + def test_batch_dcs_failure_does_not_block_cleanup(self, patch_ray_get): + dcs_failed_engine = make_mock_engine() + dcs_failed_engine.unregister_dcs.remote.side_effect = RuntimeError("DCS unavailable") + healthy_engine = make_mock_engine() + dcs_failed_group = make_engine_group(engines=[dcs_failed_engine], is_scaled_out=True, rank_offset=1) + healthy_group = make_engine_group(engines=[healthy_engine], is_scaled_out=True, rank_offset=2) + dcs_failed_group.pg = (MagicMock(), [], []) + healthy_group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[dcs_failed_group, healthy_group]) + manager = create_test_manager(servers={"default": server}) + manager.args.scale_in_drain_timeout = 0 + + with patch("ray.util.remove_placement_group"): + manager._handle_evictions([("default", dcs_failed_group, 0), ("default", healthy_group, 0)]) + + assert dcs_failed_group.lifecycle_status is EngineGroupLifecycle.REMOVED + assert healthy_group.lifecycle_status is EngineGroupLifecycle.REMOVED + assert dcs_failed_group not in server.engine_groups + assert healthy_group not in server.engine_groups + assert dcs_failed_group.all_engines[0] is None + assert healthy_group.all_engines[0] is None + dcs_failed_engine.shutdown.remote.assert_called_once() + healthy_engine.shutdown.remote.assert_called_once() + assert manager.set_weight_updating(True) is True + + @pytest.mark.parametrize("request_kind", ["scale_out", "scale_in"]) + def test_eviction_claim_blocks_new_scale_request(self, patch_ray_get, request_kind): + initial = make_engine_group(engines=[make_mock_engine()]) + engine = make_mock_engine() + group = make_engine_group(engines=[engine], is_scaled_out=True, rank_offset=1) + group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[initial, group]) + manager = create_test_manager(servers={"default": server}) + manager.args.scale_in_drain_timeout = 0 + removal_started = threading.Event() + allow_removal = threading.Event() + + async def block_removal(*_args, **_kwargs): + removal_started.set() + assert allow_removal.wait(timeout=5) + return ["group_1_engine_0"], [] + + with patch.object(manager, "_remove_live_engines", side_effect=block_removal): + eviction_thread = threading.Thread( + target=manager._handle_evictions, + args=([("default", group, 0)],), + ) + eviction_thread.start() + assert removal_started.wait(timeout=5) + + if request_kind == "scale_out": + result = manager.create_scale_out_request(num_replicas=3) + else: + result = manager.create_scale_in_request(num_replicas=1) + + assert result["status"] == "CONFLICT" + assert group.lifecycle_status is EngineGroupLifecycle.DRAINING + allow_removal.set() + eviction_thread.join(timeout=5) + assert not eviction_thread.is_alive() + + def test_inserted_scale_out_does_not_delay_eviction_fence(self, patch_ray_get): + initial = make_engine_group(engines=[make_mock_engine()]) + engine = make_mock_engine() + group = make_engine_group(engines=[engine], is_scaled_out=True, rank_offset=1) + group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[initial, group]) + manager = create_test_manager(servers={"default": server}) + manager.args.scale_in_drain_timeout = 0 + removal_started = threading.Event() + allow_removal = threading.Event() + + result = manager.create_scale_out_request(num_replicas=3) + assert result["status"] == "PENDING" + + async def block_removal(*_args, **_kwargs): + removal_started.set() + assert allow_removal.wait(timeout=5) + return ["group_1_engine_0"], [] + + with patch.object(manager, "_remove_live_engines", side_effect=block_removal): + eviction_thread = threading.Thread( + target=manager._handle_evictions, + args=([("default", group, 0)],), + ) + eviction_thread.start() + assert removal_started.wait(timeout=5) + + assert group.lifecycle_status is EngineGroupLifecycle.DRAINING + assert manager.set_weight_updating(True) is False + allow_removal.set() + eviction_thread.join(timeout=5) + assert not eviction_thread.is_alive() + + def test_pending_target_scale_in_adopts_eviction_without_extra_removal(self, patch_ray_get): + initial_engine = make_mock_engine() + initial = make_engine_group(engines=[initial_engine]) + elastic_engine = make_mock_engine() + group = make_engine_group(engines=[elastic_engine], is_scaled_out=True, rank_offset=1) + group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[initial, group]) + manager = create_test_manager(servers={"default": server}) + manager.args.scale_in_drain_timeout = 0 + + result = manager.create_scale_in_request(num_replicas=1) + assert result["status"] == "PENDING" + + manager._handle_evictions([("default", group, 0)]) + + assert group.lifecycle_status is EngineGroupLifecycle.DRAINING + elastic_engine.unregister_from_router.remote.assert_not_called() + + request = manager._scale_in_requests[result["request_id"]] + with patch("ray.util.remove_placement_group"): + asyncio.run(manager._scale_in(request)) + + assert request.status is ScaleInStatus.COMPLETED + assert request.selected_engines == ["group_1_engine_0"] + elastic_engine.unregister_from_router.remote.assert_called_once() + initial_engine.unregister_from_router.remote.assert_not_called() + + def test_scale_in_by_url_leaves_other_eviction_for_next_batch(self, patch_ray_get): + initial = make_engine_group(engines=[make_mock_engine()]) + first_engine = make_mock_engine(url="http://first:1") + second_engine = make_mock_engine(url="http://second:2") + first_group = make_engine_group(engines=[first_engine], is_scaled_out=True, rank_offset=1) + second_group = make_engine_group(engines=[second_engine], is_scaled_out=True, rank_offset=2) + first_group.pg = (MagicMock(), [], []) + second_group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[initial, first_group, second_group]) + manager = create_test_manager(servers={"default": server}) + manager.args.scale_in_drain_timeout = 0 + + result = manager.create_scale_in_request(engine_urls=["first:1"]) + assert result["status"] == "PENDING" + manager._handle_evictions([("default", first_group, 0), ("default", second_group, 0)]) + + request = manager._scale_in_requests[result["request_id"]] + with patch("ray.util.remove_placement_group"): + asyncio.run(manager._scale_in(request)) + + assert request.status is ScaleInStatus.COMPLETED + assert request.selected_engines == ["group_1_engine_0"] + assert first_group.lifecycle_status is EngineGroupLifecycle.REMOVED + assert second_group.lifecycle_status is EngineGroupLifecycle.DRAINING + second_engine.unregister_from_router.remote.assert_not_called() + + manager._handle_evictions([("default", second_group, 0)]) + + assert second_group.lifecycle_status is EngineGroupLifecycle.REMOVED + assert first_engine.unregister_from_router.remote.call_count == 1 + assert second_engine.unregister_from_router.remote.call_count == 1 + + def test_eviction_fence_stays_closed_after_weight_update_timeout(self, patch_ray_get): + engine = make_mock_engine() + group = make_engine_group(engines=[engine], is_scaled_out=True) + group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": server}) + manager._training_weight_updating = True + + with patch("relax.distributed.ray.rollout.time.monotonic", side_effect=[0.0, 91.0]): + manager._handle_evictions([("default", group, 0)]) + + assert group.lifecycle_status is EngineGroupLifecycle.DRAINING + assert manager.set_weight_updating(True) is False + engine.unregister_from_router.remote.assert_not_called() + + def test_duplicate_eviction_has_single_removal_owner(self, patch_ray_get): + engine = make_mock_engine() + group = make_engine_group(engines=[engine], is_scaled_out=True) + group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": server}) + manager.args.scale_in_drain_timeout = 0 + + manager._handle_evictions([("default", group, 0)]) + manager._handle_evictions([("default", group, 0)]) + + engine.unregister_from_router.remote.assert_called_once() + engine.unregister_dcs.remote.assert_called_once() + engine.shutdown.remote.assert_called_once() -class TestHandleSingleEviction: def test_marks_intentionally_removed(self, patch_ray_get): e = make_mock_engine() g = make_engine_group(engines=[e], is_scaled_out=True) + g.pg = (MagicMock(), [], []) srv = make_rollout_server(engine_groups=[g]) manager = create_test_manager(servers={"default": srv}) + manager.args.scale_in_drain_timeout = 0 mock_monitor = MagicMock() mock_monitor._engine_group = g manager._health_monitors.append(mock_monitor) with patch("ray.kill"): - manager._handle_single_eviction("default", g, 0) + manager._handle_evictions([("default", g, 0)]) mock_monitor.mark_intentionally_removed.assert_called_with(0) def test_sets_engine_to_none(self, patch_ray_get): e = make_mock_engine() g = make_engine_group(engines=[e], is_scaled_out=True) + g.pg = (MagicMock(), [], []) srv = make_rollout_server(engine_groups=[g]) manager = create_test_manager(servers={"default": srv}) + manager.args.scale_in_drain_timeout = 0 with patch("ray.kill"): - manager._handle_single_eviction("default", g, 0) + manager._handle_evictions([("default", g, 0)]) assert g.all_engines[0] is None def test_cleans_up_empty_groups(self, patch_ray_get): e = make_mock_engine() g = make_engine_group(engines=[e], is_scaled_out=True) + g.pg = (MagicMock(), [], []) srv = make_rollout_server(engine_groups=[g]) manager = create_test_manager(servers={"default": srv}) + manager.args.scale_in_drain_timeout = 0 with patch("ray.kill"): - manager._handle_single_eviction("default", g, 0) + manager._handle_evictions([("default", g, 0)]) # Group should be cleaned up since it's now empty assert len(srv.engine_groups) == 0 @@ -341,17 +633,17 @@ def test_multiple_groups(self, patch_ray_get): class TestSetWeightUpdating: - def test_sets_flag_and_mirrors_to_engines(self, patch_ray_get): + def test_manager_owns_flag_without_signal_handler_rpc(self, patch_ray_get): e1 = make_mock_engine() e2 = make_mock_engine() g = make_engine_group(engines=[e1, e2]) srv = make_rollout_server(engine_groups=[g]) manager = create_test_manager(servers={"default": srv}) - manager.set_weight_updating(True) + assert manager.set_weight_updating(True) is True assert manager._is_weight_updating is True - e1.set_weight_updating.remote.assert_called_with(True) - e2.set_weight_updating.remote.assert_called_with(True) + e1.set_weight_updating.remote.assert_not_called() + e2.set_weight_updating.remote.assert_not_called() def test_unsets_flag(self, patch_ray_get): e1 = make_mock_engine() diff --git a/tests/distributed/ray/test_scale_in.py b/tests/distributed/ray/test_scale_in.py index adb029439..0aaafbf66 100644 --- a/tests/distributed/ray/test_scale_in.py +++ b/tests/distributed/ray/test_scale_in.py @@ -3,13 +3,15 @@ """Tests for scale-in request creation, engine selection, draining, removal, and cleanup.""" -from unittest.mock import MagicMock, patch +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch import pytest try: from relax.distributed.ray.rollout import ( + EngineGroupLifecycle, ScaleInRequest, ScaleInStatus, ScaleOutRequest, @@ -233,7 +235,8 @@ def test_num_replicas_no_removal_needed(self, patch_ray_get): ) assert manager._select_engines_for_removal(req, srv) == [] - def test_by_engine_urls(self, patch_ray_get): + @pytest.mark.asyncio + async def test_by_engine_urls(self, patch_ray_get): """Select engines matching specific URLs.""" e1 = make_mock_engine(url="http://a:1") e2 = make_mock_engine(url="http://b:2") @@ -246,10 +249,12 @@ def test_by_engine_urls(self, patch_ray_get): status=ScaleInStatus.PENDING, engine_urls=["http://a:1"], ) - infos = manager._select_engines_for_removal(req, srv) + candidates = await manager._resolve_scale_in_url_candidates(req, srv) + infos = manager._select_engines_for_removal(req, srv, url_candidates=candidates) assert len(infos) == 1 - def test_by_engine_urls_normalization(self, patch_ray_get): + @pytest.mark.asyncio + async def test_by_engine_urls_normalization(self, patch_ray_get): """URL normalization: http://host:port matches host:port.""" e1 = make_mock_engine(url="http://a:1") g = make_engine_group(engines=[e1], is_scaled_out=True) @@ -261,10 +266,12 @@ def test_by_engine_urls_normalization(self, patch_ray_get): status=ScaleInStatus.PENDING, engine_urls=["a:1"], ) - infos = manager._select_engines_for_removal(req, srv) + candidates = await manager._resolve_scale_in_url_candidates(req, srv) + infos = manager._select_engines_for_removal(req, srv, url_candidates=candidates) assert len(infos) == 1 - def test_dead_engines_skipped(self, patch_ray_get): + @pytest.mark.asyncio + async def test_dead_engines_skipped(self, patch_ray_get): """Dead (None) engines are not candidates.""" g = make_engine_group(engines=[None, None], is_scaled_out=True) srv = make_rollout_server(engine_groups=[g]) @@ -276,9 +283,94 @@ def test_dead_engines_skipped(self, patch_ray_get): num_replicas=0, engine_urls=["a:1"], ) - infos = manager._select_engines_for_removal(req, srv) + candidates = await manager._resolve_scale_in_url_candidates(req, srv) + infos = manager._select_engines_for_removal(req, srv, url_candidates=candidates) assert len(infos) == 0 + @pytest.mark.asyncio + async def test_url_probe_does_not_block_eviction_fence(self, patch_ray_get): + probe_started = asyncio.Event() + release_probe = asyncio.Event() + + async def blocked_url(): + probe_started.set() + await release_probe.wait() + return "http://elastic:1" + + engine = make_mock_engine(url="http://elastic:1") + engine.get_url.remote.return_value = blocked_url() + group = make_engine_group(engines=[engine], is_scaled_out=True) + group.pg = (MagicMock(), [], []) + srv = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": srv}) + req = ScaleInRequest( + request_id="r-url", + status=ScaleInStatus.PENDING, + engine_urls=["elastic:1"], + ) + manager._scale_in_requests[req.request_id] = req + + resolve_task = asyncio.create_task(manager._resolve_scale_in_url_candidates(req, srv)) + await asyncio.wait_for(probe_started.wait(), timeout=1) + + manager._handle_evictions([("default", group, 0)]) + + assert group.eviction_requested is True + assert group.lifecycle_status is EngineGroupLifecycle.DRAINING + release_probe.set() + assert await resolve_task == [(group, 0, engine)] + + @pytest.mark.asyncio + @pytest.mark.parametrize("mutation", ["replace_actor", "remove_group"]) + async def test_url_candidates_are_revalidated_before_claim(self, patch_ray_get, mutation): + engine = make_mock_engine(url="http://elastic:1") + group = make_engine_group(engines=[engine], is_scaled_out=True) + srv = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": srv}) + req = ScaleInRequest( + request_id="r-url", + status=ScaleInStatus.PENDING, + engine_urls=["elastic:1"], + ) + + candidates = await manager._resolve_scale_in_url_candidates(req, srv) + if mutation == "replace_actor": + group.all_engines[0] = make_mock_engine(url="http://elastic:1") + else: + srv.engine_groups.remove(group) + + assert manager._select_engines_for_removal(req, srv, url_candidates=candidates) == [] + + @pytest.mark.asyncio + async def test_url_probes_run_concurrently(self, patch_ray_get): + first_started = asyncio.Event() + second_started = asyncio.Event() + release_probes = asyncio.Event() + + async def blocked_url(started, url): + started.set() + await release_probes.wait() + return url + + first = make_mock_engine(url="http://first:1") + second = make_mock_engine(url="http://second:2") + first.get_url.remote.return_value = blocked_url(first_started, "http://first:1") + second.get_url.remote.return_value = blocked_url(second_started, "http://second:2") + group = make_engine_group(engines=[first, second], is_scaled_out=True) + srv = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": srv}) + req = ScaleInRequest( + request_id="r-url", + status=ScaleInStatus.PENDING, + engine_urls=["first:1", "second:2"], + ) + + resolve_task = asyncio.create_task(manager._resolve_scale_in_url_candidates(req, srv)) + await asyncio.wait_for(asyncio.gather(first_started.wait(), second_started.wait()), timeout=1) + release_probes.set() + + assert await resolve_task == [(group, 0, first), (group, 1, second)] + # ========================= _drain_engines ================================== @@ -335,6 +427,125 @@ async def test_force_skips_drain_wait(self): class TestScaleInExecution: + @pytest.mark.asyncio + @pytest.mark.parametrize("force", [False, True]) + async def test_scale_in_fence_rejects_update_started_during_removal(self, force, patch_ray_get): + engine = make_mock_engine(url="http://elastic:1") + group = make_engine_group(engines=[engine], is_scaled_out=True) + server = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": server}) + request = ScaleInRequest( + request_id="r-fence", + status=ScaleInStatus.PENDING, + engine_urls=["http://elastic:1"], + timeout_secs=5, + force=force, + ) + removal_started = asyncio.Event() + finish_removal = asyncio.Event() + + async def _remove(*_args, **_kwargs): + removal_started.set() + await finish_removal.wait() + return ["group_0_engine_0"], [] + + manager._remove_live_engines = _remove + scale_in_task = asyncio.create_task(manager._scale_in(request)) + await removal_started.wait() + + assert group.lifecycle_status is EngineGroupLifecycle.DRAINING + assert manager.set_weight_updating(True) is False + assert manager._training_weight_updating is False + + finish_removal.set() + await scale_in_task + + assert request.status is ScaleInStatus.COMPLETED + assert group.lifecycle_status is EngineGroupLifecycle.ACTIVE + + @pytest.mark.asyncio + @pytest.mark.parametrize("owner", ["_training_weight_updating", "_scale_out_weight_updating"]) + async def test_scale_in_fence_waits_for_existing_weight_update(self, owner, patch_ray_get): + engine = make_mock_engine(url="http://elastic:1") + group = make_engine_group(engines=[engine], is_scaled_out=True) + manager = create_test_manager(servers={"default": make_rollout_server(engine_groups=[group])}) + setattr(manager, owner, True) + manager._remove_live_engines = AsyncMock(return_value=(["engine"], [])) + request = ScaleInRequest( + request_id="r-existing-update", + status=ScaleInStatus.PENDING, + engine_urls=["http://elastic:1"], + timeout_secs=5, + ) + + scale_in_task = asyncio.create_task(manager._scale_in(request)) + while group.lifecycle_status is not EngineGroupLifecycle.DRAINING: + await asyncio.sleep(0) + + manager._remove_live_engines.assert_not_called() + assert manager.set_weight_updating(True) is False + + if owner == "_training_weight_updating": + assert manager.set_weight_updating(False) is True + else: + manager._scale_out_weight_updating = False + await scale_in_task + + manager._remove_live_engines.assert_awaited_once() + assert request.status is ScaleInStatus.COMPLETED + + @pytest.mark.asyncio + async def test_scale_in_fence_timeout_restores_active_without_cleanup(self, patch_ray_get): + engine = make_mock_engine(url="http://elastic:1") + group = make_engine_group(engines=[engine], is_scaled_out=True) + manager = create_test_manager(servers={"default": make_rollout_server(engine_groups=[group])}) + manager._training_weight_updating = True + manager._remove_live_engines = MagicMock() + request = ScaleInRequest( + request_id="r-timeout", + status=ScaleInStatus.PENDING, + engine_urls=["http://elastic:1"], + timeout_secs=0.01, + ) + + await manager._scale_in(request) + + assert request.status is ScaleInStatus.FAILED + assert "weight-update fence" in request.error_message + assert group.lifecycle_status is EngineGroupLifecycle.ACTIVE + manager._remove_live_engines.assert_not_called() + + @pytest.mark.asyncio + async def test_sigterm_during_explicit_scale_in_timeout_keeps_persistent_fence(self, patch_ray_get): + engine = make_mock_engine(url="http://elastic:1") + group = make_engine_group(engines=[engine], is_scaled_out=True) + group.pg = (MagicMock(), [], []) + manager = create_test_manager(servers={"default": make_rollout_server(engine_groups=[group])}) + manager._training_weight_updating = True + manager._remove_live_engines = MagicMock() + request = ScaleInRequest( + request_id="r-timeout-sigterm", + status=ScaleInStatus.PENDING, + engine_urls=["http://elastic:1"], + timeout_secs=0.01, + ) + manager._scale_in_requests[request.request_id] = request + + scale_in_task = asyncio.create_task(manager._scale_in(request)) + while group.lifecycle_status is not EngineGroupLifecycle.DRAINING: + await asyncio.sleep(0) + + manager._handle_evictions([("default", group, 0)]) + assert group.eviction_requested is True + + await scale_in_task + + assert request.status is ScaleInStatus.FAILED + assert group.lifecycle_status is EngineGroupLifecycle.DRAINING + assert group.eviction_requested is True + assert manager.set_weight_updating(True) is False + manager._remove_live_engines.assert_not_called() + @pytest.mark.asyncio async def test_force_removes_only_engines_unregistered_from_router(self, patch_ray_get): removable = make_mock_engine(url="http://a:1") @@ -365,6 +576,130 @@ async def test_force_removes_only_engines_unregistered_from_router(self, patch_r monitor.mark_intentionally_removed.assert_called_once_with(0) +@pytest.mark.asyncio +async def test_live_removal_orders_router_drain_dcs_shutdown_and_pg(patch_ray_get): + events = [] + engine = make_mock_engine() + engine.unregister_from_router.remote.side_effect = lambda **_kwargs: ( + events.append("router") or AwaitableValue(True) + ) + engine.unregister_dcs.remote.side_effect = lambda: events.append("dcs") or AwaitableValue(None) + engine.shutdown.remote.side_effect = lambda: events.append("shutdown") or AwaitableValue(None) + group = make_engine_group(engines=[engine], is_scaled_out=True) + group.pg = (MagicMock(), [], []) + server = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": server}) + + async def _sleep(_seconds): + events.append("drain") + + with ( + patch("relax.distributed.ray.rollout.asyncio.sleep", side_effect=_sleep), + patch( + "relax.distributed.ray.rollout.ray.util.remove_placement_group", + side_effect=lambda _pg: events.append("pg"), + ), + ): + removed, failed = await manager._remove_live_engines( + server, + [(group, 0)], + drain_timeout=1, + shutdown_timeout=2, + force=False, + ) + + assert removed == ["group_0_engine_0"] + assert failed == [] + assert events == ["router", "drain", "dcs", "shutdown", "pg"] + + +@pytest.mark.asyncio +async def test_batch_live_removal_runs_each_phase_concurrently_and_drains_once(patch_ray_get): + events = [] + + def concurrent_phase(name): + started = 0 + both_started = asyncio.Event() + + async def run(): + nonlocal started + started += 1 + events.append(f"{name}_start_{started}") + if started == 2: + both_started.set() + await both_started.wait() + events.append(f"{name}_end") + return True + + return run + + router_phase = concurrent_phase("router") + dcs_phase = concurrent_phase("dcs") + shutdown_phase = concurrent_phase("shutdown") + groups = [] + for rank_offset in (1, 2): + engine = make_mock_engine() + engine.unregister_from_router.remote.side_effect = lambda **_kwargs: router_phase() + engine.unregister_dcs.remote.side_effect = lambda: dcs_phase() + engine.shutdown.remote.side_effect = lambda: shutdown_phase() + group = make_engine_group(engines=[engine], is_scaled_out=True, rank_offset=rank_offset) + group.pg = (MagicMock(), [], []) + groups.append(group) + + server = make_rollout_server(engine_groups=groups) + manager = create_test_manager(servers={"default": server}) + + async def drain_once(_seconds): + events.append("drain") + + with ( + patch("relax.distributed.ray.rollout.asyncio.sleep", side_effect=drain_once) as sleep, + patch("relax.distributed.ray.rollout.ray.util.remove_placement_group") as remove_pg, + ): + removed, failed = await manager._remove_live_engines( + server, + [(groups[0], 0), (groups[1], 0)], + drain_timeout=1, + shutdown_timeout=2, + force=False, + ) + + assert removed == ["group_1_engine_0", "group_2_engine_0"] + assert failed == [] + sleep.assert_awaited_once_with(1) + assert events.count("drain") == 1 + assert events.index("router_start_2") < events.index("router_end") < events.index("drain") + assert events.index("drain") < events.index("dcs_start_1") + assert events.index("dcs_start_2") < events.index("dcs_end") + assert events.index("dcs_end") < events.index("shutdown_start_1") + assert events.index("shutdown_start_2") < events.index("shutdown_end") + assert remove_pg.call_count == 2 + assert server.engine_groups == [] + + +@pytest.mark.asyncio +async def test_dcs_failure_warns_and_continues_shutdown(patch_ray_get): + engine = make_mock_engine() + engine.unregister_dcs.remote.side_effect = RuntimeError("DCS unavailable") + group = make_engine_group(engines=[engine], is_scaled_out=True) + server = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": server}) + + removed, failed = await manager._remove_live_engines( + server, + [(group, 0)], + drain_timeout=1, + shutdown_timeout=2, + force=True, + ) + + assert removed == ["group_0_engine_0"] + assert failed == [] + engine.shutdown.remote.assert_called_once() + assert group.all_engines[0] is None + assert server.engine_groups == [] + + # ========================= _remove_engine ================================== @@ -392,6 +727,20 @@ async def test_fallback_to_ray_kill_on_shutdown_failure(self): mock_kill.assert_called_once_with(e1) assert g.all_engines[0] is None + @pytest.mark.asyncio + async def test_dcs_failure_still_shuts_down_engine(self): + engine = make_mock_engine() + engine.unregister_dcs.remote.side_effect = RuntimeError("DCS unavailable") + group = make_engine_group(engines=[engine], is_scaled_out=True) + manager = create_test_manager() + + with patch("ray.kill") as kill: + await manager._remove_engine(group, 0, shutdown_timeout=1) + + assert group.all_engines[0] is None + engine.shutdown.remote.assert_called_once() + kill.assert_not_called() + @pytest.mark.asyncio async def test_multi_node_engine_removal(self, patch_ray_get): """All sub-actors of a multi-node engine are removed.""" diff --git a/tests/distributed/ray/test_scale_out.py b/tests/distributed/ray/test_scale_out.py index d17d4b0ce..e71058c9d 100644 --- a/tests/distributed/ray/test_scale_out.py +++ b/tests/distributed/ray/test_scale_out.py @@ -29,6 +29,17 @@ pytestmark = pytest.mark.skipif(not HAS_DEPS, reason="Missing ray/sglang dependencies") +def test_engine_rank_allocator_never_reuses_removed_rank(): + group = make_engine_group(engines=[make_mock_engine()], rank_offset=4) + manager = create_test_manager(servers={"default": make_rollout_server(engine_groups=[group])}) + + first = manager._reserve_engine_ranks(1) + group.all_engines[0] = None + second = manager._reserve_engine_ranks(1) + + assert (first, second) == (5, 6) + + # ==================== create_scale_out_request ============================= diff --git a/tests/distributed/ray/test_utils.py b/tests/distributed/ray/test_utils.py index 98da2ecf7..89f4051f7 100644 --- a/tests/distributed/ray/test_utils.py +++ b/tests/distributed/ray/test_utils.py @@ -16,6 +16,7 @@ try: from relax.distributed.ray.rollout import ( + EngineGroupLifecycle, RolloutManager, ScaleOutRequest, ScaleOutStatus, @@ -350,3 +351,28 @@ def test_set_num_new_engines(self): srv.num_new_engines = 0 assert g1.num_new_engines == 0 assert g2.num_new_engines == 0 + + @pytest.mark.parametrize( + ("lifecycle_status", "expected_calls"), + [ + (EngineGroupLifecycle.ACTIVE, 1), + (EngineGroupLifecycle.DRAINING, 0), + (EngineGroupLifecycle.REMOVING, 0), + ], + ) + def test_recover_scaled_group_only_when_active(self, monkeypatch, lifecycle_status, expected_calls): + group = make_engine_group(is_scaled_out=True) + group.pg = (object(), [], []) + group.lifecycle_status = lifecycle_status + calls = [] + + def start_engines(port_cursors): + calls.append(port_cursors) + return [], port_cursors + + monkeypatch.setattr(group, "start_engines", start_engines) + server = make_rollout_server(engine_groups=[group]) + + server.recover() + + assert len(calls) == expected_calls diff --git a/tests/distributed/ray/test_weight_sync.py b/tests/distributed/ray/test_weight_sync.py index 63b280a69..212828bc7 100644 --- a/tests/distributed/ray/test_weight_sync.py +++ b/tests/distributed/ray/test_weight_sync.py @@ -495,6 +495,29 @@ async def test_lock_acquired_and_released(self, patch_async_helpers): manager._weight_sync_lock.acquire.remote.assert_called() manager._weight_sync_lock.release.remote.assert_called() + @pytest.mark.asyncio + async def test_scale_out_claims_lifecycle_owner_before_weight_sync_lock(self, patch_async_helpers): + seed = make_mock_engine(url="http://seed:1", weight_version="v1") + group = make_engine_group(engines=[seed]) + server = make_rollout_server(engine_groups=[group]) + manager = create_test_manager(servers={"default": server}) + + def acquire_lock(): + assert manager._scale_out_weight_updating is True + return AwaitableValue(True) + + manager._weight_sync_lock.acquire.remote.side_effect = acquire_lock + + ok = await manager._sync_weights_from_seed_engine( + [make_mock_engine()], + timeout=60, + model_name="default", + ) + + assert ok is True + assert manager._scale_out_weight_updating is False + manager._weight_sync_lock.release.remote.assert_called_once() + @pytest.mark.asyncio async def test_sync_failure_releases_weight_sync_lock(self, patch_async_helpers): """A sync that fails on every attempt still ALWAYS releases the shared From 773bd5879b2e607e8bf214ed9559b342bf9ac8a0 Mon Sep 17 00:00:00 2001 From: yangrui6 Date: Tue, 1 Sep 2026 15:19:50 +0800 Subject: [PATCH 06/34] fix(multimodal): bound Qwen-VL image ratios MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Normalize extreme Qwen-VL images before processing - Detect Transformers Qwen-VL image processors through their inner processor MRO - Expand only the short side when an image exceeds the hard 200 aspect-ratio limit - Preserve normal images and non-Qwen processor behavior --- # ✅ Tests ## Cover Qwen-VL ratio normalization - Test horizontal, vertical, boundary, and idempotent resize cases - Verify processor workers preserve image order and isolate non-Qwen processors (cherry picked from commit 25a3bf116981c6b69048eb21e737f8341c6fd5f7) --- relax/utils/data/processor_pool.py | 25 ++++++ relax/utils/multimodal/image_utils.py | 26 +++++++ tests/utils/data/test_processor_pool.py | 90 ++++++++++++++++++++++ tests/utils/multimodal/test_image_utils.py | 20 +++++ 4 files changed, 161 insertions(+) create mode 100644 tests/utils/data/test_processor_pool.py diff --git a/relax/utils/data/processor_pool.py b/relax/utils/data/processor_pool.py index 27f3baf14..fce2caf6c 100644 --- a/relax/utils/data/processor_pool.py +++ b/relax/utils/data/processor_pool.py @@ -35,6 +35,7 @@ remap_mm_train_inputs, ) from relax.utils.logging_utils import get_logger +from relax.utils.multimodal.image_utils import resize_qwen_vl_extreme_aspect_ratio logger = get_logger(__name__) @@ -51,6 +52,20 @@ def _init_worker(model_path: str, trust_remote_code: bool) -> None: logger.info(f"ProcessorPool worker initialized (pid={os.getpid()})") +def _is_qwen_vl_processor(processor: object) -> bool: + """Return whether ``processor`` uses a Transformers Qwen-VL image + processor.""" + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + return False + for cls in type(image_processor).__mro__: + if cls.__module__.startswith( + ("transformers.models.qwen2_vl.", "transformers.models.qwen3_vl.") + ) and cls.__name__.startswith(("Qwen2VLImageProcessor", "Qwen3VLImageProcessor")): + return True + return False + + def prepare_mm_inputs_for_ipc(multimodal_inputs: dict) -> dict: """Prepare multimodal inputs for efficient cross-process transfer. @@ -104,6 +119,16 @@ def process_sample_in_worker( restored = dict(multimodal_inputs) if images := restored.get("images"): restored["images"] = [Image.fromarray(arr) for arr in images] + if _is_qwen_vl_processor(_worker_processor): + resized_images = [] + for image in restored["images"]: + resized = resize_qwen_vl_extreme_aspect_ratio(image) + if resized is not image: + logger.warning( + f"Qwen-VL image aspect ratio exceeded 200; resized from {image.size} to {resized.size}." + ) + resized_images.append(resized) + restored["images"] = resized_images # Videos arrive as shared-memory torch.Tensors — usable directly by the processor. # Audio arrives as numpy arrays — usable directly by the processor. diff --git a/relax/utils/multimodal/image_utils.py b/relax/utils/multimodal/image_utils.py index dccbaa6d5..e6f36a277 100644 --- a/relax/utils/multimodal/image_utils.py +++ b/relax/utils/multimodal/image_utils.py @@ -18,6 +18,8 @@ SPATIAL_MERGE_SIZE = 2 +QWEN_VL_MAX_ASPECT_RATIO = 200 +QWEN_VL_SAFE_ASPECT_RATIO = 199 ImageInput = Union[ @@ -28,6 +30,30 @@ ] +def resize_qwen_vl_extreme_aspect_ratio(image: Image.Image) -> Image.Image: + """Resize images that exceed Qwen-VL's hard aspect-ratio limit. + + Qwen-VL processors reject images whose long-to-short-side ratio is greater + than 200 before their normal pixel and patch-alignment resize runs. Expand + only the short side to keep the long-side pixels intact; Qwen's processor + remains responsible for its checkpoint-specific pixel and patch limits. + """ + height, width = image.height, image.width + short_side = max(min(height, width), 1) + aspect_ratio = max(height, width) / short_side + if aspect_ratio <= QWEN_VL_MAX_ASPECT_RATIO: + return image + + target_short_side = math.ceil(max(height, width) / QWEN_VL_SAFE_ASPECT_RATIO) + if height > width: + target_height = height + target_width = target_short_side + else: + target_width = width + target_height = target_short_side + return image.resize((target_width, target_height), Image.LANCZOS) + + def get_resize_height_width( max_ratio: Optional[float], height: int, diff --git a/tests/utils/data/test_processor_pool.py b/tests/utils/data/test_processor_pool.py new file mode 100644 index 000000000..879b7f563 --- /dev/null +++ b/tests/utils/data/test_processor_pool.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Tests for processor worker multimodal preprocessing.""" + +import numpy as np +from PIL import Image + +from relax.utils.data import processor_pool + + +class Qwen2VLImageProcessor: + pass + + +Qwen2VLImageProcessor.__module__ = "transformers.models.qwen2_vl.image_processing_qwen2_vl" + + +class CustomQwen2VLImageProcessor(Qwen2VLImageProcessor): + pass + + +class ProcessorWithQwenImageProcessor: + def __init__(self) -> None: + self.image_processor = CustomQwen2VLImageProcessor() + self.received_sizes: list[tuple[int, int]] = [] + + def __call__(self, **kwargs): + self.received_sizes = [image.size for image in kwargs["images"]] + return {"input_ids": [[1]]} + + +class OtherProcessor: + def __init__(self) -> None: + self.received_image: Image.Image | None = None + + def __call__(self, **kwargs): + self.received_image = kwargs["images"][0] + return {"input_ids": [[1]]} + + +class SameNamedNonTransformersImageProcessor: + pass + + +SameNamedNonTransformersImageProcessor.__name__ = "Qwen2VLImageProcessor" + + +class ProcessorWithSameNamedNonTransformersImageProcessor: + image_processor = SameNamedNonTransformersImageProcessor() + + +def test_process_sample_resizes_only_extreme_images_for_qwen_vl(monkeypatch): + processor = ProcessorWithQwenImageProcessor() + monkeypatch.setattr(processor_pool, "_worker_processor", processor) + + prompt_ids, train_inputs = processor_pool.process_sample_in_worker( + "prompt", + { + "images": [ + np.asarray(Image.new("RGB", (750, 1))), + np.asarray(Image.new("RGB", (100, 50))), + np.asarray(Image.new("RGB", (1, 750))), + ] + }, + {}, + ) + + assert prompt_ids == [1] + assert train_inputs is None + assert processor.received_sizes == [(750, 4), (100, 50), (4, 750)] + + +def test_process_sample_does_not_resize_extreme_image_for_other_processors(monkeypatch): + processor = OtherProcessor() + monkeypatch.setattr(processor_pool, "_worker_processor", processor) + + prompt_ids, train_inputs = processor_pool.process_sample_in_worker( + "prompt", {"images": [np.asarray(Image.new("RGB", (750, 1)))]}, {} + ) + + assert prompt_ids == [1] + assert train_inputs is None + assert processor.received_image is not None + assert processor.received_image.size == (750, 1) + + +def test_qwen_vl_detection_rejects_same_class_name_from_other_module(): + processor = ProcessorWithSameNamedNonTransformersImageProcessor() + + assert not processor_pool._is_qwen_vl_processor(processor) diff --git a/tests/utils/multimodal/test_image_utils.py b/tests/utils/multimodal/test_image_utils.py index 161b92516..f941dd226 100644 --- a/tests/utils/multimodal/test_image_utils.py +++ b/tests/utils/multimodal/test_image_utils.py @@ -11,6 +11,7 @@ get_resize_height_width, image_smart_resize, load_image, + resize_qwen_vl_extreme_aspect_ratio, to_rgb, ) @@ -72,6 +73,25 @@ def test_image_smart_resize_preserves_mode_and_patch_alignment(): assert 28 * 28 <= resized.width * resized.height <= 4 * 28 * 28 +@pytest.mark.parametrize("size", [(750, 1), (1, 750), (201, 1), (1, 201)]) +def test_resize_qwen_vl_extreme_aspect_ratio_makes_image_safe(size): + image = Image.new("RGB", size, (12, 34, 56)) + + resized = resize_qwen_vl_extreme_aspect_ratio(image) + + assert max(resized.size) / min(resized.size) < 200 + assert resized.mode == image.mode + assert resized.width >= resized.height if image.width >= image.height else resized.height >= resized.width + assert resize_qwen_vl_extreme_aspect_ratio(resized) is resized + + +@pytest.mark.parametrize("size", [(200, 1), (1, 200), (100, 100)]) +def test_resize_qwen_vl_extreme_aspect_ratio_leaves_valid_image_unchanged(size): + image = Image.new("RGB", size, (12, 34, 56)) + + assert resize_qwen_vl_extreme_aspect_ratio(image) is image + + def test_to_rgb_composites_rgba_over_white(): image = Image.new("RGBA", (1, 1), (255, 0, 0, 128)) From e97088adf0da787e56e84f1caa1b6409946602b7 Mon Sep 17 00:00:00 2001 From: pojun Date: Wed, 2 Sep 2026 14:19:00 +0800 Subject: [PATCH 07/34] fix(megatron-patch): restore MTP detach routing dropped in mcore upgrade The mcore upgrade in 2d5f1cd1 (20260506-85bced0ae -> 20260728-0e6ac576f) dropped the whole multi_token_prediction.py block. Most of it was correct to drop -- upstream had absorbed it -- but two capabilities regressed. 1. Three-way detach routing became a silent no-op. Upstream only has a single coarse `mtp_detach_heads`, and it defaults to False (transformer_config.py:88). Relax's configure_mtp_detach_paths (relax/backends/megatron/model_provider.py:55) still setattr's mtp_detach_embedding / _backbone / _lm_head, which nothing reads anymore. So --mtp-detach-paths (default: detach all three) and --mtp-only-training silently did the opposite of what they declare: the MTP auxiliary loss backpropagated into embedding, backbone and lm_head. No error, args still accepted, logs still emitted. Re-routes upstream's three detach sites to the per-path flags, keeping upstream's own implementations: - process_mtp_loss -> mtp_detach_lm_head - _get_embeddings -> mtp_detach_embedding - MultiTokenPredictionBlock -> mtp_detach_backbone The backbone site also regains `offset == 0` and `.requires_grad_(True)` from the old patch. Under VPP, offset > 0 means hidden_states came from a previous MTP stage rather than the main backbone, so upstream's unconditional detach severs gradient flow between MTP layers. 2. Logged MTP loss was missing the scaling factor. mtp_loss_scale is defined after save_loss_to_tracker upstream. Moved it ahead and applied it to the logged sum. Verified save_loss_to_tracker (multi_token_prediction.py:461) only normalizes -- it absorbed the old patch's safe-divide but never applied the scaling factor -- so there is no double scaling on either the per-token or microbatch-normalized branch. Note this shifts MTP loss curves by mtp_loss_scaling_factor (default 0.2) relative to runs on the current image. Deliberately NOT restored, having verified each landed upstream: the functional_call-based lm-head detach (upstream detaches output_weight directly), the _checkpointed_forward non-tensor rewrite (upstream captures via closure), the labels-is-None early return (upstream derives labels from input_ids), bridge/peft/utils.py create_peft, tensor_parallel/layers.py dgrad fold, qwen35_vl_bridge.py packed MTP experts, and the yarn position-embedding choice. Also left alone: MultiTokenPredictionBlock.__init__'s grad_norm_group tag, which is an upstream addition the old patch never had. tests/backends/megatron/test_mtp_only_training.py: 37 passed (was 36 passed, 1 failed). test_frozen_weight_dgrad.py::test_megatron_patch_carries_dgrad_fold still fails -- that one is a stale guard for a fix now shipped by upstream and is tracked separately. Co-Authored-By: Claude (cherry picked from commit ec45bc8e08bbedc637609aa78694d3637c99492a) --- .../patch/megatron/20260728-0e6ac576f.patch | 66 +++++++++++++++++++ .../megatron/test_frozen_weight_dgrad.py | 24 ++----- 2 files changed, 73 insertions(+), 17 deletions(-) diff --git a/docker/patch/megatron/20260728-0e6ac576f.patch b/docker/patch/megatron/20260728-0e6ac576f.patch index de0823ca1..750df19c9 100644 --- a/docker/patch/megatron/20260728-0e6ac576f.patch +++ b/docker/patch/megatron/20260728-0e6ac576f.patch @@ -1284,6 +1284,72 @@ index 2796bc6..71e036e 100644 def _maintain_float32_expert_bias(self): """ Maintain the expert bias in float32. +diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py +--- a/megatron/core/transformer/multi_token_prediction.py ++++ b/megatron/core/transformer/multi_token_prediction.py +@@ -1031,7 +1031,7 @@ + ) + derived_labels_from_input_ids = True + +- if config.mtp_detach_heads: ++ if getattr(config, "mtp_detach_lm_head", True): + if output_weight is not None: + output_weight = output_weight.detach() + else: +@@ -1084,6 +1084,7 @@ + mtp_logits = scale_logits_fn(mtp_logits) + mtp_loss = compute_language_model_loss(mtp_labels, mtp_logits) + mtp_loss = loss_mask * mtp_loss ++ mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers + + if is_training: + if mtp_logits is not None: +@@ -1097,7 +1098,7 @@ + total = torch.zeros((), device=mtp_loss.device, dtype=mtp_loss.dtype) + + MTPLossLoggingHelper.save_loss_to_tracker( +- torch.sum(mtp_loss), ++ mtp_loss_scale * torch.sum(mtp_loss), + num_tokens, + mtp_layer_number, + config.mtp_num_layers, +@@ -1106,7 +1107,6 @@ + avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True), + calculate_per_token_loss=config.calculate_per_token_loss, + ) +- mtp_loss_scale = config.mtp_loss_scaling_factor / config.mtp_num_layers + if config.calculate_per_token_loss: + # When calculate_per_token_loss is enabled, finalize_model_grads will + # divide all gradients by total_num_tokens (from main loss). +@@ -1383,12 +1383,12 @@ + # embedding + decoder_input = embedding(input_ids=input_ids, position_ids=position_ids) + +- if self.config.mtp_detach_heads: ++ if getattr(self.config, "mtp_detach_embedding", True): + decoder_input = decoder_input.detach() + + hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + # make_viewless_tensor no-ops when hidden_states is not a view (_base is None), +- # which happens after detach() with mtp_detach_heads. Activation ++ # which happens after detach() with mtp_detach_embedding. Activation + # checkpointing (CheckpointFunction.apply) requires at least one input tensor + # with requires_grad=True to produce a differentiable output, so we ensure it + # here to maintain gradient flow to MTP layer parameters. +@@ -2167,8 +2167,11 @@ + else: + hidden_states = hidden_states_list[offset] + +- if self.config.mtp_detach_heads: +- hidden_states = hidden_states.detach() ++ if offset == 0 and getattr(self.config, "mtp_detach_backbone", True): ++ # offset > 0 means hidden_states came from a previous MTP stage under VPP, ++ # not from the main backbone -- detaching there would sever gradient flow ++ # between MTP layers. ++ hidden_states = hidden_states.detach().requires_grad_(True) + + for iteration in range(self.config.mtp_num_layers): + layer_idx = 0 if self.mtp_use_repeated_layer else iteration diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 43e45b7..8166f94 100644 --- a/megatron/core/transformer/transformer_config.py diff --git a/tests/backends/megatron/test_frozen_weight_dgrad.py b/tests/backends/megatron/test_frozen_weight_dgrad.py index 17e1d79d0..670944e13 100644 --- a/tests/backends/megatron/test_frozen_weight_dgrad.py +++ b/tests/backends/megatron/test_frozen_weight_dgrad.py @@ -1,11 +1,11 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. -"""Alarm for the frozen-weight DGRAD fold in the Megatron image patch. +"""Alarm for the frozen-weight DGRAD fold in Megatron. -``docker/patch/megatron/20260506-85bced0ae.patch`` backports the -``LinearWithFrozenWeight.backward`` hunk from NVIDIA/Megatron-LM#5092; a -Megatron bump could drop it silently, so the first test guards the patch text -and runs in CI while the rest check the behaviour and need Megatron. +NVIDIA/Megatron-LM#5092 made ``LinearWithFrozenWeight.backward`` reshape so +that size-1 leading dims fold into a single ``mm``. mcore has carried it +upstream since 0e6ac576f, so these tests check the behaviour directly and need +Megatron on the path. """ import inspect @@ -23,8 +23,7 @@ needs_megatron = pytest.mark.skipif(LinearWithFrozenWeight is None, reason="requires Megatron") -PATCH = Path(__file__).resolve().parents[3] / "docker" / "patch" / "latest" / "megatron.patch" -HINT = f"the LinearWithFrozenWeight.backward hunk from NVIDIA/Megatron-LM#5092 is missing from {PATCH}" +HINT = "the LinearWithFrozenWeight.backward fold from NVIDIA/Megatron-LM#5092 is missing from Megatron" FOLD = "grad_output.reshape(-1, grad_output.size(-1))" @@ -44,19 +43,10 @@ def _megatron_lacks_dgrad_fold() -> bool: needs_fold = pytest.mark.skipif( _megatron_lacks_dgrad_fold(), - reason=f"the Megatron on the path predates the DGRAD fold; rebuild the image with {PATCH}", + reason="the Megatron on the path predates the DGRAD fold; rebuild the image with a newer mcore", ) -def test_megatron_patch_carries_dgrad_fold(): - """Runs without Megatron installed, so CI covers it.""" - assert PATCH.is_file(), HINT - patch = PATCH.read_text() - assert "megatron/core/tensor_parallel/layers.py" in patch, HINT - assert FOLD in patch, HINT - assert "grad_input.reshape(*grad_output.shape[:-1], weight.size(1))" in patch, HINT - - class _MatmulSpy(TorchDispatchMode): """Records which matmul-family aten op a block dispatches to.""" From 4457b1630f3e222b7ba2fc7e2b5882f8e3eb55d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=A8=E7=9D=BF?= Date: Wed, 2 Sep 2026 14:49:05 +0800 Subject: [PATCH 08/34] feat(sft): add async prepack pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ⭐ Feature ## Add opt-in asynchronous SFT prepacking - Add --sft-async-prepack to offload complete TransferQueue shard fetches, sequence-length balancing, CPU THD packing, pinned-memory staging, and H2D copies from the training thread - Preserve the standard SFT buffer-level partitioning, loss scaling, and oversized-sample behavior while reconciling the DP-wide micro-batch count before forward/backward - Prefetch one future rollout window with explicit identity and error propagation, and overlap its first H2D copy on a dedicated stream - Add PrepackedBatch and recursive move, pin, and record-stream helpers so the training path can consume prepared batches without repeating packing work - Validate the supported dynamic-batch, per-rank-fetch, THD, PP, CP, VPP, and routing-replay configuration and auto-enable the lookahead depth required for overlap - Add a Qwen3 0.6B recipe and enable async prepacking in the Qwen3-VL Pokemon recipe --- # 🐛 Bug Fix ## Make distributed sampling and SFT execution robust - Attach total_lengths custom metadata to SFT partitions and validate the producer contract before token-budget sampling - Wait for an equal-size DP-ready round before balancing streaming samples, preventing sequence partition assertions when fewer samples than DP ranks are available - Keep async SFT prepacking on SeqlenBalancedSampler and lazy-load the streaming sampler only for fully-async dynamic batching - Repartition real samples to the DP-wide maximum micro-batch count instead of executing expensive full-batch dummy work - Synchronize prefetch errors across ranks, fail fast on pinned-memory failures, and keep copy-stream tensor lifetimes valid - Weight SFT rollout metrics by local sample count and create collective statistics on the active training device - Cache the Megatron main-rank role before SFT prediction tears down reloadable process groups --- # ⚡ Performance ## Reduce SFT input-pipeline stalls - Move CPU packing and pinned-memory preparation off the training thread - Pipeline first-batch and subsequent H2D copies without racing the training copy stream - Avoid streaming tail polling for SFT by fetching one complete balanced local shard per window --- # ✅ Tests ## Cover sampler and SFT data contracts - Add equal-size DP-round and missing-custom-metadata regression tests for the streaming sampler - Verify SFT producers publish sequence lengths and collective statistics use the training device - Pass 24 targeted SFT, sampler, iterator, component, and Megatron data tests - Pass pre-commit run --all-files --show-diff-on-failure (cherry picked from commit d5d2b1b567931135cf9891d89940ac94236baf73) --- relax/backends/megatron/actor.py | 590 +++++++++++++++++- relax/backends/megatron/data.py | 126 +++- relax/components/sft.py | 2 + relax/core/controller.py | 7 +- relax/engine/sft/predict/runner.py | 9 +- relax/utils/arguments.py | 47 +- relax/utils/data/micro_batch_ring.py | 139 +++++ relax/utils/data/stream_dataloader.py | 22 +- .../training/sft/run-qwen3-0.6B-math-8xgpu.sh | 132 ++++ .../sft/run-qwen3-vl-4B-pokemon-8xgpu.sh | 1 + tests/backends/megatron/test_data_vpp.py | 50 ++ tests/backends/megatron/test_sft_prepack.py | 58 ++ tests/components/test_sft.py | 1 + .../utils/data/test_streaming_tq_iterator.py | 51 ++ tests/utils/test_arguments_sft.py | 37 ++ 15 files changed, 1224 insertions(+), 48 deletions(-) create mode 100644 relax/utils/data/micro_batch_ring.py create mode 100755 scripts/training/sft/run-qwen3-0.6B-math-8xgpu.sh create mode 100644 tests/backends/megatron/test_sft_prepack.py create mode 100644 tests/utils/test_arguments_sft.py diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index 0e93d91fa..f493793a3 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -41,6 +41,13 @@ from relax.utils import device as device_utils from relax.utils import tracking_utils from relax.utils.async_utils import run +from relax.utils.data.data import get_minimum_num_micro_batch_size +from relax.utils.data.micro_batch_ring import ( + PrefetchedSFTWindow, + SFTWindowPrefetcher, + is_sft_async_prepack_enabled, +) +from relax.utils.data.seqlen_balancing import get_seqlen_balanced_partitions from relax.utils.data.stream_dataloader import ( MicroBatchListIterator, StreamingTQIterator, @@ -93,12 +100,16 @@ ROLLOUT_MINI_GLOBAL_SAMPLE_COUNTS_KEY, ROLLOUT_MINI_LOCAL_SAMPLE_COUNTS_KEY, DataIterator, + PrepackedBatch, build_rollout_minibatch_plan, concat_rollout_batches, get_data_iterator, log_perf_data, log_perf_data_fwd, log_rollout_data, + move_tensors_to_device, + prepack_sft_micro_batch_cpu, + record_tensors_on_stream, ) from .initialize import init, is_megatron_main_rank from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values @@ -174,6 +185,114 @@ def _slice(value): return {k: _slice(v) for k, v in rollout_data.items()} +def _select_rollout_samples(rollout_data: RolloutBatch, indices: list[int]) -> RolloutBatch: + """Slice a rollout batch to the given per-sample indices.""" + num_samples = len(rollout_data["total_lengths"]) + + def _select(value): + if isinstance(value, list) and len(value) == num_samples: + return [value[index] for index in indices] + if isinstance(value, tuple) and len(value) == num_samples: + return tuple(value[index] for index in indices) + if isinstance(value, torch.Tensor) and value.ndim > 0 and value.size(0) == num_samples: + return value[indices] + return value + + return {key: _select(value) for key, value in rollout_data.items()} + + +class _SFTPrefetchMiss(RuntimeError): + """Background prefetch did not find data within its local retry budget.""" + + +def _raise_if_sft_peer_error( + local_error: BaseException | None, + *, + phase: str, + device: torch.device, + tp_group: Any, + dp_group: Any, +) -> None: + """Propagate a rank-local prepack failure across the TP x DP mesh.""" + error_flag = torch.tensor([1 if local_error is not None else 0], dtype=torch.int, device=device) + dist.all_reduce(error_flag, op=dist.ReduceOp.MAX, group=tp_group) + dist.all_reduce(error_flag, op=dist.ReduceOp.MAX, group=dp_group) + if int(error_flag.item()) == 0: + return + if local_error is not None: + raise local_error + raise RuntimeError(f"SFT prepack {phase} failed on a peer rank; aborting to avoid a distributed hang.") + + +def _should_pause_sft_prepack_lookahead(args: Namespace, rollout_id: int) -> bool: + """Avoid retaining the next prefetched window across memory-heavy step + boundaries.""" + next_rollout_id = rollout_id + 1 + is_train_done = next_rollout_id == args.num_rollout + should_save_after_step = args.save is not None and ( + args.rotate_ckpt + or (args.save_interval is not None and (next_rollout_id % args.save_interval == 0 or is_train_done)) + ) + return should_run_sft_eval(args, rollout_id) or should_run_sft_predict(args, rollout_id) or should_save_after_step + + +class _SFTPrepackedDeviceIterator: + """Two-slot H2D pipeline over fully packed pinned-CPU micro-batches. + + All H2D copies run on ``copy_stream``; the training thread waits on a per- + batch ready event before consuming a batch, and calls ``record_stream`` so + the caching allocator does not reclaim the pinned source before the + consumer stream finishes. + """ + + def __init__( + self, + packed_cpu: list[tuple[PrepackedBatch, Any]], + first_device_micro_batch: tuple[PrepackedBatch, Any], + first_ready_event: Any, + copy_stream: Any, + device: torch.device, + ) -> None: + self._packed_cpu = packed_cpu + self._next_device_micro_batch: tuple[PrepackedBatch, Any] | None = first_device_micro_batch + self._next_ready_event = first_ready_event + self._copy_stream = copy_stream + self._device = device + self._offset = 0 + + def __iter__(self) -> "_SFTPrepackedDeviceIterator": + return self + + def __next__(self) -> tuple[PrepackedBatch, Any]: + if self._offset >= len(self._packed_cpu): + raise StopIteration + + assert self._next_device_micro_batch is not None + current_batch, current_meta = self._next_device_micro_batch + train_stream = device_utils.current_stream(self._device) + train_stream.wait_event(self._next_ready_event) + record_tensors_on_stream(current_batch, train_stream) + + self._offset += 1 + if self._offset < len(self._packed_cpu): + next_cpu_batch, next_meta = self._packed_cpu[self._offset] + with device_utils.stream_context(self._copy_stream): + next_device_batch = move_tensors_to_device(next_cpu_batch, self._device, non_blocking=True) + next_ready_event = device_utils.Event() + next_ready_event.record(self._copy_stream) + # move_tensors_to_device returns a plain dict; re-wrap so + # get_batch() short-circuits on the PrepackedBatch marker. + self._next_device_micro_batch = (PrepackedBatch(next_device_batch), next_meta) + self._next_ready_event = next_ready_event + + return current_batch, current_meta + + def close(self) -> None: + self._packed_cpu = [] + self._next_device_micro_batch = None + self._next_ready_event = None + + class MegatronTrainRayActor(TrainRayActor): @property def _per_step_rollout(self) -> bool: @@ -424,6 +543,22 @@ def _init( self.prof.on_init_end() self.data_iterator = None + self._sft_window_prefetcher = SFTWindowPrefetcher() if is_sft_async_prepack_enabled(self.args) else None + self._sft_prefetch_stream = None + self._sft_copy_stream = None + self._sft_device = None + if self._sft_window_prefetcher is not None: + self._validate_sft_prepack_pipeline() + self._sft_device = device_utils.make_current_torch_device() + # Two streams, one owner each: _sft_prefetch_stream is written to + # by the background prefetch worker for the first H2D of the next + # window; _sft_copy_stream is written to by the training thread + # for the overlapped H2D of the current window's later batches. + # Sharing one stream would let both host threads race on the + # enqueue order and let the next window's copy jump ahead of the + # current step's remaining copies. + self._sft_prefetch_stream = device_utils.Stream(device=self._sft_device) + self._sft_copy_stream = device_utils.Stream(device=self._sft_device) if role == "actor": capture_hooks.maybe_enable_for_actor() @@ -487,6 +622,28 @@ def _zero_expert_lora_adapters(self) -> None: zeroed += 1 logger.info("[LoRA] zeroed %d grouped-expert adapter params on role=%s", zeroed, self.role) + def _validate_sft_prepack_pipeline(self) -> None: + unsupported = [] + if mpu.get_pipeline_model_parallel_world_size() != 1: + unsupported.append("pipeline parallelism") + if mpu.get_context_parallel_world_size() != 1: + unsupported.append("context parallelism") + if (mpu.get_virtual_pipeline_model_parallel_world_size() or 1) != 1: + unsupported.append("virtual pipeline parallelism") + if getattr(self.args, "dynamic_context_parallel", False): + unsupported.append("dynamic context parallelism") + if getattr(self.args, "calculate_per_token_loss", False): + unsupported.append("per-token loss") + if device_utils.get_device_name() != "cuda": + unsupported.append(f"device={device_utils.get_device_name()}") + if self.args.qkv_format != "thd": + unsupported.append(f"qkv_format={self.args.qkv_format}") + if unsupported: + raise ValueError( + "--sft-async-prepack currently supports THD with PP=1, CP=1 and VPP disabled; " + f"unsupported: {', '.join(unsupported)}." + ) + @timer def sleep(self) -> None: assert self.args.offload_train @@ -792,10 +949,23 @@ def train(self, rollout_id: int) -> None: task_name = f"{base_task_name}_critic" else: task_name = base_task_name + if is_sft_mode(self.args) and self._sft_window_prefetcher is not None: + data_fields = build_data_fields(self.args, consumer="actor") + rollout_data, prepared_iterator, num_microbatches = self._get_prefetched_sft_window( + task_name, + rollout_id, + data_fields, + ) + return self.train_actor( + rollout_id, + rollout_data, + prepared_data_iterator=[prepared_iterator], + prepared_num_microbatches=[num_microbatches], + ) + empty_poll_sleep_s = Envs.RELAX_EMPTY_POLL_SLEEP_MS / 1000.0 rollout_mini_batches: list[RolloutBatch] = [] rollout_mini_batch_metas: list = [] rollout_mini_local_sample_counts: list[int] = [] - empty_poll_sleep_s = Envs.RELAX_EMPTY_POLL_SLEEP_MS / 1000.0 fetch_iter = 0 while batch_index < num_rollout_minis and not self.all_consumed(task_name, rollout_id): consumer = "critic" if self.role == "critic" else "actor" @@ -841,6 +1011,382 @@ def train(self, rollout_id: int) -> None: else: return self.train_actor(rollout_id, rollout_data) + def _get_prefetched_sft_window( + self, + task_name: str, + rollout_id: int, + data_fields: list[str], + ) -> tuple[RolloutBatch, _SFTPrepackedDeviceIterator, int]: + assert self._sft_window_prefetcher is not None + identity = (task_name, tuple(data_fields)) + self._sft_window_prefetcher.prefetch( + rollout_id, + partial(self._fetch_and_prepack_sft_window, task_name, rollout_id, data_fields), + identity=identity, + ) + # Capture prefetch failure locally so all ranks can synchronize on + # error status before touching any collective. A background TQ / + # packing / pin failure on one rank must not leave peer ranks stuck + # in the first all_reduce below. + device = device_utils.make_current_torch_device() + tp_group = mpu.get_tensor_model_parallel_group() + dp_group = mpu.get_data_parallel_group(with_context_parallel=False) + local_error: BaseException | None = None + window = None + with timer("sft_prefetch_wait"): + try: + window = self._sft_window_prefetcher.get(rollout_id) + except _SFTPrefetchMiss: + window = None + except BaseException as exc: # noqa: BLE001 + local_error = exc + + window = self._agree_sft_prefetch_window(task_name, rollout_id, data_fields, window, local_error) + + local_k = len(window.packed_micro_batches) + local_samples = len(window.rollout_data["total_lengths"]) + + # Run every collective unconditionally, aggregate failure into a + # single global bit, then raise together. Raising between collectives + # would leave peers stuck in the next all_reduce. + # tp_bounds: [max_k, -min_k, max_samples, -min_samples] MAX-reduced + # across TP replicas. Only tp_max_k / TP sample-count agreement are + # used downstream; tp_min_k folds into DP-wide MAX below. + tp_bounds = torch.tensor([local_k, -local_k, local_samples, -local_samples], dtype=torch.int, device=device) + dist.all_reduce(tp_bounds, op=dist.ReduceOp.MAX, group=tp_group) + tp_max_k = int(tp_bounds[0].item()) + tp_max_samples, tp_min_samples = int(tp_bounds[2].item()), -int(tp_bounds[3].item()) + + dp_k_bounds = torch.tensor([tp_max_k], dtype=torch.int, device=device) + dist.all_reduce(dp_k_bounds, op=dist.ReduceOp.MAX, group=dp_group) + max_k = int(dp_k_bounds[0].item()) + + global_samples = torch.tensor([local_samples], dtype=torch.int, device=device) + dist.all_reduce(global_samples, op=dist.ReduceOp.SUM, group=dp_group) + total_samples = int(global_samples.item()) + + errors: list[str] = [] + if tp_max_samples != tp_min_samples: + errors.append(f"real sample count differs across TP ranks (min={tp_min_samples}, max={tp_max_samples})") + if total_samples != self.args.global_batch_size: + errors.append( + f"sample count mismatch (expected global_batch_size={self.args.global_batch_size}, got {total_samples})" + ) + validation_error = None + if errors: + validation_error = RuntimeError( + f"SFT prefetch validation failed for rollout_id={rollout_id}: " + "; ".join(errors) + "." + ) + _raise_if_sft_peer_error( + validation_error, + phase=f"validation for rollout_id={rollout_id}", + device=device, + tp_group=tp_group, + dp_group=dp_group, + ) + + packed_micro_batches = window.packed_micro_batches + first_device_micro_batch = window.first_device_micro_batch + first_ready_event = window.first_ready_event + + # Match get_data_iterator: first agree on DP-wide K, then repartition + # every rank's real samples into that K. Repeating a complete local + # micro-batch as a zero-loss dummy still runs its full forward/backward + # and made long-sequence SFT substantially slower. + repack_error: BaseException | None = None + if local_k < max_k: + try: + first_ready_event.synchronize() + micro_batch_indices = get_seqlen_balanced_partitions( + window.rollout_data["total_lengths"], max_k, equal_size=False + ) + packed_micro_batches = [ + ( + prepack_sft_micro_batch_cpu( + self.args, + _select_rollout_samples(window.rollout_data, indices), + ), + None, + ) + for indices in micro_batch_indices + ] + + assert self._sft_copy_stream is not None + assert self._sft_device is not None + with device_utils.stream_context(self._sft_copy_stream): + first_cpu_batch, first_meta = packed_micro_batches[0] + first_device_batch = PrepackedBatch( + move_tensors_to_device(first_cpu_batch, self._sft_device, non_blocking=True) + ) + first_ready_event = device_utils.Event() + first_ready_event.record(self._sft_copy_stream) + first_device_micro_batch = (first_device_batch, first_meta) + except BaseException as exc: # noqa: BLE001 + repack_error = exc + _raise_if_sft_peer_error( + repack_error, + phase=f"repack for rollout_id={rollout_id}", + device=device, + tp_group=tp_group, + dp_group=dp_group, + ) + local_k = max_k + + next_rollout_id = rollout_id + 1 + should_pause_lookahead = _should_pause_sft_prepack_lookahead(self.args, rollout_id) + if next_rollout_id < self.args.num_rollout and self.args.max_staleness >= 1 and not should_pause_lookahead: + self._sft_window_prefetcher.prefetch( + next_rollout_id, + partial(self._fetch_and_prepack_sft_window, task_name, next_rollout_id, data_fields), + identity=(task_name, tuple(data_fields)), + ) + + assert self._sft_copy_stream is not None + assert self._sft_device is not None + iterator = _SFTPrepackedDeviceIterator( + packed_micro_batches, + first_device_micro_batch, + first_ready_event, + self._sft_copy_stream, + self._sft_device, + ) + return window.rollout_data, iterator, local_k + + def _sft_prepack_local_batch_size(self) -> int: + dp_size = mpu.get_data_parallel_world_size(with_context_parallel=False) + if self.args.global_batch_size % dp_size != 0: + raise ValueError( + "--sft-async-prepack requires global_batch_size divisible by data parallel size, " + f"got global_batch_size={self.args.global_batch_size}, dp_size={dp_size}." + ) + return self.args.global_batch_size // dp_size + + def _fetch_sft_prepack_rollout_once( + self, + task_name: str, + rollout_id: int, + data_fields: list[str], + ) -> RolloutBatch | None: + # The producer writes exactly one global SFT batch per partition. Use + # the regular seqlen-balanced sampler to fetch the complete local shard + # in one request, just like the stable non-prepack path. Streaming + # token-budget rounds wait on straggling tail samples and their polling + # backoff can turn a small producer skew into multi-second stalls. + dp_rank = mpu.get_data_parallel_rank(with_context_parallel=False) + batch_size = self._sft_prepack_local_batch_size() + partition_id = sft_partition_id(self.args, rollout_id) + sampling_config = {"dp_rank": dp_rank, "task_name": task_name} + rollout_data, _batch_meta = get_data_from_transfer_queue( + args=self.args, + tq_client=self.data_system_client, + data_fields=data_fields, + batch_size=batch_size, + partition_id=partition_id, + task_name=task_name, + sampling_config=sampling_config, + batch_index=0, + broadcast_pp=False, + per_rank_fetch=True, + post_process=False, + synchronize_per_rank_fetch=False, + ) + return rollout_data + + def _pack_sft_prepack_window( + self, + rollout_id: int, + rollout_data: RolloutBatch, + prefetch_started: float, + drain_finished: float, + ) -> PrefetchedSFTWindow: + batch_size = self._sft_prepack_local_batch_size() + dp_rank = mpu.get_data_parallel_rank(with_context_parallel=False) + if len(rollout_data["total_lengths"]) != batch_size: + raise RuntimeError( + f"SFT prefetch local shard size mismatch for rollout_id={rollout_id}: " + f"expected {batch_size}, got {len(rollout_data['total_lengths'])}." + ) + + # K_local is speculative so CPU packing can stay in the background. + # The training thread performs the DP-wide MAX and only short ranks + # repack their real samples to K_global. + samples = rollout_data["total_lengths"] + # CP=1 by _validate_sft_prepack_pipeline; cp_size factor is a no-op + # but kept explicit to mirror get_data_iterator's formula. + cp_size = mpu.get_context_parallel_world_size() + max_tokens = self.args.max_tokens_per_gpu * cp_size + k_local = get_minimum_num_micro_batch_size(samples, max_tokens) + micro_batch_indices = get_seqlen_balanced_partitions(samples, k_local, equal_size=False) + packed_cpu = [ + (prepack_sft_micro_batch_cpu(self.args, _select_rollout_samples(rollout_data, indices)), None) + for indices in micro_batch_indices + ] + packing_finished = time.monotonic() + + assert self._sft_prefetch_stream is not None + assert self._sft_device is not None + device_utils.set_device(self._sft_device) + # First H2D goes on the prefetch-owned stream so the training thread's + # _sft_copy_stream (used by _SFTPrepackedDeviceIterator) never has + # background enqueues racing in front of the current step's copies. + with device_utils.stream_context(self._sft_prefetch_stream): + first_cpu_batch, first_meta = packed_cpu[0] + first_device_batch = PrepackedBatch( + move_tensors_to_device(first_cpu_batch, self._sft_device, non_blocking=True) + ) + first_ready_event = device_utils.Event() + first_ready_event.record(self._sft_prefetch_stream) + + logger.info( + "[sft-prepack] rollout=%d dp=%d real_samples=%d micro_batches=%d " + "tq_drain=%.3fs cpu_pack_pin=%.3fs first_h2d_enqueue=%.3fs", + rollout_id, + dp_rank, + len(rollout_data["total_lengths"]), + len(packed_cpu), + drain_finished - prefetch_started, + packing_finished - drain_finished, + time.monotonic() - packing_finished, + ) + + return PrefetchedSFTWindow( + rollout_id=rollout_id, + rollout_data=rollout_data, + packed_micro_batches=packed_cpu, + first_device_micro_batch=(first_device_batch, first_meta), + first_ready_event=first_ready_event, + ) + + def _fetch_sft_prepack_window_once( + self, + task_name: str, + rollout_id: int, + data_fields: list[str], + ) -> PrefetchedSFTWindow | None: + fetch_started = time.monotonic() + rollout_data = self._fetch_sft_prepack_rollout_once(task_name, rollout_id, data_fields) + if rollout_data is None: + return None + return self._pack_sft_prepack_window(rollout_id, rollout_data, fetch_started, time.monotonic()) + + def _agree_sft_prefetch_window( + self, + task_name: str, + rollout_id: int, + data_fields: list[str], + window: PrefetchedSFTWindow | None, + local_error: BaseException | None, + ) -> PrefetchedSFTWindow: + device = device_utils.make_current_torch_device() + tp_group = mpu.get_tensor_model_parallel_group() + dp_group = mpu.get_data_parallel_group(with_context_parallel=False) + + error_flag = torch.tensor([1 if local_error is not None else 0], dtype=torch.int, device=device) + dist.all_reduce(error_flag, op=dist.ReduceOp.MAX, group=tp_group) + dist.all_reduce(error_flag, op=dist.ReduceOp.MAX, group=dp_group) + if int(error_flag.item()) != 0: + if local_error is not None: + raise local_error + raise RuntimeError( + f"SFT prefetch failed on a peer rank for rollout_id={rollout_id}; " + "aborting to avoid a distributed hang." + ) + + # Any recovery collectives run in the training thread. The background + # worker is deliberately limited to rank-local TQ RPC, CPU packing, and + # H2D enqueue so it cannot interleave NCCL with model fwd/bwd collectives. + max_retries = Envs.RELAX_FETCH_SPLIT_MAX_RETRIES + split_attempt = 0 + empty_attempt = 0 + recovery_error: BaseException | None = None + while True: + has_window = window is not None + status = torch.tensor( + [1 if has_window else 0, 0 if has_window else 1, 1 if recovery_error is not None else 0], + dtype=torch.int, + device=device, + ) + dist.all_reduce(status, op=dist.ReduceOp.SUM, group=tp_group) + got, missing, failed = status.tolist() + + if failed: + if recovery_error is None: + recovery_error = RuntimeError( + f"SFT prefetch recovery failed on a peer TP rank for rollout_id={rollout_id}." + ) + break + if missing == 0: + break + + if got: + split_attempt += 1 + if split_attempt > max_retries: + recovery_error = RuntimeError( + f"[sft-prepack] rollout={rollout_id}: TP ranks still split after {max_retries} " + f"foreground retries ({got}/{got + missing} ranks have data)." + ) + continue + if not has_window: + logger.warning( + "[sft-prepack] rollout=%d: %d/%d TP ranks have data, foreground re-fetch attempt %d", + rollout_id, + got, + got + missing, + split_attempt, + ) + retry_sleep_s = 0.25 + else: + empty_attempt += 1 + if empty_attempt % 100 == 0 and not has_window: + logger.info( + "[sft-prepack] rollout=%d waiting for complete balanced shard; attempts=%d", + rollout_id, + empty_attempt, + ) + retry_sleep_s = Envs.RELAX_EMPTY_POLL_SLEEP_MS / 1000.0 + + if not has_window: + time.sleep(retry_sleep_s) + try: + window = self._fetch_sft_prepack_window_once(task_name, rollout_id, data_fields) + except BaseException as exc: # noqa: BLE001 + recovery_error = exc + + recovery_error_flag = torch.tensor([1 if recovery_error is not None else 0], dtype=torch.int, device=device) + dist.all_reduce(recovery_error_flag, op=dist.ReduceOp.MAX, group=dp_group) + if int(recovery_error_flag.item()) != 0: + if recovery_error is not None: + raise recovery_error + raise RuntimeError( + f"SFT prefetch recovery failed on a peer DP rank for rollout_id={rollout_id}; " + "aborting to avoid a distributed hang." + ) + + assert window is not None + return window + + def _fetch_and_prepack_sft_window( + self, + task_name: str, + rollout_id: int, + data_fields: list[str], + ) -> PrefetchedSFTWindow: + """Fetch, fully pack, pin, and enqueue H2D for one SFT window.""" + dp_rank = mpu.get_data_parallel_rank(with_context_parallel=False) + window = None + fetch_attempt = 0 + while window is None and fetch_attempt <= Envs.RELAX_FETCH_SPLIT_MAX_RETRIES: + window = self._fetch_sft_prepack_window_once(task_name, rollout_id, data_fields) + if window is None: + fetch_attempt += 1 + time.sleep(Envs.RELAX_EMPTY_POLL_SLEEP_MS / 1000.0) + if window is None: + raise _SFTPrefetchMiss( + f"SFT prefetch rollout={rollout_id} dp={dp_rank} found no data after " + f"{fetch_attempt} rank-local attempts." + ) + return window + def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None: # Create data iterator for log_probs and train. data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) @@ -887,7 +1433,13 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None: if self._per_step_rollout and self.args.offload_train: self.sleep() - def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: + def train_actor( + self, + rollout_id: int, + rollout_data: RolloutBatch, + prepared_data_iterator: list[_SFTPrepackedDeviceIterator] | None = None, + prepared_num_microbatches: list[int] | None = None, + ) -> None: # PPO colocate: ``values`` and ``loss_masks`` reach us via TransferQueue # and land on CPU (critic ``.cpu()`` s ``values`` before PUT). Inline # GAE + normalize_advantages need GPU tensors — dispatch here so the @@ -903,9 +1455,17 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: ] # Create data iterator for actor forward + routing replay + train. - data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) + uses_prepared_sft_window = prepared_data_iterator is not None and prepared_num_microbatches is not None + if not uses_prepared_sft_window: + data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data) + else: + data_iterator, num_microbatches = prepared_data_iterator, prepared_num_microbatches # Create a separate iterator with a larger token budget for ref/teacher log-probs - if self.args.use_dynamic_batch_size and self.args.log_probs_max_tokens_per_gpu != self.args.max_tokens_per_gpu: + if ( + not is_sft_mode(self.args) + and self.args.use_dynamic_batch_size + and self.args.log_probs_max_tokens_per_gpu != self.args.max_tokens_per_gpu + ): data_iterator_logprobs, num_microbatches_logprobs = get_data_iterator( self.args, self.model, @@ -998,14 +1558,20 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None: if self.args.use_routing_replay: os.environ["ROUTING_REPLAY_STAGE"] = "replay_backward" with timer("actor_train"): - train( - rollout_id, - self.model, - self.optimizer, - self.opt_param_scheduler, - data_iterator, - num_microbatches, - ) + try: + train( + rollout_id, + self.model, + self.optimizer, + self.opt_param_scheduler, + data_iterator, + num_microbatches, + ) + finally: + for iterator in data_iterator: + close = getattr(iterator, "close", None) + if close is not None: + close() self.prof.step(rollout_id=rollout_id) diff --git a/relax/backends/megatron/data.py b/relax/backends/megatron/data.py index be2346af3..02dccdb70 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -1,8 +1,8 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. from argparse import Namespace -from collections.abc import Sequence -from copy import deepcopy +from collections.abc import Callable, Sequence +from copy import copy, deepcopy from dataclasses import dataclass from typing import Any @@ -43,6 +43,18 @@ ROLLOUT_MINI_PROMPT_GROUP_COUNTS_KEY = "rollout_mini_prompt_group_counts" +class PrepackedBatch(dict): + """Marker dict for SFT micro-batches that have already gone through + ``get_batch`` on the prefetch worker. + + ``get_batch`` short-circuits on ``isinstance(batch, PrepackedBatch)`` so + the second pass on the training thread only moves data to device without + re-doing CP splitting / packed-seq bookkeeping. + """ + + __slots__ = () + + @dataclass(frozen=True) class RolloutMiniBatchPlan: num_rollout_minis: int @@ -272,6 +284,7 @@ def get_batch( qkv_format: str = "thd", allgather_cp: bool = False, is_vl_model: bool = False, + pack_device: torch.device | None = None, ) -> dict[str, torch.Tensor | PackedSeqParams | list[torch.Tensor] | None]: """Generate a CP-ready micro-batch with packed sequence parameters. @@ -311,7 +324,18 @@ def get_batch( else: batch, _ = next(data_iterator) + if isinstance(batch, PrepackedBatch): + return batch + + if pack_device is not None: + for key, dtype in (("tokens", torch.long), ("loss_masks", torch.int)): + values = batch.get(key) + if values is not None: + batch[key] = [torch.as_tensor(value, dtype=dtype, device=pack_device) for value in values] + use_dynamic_context_parallel = getattr(get_args(), "dynamic_context_parallel", False) + if use_dynamic_context_parallel and pack_device is not None and pack_device.type == "cpu": + raise ValueError("CPU prepacking does not support dynamic context parallelism.") if use_dynamic_context_parallel: # Pick this mb's CP size with the SAME per-GPU token budget the iterator was packed # with: forward-only iterators carry log_probs_max_tokens_per_gpu, the training @@ -332,6 +356,7 @@ def get_batch( cp_rank = mpu.get_context_parallel_rank() tokens = batch["tokens"] + batch_device = pack_device or tokens[0].device # use 0 as the pad token id should be fine? pad_token_id = 0 pad_size = mpu.get_tensor_model_parallel_world_size() * pad_multiplier @@ -371,7 +396,7 @@ def get_batch( if needs_unsplit_input and cp_size > 1: tp_size = mpu.get_tensor_model_parallel_world_size() align_size = tp_size * cp_size * 2 - device = device_utils.make_current_torch_device() + device = batch_device seqlens = torch.tensor([t.size(0) for t in tokens], dtype=torch.int32, device=device) seqlens_padded = (seqlens + align_size - 1) // align_size * align_size @@ -420,9 +445,7 @@ def get_batch( tokens = F.pad(tokens, (0, pad), value=pad_token_id) cu_seqlens_list.append(cu_seqlens_list[-1] + pad) - cu_seqlens = torch.tensor( - cu_seqlens_list, dtype=torch.int, device=device_utils.make_current_torch_device() - ) + cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=batch_device) tokens = tokens.chunk(cp_size, dim=0)[cp_rank] else: tokens = [ @@ -443,9 +466,7 @@ def get_batch( cu_seqlens.append(cu_seqlens[-1] + pad) # thd requires the cu_seqlens to be of the origin length - cu_seqlens = ( - torch.tensor(cu_seqlens, dtype=torch.int).to(device_utils.make_current_torch_device()) * cp_size - ) + cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int, device=batch_device) * cp_size max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() packed_seq_params = PackedSeqParams( @@ -602,9 +623,27 @@ def get_batch( ] batch = move_tensors_to_device(batch, batch["tokens"].device) + if pack_device is not None: + batch = PrepackedBatch(batch) return batch +def prepack_sft_micro_batch_cpu(args: Namespace, batch: RolloutBatch) -> PrepackedBatch: + """Build one complete SFT micro-batch on CPU for background prefetch.""" + iterator = iter([(batch, None)]) + packed = get_batch( + iterator, + ["tokens"], + args.data_pad_size_multiplier, + args.qkv_format, + args.allgather_cp, + getattr(args, "is_vl_model", False), + pack_device=torch.device("cpu"), + ) + pinned = pin_tensors(packed) + return pinned if isinstance(pinned, PrepackedBatch) else PrepackedBatch(pinned) + + def gather_log_data( metric_name: str, args: Namespace, @@ -887,10 +926,11 @@ def _generate_data_iterator( end = step_offsets[i + 1] num_microbatches.append(get_minimum_num_micro_batch_size(samples[start:end], _max_tokens * cp_size)) - num_microbatches = torch.tensor( + required_num_microbatches = torch.tensor( num_microbatches, dtype=torch.int, device=device_utils.make_current_torch_device() ) - dist.all_reduce(num_microbatches, op=dist.ReduceOp.MAX, group=dp_group) + dist.all_reduce(required_num_microbatches, op=dist.ReduceOp.MAX, group=dp_group) + num_microbatches = required_num_microbatches if vpp_size > 1: # vpp requires the number of microbatches to be divisible by vpp_size @@ -1071,7 +1111,7 @@ def log_rollout_data( stats = torch.tensor( [max(total_lengths), -min(total_lengths)], dtype=torch.int64, - device=loss_masks[0].device, + device=device_utils.make_current_torch_device(), ) dist.all_reduce(stats, op=dist.ReduceOp.MAX, group=dp_group) log_dict["total_lengths/max"] = int(stats[0].item()) @@ -1222,17 +1262,65 @@ def sync_actor_critic_data( ) -def move_tensors_to_device(data, device): +def _map_packed_seq_params(data: PackedSeqParams, map_value: Callable[[Any], Any]) -> PackedSeqParams: + mapped = copy(data) + for name, value in vars(data).items(): + setattr(mapped, name, map_value(value)) + return mapped + + +def move_tensors_to_device(data: Any, device: torch.device, *, non_blocking: bool = False) -> Any: """Recursively move tensors in a (nested) dict/list to the specified device. Non-tensor values are left unchanged. """ if isinstance(data, dict): - return {k: move_tensors_to_device(v, device) for k, v in data.items()} - elif isinstance(data, list): - return [move_tensors_to_device(v, device) for v in data] + return {k: move_tensors_to_device(v, device, non_blocking=non_blocking) for k, v in data.items()} + if isinstance(data, list): + return [move_tensors_to_device(v, device, non_blocking=non_blocking) for v in data] + if isinstance(data, tuple): + return tuple(move_tensors_to_device(v, device, non_blocking=non_blocking) for v in data) + if isinstance(data, PackedSeqParams): + return _map_packed_seq_params( + data, + lambda value: move_tensors_to_device(value, device, non_blocking=non_blocking), + ) + if isinstance(data, torch.Tensor): + return data.to(device, non_blocking=non_blocking) + return data # e.g., int, str, None, etc. + + +def pin_tensors(data: Any) -> Any: + """Recursively pin CPU tensors in a pre-packed micro-batch.""" + if isinstance(data, dict): + return {k: pin_tensors(v) for k, v in data.items()} + if isinstance(data, list): + return [pin_tensors(v) for v in data] + if isinstance(data, tuple): + return tuple(pin_tensors(v) for v in data) + if isinstance(data, PackedSeqParams): + return _map_packed_seq_params(data, pin_tensors) + if isinstance(data, torch.Tensor) and data.device.type == "cpu": + try: + return data.contiguous().pin_memory() + except RuntimeError as exc: + from relax.utils.data.micro_batch_ring import _raise_pin_memory_failure + + _raise_pin_memory_failure("pin_tensors", exc) + return data + + +def record_tensors_on_stream(data: Any, stream: Any) -> None: + """Keep asynchronously copied tensors alive on the consuming stream.""" + if isinstance(data, dict): + for value in data.values(): + record_tensors_on_stream(value, stream) + elif isinstance(data, (list, tuple)): + for value in data: + record_tensors_on_stream(value, stream) + elif isinstance(data, PackedSeqParams): + for value in vars(data).values(): + record_tensors_on_stream(value, stream) elif isinstance(data, torch.Tensor): - return data.to(device) - else: - return data # e.g., int, str, None, etc. + data.record_stream(stream) diff --git a/relax/components/sft.py b/relax/components/sft.py index e9deb4e80..e4145e6e1 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -370,6 +370,7 @@ async def _produce_one_step(self) -> None: await self.data_system_client.async_put( data=dict_to_tensordict(backend_batch, batch_size=len(backend_batch["tokens"])), partition_id=f"sft_{self.step}", + custom_meta=[{"total_lengths": int(length)} for length in backend_batch["total_lengths"]], ) if crossed_epoch: self._logger.info( @@ -498,6 +499,7 @@ async def _maybe_produce_eval(self) -> None: await self.data_system_client.async_put( data=dict_to_tensordict(chunk, batch_size=len(chunk["tokens"])), partition_id=partition_id, + custom_meta=[{"total_lengths": int(length)} for length in chunk["total_lengths"]], ) drained = await self._wait_for_partition_drained(partition_id, timeout_sec=chunk_drain_timeout) if not drained: diff --git a/relax/core/controller.py b/relax/core/controller.py index 5f0405bda..4c8cff10e 100644 --- a/relax/core/controller.py +++ b/relax/core/controller.py @@ -292,7 +292,12 @@ def _cleanup_s3_model_weights_after_init(self, *, force: bool = False) -> None: def _initialize_data_system(self): algo_key = resolve_sft_algo_key(self.config) dp_size = compute_dp_size(self.config) - if getattr(self.config, "fully_async", False) and getattr(self.config, "use_dynamic_batch_size", False): + use_sft_prepack = algo_key == "sft" and getattr(self.config, "sft_async_prepack", False) + if ( + getattr(self.config, "fully_async", False) + and getattr(self.config, "use_dynamic_batch_size", False) + and not use_sft_prepack + ): sampler = IdentityWindowSampler( dp_size=dp_size, placement="streaming", diff --git a/relax/engine/sft/predict/runner.py b/relax/engine/sft/predict/runner.py index de04a04f3..405143652 100644 --- a/relax/engine/sft/predict/runner.py +++ b/relax/engine/sft/predict/runner.py @@ -43,6 +43,11 @@ def run_sft_predict(actor, rollout_id: int) -> None: from relax.backends.megatron.initialize import is_megatron_main_rank args = actor.args + # Cache the topology-derived role while Megatron process groups are live. + # sleep() and update_weights() intentionally destroy the reloadable NCCL + # groups until wake_up(), so querying mpu ranks during predict would access + # a ReloadableProcessGroup whose inner group is None. + is_main_rank = is_megatron_main_rank() dist.barrier(group=get_gloo_group()) _t_predict_start = time.monotonic() _t_sleep = _t_update = _t_http = _t_wake = 0.0 @@ -54,7 +59,7 @@ def run_sft_predict(actor, rollout_id: int) -> None: actor.update_weights() _t_update = time.monotonic() - _t dist.barrier(group=get_gloo_group()) - if is_megatron_main_rank(): + if is_main_rank: _t = time.monotonic() try: # Without a timeout, a hung rollout service blocks this rank @@ -76,7 +81,7 @@ def run_sft_predict(actor, rollout_id: int) -> None: _t = time.monotonic() actor.wake_up() _t_wake = time.monotonic() - _t - if is_megatron_main_rank(): + if is_main_rank: step = compute_rollout_step(args, rollout_id) metrics = { "perf/sft_predict_time": time.monotonic() - _t_predict_start, diff --git a/relax/utils/arguments.py b/relax/utils/arguments.py index 0848c4f45..080c72988 100644 --- a/relax/utils/arguments.py +++ b/relax/utils/arguments.py @@ -621,6 +621,22 @@ def add_train_arguments(parser): default=4, help="Worker threads inside the SFT PrefetchBuffer for I/O-bound media decoding.", ) + parser.add_argument( + "--sft-async-prepack", + action="store_true", + default=False, + help=( + "Enable the SFT prepack pipeline: TQ fetch + seqlen-balanced " + "micro-batch partitioning + THD packing + pinned-memory H2D " + "are all offloaded to a background worker, keeping only " + "fwd/bwd on the training thread. Data / batch / loss-scaling " + "semantics match the standard SFT path exactly (same K, same " + "get_seqlen_balanced_partitions, same __loss_scale__). Requires " + "--per-rank-fetch and at least two in-flight steps " + "(--max-staleness >= 1 or --sft-max-in-flight-steps >= 2); " + "PP=1, CP=1, VPP=1 and THD qkv format only." + ), + ) parser.add_argument( "--sft-oversize-strategy", type=str, @@ -1372,7 +1388,7 @@ def add_data_arguments(parser): help=( "Balance the number of tokens between data parallel ranks with `karmarkar_karp` for verl. " "Note that this may allocate the different response of the same prompt into different training steps. " - "In fully-async + --use-dynamic-batch-size mode this is effectively always on: the " + "In streaming dynamic-batch mode this is effectively always on: the " "StreamingTokenBudgetSampler already balances tokens across DP ranks per sample, so the " "flag is accepted but has no additional effect there." ), @@ -2844,8 +2860,7 @@ def _parse_args_impl(add_custom_arguments=None, *, model_source=None): if not args.debug_train_only: sglang_validate_args(args) - # Only fully-async mode relies on the newer TransferQueue (e.g. - # StreamingTokenBudgetSampler), so gate the version requirement on it. + # Only fully-async mode relies on the newer TransferQueue streaming sampler. if getattr(args, "fully_async", False): check_transfer_queue_version() @@ -2957,11 +2972,16 @@ def _normalize_mtp_only_training_args(args) -> None: def _normalize_sft_max_in_flight_steps(args, is_sft: bool) -> None: sft_max_in_flight_steps = getattr(args, "sft_max_in_flight_steps", None) if sft_max_in_flight_steps is None: + if is_sft and getattr(args, "sft_async_prepack", False) and args.max_staleness < 1: + raise ValueError("--sft-async-prepack requires --max-staleness >= 1 or --sft-max-in-flight-steps >= 2.") return if not is_sft: raise ValueError("--sft-max-in-flight-steps is only meaningful under --loss-type sft.") - if sft_max_in_flight_steps < 1: + minimum_steps = 2 if getattr(args, "sft_async_prepack", False) else 1 + if sft_max_in_flight_steps < minimum_steps: + if minimum_steps == 2: + raise ValueError("--sft-async-prepack requires --sft-max-in-flight-steps >= 2.") raise ValueError("--sft-max-in-flight-steps must be >= 1.") args.max_staleness = sft_max_in_flight_steps - 1 @@ -3562,6 +3582,16 @@ def slime_validate_args(args): if args.loss_type == "sft": if not args.custom_dataset_class_path and not args.prompt_data: raise ValueError("--loss-type sft requires --prompt-data.") + if getattr(args, "sft_async_prepack", False): + if not args.per_rank_fetch: + raise ValueError( + "--sft-async-prepack enables background prepacking and requires --per-rank-fetch; " + "background prefetch workers must not execute CP/TP/PP collectives." + ) + if args.use_routing_replay or args.use_rollout_routing_replay: + raise ValueError( + "--sft-async-prepack does not support routing replay because its iterator is single-pass." + ) if args.sft_oversize_strategy == "custom" and not args.sft_oversize_custom_function_path: raise ValueError("--sft-oversize-strategy custom requires --sft-oversize-custom-function-path.") # SFT does not use advantages / reference; force-disable to avoid wasted compute. @@ -3576,14 +3606,13 @@ def slime_validate_args(args): "SFT relies on dynamic batching to bound per-GPU tokens (CP-aware) and to filter " "samples that cannot fit on a single GPU." ) - # The controller always installs SeqlenBalancedSampler for SFT (see - # `core/controller.py:_initialize_data_system`). That sampler can hand - # different sample counts to each DP rank, which the Megatron data - # path only handles correctly when args.balance_data is True. Force it - # on so the two layers stay consistent. + # The controller installs SeqlenBalancedSampler for SFT, so keep the + # Megatron data path in DP-balanced mode as well. if not args.balance_data: logger.info("--loss-type sft: auto-enabling --balance-data for DP-balanced batching.") args.balance_data = True + elif getattr(args, "sft_async_prepack", False): + raise ValueError("--sft-async-prepack is only meaningful under --loss-type sft.") task_type = getattr(args, "task_type", "causal_lm") if task_type == "seq_cls": diff --git a/relax/utils/data/micro_batch_ring.py b/relax/utils/data/micro_batch_ring.py new file mode 100644 index 000000000..c084a5306 --- /dev/null +++ b/relax/utils/data/micro_batch_ring.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""CPU-only prefetch helpers for the SFT prepack pipeline. + +Enabled by ``--sft-async-prepack``. The prefetch worker fetches the whole +rollout partition from TransferQueue, runs the same seqlen-balanced K-way +partition dev uses, CPU-packs each micro-batch, pins CPU memory, and enqueues +the first H2D copy — all off the training thread. The training thread only pays +for one H2D wait per micro-batch. +""" + +from __future__ import annotations + +import threading +from collections.abc import Callable, Hashable +from dataclasses import dataclass +from typing import Any + + +def _raise_pin_memory_failure(context: str, exc: BaseException) -> None: + """Fail fast on pin_memory errors. + + Falling back to a pageable CPU tensor would push the H2D cost onto the + training thread (``.to(..., non_blocking=True)`` implicitly stages via a + pinned buffer or blocks) and silently erase the whole point of the prepack + pipeline. Surface the failure instead so the user can tune pinned memory + limits or drop --sft-async-prepack. + """ + raise RuntimeError( + f"pin_memory failed in {context} ({exc}); the SFT prepack pipeline requires pinned CPU " + "tensors so the training thread never issues pageable H2D copies. Increase the pinned " + "memory limit (e.g. ulimit -l) or disable --sft-async-prepack." + ) from exc + + +def is_sft_async_prepack_enabled(args: Any) -> bool: + """The SFT prepack pipeline is opt-in via --sft-async-prepack.""" + return bool(getattr(args, "sft_async_prepack", False)) + + +@dataclass(frozen=True) +class PrefetchedSFTWindow: + rollout_id: int + rollout_data: dict[str, Any] + packed_micro_batches: list[tuple[dict[str, Any], Any]] + first_device_micro_batch: tuple[dict[str, Any], Any] + first_ready_event: Any + + +class SFTWindowPrefetcher: + """Keep at most one future SFT window in a daemon worker. + + Re-invoking ``prefetch`` with the same ``rollout_id`` is idempotent: the + already-running worker's result will be returned by the next ``get``. Pass + an ``identity`` key so a lookahead ``prefetch`` (typically triggered after + a successful ``get``) cannot silently reuse a stale closure. + """ + + def __init__(self) -> None: + self._condition = threading.Condition() + self._rollout_id: int | None = None + self._identity: Hashable | None = None + self._result: PrefetchedSFTWindow | None = None + self._error: BaseException | None = None + self._thread: threading.Thread | None = None + + def prefetch( + self, + rollout_id: int, + fetch_fn: Callable[[], PrefetchedSFTWindow], + *, + identity: Hashable | None = None, + ) -> None: + with self._condition: + if self._rollout_id is not None: + if self._rollout_id != rollout_id: + raise RuntimeError( + f"SFT prefetcher already owns rollout {self._rollout_id}; cannot start rollout {rollout_id}" + ) + if identity is not None and self._identity is not None and identity != self._identity: + raise RuntimeError( + f"SFT prefetcher rollout {rollout_id} re-invoked with different identity " + f"{identity!r} (previous={self._identity!r}); the second fetch_fn would " + "silently be ignored." + ) + return + self._rollout_id = rollout_id + self._identity = identity + self._result = None + self._error = None + + def _run() -> None: + try: + result = fetch_fn() + if result.rollout_id != rollout_id: + raise RuntimeError( + f"SFT prefetch worker returned rollout {result.rollout_id}; expected {rollout_id}" + ) + with self._condition: + self._result = result + self._condition.notify_all() + except BaseException as exc: # noqa: BLE001 + with self._condition: + self._error = exc + self._condition.notify_all() + + self._thread = threading.Thread(target=_run, name=f"sft-window-{rollout_id}", daemon=True) + self._thread.start() + + def get(self, rollout_id: int) -> PrefetchedSFTWindow: + with self._condition: + if self._rollout_id != rollout_id: + raise RuntimeError(f"SFT prefetcher has rollout {self._rollout_id}, requested {rollout_id}") + while self._result is None and self._error is None: + self._condition.wait() + if self._error is not None: + error = self._error + self._reset_locked() + raise error + assert self._result is not None + result = self._result + self._reset_locked() + return result + + def _reset_locked(self) -> None: + self._rollout_id = None + self._identity = None + self._result = None + self._error = None + # The daemon thread has already completed by the time get() returns; + # dropping the reference is safe. + self._thread = None + + +__all__ = [ + "PrefetchedSFTWindow", + "SFTWindowPrefetcher", + "is_sft_async_prepack_enabled", +] diff --git a/relax/utils/data/stream_dataloader.py b/relax/utils/data/stream_dataloader.py index a9406f2f3..9e22a8f0b 100644 --- a/relax/utils/data/stream_dataloader.py +++ b/relax/utils/data/stream_dataloader.py @@ -644,6 +644,8 @@ def get_data_from_transfer_queue( per_rank_fetch: bool = False, token_budget: int | None = None, allow_underfill: bool = True, + post_process: bool = True, + synchronize_per_rank_fetch: bool = True, ): """Fetch a batch from the transfer queue and broadcast it across tensor- parallel and optionally pipeline-parallel ranks. @@ -681,11 +683,21 @@ def get_data_from_transfer_queue( dominates ``tgd_bcast_tp_time``. Caller must ensure ``rollout_routed_experts`` is not in ``data_fields`` (its bcast path is incompatible) — actor.py guards this. + post_process: Move and reshape rollout fields for Megatron. Set to + False only for CPU-only background prefetch; the training thread + must materialize each micro-batch before use. + synchronize_per_rank_fetch: When True, per-rank fetches agree on + empty-vs-data state across the model-parallel replica before + returning. Set to False only when the caller has its own foreground + synchronization point and must keep this call collective-free. Returns: Tuple[Optional[dict], Optional[Any]]: A tuple of (rollout_data, batch_meta). If no data is available, both elements are None. """ + if not synchronize_per_rank_fetch and not per_rank_fetch: + raise ValueError("synchronize_per_rank_fetch=False requires per_rank_fetch=True") + # Compose request configuration and ask the queue for metadata. config = {**sampling_config, "batch_index": batch_index, "partition_id": partition_id} if token_budget is not None: @@ -765,7 +777,7 @@ def _fetch_once() -> list: # will receive the real data via broadcast. rollout_data = [None, None] - if per_rank_fetch: + if per_rank_fetch and synchronize_per_rank_fetch: # No broadcast follows, so a producer race can split this logical DP # rank into "got data" and "empty meta" model-parallel subsets. rollout_data = _agree_on_fetch( @@ -777,7 +789,7 @@ def _fetch_once() -> list: # Use an explicit device so the communication backend (e.g. NCCL) # can bind to a known device context. - cuda_dev = device_utils.make_current_torch_device() + cuda_dev = None if per_rank_fetch else device_utils.make_current_torch_device() # --- Extract rollout_routed_experts BEFORE broadcast_object_list --- # broadcast_object_list uses pickle for the entire payload. When @@ -959,7 +971,8 @@ def _fetch_once() -> list: if has_multimodal and mm_inputs is not None: rollout_data["multimodal_train_inputs"] = mm_inputs - post_process_rollout_data(args, rollout_data) + if post_process: + post_process_rollout_data(args, rollout_data) return rollout_data, batch_meta @@ -1433,10 +1446,9 @@ def __next__(self) -> Tuple[Dict[str, Any], Any]: ) return self._make_dummy_batch() - partition_id = f"train_{self.rollout_id}" - t0 = time.monotonic() empty_streak = 0 + partition_id = f"train_{self.rollout_id}" while True: sampling_config = self._sampling_config(self._batch_index) diff --git a/scripts/training/sft/run-qwen3-0.6B-math-8xgpu.sh b/scripts/training/sft/run-qwen3-0.6B-math-8xgpu.sh new file mode 100755 index 000000000..af393163e --- /dev/null +++ b/scripts/training/sft/run-qwen3-0.6B-math-8xgpu.sh @@ -0,0 +1,132 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# Qwen3-0.6B SFT on OpenMathReasoning-mini, 8xGPU single-node, ray-submit launch. +# +# Uses the SFT prepack pipeline: --sft-async-prepack offloads TQ fetch + +# seqlen-balanced packing + pinned-memory H2D to a background worker, +# keeping only fwd/bwd on the training thread. Data / batch / loss-scaling +# semantics match dev exactly (K auto-aligns across DP by repartitioning real samples). +# +# --sft-max-in-flight-steps N: TQ buffer depth = N. Also maps to +# max_staleness = N - 1, which controls how many rollouts ahead the +# prefetch worker can look. Bigger N = better overlap between train step +# and prefetch, at the cost of more pinned CPU memory. N=4 is a good +# default; --sft-async-prepack rejects N=1 because it cannot overlap the +# next window with the current training step. +# +# 0.6B fits on a single GPU so TP/PP/CP are all 1 and DP=8. +# +# Usage: +# bash scripts/training/sft/run-qwen3-0.6B-math-8xgpu.sh + +set -ex +set -o pipefail + +now=$(date "+%Y-%m-%d-%H:%M:%S") +echo "当前时间: $now" + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +# Auto-source local environment when not launched via an external entrypoint +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../entrypoint/local.sh" +fi +source "${MODEL_CONFIG_DIR}/qwen3-0.6B.sh" + +PROJECT_NAME="${PROJECT_NAME:=Relax/sft/math}" +EXP_NAME=qwen3-0.6b-sft-math-gpu8 +EXP_DIR="${MODEL_DIR:=${SCRIPT_DIR}/../../../../exps}" +DATA_DIR="${DATA_DIR:=${SCRIPT_DIR}/data}" +PROMPT_DATA="${PROMPT_DATA:=${DATA_DIR}/sft/data/OpenMathReasoning-mini/data/cot-00000-of-00001.parquet}" +SAVE_DIR="${SAVE_DIR:=${SCRIPT_DIR}/../../../checkpoints/qwen3-0.6B-math-sft}" + +CKPT_ARGS=( + --hf-checkpoint ${EXP_DIR}/Qwen3-0.6B + --ref-load ${EXP_DIR}/Qwen3-0.6B + + --megatron-to-hf-mode bridge + --save ${SAVE_DIR}/sft/${EXP_NAME} + --load ${SAVE_DIR}/sft/${EXP_NAME} + --save-interval 1000 + --num-epoch 10 +) + +SFT_ARGS=( + --loss-type sft + --prompt-data "${PROMPT_DATA}" + --input-key problem + --label-key generated_solution + --global-batch-size 32 + --use-dynamic-batch-size + --max-tokens-per-gpu 20480 + --sft-oversize-strategy skip + --balance-data + --per-rank-fetch + --sft-async-prepack + --sft-prefetch-num-workers 8 + --sft-prefetch-buffer-size 256 +) + +EVAL_ARGS=( + --eval-size 0.01 + --eval-interval 1000 +) + +PERF_ARGS=( + --tensor-model-parallel-size 1 + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --no-rope-fusion + + --colocate + --cross-entropy-loss-fusion +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-5 + --lr-decay-style cosine + --min-lr 1e-6 + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 + --clip-grad 1.0 +) + +WANDB_ARGS=( + --use-clearml + --use-metrics-service + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name ${EXP_NAME}-${now} +) + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash + --use-health-check +) + +mkdir -p log + +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource '{"sft": [1, 0], "actor": [1, 8], "rollout": [1, 8]}' \ + --sft-max-in-flight-steps 4 \ + --num-data-storage-units 8 \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${SFT_ARGS[@]}" \ + "${EVAL_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${WANDB_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${MISC_ARGS[@]}" 2>&1 | tee log/qwen3-0.6b-sft-math-gpu8-${now}.log diff --git a/scripts/training/sft/run-qwen3-vl-4B-pokemon-8xgpu.sh b/scripts/training/sft/run-qwen3-vl-4B-pokemon-8xgpu.sh index c2ce39971..0ea644e96 100755 --- a/scripts/training/sft/run-qwen3-vl-4B-pokemon-8xgpu.sh +++ b/scripts/training/sft/run-qwen3-vl-4B-pokemon-8xgpu.sh @@ -60,6 +60,7 @@ SFT_ARGS=( --max-tokens-per-gpu 20480 --balance-data --per-rank-fetch + --sft-async-prepack --sft-prefetch-num-workers 16 --sft-prefetch-buffer-size 512 ) diff --git a/tests/backends/megatron/test_data_vpp.py b/tests/backends/megatron/test_data_vpp.py index c2f8980aa..8eeb706e5 100644 --- a/tests/backends/megatron/test_data_vpp.py +++ b/tests/backends/megatron/test_data_vpp.py @@ -94,6 +94,56 @@ def test_rollout_minibatch_plan_rejects_non_divisible_prompt_groups(monkeypatch) data_module.build_rollout_minibatch_plan(args, dp_size=2) +def test_log_rollout_data_creates_collective_stats_on_training_device(monkeypatch): + data_module = _load_data_module(monkeypatch) + loss_masks = [torch.tensor([0, 1]), torch.tensor([0, 1, 1])] + requested_devices = [] + original_tensor = torch.tensor + + monkeypatch.setattr(data_module.mpu, "get_tensor_model_parallel_rank", lambda: 0, raising=False) + monkeypatch.setattr(data_module.mpu, "is_pipeline_last_stage", lambda: True, raising=False) + monkeypatch.setattr(data_module.mpu, "get_context_parallel_world_size", lambda: 1, raising=False) + monkeypatch.setattr( + data_module.mpu, + "get_data_parallel_group", + lambda with_context_parallel=True: object(), + raising=False, + ) + monkeypatch.setattr(data_module.device_utils, "make_current_torch_device", lambda: "training-device") + + def capture_tensor(*args, **kwargs): + requested_devices.append(kwargs.get("device")) + kwargs["device"] = "cpu" + return original_tensor(*args, **kwargs) + + monkeypatch.setattr(data_module.torch, "tensor", capture_tensor) + monkeypatch.setattr(data_module.dist, "all_reduce", lambda *_args, **_kwargs: None) + monkeypatch.setattr(data_module, "gather_log_data", lambda *_args, **_kwargs: None) + monkeypatch.setattr(data_module, "maybe_padded_total_lengths", lambda *_args, **_kwargs: None) + + args = Namespace( + qkv_format="thd", + is_vl_model=False, + uses_unsplit_forward=False, + dynamic_context_parallel=False, + use_opd=False, + rollout_batch_size=2, + n_samples_per_prompt=1, + ci_test=False, + log_multi_turn=False, + log_correct_samples=False, + ) + rollout_data = { + "total_lengths": [4, 6], + "response_lengths": [2, 3], + "loss_masks": loss_masks, + } + + data_module.log_rollout_data(rollout_id=0, args=args, rollout_data=rollout_data) + + assert requested_devices == ["training-device"] + + def test_concat_rollout_batches_preserves_order_and_scalar_metadata(monkeypatch): data_module = _load_data_module(monkeypatch) diff --git a/tests/backends/megatron/test_sft_prepack.py b/tests/backends/megatron/test_sft_prepack.py new file mode 100644 index 000000000..201cc41d9 --- /dev/null +++ b/tests/backends/megatron/test_sft_prepack.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from argparse import Namespace + +import pytest +import torch + + +try: + from relax.backends.megatron import actor as actor_module +except (ImportError, AssertionError) as exc: + pytest.skip(f"relax.backends.megatron.actor unavailable: {exc}", allow_module_level=True) + + +def test_sft_prepacked_iterator_close_releases_device_batch_and_event(): + iterator = actor_module._SFTPrepackedDeviceIterator( + packed_cpu=[], + first_device_micro_batch=(actor_module.PrepackedBatch(), None), + first_ready_event=object(), + copy_stream=object(), + device=torch.device("cpu"), + ) + + iterator.close() + + assert iterator._packed_cpu == [] + assert iterator._next_device_micro_batch is None + assert iterator._next_ready_event is None + + +def test_sft_peer_error_agreement_propagates_peer_failure(monkeypatch): + groups = [] + + def mark_peer_error(error_flag, *, op, group): + groups.append(group) + error_flag.fill_(1) + + monkeypatch.setattr(actor_module.dist, "all_reduce", mark_peer_error) + + with pytest.raises(RuntimeError, match="failed on a peer rank"): + actor_module._raise_if_sft_peer_error( + None, + phase="validation", + device=torch.device("cpu"), + tp_group="tp", + dp_group="dp", + ) + + assert groups == ["tp", "dp"] + + +def test_sft_lookahead_pauses_on_checkpoint_boundary(monkeypatch): + monkeypatch.setattr(actor_module, "should_run_sft_eval", lambda *_args: False) + monkeypatch.setattr(actor_module, "should_run_sft_predict", lambda *_args: False) + args = Namespace(num_rollout=100, save="/checkpoint", rotate_ckpt=False, save_interval=20) + + assert actor_module._should_pause_sft_prepack_lookahead(args, rollout_id=19) is True + assert actor_module._should_pause_sft_prepack_lookahead(args, rollout_id=18) is False diff --git a/tests/components/test_sft.py b/tests/components/test_sft.py index cdd854570..eacf1261c 100644 --- a/tests/components/test_sft.py +++ b/tests/components/test_sft.py @@ -164,6 +164,7 @@ async def test_sft_step_pushes_one_batch_to_tq(monkeypatch): assert "total_lengths" in pushed_data assert "response_lengths" in pushed_data assert kwargs_call.get("partition_id") == "sft_0" + assert kwargs_call.get("custom_meta") == [{"total_lengths": 8}] * 4 @pytest.mark.asyncio diff --git a/tests/utils/data/test_streaming_tq_iterator.py b/tests/utils/data/test_streaming_tq_iterator.py index 8b594bb14..256afa631 100644 --- a/tests/utils/data/test_streaming_tq_iterator.py +++ b/tests/utils/data/test_streaming_tq_iterator.py @@ -91,6 +91,57 @@ def test_streaming_tq_iterator_sampling_config_carries_window_quota(monkeypatch) assert "consumed_samples" not in config +def test_get_data_from_transfer_queue_can_skip_per_rank_agreement(monkeypatch): + stream_module = _load_stream_module(monkeypatch) + + class _Meta: + size = 0 + + class _Client: + def get_meta(self, **kwargs): + return _Meta() + + def fail_agreement(*args, **kwargs): + raise AssertionError("_agree_on_fetch should not run") + + monkeypatch.setattr(stream_module, "_agree_on_fetch", fail_agreement) + + data, meta = stream_module.get_data_from_transfer_queue( + args=Namespace(), + tq_client=_Client(), + data_fields=["tokens"], + batch_size=1, + partition_id="partition", + task_name="task", + sampling_config={}, + batch_index=0, + broadcast_pp=False, + per_rank_fetch=True, + synchronize_per_rank_fetch=False, + ) + + assert data is None + assert meta.size == 0 + + +def test_get_data_from_transfer_queue_rejects_disabled_agreement_without_per_rank_fetch(monkeypatch): + stream_module = _load_stream_module(monkeypatch) + + with pytest.raises(ValueError, match="requires per_rank_fetch=True"): + stream_module.get_data_from_transfer_queue( + args=Namespace(), + tq_client=object(), + data_fields=["tokens"], + batch_size=1, + partition_id="partition", + task_name="task", + sampling_config={}, + batch_index=0, + per_rank_fetch=False, + synchronize_per_rank_fetch=False, + ) + + def test_streaming_tq_iterator_finishes_on_window_drained_with_underfill(monkeypatch): """Regression for the fully-async DP-imbalance deadlock. diff --git a/tests/utils/test_arguments_sft.py b/tests/utils/test_arguments_sft.py new file mode 100644 index 000000000..74cd283ee --- /dev/null +++ b/tests/utils/test_arguments_sft.py @@ -0,0 +1,37 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from types import SimpleNamespace + +import pytest + +from relax.utils.arguments import _normalize_sft_max_in_flight_steps + + +def test_sft_async_prepack_rejects_single_in_flight_step(): + args = SimpleNamespace(sft_max_in_flight_steps=1, sft_async_prepack=True, max_staleness=0) + + with pytest.raises(ValueError, match="requires --sft-max-in-flight-steps >= 2"): + _normalize_sft_max_in_flight_steps(args, is_sft=True) + + +def test_sft_async_prepack_rejects_zero_max_staleness_without_alias(): + args = SimpleNamespace(sft_max_in_flight_steps=None, sft_async_prepack=True, max_staleness=0) + + with pytest.raises(ValueError, match="requires --max-staleness >= 1"): + _normalize_sft_max_in_flight_steps(args, is_sft=True) + + +def test_sft_async_prepack_maps_two_in_flight_steps_to_one_stale_step(): + args = SimpleNamespace(sft_max_in_flight_steps=2, sft_async_prepack=True, max_staleness=0) + + _normalize_sft_max_in_flight_steps(args, is_sft=True) + + assert args.max_staleness == 1 + + +def test_sft_without_async_prepack_allows_one_in_flight_step(): + args = SimpleNamespace(sft_max_in_flight_steps=1, sft_async_prepack=False, max_staleness=0) + + _normalize_sft_max_in_flight_steps(args, is_sft=True) + + assert args.max_staleness == 0 From 29e43c4c714829788c6b81b87ae0ceb0c783f0a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E4=BD=B3=E5=85=B4?= Date: Wed, 2 Sep 2026 15:38:35 +0800 Subject: [PATCH 09/34] feat: adapt new megatron (cherry picked from commit d2713039635dfd614e438c5cff33e3f1814174c8) --- relax/backends/megatron/model.py | 24 ++++++++++++++++--- relax/backends/megatron/model_provider.py | 1 + .../hf_weight_iterator_bridge.py | 9 ++++++- .../backends/device_direct.py | 7 +++++- relax/models/dots_ocr/megatron/bridge.py | 4 ++-- relax/models/qwen_omni/qwen3_omni_bridge.py | 4 ++-- 6 files changed, 40 insertions(+), 9 deletions(-) diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 55ebe5e4d..3a00fd1db 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -1483,14 +1483,32 @@ def train( mtp_loss_scale = 1 / num_microbatches[step_id] tracker = MTPLossLoggingHelper.tracker - if "values" in tracker: + # mcore >= 0.19 renamed the tracker payload: the per-microbatch + # accumulator is now "loss_sums" (plus "num_tokens" in per-token mode) + # instead of "values", and the cross-rank reduction moved into + # MTPLossLoggingHelper.reduce_loss_in_tracker(), which repopulates + # "values" and handles both normalization modes. Older mcore exposes + # "values" directly with no such helper, so reduce by hand there. + # Mirrors upstream MTPLossLoggingHelper.track_mtp_metrics. + mtp_losses = None + reduce_in_tracker = getattr(MTPLossLoggingHelper, "reduce_loss_in_tracker", None) + if reduce_in_tracker is not None: + reduce_in_tracker() + elif "values" in tracker: values = tracker["values"] if tracker.get("reduce_group") is not None: torch.distributed.all_reduce(values, group=tracker.get("reduce_group")) if tracker.get("avg_group") is not None: torch.distributed.all_reduce(values, group=tracker["avg_group"], op=torch.distributed.ReduceOp.AVG) + + # "values" is the reduced payload on both old and new mcore; + # "loss_values" is the compat slot filled by save_loss_and_metrics_to_tracker. + mtp_values = tracker.get("values") + if mtp_values is None: + mtp_values = tracker.get("loss_values") + if mtp_values is not None: # here we assume only one mtp layer - mtp_losses = (tracker["values"] * mtp_loss_scale).item() + mtp_losses = (mtp_values * mtp_loss_scale).item() MTPLossLoggingHelper.clean_loss_in_tracker() # CI check: verify MTP loss is within expected bounds @@ -1516,7 +1534,7 @@ def train( log_dict[f"train/{role_tag}grad_norm"] = ( grad_norm.item() if isinstance(grad_norm, torch.Tensor) else grad_norm ) - if args.enable_mtp_training: + if args.enable_mtp_training and mtp_losses is not None: log_dict[f"train/{role_tag}mtp_loss"] = mtp_losses log_dict[f"train/{role_tag}global_batch_size"] = global_batch_sizes[step_id] diff --git a/relax/backends/megatron/model_provider.py b/relax/backends/megatron/model_provider.py index fff4e85aa..f09516a8e 100644 --- a/relax/backends/megatron/model_provider.py +++ b/relax/backends/megatron/model_provider.py @@ -264,6 +264,7 @@ def wrapped_model_provider( "recompute_granularity", "recompute_method", "recompute_num_layers", + "recompute_modules", "distribute_saved_activations", "moe_router_load_balancing_type", "moe_router_dtype", diff --git a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py index 3504d99be..0ac4de892 100644 --- a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py +++ b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py @@ -7,9 +7,16 @@ import torch import torch.distributed as dist -from megatron.bridge.peft.lora import LoRAMerge from megatron.core import mpu + +try: + # Megatron-Bridge >= 0.6.0 moved LoRAMerge out of peft.lora into its own module. + from megatron.bridge.peft.lora_merge import LoRAMerge +except ImportError: # bridge <= 0.5.x + from megatron.bridge.peft.lora import LoRAMerge + + from relax.utils import device as device_utils from relax.utils.logging_utils import get_logger from relax.utils.megatron_peft_utils import ( diff --git a/relax/distributed/checkpoint_service/backends/device_direct.py b/relax/distributed/checkpoint_service/backends/device_direct.py index fa7fdb0bf..b23dddf1a 100644 --- a/relax/distributed/checkpoint_service/backends/device_direct.py +++ b/relax/distributed/checkpoint_service/backends/device_direct.py @@ -1151,7 +1151,12 @@ def _merge_full_base(self, name: str, param: torch.Tensor) -> torch.Tensor: slot = self._lora_adapter_full.get(_base_param_prefix(name)) if self._lora_adapter_full else None if not slot or "in" not in slot or "out" not in slot: return param - from megatron.bridge.peft.lora import LoRAMerge + + try: + # Megatron-Bridge >= 0.6.0 moved LoRAMerge into its own module. + from megatron.bridge.peft.lora_merge import LoRAMerge + except ImportError: # bridge <= 0.5.x + from megatron.bridge.peft.lora import LoRAMerge linear_in = slot["in"].float() linear_out = slot["out"].float() diff --git a/relax/models/dots_ocr/megatron/bridge.py b/relax/models/dots_ocr/megatron/bridge.py index 9f0fa09ee..1d55768fc 100644 --- a/relax/models/dots_ocr/megatron/bridge.py +++ b/relax/models/dots_ocr/megatron/bridge.py @@ -5,7 +5,7 @@ from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping, QKVMapping, ReplicatedMapping from megatron.bridge.models.conversion.transformers_compat import rope_theta_from_hf -from megatron.bridge.models.hf_pretrained.vlm import PreTrainedVLM +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM from relax.models.dots_ocr.configuration import DotsVisionConfig from relax.models.dots_ocr.megatron.model import DotsOCRModel @@ -17,7 +17,7 @@ target=DotsOCRModel, ) class DotsOCRBridge(MegatronModelBridge): - def provider_bridge(self, hf_pretrained: PreTrainedVLM) -> DotsOCRModelProvider: + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> DotsOCRModelProvider: hf_config = hf_pretrained.config model_dtype = self.dtype_from_hf(hf_config, default=torch.float32) vision_config = hf_config.vision_config diff --git a/relax/models/qwen_omni/qwen3_omni_bridge.py b/relax/models/qwen_omni/qwen3_omni_bridge.py index c3e330c2d..380767576 100644 --- a/relax/models/qwen_omni/qwen3_omni_bridge.py +++ b/relax/models/qwen_omni/qwen3_omni_bridge.py @@ -6,7 +6,7 @@ from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping, QKVMapping, ReplicatedMapping -from megatron.bridge.models.hf_pretrained.vlm import PreTrainedVLM +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM from transformers import Qwen3OmniMoeForConditionalGeneration from relax.models.qwen_omni.modeling_qwen3_omni.model import Qwen3OmniMoeModel @@ -42,7 +42,7 @@ def __init__(self): super().__init__() self.hf_weights_cache = {} - def provider_bridge(self, hf_pretrained: PreTrainedVLM) -> Qwen3OmniModelProvider: + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM) -> Qwen3OmniModelProvider: """Create a Qwen3OmniModelProvider from a HuggingFace pretrained MoE model. From 97ba6f829da99903719d7d358aa6f62d1178ef42 Mon Sep 17 00:00:00 2001 From: root Date: Fri, 14 Aug 2026 00:11:46 +0800 Subject: [PATCH 10/34] feat(sft): add sharded TQ producers - Add RELAX_SFT_TQ_SHARDS helpers and SFT shard partition naming. - Run multiple remote shard producers for async prepack when more than one shard is requested. - Keep the single-shard path local so existing behavior stays unchanged by default. - Fetch shard partitions rank-locally in the Megatron prepack path and concatenate them before training consumption. - Patch Qwen chat templates with generation markers so tokenizer-aware assistant masks avoid the slow fallback path. - Keep async prefetch from falling back to foreground work and support tensor-list conversion into NestedTensor batches. --- - Add coverage for shard partition helpers, SFT remote shard producer control flow, Qwen template patching, async prefetch behavior, and tensor-list NestedTensor conversion. (cherry picked from commit ef126e597835021c4dc1f8dd8dadf4279d8c8c72) --- relax/backends/megatron/actor.py | 94 ++- relax/components/actor.py | 37 +- relax/components/sft.py | 568 +++++++++++++++--- relax/engine/sft/dataset/chat_template.py | 21 + .../sft/dataset/qwen_chat_template_patch.py | 108 ++++ relax/engine/sft/dataset/streaming.py | 68 ++- relax/engine/sft/runtime.py | 31 + relax/utils/data/streaming_dataset.py | 98 ++- relax/utils/utils.py | 14 +- tests/components/test_sft.py | 81 ++- .../engine/sft/dataset/test_chat_template.py | 2 + .../dataset/test_qwen_chat_template_patch.py | 177 +++++- tests/engine/sft/dataset/test_streaming.py | 135 ++++- tests/utils/data/test_streaming_dataset.py | 67 +++ 14 files changed, 1342 insertions(+), 159 deletions(-) diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py index f493793a3..7b0a8f258 100644 --- a/relax/backends/megatron/actor.py +++ b/relax/backends/megatron/actor.py @@ -33,7 +33,9 @@ from relax.engine.sft.runtime import ( is_sft_mode, sft_partition_id, + sft_partition_ids, sft_task_name, + sft_tq_num_shards, should_run_sft_eval, should_run_sft_predict, should_skip_mtp_only_weight_management, @@ -830,11 +832,12 @@ def _run_step_evaluation(self, rollout_id: int, *, end_update_weight: bool = Fal try: if should_run_eval: if dist.get_rank() == 0: - run( - self.data_system_client.async_clear_partition( - partition_id=sft_partition_id(self.args, rollout_id) + for partition_id in sft_partition_ids(self.args, rollout_id): + run( + self.data_system_client.async_clear_partition( + partition_id=partition_id, + ) ) - ) dist.barrier(group=get_gloo_group()) run_sft_eval(self, rollout_id) @@ -1174,23 +1177,72 @@ def _fetch_sft_prepack_rollout_once( # backoff can turn a small producer skew into multi-second stalls. dp_rank = mpu.get_data_parallel_rank(with_context_parallel=False) batch_size = self._sft_prepack_local_batch_size() - partition_id = sft_partition_id(self.args, rollout_id) - sampling_config = {"dp_rank": dp_rank, "task_name": task_name} - rollout_data, _batch_meta = get_data_from_transfer_queue( - args=self.args, - tq_client=self.data_system_client, - data_fields=data_fields, - batch_size=batch_size, - partition_id=partition_id, - task_name=task_name, - sampling_config=sampling_config, - batch_index=0, - broadcast_pp=False, - per_rank_fetch=True, - post_process=False, - synchronize_per_rank_fetch=False, - ) - return rollout_data + partition_ids = sft_partition_ids(self.args, rollout_id) + num_shards = sft_tq_num_shards(self.args) + if len(partition_ids) != num_shards: + raise RuntimeError( + f"SFT shard partition mismatch for rollout_id={rollout_id}: " + f"partition_ids={partition_ids}, num_shards={num_shards}." + ) + + if num_shards <= 1: + partition_id = partition_ids[0] + sampling_config = {"dp_rank": dp_rank, "task_name": task_name} + rollout_data, _batch_meta = get_data_from_transfer_queue( + args=self.args, + tq_client=self.data_system_client, + data_fields=data_fields, + batch_size=batch_size, + partition_id=partition_id, + task_name=task_name, + sampling_config=sampling_config, + batch_index=0, + broadcast_pp=False, + per_rank_fetch=True, + post_process=False, + synchronize_per_rank_fetch=False, + ) + return rollout_data + + if batch_size % num_shards != 0: + raise ValueError( + "RELAX_SFT_TQ_SHARDS requires each DP local SFT batch to be divisible by shard count, " + f"got local_batch_size={batch_size}, num_shards={num_shards}." + ) + + # Avoid partially consuming shard 0 while shard N is not produced yet. + partitions = run(self.data_system_client.async_get_partition_list()) + if partitions is None or any(partition_id not in partitions for partition_id in partition_ids): + return None + + shard_batch_size = batch_size // num_shards + shard_batches: list[RolloutBatch] = [] + for shard_id, partition_id in enumerate(partition_ids): + sampling_config = {"dp_rank": dp_rank, "task_name": task_name} + rollout_data, _batch_meta = get_data_from_transfer_queue( + args=self.args, + tq_client=self.data_system_client, + data_fields=data_fields, + batch_size=shard_batch_size, + partition_id=partition_id, + task_name=task_name, + sampling_config=sampling_config, + batch_index=0, + broadcast_pp=False, + per_rank_fetch=True, + post_process=False, + synchronize_per_rank_fetch=False, + ) + if rollout_data is None: + if shard_id > 0: + raise RuntimeError( + f"SFT shard fetch split for rollout_id={rollout_id}: shard {shard_id} returned no data " + "after earlier shards were consumed. Check producer partition readiness." + ) + return None + shard_batches.append(rollout_data) + + return concat_rollout_batches(shard_batches) def _pack_sft_prepack_window( self, diff --git a/relax/components/actor.py b/relax/components/actor.py index c915cf90f..20529fc7b 100644 --- a/relax/components/actor.py +++ b/relax/components/actor.py @@ -14,7 +14,7 @@ from relax.components.base import Base from relax.distributed.coordination import PeerStepBarrier, RolloutOffloadBarrier from relax.distributed.ray.placement_group import allocate_train_group -from relax.engine.sft.runtime import is_sft_mode, sft_partition_id, sft_task_name +from relax.engine.sft.runtime import is_sft_mode, sft_partition_ids, sft_task_name from relax.utils.async_utils import run from relax.utils.opd.opd_utils import set_managed_opd_teacher_on_train_group @@ -163,10 +163,11 @@ async def run(self) -> None: if self._done_event is not None: await self._done_event.wait() return - self.data_system_client.reset_consumption( - partition_id=sft_partition_id(self.config, self.step), - task_name=sft_task_name(self.config, component="actor"), - ) + for partition_id in sft_partition_ids(self.config, self.step): + self.data_system_client.reset_consumption( + partition_id=partition_id, + task_name=sft_task_name(self.config, component="actor"), + ) # Create an asyncio.Event bound to the current event loop so the # background thread can signal completion without blocking the loop. loop = asyncio.get_running_loop() @@ -219,11 +220,8 @@ def _background_run(self) -> None: self._logger.info(f"Actor training completed step {local_step}/{self.config.num_rollout}") if did_train: - run( - self.data_system_client.async_clear_partition( - partition_id=sft_partition_id(self.config, local_step) - ) - ) + for partition_id in sft_partition_ids(self.config, local_step): + run(self.data_system_client.async_clear_partition(partition_id=partition_id)) self._logger.info(f"Actor cleared data for step {local_step}/{self.config.num_rollout}") try: @@ -249,9 +247,9 @@ def _wait_for_rollout_data(self) -> bool: True if data is ready and training can proceed, False if should continue waiting (caller should skip this iteration) """ - partition_id = sft_partition_id(self.config, self.step) + partition_ids = sft_partition_ids(self.config, self.step) partition_list = run(self.data_system_client.async_get_partition_list()) - if partition_list is None or partition_id not in partition_list: + if partition_list is None or any(partition_id not in partition_list for partition_id in partition_ids): time.sleep(1) return False @@ -354,10 +352,10 @@ def train(self, step: int, clear_data: bool = True) -> Dict[str, Any]: try: # Check if rollout data is available for this step - partition_id = sft_partition_id(self.config, step) + partition_ids = sft_partition_ids(self.config, step) partition_list = run(self.data_system_client.async_get_partition_list()) - if partition_list is not None and partition_id in partition_list: + if partition_list is not None and all(partition_id in partition_list for partition_id in partition_ids): self._logger.info(f"Data available for step {step}, executing training") # Execute training @@ -365,12 +363,13 @@ def train(self, step: int, clear_data: bool = True) -> Dict[str, Any]: # Only clear partition data if clear_data is True if clear_data and did_train: - run(self.data_system_client.async_clear_partition(partition_id=partition_id)) - self._logger.info(f"Cleared data partition: {partition_id}") + for partition_id in partition_ids: + run(self.data_system_client.async_clear_partition(partition_id=partition_id)) + self._logger.info(f"Cleared data partitions: {partition_ids}") elif clear_data: - self._logger.info(f"Skipped clearing partition after skipped actor step: {partition_id}") + self._logger.info(f"Skipped clearing partitions after skipped actor step: {partition_ids}") else: - self._logger.info(f"Keeping data partition (clear_data=False): {partition_id}") + self._logger.info(f"Keeping data partitions (clear_data=False): {partition_ids}") metrics["data_consumed"] = True metrics["elapsed_time"] = time.time() - start_time @@ -379,7 +378,7 @@ def train(self, step: int, clear_data: bool = True) -> Dict[str, Any]: self._logger.warning(f"No data available for step {step}, skipping training") metrics["data_consumed"] = False metrics["success"] = True - metrics["message"] = f"No data in partition {partition_id}" + metrics["message"] = f"No data in partitions {partition_ids}" except Exception as e: self._logger.error(f"Training failed at step {step}: {e}") diff --git a/relax/components/sft.py b/relax/components/sft.py index e4145e6e1..9b10a34bf 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -29,6 +29,7 @@ import random from typing import Any +import ray import torch.nn.functional as F import transfer_queue as tq from ray import serve @@ -37,8 +38,14 @@ from relax.components.base import Base from relax.engine.sft.dataset.streaming import ProcessedSample, SFTStreamingDataset, pack_samples_for_tq from relax.engine.sft.debug_print import print_first_sample -from relax.engine.sft.runtime import resolve_sft_split_indices +from relax.engine.sft.runtime import ( + resolve_sft_split_indices, + sft_logical_partition_id, + sft_partition_ids, + sft_tq_num_shards, +) from relax.utils.data.processor_pool import ProcessorPool +from relax.utils.logging_utils import get_logger from relax.utils.misc import load_function from relax.utils.s3_model_loader import prepare_model_maybe_update_args from relax.utils.training.eval_config import build_named_prompt_data_configs @@ -82,12 +89,370 @@ def _resolve_classification_sentinel_token_id(tokenizer) -> int: raise ValueError("--task-type seq_cls requires a tokenizer with a valid EOS or PAD token id.") +def _resolve_sft_dataset_options(config: Any, logger: Any | None = None) -> dict[str, Any]: + oversize_strategy = getattr(config, "sft_oversize_strategy", "keep") + invalid_multimodal_strategy = getattr(config, "sft_invalid_multimodal_strategy", "error") + oversize_custom_path = getattr(config, "sft_oversize_custom_function_path", None) + oversize_custom_fn = None + if oversize_strategy == "custom": + if not oversize_custom_path: + raise ValueError("--sft-oversize-strategy custom requires --sft-oversize-custom-function-path.") + oversize_custom_fn = load_function(oversize_custom_path) + if logger is not None: + logger.info(f"SFT oversize strategy: custom (loaded {oversize_custom_path})") + elif logger is not None: + logger.info(f"SFT oversize strategy: {oversize_strategy}") + if logger is not None: + logger.info(f"SFT invalid multimodal strategy: {invalid_multimodal_strategy}") + return { + "oversize_strategy": oversize_strategy, + "invalid_multimodal_strategy": invalid_multimodal_strategy, + "oversize_custom_fn": oversize_custom_fn, + } + + +def _create_sft_train_dataset( + config: Any, + *, + tokenizer: Any, + processor_pool: ProcessorPool | None, + capacity: int, + prefetch_buffer_size: int, + prefetch_chunk_size: int, + prefetch_num_workers: int, + pad_token_ids: frozenset[int], + oversize_strategy: str, + oversize_custom_fn: Any, + invalid_multimodal_strategy: str, + task_type: str, + classification_sentinel_token_id: int | None, +) -> Any: + dataset_cls = _load_custom_dataset_class(getattr(config, "custom_dataset_class_path", None)) + if dataset_cls is None: + return SFTStreamingDataset( + path=config.prompt_data, + tokenizer=tokenizer, + processor_pool=processor_pool, + capacity=capacity, + prompt_key=config.input_key, + label_key=config.label_key, + multimodal_keys=config.multimodal_keys, + conversation_key_map=getattr(config, "conversation_key_map", None), + metadata_key=config.metadata_key, + tool_key=config.tool_key, + system_prompt=config.system_prompt, + seed=getattr(config, "seed", 42), + prefetch_max_cached=prefetch_buffer_size, + prefetch_chunk_size=prefetch_chunk_size, + prefetch_num_workers=prefetch_num_workers, + pad_token_ids=pad_token_ids, + oversize_strategy=oversize_strategy, + oversize_custom_fn=oversize_custom_fn, + invalid_multimodal_strategy=invalid_multimodal_strategy, + apply_chat_template_kwargs=getattr(config, "apply_chat_template_kwargs", None), + require_response=task_type != "seq_cls", + task_type=task_type, + num_labels=getattr(config, "num_labels", None), + problem_type=getattr(config, "problem_type", "single_label_classification"), + classification_sentinel_token_id=classification_sentinel_token_id, + ) + return dataset_cls.from_args( + config, + tokenizer=tokenizer, + processor_pool=processor_pool, + pad_token_ids=pad_token_ids, + ) + + +def _prepare_sft_tq_payload(samples: list[ProcessedSample], *, force_multimodal_field: bool) -> dict[str, Any]: + backend_batch = pack_samples_for_tq(samples, force_multimodal_field=force_multimodal_field) + assert backend_batch is not None + return { + "data": dict_to_tensordict(backend_batch, batch_size=len(backend_batch["tokens"])), + "custom_meta": [{"total_lengths": int(length)} for length in backend_batch["total_lengths"]], + } + + +def _split_sft_samples_for_shards(samples: list[ProcessedSample], num_shards: int) -> list[list[ProcessedSample]]: + if num_shards <= 1: + return [samples] + if len(samples) % num_shards != 0: + raise ValueError( + f"RELAX_SFT_TQ_SHARDS={num_shards} requires global_batch_size divisible by shard count; " + f"got {len(samples)} samples." + ) + shard_size = len(samples) // num_shards + return [samples[i * shard_size : (i + 1) * shard_size] for i in range(num_shards)] + + +def _sft_train_partitions_in_flight(partitions: list[str]) -> int: + logical_partitions = { + sft_logical_partition_id(partition) + for partition in partitions + if partition.startswith("sft_") and not partition.startswith("sft_eval_") + } + return len(logical_partitions) + + +def _sft_prefetch_workers_per_shard(config: Any, num_shards: int) -> int: + prefetch_num_workers = max(1, int(getattr(config, "sft_prefetch_num_workers", 4) or 1)) + if num_shards <= 1: + return prefetch_num_workers + return max(1, (prefetch_num_workers + num_shards - 1) // num_shards) + + +def _validate_sft_shard_count(config: Any, num_shards: int) -> None: + if num_shards <= 1: + return + global_batch_size = int(getattr(config, "global_batch_size", 0) or 0) + if global_batch_size % num_shards != 0: + raise ValueError( + f"RELAX_SFT_TQ_SHARDS={num_shards} requires global_batch_size divisible by shard count; " + f"got global_batch_size={global_batch_size}." + ) + + +def _sft_batch_producer_actor_options(config: Any, runtime_env: Any | None, num_shards: int) -> dict[str, Any]: + options: dict[str, Any] = {"num_cpus": _sft_prefetch_workers_per_shard(config, num_shards)} + if runtime_env is not None: + options["runtime_env"] = runtime_env + return options + + +async def _ray_get_many_async(refs: list[Any]) -> list[Any]: + if not refs: + return [] + loop = asyncio.get_running_loop() + return await loop.run_in_executor(None, ray.get, refs) + + +def _set_sft_dataset_position(dataset: Any, epoch: int, position: int, prefetch_limit: int | None = None) -> None: + index_manager = getattr(dataset, "index_manager", None) + if index_manager is None: + dataset.shuffle(epoch, position=position) + return + total_size = getattr(index_manager, "total_size", None) + if not isinstance(total_size, int) or total_size <= 0: + dataset.shuffle(epoch, position=position) + return + index_manager.shuffle(epoch) + target_position = min(position, total_size) + index_manager.position = target_position + prefetch = getattr(dataset, "_prefetch", None) + indices = getattr(index_manager, "indices", None) + if prefetch is not None and indices is not None: + remaining = list(indices[target_position:]) + if prefetch_limit is not None: + remaining = remaining[: max(0, prefetch_limit)] + prefetch.set_index_order(remaining) + + +def _has_sft_eval_work(config: Any) -> bool: + return getattr(config, "eval_size", None) is not None or bool( + build_named_prompt_data_configs(getattr(config, "eval_prompt_data", None)) + ) + + +@ray.remote +class _SFTBatchProducerActor: + """Own one SFT data shard and write it directly to TransferQueue.""" + + def __init__( + self, + config: Any, + shard_id: int, + num_shards: int, + prefetch_num_workers: int | None = None, + ): + self.config = config + self._shard_id = shard_id + self._num_shards = max(1, num_shards) + self._prefetch_num_workers = prefetch_num_workers + self._logger = get_logger(__name__) + self._dataset: Any | None = None + self.data_system_client: Any | None = None + self._tokenizer = None + self._processor_pool: ProcessorPool | None = None + self._train_size = 0 + + def _shard_batch_size(self, global_batch_size: int) -> int: + if global_batch_size % self._num_shards != 0: + raise ValueError( + f"RELAX_SFT_TQ_SHARDS={self._num_shards} requires global_batch_size divisible by shard count; " + f"got global_batch_size={global_batch_size}." + ) + return global_batch_size // self._num_shards + + def _seek_shard_for_step(self, step: int, global_batch_size: int) -> None: + assert self._dataset is not None + if self._train_size <= 0: + return + shard_batch_size = self._shard_batch_size(global_batch_size) + consumed = step * global_batch_size + self._shard_id * shard_batch_size + epoch = consumed // self._train_size + position = consumed % self._train_size + index_manager = getattr(self._dataset, "index_manager", None) + if ( + index_manager is not None + and getattr(index_manager, "current_epoch", None) == epoch + and getattr(index_manager, "position", None) == position + ): + return + _set_sft_dataset_position(self._dataset, epoch, position, prefetch_limit=shard_batch_size) + + def initialize(self, start_step: int) -> dict[str, Any]: + if self._dataset is not None: + return self._state() + + tq.init(self.config.tq_config) + self.data_system_client = tq.get_client() + + prepare_model_maybe_update_args(self.config, completeness="metadata") + prefetch_num_workers = ( + self._prefetch_num_workers + if self._prefetch_num_workers is not None + else _sft_prefetch_workers_per_shard(self.config, self._num_shards) + ) + self._tokenizer = AutoTokenizer.from_pretrained(self.config.hf_checkpoint, trust_remote_code=True) + try: + self._processor_pool = ProcessorPool( + self.config.hf_checkpoint, + pool_size=prefetch_num_workers, + trust_remote_code=True, + ) + except Exception as exc: + self._logger.warning(f"Could not init ProcessorPool ({exc}); multimodal samples will fail at push.") + self._processor_pool = None + + pad_token_ids = _resolve_pad_token_ids_from_config(self.config.hf_checkpoint) + cp_size = max(1, getattr(self.config, "context_parallel_size", 1) or 1) + capacity = self.config.max_tokens_per_gpu * cp_size + dataset_options = _resolve_sft_dataset_options(self.config) + task_type = getattr(self.config, "task_type", "causal_lm") + classification_sentinel_token_id = ( + _resolve_classification_sentinel_token_id(self._tokenizer) if task_type == "seq_cls" else None + ) + + self._dataset = _create_sft_train_dataset( + self.config, + tokenizer=self._tokenizer, + processor_pool=self._processor_pool, + capacity=capacity, + prefetch_buffer_size=getattr(self.config, "sft_prefetch_buffer_size", 256), + prefetch_chunk_size=getattr(self.config, "sft_prefetch_chunk_size", 32), + prefetch_num_workers=prefetch_num_workers, + pad_token_ids=pad_token_ids, + task_type=task_type, + classification_sentinel_token_id=classification_sentinel_token_id, + **dataset_options, + ) + n_avail = len(self._dataset) + self._train_size = n_avail + eval_size_arg = getattr(self.config, "eval_size", None) + if eval_size_arg is not None: + if eval_size_arg < 1: + n_eval = max(1, int(n_avail * eval_size_arg)) + else: + n_eval = int(eval_size_arg) + n_eval = min(n_eval, max(n_avail - 1, 0)) + self._train_size = n_avail - n_eval if n_eval > 0 else n_avail + + shard_batch_size = self._shard_batch_size(self.config.global_batch_size) + if self._train_size > 0: + consumed = start_step * self.config.global_batch_size + self._shard_id * shard_batch_size + start_epoch = consumed // self._train_size + position = consumed % self._train_size + else: + start_epoch, position = 0, 0 + _set_sft_dataset_position(self._dataset, start_epoch, position, prefetch_limit=shard_batch_size) + self._logger.info( + f"SFT remote batch producer initialized: train_size={self._train_size} " + f"prefetch_num_workers={prefetch_num_workers} shard={self._shard_id}/{self._num_shards}" + ) + return self._state() + + def _state(self) -> dict[str, Any]: + return { + "train_size": self._train_size, + "shard_id": self._shard_id, + "num_shards": self._num_shards, + "dataset": f"{self._dataset.__class__.__module__}.{self._dataset.__class__.__name__}" + if self._dataset is not None + else None, + } + + async def produce_partition( + self, + step: int, + partition_id: str, + global_batch_size: int, + force_multimodal_field: bool, + ) -> dict[str, Any]: + assert self._dataset is not None and self._tokenizer is not None + assert self.data_system_client is not None + if self._train_size == 0: + raise RuntimeError("SFT train pool is empty (check --eval-size relative to dataset size).") + + self._seek_shard_for_step(step, global_batch_size) + batch_size = self._shard_batch_size(global_batch_size) + samples, crossed_epoch = await self._dataset.get_batch_async(batch_size) + if len(samples) != batch_size: + raise RuntimeError( + f"SFT step {step} shard {self._shard_id}/{self._num_shards}: " + f"dataset returned {len(samples)}/{batch_size} samples " + "after bounded refill attempts. Refusing to push a partial TQ partition because the Megatron " + "consumer requires a full global batch. Check invalid-multimodal and oversize skip warnings." + ) + + if step == 0 and self._shard_id == 0 and samples: + s = samples[0] + try: + print_first_sample( + step=step, + sample_idx=s.source_idx, + input_ids=s.tokens, + loss_mask=s.loss_mask, + multimodal_train_inputs=s.multimodal_train_inputs, + tokenizer=self._tokenizer, + ) + except Exception as exc: + self._logger.warning(f"print_first_sample failed: {exc}") + + payload = _prepare_sft_tq_payload(samples, force_multimodal_field=force_multimodal_field) + await self.data_system_client.async_put( + data=payload["data"], + partition_id=partition_id, + custom_meta=payload["custom_meta"], + ) + return { + "status": "ok", + "partition_id": partition_id, + "shard_id": self._shard_id, + "num_shards": self._num_shards, + "crossed_epoch": crossed_epoch, + "epoch": getattr(getattr(self._dataset, "index_manager", None), "current_epoch", None), + "samples": len(samples), + } + + async def stop(self) -> None: + if self._dataset is not None: + self._dataset.stop() + if self._processor_pool is not None: + close = getattr(self._processor_pool, "close", None) + shutdown = getattr(self._processor_pool, "shutdown", None) + if callable(close): + close() + elif callable(shutdown): + shutdown() + + @serve.deployment class SFT(Base): def __init__(self, healthy, pgs, num_gpus, config, role, runtime_env=None): # noqa: ARG002 super().__init__() self.config = config self.role = role + self._runtime_env = runtime_env self.healthy = healthy self.step = getattr(config, "start_rollout_id", 0) @@ -97,14 +462,75 @@ def __init__(self, healthy, pgs, num_gpus, config, role, runtime_env=None): # n self._dataset: Any | None = None self._eval_dataset: Any | None = None self._eval_indices: tuple[int, ...] | None = None + self._batch_producers: list[Any] = [] self._train_size: int = 0 self._tokenizer = None self._processor_pool: ProcessorPool | None = None self._stop_event = asyncio.Event() self._run_task: asyncio.Task | None = None + def _should_use_remote_batch_producer(self) -> bool: + num_shards = sft_tq_num_shards(self.config) + _validate_sft_shard_count(self.config, num_shards) + if num_shards <= 1: + return False + if not getattr(self.config, "sft_async_prepack", False): + return False + if getattr(self.config, "task_type", "causal_lm") == "seq_cls": + self._logger.info( + "SFT remote shard producer disabled: sequence classification uses the local producer path." + ) + return False + if not ray.is_initialized(): + self._logger.info("SFT remote shard producer disabled: Ray is not initialized.") + return False + if _has_sft_eval_work(self.config): + self._logger.info("SFT remote shard producer disabled: eval is configured; using local producer path.") + return False + if getattr(self.config, "custom_dataset_class_path", None): + self._logger.info( + "SFT remote shard producer disabled: custom dataset is configured; using local producer path." + ) + return False + oversize_strategy = getattr(self.config, "sft_oversize_strategy", "keep") + invalid_multimodal_strategy = getattr(self.config, "sft_invalid_multimodal_strategy", "error") + if oversize_strategy in {"skip", "custom"} or invalid_multimodal_strategy == "skip": + self._logger.info( + "SFT remote shard producer disabled: skip-capable data filtering is configured " + f"(oversize_strategy={oversize_strategy}, invalid_multimodal_strategy={invalid_multimodal_strategy}). " + "Using local coordinator path to preserve sample order." + ) + return False + return True + + def _init_remote_batch_producers(self) -> None: + if self._batch_producers: + return + num_shards = sft_tq_num_shards(self.config) + prefetch_num_workers = _sft_prefetch_workers_per_shard(self.config, num_shards) + options = _sft_batch_producer_actor_options(self.config, self._runtime_env, num_shards) + self._batch_producers = [ + _SFTBatchProducerActor.options(**options).remote( + self.config, + shard_id, + num_shards, + prefetch_num_workers, + ) + for shard_id in range(num_shards) + ] + states = ray.get([producer.initialize.remote(self.step) for producer in self._batch_producers]) + self._train_size = int(states[0].get("train_size") or 0) + self._logger.info( + f"SFT remote shard producer enabled: dataset={states[0].get('dataset')} " + f"train_size={self._train_size} shards={num_shards} " + f"prefetch_workers_per_shard={prefetch_num_workers}" + ) + def _init_data_pipeline(self) -> None: - if self._dataset is not None: + if self._dataset is not None or getattr(self, "_batch_producers", None): + return + if self._should_use_remote_batch_producer(): + self._init_remote_batch_producers() return prepare_model_maybe_update_args(self.config, completeness="metadata") self._tokenizer = AutoTokenizer.from_pretrained(self.config.hf_checkpoint, trust_remote_code=True) @@ -127,55 +553,20 @@ def _init_data_pipeline(self) -> None: _resolve_classification_sentinel_token_id(self._tokenizer) if task_type == "seq_cls" else None ) - oversize_strategy = getattr(self.config, "sft_oversize_strategy", "keep") - invalid_multimodal_strategy = getattr(self.config, "sft_invalid_multimodal_strategy", "error") - oversize_custom_path = getattr(self.config, "sft_oversize_custom_function_path", None) - oversize_custom_fn = None - if oversize_strategy == "custom": - if not oversize_custom_path: - raise ValueError("--sft-oversize-strategy custom requires --sft-oversize-custom-function-path.") - oversize_custom_fn = load_function(oversize_custom_path) - self._logger.info(f"SFT oversize strategy: custom (loaded {oversize_custom_path})") - else: - self._logger.info(f"SFT oversize strategy: {oversize_strategy}") - self._logger.info(f"SFT invalid multimodal strategy: {invalid_multimodal_strategy}") - - dataset_cls = _load_custom_dataset_class(getattr(self.config, "custom_dataset_class_path", None)) - if dataset_cls is None: - self._dataset = SFTStreamingDataset( - path=self.config.prompt_data, - tokenizer=self._tokenizer, - processor_pool=self._processor_pool, - capacity=capacity, - prompt_key=self.config.input_key, - label_key=self.config.label_key, - multimodal_keys=self.config.multimodal_keys, - conversation_key_map=getattr(self.config, "conversation_key_map", None), - metadata_key=self.config.metadata_key, - tool_key=self.config.tool_key, - system_prompt=self.config.system_prompt, - seed=seed, - prefetch_max_cached=prefetch_buffer_size, - prefetch_chunk_size=prefetch_chunk_size, - prefetch_num_workers=prefetch_num_workers, - pad_token_ids=pad_token_ids, - oversize_strategy=oversize_strategy, - oversize_custom_fn=oversize_custom_fn, - invalid_multimodal_strategy=invalid_multimodal_strategy, - apply_chat_template_kwargs=getattr(self.config, "apply_chat_template_kwargs", None), - require_response=task_type != "seq_cls", - task_type=task_type, - num_labels=getattr(self.config, "num_labels", None), - problem_type=getattr(self.config, "problem_type", "single_label_classification"), - classification_sentinel_token_id=classification_sentinel_token_id, - ) - else: - self._dataset = dataset_cls.from_args( - self.config, - tokenizer=self._tokenizer, - processor_pool=self._processor_pool, - pad_token_ids=pad_token_ids, - ) + dataset_options = _resolve_sft_dataset_options(self.config, self._logger) + self._dataset = _create_sft_train_dataset( + self.config, + tokenizer=self._tokenizer, + processor_pool=self._processor_pool, + capacity=capacity, + prefetch_buffer_size=prefetch_buffer_size, + prefetch_chunk_size=prefetch_chunk_size, + prefetch_num_workers=prefetch_num_workers, + pad_token_ids=pad_token_ids, + task_type=task_type, + classification_sentinel_token_id=classification_sentinel_token_id, + **dataset_options, + ) n_avail = len(self._dataset) self._train_size = n_avail @@ -228,9 +619,9 @@ def _init_data_pipeline(self) -> None: seed=seed, prefetch_max_cached=0, pad_token_ids=pad_token_ids, - oversize_strategy=oversize_strategy, - oversize_custom_fn=oversize_custom_fn, - invalid_multimodal_strategy=invalid_multimodal_strategy, + oversize_strategy=dataset_options["oversize_strategy"], + oversize_custom_fn=dataset_options["oversize_custom_fn"], + invalid_multimodal_strategy=dataset_options["invalid_multimodal_strategy"], apply_chat_template_kwargs=getattr(self.config, "apply_chat_template_kwargs", None), require_response=task_type != "seq_cls", task_type=task_type, @@ -305,7 +696,7 @@ async def _wait_for_buffer_capacity(self) -> None: partitions = await self.data_system_client.async_get_partition_list() if partitions is None: return - in_flight = sum(1 for p in partitions if p.startswith("sft_") and not p.startswith("sft_eval_")) + in_flight = _sft_train_partitions_in_flight(partitions) if in_flight < max_in_flight: if wait_count > 0: self._logger.info( @@ -349,11 +740,39 @@ def _maybe_print_first_sample(self, samples: list[ProcessedSample]) -> None: self._logger.warning(f"print_first_sample failed: {exc}") async def _produce_one_step(self) -> None: - assert self._dataset is not None and self._tokenizer is not None + batch_producers = getattr(self, "_batch_producers", []) + assert batch_producers or (self._dataset is not None and self._tokenizer is not None) await self._wait_for_buffer_capacity() if self._train_size == 0: raise RuntimeError("SFT train pool is empty (check --eval-size relative to dataset size).") + partition_ids = sft_partition_ids(self.config, self.step) + num_shards = len(partition_ids) + if batch_producers: + if len(batch_producers) != num_shards: + raise RuntimeError( + f"SFT remote producer shard count mismatch: actors={len(batch_producers)}, " + f"partitions={num_shards}." + ) + payloads = await _ray_get_many_async( + [ + producer.produce_partition.remote( + self.step, + partition_id, + self.config.global_batch_size, + self.config.multimodal_keys is not None, + ) + for producer, partition_id in zip(batch_producers, partition_ids, strict=True) + ] + ) + crossed_epoch = any(bool(payload["crossed_epoch"]) for payload in payloads) + current_epoch = max(payload.get("epoch") or 0 for payload in payloads) + if crossed_epoch: + self._logger.info(f"SFT step {self.step}: epoch boundary crossed (epoch={current_epoch})") + self.step += 1 + return + + assert self._dataset is not None # When prefetch is on, get_batch_async delegates to the sync prefetch # path (already parallel via background threads). When prefetch is off, # it parallelises multimodal preprocess via asyncio.gather over the pool. @@ -365,13 +784,20 @@ async def _produce_one_step(self) -> None: "consumer requires a full global batch. Check invalid-multimodal and oversize skip warnings." ) self._maybe_print_first_sample(samples) - backend_batch = pack_samples_for_tq(samples, force_multimodal_field=self.config.multimodal_keys is not None) - assert backend_batch is not None - await self.data_system_client.async_put( - data=dict_to_tensordict(backend_batch, batch_size=len(backend_batch["tokens"])), - partition_id=f"sft_{self.step}", - custom_meta=[{"total_lengths": int(length)} for length in backend_batch["total_lengths"]], - ) + for partition_id, shard_samples in zip( + partition_ids, + _split_sft_samples_for_shards(samples, num_shards), + strict=True, + ): + payload = _prepare_sft_tq_payload( + shard_samples, + force_multimodal_field=self.config.multimodal_keys is not None, + ) + await self.data_system_client.async_put( + data=payload["data"], + partition_id=partition_id, + custom_meta=payload["custom_meta"], + ) if crossed_epoch: self._logger.info( f"SFT step {self.step}: epoch boundary crossed (epoch={self._dataset.index_manager.current_epoch})" @@ -456,9 +882,10 @@ async def _maybe_produce_eval(self) -> None: assert backend_batch is not None n_samples = len(backend_batch["tokens"]) - # Drain the current train partition so the eval chunks have the full - # TQ capacity to themselves. - await self._wait_for_partition_drained(f"sft_{self.step}") + # Drain the current train partition(s) so the eval chunks have the + # full TQ capacity to themselves. + for partition_id in sft_partition_ids(self.config, self.step): + await self._wait_for_partition_drained(partition_id) chunk_size = self.config.global_batch_size # Causal-LM eval drops trailing samples that don't fill a full chunk; @@ -468,8 +895,7 @@ async def _maybe_produce_eval(self) -> None: # which returns size=0 when the partition has fewer than batch_size # samples, so a partial last chunk would never be marked consumed and # the actor's `while not all_consumed` loop would spin forever (it - # already burned a full eval round in the wild — see the - # `[get_data_profile] samples=0` log spam). + # already burned a full eval round in the wild). n_chunks = n_samples // chunk_size n_dropped = n_samples - n_chunks * chunk_size if n_chunks == 0: @@ -533,6 +959,12 @@ async def _async_run(self) -> None: async def stop(self) -> None: self._stop_event.set() + for producer in getattr(self, "_batch_producers", []): + try: + await _ray_get_many_async([producer.stop.remote()]) + except Exception as exc: + self._logger.warning(f"SFT remote batch producer stop failed: {exc}") + self._batch_producers = [] if self._dataset is not None: self._dataset.stop() if self._eval_dataset is not None: diff --git a/relax/engine/sft/dataset/chat_template.py b/relax/engine/sft/dataset/chat_template.py index 0c4cf6884..f1f027e03 100644 --- a/relax/engine/sft/dataset/chat_template.py +++ b/relax/engine/sft/dataset/chat_template.py @@ -12,6 +12,7 @@ import hashlib import re +import threading from collections.abc import Mapping from typing import Any @@ -60,6 +61,10 @@ def _render_with_assistant_mask( apply_chat_template_kwargs: dict | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Path 1: ask tokenizer for the assistant-only mask directly.""" + apply_chat_template_kwargs = _thread_local_chat_template_kwargs( + tokenizer, + apply_chat_template_kwargs, + ) result = tokenizer.apply_chat_template( _to_chat_messages(sample), tools=sample.tools, @@ -80,6 +85,22 @@ def _render_with_assistant_mask( return input_ids.long(), masks.long() +def _thread_local_chat_template_kwargs(tokenizer, apply_chat_template_kwargs: dict | None) -> dict: + """Avoid sharing HF AssistantTracker state across prefetch threads. + + Transformers caches compiled Jinja templates by the template string. The + compiled environment owns the ``AssistantTracker`` used by + ``return_assistant_tokens_mask=True``, and that tracker is not thread-safe. + Appending a Jinja comment makes each prefetch thread use a distinct + compiled environment without changing rendered text. + """ + kwargs = dict(apply_chat_template_kwargs or {}) + template = kwargs.get("chat_template") or getattr(tokenizer, "chat_template", None) + if isinstance(template, str): + kwargs["chat_template"] = f"{template}{{# relax_thread={threading.get_ident()} #}}" + return kwargs + + _THINK_OPEN = "\n" _IM_END = "<|im_end|>" # Qwen3.5 wraps tool messages inside a user block as diff --git a/relax/engine/sft/dataset/qwen_chat_template_patch.py b/relax/engine/sft/dataset/qwen_chat_template_patch.py index 6164615bc..880725732 100644 --- a/relax/engine/sft/dataset/qwen_chat_template_patch.py +++ b/relax/engine/sft/dataset/qwen_chat_template_patch.py @@ -3,6 +3,7 @@ """Qwen chat-template compatibility patches for SFT.""" import hashlib +import re from collections.abc import Mapping from functools import lru_cache from typing import Any @@ -23,6 +24,72 @@ _QWEN38_PRESERVE_HISTORY_GATE = ( "{%- if preserve_thinking is undefined or preserve_thinking is true or loop.index0 > ns.last_query_index %}" ) +_QWEN_VISIBLE_THINKING_SET = ( + "{%- set relax_has_visible_thinking = (preserve_thinking is defined and preserve_thinking is true) " + "or (loop.index0 > ns.last_query_index) %}" +) +_QWEN_ASSISTANT_RENDER_BLOCK = "\n".join( + ( + f" {_QWEN_PRESERVE_HISTORY_GATE}", + " {{- '<|im_start|>' + message.role + '\\n\\n' + reasoning_content + '\\n\\n\\n' + content }}", + " {%- else %}", + " {{- '<|im_start|>' + message.role + '\\n' + content }}", + " {%- endif %}", + ) +) +_QWEN_ASSISTANT_GENERATION_RENDER_BLOCK = "\n".join( + ( + f" {_QWEN_VISIBLE_THINKING_SET}", + " {%- if relax_has_visible_thinking %}", + " {{- '<|im_start|>' + message.role + '\\n\\n' }}", + " {%- else %}", + " {{- '<|im_start|>' + message.role + '\\n' }}", + " {%- endif %}", + " {%- generation %}", + " {%- if relax_has_visible_thinking %}", + " {{- reasoning_content + '\\n\\n\\n' + content }}", + " {%- else %}", + " {{- content }}", + " {%- endif %}", + ) +) +_QWEN_ASSISTANT_END = " {{- '<|im_end|>\\n' }}" +_QWEN_ASSISTANT_GENERATION_END = "\n".join((_QWEN_ASSISTANT_END, " {%- endgeneration %}")) +_QWEN_COMPACT_VISIBLE_THINKING_SET = ( + "{%- set relax_has_visible_thinking = " + "((preserve_thinking is defined and preserve_thinking is true) or (loop.index0 > ns.last_query_index)) " + "and (loop.last or (not loop.last and reasoning_content)) %}" +) +_QWEN_COMPACT_ASSISTANT_RENDER_BLOCK = "\n".join( + ( + f" {_QWEN_PRESERVE_HISTORY_GATE}", + " {%- if loop.last or (not loop.last and reasoning_content) %}", + " {{- '<|im_start|>' + message.role + '\\n\\n' + reasoning_content.strip('\\n') + '\\n\\n\\n' + content.lstrip('\\n') }}", + " {%- else %}", + " {{- '<|im_start|>' + message.role + '\\n' + content }}", + " {%- endif %}", + " {%- else %}", + " {{- '<|im_start|>' + message.role + '\\n' + content }}", + " {%- endif %}", + ) +) +_QWEN_COMPACT_ASSISTANT_GENERATION_RENDER_BLOCK = "\n".join( + ( + f" {_QWEN_COMPACT_VISIBLE_THINKING_SET}", + " {%- if relax_has_visible_thinking %}", + " {{- '<|im_start|>' + message.role + '\\n\\n' }}", + " {%- else %}", + " {{- '<|im_start|>' + message.role + '\\n' }}", + " {%- endif %}", + " {%- generation %}", + " {%- if relax_has_visible_thinking %}", + " {{- reasoning_content.strip('\\n') + '\\n\\n\\n' + content.lstrip('\\n') }}", + " {%- else %}", + " {{- content }}", + " {%- endif %}", + ) +) +_GENERATION_MARKER_RE = re.compile(r"{%-?\s*generation\s*-?%}") _PATCH_NAME = "qwen_history_thinking" @@ -65,6 +132,17 @@ def _has_learnable_historical_thinking(sample: CanonicalSample) -> bool: ) +def _assistant_generation_mask_matches_sample(sample: CanonicalSample) -> bool: + """Whether marking every rendered assistant block preserves learn flags.""" + for message in sample.messages: + if message.role == "assistant": + if not message.learn: + return False + elif message.learn: + return False + return True + + @lru_cache(maxsize=32) def _patch_qwen_history_gate(template: str) -> tuple[str, bool] | None: """Backport Qwen3.6's preserve gate to the exact Qwen3.5 gate.""" @@ -90,6 +168,32 @@ def _patch_qwen_history_gate(template: str) -> tuple[str, bool] | None: ) +@lru_cache(maxsize=32) +def _patch_qwen_generation_markers(template: str) -> tuple[str, bool]: + """Add HF generation markers around Qwen assistant output when safe.""" + if _GENERATION_MARKER_RE.search(template): + return template, False + + render_blocks = ( + (_QWEN_ASSISTANT_RENDER_BLOCK, _QWEN_ASSISTANT_GENERATION_RENDER_BLOCK), + (_QWEN_COMPACT_ASSISTANT_RENDER_BLOCK, _QWEN_COMPACT_ASSISTANT_GENERATION_RENDER_BLOCK), + ) + for render_block, generation_render_block in render_blocks: + render_count = template.count(render_block) + if render_count != 1: + continue + + render_pos = template.find(render_block) + patched = template[:render_pos] + generation_render_block + template[render_pos + len(render_block) :] + end_pos = patched.find(_QWEN_ASSISTANT_END, render_pos + len(generation_render_block)) + if end_pos < 0: + return template, False + patched = patched[:end_pos] + _QWEN_ASSISTANT_GENERATION_END + patched[end_pos + len(_QWEN_ASSISTANT_END) :] + return patched, patched != template + + return template, False + + def try_patch_qwen_chat_template( sample: CanonicalSample, template: str | None, @@ -112,6 +216,10 @@ def try_patch_qwen_chat_template( if preserve_thinking is None and _has_learnable_historical_thinking(sample): resolved_kwargs["preserve_thinking"] = True + if _assistant_generation_mask_matches_sample(sample): + patched_template, generation_changed = _patch_qwen_generation_markers(patched_template) + changed = changed or generation_changed + return TemplatePatchResult( template=patched_template, kwargs=resolved_kwargs, diff --git a/relax/engine/sft/dataset/streaming.py b/relax/engine/sft/dataset/streaming.py index 89ae5eecc..926c3a9d4 100644 --- a/relax/engine/sft/dataset/streaming.py +++ b/relax/engine/sft/dataset/streaming.py @@ -5,6 +5,7 @@ import asyncio import json import threading +import time from dataclasses import dataclass from numbers import Integral from typing import Any, Callable, Iterable, Optional @@ -474,7 +475,7 @@ def get_batch(self, n: int) -> tuple[list[ProcessedSample], bool]: async def get_batch_async(self, n: int) -> tuple[list[ProcessedSample], bool]: if self._prefetch is not None: - return self._get_batch_prefetch(n) + return await self._get_batch_prefetch_async(n) return await self._get_batch_async_gather(n) def get_batch_in_order(self, start: int, n: int) -> list[ProcessedSample]: @@ -626,6 +627,59 @@ def _get_batch_prefetch(self, n: int) -> tuple[list[ProcessedSample], bool]: ) return samples, crossed_epoch + async def _get_batch_prefetch_async(self, n: int) -> tuple[list[ProcessedSample], bool]: + self._raise_if_failed() + samples: list[ProcessedSample] = [] + crossed_epoch = False + max_attempts = max(n * 10, 32) + attempts = 0 + prefetch_wait_timeout_s = 300.0 + prefetch_wait_poll_s = 0.02 + assert self._prefetch is not None + while len(samples) < n and attempts < max_attempts: + indices, epoch_crossed = self.index_manager.get_next_indices(1) + attempts += 1 + if epoch_crossed and not crossed_epoch: + crossed_epoch = True + remaining = self.index_manager.indices[self.index_manager.position :] + self._prefetch.set_index_order(list(remaining)) + logger.info( + f"SFTStreamingDataset: epoch boundary crossed, prefetch re-primed " + f"(epoch={self.index_manager.current_epoch}, remaining={len(remaining)})" + ) + idx = indices[0] + found, sample = self._prefetch.get_cached(idx) + wait_started = time.monotonic() + while not found: + self._raise_if_failed() + if not self._prefetch.is_alive: + raise RuntimeError( + f"SFTStreamingDataset: prefetch worker exited before sample idx={idx} was cached " + f"(cache_size={self._prefetch.cache_size})" + ) + if time.monotonic() - wait_started >= prefetch_wait_timeout_s: + raise TimeoutError( + f"SFTStreamingDataset: timed out waiting for prefetched sample idx={idx} " + f"after {prefetch_wait_timeout_s:.1f}s " + f"(cache_size={self._prefetch.cache_size}, prefetch_alive={self._prefetch.is_alive})" + ) + await asyncio.sleep(prefetch_wait_poll_s) + found, sample = self._prefetch.get_cached(idx, record_miss=False) + if sample is None: + # Cached None means the background worker processed this index + # and decided it should be skipped. Surface any latched error; + # otherwise keep refilling the batch. + self._raise_if_failed() + else: + samples.append(sample) + self._raise_if_failed() + if len(samples) < n: + logger.warning( + f"SFTStreamingDataset.get_batch_async (prefetch): returned {len(samples)}/{n} samples " + f"after {attempts} attempts." + ) + return samples, crossed_epoch + def _get_batch_inline(self, n: int) -> tuple[list[ProcessedSample], bool]: self._raise_if_failed() samples: list[ProcessedSample] = [] @@ -723,15 +777,13 @@ def _render_one(self, idx: int) -> "_RenderedSample | None": f"exceeds per-GPU capacity {self.capacity}; skipping." ) return None - rendered_text = ( - render_to_text( + rendered_text = None + if has_multimodal_content(sample): + rendered_text = render_to_text( sample, tokenizer=self.tokenizer, apply_chat_template_kwargs=self.apply_chat_template_kwargs, ) - if has_multimodal_content(sample) - else None - ) return _RenderedSample( idx=idx, sample=sample, @@ -959,8 +1011,8 @@ def pack_samples_for_tq( ) -> Optional[dict]: if not samples: return None - tokens = [s.tokens.tolist() for s in samples] - loss_masks = [s.loss_mask.tolist() for s in samples] + tokens = [s.tokens for s in samples] + loss_masks = [s.loss_mask for s in samples] total_lengths = [s.total_length for s in samples] has_mm = force_multimodal_field or any(s.multimodal_train_inputs is not None for s in samples) is_classification = any(s.classification_label is not None for s in samples) diff --git a/relax/engine/sft/runtime.py b/relax/engine/sft/runtime.py index fa2b27d12..0104c10af 100644 --- a/relax/engine/sft/runtime.py +++ b/relax/engine/sft/runtime.py @@ -7,6 +7,7 @@ here keeps the dispatchers in those files to one-line calls. """ +import os import random from argparse import Namespace @@ -85,6 +86,36 @@ def sft_partition_id(args: Namespace, step: int) -> str: return f"sft_{step}" if is_sft_mode(args) else f"train_{step}" +def sft_tq_num_shards(args: Namespace) -> int: + """Number of TQ shard partitions per SFT step. + + Kept as an env knob while this path is experimental so launch scripts can + do A/B tests without adding a public CLI surface. + """ + if not is_sft_mode(args) or not getattr(args, "sft_async_prepack", False): + return 1 + raw_value = os.environ.get("RELAX_SFT_TQ_SHARDS", "1") + try: + return max(1, int(raw_value)) + except ValueError: + return 1 + + +def sft_partition_ids(args: Namespace, step: int) -> list[str]: + base_partition_id = sft_partition_id(args, step) + num_shards = sft_tq_num_shards(args) + if num_shards == 1: + return [base_partition_id] + return [f"{base_partition_id}_shard_{shard_id}_of_{num_shards}" for shard_id in range(num_shards)] + + +def sft_logical_partition_id(partition_id: str) -> str: + marker = "_shard_" + if partition_id.startswith("sft_") and marker in partition_id: + return partition_id.split(marker, 1)[0] + return partition_id + + def sft_task_name(args: Namespace, *, component: str = "actor") -> str: """Return the TransferQueue task name. diff --git a/relax/utils/data/streaming_dataset.py b/relax/utils/data/streaming_dataset.py index 404475303..1f7605b59 100644 --- a/relax/utils/data/streaming_dataset.py +++ b/relax/utils/data/streaming_dataset.py @@ -466,6 +466,8 @@ class PrefetchBuffer: fetching immediately, well before ``get_batch`` is called. - ``get(idx)`` pops from the cache (near-zero latency on hit) or falls back to a synchronous single-sample fetch on miss. + - ``get_cached(idx)`` lets async producers avoid that fallback and wait + for the background worker instead. - The cache is bounded by ``max_cached``; when full the prefetch thread pauses until consumers free space via ``get()`` calls. - A ``ThreadPoolExecutor`` is used inside the prefetch thread to @@ -499,13 +501,15 @@ def __init__( ``ThreadPoolExecutor`` for I/O-bound decoding. """ self._process_fn = process_fn - self._chunk_size = chunk_size + self._chunk_size = max(1, min(chunk_size, max_cached)) if max_cached > 0 else max(1, chunk_size) self._max_cached = max_cached self._num_workers = num_workers # Thread-safe cache: idx -> Optional[Sample] self._cache: dict[int, Optional[Sample]] = {} + self._fallback_fetched_indices: set[int] = set() self._lock = threading.Lock() + self._cache_updated = threading.Condition(self._lock) # Ordered index sequence set by set_index_order self._indices: list[int] = [] @@ -518,13 +522,15 @@ def __init__( # Stats self._prefetch_hits = 0 self._prefetch_misses = 0 + self._prefetch_stale_drops = 0 # Thread lifecycle self._stop = threading.Event() self._thread: Optional[threading.Thread] = None logger.info( - f"PrefetchBuffer created: max_cached={max_cached}, chunk_size={chunk_size}, num_workers={num_workers}" + f"PrefetchBuffer created: max_cached={max_cached}, chunk_size={chunk_size}, " + f"effective_chunk_size={self._chunk_size}, num_workers={num_workers}" ) # -- Public API -------------------------------------------------------- @@ -544,10 +550,12 @@ def set_index_order(self, indices: list[int]) -> None: if self._thread.is_alive(): logger.warning("Previous prefetch thread did not stop within 10s; it will exit on its own stop-event") - with self._lock: + with self._cache_updated: self._cache.clear() + self._fallback_fetched_indices.clear() self._indices = list(indices) self._pos = 0 + self._cache_updated.notify_all() # Create a fresh stop-event for the new thread so the old thread # (if still draining) keeps seeing its own set() signal and exits. @@ -571,20 +579,62 @@ def get(self, idx: int) -> Optional[Sample]: # Signal prefetch thread that space is available self._space_available.set() return sample + self._prefetch_misses += 1 + self._fallback_fetched_indices.add(idx) # Cache miss — synchronous fallback - self._prefetch_misses += 1 try: return self._process_fn(idx) except Exception: logger.exception(f"Prefetch fallback failed for index {idx}") return None + def get_cached(self, idx: int, *, record_miss: bool = True) -> tuple[bool, Optional[Sample]]: + """Pop a prefetched sample without synchronous fallback. + + Returns ``(available, sample)`` so callers can distinguish a missing + cache entry from a cached ``None`` sample that should be skipped. + """ + with self._lock: + if idx in self._cache: + sample = self._cache.pop(idx) + self._prefetch_hits += 1 + self._space_available.set() + return True, sample + if record_miss: + self._prefetch_misses += 1 + return False, None + + def wait_for(self, idx: int, timeout: float | None = None) -> bool: + """Wait until *idx* is present in the cache, or timeout/stop occurs.""" + deadline = None if timeout is None else time.monotonic() + timeout + with self._cache_updated: + while idx not in self._cache and not self._stop.is_set(): + thread = self._thread + if thread is not None and not thread.is_alive(): + break + if deadline is None: + self._cache_updated.wait() + continue + remaining = deadline - time.monotonic() + if remaining <= 0: + break + self._cache_updated.wait(timeout=remaining) + return idx in self._cache + + @property + def is_alive(self) -> bool: + """Return whether the current prefetch thread is alive.""" + thread = self._thread + return bool(thread is not None and thread.is_alive()) + def stop(self) -> None: """Signal the prefetch thread to stop.""" self._stop.set() # Unblock if waiting on space self._space_available.set() + with self._cache_updated: + self._cache_updated.notify_all() if self._thread is not None and self._thread.is_alive(): self._thread.join(timeout=15) if self._thread.is_alive(): @@ -592,8 +642,9 @@ def stop(self) -> None: def clear(self) -> None: """Clear the cache and reset position (without stopping the thread).""" - with self._lock: + with self._cache_updated: self._cache.clear() + self._cache_updated.notify_all() self._space_available.set() @property @@ -637,13 +688,33 @@ def _run(self, stop_event: threading.Event) -> None: # 2. Filter out indices already in cache with self._lock: - to_fetch = [i for i in chunk if i not in self._cache] + to_fetch = [] + for idx in chunk: + if idx in self._cache: + continue + if idx in self._fallback_fetched_indices: + self._fallback_fetched_indices.discard(idx) + self._prefetch_stale_drops += 1 + continue + to_fetch.append(idx) if not to_fetch: continue # 3. Wait until cache has room for this chunk while not stop_event.is_set(): with self._lock: + pending = [] + for idx in to_fetch: + if idx in self._cache: + continue + if idx in self._fallback_fetched_indices: + self._fallback_fetched_indices.discard(idx) + self._prefetch_stale_drops += 1 + continue + pending.append(idx) + to_fetch = pending + if not to_fetch: + break if len(self._cache) + len(to_fetch) <= self._max_cached: break self._space_available.clear() @@ -651,6 +722,9 @@ def _run(self, stop_event: threading.Event) -> None: if not self._space_available.wait(timeout=0.1): continue + if not to_fetch: + continue + if stop_event.is_set(): return @@ -679,13 +753,21 @@ def _run(self, stop_event: threading.Event) -> None: continue # 5. Store results in cache - with self._lock: + with self._cache_updated: for idx, sample in results.items(): + if idx in self._fallback_fetched_indices: + self._fallback_fetched_indices.discard(idx) + self._prefetch_stale_drops += 1 + continue self._cache[idx] = sample + self._cache_updated.notify_all() + with self._cache_updated: + self._cache_updated.notify_all() logger.info( f"Prefetch thread finished. Hit rate: {self.hit_rate:.1%} " - f"(hits={self._prefetch_hits}, misses={self._prefetch_misses})" + f"(hits={self._prefetch_hits}, misses={self._prefetch_misses}, " + f"stale_drops={self._prefetch_stale_drops})" ) diff --git a/relax/utils/utils.py b/relax/utils/utils.py index c367bd62a..2eaf1023b 100644 --- a/relax/utils/utils.py +++ b/relax/utils/utils.py @@ -245,10 +245,11 @@ def dict_to_tensordict( batch_size: Union[int, torch.Size, None] = None, device: Optional[torch.device] = None, ) -> TensorDict: - """Convert a nested-list dictionary to a TensorDict. + """Convert a nested-list / tensor-list dictionary to a TensorDict. Args: - data: Mapping of keys to nested lists (supports depth 1 or 2). + data: Mapping of keys to nested lists (supports depth 1 or 2) or + lists of per-sample tensors. batch_size: Optional batch size. If None, caller may set an appropriate batch size (TensorDict accepts None or an int/torch.Size). device: Optional target torch.device for created tensors. @@ -283,6 +284,12 @@ def _to_tensor_2d(lst): tensors = [torch.tensor(seq, dtype=dtype, device=device) for seq in lst] return torch.nested.as_nested_tensor(tensors, layout=torch.jagged) + def _to_nested_tensor_from_tensor_list(lst): + if not all(isinstance(item, torch.Tensor) for item in lst): + raise TypeError("Mixed tensor and non-tensor values are not supported") + tensors = [tensor.to(device=device) for tensor in lst] if device is not None else lst + return torch.nested.as_nested_tensor(tensors, layout=torch.jagged) + result = {} for key, value in data.items(): @@ -310,6 +317,9 @@ def _to_tensor_2d(lst): ] result[key] = torch.nested.as_nested_tensor(tensors, layout=torch.jagged) continue + if value and isinstance(value[0], torch.Tensor): + result[key] = _to_nested_tensor_from_tensor_list(value) + continue depth = _nesting_depth(value) if depth == 0: # empty list [] tensor = torch.empty(0) diff --git a/tests/components/test_sft.py b/tests/components/test_sft.py index eacf1261c..e6078b5ed 100644 --- a/tests/components/test_sft.py +++ b/tests/components/test_sft.py @@ -2,6 +2,7 @@ """Unit tests for SFT producer component (loop-only, no Ray runtime).""" +import asyncio import sys from types import ModuleType, SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -54,6 +55,8 @@ def _make_args(global_batch_size=4, max_tokens_per_gpu=128, num_rollout=1): start_rollout_id=0, seed=42, max_staleness=0, + sft_async_prepack=False, + custom_dataset_class_path=None, ) @@ -125,8 +128,7 @@ def test_sft_component_imports_without_ray(): from relax.components.sft import SFT # noqa: F401 -@pytest.mark.asyncio -async def test_sft_step_pushes_one_batch_to_tq(monkeypatch): +def test_sft_step_pushes_one_batch_to_tq(monkeypatch): from relax.components.sft import SFT _patch_pipeline_dependencies(monkeypatch) @@ -155,7 +157,7 @@ async def test_sft_step_pushes_one_batch_to_tq(monkeypatch): sft._stop_event.is_set = MagicMock(return_value=False) sft._init_data_pipeline() - await sft._produce_one_step() + asyncio.run(sft._produce_one_step()) assert fake_client.async_put.await_count == 1 args_call, kwargs_call = fake_client.async_put.call_args pushed_data = kwargs_call.get("data") @@ -167,9 +169,51 @@ async def test_sft_step_pushes_one_batch_to_tq(monkeypatch): assert kwargs_call.get("custom_meta") == [{"total_lengths": 8}] * 4 -@pytest.mark.asyncio +def test_sft_step_pushes_sharded_batches_to_tq(monkeypatch): + from relax.components.sft import SFT, _sft_train_partitions_in_flight + + _patch_pipeline_dependencies(monkeypatch) + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "2") + monkeypatch.setattr("relax.components.sft.ray.is_initialized", lambda: False) + + fake_client = MagicMock() + fake_client.async_put = AsyncMock(return_value=None) + monkeypatch.setattr("relax.components.sft.tq.init", lambda *a, **kw: None) + monkeypatch.setattr("relax.components.sft.tq.get_client", lambda: fake_client) + + args = _make_args(global_batch_size=4) + args.sft_async_prepack = True + SFTCls = SFT.func_or_class + sft = SFTCls.__new__(SFTCls) + sft.config = args + sft.role = "sft" + sft._healthy = True + sft.step = 0 + sft.data_system_client = fake_client + sft._dataset = None + sft._eval_dataset = None + sft._eval_indices = None + sft._train_size = 0 + sft._tokenizer = None + sft._processor_pool = None + sft._logger_instance = None + sft._stop_event = MagicMock() + sft._stop_event.is_set = MagicMock(return_value=False) + sft._runtime_env = None + + sft._init_data_pipeline() + asyncio.run(sft._produce_one_step()) + + assert fake_client.async_put.await_count == 2 + seen_partitions = [c.kwargs.get("partition_id") for c in fake_client.async_put.call_args_list] + assert seen_partitions == ["sft_0_shard_0_of_2", "sft_0_shard_1_of_2"] + seen_meta = [c.kwargs.get("custom_meta") for c in fake_client.async_put.call_args_list] + assert seen_meta == [[{"total_lengths": 8}] * 2, [{"total_lengths": 8}] * 2] + assert _sft_train_partitions_in_flight(seen_partitions) == 1 + + @pytest.mark.parametrize("returned_count", [0, 3]) -async def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_count): +def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_count): from relax.components.sft import SFT fake_ds, _ = _patch_pipeline_dependencies(monkeypatch) @@ -200,14 +244,13 @@ async def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_cou sft._init_data_pipeline() with pytest.raises(RuntimeError, match=rf"dataset returned {returned_count}/4 samples"): - await sft._produce_one_step() + asyncio.run(sft._produce_one_step()) fake_client.async_put.assert_not_awaited() assert sft.step == 0 -@pytest.mark.asyncio -async def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): +def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): from relax.components.sft import SFT _patch_pipeline_dependencies(monkeypatch) @@ -234,14 +277,13 @@ async def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): sft._stop_event.is_set = MagicMock(return_value=False) with pytest.raises(RuntimeError, match="source produced 0 valid samples"): - await sft._maybe_produce_eval() + asyncio.run(sft._maybe_produce_eval()) fake_client.async_put.assert_not_awaited() -@pytest.mark.asyncio @pytest.mark.parametrize("n_real", [1, 3, 4, 5, 8]) -async def test_classification_eval_pads_without_dropping_real_samples(n_real): +def test_classification_eval_pads_without_dropping_real_samples(n_real): from relax.components.sft import SFT samples = [ @@ -273,7 +315,7 @@ async def test_classification_eval_pads_without_dropping_real_samples(n_real): sft._build_eval_batches = MagicMock(return_value=samples) sft._wait_for_partition_drained = AsyncMock(return_value=True) - await sft._maybe_produce_eval() + asyncio.run(sft._maybe_produce_eval()) expected_chunks = (n_real + 3) // 4 assert fake_client.async_put.await_count == expected_chunks @@ -283,8 +325,7 @@ async def test_classification_eval_pads_without_dropping_real_samples(n_real): assert partition_ids == [f"sft_eval_0_n{expected_chunks}_{idx}" for idx in range(expected_chunks)] -@pytest.mark.asyncio -async def test_sft_loop_advances_step(monkeypatch): +def test_sft_loop_advances_step(monkeypatch): from relax.components.sft import SFT _patch_pipeline_dependencies(monkeypatch) @@ -315,15 +356,14 @@ async def test_sft_loop_advances_step(monkeypatch): sft._init_data_pipeline() for _ in range(3): - await sft._produce_one_step() + asyncio.run(sft._produce_one_step()) assert sft.step == 3 assert fake_client.async_put.await_count == 3 seen_partitions = [c.kwargs.get("partition_id") for c in fake_client.async_put.call_args_list] assert seen_partitions == ["sft_0", "sft_1", "sft_2"] -@pytest.mark.asyncio -async def test_sft_resume_only_produces_remaining_steps(): +def test_sft_resume_only_produces_remaining_steps(): from relax.components.sft import SFT SFTCls = SFT.func_or_class @@ -338,15 +378,14 @@ async def _produce_one_step(): sft._produce_one_step = AsyncMock(side_effect=_produce_one_step) - await sft._async_run() + asyncio.run(sft._async_run()) assert sft.step == 5 assert sft._produce_one_step.await_count == 3 -@pytest.mark.asyncio @pytest.mark.parametrize("start_step", [5, 6]) -async def test_sft_resume_at_or_after_end_produces_nothing(start_step): +def test_sft_resume_at_or_after_end_produces_nothing(start_step): from relax.components.sft import SFT SFTCls = SFT.func_or_class @@ -357,6 +396,6 @@ async def test_sft_resume_at_or_after_end_produces_nothing(start_step): sft._stop_event.is_set = MagicMock(return_value=False) sft._produce_one_step = AsyncMock() - await sft._async_run() + asyncio.run(sft._async_run()) sft._produce_one_step.assert_not_awaited() diff --git a/tests/engine/sft/dataset/test_chat_template.py b/tests/engine/sft/dataset/test_chat_template.py index cf177642d..53b5225a0 100644 --- a/tests/engine/sft/dataset/test_chat_template.py +++ b/tests/engine/sft/dataset/test_chat_template.py @@ -109,6 +109,8 @@ def test_render_uses_assistant_mask_when_template_supports_it(): # Verify return_assistant_tokens_mask was requested call_kwargs = tok.apply_chat_template.call_args.kwargs assert call_kwargs.get("return_assistant_tokens_mask") is True + assert call_kwargs["chat_template"].startswith(tok.chat_template) + assert "relax_thread=" in call_kwargs["chat_template"] def test_render_falls_back_when_no_generation_marker(capsys): diff --git a/tests/engine/sft/dataset/test_qwen_chat_template_patch.py b/tests/engine/sft/dataset/test_qwen_chat_template_patch.py index 4077e8ce3..2661d673c 100644 --- a/tests/engine/sft/dataset/test_qwen_chat_template_patch.py +++ b/tests/engine/sft/dataset/test_qwen_chat_template_patch.py @@ -21,6 +21,71 @@ _QWEN36_TEMPLATE = _QWEN35_TEMPLATE.replace(_QWEN_HISTORY_GATE, _QWEN_PRESERVE_HISTORY_GATE) # Qwen3.8 ships a preserve-by-default gate that also references reasoning_content. _QWEN38_TEMPLATE = "\n".join(("template-start reasoning_content", _QWEN38_PRESERVE_HISTORY_GATE, "template-end")) +_QWEN35_RENDER_TEMPLATE = "\n".join( + ( + "template-start", + "{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}", + "{%- for message in messages %}", + " {%- set content = message.content|trim %}", + ' {%- if message.role == "assistant" %}', + " {%- set reasoning_content = '' %}", + " {%- if message.reasoning_content is string %}", + " {%- set reasoning_content = message.reasoning_content %}", + " {%- else %}", + " {%- if '' in content %}", + " {%- set reasoning_content = content.split('')[0].rstrip('\\n').split('')[-1].lstrip('\\n') %}", + " {%- set content = content.split('')[-1].lstrip('\\n') %}", + " {%- endif %}", + " {%- endif %}", + " {%- set reasoning_content = reasoning_content|trim %}", + f" {_QWEN_HISTORY_GATE}", + " {{- '<|im_start|>' + message.role + '\\n\\n' + reasoning_content + '\\n\\n\\n' + content }}", + " {%- else %}", + " {{- '<|im_start|>' + message.role + '\\n' + content }}", + " {%- endif %}", + " {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}", + " {{- 'dummy' }}", + " {%- endif %}", + " {{- '<|im_end|>\\n' }}", + " {%- endif %}", + "{%- endfor %}", + "template-end", + ) +) +_QWEN3_COMPACT_RENDER_TEMPLATE = "\n".join( + ( + "template-start", + "{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}", + "{%- for message in messages %}", + " {%- set content = message.content|trim %}", + ' {%- if message.role == "assistant" %}', + " {%- set reasoning_content = '' %}", + " {%- if message.reasoning_content is string %}", + " {%- set reasoning_content = message.reasoning_content %}", + " {%- else %}", + " {%- if '' in content %}", + " {%- set reasoning_content = content.split('')[0].rstrip('\\n').split('')[-1].lstrip('\\n') %}", + " {%- set content = content.split('')[-1].lstrip('\\n') %}", + " {%- endif %}", + " {%- endif %}", + f" {_QWEN_HISTORY_GATE}", + " {%- if loop.last or (not loop.last and reasoning_content) %}", + " {{- '<|im_start|>' + message.role + '\\n\\n' + reasoning_content.strip('\\n') + '\\n\\n\\n' + content.lstrip('\\n') }}", + " {%- else %}", + " {{- '<|im_start|>' + message.role + '\\n' + content }}", + " {%- endif %}", + " {%- else %}", + " {{- '<|im_start|>' + message.role + '\\n' + content }}", + " {%- endif %}", + " {%- if message.tool_calls %}", + " {{- 'dummy' }}", + " {%- endif %}", + " {{- '<|im_end|>\\n' }}", + " {%- endif %}", + "{%- endfor %}", + "template-end", + ) +) def _make_sample(*, historical_learn: bool = True) -> CanonicalSample: @@ -62,6 +127,7 @@ def __init__(self, chat_template: str | None = None): if chat_template is not None: self.chat_template = chat_template self.last_template = self.chat_template + self.used_assistant_mask = False @staticmethod def _tokenize(text): @@ -78,9 +144,50 @@ def _render_tool_calls(tool_calls): text += "\n\n" return text - def apply_chat_template(self, messages, *, tools=None, tokenize=True, **kwargs): # noqa: ARG002 + @staticmethod + def _assistant_mask(messages, rendered): + mask = [0] * len(rendered) + cursor = 0 + for message in messages: + role = message["role"] + if role == "tool": + open_pos = rendered.find("\n\n", cursor) + close_pos = rendered.find("\n", open_pos + 1) + cursor = close_pos + len("\n") + continue + header = f"<|im_start|>{role}\n" + header_pos = rendered.find(header, cursor) + content_start = header_pos + len(header) + end_pos = rendered.find("<|im_end|>", content_start) + span_end = end_pos + len("<|im_end|>") + if span_end < len(rendered) and rendered[span_end] == "\n": + span_end += 1 + cursor = span_end + if role != "assistant": + continue + mask_start = content_start + if rendered[content_start : content_start + len("\n")] == "\n": + mask_start += len("\n") + for pos in range(mask_start, span_end): + mask[pos] = 1 + return mask + + def apply_chat_template( + self, + messages, + *, + tools=None, + tokenize=True, + return_tensors=None, + return_dict=False, + return_assistant_tokens_mask=False, + **kwargs, + ): # noqa: ARG002 self.last_template = kwargs.get("chat_template", self.chat_template) - preserve = kwargs.get("preserve_thinking") is True and _QWEN_PRESERVE_HISTORY_GATE in self.last_template + has_preserve_gate = ( + _QWEN_PRESERVE_HISTORY_GATE in self.last_template or "relax_has_visible_thinking" in self.last_template + ) + preserve = kwargs.get("preserve_thinking") is True and has_preserve_gate last_user_index = max( (index for index, message in enumerate(messages) if message["role"] == "user"), default=-1, @@ -114,6 +221,16 @@ def apply_chat_template(self, messages, *, tools=None, tokenize=True, **kwargs): if not tokenize: return rendered ids, _ = self._tokenize(rendered) + if return_assistant_tokens_mask: + self.used_assistant_mask = True + result_ids = torch.tensor([ids], dtype=torch.long) if return_tensors == "pt" else [ids] + return { + "input_ids": result_ids, + "assistant_masks": [self._assistant_mask(messages, rendered)], + } + if return_dict: + result_ids = torch.tensor([ids], dtype=torch.long) if return_tensors == "pt" else [ids] + return {"input_ids": result_ids} return ids def __call__(self, text, *, add_special_tokens=False, return_offsets_mapping=False, **kwargs): # noqa: ARG002 @@ -220,6 +337,46 @@ def test_qwen_patch_allows_compression_of_unlearned_history(): assert result.kwargs["preserve_thinking"] is False +def test_qwen_patch_adds_generation_markers_for_all_learned_assistants(): + result = try_patch_qwen_chat_template(_make_sample(), _QWEN35_RENDER_TEMPLATE, {}) + assert result is not None + assert result.changed + assert "{%- generation %}" in result.template + assert "{%- endgeneration %}" in result.template + + +def test_qwen_patch_adds_generation_markers_for_compact_template(): + result = try_patch_qwen_chat_template(_make_sample(), _QWEN3_COMPACT_RENDER_TEMPLATE, {}) + assert result is not None + assert result.changed + assert "{%- generation %}" in result.template + assert "{%- endgeneration %}" in result.template + assert "reasoning_content.strip('\\n')" in result.template + assert "content.lstrip('\\n')" in result.template + + +def test_qwen_patch_keeps_fallback_when_assistant_learn_flags_need_custom_mask(): + result = try_patch_qwen_chat_template( + _make_sample(historical_learn=False), + _QWEN35_RENDER_TEMPLATE, + {"preserve_thinking": False}, + ) + assert result is not None + assert result.changed + assert "{%- generation %}" not in result.template + + +def test_qwen_compact_patch_keeps_fallback_when_assistant_learn_flags_need_custom_mask(): + result = try_patch_qwen_chat_template( + _make_sample(historical_learn=False), + _QWEN3_COMPACT_RENDER_TEMPLATE, + {"preserve_thinking": False}, + ) + assert result is not None + assert result.changed + assert "{%- generation %}" not in result.template + + def test_qwen_patch_explicit_true_preserves_unlearned_history(): result = try_patch_qwen_chat_template( _make_sample(historical_learn=False), @@ -254,12 +411,13 @@ def test_qwen_patch_excludes_wrapped_tool_response_from_last_user_boundary(): def test_qwen35_render_with_loss_mask_preserves_think_before_tool_call(): - tokenizer = _FakeQwenHistoryTokenizer() + tokenizer = _FakeQwenHistoryTokenizer(_QWEN35_RENDER_TEMPLATE) input_ids, loss_mask = render_with_loss_mask(_make_sample(), tokenizer=tokenizer) learned = _learned_text(input_ids, loss_mask) - assert tokenizer.last_template == _QWEN36_TEMPLATE - assert tokenizer.chat_template == _QWEN35_TEMPLATE + assert tokenizer.used_assistant_mask + assert tokenizer.last_template != tokenizer.chat_template + assert "{%- generation %}" in tokenizer.last_template assert re.search(r"NEED_SKILL\n\n+", learned) assert "activate_skill" in learned assert "skill loaded" not in learned @@ -267,7 +425,7 @@ def test_qwen35_render_with_loss_mask_preserves_think_before_tool_call(): def test_qwen35_render_with_loss_mask_explicit_false_compresses_history(): - tokenizer = _FakeQwenHistoryTokenizer() + tokenizer = _FakeQwenHistoryTokenizer(_QWEN35_RENDER_TEMPLATE) input_ids, loss_mask = render_with_loss_mask( _make_sample(), tokenizer=tokenizer, @@ -275,14 +433,15 @@ def test_qwen35_render_with_loss_mask_explicit_false_compresses_history(): ) learned = _learned_text(input_ids, loss_mask) - assert tokenizer.last_template == _QWEN36_TEMPLATE + assert tokenizer.used_assistant_mask + assert "{%- generation %}" in tokenizer.last_template assert "NEED_SKILL" not in learned assert "activate_skill" in learned def test_qwen35_render_to_text_uses_same_patch_dispatcher(): - tokenizer = _FakeQwenHistoryTokenizer() + tokenizer = _FakeQwenHistoryTokenizer(_QWEN35_RENDER_TEMPLATE) text = render_to_text(_make_sample(), tokenizer=tokenizer) first_assistant = text.split("<|im_start|>assistant\n", 1)[1].split("<|im_end|>", 1)[0] assert "\nNEED_SKILL\n" in first_assistant - assert tokenizer.last_template == _QWEN36_TEMPLATE + assert tokenizer.last_template != tokenizer.chat_template diff --git a/tests/engine/sft/dataset/test_streaming.py b/tests/engine/sft/dataset/test_streaming.py index d0b9d1bff..d58d97783 100644 --- a/tests/engine/sft/dataset/test_streaming.py +++ b/tests/engine/sft/dataset/test_streaming.py @@ -18,6 +18,7 @@ _expand_loss_mask_via_alignment, pack_samples_for_tq, ) +from relax.utils.utils import dict_to_tensordict def _write_jsonl(path: Path, rows: list[dict]) -> None: @@ -216,8 +217,8 @@ def test_pack_samples_for_tq_marks_samples_as_sft(tmp_path: Path): assert batch is not None assert batch["response_lengths"] == batch["total_lengths"] - assert batch["response_lengths"][0] == len(batch["tokens"][0]) - assert sum(batch["loss_masks"][0]) == len("A") + assert batch["response_lengths"][0] == batch["tokens"][0].numel() + assert int(batch["loss_masks"][0].sum().item()) == len("A") ds.stop() @@ -457,12 +458,32 @@ def _make_text_only_sample() -> ProcessedSample: def test_pack_samples_for_tq_omits_multimodal_field_for_text_only_batch(): # Default behaviour: an all-text batch carries no multimodal key. - batch = pack_samples_for_tq([_make_text_only_sample()]) + sample = _make_text_only_sample() + batch = pack_samples_for_tq([sample]) assert batch is not None + assert batch["tokens"][0] is sample.tokens + assert batch["loss_masks"][0] is sample.loss_mask assert "multimodal_train_inputs" not in batch +def test_dict_to_tensordict_preserves_tensor_list_as_nested_tensor(): + batch = { + "tokens": [torch.tensor([1, 2, 3]), torch.tensor([4, 5])], + "loss_masks": [torch.tensor([0, 1, 1]), torch.tensor([1, 1])], + "total_lengths": [3, 2], + "response_lengths": [3, 2], + } + + td = dict_to_tensordict(batch, batch_size=2) + + assert td["tokens"].is_nested + assert td["loss_masks"].is_nested + assert [tensor.tolist() for tensor in td["tokens"]] == [[1, 2, 3], [4, 5]] + assert [tensor.tolist() for tensor in td["loss_masks"]] == [[0, 1, 1], [1, 1]] + assert td["total_lengths"].tolist() == [3, 2] + + def test_pack_samples_for_tq_forces_multimodal_field_for_text_only_batch(): # A VL run (multimodal_keys configured) must always emit the field so the # consumer's fixed TQ field list stays satisfied even for text-only batches. @@ -794,6 +815,114 @@ def test_streaming_dataset_invalid_multimodal_skip_refills_batch(tmp_path: Path, ds.stop() +def test_streaming_dataset_async_prefetch_waits_without_foreground_fallback(tmp_path: Path): + path = tmp_path / "train.jsonl" + _write_jsonl( + path, + [ + {"messages": [{"role": "assistant", "content": "A"}]}, + {"messages": [{"role": "assistant", "content": "B"}]}, + ], + ) + ds = SFTStreamingDataset( + path=str(path), + tokenizer=_FakeTokenizer(), + processor_pool=None, + capacity=None, + prompt_key="messages", + seed=0, + prefetch_max_cached=0, + ) + ds.shuffle(0) + ds.index_manager.position = 0 + first_idx = ds.index_manager.indices[0] + + class _FakePrefetch: + cache_size = 1 + is_alive = True + + def __init__(self) -> None: + self.get_called = False + self.get_cached_calls = 0 + + def get(self, idx: int): # noqa: ARG002 + self.get_called = True + raise AssertionError("async prefetch path must not use synchronous fallback") + + def get_cached(self, idx: int, *, record_miss: bool = True): # noqa: ARG002 + self.get_cached_calls += 1 + if self.get_cached_calls == 1: + return False, None + return True, ProcessedSample( + tokens=torch.tensor([1], dtype=torch.long), + loss_mask=torch.tensor([1], dtype=torch.long), + total_length=1, + multimodal_train_inputs=None, + source_idx=idx, + ) + + def set_index_order(self, indices: list[int]) -> None: # noqa: ARG002 + raise AssertionError("test should not re-prime prefetch") + + def stop(self) -> None: + pass + + prefetch = _FakePrefetch() + ds._prefetch = prefetch + + try: + samples, crossed = asyncio.run(ds.get_batch_async(1)) + assert crossed is False + assert [item.source_idx for item in samples] == [first_idx] + assert prefetch.get_called is False + assert prefetch.get_cached_calls == 2 + finally: + ds.stop() + + +def test_streaming_dataset_async_prefetch_raises_when_worker_exits(tmp_path: Path): + path = tmp_path / "train.jsonl" + _write_jsonl( + path, + [ + {"messages": [{"role": "assistant", "content": "A"}]}, + {"messages": [{"role": "assistant", "content": "B"}]}, + ], + ) + ds = SFTStreamingDataset( + path=str(path), + tokenizer=_FakeTokenizer(), + processor_pool=None, + capacity=None, + prompt_key="messages", + seed=0, + prefetch_max_cached=0, + ) + ds.shuffle(0) + ds.index_manager.position = 0 + + class _DeadPrefetch: + cache_size = 0 + is_alive = False + + def get_cached(self, idx: int, *, record_miss: bool = True): # noqa: ARG002 + return False, None + + def wait_for(self, idx: int, timeout: float | None = None): # noqa: ARG002 + return False + + def stop(self) -> None: + pass + + ds._prefetch = _DeadPrefetch() + + try: + with pytest.raises(RuntimeError, match="prefetch worker exited"): + asyncio.run(ds.get_batch_async(1)) + finally: + ds.stop() + + def test_streaming_dataset_rejects_unknown_invalid_multimodal_strategy(tmp_path: Path): path = tmp_path / "train.jsonl" _write_jsonl(path, _invalid_image_url_rows()) diff --git a/tests/utils/data/test_streaming_dataset.py b/tests/utils/data/test_streaming_dataset.py index c1426a5a9..1c633bccb 100644 --- a/tests/utils/data/test_streaming_dataset.py +++ b/tests/utils/data/test_streaming_dataset.py @@ -7,6 +7,8 @@ import json import os import tempfile +import threading +import time from unittest.mock import MagicMock import pytest @@ -282,6 +284,71 @@ def test_buffer_clear(self): assert buffer.get(0) is None +class TestPrefetchBuffer: + """Tests for PrefetchBuffer race behavior.""" + + def test_chunk_size_is_clamped_to_cache_capacity(self): + from relax.utils.data.streaming_dataset import PrefetchBuffer + + def process_fn(idx: int) -> str: + return f"sample-{idx}" + + buffer = PrefetchBuffer(process_fn, chunk_size=8, max_cached=1, num_workers=1) + buffer.set_index_order([0]) + try: + assert buffer.wait_for(0, timeout=2) + found, sample = buffer.get_cached(0) + assert found is True + assert sample == "sample-0" + finally: + buffer.stop() + + def test_missed_inflight_index_is_not_stored_as_stale_cache(self): + from relax.utils.data.streaming_dataset import PrefetchBuffer + + background_started = threading.Event() + release_background = threading.Event() + calls = [] + calls_lock = threading.Lock() + + def process_fn(idx: int) -> str: + thread_name = threading.current_thread().name + with calls_lock: + calls.append((idx, thread_name)) + if idx == 0 and thread_name.startswith("pf"): + background_started.set() + assert release_background.wait(timeout=2) + return f"sample-{idx}-{thread_name}" + + buffer = PrefetchBuffer(process_fn, chunk_size=1, max_cached=4, num_workers=1) + buffer.set_index_order([0, 1]) + try: + assert background_started.wait(timeout=2) + + sample = buffer.get(0) + assert sample.startswith("sample-0-") + + release_background.set() + for _ in range(200): + with buffer._lock: + cached_keys = set(buffer._cache) + stale_drops = buffer._prefetch_stale_drops + if 1 in cached_keys and stale_drops: + break + time.sleep(0.01) + + with buffer._lock: + assert 0 not in buffer._cache + assert 1 in buffer._cache + assert buffer._prefetch_stale_drops == 1 + + assert any(idx == 0 and name.startswith("pf") for idx, name in calls) + assert any(idx == 0 and not name.startswith("pf") for idx, name in calls) + finally: + release_background.set() + buffer.stop() + + class TestIndexManager: """Tests for IndexManager class.""" From 798483a0786afc0cff88b54380d1b9dc84f447b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=B0=8F=E5=87=89=28=E8=96=9B=E5=B0=8A=E5=B0=A7=29?= Date: Thu, 20 Aug 2026 11:10:48 +0800 Subject: [PATCH 11/34] fix(sft): prevent prefetch boundary hangs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Re-prime current index across epoch boundaries - Include the current epoch-boundary index when resetting SFT prefetch order. - Re-prime on every crossed epoch so async prefetch does not wait for an index that was omitted from the worker queue. - Return immediately from PrefetchBuffer.wait_for when no prefetch thread has been started. --- # ✅ Tests ## Cover prefetch wait edge cases - Add regression coverage for async SFT prefetch epoch-boundary re-priming. - Add coverage for wait_for(timeout=None) before set_index_order starts a worker. (cherry picked from commit 9635f298dc42fc380f262b3b501d1ec73bfd1d70) --- relax/engine/sft/dataset/streaming.py | 16 +++---- relax/utils/data/streaming_dataset.py | 2 +- tests/engine/sft/dataset/test_streaming.py | 55 ++++++++++++++++++++++ tests/utils/data/test_streaming_dataset.py | 15 ++++++ 4 files changed, 79 insertions(+), 9 deletions(-) diff --git a/relax/engine/sft/dataset/streaming.py b/relax/engine/sft/dataset/streaming.py index 926c3a9d4..f34d78992 100644 --- a/relax/engine/sft/dataset/streaming.py +++ b/relax/engine/sft/dataset/streaming.py @@ -601,16 +601,16 @@ def _get_batch_prefetch(self, n: int) -> tuple[list[ProcessedSample], bool]: while len(samples) < n and attempts < max_attempts: indices, epoch_crossed = self.index_manager.get_next_indices(1) attempts += 1 - if epoch_crossed and not crossed_epoch: + idx = indices[0] + if epoch_crossed: crossed_epoch = True - remaining = self.index_manager.indices[self.index_manager.position :] + remaining = [idx, *self.index_manager.indices[self.index_manager.position :]] assert self._prefetch is not None - self._prefetch.set_index_order(list(remaining)) + self._prefetch.set_index_order(remaining) logger.info( f"SFTStreamingDataset: epoch boundary crossed, prefetch re-primed " f"(epoch={self.index_manager.current_epoch}, remaining={len(remaining)})" ) - idx = indices[0] assert self._prefetch is not None sample = self._prefetch.get(idx) if sample is None: @@ -639,15 +639,15 @@ async def _get_batch_prefetch_async(self, n: int) -> tuple[list[ProcessedSample] while len(samples) < n and attempts < max_attempts: indices, epoch_crossed = self.index_manager.get_next_indices(1) attempts += 1 - if epoch_crossed and not crossed_epoch: + idx = indices[0] + if epoch_crossed: crossed_epoch = True - remaining = self.index_manager.indices[self.index_manager.position :] - self._prefetch.set_index_order(list(remaining)) + remaining = [idx, *self.index_manager.indices[self.index_manager.position :]] + self._prefetch.set_index_order(remaining) logger.info( f"SFTStreamingDataset: epoch boundary crossed, prefetch re-primed " f"(epoch={self.index_manager.current_epoch}, remaining={len(remaining)})" ) - idx = indices[0] found, sample = self._prefetch.get_cached(idx) wait_started = time.monotonic() while not found: diff --git a/relax/utils/data/streaming_dataset.py b/relax/utils/data/streaming_dataset.py index 1f7605b59..72f0196c2 100644 --- a/relax/utils/data/streaming_dataset.py +++ b/relax/utils/data/streaming_dataset.py @@ -611,7 +611,7 @@ def wait_for(self, idx: int, timeout: float | None = None) -> bool: with self._cache_updated: while idx not in self._cache and not self._stop.is_set(): thread = self._thread - if thread is not None and not thread.is_alive(): + if thread is None or not thread.is_alive(): break if deadline is None: self._cache_updated.wait() diff --git a/tests/engine/sft/dataset/test_streaming.py b/tests/engine/sft/dataset/test_streaming.py index d58d97783..6f238e228 100644 --- a/tests/engine/sft/dataset/test_streaming.py +++ b/tests/engine/sft/dataset/test_streaming.py @@ -880,6 +880,61 @@ def stop(self) -> None: ds.stop() +def test_streaming_dataset_async_prefetch_reprimes_current_epoch_boundary_index(tmp_path: Path): + path = tmp_path / "train.jsonl" + _write_jsonl( + path, + [ + {"messages": [{"role": "assistant", "content": "A"}]}, + {"messages": [{"role": "assistant", "content": "B"}]}, + ], + ) + ds = SFTStreamingDataset( + path=str(path), + tokenizer=_FakeTokenizer(), + processor_pool=None, + capacity=None, + prompt_key="messages", + seed=0, + prefetch_max_cached=0, + ) + ds.shuffle(0) + ds.index_manager.position = len(ds.index_manager.indices) + + class _FakePrefetch: + cache_size = 0 + is_alive = True + + def __init__(self) -> None: + self.index_orders: list[list[int]] = [] + + def set_index_order(self, indices: list[int]) -> None: + self.index_orders.append(list(indices)) + + def get_cached(self, idx: int, *, record_miss: bool = True): # noqa: ARG002 + return True, ProcessedSample( + tokens=torch.tensor([1], dtype=torch.long), + loss_mask=torch.tensor([1], dtype=torch.long), + total_length=1, + multimodal_train_inputs=None, + source_idx=idx, + ) + + def stop(self) -> None: + pass + + prefetch = _FakePrefetch() + ds._prefetch = prefetch + + try: + samples, crossed = asyncio.run(ds.get_batch_async(3)) + assert crossed is True + assert len(prefetch.index_orders) == 2 + assert [order[0] for order in prefetch.index_orders] == [samples[0].source_idx, samples[2].source_idx] + finally: + ds.stop() + + def test_streaming_dataset_async_prefetch_raises_when_worker_exits(tmp_path: Path): path = tmp_path / "train.jsonl" _write_jsonl( diff --git a/tests/utils/data/test_streaming_dataset.py b/tests/utils/data/test_streaming_dataset.py index 1c633bccb..f6560e744 100644 --- a/tests/utils/data/test_streaming_dataset.py +++ b/tests/utils/data/test_streaming_dataset.py @@ -348,6 +348,21 @@ def process_fn(idx: int) -> str: release_background.set() buffer.stop() + def test_wait_for_returns_false_when_no_thread_started(self): + from relax.utils.data.streaming_dataset import PrefetchBuffer + + buffer = PrefetchBuffer(lambda idx: f"sample-{idx}", chunk_size=1, max_cached=1, num_workers=1) + result = [] + + waiter = threading.Thread(target=lambda: result.append(buffer.wait_for(0, timeout=None)), daemon=True) + waiter.start() + try: + waiter.join(timeout=0.5) + assert not waiter.is_alive() + assert result == [False] + finally: + buffer.stop() + class TestIndexManager: """Tests for IndexManager class.""" From 7cc2cde4cf99880bd077fe99ac5c21634162e07e Mon Sep 17 00:00:00 2001 From: root Date: Tue, 1 Sep 2026 20:32:49 +0800 Subject: [PATCH 12/34] fix(sft): harden sharded prepack runtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Synchronize sharded prepack runtime - Register and propagate RELAX_SFT_TQ_SHARDS to every Ray actor. - Recheck the final cached sample after the bounded prefetch worker exits. --- # ✅ Tests ## Cover rebase-sensitive SFT paths - Verify target-only fallback and per-node model preparation for remote producers. - Exercise producer re-priming across steps and all-shard Megatron consumption. - Cover typed shard configuration propagation and the worker-exit cache race. (cherry picked from commit c1640dc27e31362a91092f76032ef7e75629e77d) --- relax/engine/sft/dataset/streaming.py | 3 + relax/engine/sft/runtime.py | 9 +- relax/utils/env.py | 1 + relax/utils/utils.py | 5 + tests/backends/megatron/test_sft_prepack.py | 57 +++++++++ tests/components/test_sft.py | 126 ++++++++++++++++++++ tests/engine/sft/dataset/test_streaming.py | 56 +++++++++ tests/utils/test_sft_runtime_env.py | 46 +++++++ 8 files changed, 297 insertions(+), 6 deletions(-) create mode 100644 tests/utils/test_sft_runtime_env.py diff --git a/relax/engine/sft/dataset/streaming.py b/relax/engine/sft/dataset/streaming.py index f34d78992..08aca86bf 100644 --- a/relax/engine/sft/dataset/streaming.py +++ b/relax/engine/sft/dataset/streaming.py @@ -653,6 +653,9 @@ async def _get_batch_prefetch_async(self, n: int) -> tuple[list[ProcessedSample] while not found: self._raise_if_failed() if not self._prefetch.is_alive: + found, sample = self._prefetch.get_cached(idx, record_miss=False) + if found: + break raise RuntimeError( f"SFTStreamingDataset: prefetch worker exited before sample idx={idx} was cached " f"(cache_size={self._prefetch.cache_size})" diff --git a/relax/engine/sft/runtime.py b/relax/engine/sft/runtime.py index 0104c10af..35206539c 100644 --- a/relax/engine/sft/runtime.py +++ b/relax/engine/sft/runtime.py @@ -7,10 +7,11 @@ here keeps the dispatchers in those files to one-line calls. """ -import os import random from argparse import Namespace +from relax.utils.env import Envs + def resolve_sft_eval_split(total_size: int, eval_size: float | int | None) -> tuple[int, int]: """Return ``(train_size, eval_size)`` for an SFT split.""" @@ -94,11 +95,7 @@ def sft_tq_num_shards(args: Namespace) -> int: """ if not is_sft_mode(args) or not getattr(args, "sft_async_prepack", False): return 1 - raw_value = os.environ.get("RELAX_SFT_TQ_SHARDS", "1") - try: - return max(1, int(raw_value)) - except ValueError: - return 1 + return max(1, Envs.RELAX_SFT_TQ_SHARDS) def sft_partition_ids(args: Namespace, step: int) -> list[str]: diff --git a/relax/utils/env.py b/relax/utils/env.py index 94bb6dd5d..7651dede6 100644 --- a/relax/utils/env.py +++ b/relax/utils/env.py @@ -202,6 +202,7 @@ class Envs(metaclass=_EnvsMeta): RELAX_DEVICE_TYPE = EnvProperty("RELAX_DEVICE_TYPE", str, "") RELAX_EMPTY_POLL_SLEEP_MS = EnvProperty("RELAX_EMPTY_POLL_SLEEP_MS", float, 50.0) RELAX_FETCH_SPLIT_MAX_RETRIES = EnvProperty("RELAX_FETCH_SPLIT_MAX_RETRIES", int, 20) + RELAX_SFT_TQ_SHARDS = EnvProperty("RELAX_SFT_TQ_SHARDS", int, 1) # ------------- S3 model cache cleanup ------------- RELAX_S3_MODEL_CLEANUP_TASK_TIMEOUT_S = EnvProperty("RELAX_S3_MODEL_CLEANUP_TASK_TIMEOUT_S", float, 600.0) diff --git a/relax/utils/utils.py b/relax/utils/utils.py index 2eaf1023b..24b14c2de 100644 --- a/relax/utils/utils.py +++ b/relax/utils/utils.py @@ -419,6 +419,11 @@ def post_process_env(args, env): if extra_modules and "RELAX_EXTRA_MODULES" not in env["env_vars"]: env["env_vars"]["RELAX_EXTRA_MODULES"] = extra_modules + # Producer and consumer derive the same TransferQueue partition names from + # this value, so it must be identical in every Serve and Megatron actor. + if "RELAX_SFT_TQ_SHARDS" not in env["env_vars"]: + env["env_vars"]["RELAX_SFT_TQ_SHARDS"] = str(Envs.RELAX_SFT_TQ_SHARDS) + # Generic env-var passthrough for overlay packages. Comma-separated list # of env-var names the driver wants forwarded to every Ray actor. Each # name is copied from the driver's os.environ; missing names are diff --git a/tests/backends/megatron/test_sft_prepack.py b/tests/backends/megatron/test_sft_prepack.py index 201cc41d9..9781fe19b 100644 --- a/tests/backends/megatron/test_sft_prepack.py +++ b/tests/backends/megatron/test_sft_prepack.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 Relax Authors. All Rights Reserved. from argparse import Namespace +from unittest.mock import MagicMock import pytest import torch @@ -56,3 +57,59 @@ def test_sft_lookahead_pauses_on_checkpoint_boundary(monkeypatch): assert actor_module._should_pause_sft_prepack_lookahead(args, rollout_id=19) is True assert actor_module._should_pause_sft_prepack_lookahead(args, rollout_id=18) is False + + +def test_sft_prepack_fetch_concatenates_ready_tq_shards(monkeypatch): + partition_ids = ["sft_3_shard_0_of_2", "sft_3_shard_1_of_2"] + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "2") + monkeypatch.setattr(actor_module.mpu, "get_data_parallel_rank", lambda **_kwargs: 1) + monkeypatch.setattr(actor_module.mpu, "get_data_parallel_world_size", lambda **_kwargs: 2) + monkeypatch.setattr(actor_module, "run", lambda value: value) + + calls = [] + + def _get_data_from_transfer_queue(**kwargs): + calls.append(kwargs) + shard_id = len(calls) - 1 + base = shard_id * 10 + return ( + { + "tokens": [[base], [base + 1]], + "total_lengths": [base, base + 1], + }, + None, + ) + + monkeypatch.setattr(actor_module, "get_data_from_transfer_queue", _get_data_from_transfer_queue) + + actor = object.__new__(actor_module.MegatronTrainRayActor) + actor.args = Namespace(global_batch_size=8, loss_type="sft", sft_async_prepack=True) + actor.data_system_client = MagicMock() + actor.data_system_client.async_get_partition_list.return_value = partition_ids + + batch = actor._fetch_sft_prepack_rollout_once("sft_train", rollout_id=3, data_fields=["tokens"]) + + assert batch == { + "tokens": [[0], [1], [10], [11]], + "total_lengths": [0, 1, 10, 11], + } + assert [call["partition_id"] for call in calls] == partition_ids + assert [call["batch_size"] for call in calls] == [2, 2] + assert all(call["sampling_config"]["dp_rank"] == 1 for call in calls) + + +def test_sft_prepack_fetch_waits_until_all_tq_shards_are_ready(monkeypatch): + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "2") + monkeypatch.setattr(actor_module.mpu, "get_data_parallel_rank", lambda **_kwargs: 0) + monkeypatch.setattr(actor_module.mpu, "get_data_parallel_world_size", lambda **_kwargs: 2) + monkeypatch.setattr(actor_module, "run", lambda value: value) + fetch = MagicMock() + monkeypatch.setattr(actor_module, "get_data_from_transfer_queue", fetch) + + actor = object.__new__(actor_module.MegatronTrainRayActor) + actor.args = Namespace(global_batch_size=8, loss_type="sft", sft_async_prepack=True) + actor.data_system_client = MagicMock() + actor.data_system_client.async_get_partition_list.return_value = ["sft_3_shard_0_of_2"] + + assert actor._fetch_sft_prepack_rollout_once("sft_train", rollout_id=3, data_fields=["tokens"]) is None + fetch.assert_not_called() diff --git a/tests/components/test_sft.py b/tests/components/test_sft.py index e6078b5ed..dd9601589 100644 --- a/tests/components/test_sft.py +++ b/tests/components/test_sft.py @@ -212,6 +212,132 @@ def test_sft_step_pushes_sharded_batches_to_tq(monkeypatch): assert _sft_train_partitions_in_flight(seen_partitions) == 1 +@pytest.mark.parametrize( + ("attribute", "value"), + [ + ("task_type", "seq_cls"), + ("eval_size", 0.25), + ], +) +def test_sft_remote_batch_producer_falls_back_for_target_only_modes(monkeypatch, attribute, value): + from relax.components.sft import SFT + + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "2") + monkeypatch.setattr("relax.components.sft.ray.is_initialized", lambda: True) + + args = _make_args(global_batch_size=4) + args.sft_async_prepack = True + setattr(args, attribute, value) + SFTCls = SFT.func_or_class + sft = SFTCls.__new__(SFTCls) + sft.config = args + sft._logger_instance = MagicMock() + + assert sft._should_use_remote_batch_producer() is False + + +def test_sft_remote_batch_producer_resolves_model_on_its_node(monkeypatch): + from relax.components import sft as sft_module + + call_order = [] + fake_client = MagicMock() + fake_tokenizer = MagicMock() + + def _prepare_model(config, *, completeness): + call_order.append("prepare") + assert completeness == "metadata" + config.hf_checkpoint = "/dev/shm/resolved-sft-model" + + def _load_tokenizer(path, **kwargs): + call_order.append("tokenizer") + assert path == "/dev/shm/resolved-sft-model" + assert kwargs == {"trust_remote_code": True} + return fake_tokenizer + + class _FakeIndexManager: + total_size = 4 + + def __init__(self): + self.current_epoch = -1 + self.position = 0 + self.indices = list(range(self.total_size)) + + def shuffle(self, epoch): + self.current_epoch = epoch + self.position = 0 + self.indices = list(range(self.total_size)) + + fake_dataset = MagicMock() + fake_dataset.__len__ = MagicMock(return_value=4) + fake_dataset.index_manager = _FakeIndexManager() + fake_dataset._prefetch = None + + monkeypatch.setattr(sft_module.tq, "init", MagicMock()) + monkeypatch.setattr(sft_module.tq, "get_client", MagicMock(return_value=fake_client)) + monkeypatch.setattr(sft_module, "prepare_model_maybe_update_args", _prepare_model) + monkeypatch.setattr(sft_module.AutoTokenizer, "from_pretrained", _load_tokenizer) + monkeypatch.setattr(sft_module, "ProcessorPool", MagicMock()) + monkeypatch.setattr(sft_module, "_resolve_pad_token_ids_from_config", MagicMock(return_value=frozenset())) + create_dataset = MagicMock(return_value=fake_dataset) + monkeypatch.setattr(sft_module, "_create_sft_train_dataset", create_dataset) + + args = _make_args(global_batch_size=2) + producer_cls = sft_module._SFTBatchProducerActor.__ray_metadata__.modified_class + producer = producer_cls(args, shard_id=0, num_shards=2, prefetch_num_workers=1) + + state = producer.initialize(start_step=0) + + assert call_order == ["prepare", "tokenizer"] + assert args.hf_checkpoint == "/dev/shm/resolved-sft-model" + assert state["train_size"] == 4 + assert create_dataset.call_args.kwargs["task_type"] == "causal_lm" + + +def test_sft_remote_batch_producer_reprimes_before_second_step(monkeypatch): + from relax.components import sft as sft_module + + class _FakeIndexManager: + total_size = 16 + + def __init__(self): + self.current_epoch = 0 + self.position = 0 + self.indices = list(range(self.total_size)) + + def shuffle(self, epoch): + if epoch != self.current_epoch: + self.current_epoch = epoch + self.position = 0 + self.indices = list(range(self.total_size)) + + class _FakeDataset: + def __init__(self): + self.index_manager = _FakeIndexManager() + self._prefetch = MagicMock() + + async def get_batch_async(self, batch_size): + start = self.index_manager.position + self.index_manager.position += batch_size + return [_make_processed(idx) for idx in range(start, start + batch_size)], False + + monkeypatch.setattr(sft_module, "print_first_sample", MagicMock()) + args = _make_args(global_batch_size=4) + producer_cls = sft_module._SFTBatchProducerActor.__ray_metadata__.modified_class + producer = producer_cls(args, shard_id=0, num_shards=2, prefetch_num_workers=1) + producer._dataset = _FakeDataset() + producer._tokenizer = MagicMock() + producer.data_system_client = MagicMock() + producer.data_system_client.async_put = AsyncMock() + producer._train_size = 16 + + asyncio.run(producer.produce_partition(0, "sft_0_shard_0_of_2", 4, False)) + producer._dataset._prefetch.set_index_order.assert_not_called() + + asyncio.run(producer.produce_partition(1, "sft_1_shard_0_of_2", 4, False)) + producer._dataset._prefetch.set_index_order.assert_called_once_with([4, 5]) + assert producer.data_system_client.async_put.await_count == 2 + + @pytest.mark.parametrize("returned_count", [0, 3]) def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_count): from relax.components.sft import SFT diff --git a/tests/engine/sft/dataset/test_streaming.py b/tests/engine/sft/dataset/test_streaming.py index 6f238e228..96e835c0e 100644 --- a/tests/engine/sft/dataset/test_streaming.py +++ b/tests/engine/sft/dataset/test_streaming.py @@ -880,6 +880,62 @@ def stop(self) -> None: ds.stop() +def test_streaming_dataset_async_prefetch_rechecks_cache_after_worker_exit(tmp_path: Path): + path = tmp_path / "train.jsonl" + _write_jsonl( + path, + [ + {"messages": [{"role": "assistant", "content": "A"}]}, + {"messages": [{"role": "assistant", "content": "B"}]}, + ], + ) + ds = SFTStreamingDataset( + path=str(path), + tokenizer=_FakeTokenizer(), + processor_pool=None, + capacity=None, + prompt_key="messages", + seed=0, + prefetch_max_cached=0, + ) + ds.shuffle(0) + ds.index_manager.position = 0 + first_idx = ds.index_manager.indices[0] + + class _ExitedPrefetch: + cache_size = 1 + is_alive = False + + def __init__(self) -> None: + self.get_cached_calls = 0 + + def get_cached(self, idx: int, *, record_miss: bool = True): # noqa: ARG002 + self.get_cached_calls += 1 + if self.get_cached_calls == 1: + return False, None + return True, ProcessedSample( + tokens=torch.tensor([1], dtype=torch.long), + loss_mask=torch.tensor([1], dtype=torch.long), + total_length=1, + multimodal_train_inputs=None, + source_idx=idx, + ) + + def stop(self) -> None: + pass + + prefetch = _ExitedPrefetch() + ds._prefetch = prefetch + + try: + samples, crossed = asyncio.run(ds.get_batch_async(1)) + assert crossed is False + assert [item.source_idx for item in samples] == [first_idx] + assert prefetch.get_cached_calls == 2 + finally: + ds.stop() + + def test_streaming_dataset_async_prefetch_reprimes_current_epoch_boundary_index(tmp_path: Path): path = tmp_path / "train.jsonl" _write_jsonl( diff --git a/tests/utils/test_sft_runtime_env.py b/tests/utils/test_sft_runtime_env.py new file mode 100644 index 000000000..7dac09fb8 --- /dev/null +++ b/tests/utils/test_sft_runtime_env.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from types import SimpleNamespace + +from relax.engine.sft.runtime import sft_tq_num_shards +from relax.utils.env import Envs, known_env_names +from relax.utils.utils import post_process_env + + +def _make_args() -> SimpleNamespace: + return SimpleNamespace( + fully_async=False, + use_dynamic_batch_size=False, + rollout_batch_size=1, + n_samples_per_prompt=1, + partial_rollout=False, + use_dynamic_global_batch_size=False, + over_sampling_batch_size=1, + ) + + +def test_sft_tq_shards_is_registered_and_typed(monkeypatch): + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "3") + args = SimpleNamespace(loss_type="sft", sft_async_prepack=True) + + assert "RELAX_SFT_TQ_SHARDS" in known_env_names() + assert Envs.RELAX_SFT_TQ_SHARDS == 3 + assert sft_tq_num_shards(args) == 3 + + +def test_post_process_env_propagates_sft_tq_shards(monkeypatch): + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "4") + monkeypatch.setattr("relax.utils.utils._resolve_to_ip", lambda _addr: "127.0.0.1") + + runtime_env = post_process_env(_make_args(), {"env_vars": {}}) + + assert runtime_env["env_vars"]["RELAX_SFT_TQ_SHARDS"] == "4" + + +def test_post_process_env_preserves_configured_sft_tq_shards(monkeypatch): + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "4") + monkeypatch.setattr("relax.utils.utils._resolve_to_ip", lambda _addr: "127.0.0.1") + + runtime_env = post_process_env(_make_args(), {"env_vars": {"RELAX_SFT_TQ_SHARDS": "2"}}) + + assert runtime_env["env_vars"]["RELAX_SFT_TQ_SHARDS"] == "2" From 6f0ad65f0d7b80876fadb29a64834b5e4d6e8987 Mon Sep 17 00:00:00 2001 From: root Date: Tue, 1 Sep 2026 21:10:52 +0800 Subject: [PATCH 13/34] test(sft): restore native async tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ✅ Tests ## Restore pytest-asyncio coverage - Replace per-call asyncio.run wrappers with native await expressions. - Keep sharded producer steps on one event loop and align with the target branch test style. (cherry picked from commit 5a8d24a178b8f343cf84dd2f34b40be9c23ba91f) --- tests/components/test_sft.py | 48 +++++++++++++++++++++--------------- 1 file changed, 28 insertions(+), 20 deletions(-) diff --git a/tests/components/test_sft.py b/tests/components/test_sft.py index dd9601589..9460836fd 100644 --- a/tests/components/test_sft.py +++ b/tests/components/test_sft.py @@ -2,7 +2,6 @@ """Unit tests for SFT producer component (loop-only, no Ray runtime).""" -import asyncio import sys from types import ModuleType, SimpleNamespace from unittest.mock import AsyncMock, MagicMock @@ -128,7 +127,8 @@ def test_sft_component_imports_without_ray(): from relax.components.sft import SFT # noqa: F401 -def test_sft_step_pushes_one_batch_to_tq(monkeypatch): +@pytest.mark.asyncio +async def test_sft_step_pushes_one_batch_to_tq(monkeypatch): from relax.components.sft import SFT _patch_pipeline_dependencies(monkeypatch) @@ -157,7 +157,7 @@ def test_sft_step_pushes_one_batch_to_tq(monkeypatch): sft._stop_event.is_set = MagicMock(return_value=False) sft._init_data_pipeline() - asyncio.run(sft._produce_one_step()) + await sft._produce_one_step() assert fake_client.async_put.await_count == 1 args_call, kwargs_call = fake_client.async_put.call_args pushed_data = kwargs_call.get("data") @@ -169,7 +169,8 @@ def test_sft_step_pushes_one_batch_to_tq(monkeypatch): assert kwargs_call.get("custom_meta") == [{"total_lengths": 8}] * 4 -def test_sft_step_pushes_sharded_batches_to_tq(monkeypatch): +@pytest.mark.asyncio +async def test_sft_step_pushes_sharded_batches_to_tq(monkeypatch): from relax.components.sft import SFT, _sft_train_partitions_in_flight _patch_pipeline_dependencies(monkeypatch) @@ -202,7 +203,7 @@ def test_sft_step_pushes_sharded_batches_to_tq(monkeypatch): sft._runtime_env = None sft._init_data_pipeline() - asyncio.run(sft._produce_one_step()) + await sft._produce_one_step() assert fake_client.async_put.await_count == 2 seen_partitions = [c.kwargs.get("partition_id") for c in fake_client.async_put.call_args_list] @@ -293,7 +294,8 @@ def shuffle(self, epoch): assert create_dataset.call_args.kwargs["task_type"] == "causal_lm" -def test_sft_remote_batch_producer_reprimes_before_second_step(monkeypatch): +@pytest.mark.asyncio +async def test_sft_remote_batch_producer_reprimes_before_second_step(monkeypatch): from relax.components import sft as sft_module class _FakeIndexManager: @@ -330,16 +332,17 @@ async def get_batch_async(self, batch_size): producer.data_system_client.async_put = AsyncMock() producer._train_size = 16 - asyncio.run(producer.produce_partition(0, "sft_0_shard_0_of_2", 4, False)) + await producer.produce_partition(0, "sft_0_shard_0_of_2", 4, False) producer._dataset._prefetch.set_index_order.assert_not_called() - asyncio.run(producer.produce_partition(1, "sft_1_shard_0_of_2", 4, False)) + await producer.produce_partition(1, "sft_1_shard_0_of_2", 4, False) producer._dataset._prefetch.set_index_order.assert_called_once_with([4, 5]) assert producer.data_system_client.async_put.await_count == 2 +@pytest.mark.asyncio @pytest.mark.parametrize("returned_count", [0, 3]) -def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_count): +async def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_count): from relax.components.sft import SFT fake_ds, _ = _patch_pipeline_dependencies(monkeypatch) @@ -370,13 +373,14 @@ def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_count): sft._init_data_pipeline() with pytest.raises(RuntimeError, match=rf"dataset returned {returned_count}/4 samples"): - asyncio.run(sft._produce_one_step()) + await sft._produce_one_step() fake_client.async_put.assert_not_awaited() assert sft.step == 0 -def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): +@pytest.mark.asyncio +async def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): from relax.components.sft import SFT _patch_pipeline_dependencies(monkeypatch) @@ -403,13 +407,14 @@ def test_sft_eval_rejects_source_with_no_valid_samples(monkeypatch): sft._stop_event.is_set = MagicMock(return_value=False) with pytest.raises(RuntimeError, match="source produced 0 valid samples"): - asyncio.run(sft._maybe_produce_eval()) + await sft._maybe_produce_eval() fake_client.async_put.assert_not_awaited() +@pytest.mark.asyncio @pytest.mark.parametrize("n_real", [1, 3, 4, 5, 8]) -def test_classification_eval_pads_without_dropping_real_samples(n_real): +async def test_classification_eval_pads_without_dropping_real_samples(n_real): from relax.components.sft import SFT samples = [ @@ -441,7 +446,7 @@ def test_classification_eval_pads_without_dropping_real_samples(n_real): sft._build_eval_batches = MagicMock(return_value=samples) sft._wait_for_partition_drained = AsyncMock(return_value=True) - asyncio.run(sft._maybe_produce_eval()) + await sft._maybe_produce_eval() expected_chunks = (n_real + 3) // 4 assert fake_client.async_put.await_count == expected_chunks @@ -451,7 +456,8 @@ def test_classification_eval_pads_without_dropping_real_samples(n_real): assert partition_ids == [f"sft_eval_0_n{expected_chunks}_{idx}" for idx in range(expected_chunks)] -def test_sft_loop_advances_step(monkeypatch): +@pytest.mark.asyncio +async def test_sft_loop_advances_step(monkeypatch): from relax.components.sft import SFT _patch_pipeline_dependencies(monkeypatch) @@ -482,14 +488,15 @@ def test_sft_loop_advances_step(monkeypatch): sft._init_data_pipeline() for _ in range(3): - asyncio.run(sft._produce_one_step()) + await sft._produce_one_step() assert sft.step == 3 assert fake_client.async_put.await_count == 3 seen_partitions = [c.kwargs.get("partition_id") for c in fake_client.async_put.call_args_list] assert seen_partitions == ["sft_0", "sft_1", "sft_2"] -def test_sft_resume_only_produces_remaining_steps(): +@pytest.mark.asyncio +async def test_sft_resume_only_produces_remaining_steps(): from relax.components.sft import SFT SFTCls = SFT.func_or_class @@ -504,14 +511,15 @@ async def _produce_one_step(): sft._produce_one_step = AsyncMock(side_effect=_produce_one_step) - asyncio.run(sft._async_run()) + await sft._async_run() assert sft.step == 5 assert sft._produce_one_step.await_count == 3 +@pytest.mark.asyncio @pytest.mark.parametrize("start_step", [5, 6]) -def test_sft_resume_at_or_after_end_produces_nothing(start_step): +async def test_sft_resume_at_or_after_end_produces_nothing(start_step): from relax.components.sft import SFT SFTCls = SFT.func_or_class @@ -522,6 +530,6 @@ def test_sft_resume_at_or_after_end_produces_nothing(start_step): sft._stop_event.is_set = MagicMock(return_value=False) sft._produce_one_step = AsyncMock() - asyncio.run(sft._async_run()) + await sft._async_run() sft._produce_one_step.assert_not_awaited() From 5781324bef19aadec466713e4a2f1ddecf9a9db0 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 2 Sep 2026 11:40:43 +0800 Subject: [PATCH 14/34] support eval when multi producer sft (cherry picked from commit 7744c284a7c685249fa0de2de6f6bfb269432bb4) --- docs/en/guide/configuration.md | 35 +++++++- docs/en/guide/sft-training.md | 31 +++++++ docs/zh/guide/configuration.md | 35 +++++++- docs/zh/guide/sft-training.md | 31 +++++++ relax/components/sft.py | 128 ++++++++++++++++++++++++--- tests/components/test_sft.py | 155 ++++++++++++++++++++++++++++++++- 6 files changed, 402 insertions(+), 13 deletions(-) diff --git a/docs/en/guide/configuration.md b/docs/en/guide/configuration.md index 083d4f940..d37b1b27f 100644 --- a/docs/en/guide/configuration.md +++ b/docs/en/guide/configuration.md @@ -404,7 +404,7 @@ These flags only apply under `--loss-type sft`. The SFT pipeline runs an `SFTStr | Parameter | Type | Default | Description | |-----------|------|---------|-------------| -| `--eval-size` | float | None | Carve a held-out eval split from `--prompt-data` instead of supplying a separate `--eval-prompt-data`. A value <1 is treated as a fraction of the train dataset (e.g. `0.05` → last 5%); a value ≥1 is treated as an absolute sample count. The reserved tail is removed from the train pool so train and eval samples never overlap. Mutually exclusive with `--eval-prompt-data`. | +| `--eval-size` | float | None | Carve a held-out eval split from `--prompt-data` instead of supplying a separate `--eval-prompt-data`. A value <1 is treated as a fraction of the train dataset (e.g. `0.05` → 5%); a value ≥1 is treated as an absolute sample count. Rows are randomly split once using `--seed`, and the held-out rows are removed from the train pool so train and eval samples never overlap. Mutually exclusive with `--eval-prompt-data`. | | `--sft-predict-interval` | int | None | Every N rollout steps run a generative predict pass over the eval set and write completions to `/predict/predictions_step_.jsonl`. Setting this flag implicitly spins up the Rollout role under SFT (SGLang must be online). Controls the generative complement to the always-on PPL eval (`--eval-interval`). **Requires** `--save` (writes under `/predict/`) and at least one eval source (`--eval-prompt-data` / `--eval-config` / `--eval-size`). | ### Streaming Dataset Prefetch @@ -417,6 +417,39 @@ The SFT producer uses its own `PrefetchBuffer` independent from the rollout data | `--sft-prefetch-chunk-size` | int | 32 | Chunk size dispatched to the SFT prefetch thread-pool per round. | | `--sft-prefetch-num-workers` | int | 4 | Worker threads inside the SFT PrefetchBuffer for I/O-bound media decoding (video/image). | +### Sharded TransferQueue Producers + +`RELAX_SFT_TQ_SHARDS` controls how many TransferQueue partitions are produced for each async-prepacked SFT train step. This is an experimental environment variable so shard counts can be A/B tested without adding a public CLI flag. + +| Environment variable | Type | Default | Description | +|----------------------|------|---------|-------------| +| `RELAX_SFT_TQ_SHARDS` | int | 1 | Number of SFT TransferQueue shards. Values less than or equal to 0 are treated as 1. | + +::: warning Activation requirement +This variable does not enable prepacking. It is effective only with `--loss-type sft --sft-async-prepack`; otherwise Relax uses one partition. Async prepacking also requires `--per-rank-fetch`, at least two in-flight steps (`--max-staleness >= 1` or `--sft-max-in-flight-steps >= 2`), PP=1, CP=1, VPP=1, and THD QKV format. +::: + +With `N > 1`, step `K` uses partitions `sft_K_shard_0_of_N` through `sft_K_shard__of_N`; one shard keeps the existing `sft_K` name. The consumer waits until all shard partitions are ready, then reads an equal slice from each. Consequently, both `global_batch_size` and each DP-local batch (`global_batch_size / data_parallel_size`) must be divisible by `N`. + +When eligible, Relax starts `N` remote `_SFTBatchProducerActor` instances for train batches. The configured `--sft-prefetch-num-workers` is distributed as `ceil(workers / N)` per shard, with at least one worker per shard. Eval is still coordinated locally: `--eval-size` uses the same deterministic train/eval split inside every remote producer, while the coordinator renders the held-out samples (or `--eval-prompt-data`) and pushes `sft_eval__n_` partitions at eval intervals. + +Remote train producers fall back to the local coordinator path in any of these cases: + +- Ray is not initialized. +- `--task-type seq_cls` is used. +- `--custom-dataset-class-path` is set. +- `--sft-oversize-strategy` is `skip` or `custom`, or `--sft-invalid-multimodal-strategy` is `skip`. + +The local fallback still splits the batch into `N` TransferQueue partitions, but it does not create `N` remote producer actors. Look for `SFT remote shard producer enabled: ... shards=N ...` to confirm that remote producer parallelism is active; fallback paths log `SFT remote shard producer disabled: ...` with the reason. + +Configure the value in the Ray runtime environment so the producer and consumer derive identical partition names: + +```yaml +# configs/env.yaml +env_vars: + RELAX_SFT_TQ_SHARDS: "2" +``` + ### Oversize Sample Handling How the SFT producer handles samples whose tokenized + media-expanded length exceeds the per-GPU capacity (`--max-tokens-per-gpu × --context-parallel-size`). All branches log a WARNING per oversized sample. diff --git a/docs/en/guide/sft-training.md b/docs/en/guide/sft-training.md index c745a1e65..35cbc37b3 100644 --- a/docs/en/guide/sft-training.md +++ b/docs/en/guide/sft-training.md @@ -419,6 +419,36 @@ Pokemon 1 GPU script: `"sft": [1, 0]` means the SFT producer is CPU-only. The Actor owns training GPUs. Rollout GPUs are needed when periodic predict is enabled. +### Sharded SFT Producers + +`RELAX_SFT_TQ_SHARDS` is an experimental throughput knob for SFT async prepacking. It splits each global SFT batch across multiple TransferQueue partitions. Configure it in the Ray runtime environment: + +```yaml +# configs/env.yaml +env_vars: + RELAX_SFT_TQ_SHARDS: "2" +``` + +Enable async prepacking in the training arguments as well: + +```bash +--per-rank-fetch +--sft-async-prepack +--sft-max-in-flight-steps 4 +``` + +`--sft-max-in-flight-steps 4` is an example; async prepacking requires at least 2, or equivalently `--max-staleness >= 1`. The environment variable does not enable prepacking by itself and is ignored without `--loss-type sft --sft-async-prepack`. + +For `N` shards, both `--global-batch-size` and the per-DP-rank local batch (`global_batch_size / data_parallel_size`) must be divisible by `N`. For example, with global batch size 32 and DP size 8, the local batch is 4, so 2 or 4 shards are valid but 3 is not. + +When the remote path is eligible, `N` shards launch `N` Ray producer actors for train batches. Eval can be enabled at the same time: train shards are produced remotely, while the coordinator renders the eval split or eval prompt data and pushes the usual `sft_eval__n_` partitions at eval intervals. Some configurations, including sequence classification, custom datasets, and skip-capable sample filtering, use one local producer that still writes `N` partitions. Therefore, `RELAX_SFT_TQ_SHARDS=2` does not always mean two producer actors. Confirm the remote path from this log: + +```text +SFT remote shard producer enabled: ... shards=2 ... +``` + +When using `--eval-size`, every remote train producer applies the same seed-based held-out split so train and eval rows do not overlap. When launching Python directly, `export RELAX_SFT_TQ_SHARDS=2` is also supported. For a Ray Job, putting the value in `configs/env.yaml` ensures that producers and consumers receive the same value. See [Sharded TransferQueue Producers](./configuration.md#sharded-transferqueue-producers) for partition naming, fallback conditions, and worker allocation. + ## Launch ### Single Node @@ -503,6 +533,7 @@ If GPUs wait on SFT data: | `--sft-prefetch-chunk-size` | Increase | Dispatches larger prefetch chunks, with higher memory pressure. | | `--per-rank-fetch` | Enable for multi-GPU | Lets TP/PP ranks pull from TransferQueue directly. Pair with enough `--num-data-storage-units`. | | `--max-staleness` | Increase for I/O-heavy SFT | Lets the producer run ahead. The Pokemon 8 GPU script uses `--max-staleness 4`. | +| `RELAX_SFT_TQ_SHARDS` | Start from 2 | Parallelizes eligible async-prepack producers. Increase only when data preparation is the bottleneck and batch divisibility constraints are met. | For text-only math, prefetch usually matters less than sequence length and model parallelism. For Pokemon, image loading and processor work are common bottlenecks. diff --git a/docs/zh/guide/configuration.md b/docs/zh/guide/configuration.md index 4436d93ad..38948a3db 100644 --- a/docs/zh/guide/configuration.md +++ b/docs/zh/guide/configuration.md @@ -404,7 +404,7 @@ PPO 当前支持同步 colocate 模式，并要求在 `--resource` 中包含 `cr | 参数 | 类型 | 默认值 | 说明 | |------|------|--------|------| -| `--eval-size` | float | None | 从 `--prompt-data` 切出一份 holdout eval 集，而不是另外指定 `--eval-prompt-data`。值 <1 视为训练集的占比（例如 `0.05` → 末尾 5%）；值 ≥1 视为绝对样本数。被预留的尾部会从训练池里移除，所以训练样本和 eval 样本永不重叠。与 `--eval-prompt-data` 互斥。 | +| `--eval-size` | float | None | 从 `--prompt-data` 切出一份 holdout eval 集，而不是另外指定 `--eval-prompt-data`。值 <1 视为训练集的占比（例如 `0.05` → 5%）；值 ≥1 视为绝对样本数。行 ID 会用 `--seed` 随机切分一次，被预留的行会从训练池里移除，所以训练样本和 eval 样本永不重叠。与 `--eval-prompt-data` 互斥。 | | `--sft-predict-interval` | int | None | 每 N 个 rollout step 在 eval 集上跑一次生成式 predict，把生成结果写到 `/predict/predictions_step_.jsonl`。设置该参数后会自动拉起 Rollout 角色（SGLang 必须在线）。它是 always-on 的 PPL eval（`--eval-interval`）的生成式补充。**必需**：`--save`（写到 `/predict/` 下）以及至少一个 eval 数据源（`--eval-prompt-data` / `--eval-config` / `--eval-size`）。 | ### 流式数据集预取 @@ -417,6 +417,39 @@ SFT producer 用自己的 `PrefetchBuffer`，跟 rollout 数据源的 `--prefetc | `--sft-prefetch-chunk-size` | int | 32 | 每轮派发给 SFT 预取线程池的 chunk 大小。 | | `--sft-prefetch-num-workers` | int | 4 | SFT PrefetchBuffer 内部用于 I/O 密集型媒体解码（视频/图像）的工作线程数。 | +### TransferQueue 分片 Producer + +`RELAX_SFT_TQ_SHARDS` 控制每个 async-prepacked SFT 训练 step 生成多少个 TransferQueue 分区。这是一个实验性环境变量，便于在不增加公开 CLI 参数的情况下对 shard 数做 A/B 测试。 + +| 环境变量 | 类型 | 默认值 | 说明 | +|----------|------|--------|------| +| `RELAX_SFT_TQ_SHARDS` | int | 1 | SFT TransferQueue shard 数。小于等于 0 的值按 1 处理。 | + +::: warning 生效条件 +这个变量本身不会开启 prepack。它仅在使用 `--loss-type sft --sft-async-prepack` 时生效，否则 Relax 只使用一个分区。Async prepack 还要求开启 `--per-rank-fetch`、至少允许两个 in-flight step（`--max-staleness >= 1` 或 `--sft-max-in-flight-steps >= 2`）、PP=1、CP=1、VPP=1，并使用 THD QKV 格式。 +::: + +当 `N > 1` 时，step `K` 使用从 `sft_K_shard_0_of_N` 到 `sft_K_shard__of_N` 的分区；单 shard 仍使用原有的 `sft_K` 名称。Consumer 会等待所有 shard 分区就绪，然后从每个分区读取相同数量的样本。因此，`global_batch_size` 和每个 DP rank 的本地 batch（`global_batch_size / data_parallel_size`）都必须能被 `N` 整除。 + +满足条件时，Relax 会为 train batch 启动 `N` 个远端 `_SFTBatchProducerActor`。配置的 `--sft-prefetch-num-workers` 会按每个 shard `ceil(workers / N)` 分配，且每个 shard 至少有一个 worker。Eval 仍由 coordinator 本地协调：`--eval-size` 会在每个远端 producer 内使用同一份确定性的 train/eval split，coordinator 负责渲染 holdout 样本（或 `--eval-prompt-data`）并在 eval interval 推送 `sft_eval__n_` 分区。 + +遇到以下任一情况时，remote train producer 会 fallback 到本地 coordinator 路径： + +- Ray 未初始化。 +- 使用 `--task-type seq_cls`。 +- 设置了 `--custom-dataset-class-path`。 +- `--sft-oversize-strategy` 为 `skip` 或 `custom`，或者 `--sft-invalid-multimodal-strategy` 为 `skip`。 + +本地 fallback 仍会把 batch 拆成 `N` 个 TransferQueue 分区，但不会创建 `N` 个远端 producer actor。日志出现 `SFT remote shard producer enabled: ... shards=N ...` 才表示 remote producer 并行已生效；fallback 路径会记录 `SFT remote shard producer disabled: ...` 及具体原因。 + +请在 Ray 运行时环境中配置这个值，确保 producer 和 consumer 推导出相同的分区名： + +```yaml +# configs/env.yaml +env_vars: + RELAX_SFT_TQ_SHARDS: "2" +``` + ### 超长样本处理 SFT producer 如何处理 tokenize + 多模态展开后长度超过单卡容量（`--max-tokens-per-gpu × --context-parallel-size`）的样本。所有分支都会给每个超长样本打一条 WARNING 日志。 diff --git a/docs/zh/guide/sft-training.md b/docs/zh/guide/sft-training.md index 07961b314..39471562e 100644 --- a/docs/zh/guide/sft-training.md +++ b/docs/zh/guide/sft-training.md @@ -419,6 +419,36 @@ Pokemon 1 GPU 脚本： `"sft": [1, 0]` 表示 SFT producer 是 CPU-only。Actor 使用训练 GPU。开启周期性 predict 时，Rollout 也需要 GPU 资源。 +### SFT 分片 Producer + +`RELAX_SFT_TQ_SHARDS` 是 SFT async prepack 的实验性吞吐调优开关，用于把每个全局 SFT batch 拆到多个 TransferQueue 分区。推荐在 Ray 运行时环境中配置： + +```yaml +# configs/env.yaml +env_vars: + RELAX_SFT_TQ_SHARDS: "2" +``` + +同时在训练参数中开启 async prepack： + +```bash +--per-rank-fetch +--sft-async-prepack +--sft-max-in-flight-steps 4 +``` + +`--sft-max-in-flight-steps 4` 只是示例；async prepack 要求该值至少为 2，或者等价地设置 `--max-staleness >= 1`。这个环境变量本身不会开启 prepack；未使用 `--loss-type sft --sft-async-prepack` 时会被忽略。 + +使用 `N` 个 shard 时，`--global-batch-size` 和每个 DP rank 的本地 batch（`global_batch_size / data_parallel_size`）都必须能被 `N` 整除。例如 global batch size 为 32、DP size 为 8 时，本地 batch 为 4，因此可以设置 2 或 4 个 shard，不能设置 3 个。 + +满足 remote 路径条件时，`N` 个 shard 会为 train batch 启动 `N` 个 Ray producer actor。Eval 可以同时开启：train shard 由远端 producer 生产，coordinator 负责渲染 eval split 或 eval prompt data，并在 eval interval 推送原有的 `sft_eval__n_` 分区。部分配置（包括序列分类、自定义数据集和允许跳过样本的过滤策略）会使用单个本地 producer，但仍然写入 `N` 个分区。因此，`RELAX_SFT_TQ_SHARDS=2` 并不一定表示有两个 producer actor。可以通过下面的日志确认 remote 路径是否真正启用： + +```text +SFT remote shard producer enabled: ... shards=2 ... +``` + +使用 `--eval-size` 时，每个远端 train producer 都会应用同一份基于 `--seed` 的 holdout split，保证 train 和 eval 行不重叠。直接启动 Python 时也可以使用 `export RELAX_SFT_TQ_SHARDS=2`。通过 Ray Job 启动时，将它写入 `configs/env.yaml` 可以确保 producer 和 consumer 收到相同的值。分区命名、fallback 条件和 worker 分配方式见 [TransferQueue 分片 Producer](./configuration.md#transferqueue-分片-producer)。 + ## 启动 ### 单机 @@ -503,6 +533,7 @@ bash scripts/entrypoint/spmd-multinode.sh \ | `--sft-prefetch-chunk-size` | 调大 | 一次派发更多预取样本，但会增加内存压力。 | | `--per-rank-fetch` | 多 GPU 时开启 | 让 TP/PP rank 直接从 TransferQueue 拉数据，需配足 `--num-data-storage-units`。 | | `--max-staleness` | I/O 重时调大 | 允许 producer 提前生产。Pokemon 8 GPU 脚本使用 `--max-staleness 4`。 | +| `RELAX_SFT_TQ_SHARDS` | 从 2 开始 | 并行化满足条件的 async-prepack producer。仅当数据准备是瓶颈且 batch 满足整除约束时再增加。 | 纯文本 math 任务通常更受序列长度和模型并行影响；Pokemon 任务更容易被图片读取和 processor 工作拖慢。 diff --git a/relax/components/sft.py b/relax/components/sft.py index 9b10a34bf..5d4126741 100644 --- a/relax/components/sft.py +++ b/relax/components/sft.py @@ -350,12 +350,25 @@ def initialize(self, start_step: int) -> dict[str, Any]: self._train_size = n_avail eval_size_arg = getattr(self.config, "eval_size", None) if eval_size_arg is not None: - if eval_size_arg < 1: - n_eval = max(1, int(n_avail * eval_size_arg)) - else: - n_eval = int(eval_size_arg) - n_eval = min(n_eval, max(n_avail - 1, 0)) - self._train_size = n_avail - n_eval if n_eval > 0 else n_avail + train_indices, eval_indices = resolve_sft_split_indices( + n_avail, + eval_size_arg, + getattr(self.config, "seed", 42), + ) + self._train_size = len(train_indices) + if eval_indices: + restrict_training_indices = getattr(self._dataset, "restrict_training_indices", None) + if not callable(restrict_training_indices): + raise TypeError( + "--eval-size requires the SFT dataset to implement restrict_training_indices(indices) " + "for a deterministic remote-producer split." + ) + restrict_training_indices(train_indices) + self._logger.info( + f"--eval-size randomly held out {len(eval_indices)} samples with seed=" + f"{getattr(self.config, 'seed', 42)}; remote shard {self._shard_id}/{self._num_shards} " + f"train pool size now {self._train_size}." + ) shard_batch_size = self._shard_batch_size(self.config.global_batch_size) if self._train_size > 0: @@ -484,9 +497,6 @@ def _should_use_remote_batch_producer(self) -> bool: if not ray.is_initialized(): self._logger.info("SFT remote shard producer disabled: Ray is not initialized.") return False - if _has_sft_eval_work(self.config): - self._logger.info("SFT remote shard producer disabled: eval is configured; using local producer path.") - return False if getattr(self.config, "custom_dataset_class_path", None): self._logger.info( "SFT remote shard producer disabled: custom dataset is configured; using local producer path." @@ -519,18 +529,115 @@ def _init_remote_batch_producers(self) -> None: for shard_id in range(num_shards) ] states = ray.get([producer.initialize.remote(self.step) for producer in self._batch_producers]) - self._train_size = int(states[0].get("train_size") or 0) + train_sizes = {int(state.get("train_size") or 0) for state in states} + if len(train_sizes) != 1: + raise RuntimeError(f"SFT remote shard producer train-size mismatch: states={states}") + self._train_size = train_sizes.pop() self._logger.info( f"SFT remote shard producer enabled: dataset={states[0].get('dataset')} " f"train_size={self._train_size} shards={num_shards} " f"prefetch_workers_per_shard={prefetch_num_workers}" ) + def _init_remote_eval_pipeline(self) -> None: + eval_prompt_data = build_named_prompt_data_configs(getattr(self.config, "eval_prompt_data", None)) + eval_size_arg = getattr(self.config, "eval_size", None) + if eval_size_arg is None and not eval_prompt_data: + return + + prepare_model_maybe_update_args(self.config, completeness="metadata") + self._tokenizer = AutoTokenizer.from_pretrained(self.config.hf_checkpoint, trust_remote_code=True) + try: + self._processor_pool = ProcessorPool(self.config.hf_checkpoint, pool_size=None, trust_remote_code=True) + except Exception as exc: + self._logger.warning(f"Could not init ProcessorPool ({exc}); multimodal eval samples will fail at push.") + self._processor_pool = None + pad_token_ids = _resolve_pad_token_ids_from_config(self.config.hf_checkpoint) + self._logger.info(f"Resolved multimodal pad token ids from model config: {sorted(pad_token_ids)}") + + cp_size = max(1, getattr(self.config, "context_parallel_size", 1) or 1) + capacity = self.config.max_tokens_per_gpu * cp_size + seed = getattr(self.config, "seed", 42) + task_type = getattr(self.config, "task_type", "causal_lm") + classification_sentinel_token_id = ( + _resolve_classification_sentinel_token_id(self._tokenizer) if task_type == "seq_cls" else None + ) + dataset_options = _resolve_sft_dataset_options(self.config, self._logger) + + if eval_size_arg is not None: + self._dataset = _create_sft_train_dataset( + self.config, + tokenizer=self._tokenizer, + processor_pool=self._processor_pool, + capacity=capacity, + prefetch_buffer_size=0, + prefetch_chunk_size=getattr(self.config, "sft_prefetch_chunk_size", 32), + prefetch_num_workers=1, + pad_token_ids=pad_token_ids, + task_type=task_type, + classification_sentinel_token_id=classification_sentinel_token_id, + **dataset_options, + ) + n_avail = len(self._dataset) + train_indices, eval_indices = resolve_sft_split_indices(n_avail, eval_size_arg, seed) + if self._train_size != len(train_indices): + raise RuntimeError( + "SFT remote eval split mismatch between coordinator and shard producers: " + f"coordinator train_size={len(train_indices)}, remote train_size={self._train_size}." + ) + n_eval = len(eval_indices) + if n_eval == 0: + self._logger.warning( + f"--eval-size {eval_size_arg} resolves to 0 samples on a dataset of size {n_avail}; " + "eval will be skipped." + ) + else: + get_batch_by_indices = getattr(self._dataset, "get_batch_by_indices", None) + if not callable(get_batch_by_indices): + raise TypeError( + "--eval-size requires the SFT dataset to implement get_batch_by_indices(indices) " + "for deterministic remote-producer eval." + ) + self._eval_indices = eval_indices + self._logger.info(f"SFT remote eval split initialized: held out {n_eval} samples with seed={seed}.") + return + + eval_input_key = getattr(self.config, "eval_input_key", None) or self.config.input_key + eval_label_key = getattr(self.config, "eval_label_key", None) or self.config.label_key + eval_tool_key = getattr(self.config, "eval_tool_key", None) or self.config.tool_key + self._eval_dataset = SFTStreamingDataset( + path=[d.path for d in eval_prompt_data], + tokenizer=self._tokenizer, + processor_pool=self._processor_pool, + capacity=capacity, + prompt_key=eval_input_key, + label_key=eval_label_key, + multimodal_keys=self.config.multimodal_keys, + conversation_key_map=getattr(self.config, "conversation_key_map", None), + metadata_key=self.config.metadata_key, + tool_key=eval_tool_key, + system_prompt=self.config.system_prompt, + source_name="+".join(d.name for d in eval_prompt_data), + seed=seed, + prefetch_max_cached=0, + pad_token_ids=pad_token_ids, + oversize_strategy=dataset_options["oversize_strategy"], + oversize_custom_fn=dataset_options["oversize_custom_fn"], + invalid_multimodal_strategy=dataset_options["invalid_multimodal_strategy"], + apply_chat_template_kwargs=getattr(self.config, "apply_chat_template_kwargs", None), + require_response=task_type != "seq_cls", + task_type=task_type, + num_labels=getattr(self.config, "num_labels", None), + problem_type=getattr(self.config, "problem_type", "single_label_classification"), + classification_sentinel_token_id=classification_sentinel_token_id, + ) + def _init_data_pipeline(self) -> None: if self._dataset is not None or getattr(self, "_batch_producers", None): return if self._should_use_remote_batch_producer(): self._init_remote_batch_producers() + self._init_remote_eval_pipeline() return prepare_model_maybe_update_args(self.config, completeness="metadata") self._tokenizer = AutoTokenizer.from_pretrained(self.config.hf_checkpoint, trust_remote_code=True) @@ -769,6 +876,7 @@ async def _produce_one_step(self) -> None: current_epoch = max(payload.get("epoch") or 0 for payload in payloads) if crossed_epoch: self._logger.info(f"SFT step {self.step}: epoch boundary crossed (epoch={current_epoch})") + await self._maybe_produce_eval() self.step += 1 return diff --git a/tests/components/test_sft.py b/tests/components/test_sft.py index 9460836fd..16e404857 100644 --- a/tests/components/test_sft.py +++ b/tests/components/test_sft.py @@ -217,7 +217,6 @@ async def test_sft_step_pushes_sharded_batches_to_tq(monkeypatch): ("attribute", "value"), [ ("task_type", "seq_cls"), - ("eval_size", 0.25), ], ) def test_sft_remote_batch_producer_falls_back_for_target_only_modes(monkeypatch, attribute, value): @@ -237,6 +236,23 @@ def test_sft_remote_batch_producer_falls_back_for_target_only_modes(monkeypatch, assert sft._should_use_remote_batch_producer() is False +def test_sft_remote_batch_producer_allows_eval_size(monkeypatch): + from relax.components.sft import SFT + + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "2") + monkeypatch.setattr("relax.components.sft.ray.is_initialized", lambda: True) + + args = _make_args(global_batch_size=4) + args.sft_async_prepack = True + args.eval_size = 0.25 + SFTCls = SFT.func_or_class + sft = SFTCls.__new__(SFTCls) + sft.config = args + sft._logger_instance = MagicMock() + + assert sft._should_use_remote_batch_producer() is True + + def test_sft_remote_batch_producer_resolves_model_on_its_node(monkeypatch): from relax.components import sft as sft_module @@ -294,6 +310,47 @@ def shuffle(self, epoch): assert create_dataset.call_args.kwargs["task_type"] == "causal_lm" +def test_sft_remote_batch_producer_restricts_train_pool_for_eval_size(monkeypatch): + from relax.components import sft as sft_module + + class _FakeIndexManager: + total_size = 10 + + def __init__(self): + self.current_epoch = -1 + self.position = 0 + self.indices = list(range(self.total_size)) + + def shuffle(self, epoch): + self.current_epoch = epoch + self.position = 0 + self.indices = list(range(self.total_size)) + + fake_dataset = MagicMock() + fake_dataset.__len__ = MagicMock(return_value=10) + fake_dataset.index_manager = _FakeIndexManager() + fake_dataset._prefetch = None + + monkeypatch.setattr(sft_module.tq, "init", MagicMock()) + monkeypatch.setattr(sft_module.tq, "get_client", MagicMock()) + monkeypatch.setattr(sft_module, "prepare_model_maybe_update_args", MagicMock()) + monkeypatch.setattr(sft_module.AutoTokenizer, "from_pretrained", MagicMock(return_value=MagicMock())) + monkeypatch.setattr(sft_module, "ProcessorPool", MagicMock()) + monkeypatch.setattr(sft_module, "_resolve_pad_token_ids_from_config", MagicMock(return_value=frozenset())) + monkeypatch.setattr(sft_module, "_create_sft_train_dataset", MagicMock(return_value=fake_dataset)) + + args = _make_args(global_batch_size=2) + args.eval_size = 0.2 + producer_cls = sft_module._SFTBatchProducerActor.__ray_metadata__.modified_class + producer = producer_cls(args, shard_id=0, num_shards=2, prefetch_num_workers=1) + + state = producer.initialize(start_step=0) + + train_indices, _eval_indices = resolve_sft_split_indices(10, 0.2, seed=args.seed) + assert state["train_size"] == 8 + fake_dataset.restrict_training_indices.assert_called_once_with(train_indices) + + @pytest.mark.asyncio async def test_sft_remote_batch_producer_reprimes_before_second_step(monkeypatch): from relax.components import sft as sft_module @@ -340,6 +397,102 @@ async def get_batch_async(self, batch_size): assert producer.data_system_client.async_put.await_count == 2 +def test_sft_remote_eval_size_initializes_eval_dataset_without_train_overlap(monkeypatch): + from relax.components.sft import SFT + + fake_ds, _ = _patch_pipeline_dependencies(monkeypatch, n_samples=10) + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "2") + monkeypatch.setattr("relax.components.sft.ray.is_initialized", lambda: True) + monkeypatch.setattr("relax.components.sft.tq.init", lambda *a, **kw: None) + monkeypatch.setattr("relax.components.sft.tq.get_client", MagicMock()) + + args = _make_args(global_batch_size=2) + args.sft_async_prepack = True + args.eval_size = 0.2 + SFTCls = SFT.func_or_class + sft = SFTCls.__new__(SFTCls) + sft.config = args + sft.role = "sft" + sft.step = 0 + sft._dataset = None + sft._eval_dataset = None + sft._eval_indices = None + sft._batch_producers = [] + sft._train_size = 0 + sft._tokenizer = None + sft._processor_pool = None + sft._logger_instance = MagicMock() + sft._runtime_env = None + sft._init_remote_batch_producers = MagicMock(side_effect=lambda: setattr(sft, "_train_size", 8)) + + sft._init_data_pipeline() + + train_indices, eval_indices = resolve_sft_split_indices(10, 0.2, seed=args.seed) + assert sft._train_size == len(train_indices) + assert sft._eval_indices == eval_indices + fake_ds.restrict_training_indices.assert_not_called() + assert [sample.source_idx for sample in sft._build_eval_batches()] == list(eval_indices) + + +@pytest.mark.asyncio +async def test_sft_remote_step_produces_eval_before_advancing_step(monkeypatch): + from relax.components import sft as sft_module + from relax.components.sft import SFT + + class _RemoteProduce: + def __init__(self): + self.calls = [] + + def remote(self, step, partition_id, global_batch_size, force_multimodal_field): + self.calls.append((step, partition_id, global_batch_size, force_multimodal_field)) + return object() + + class _RemoteProducer: + def __init__(self): + self.produce_partition = _RemoteProduce() + + monkeypatch.setenv("RELAX_SFT_TQ_SHARDS", "2") + monkeypatch.setattr( + sft_module, + "_ray_get_many_async", + AsyncMock( + return_value=[ + {"crossed_epoch": False, "epoch": 0}, + {"crossed_epoch": False, "epoch": 0}, + ] + ), + ) + + args = _make_args(global_batch_size=4) + args.sft_async_prepack = True + SFTCls = SFT.func_or_class + sft = SFTCls.__new__(SFTCls) + sft.config = args + sft.role = "sft" + sft.step = 0 + sft.data_system_client = MagicMock() + sft._dataset = None + sft._eval_dataset = MagicMock() + sft._eval_indices = None + sft._batch_producers = [_RemoteProducer(), _RemoteProducer()] + sft._train_size = 8 + sft._tokenizer = None + sft._processor_pool = None + sft._logger_instance = MagicMock() + sft._stop_event = MagicMock() + sft._stop_event.is_set = MagicMock(return_value=False) + + async def _maybe_produce_eval(): + assert sft.step == 0 + + sft._maybe_produce_eval = AsyncMock(side_effect=_maybe_produce_eval) + + await sft._produce_one_step() + + sft._maybe_produce_eval.assert_awaited_once() + assert sft.step == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("returned_count", [0, 3]) async def test_sft_step_rejects_empty_or_partial_batch(monkeypatch, returned_count): From 6b19d6ed34dcffafc1a67b593b66e2254a9f8c53 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 2 Sep 2026 16:13:47 +0800 Subject: [PATCH 15/34] modify some docs (cherry picked from commit 0c806f6f83144fbdeb453dcf2da58144117f1b52) --- docs/en/guide/configuration.md | 4 ++-- docs/en/guide/sft-training.md | 2 +- docs/zh/guide/configuration.md | 4 ++-- docs/zh/guide/sft-training.md | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/en/guide/configuration.md b/docs/en/guide/configuration.md index d37b1b27f..21c679813 100644 --- a/docs/en/guide/configuration.md +++ b/docs/en/guide/configuration.md @@ -429,7 +429,7 @@ The SFT producer uses its own `PrefetchBuffer` independent from the rollout data This variable does not enable prepacking. It is effective only with `--loss-type sft --sft-async-prepack`; otherwise Relax uses one partition. Async prepacking also requires `--per-rank-fetch`, at least two in-flight steps (`--max-staleness >= 1` or `--sft-max-in-flight-steps >= 2`), PP=1, CP=1, VPP=1, and THD QKV format. ::: -With `N > 1`, step `K` uses partitions `sft_K_shard_0_of_N` through `sft_K_shard__of_N`; one shard keeps the existing `sft_K` name. The consumer waits until all shard partitions are ready, then reads an equal slice from each. Consequently, both `global_batch_size` and each DP-local batch (`global_batch_size / data_parallel_size`) must be divisible by `N`. +With `N > 1`, step `K` uses partitions `sft_K_shard_0_of_N` through `sft_K_shard__of_N`; the existing `sft_K` name is used only when `N == 1`. The consumer waits until all shard partitions are ready, then reads an equal slice from each. Consequently, both `global_batch_size` and each DP-local batch (`global_batch_size / data_parallel_size`) must be divisible by `N`. When eligible, Relax starts `N` remote `_SFTBatchProducerActor` instances for train batches. The configured `--sft-prefetch-num-workers` is distributed as `ceil(workers / N)` per shard, with at least one worker per shard. Eval is still coordinated locally: `--eval-size` uses the same deterministic train/eval split inside every remote producer, while the coordinator renders the held-out samples (or `--eval-prompt-data`) and pushes `sft_eval__n_` partitions at eval intervals. @@ -437,7 +437,7 @@ Remote train producers fall back to the local coordinator path in any of these c - Ray is not initialized. - `--task-type seq_cls` is used. -- `--custom-dataset-class-path` is set. +- `--custom-dataset-class` / `--custom-dataset-class-path` is set. - `--sft-oversize-strategy` is `skip` or `custom`, or `--sft-invalid-multimodal-strategy` is `skip`. The local fallback still splits the batch into `N` TransferQueue partitions, but it does not create `N` remote producer actors. Look for `SFT remote shard producer enabled: ... shards=N ...` to confirm that remote producer parallelism is active; fallback paths log `SFT remote shard producer disabled: ...` with the reason. diff --git a/docs/en/guide/sft-training.md b/docs/en/guide/sft-training.md index 35cbc37b3..7b85b5227 100644 --- a/docs/en/guide/sft-training.md +++ b/docs/en/guide/sft-training.md @@ -441,7 +441,7 @@ Enable async prepacking in the training arguments as well: For `N` shards, both `--global-batch-size` and the per-DP-rank local batch (`global_batch_size / data_parallel_size`) must be divisible by `N`. For example, with global batch size 32 and DP size 8, the local batch is 4, so 2 or 4 shards are valid but 3 is not. -When the remote path is eligible, `N` shards launch `N` Ray producer actors for train batches. Eval can be enabled at the same time: train shards are produced remotely, while the coordinator renders the eval split or eval prompt data and pushes the usual `sft_eval__n_` partitions at eval intervals. Some configurations, including sequence classification, custom datasets, and skip-capable sample filtering, use one local producer that still writes `N` partitions. Therefore, `RELAX_SFT_TQ_SHARDS=2` does not always mean two producer actors. Confirm the remote path from this log: +When the remote path is eligible, `N` shards launch `N` Ray producer actors for train batches. Eval can be enabled at the same time: train shards are produced remotely, while the coordinator renders the eval split or eval prompt data and pushes the usual `sft_eval__n_` partitions at eval intervals. Some configurations, including sequence classification, custom datasets (`--custom-dataset-class` / `--custom-dataset-class-path`), and skip-capable sample filtering, use one local producer that still writes `N` partitions. Therefore, `RELAX_SFT_TQ_SHARDS=2` does not always mean two producer actors. Confirm the remote path from this log: ```text SFT remote shard producer enabled: ... shards=2 ... diff --git a/docs/zh/guide/configuration.md b/docs/zh/guide/configuration.md index 38948a3db..0561a77b4 100644 --- a/docs/zh/guide/configuration.md +++ b/docs/zh/guide/configuration.md @@ -429,7 +429,7 @@ SFT producer 用自己的 `PrefetchBuffer`，跟 rollout 数据源的 `--prefetc 这个变量本身不会开启 prepack。它仅在使用 `--loss-type sft --sft-async-prepack` 时生效，否则 Relax 只使用一个分区。Async prepack 还要求开启 `--per-rank-fetch`、至少允许两个 in-flight step（`--max-staleness >= 1` 或 `--sft-max-in-flight-steps >= 2`）、PP=1、CP=1、VPP=1，并使用 THD QKV 格式。 ::: -当 `N > 1` 时，step `K` 使用从 `sft_K_shard_0_of_N` 到 `sft_K_shard__of_N` 的分区；单 shard 仍使用原有的 `sft_K` 名称。Consumer 会等待所有 shard 分区就绪，然后从每个分区读取相同数量的样本。因此，`global_batch_size` 和每个 DP rank 的本地 batch（`global_batch_size / data_parallel_size`）都必须能被 `N` 整除。 +当 `N > 1` 时，step `K` 使用从 `sft_K_shard_0_of_N` 到 `sft_K_shard__of_N` 的分区；原有的 `sft_K` 名称只会在 `N == 1` 时使用。Consumer 会等待所有 shard 分区就绪，然后从每个分区读取相同数量的样本。因此，`global_batch_size` 和每个 DP rank 的本地 batch（`global_batch_size / data_parallel_size`）都必须能被 `N` 整除。 满足条件时，Relax 会为 train batch 启动 `N` 个远端 `_SFTBatchProducerActor`。配置的 `--sft-prefetch-num-workers` 会按每个 shard `ceil(workers / N)` 分配，且每个 shard 至少有一个 worker。Eval 仍由 coordinator 本地协调：`--eval-size` 会在每个远端 producer 内使用同一份确定性的 train/eval split，coordinator 负责渲染 holdout 样本（或 `--eval-prompt-data`）并在 eval interval 推送 `sft_eval__n_` 分区。 @@ -437,7 +437,7 @@ SFT producer 用自己的 `PrefetchBuffer`，跟 rollout 数据源的 `--prefetc - Ray 未初始化。 - 使用 `--task-type seq_cls`。 -- 设置了 `--custom-dataset-class-path`。 +- 设置了 `--custom-dataset-class` / `--custom-dataset-class-path`。 - `--sft-oversize-strategy` 为 `skip` 或 `custom`，或者 `--sft-invalid-multimodal-strategy` 为 `skip`。 本地 fallback 仍会把 batch 拆成 `N` 个 TransferQueue 分区，但不会创建 `N` 个远端 producer actor。日志出现 `SFT remote shard producer enabled: ... shards=N ...` 才表示 remote producer 并行已生效；fallback 路径会记录 `SFT remote shard producer disabled: ...` 及具体原因。 diff --git a/docs/zh/guide/sft-training.md b/docs/zh/guide/sft-training.md index 39471562e..7bfedcd94 100644 --- a/docs/zh/guide/sft-training.md +++ b/docs/zh/guide/sft-training.md @@ -441,7 +441,7 @@ env_vars: 使用 `N` 个 shard 时，`--global-batch-size` 和每个 DP rank 的本地 batch（`global_batch_size / data_parallel_size`）都必须能被 `N` 整除。例如 global batch size 为 32、DP size 为 8 时，本地 batch 为 4，因此可以设置 2 或 4 个 shard，不能设置 3 个。 -满足 remote 路径条件时，`N` 个 shard 会为 train batch 启动 `N` 个 Ray producer actor。Eval 可以同时开启：train shard 由远端 producer 生产，coordinator 负责渲染 eval split 或 eval prompt data，并在 eval interval 推送原有的 `sft_eval__n_` 分区。部分配置（包括序列分类、自定义数据集和允许跳过样本的过滤策略）会使用单个本地 producer，但仍然写入 `N` 个分区。因此，`RELAX_SFT_TQ_SHARDS=2` 并不一定表示有两个 producer actor。可以通过下面的日志确认 remote 路径是否真正启用： +满足 remote 路径条件时，`N` 个 shard 会为 train batch 启动 `N` 个 Ray producer actor。Eval 可以同时开启：train shard 由远端 producer 生产，coordinator 负责渲染 eval split 或 eval prompt data，并在 eval interval 推送原有的 `sft_eval__n_` 分区。部分配置（包括序列分类、自定义数据集（`--custom-dataset-class` / `--custom-dataset-class-path`）和允许跳过样本的过滤策略）会使用单个本地 producer，但仍然写入 `N` 个分区。因此，`RELAX_SFT_TQ_SHARDS=2` 并不一定表示有两个 producer actor。可以通过下面的日志确认 remote 路径是否真正启用： ```text SFT remote shard producer enabled: ... shards=2 ... From 2f8e183a1c6eb8e45d49649447ded12678aa255f Mon Sep 17 00:00:00 2001 From: yangrui6 Date: Wed, 2 Sep 2026 20:14:00 +0800 Subject: [PATCH 16/34] fix(multimodal): load truncated images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 🐛 Bug Fix ## Allow Pillow to decode truncated image payloads - Enable Pillow's process-wide truncated image compatibility for multimodal loaders - Prevent SFT producers from failing on recoverable image tail truncation --- # ✅ Tests ## Cover truncated JPEG loading - Verify a JPEG with four missing trailing bytes is decoded successfully (cherry picked from commit 0ec114c897e4f17b3f06182f775d55ec624de511) --- relax/utils/multimodal/image_utils.py | 5 ++++- tests/utils/multimodal/test_image_utils.py | 10 ++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/relax/utils/multimodal/image_utils.py b/relax/utils/multimodal/image_utils.py index e6f36a277..de29a9d09 100644 --- a/relax/utils/multimodal/image_utils.py +++ b/relax/utils/multimodal/image_utils.py @@ -10,7 +10,7 @@ import numpy as np import requests -from PIL import Image +from PIL import Image, ImageFile from relax.utils.env import Envs @@ -22,6 +22,9 @@ QWEN_VL_SAFE_ASPECT_RATIO = 199 +ImageFile.LOAD_TRUNCATED_IMAGES = True + + ImageInput = Union[ Image.Image, np.ndarray, diff --git a/tests/utils/multimodal/test_image_utils.py b/tests/utils/multimodal/test_image_utils.py index f941dd226..06d449475 100644 --- a/tests/utils/multimodal/test_image_utils.py +++ b/tests/utils/multimodal/test_image_utils.py @@ -164,6 +164,16 @@ def test_load_image_supports_local_path_and_file_uri(tmp_path): assert from_uri.getpixel((1, 1)) == (9, 8, 7) +def test_load_image_supports_truncated_jpeg(): + source = Image.new("RGB", (8, 8), (1, 2, 3)) + buffer = BytesIO() + source.save(buffer, format="JPEG") + + loaded = load_image(buffer.getvalue()[:-4]) + + assert loaded.size == (8, 8) + + def test_load_image_rejects_unsupported_type(): with pytest.raises(NotImplementedError, match="Unsupported image input type"): load_image(123) From 1d2f98c1e767b95d3a0c9525f3e9dfe6833e5975 Mon Sep 17 00:00:00 2001 From: pojun Date: Thu, 3 Sep 2026 14:55:12 +0800 Subject: [PATCH 17/34] feat(sft): make the loss-mask fallback delimiter-agnostic MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # ✨ Features `_render_per_message_fallback` hardcoded ChatML (`<|im_start|>{role}\n` … `<|im_end|>`). That is fine for Qwen, but any non-ChatML model dies on the very first sample with RuntimeError: could not locate 'user' message after cursor 0 in rendered chat template output because the scan finds no headers. Path 1 (`{% generation %}`) is not an escape hatch either -- Qwen3's own template has no such marker, so ChatML models rely on this fallback too. Extract the delimiters into a `_Dialect` record and select it from what the template actually rendered. Behaviour for ChatML is byte-for-byte unchanged. ## Added: gemma-4 Frames turns as `<|turn>model\n … ` (note: the assistant role renders as `model`) and delimits reasoning with `<|channel>thought … `. Unlike ChatML's fixed-length `\n` opener, the reasoning span has a closing marker, so the mask resumes after `` rather than after a fixed offset. Mirrors THUDM/slime's `gen_multi_turn_loss_mask_gemma4`. Tool messages are gated per dialect: gemma-4's tool-call framing differs and is not implemented, so it raises rather than silently mis-masking. # ✅ Verification Against gemma-4-31B-it's stock template: - single turn -- loss covers exactly `'The weather is nice today.\n'`; bos, `<|turn>user\n…`, and the `<|turn>model\n` header are all excluded - multi turn -- `'a1\na2\n'`, both replies and neither prompt - reasoning block detected and the mask resumes after `` Qwen3 regression: dialect still resolves to chatml, loss still covers `'…\n\nworld<|im_end|>\n'` with the user turn excluded. An end-to-end 40-step SFT run on gemma-4-31B produced an identical loss curve to the previous approach (4.3340 -> 1.1633 vs 1.1602, bf16 noise), confirming the rendered training sequence is unchanged. (cherry picked from commit c6ca8d385221170d37e914e32195678adc3c353c) (cherry picked from commit 199b956a4b45553f849fde27828c7011efa3932f) --- relax/backends/megatron/data.py | 9 +- relax/backends/megatron/model.py | 4 +- relax/engine/sft/dataset/chat_template.py | 98 +++++++- .../sft/dataset/gemma4_chat_template_patch.py | 85 +++++++ relax/models/__init__.py | 17 +- relax/models/gemma4/__init__.py | 8 + relax/models/gemma4/attention.py | 196 ++++++++++++++++ relax/models/gemma4/gemma4_bridge.py | 56 +++++ relax/models/gemma4/gemma4_provider.py | 168 +++++++++++++ scripts/models/gemma4-26B.sh | 59 +++++ scripts/models/gemma4-31B.sh | 35 +++ .../training/sft/run-gemma4-26B-sft-8xgpu.sh | 204 ++++++++++++++++ .../sft/run-gemma4-26B-sft-hf-8xgpu.sh | 222 ++++++++++++++++++ .../training/sft/run-gemma4-31B-sft-8xgpu.sh | 187 +++++++++++++++ .../megatron/test_save_hf_strictness.py | 61 +++++ .../weight_update/test_lora_weight_sync.py | 6 +- tests/models/gemma4/test_attention.py | 88 +++++++ 17 files changed, 1484 insertions(+), 19 deletions(-) create mode 100644 relax/engine/sft/dataset/gemma4_chat_template_patch.py create mode 100644 relax/models/gemma4/__init__.py create mode 100644 relax/models/gemma4/attention.py create mode 100644 relax/models/gemma4/gemma4_bridge.py create mode 100644 relax/models/gemma4/gemma4_provider.py create mode 100644 scripts/models/gemma4-26B.sh create mode 100644 scripts/models/gemma4-31B.sh create mode 100644 scripts/training/sft/run-gemma4-26B-sft-8xgpu.sh create mode 100755 scripts/training/sft/run-gemma4-26B-sft-hf-8xgpu.sh create mode 100644 scripts/training/sft/run-gemma4-31B-sft-8xgpu.sh create mode 100644 tests/models/gemma4/test_attention.py diff --git a/relax/backends/megatron/data.py b/relax/backends/megatron/data.py index 02dccdb70..39cc70a67 100644 --- a/relax/backends/megatron/data.py +++ b/relax/backends/megatron/data.py @@ -446,6 +446,7 @@ def get_batch( cu_seqlens_list.append(cu_seqlens_list[-1] + pad) cu_seqlens = torch.tensor(cu_seqlens_list, dtype=torch.int, device=batch_device) + cu_seqlens_cpu = cu_seqlens_list tokens = tokens.chunk(cp_size, dim=0)[cp_rank] else: tokens = [ @@ -466,9 +467,10 @@ def get_batch( cu_seqlens.append(cu_seqlens[-1] + pad) # thd requires the cu_seqlens to be of the origin length - cu_seqlens = torch.tensor(cu_seqlens, dtype=torch.int, device=batch_device) * cp_size + cu_seqlens_cpu = [offset * cp_size for offset in cu_seqlens] + cu_seqlens = torch.tensor(cu_seqlens_cpu, dtype=torch.int, device=batch_device) - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + max_seqlen = max(end - start for start, end in zip(cu_seqlens_cpu, cu_seqlens_cpu[1:])) packed_seq_params = PackedSeqParams( cu_seqlens_q=cu_seqlens, cu_seqlens_kv=cu_seqlens, @@ -476,6 +478,9 @@ def get_batch( max_seqlen_kv=max_seqlen, qkv_format="thd", ) + # Python boundaries let attention implementations iterate packed + # subsequences without synchronizing individual accelerator scalars. + packed_seq_params.cu_seqlens_q_cpu = cu_seqlens_cpu if use_dynamic_context_parallel: packed_seq_params.local_cp_size = cp_size packed_seq_params.cp_group = cp_group diff --git a/relax/backends/megatron/model.py b/relax/backends/megatron/model.py index 3a00fd1db..8e7bbd00d 100644 --- a/relax/backends/megatron/model.py +++ b/relax/backends/megatron/model.py @@ -1507,8 +1507,8 @@ def train( if mtp_values is None: mtp_values = tracker.get("loss_values") if mtp_values is not None: - # here we assume only one mtp layer - mtp_losses = (mtp_values * mtp_loss_scale).item() + # Sum across MTP prediction depths. + mtp_losses = (mtp_values * mtp_loss_scale).sum().item() MTPLossLoggingHelper.clean_loss_in_tracker() # CI check: verify MTP loss is within expected bounds diff --git a/relax/engine/sft/dataset/chat_template.py b/relax/engine/sft/dataset/chat_template.py index f1f027e03..5ca62e448 100644 --- a/relax/engine/sft/dataset/chat_template.py +++ b/relax/engine/sft/dataset/chat_template.py @@ -11,14 +11,17 @@ """ import hashlib +import os import re import threading from collections.abc import Mapping +from dataclasses import dataclass, replace from typing import Any import torch from relax.engine.sft.dataset.chat_template_patch import TemplatePatchResult, apply_chat_template_patchers +from relax.engine.sft.dataset.gemma4_chat_template_patch import try_patch_gemma4_thinking from relax.engine.sft.dataset.qwen_chat_template_patch import try_patch_qwen_chat_template from relax.engine.sft.dataset.sample import CanonicalSample from relax.utils.logging_utils import get_logger @@ -31,7 +34,7 @@ # form for the purpose of marking assistant-token spans, so they should be # recognised as the same marker. _GENERATION_MARKER_RE = re.compile(r"{%-?\s*generation\s*-?%}") -_CHAT_TEMPLATE_PATCHERS = (try_patch_qwen_chat_template,) +_CHAT_TEMPLATE_PATCHERS = (try_patch_qwen_chat_template, try_patch_gemma4_thinking) _FALLBACK_WARNED: set[int] = set() # tokenizer id → warned once _TEMPLATE_LOGGED: set[tuple[int, int, str]] = set() # tokenizer id + template hash + preserve mode @@ -110,6 +113,63 @@ def _thread_local_chat_template_kwargs(tokenizer, apply_chat_template_kwargs: di _TOOL_RESPONSE_CLOSE = "\n" +@dataclass(frozen=True) +class _Dialect: + """Turn delimiters for the fallback's text scan. + + ChatML and gemma-4 frame turns differently; same scan, other delimiters. + """ + + name: str + role_names: dict # canonical role -> the name the template renders + header_fmt: str # "{role}" placeholder + end: str + think_open: str | None # None: nothing to exclude — the whole reply is learned + think_close: str | None # None: skip only the opener (ChatML \n) + supports_tools: bool + + def header(self, role: str) -> str: + return self.header_fmt.format(role=self.role_names.get(role, role)) + + +_CHATML = _Dialect( + name="chatml", + role_names={}, + header_fmt="<|im_start|>{role}\n", + end=_IM_END, + think_open=_THINK_OPEN, + think_close=None, + supports_tools=True, +) + +# gemma-4 renders the assistant role as "model". Reasoning is a delimited block, +# so the mask must resume after rather than after a fixed-length +# opener. Matches THUDM/slime's gen_multi_turn_loss_mask_gemma4. +_GEMMA4 = _Dialect( + name="gemma4", + role_names={"assistant": "model"}, + header_fmt="<|turn>{role}\n", + end="", + think_open="<|channel>thought\n", + think_close="", + supports_tools=False, +) + +# Same delimiters, but the reasoning block stays IN the loss. Only reachable via +# GEMMA4_SFT_THINKING=1, which patches the Jinja to emit an empty thought block +# on every assistant turn -- see gemma4_chat_template_patch.py. +_GEMMA4_THINKING = replace(_GEMMA4, name="gemma4_thinking", think_open=None, think_close=None) + + +def _detect_dialect(rendered_text: str) -> _Dialect: + """Pick delimiters from what the template actually emitted.""" + if "<|turn>" in rendered_text and "" in rendered_text: + if os.environ.get("GEMMA4_SFT_THINKING", "0") in ("1", "true", "True"): + return _GEMMA4_THINKING + return _GEMMA4 + return _CHATML + + def _render_per_message_fallback( sample: CanonicalSample, *, @@ -170,9 +230,15 @@ def _render_per_message_fallback( ) char_mask = bytearray(len(rendered_text)) # zeros + dialect = _detect_dialect(rendered_text) cursor = 0 for msg in sample.messages: if msg.role == "tool": + if not dialect.supports_tools: + raise RuntimeError( + f"tool messages are not supported by the {dialect.name!r} loss-mask dialect; " + f"add its tool-call delimiters to _Dialect first" + ) open_pos = rendered_text.find(_TOOL_RESPONSE_OPEN, cursor) if open_pos < 0: raise RuntimeError( @@ -186,17 +252,20 @@ def _render_per_message_fallback( span_end = close_pos cursor = close_pos + len(_TOOL_RESPONSE_CLOSE) else: - header = f"<|im_start|>{msg.role}\n" + header = dialect.header(msg.role) header_pos = rendered_text.find(header, cursor) if header_pos < 0: raise RuntimeError( - f"could not locate {msg.role!r} message after cursor {cursor} in rendered chat template output" + f"could not locate {msg.role!r} message after cursor {cursor} in rendered chat " + f"template output (dialect={dialect.name!r}, header={header!r})" ) content_start = header_pos + len(header) - end_pos = rendered_text.find(_IM_END, content_start) + end_pos = rendered_text.find(dialect.end, content_start) if end_pos < 0: - raise RuntimeError(f"could not locate <|im_end|> for {msg.role!r} message") - span_end = end_pos + len(_IM_END) + raise RuntimeError( + f"could not locate {dialect.end!r} for {msg.role!r} message (dialect={dialect.name!r})" + ) + span_end = end_pos + len(dialect.end) if span_end < len(rendered_text) and rendered_text[span_end] == "\n": span_end += 1 cursor = span_end @@ -205,8 +274,21 @@ def _render_per_message_fallback( continue mask_start = content_start - if msg.role == "assistant" and rendered_text[content_start : content_start + len(_THINK_OPEN)] == _THINK_OPEN: - mask_start += len(_THINK_OPEN) + if ( + msg.role == "assistant" + and dialect.think_open is not None + and rendered_text.startswith(dialect.think_open, content_start) + ): + if dialect.think_close is None: + # ChatML: only the opener is excluded; the reasoning body is learned. + mask_start += len(dialect.think_open) + else: + # Delimited block (gemma-4): the whole reasoning span stays out of + # the loss, so training targets only the visible reply. + close_pos = rendered_text.find(dialect.think_close, content_start) + if close_pos < 0: + raise RuntimeError(f"found {dialect.think_open!r} without a matching {dialect.think_close!r}") + mask_start = close_pos + len(dialect.think_close) for pos in range(mask_start, span_end): char_mask[pos] = 1 diff --git a/relax/engine/sft/dataset/gemma4_chat_template_patch.py b/relax/engine/sft/dataset/gemma4_chat_template_patch.py new file mode 100644 index 000000000..6972e2d86 --- /dev/null +++ b/relax/engine/sft/dataset/gemma4_chat_template_patch.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""gemma-4 thinking-structure chat-template patch for SFT.""" + +import os +from collections.abc import Mapping +from functools import lru_cache +from typing import Any + +from relax.engine.sft.dataset.chat_template_patch import TemplatePatchResult +from relax.engine.sft.dataset.sample import CanonicalSample + + +_PATCH_NAME = "gemma4_thinking" +_ENV_FLAG = "GEMMA4_SFT_THINKING" + +# The empty thought block ms-swift's `gemma4` template counts in the loss. +# Unreachable in the official Jinja, which is why this patch exists. +GEMMA4_EMPTY_THOUGHT = "<|channel>thought\n" + +_ANCHOR = "{{- strip_thinking(message['content']) -}}" +_REPLACEMENT = ( + "{%- if enable_thinking and not thinking_text -%}" + "{{- '<|channel>thought\\n' -}}" + "{%- endif -%}" + "{{- strip_thinking(message['content']) -}}" +) + + +def _is_gemma4_template(template: str) -> bool: + """gemma-4 by its delimiters, not by model name (the caller may not know + it).""" + return "<|turn>" in template and "" in template and "strip_thinking" in template + + +@lru_cache(maxsize=32) +def _patch_template(template: str) -> tuple[str, bool] | None: + """Insert the empty-thought emission. + + None = not gemma-4, don't touch. + """ + if not _is_gemma4_template(template): + return None + # Fail loud rather than fuzzy-match. + count = template.count(_ANCHOR) + if count == 0: + raise RuntimeError( + f"{_PATCH_NAME}: anchor not found in the gemma-4 chat template. " + "Upstream changed it; re-derive the patch instead of letting it " + "silently no-op." + ) + if count > 1: + raise RuntimeError(f"{_PATCH_NAME}: anchor appears {count}x; refusing to guess which one to patch.") + return template.replace(_ANCHOR, _REPLACEMENT, 1), True + + +def try_patch_gemma4_thinking( + sample: CanonicalSample, + template: str | None, + kwargs: Mapping[str, Any], +) -> TemplatePatchResult | None: + """Emit ms-swift's `gemma4` thinking scaffold. Off unless + GEMMA4_SFT_THINKING=1. + + The scaffold adds a `<|think|>` system turn and an empty thought block that + counts toward the loss -- 4 identical tokens on every sample. Only worth + turning on to match a baseline that has them. + """ + if os.environ.get(_ENV_FLAG, "0") not in ("1", "true", "True"): + return None + if not template: + return None + patched = _patch_template(template) + if patched is None: + return None + new_template, _ = patched + new_kwargs = dict(kwargs) + # Drives the `<|turn>system\n<|think|>\n\n` block (template line ~189). + new_kwargs["enable_thinking"] = True + return TemplatePatchResult( + template=new_template, + kwargs=new_kwargs, + patch_name=_PATCH_NAME, + changed=True, + ) diff --git a/relax/models/__init__.py b/relax/models/__init__.py index baa760c5f..48b16fd0f 100644 --- a/relax/models/__init__.py +++ b/relax/models/__init__.py @@ -11,11 +11,8 @@ from relax.models.qwen_omni.qwen3_omni_bridge import Qwen3OmniMoEBridge # noqa: F811 from relax.models.qwen_omni.qwen3_omni_provider import Qwen3OmniModelProvider # noqa: F811 -# Import glm_moe_dsa in its own try/except so a failure above does not block -# the GLM5Bridge @register_bridge decorator from running. Without this, an -# unrelated qwen_omni circular-import error prevents GLM5Bridge from being -# registered, and AutoBridge silently falls back to the generic MLA bridge, -# bypassing the fused DSAMLASelfAttention spec. +# Own try/except so a failure above still lets @register_bridge run -- otherwise +# AutoBridge silently falls back to the generic MLA bridge. try: from relax.models import glm_moe_dsa # noqa: F401 except Exception as _e: @@ -33,6 +30,16 @@ _logging.getLogger(__name__).warning("Failed to import relax.models.dots_ocr.megatron: %s", _e) +# Register the gemma-4 bridge, overriding upstream's entry so gemma-4 builds +# with Relax's packed-safe core attention. +try: + from relax.models import gemma4 # noqa: F401 +except Exception as _e: + import logging as _logging + + _logging.getLogger(__name__).warning("Failed to import relax.models.gemma4: %s", _e) + + __all__ = [ "Qwen3OmniMoEBridge", "Qwen3OmniMoeModel", diff --git a/relax/models/gemma4/__init__.py b/relax/models/gemma4/__init__.py new file mode 100644 index 000000000..fb9250a32 --- /dev/null +++ b/relax/models/gemma4/__init__.py @@ -0,0 +1,8 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +from relax.models.gemma4.gemma4_bridge import Gemma4DenseBridge + + +__all__ = [ + "Gemma4DenseBridge", +] diff --git a/relax/models/gemma4/attention.py b/relax/models/gemma4/attention.py new file mode 100644 index 000000000..231290fd6 --- /dev/null +++ b/relax/models/gemma4/attention.py @@ -0,0 +1,196 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Gemma-4 core attention that survives the packed (THD) backward. + +TE's cuDNN FusedAttention emits NaN gradients in the THD backward on gemma-4's +sliding layers. The forward is bit-identical to the unpacked run, so the defect +is that kernel's backward alone; the originating layer varies between otherwise +identical runs. Upstream report (context-parallel variant of the same failure): +https://github.com/NVIDIA/TransformerEngine/issues/2186 + +Ported from THUDM/slime ``slime_plugins/models/gemma4.py::SDPACoreAttention``. + +Dispatch, CP == 1 only: + + thd + cu_seqlens + head_dim <= 256 -> flash_attn_varlen_func, window (sw-1, 0) + thd + cu_seqlens + head_dim > 256 -> per-sub-sequence causal SDPA + thd without cu_seqlens -> raises + sbhd -> raises +""" + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + + +def _is_sliding_layer(config, layer_number: int) -> bool: + """Is this layer a sliding-window layer? + + Dense providers encode the pattern in ``window_attn_skip_freq``, MoE ones + in ``interleaved_attn_pattern``. + """ + window_size = getattr(config, "window_size", None) + if not window_size: + return False + skip_freq = getattr(config, "window_attn_skip_freq", None) + if isinstance(skip_freq, list): + layer_type = skip_freq[layer_number - 1] + if isinstance(layer_type, str): + return layer_type == "sliding_attention" + return bool(layer_type) + if skip_freq is None: + # MoE: pattern lives in interleaved_attn_pattern. Returning True here + # unconditionally would put a window on the global layers. + pattern = getattr(config, "interleaved_attn_pattern", None) + if pattern: + return layer_number % sum(pattern) != 0 + return True + if isinstance(skip_freq, int): + return layer_number % skip_freq != 0 + return False + + +def _window_pair(window_size) -> tuple: + """Normalise gemma-4's two window conventions to flash-attn's ``(left, + right)``. + + Dense stores the inclusive pair, MoE the raw span. Do not subtract from the + dense tuple -- the -1 is already baked in and doing it twice shifts the + window by a token. + """ + if not window_size: + return (-1, -1) + if isinstance(window_size, int): + return (window_size - 1, 0) + return tuple(window_size) + + +class Gemma4CoreAttention(nn.Module): + """Drop-in replacement for TEDotProductAttention on gemma-4 dense + layers.""" + + def __init__( + self, + config, + layer_number: int, + attn_mask_type=None, + attention_type: str = "self", + attention_dropout: Optional[float] = None, + softmax_scale: Optional[float] = None, + **kwargs, + ): + super().__init__() + # Megatron/TE hand core_attention a moving set of kwargs; accept and drop. + del kwargs + self.config = config + self.layer_number = layer_number + self.attention_type = attention_type + # gemma-4 sets softmax_scale to 1.0, not head_dim**-0.5, because the query + # is pre-scaled upstream. Do not "fix" it -- doing so made per-layer + # divergence go 0.25 -> 14.6. + self.softmax_scale = softmax_scale if softmax_scale is not None else getattr(config, "softmax_scale", None) + self.dropout_p = config.attention_dropout if attention_dropout is None else attention_dropout + self.is_sliding = _is_sliding_layer(config, layer_number) + self.window_size = _window_pair(getattr(config, "window_size", None)) if self.is_sliding else (-1, -1) + + def _scale(self, head_dim: int) -> float: + return self.softmax_scale if self.softmax_scale is not None else head_dim**-0.5 + + # ---------------- THD ---------------- + + def _thd_flash(self, query, key, value, cu_seqlens, max_seqlen: int): + """head_dim <= 256: flash-attn's variable-length kernel, replacing + cuDNN FusedAttention.""" + from flash_attn import flash_attn_varlen_func + + cu = cu_seqlens.to(torch.int32) + out = flash_attn_varlen_func( + query.contiguous(), + key.contiguous(), + value.contiguous(), + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=max_seqlen, + max_seqlen_k=max_seqlen, + dropout_p=self.dropout_p if self.training else 0.0, + softmax_scale=self._scale(query.shape[2]), + causal=True, + window_size=self.window_size, + ) + return out.reshape(query.shape[0], -1) + + def _thd_sdpa_per_subseq(self, query, key, value, boundaries): + """head_dim > 256 (gemma-4's global layers), which flash-attn cannot + take. + + Loops sub-sequences rather than materialising a [T, T] block-diagonal + mask. + """ + n_q, head_dim = query.shape[1], query.shape[2] + n_kv = key.shape[1] + scale = self._scale(head_dim) + out = torch.empty(query.shape[0], n_q * head_dim, dtype=query.dtype, device=query.device) + for s, e in zip(boundaries, boundaries[1:]): + o = F.scaled_dot_product_attention( + query[s:e].unsqueeze(0).transpose(1, 2), + key[s:e].unsqueeze(0).transpose(1, 2), + value[s:e].unsqueeze(0).transpose(1, 2), + dropout_p=self.dropout_p if self.training else 0.0, + scale=scale, + is_causal=True, + enable_gqa=(n_q != n_kv), + ) + out[s:e] = o.transpose(1, 2).reshape(e - s, -1) + return out + + # ---------------- entry ---------------- + + def forward( + self, + query, + key, + value, + attention_mask=None, + attn_mask_type=None, + packed_seq_params=None, + **kwargs, + ): + cp_size = getattr(self.config, "context_parallel_size", 1) or 1 + if cp_size > 1: + raise NotImplementedError( + "Gemma4CoreAttention does not implement context parallelism. slime's CP path " + "relies on its own zig-zag CP layout, which is not Relax's convention; porting " + "it unverified would silently compute the wrong attention. Run with " + "--context-parallel-size 1, or implement CP against Relax's cp_utils first." + ) + + if query.dim() == 3: # thd: [T, np, hn] + cu_seqlens = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None + if cu_seqlens is not None: + if query.shape[2] <= 256: + max_seqlen = packed_seq_params.max_seqlen_q + if max_seqlen is None: + raise ValueError("packed_seq_params.max_seqlen_q is required for Gemma4 flash attention") + return self._thd_flash(query, key, value, cu_seqlens, max_seqlen) + boundaries = getattr(packed_seq_params, "cu_seqlens_q_cpu", None) + if boundaries is None: + # PackedSeqParams built outside Relax's data path. Copy once. + boundaries = cu_seqlens.detach().cpu().tolist() + return self._thd_sdpa_per_subseq(query, key, value, boundaries) + raise NotImplementedError( + "Gemma4CoreAttention requires packed THD inputs with cu_seqlens. " + "Treating THD input without sequence boundaries as one sequence would " + "silently lose Gemma-4's sliding-window and padding-mask semantics." + ) + + raise NotImplementedError( + "Gemma4CoreAttention currently supports packed THD attention only. " + "The SBHD path needs an explicit causal sliding-window mask before it can be enabled safely." + ) diff --git a/relax/models/gemma4/gemma4_bridge.py b/relax/models/gemma4/gemma4_bridge.py new file mode 100644 index 000000000..8a71daa37 --- /dev/null +++ b/relax/models/gemma4/gemma4_bridge.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Relax's gemma-4 bridge: upstream's, with two variant fixes. + +Subclasses upstream's bridge and only re-points the provider at its Relax +subclass in :data:`~relax.models.gemma4.gemma4_provider.RELAX_PROVIDERS`: +dense needs a packed-safe attention, MoE needs its ``rotary_base`` tuple kept +out of reach of Relax's ``bridge_keys`` override. +""" + +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.gemma_vl.gemma4_vl_bridge import Gemma4VLBridge +from megatron.bridge.models.gemma_vl.gemma4_vl_provider import Gemma4VLModelProvider +from megatron.bridge.models.gemma_vl.modeling_gemma4_vl import Gemma4VLModel +from megatron.bridge.models.hf_pretrained.causal_lm import PreTrainedCausalLM + +from relax.models.gemma4.gemma4_provider import RELAX_PROVIDERS, stash_dual_rope +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + + +@MegatronModelBridge.register_bridge( + source="Gemma4ForConditionalGeneration", + target=Gemma4VLModel, + provider=Gemma4VLModelProvider, + model_type="gemma4_vl", +) +class Gemma4DenseBridge(Gemma4VLBridge): + """Bridge for gemma-4 full-parameter SFT -- dense and MoE alike, despite + the ``Dense`` in the name. + + Example: + >>> from megatron.bridge import AutoBridge + >>> bridge = AutoBridge.from_hf_pretrained("google/gemma-4-31B-it") + >>> provider = bridge.to_megatron_provider() + + ``GEMMA4_CONVERSION_MODE=text`` loads the VL checkpoint as text-only. + """ + + def provider_bridge(self, hf_pretrained: PreTrainedCausalLM): + """Return upstream's provider, re-pointed at its Relax subclass.""" + provider = super().provider_bridge(hf_pretrained) + + # Exact type, not isinstance -- each variant has its own subclass. + replacement = RELAX_PROVIDERS.get(type(provider)) + if replacement is None: + logger.debug("%s is not a known gemma-4 provider; leaving it alone", type(provider).__name__) + return provider + + # Must run before Relax's bridge_keys override loop, which fires as soon + # as this returns. No-op on dense. + stash_dual_rope(provider) + provider.__class__ = replacement + return provider diff --git a/relax/models/gemma4/gemma4_provider.py b/relax/models/gemma4/gemma4_provider.py new file mode 100644 index 000000000..f3f43fcd1 --- /dev/null +++ b/relax/models/gemma4/gemma4_provider.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Gemma-4 providers adapted for Relax's training path. + +Two unrelated upstream assumptions break under Relax, one per variant: + +* **dense** -- Megatron-Bridge wires it to TE's attention, which emits NaN + gradients in the packed (THD) backward. See :mod:`relax.models.gemma4.attention`. +* **MoE** -- its ``rotary_base`` is a tuple, and Relax's ``bridge_keys`` override + flattens it to a scalar. See :class:`_Gemma4MoEProvideMixin`. + +Both are delivered by :class:`~relax.models.gemma4.gemma4_bridge.Gemma4DenseBridge` +re-pointing the provider instance at its counterpart in :data:`RELAX_PROVIDERS`. + +Keep the subclasses declared statically -- Relax pickles the provider to reach +the Ray train actors, and a class built by ``type(name, (mixin, base), {})`` is +not importable, so pickling it raises ``PicklingError``. +""" + +from contextlib import contextmanager +from functools import partial + +from megatron.bridge.models.gemma.gemma4_provider import Gemma4DenseProvider, Gemma4ModelProvider +from megatron.bridge.models.gemma_vl.gemma4_vl_provider import Gemma4DenseVLProvider, Gemma4VLModelProvider + +from relax.utils.logging_utils import get_logger + + +logger = get_logger(__name__) + + +@contextmanager +def _relax_core_attention(): + """Rebind ``get_gemma4_layer_spec`` for the duration of one ``build()``. + + ``Gemma4DenseProvider.build()`` ignores its ``transformer_layer_spec`` + field and calls the module-level ``get_gemma4_layer_spec(config)`` + directly, so assigning that field does nothing here. Restored in + ``finally``. + """ + from megatron.bridge.models.gemma import gemma4_provider + + from relax.models.gemma4.attention import Gemma4CoreAttention + + original = gemma4_provider.get_gemma4_layer_spec + + def get_gemma4_layer_spec(config=None): + spec = original(config) + # AttributeError here is intentional if upstream restructures the spec. + spec.submodules.self_attention.submodules.core_attention = Gemma4CoreAttention + return spec + + gemma4_provider.get_gemma4_layer_spec = get_gemma4_layer_spec + try: + yield + finally: + gemma4_provider.get_gemma4_layer_spec = original + + +class _PackedSafeBuildMixin: + """Builds the model with Relax's core attention in place of TE's.""" + + def build(self, *args, **kwargs): + with _relax_core_attention(): + model = super().build(*args, **kwargs) + logger.info("gemma-4 dense: core_attention -> Gemma4CoreAttention (packed-safe)") + return model + + +class PackedSafeGemma4DenseProvider(_PackedSafeBuildMixin, Gemma4DenseProvider): + """``Gemma4DenseProvider`` for the text-only path + (GEMMA4_CONVERSION_MODE=text).""" + + +class PackedSafeGemma4DenseVLProvider(_PackedSafeBuildMixin, Gemma4DenseVLProvider): + """``Gemma4DenseVLProvider`` for the VL path. + + Untested -- Relax's gemma-4 + work is text-only SFT so far. + """ + + +#: Where the bridge parks the ``(local, global)`` rope-theta pair. Deliberately +#: not a ``bridge_keys`` name, so Relax's override loop leaves it alone. +_ROPE_PAIR_ATTR = "_relax_gemma4_rotary_base_pair" + +#: Guards against double-wrapping ``transformer_layer_spec`` if ``provide()`` runs +#: more than once on the same provider (e.g. virtual pipeline stages). +_PACKED_SAFE_ATTR = "_relax_gemma4_packed_safe_installed" + + +def stash_dual_rope(provider) -> None: + """Park a MoE provider's ``rotary_base`` tuple out of ``bridge_keys``' + reach. + + Must run before Relax's arg-override loop; read back in ``provide()``. + """ + if isinstance(provider.rotary_base, tuple): + setattr(provider, _ROPE_PAIR_ATTR, provider.rotary_base) + + +def _packed_safe_moe_block_spec(inner_spec_fn, config): + """Build upstream's MoE block spec, then swap the attention on every layer. + + Module level and bound through ``functools.partial`` rather than a closure, + for the pickling reason in the module docstring. Upstream's per-layer + ``Gemma4TEDotProductAttention`` lands on the cuDNN kernel that NaNs in the + THD backward -- see :mod:`relax.models.gemma4.attention`. + + The signature must not grow a ``vp_stage`` parameter: ``provide()`` inspects + for one and would pass it to a builder that does not take it. + """ + from relax.models.gemma4.attention import Gemma4CoreAttention + + block_spec = inner_spec_fn(config) + # AttributeError here is intentional if upstream restructures the spec. + for layer_spec in block_spec.layer_specs: + layer_spec.submodules.self_attention.submodules.core_attention = Gemma4CoreAttention + return block_spec + + +class _Gemma4MoEProvideMixin: + """Fix-ups that must run before upstream's MoE ``provide()``. + + 1. Restore the ``(local, global)`` rope thetas that ``bridge_keys`` flattened + to a scalar -- ``provide()`` unpacks the tuple on its first line, so a + scalar raises ``TypeError``. No launch-script flag avoids it: one int + cannot express two thetas. + 2. Install the packed-safe attention. See :func:`_packed_safe_moe_block_spec`. + """ + + def provide(self, *args, **kwargs): + pair = getattr(self, _ROPE_PAIR_ATTR, None) + if pair is not None and not isinstance(self.rotary_base, tuple): + logger.info( + "gemma-4 MoE: restoring rotary_base %r -> %r (bridge_keys had flattened it)", self.rotary_base, pair + ) + self.rotary_base = pair + + if not getattr(self, _PACKED_SAFE_ATTR, False): + self.transformer_layer_spec = partial(_packed_safe_moe_block_spec, self.transformer_layer_spec) + setattr(self, _PACKED_SAFE_ATTR, True) + logger.info("gemma-4 MoE: core_attention -> Gemma4CoreAttention (packed-safe)") + + return super().provide(*args, **kwargs) + + +class RelaxGemma4MoEProvider(_Gemma4MoEProvideMixin, Gemma4ModelProvider): + """``Gemma4ModelProvider`` for the text-only MoE path.""" + + +class RelaxGemma4MoEVLProvider(_Gemma4MoEProvideMixin, Gemma4VLModelProvider): + """``Gemma4VLModelProvider`` for the MoE VL path. + + The VL model reaches the language model through + ``provide_language_model()``, which delegates to the ``provide()`` hooked + here. + """ + + +# No subclass adds fields, so an instance can be re-pointed in place and keep +# the attributes the bridge set outside the dataclass. +RELAX_PROVIDERS = { + Gemma4DenseProvider: PackedSafeGemma4DenseProvider, + Gemma4DenseVLProvider: PackedSafeGemma4DenseVLProvider, + Gemma4ModelProvider: RelaxGemma4MoEProvider, + Gemma4VLModelProvider: RelaxGemma4MoEVLProvider, +} diff --git a/scripts/models/gemma4-26B.sh b/scripts/models/gemma4-26B.sh new file mode 100644 index 000000000..8271eda46 --- /dev/null +++ b/scripts/models/gemma4-26B.sh @@ -0,0 +1,59 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +# google/gemma-4-26B-A4B-it -- MoE text path. +# +# Values come from gemma-4-26B-A4B-it/config.json `text_config`, not the +# top-level multimodal config. Needs GEMMA4_CONVERSION_MODE=text (the launch +# script sets it in the Ray runtime env) and megatron-dev >= a8d06a1. + +NLAYERS=30 +NHIDDEN=2816 +NHEADS=16 +NUM_QUERY_GROUPS=8 # num_key_value_heads +HEAD_DIM=256 # head_dim (sliding layers; global layers use 512) +VOCAB=262144 + +MOE_ROUTED_EXPERTS=128 # num_experts +MOE_ACTIVE_ROUTED_EXPERTS=8 # top_k_experts +MOE_FFN_HIDDEN=704 # moe_intermediate_size +SHARED_EXPERT_FFN_HIDDEN=2112 # intermediate_size (the shared expert, not a dense FFN) + +# Most --moe-* flags below are in bridge_keys +# (relax/backends/megatron/model_provider.py): omitting one overwrites the +# checkpoint value with megatron's default instead of keeping it. +MODEL_ARGS=( + --num-layers ${NLAYERS} + --hidden-size ${NHIDDEN} + --ffn-hidden-size ${SHARED_EXPERT_FFN_HIDDEN} + --num-attention-heads ${NHEADS} + --group-query-attention + --num-query-groups ${NUM_QUERY_GROUPS} + --kv-channels ${HEAD_DIM} + --vocab-size ${VOCAB} + + --normalization RMSNorm + --norm-epsilon 1e-6 + --position-embedding-type rope + --disable-bias-linear + --qk-layernorm + + --num-experts ${MOE_ROUTED_EXPERTS} + --moe-router-topk ${MOE_ACTIVE_ROUTED_EXPERTS} + --moe-ffn-hidden-size ${MOE_FFN_HIDDEN} + --moe-shared-expert-intermediate-size ${SHARED_EXPERT_FFN_HIDDEN} + --moe-layer-freq 1 + --moe-grouped-gemm + --moe-permute-fusion + --moe-router-pre-softmax + --moe-router-dtype fp32 + --moe-token-dispatcher-type alltoall + + --moe-aux-loss-coeff 0 + --moe-router-load-balancing-type none + + --no-rope-fusion +) + +# Do not add --rotary-base, --swiglu, --untie-embeddings-and-output-weights or +# --make-vocab-size-divisible-by: the provider owns them and some fail +# validation. Expert-parallel sizes belong to the launch script's PERF_ARGS. diff --git a/scripts/models/gemma4-31B.sh b/scripts/models/gemma4-31B.sh new file mode 100644 index 000000000..781633906 --- /dev/null +++ b/scripts/models/gemma4-31B.sh @@ -0,0 +1,35 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +# google/gemma-4-31B-it -- DENSE text path. +# +# Needs GEMMA4_CONVERSION_MODE=text (the launch script sets it in the Ray +# runtime env) so Gemma4VLBridge drops the vision tower. +# +# Most args here only feed Relax's _hf_validate_args; the real config comes from +# Gemma4DenseProvider. The exception is `bridge_keys` +# (relax/backends/megatron/model_provider.py), which DOES overwrite the +# provider -- num-layers and rotary-base are in that list, so a wrong value +# silently builds the wrong model instead of erroring. + +MODEL_ARGS=( + --num-layers 60 + --hidden-size 5376 + --ffn-hidden-size 21504 + --num-attention-heads 32 + --group-query-attention + --num-query-groups 16 + --kv-channels 256 + --vocab-size 262144 + + --normalization RMSNorm + --norm-epsilon 1e-6 + --position-embedding-type rope + --disable-bias-linear + --qk-layernorm + + --no-rope-fusion +) + +# Do not add --swiglu, --untie-embeddings-and-output-weights or --rotary-base: +# the provider owns them, and --rotary-base 1000000 would silently corrupt every +# sliding layer rather than fail. diff --git a/scripts/training/sft/run-gemma4-26B-sft-8xgpu.sh b/scripts/training/sft/run-gemma4-26B-sft-8xgpu.sh new file mode 100644 index 000000000..ec63baf65 --- /dev/null +++ b/scripts/training/sft/run-gemma4-26B-sft-8xgpu.sh @@ -0,0 +1,204 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# gemma-4-26B-A4B-it (MoE) full-parameter SFT with packing, 8xGPU, TP=8 EP=8 PP=1, +# ray-submit launch. +# +# Usage: +# MEGATRON= MODEL_DIR= \ +# PROMPT_DATA= bash scripts/training/sft/run-gemma4-26B-sft-8xgpu.sh + +set -ex +set -o pipefail + +unset NCCL_NVLS_ENABLE + +now=$(date "+%Y-%m-%d-%H:%M:%S") + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${RELAX:-${SCRIPT_DIR}/../../..}" &>/dev/null && pwd)" +export RELAX="${RELAX_ROOT}" + +export MEGATRON="${MEGATRON:-/root/Megatron-LM/}" + +if ! PYTHONPATH="${RELAX_ROOT}:${MEGATRON}:${PYTHONPATH:-}" python3 -c \ + 'import inspect; from megatron.bridge.models.gemma.gemma4_provider import Gemma4DenseProvider, Gemma4ModelProvider; from megatron.bridge.models.gemma_vl.gemma4_vl_bridge import Gemma4VLBridge; from relax.models.gemma4.gemma4_bridge import Gemma4DenseBridge; from relax.models.gemma4.gemma4_provider import RELAX_PROVIDERS, PackedSafeGemma4DenseProvider, RelaxGemma4MoEProvider; source = inspect.getsource(Gemma4VLBridge.provider_bridge); ok = source.count("_conversion_mode()") >= 2 and RELAX_PROVIDERS.get(Gemma4DenseProvider) is PackedSafeGemma4DenseProvider and RELAX_PROVIDERS.get(Gemma4ModelProvider) is RelaxGemma4MoEProvider +if not ok: + raise RuntimeError("Gemma4 packed-safe provider mapping is unavailable")' \ + >/dev/null 2>&1; then + echo "ERROR: Gemma4 packed-safe integration failed its startup probe." >&2 + echo " ${MEGATRON} must provide the required providers and MoE text mode," >&2 + echo " and Relax's Gemma4 bridge/provider replacements must import cleanly." >&2 + exit 1 +fi + +MODEL_CONFIG_DIR="${MODEL_CONFIG_DIR:-${SCRIPT_DIR}/../../models}" +source "${MODEL_CONFIG_DIR}/gemma4-26B.sh" + +TP_SIZE="${TP_SIZE:-8}" +EP_SIZE="${EP_SIZE:-8}" +ACTOR_GPUS="${ACTOR_GPUS:-8}" + +if [ "$((ACTOR_GPUS % TP_SIZE))" -ne 0 ]; then + echo "ERROR: ACTOR_GPUS=${ACTOR_GPUS} is not divisible by TP_SIZE=${TP_SIZE}." >&2 + exit 1 +fi +if [ "$((NHEADS % TP_SIZE))" -ne 0 ] || [ "$((NUM_QUERY_GROUPS % TP_SIZE))" -ne 0 ]; then + echo "ERROR: TP_SIZE=${TP_SIZE} must divide NHEADS=${NHEADS} and NUM_QUERY_GROUPS=${NUM_QUERY_GROUPS}." >&2 + exit 1 +fi +if [ "$((MOE_ROUTED_EXPERTS % EP_SIZE))" -ne 0 ]; then + echo "ERROR: EP_SIZE=${EP_SIZE} must divide MOE_ROUTED_EXPERTS=${MOE_ROUTED_EXPERTS}." >&2 + exit 1 +fi +if [ "${TP_SIZE}" -lt 1 ] || [ "${TP_SIZE}" -gt 16 ] || [ "$((TP_SIZE & (TP_SIZE - 1)))" -ne 0 ]; then + echo "ERROR: TP_SIZE=${TP_SIZE} must be a power of two between 1 and 16." >&2 + exit 1 +fi +if [ "${EP_SIZE}" -ne "${ACTOR_GPUS}" ]; then + echo "ERROR: EP_SIZE must equal ACTOR_GPUS (= TP * CP * DP with PP=CP=1)." >&2 + echo " Got EP_SIZE=${EP_SIZE}, ACTOR_GPUS=${ACTOR_GPUS}, TP_SIZE=${TP_SIZE}." >&2 + echo " EP smaller than that replicates experts on every rank and will OOM;" >&2 + echo " larger is rejected by megatron." >&2 + exit 1 +fi +echo "parallelism: TP=${TP_SIZE} EP=${EP_SIZE} PP=1 CP=1 on ${ACTOR_GPUS} GPU(s); num_layers=${NUM_LAYERS:-30}" + +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../entrypoint/local.sh" +fi + +RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" \ + | jq -c '.env_vars.GEMMA4_CONVERSION_MODE = "text"') + +GEMMA4_SFT_THINKING="${GEMMA4_SFT_THINKING:-1}" +if [ -n "${GEMMA4_SFT_THINKING}" ]; then + RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" \ + | jq -c --arg v "${GEMMA4_SFT_THINKING}" '.env_vars.GEMMA4_SFT_THINKING = $v') +fi + +if [ -n "${NVTE_DEBUG:-}" ]; then + RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" | jq -c \ + --arg d "${NVTE_DEBUG}" --arg l "${NVTE_DEBUG_LEVEL:-2}" \ + '.env_vars.NVTE_DEBUG = $d | .env_vars.NVTE_DEBUG_LEVEL = $l') +fi +export RUNTIME_ENV_JSON + +MODEL_DIR="${MODEL_DIR:?set MODEL_DIR to the dir containing gemma-4-26B-A4B-it}" +CKPT="${CKPT:-${MODEL_DIR}/gemma-4-26B-A4B-it}" +PROMPT_DATA="${PROMPT_DATA:?set PROMPT_DATA to an SFT jsonl}" + +PROJECT_NAME="${PROJECT_NAME:-Relax/sft/gemma4}" +EXP_NAME="${EXP_NAME:-gemma4-26b-moe-sft-gpu8}" +CKPT_ROOT="${CKPT_ROOT:-/data/temp}" +SAVE_DIR="${SAVE_DIR:-${CKPT_ROOT}/gemma4-26B-sft}" + +RAY_ADDRESS="${RAY_ADDRESS:-http://${MASTER_ADDR:-127.0.0.1}:${RAY_DASHBOARD_PORT:-8265}}" + +CKPT_ARGS=( + --hf-checkpoint "${CKPT}" + --ref-load "${CKPT}" + --load "${CKPT}" + --save "${SAVE_DIR}" + + --megatron-to-hf-mode bridge + --warm-hf-checkpoint-page-cache + + --save-interval 1000000 + --num-epoch 1 +) + +SFT_ARGS=( + --loss-type sft + --prompt-data "${PROMPT_DATA}" + --input-key "${INPUT_KEY:-instruction}" + --label-key "${LABEL_KEY:-output}" + + --use-dynamic-batch-size + --max-tokens-per-gpu ${MAX_TOKENS_PER_GPU:-4096} + --balance-data + + --global-batch-size ${GLOBAL_BATCH_SIZE:-512} + --seq-length ${SEQ_LENGTH:-4096} + --num-rollout ${NUM_ROLLOUT:-40} +) + +PERF_ARGS=( + --num-layers ${NUM_LAYERS:-30} + + --tensor-model-parallel-size ${TP_SIZE} + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + + --expert-model-parallel-size ${EP_SIZE} + --expert-tensor-parallel-size 1 + --sequence-parallel + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --data-parallel-sharding-strategy optim +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 5e-6 + --min-lr 5e-7 + --lr-decay-style ${LR_DECAY_STYLE:-cosine} + --lr-warmup-fraction 0.1 + + --weight-decay ${WEIGHT_DECAY:-0.1} + --adam-beta1 0.9 + --adam-beta2 ${ADAM_BETA2:-0.95} + + --adam-eps 1e-8 + --clip-grad 1.0 +) + +PER_TOKEN_LOSS_FLAG=() +[ "${CALC_PER_TOKEN_LOSS:-1}" != "0" ] && PER_TOKEN_LOSS_FLAG=(--calculate-per-token-loss) + +MISC_ARGS=( + --bf16 + --attention-backend ${ATTENTION_BACKEND:-auto} + --cross-entropy-fusion-impl te + --log-interval 1 + --distributed-timeout-minutes ${DISTRIBUTED_TIMEOUT_MINUTES:-30} + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --recompute-loss-function + --use-health-check + --no-save-rng + + --seed ${SEED:-42} + + "${PER_TOKEN_LOSS_FLAG[@]}" +) + +WANDB_ARGS=( + --use-clearml + --use-metrics-service + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name ${EXP_NAME}-${now} +) + +mkdir -p "${SAVE_DIR}" log + +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="${RAY_ADDRESS}" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource "{\"sft\": [1, 0], \"actor\": [1, ${ACTOR_GPUS}]}" \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${SFT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "${WANDB_ARGS[@]}" 2>&1 | tee log/${EXP_NAME}-${now}.log diff --git a/scripts/training/sft/run-gemma4-26B-sft-hf-8xgpu.sh b/scripts/training/sft/run-gemma4-26B-sft-hf-8xgpu.sh new file mode 100755 index 000000000..0d3d581ea --- /dev/null +++ b/scripts/training/sft/run-gemma4-26B-sft-hf-8xgpu.sh @@ -0,0 +1,222 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# gemma-4-26B-A4B-it (MoE) full-parameter SFT with packing, 8xGPU, TP=8 EP=8 PP=1, +# ray-submit launch. Same as run-gemma4-26B-sft-8xgpu.sh, plus a HuggingFace-format +# export written alongside the native torch_dist checkpoint. +# +# Usage: +# MEGATRON= MODEL_DIR= \ +# PROMPT_DATA= bash scripts/training/sft/run-gemma4-26B-sft-hf-8xgpu.sh +# +# # fp8 export instead of bf16 (~half the size; needs a safetensors --hf-checkpoint): +# SAVE_HF_DTYPE=fp8 MEGATRON=... bash scripts/training/sft/run-gemma4-26B-sft-hf-8xgpu.sh + +set -ex +set -o pipefail + +unset NCCL_NVLS_ENABLE + +now=$(date "+%Y-%m-%d-%H:%M:%S") + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${RELAX:-${SCRIPT_DIR}/../../..}" &>/dev/null && pwd)" +export RELAX="${RELAX_ROOT}" + +export MEGATRON="${MEGATRON:-/root/Megatron-LM/}" + +if ! PYTHONPATH="${RELAX_ROOT}:${MEGATRON}:${PYTHONPATH:-}" python3 -c \ + 'import inspect; from megatron.bridge.models.gemma.gemma4_provider import Gemma4DenseProvider, Gemma4ModelProvider; from megatron.bridge.models.gemma_vl.gemma4_vl_bridge import Gemma4VLBridge; from relax.models.gemma4.gemma4_bridge import Gemma4DenseBridge; from relax.models.gemma4.gemma4_provider import RELAX_PROVIDERS, PackedSafeGemma4DenseProvider, RelaxGemma4MoEProvider; source = inspect.getsource(Gemma4VLBridge.provider_bridge); ok = source.count("_conversion_mode()") >= 2 and RELAX_PROVIDERS.get(Gemma4DenseProvider) is PackedSafeGemma4DenseProvider and RELAX_PROVIDERS.get(Gemma4ModelProvider) is RelaxGemma4MoEProvider +if not ok: + raise RuntimeError("Gemma4 packed-safe provider mapping is unavailable")' \ + >/dev/null 2>&1; then + echo "ERROR: Gemma4 packed-safe integration failed its startup probe." >&2 + echo " ${MEGATRON} must provide the required providers and MoE text mode," >&2 + echo " and Relax's Gemma4 bridge/provider replacements must import cleanly." >&2 + exit 1 +fi + +MODEL_CONFIG_DIR="${MODEL_CONFIG_DIR:-${SCRIPT_DIR}/../../models}" +source "${MODEL_CONFIG_DIR}/gemma4-26B.sh" + +TP_SIZE="${TP_SIZE:-8}" +EP_SIZE="${EP_SIZE:-8}" +ACTOR_GPUS="${ACTOR_GPUS:-8}" + +if [ "$((ACTOR_GPUS % TP_SIZE))" -ne 0 ]; then + echo "ERROR: ACTOR_GPUS=${ACTOR_GPUS} is not divisible by TP_SIZE=${TP_SIZE}." >&2 + exit 1 +fi +if [ "$((NHEADS % TP_SIZE))" -ne 0 ] || [ "$((NUM_QUERY_GROUPS % TP_SIZE))" -ne 0 ]; then + echo "ERROR: TP_SIZE=${TP_SIZE} must divide NHEADS=${NHEADS} and NUM_QUERY_GROUPS=${NUM_QUERY_GROUPS}." >&2 + exit 1 +fi +if [ "$((MOE_ROUTED_EXPERTS % EP_SIZE))" -ne 0 ]; then + echo "ERROR: EP_SIZE=${EP_SIZE} must divide MOE_ROUTED_EXPERTS=${MOE_ROUTED_EXPERTS}." >&2 + exit 1 +fi +if [ "${TP_SIZE}" -lt 1 ] || [ "${TP_SIZE}" -gt 16 ] || [ "$((TP_SIZE & (TP_SIZE - 1)))" -ne 0 ]; then + echo "ERROR: TP_SIZE=${TP_SIZE} must be a power of two between 1 and 16." >&2 + exit 1 +fi +if [ "${EP_SIZE}" -ne "${ACTOR_GPUS}" ]; then + echo "ERROR: EP_SIZE must equal ACTOR_GPUS (= TP * CP * DP with PP=CP=1)." >&2 + echo " Got EP_SIZE=${EP_SIZE}, ACTOR_GPUS=${ACTOR_GPUS}, TP_SIZE=${TP_SIZE}." >&2 + echo " EP smaller than that replicates experts on every rank and will OOM;" >&2 + echo " larger is rejected by megatron." >&2 + exit 1 +fi +echo "parallelism: TP=${TP_SIZE} EP=${EP_SIZE} PP=1 CP=1 on ${ACTOR_GPUS} GPU(s); num_layers=${NUM_LAYERS:-30}" + +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../entrypoint/local.sh" +fi + +RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" \ + | jq -c '.env_vars.GEMMA4_CONVERSION_MODE = "text"') + +GEMMA4_SFT_THINKING="${GEMMA4_SFT_THINKING:-1}" +if [ -n "${GEMMA4_SFT_THINKING}" ]; then + RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" \ + | jq -c --arg v "${GEMMA4_SFT_THINKING}" '.env_vars.GEMMA4_SFT_THINKING = $v') +fi + +if [ -n "${NVTE_DEBUG:-}" ]; then + RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" | jq -c \ + --arg d "${NVTE_DEBUG}" --arg l "${NVTE_DEBUG_LEVEL:-2}" \ + '.env_vars.NVTE_DEBUG = $d | .env_vars.NVTE_DEBUG_LEVEL = $l') +fi +export RUNTIME_ENV_JSON + +MODEL_DIR="${MODEL_DIR:?set MODEL_DIR to the dir containing gemma-4-26B-A4B-it}" +CKPT="${CKPT:-${MODEL_DIR}/gemma-4-26B-A4B-it}" +PROMPT_DATA="${PROMPT_DATA:?set PROMPT_DATA to an SFT jsonl}" + +PROJECT_NAME="${PROJECT_NAME:-Relax/sft/gemma4}" +EXP_NAME="${EXP_NAME:-gemma4-26b-moe-sft-gpu8}" + +CKPT_ROOT="${CKPT_ROOT:-/data/temp}" +SAVE_DIR="${SAVE_DIR:-${CKPT_ROOT}/gemma4-26B-sft}" + +SAVE_HF_DIR="${SAVE_HF_DIR:-${SAVE_DIR}/hf_output/${EXP_NAME}}" +SAVE_HF_DTYPE="${SAVE_HF_DTYPE:-bf16}" + +RAY_ADDRESS="${RAY_ADDRESS:-http://${MASTER_ADDR:-127.0.0.1}:${RAY_DASHBOARD_PORT:-8265}}" + +CKPT_ARGS=( + --hf-checkpoint "${CKPT}" + --ref-load "${CKPT}" + --load "${CKPT}" + --save "${SAVE_DIR}" + + --megatron-to-hf-mode bridge + --warm-hf-checkpoint-page-cache + + --save-interval 1000000 + --num-epoch 1 + + --save-hf "${SAVE_HF_DIR}/iter_{rollout_id}" + --save-hf-dtype "${SAVE_HF_DTYPE}" +) + +if [ "${SAVE_HF_DTYPE}" = "fp8" ]; then + CKPT_ARGS+=( + --save-hf-fp8-quant-mode "${SAVE_HF_FP8_QUANT_MODE:-block}" + --save-hf-fp8-block-size ${SAVE_HF_FP8_BLOCK_SIZE:-128 128} + ) +fi + +SFT_ARGS=( + --loss-type sft + --prompt-data "${PROMPT_DATA}" + --input-key "${INPUT_KEY:-instruction}" + --label-key "${LABEL_KEY:-output}" + + --use-dynamic-batch-size + --max-tokens-per-gpu ${MAX_TOKENS_PER_GPU:-4096} + --balance-data + + --global-batch-size ${GLOBAL_BATCH_SIZE:-512} + --seq-length ${SEQ_LENGTH:-4096} + --num-rollout ${NUM_ROLLOUT:-40} +) + +PERF_ARGS=( + --num-layers ${NUM_LAYERS:-30} + + --tensor-model-parallel-size ${TP_SIZE} + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + + --expert-model-parallel-size ${EP_SIZE} + --expert-tensor-parallel-size 1 + --sequence-parallel + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --data-parallel-sharding-strategy optim +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 5e-6 + --min-lr 5e-7 + --lr-decay-style ${LR_DECAY_STYLE:-cosine} + --lr-warmup-fraction 0.1 + + --weight-decay ${WEIGHT_DECAY:-0.1} + --adam-beta1 0.9 + --adam-beta2 ${ADAM_BETA2:-0.95} + + --adam-eps 1e-8 + --clip-grad 1.0 +) + +PER_TOKEN_LOSS_FLAG=() +[ "${CALC_PER_TOKEN_LOSS:-1}" != "0" ] && PER_TOKEN_LOSS_FLAG=(--calculate-per-token-loss) + +MISC_ARGS=( + --bf16 + --attention-backend ${ATTENTION_BACKEND:-auto} + --cross-entropy-fusion-impl te + --log-interval 1 + --distributed-timeout-minutes ${DISTRIBUTED_TIMEOUT_MINUTES:-30} + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --recompute-loss-function + --use-health-check + --no-save-rng + + --seed ${SEED:-42} + + "${PER_TOKEN_LOSS_FLAG[@]}" +) + +WANDB_ARGS=( + --use-clearml + --use-metrics-service + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name ${EXP_NAME}-${now} +) + +mkdir -p "${SAVE_DIR}" log + +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="${RAY_ADDRESS}" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource "{\"sft\": [1, 0], \"actor\": [1, ${ACTOR_GPUS}]}" \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${SFT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "${WANDB_ARGS[@]}" 2>&1 | tee log/${EXP_NAME}-${now}.log diff --git a/scripts/training/sft/run-gemma4-31B-sft-8xgpu.sh b/scripts/training/sft/run-gemma4-31B-sft-8xgpu.sh new file mode 100644 index 000000000..7c9e1ed9c --- /dev/null +++ b/scripts/training/sft/run-gemma4-31B-sft-8xgpu.sh @@ -0,0 +1,187 @@ +#!/bin/bash + +# Copyright (c) 2026 Relax Authors. All Rights Reserved. +# +# gemma-4-31B-it full-parameter SFT, TP=8 PP=1, ray-submit launch. +# +# Usage: +# MEGATRON= MODEL_DIR= \ +# PROMPT_DATA= bash scripts/training/sft/run-gemma4-31B-sft-8xgpu.sh +# +# # multi-node (QS sets MASTER_ADDR/POD_NAME/WORLD_SIZE): +# ACTOR_GPUS=16 bash scripts/entrypoint/spmd-multinode.sh \ +# scripts/training/sft/run-gemma4-31B-sft-8xgpu.sh + +set -ex +set -o pipefail + +unset NCCL_NVLS_ENABLE + +now=$(date "+%Y-%m-%d-%H:%M:%S") + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)" +RELAX_ROOT="$(cd -- "${RELAX:-${SCRIPT_DIR}/../../..}" &>/dev/null && pwd)" +export RELAX="${RELAX_ROOT}" + +export MEGATRON="${MEGATRON:-/root/Megatron-LM/}" + +if ! PYTHONPATH="${RELAX_ROOT}:${MEGATRON}:${PYTHONPATH:-}" python3 -c \ + 'from megatron.bridge.models.gemma.gemma4_provider import Gemma4DenseProvider; from relax.models.gemma4.gemma4_bridge import Gemma4DenseBridge; from relax.models.gemma4.gemma4_provider import RELAX_PROVIDERS, PackedSafeGemma4DenseProvider; ok = RELAX_PROVIDERS.get(Gemma4DenseProvider) is PackedSafeGemma4DenseProvider +if not ok: + raise RuntimeError("Gemma4 packed-safe provider mapping is unavailable")' \ + >/dev/null 2>&1; then + echo "ERROR: Gemma4 packed-safe integration failed its startup probe." >&2 + echo " ${MEGATRON} must provide Gemma4DenseProvider, and Relax's" >&2 + echo " Gemma4 bridge/provider replacement must import cleanly." >&2 + exit 1 +fi + +TP_SIZE="${TP_SIZE:-8}" +ACTOR_GPUS="${ACTOR_GPUS:-8}" + +if [ "$((16 % TP_SIZE))" -ne 0 ]; then + echo "ERROR: TP_SIZE=${TP_SIZE} does not divide the 16 KV heads." >&2 + exit 1 +fi +if [ "$((ACTOR_GPUS % TP_SIZE))" -ne 0 ]; then + echo "ERROR: ACTOR_GPUS=${ACTOR_GPUS} is not divisible by TP_SIZE=${TP_SIZE}." >&2 + exit 1 +fi +echo "parallelism: TP=${TP_SIZE} DP=$((ACTOR_GPUS / TP_SIZE)) PP=1 CP=1 on ${ACTOR_GPUS} GPU(s)" + +if [ -z "${RELAX_ENTRYPOINT_MODE:-}" ]; then + source "${SCRIPT_DIR}/../../entrypoint/local.sh" +fi + +RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" \ + | jq -c '.env_vars.GEMMA4_CONVERSION_MODE = "text"') + +GEMMA4_SFT_THINKING="${GEMMA4_SFT_THINKING:-1}" +if [ -n "${GEMMA4_SFT_THINKING}" ]; then + RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" \ + | jq -c --arg v "${GEMMA4_SFT_THINKING}" '.env_vars.GEMMA4_SFT_THINKING = $v') +fi + +if [ -n "${NVTE_DEBUG:-}" ]; then + RUNTIME_ENV_JSON=$(printf '%s' "${RUNTIME_ENV_JSON}" | jq -c \ + --arg d "${NVTE_DEBUG}" --arg l "${NVTE_DEBUG_LEVEL:-2}" \ + '.env_vars.NVTE_DEBUG = $d | .env_vars.NVTE_DEBUG_LEVEL = $l') +fi +export RUNTIME_ENV_JSON + +source "${MODEL_CONFIG_DIR}/gemma4-31B.sh" + +MODEL_DIR="${MODEL_DIR:?set MODEL_DIR to the dir containing gemma-4-31B-it}" +CKPT="${CKPT:-${MODEL_DIR}/gemma-4-31B-it}" +PROMPT_DATA="${PROMPT_DATA:?set PROMPT_DATA to an SFT jsonl}" + +PROJECT_NAME="${PROJECT_NAME:-Relax/sft/gemma4}" +EXP_NAME="${EXP_NAME:-gemma4-31b-sft-smoke-gpu8}" +CKPT_ROOT="${CKPT_ROOT:-/data/temp}" +SAVE_DIR="${SAVE_DIR:-${CKPT_ROOT}/gemma4-31B-sft}" + +RAY_ADDRESS="${RAY_ADDRESS:-http://${MASTER_ADDR:-127.0.0.1}:${RAY_DASHBOARD_PORT:-8265}}" + +CKPT_ARGS=( + --hf-checkpoint "${CKPT}" + --ref-load "${CKPT}" + --load "${CKPT}" + --save "${SAVE_DIR}" + + --megatron-to-hf-mode bridge + --warm-hf-checkpoint-page-cache + + --save-interval 1000000 + --num-epoch 1 +) + +SFT_ARGS=( + --loss-type sft + --prompt-data "${PROMPT_DATA}" + --input-key "${INPUT_KEY:-instruction}" + --label-key "${LABEL_KEY:-output}" + + --use-dynamic-batch-size + --max-tokens-per-gpu ${MAX_TOKENS_PER_GPU:-4096} + --balance-data + + --global-batch-size ${GLOBAL_BATCH_SIZE:-256} + --seq-length ${SEQ_LENGTH:-4096} + ${NUM_ROLLOUT:+--num-rollout ${NUM_ROLLOUT}} +) + +PERF_ARGS=( + --tensor-model-parallel-size ${TP_SIZE} + --pipeline-model-parallel-size 1 + --context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + --sequence-parallel + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + --data-parallel-sharding-strategy optim +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 5e-6 + --min-lr 5e-7 + --lr-decay-style ${LR_DECAY_STYLE:-cosine} + --lr-warmup-fraction 0.1 + + --weight-decay ${WEIGHT_DECAY:-0.1} + --adam-beta1 0.9 + --adam-beta2 ${ADAM_BETA2:-0.95} + + --adam-eps 1e-8 + --clip-grad 1.0 +) + +PER_TOKEN_LOSS_FLAG=() +[ "${CALC_PER_TOKEN_LOSS:-1}" != "0" ] && PER_TOKEN_LOSS_FLAG=(--calculate-per-token-loss) + +MISC_ARGS=( + --bf16 + --attention-backend auto + --cross-entropy-fusion-impl te + --log-interval 1 + --distributed-timeout-minutes ${DISTRIBUTED_TIMEOUT_MINUTES:-30} + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --recompute-loss-function + --use-health-check + --no-save-rng + + --seed ${SEED:-42} + + "${PER_TOKEN_LOSS_FLAG[@]}" +) + +WANDB_ARGS=( + --use-clearml + --use-metrics-service + --tb-project-name ${PROJECT_NAME} + --tb-experiment-name ${EXP_NAME}-${now} +) + +mkdir -p "${SAVE_DIR}" log + +ray job submit ${RAY_NO_WAIT:+--no-wait} --address="${RAY_ADDRESS}" \ + ${WORKING_DIR:+--working-dir "${WORKING_DIR}"} \ + --runtime-env-json="${RUNTIME_ENV_JSON}" \ + -- python3 -m relax.entrypoints.train \ + --resource "{\"sft\": [1, 0], \"actor\": [1, ${ACTOR_GPUS}]}" \ + --max-staleness 0 \ + --num-data-storage-units 1 \ + "${MODEL_ARGS[@]}" \ + "${CKPT_ARGS[@]}" \ + "${SFT_ARGS[@]}" \ + "${OPTIMIZER_ARGS[@]}" \ + "${PERF_ARGS[@]}" \ + "${MISC_ARGS[@]}" \ + "${WANDB_ARGS[@]}" 2>&1 | tee log/${EXP_NAME}-${now}.log diff --git a/tests/backends/megatron/test_save_hf_strictness.py b/tests/backends/megatron/test_save_hf_strictness.py index c333c2568..a42f06307 100644 --- a/tests/backends/megatron/test_save_hf_strictness.py +++ b/tests/backends/megatron/test_save_hf_strictness.py @@ -120,3 +120,64 @@ def test_mtp_reference_without_mtp_model_relaxes_strict(monkeypatch, tmp_path): model_mod.save_hf_model(_args(tmp_path), rollout_id=4, model=[]) assert recorded["strict"] is False + + +def _record_reconcile(monkeypatch): + """Capture how save_hf_model calls the reconciler instead of running it.""" + calls = [] + + def _fake(path, reference_hf_dir=None, supplement_mtp=True, **kwargs): + calls.append({"path": path, "reference_hf_dir": reference_hf_dir, "supplement_mtp": supplement_mtp}) + + monkeypatch.setattr(hf_export, "reconcile_hf_export_index", _fake) + return calls + + +def test_reconcile_runs_for_a_vision_only_relaxation_and_leaves_mtp_alone(monkeypatch, tmp_path): + """Ghost entries can come from either relaxation, MTP supplementation + cannot. + + Reconcile has to run whenever the save was non-strict, but pulling mtp.* + weights out of the base is only ever right when MTP was the relaxed group + -- otherwise a genuinely missing MTP tensor gets papered over with base + weights. + """ + recorded = {} + _install_fakes(monkeypatch, recorded) + calls = _record_reconcile(monkeypatch) + monkeypatch.setattr(hf_export, "reference_expects_mtp", lambda path: False) + monkeypatch.setattr(hf_export, "reference_expects_vision", lambda path: True) + + model_mod.save_hf_model(_args(tmp_path), rollout_id=6, model=_model(vision=False)) + + assert recorded["strict"] is False + assert len(calls) == 1 + assert calls[0]["supplement_mtp"] is False + + +def test_reconcile_supplements_mtp_for_the_mtp_relaxation(monkeypatch, tmp_path): + recorded = {} + _install_fakes(monkeypatch, recorded) + calls = _record_reconcile(monkeypatch) + monkeypatch.setattr(hf_export, "reference_expects_mtp", lambda path: True) + monkeypatch.setattr(hf_export, "reference_expects_vision", lambda path: False) + + model_mod.save_hf_model(_args(tmp_path), rollout_id=7, model=[]) + + assert recorded["strict"] is False + assert len(calls) == 1 + assert calls[0]["supplement_mtp"] is True + + +def test_reconcile_skipped_when_the_save_was_strict(monkeypatch, tmp_path): + """Nothing was relaxed, so there can be no ghost entries to reconcile.""" + recorded = {} + _install_fakes(monkeypatch, recorded) + calls = _record_reconcile(monkeypatch) + monkeypatch.setattr(hf_export, "reference_expects_mtp", lambda path: False) + monkeypatch.setattr(hf_export, "reference_expects_vision", lambda path: False) + + model_mod.save_hf_model(_args(tmp_path), rollout_id=8, model=[]) + + assert recorded["strict"] is True + assert calls == [] diff --git a/tests/backends/megatron/weight_update/test_lora_weight_sync.py b/tests/backends/megatron/weight_update/test_lora_weight_sync.py index 1ee816da7..fadbd7e29 100644 --- a/tests/backends/megatron/weight_update/test_lora_weight_sync.py +++ b/tests/backends/megatron/weight_update/test_lora_weight_sync.py @@ -319,10 +319,12 @@ def test_tp1_formula_reference(self): assert expected.shape == base.shape def test_tp1_matches_real_loramerge(self): - pytest.importorskip("megatron.bridge.peft.lora") from inspect import signature - from megatron.bridge.peft.lora import LoRAMerge + lora = pytest.importorskip("megatron.bridge.peft.lora") + if not hasattr(lora, "LoRAMerge"): + pytest.skip("this Megatron-Bridge build does not expose LoRAMerge") + LoRAMerge = lora.LoRAMerge if "tp_size" not in signature(LoRAMerge().merge).parameters: pytest.skip("installed megatron bridge LoRAMerge.merge lacks tp_size support") diff --git a/tests/models/gemma4/test_attention.py b/tests/models/gemma4/test_attention.py new file mode 100644 index 000000000..c613c5ca4 --- /dev/null +++ b/tests/models/gemma4/test_attention.py @@ -0,0 +1,88 @@ +# Copyright (c) 2026 Relax Authors. All Rights Reserved. + +"""Safety checks for Gemma-4's packed-only core attention.""" + +import importlib.util +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +# The unit-test image intentionally carries an older Megatron tree without the +# Gemma-4 bridge classes. Load the leaf module directly so these attention tests +# do not execute relax.models.gemma4.__init__ and fail on an unrelated optional +# integration dependency. +_MODULE_PATH = Path(__file__).parents[3] / "relax/models/gemma4/attention.py" +_SPEC = importlib.util.spec_from_file_location("_relax_gemma4_attention_test_target", _MODULE_PATH) +assert _SPEC is not None and _SPEC.loader is not None +_MODULE = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_MODULE) +Gemma4CoreAttention = _MODULE.Gemma4CoreAttention + + +def _attention(): + config = SimpleNamespace( + attention_dropout=0.0, + context_parallel_size=1, + softmax_scale=1.0, + window_size=(1023, 0), + window_attn_skip_freq=6, + ) + return Gemma4CoreAttention(config, layer_number=1) + + +def test_thd_without_sequence_boundaries_fails_loudly(): + attention = _attention() + q = torch.zeros(8, 2, 4) + with pytest.raises(NotImplementedError, match="requires packed THD inputs with cu_seqlens"): + attention(q, q, q) + + +def test_sbhd_fails_instead_of_silently_dropping_sliding_window(): + attention = _attention() + q = torch.zeros(8, 1, 2, 4) + with pytest.raises(NotImplementedError, match="supports packed THD attention only"): + attention(q, q, q) + + +def test_flash_path_reuses_packed_max_seqlen(monkeypatch): + captured = {} + + def fake_flash_attn_varlen_func(query, key, value, **kwargs): + captured.update(kwargs) + return torch.zeros_like(query) + + monkeypatch.setitem( + sys.modules, + "flash_attn", + SimpleNamespace(flash_attn_varlen_func=fake_flash_attn_varlen_func), + ) + attention = _attention() + q = torch.zeros(8, 2, 4) + packed = SimpleNamespace(cu_seqlens_q=torch.tensor([0, 3, 8]), max_seqlen_q=123) + + attention(q, q, q, packed_seq_params=packed) + + assert captured["max_seqlen_q"] == 123 + assert captured["max_seqlen_k"] == 123 + + +def test_sdpa_path_reuses_cpu_sequence_boundaries(): + class BoundarySpy: + def detach(self): + raise AssertionError("cached CPU boundaries should avoid accelerator synchronization") + + attention = _attention() + q = torch.zeros(8, 2, 257) + packed = SimpleNamespace( + cu_seqlens_q=BoundarySpy(), + cu_seqlens_q_cpu=[0, 3, 8], + max_seqlen_q=5, + ) + + output = attention(q, q, q, packed_seq_params=packed) + + assert output.shape == (8, 514) From 0a91dd1db44022920f61f05948e89a5cb83f8b38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=AE=81=E6=9C=AC=E5=93=B2?= Date: Thu, 3 Sep 2026 11:46:02 +0000 Subject: [PATCH 18/34] feat(diffusion): add native generative RL - Add FSDP2 full and LoRA training backends for Qwen-Image - Add native SGLang diffusion rollout, reward, and in-memory trajectory transport - Add transactional weight synchronization, checkpointing, offload, and resume support - Add launch recipes, data tools, evaluation assets, and bilingual documentation --- - Validate trajectory contracts and replay the final valid diffusion transition - Coordinate distributed failures, RNG persistence, checkpoint publication, and rotation - Preserve optional Megatron imports and existing text-RL defaults --- - Keep trajectories and local reward images in memory - Stream bucketed weights directly to rollout engines - Prune stale rollout objects and artifact directories --- - Add unit coverage for FSDP lifecycle, LoRA, checkpointing, rollout, rewards, and weight updates - Add SGLang patch and Qwen-Image replay contract tests (cherry picked from commit cff14b1057edb65c89051e912a9f6822b4dcfdbe) --- docker/patch/sglang/v0.5.17.patch | 834 ++++++ docs/.vitepress/config.mts | 6 +- docs/draft/multimodal-gen-rl-design.md | 1167 ++++++++ docs/en/guide/diffusion-generative-rl.md | 736 +++++ docs/zh/guide/diffusion-generative-rl.md | 735 +++++ examples/diffusion/README.md | 332 +++ ...wen-image-ft-lora-pickscore-comparison.csv | 334 +++ ...wen-image-ft-lora-pickscore-comparison.png | Bin 0 -> 495493 bytes .../qwen-image-ft-lora-pickscore-summary.json | 38 + .../assets/qwen-image-lora-eval-curve.png | Bin 0 -> 100899 bytes examples/diffusion/common/full_ft.yaml | 41 + examples/diffusion/common/lora.yaml | 66 + examples/diffusion/curate_data.py | 186 ++ examples/diffusion/evaluate.py | 138 + examples/diffusion/export_checkpoint.py | 165 ++ examples/diffusion/inspect_data.py | 121 + examples/diffusion/prepare_data.py | 162 ++ examples/diffusion/qwen_image/t2i_full.yaml | 60 + examples/diffusion/qwen_image/t2i_lora.yaml | 73 + pyproject.toml | 2 + relax/backends/fsdp/__init__.py | 3 + relax/backends/fsdp/actor.py | 2416 +++++++++++++++++ relax/backends/fsdp/arguments.py | 621 +++++ relax/backends/fsdp/checkpoint.py | 701 +++++ relax/backends/fsdp/lora.py | 467 ++++ relax/backends/fsdp/runtime.py | 319 +++ relax/backends/fsdp/weight_update.py | 246 ++ relax/backends/sglang/diffusion_engine.py | 906 +++++++ .../backends/device_direct.py | 101 +- relax/distributed/ray/actor_group.py | 20 +- relax/distributed/ray/generative_reward.py | 95 + relax/distributed/ray/rollout.py | 18 +- relax/engine/rewards/generative.py | 477 ++++ relax/engine/rewards/pickscore.py | 137 + relax/engine/rollout/data_source.py | 15 +- relax/engine/rollout/native_generation.py | 1282 +++++++++ relax/models/__init__.py | 37 +- relax/models/flow_grpo.py | 319 +++ relax/models/generative.py | 338 +++ relax/models/qwen_image/__init__.py | 7 + relax/models/qwen_image/adapter.py | 385 +++ relax/utils/arguments.py | 33 +- relax/utils/megatron_peft_utils.py | 9 +- relax/utils/rotate_ckpt.py | 28 +- relax/utils/utils.py | 16 + requirements.txt | 5 + .../diffusion/run-qwen-image-t2i-8xgpu.sh | 370 +++ .../run-qwen-image-t2i-lora-8xgpu.sh | 430 +++ tests/backends/fsdp/__init__.py | 0 tests/backends/fsdp/test_actor_lifecycle.py | 320 +++ tests/backends/fsdp/test_checkpoint.py | 464 ++++ .../backends/fsdp/test_full_weight_update.py | 116 + tests/backends/fsdp/test_lora.py | 490 ++++ tests/backends/fsdp/test_offload.py | 133 + .../backends/fsdp/test_save_and_validation.py | 590 ++++ .../backends/fsdp/test_weight_transaction.py | 284 ++ .../backends/sglang/test_diffusion_engine.py | 503 ++++ .../test_optional_megatron_import.py | 53 + .../distributed/ray/test_generative_reward.py | 94 + tests/engine/rewards/test_generative.py | 304 +++ tests/engine/rewards/test_pickscore.py | 75 + .../engine/rollout/test_native_generation.py | 658 +++++ tests/examples/test_diffusion_data.py | 97 + tests/models/__init__.py | 0 tests/models/qwen_image/__init__.py | 0 tests/models/qwen_image/test_adapter.py | 247 ++ tests/models/test_flow_grpo.py | 180 ++ tests/models/test_generative_contract.py | 119 + tests/test_model_source.py | 1 + tests/utils/test_rotate_ckpt.py | 30 + 70 files changed, 19691 insertions(+), 64 deletions(-) create mode 100644 docs/draft/multimodal-gen-rl-design.md create mode 100644 docs/en/guide/diffusion-generative-rl.md create mode 100644 docs/zh/guide/diffusion-generative-rl.md create mode 100644 examples/diffusion/README.md create mode 100644 examples/diffusion/assets/qwen-image-ft-lora-pickscore-comparison.csv create mode 100644 examples/diffusion/assets/qwen-image-ft-lora-pickscore-comparison.png create mode 100644 examples/diffusion/assets/qwen-image-ft-lora-pickscore-summary.json create mode 100644 examples/diffusion/assets/qwen-image-lora-eval-curve.png create mode 100644 examples/diffusion/common/full_ft.yaml create mode 100644 examples/diffusion/common/lora.yaml create mode 100755 examples/diffusion/curate_data.py create mode 100644 examples/diffusion/evaluate.py create mode 100755 examples/diffusion/export_checkpoint.py create mode 100755 examples/diffusion/inspect_data.py create mode 100755 examples/diffusion/prepare_data.py create mode 100644 examples/diffusion/qwen_image/t2i_full.yaml create mode 100644 examples/diffusion/qwen_image/t2i_lora.yaml create mode 100644 relax/backends/fsdp/__init__.py create mode 100644 relax/backends/fsdp/actor.py create mode 100644 relax/backends/fsdp/arguments.py create mode 100644 relax/backends/fsdp/checkpoint.py create mode 100644 relax/backends/fsdp/lora.py create mode 100644 relax/backends/fsdp/runtime.py create mode 100644 relax/backends/fsdp/weight_update.py create mode 100644 relax/backends/sglang/diffusion_engine.py create mode 100644 relax/distributed/ray/generative_reward.py create mode 100644 relax/engine/rewards/generative.py create mode 100644 relax/engine/rewards/pickscore.py create mode 100644 relax/engine/rollout/native_generation.py create mode 100644 relax/models/flow_grpo.py create mode 100644 relax/models/generative.py create mode 100644 relax/models/qwen_image/__init__.py create mode 100644 relax/models/qwen_image/adapter.py create mode 100644 scripts/training/diffusion/run-qwen-image-t2i-8xgpu.sh create mode 100755 scripts/training/diffusion/run-qwen-image-t2i-lora-8xgpu.sh create mode 100644 tests/backends/fsdp/__init__.py create mode 100644 tests/backends/fsdp/test_actor_lifecycle.py create mode 100644 tests/backends/fsdp/test_checkpoint.py create mode 100644 tests/backends/fsdp/test_full_weight_update.py create mode 100644 tests/backends/fsdp/test_lora.py create mode 100644 tests/backends/fsdp/test_offload.py create mode 100644 tests/backends/fsdp/test_save_and_validation.py create mode 100644 tests/backends/fsdp/test_weight_transaction.py create mode 100644 tests/backends/sglang/test_diffusion_engine.py create mode 100644 tests/distributed/checkpoint_service/test_optional_megatron_import.py create mode 100644 tests/distributed/ray/test_generative_reward.py create mode 100644 tests/engine/rewards/test_generative.py create mode 100644 tests/engine/rewards/test_pickscore.py create mode 100644 tests/engine/rollout/test_native_generation.py create mode 100644 tests/examples/test_diffusion_data.py create mode 100644 tests/models/__init__.py create mode 100644 tests/models/qwen_image/__init__.py create mode 100644 tests/models/qwen_image/test_adapter.py create mode 100644 tests/models/test_flow_grpo.py create mode 100644 tests/models/test_generative_contract.py create mode 100644 tests/utils/test_rotate_ckpt.py diff --git a/docker/patch/sglang/v0.5.17.patch b/docker/patch/sglang/v0.5.17.patch index e08a5000a..2b26ade6c 100644 --- a/docker/patch/sglang/v0.5.17.patch +++ b/docker/patch/sglang/v0.5.17.patch @@ -1,3 +1,837 @@ +diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +index 610267b..c592353 100644 +--- a/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py ++++ b/python/sglang/multimodal_gen/configs/pipeline_configs/qwen_image.py +@@ -15,6 +15,7 @@ from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConf + from sglang.multimodal_gen.configs.pipeline_configs.base import ( + ImagePipelineConfig, + ModelTaskType, ++ TextConditioningOutput, + maybe_unpad_latents, + pad_text_embeddings_with_mask, + shard_rotary_emb_for_sp, +@@ -28,6 +29,9 @@ from sglang.multimodal_gen.runtime.utils.condition_expansion import ( + from sglang.multimodal_gen.runtime.utils.vision import resize + from sglang.multimodal_gen.utils import calculate_dimensions + ++_QWEN_PROMPT_TEMPLATE_START_IDX = 34 ++_QWEN_MAX_SEQUENCE_LENGTH = 512 ++ + + def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor): + bool_mask = mask.bool() +@@ -47,7 +51,10 @@ def qwen_image_preprocess_text(prompt): + + + def qwen_image_postprocess_text( +- outputs, _text_inputs, drop_idx=34, return_attention_mask=False ++ outputs, ++ _text_inputs, ++ drop_idx=_QWEN_PROMPT_TEMPLATE_START_IDX, ++ return_attention_mask=False, + ): + """Postprocess Qwen text embeddings. + +@@ -61,6 +68,25 @@ def qwen_image_postprocess_text( + ) + split_hidden_states = [e[drop_idx:] for e in split_hidden_states] + conditioning = pad_text_embeddings_with_mask(split_hidden_states) ++ prompt_embeds = conditioning.prompt_embeds[:, :_QWEN_MAX_SEQUENCE_LENGTH] ++ prompt_embeds_mask = ( ++ conditioning.prompt_embeds_mask[:, :_QWEN_MAX_SEQUENCE_LENGTH] ++ if conditioning.prompt_embeds_mask is not None ++ else None ++ ) ++ prompt_seq_lens = ( ++ [ ++ min(int(seq_len), _QWEN_MAX_SEQUENCE_LENGTH) ++ for seq_len in conditioning.prompt_seq_lens ++ ] ++ if conditioning.prompt_seq_lens is not None ++ else None ++ ) ++ conditioning = TextConditioningOutput( ++ prompt_embeds, ++ prompt_embeds_mask, ++ prompt_seq_lens, ++ ) + if return_attention_mask: + return conditioning + return conditioning.prompt_embeds +@@ -179,6 +205,7 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig + dict( + padding=True, + truncation=True, ++ max_length=_QWEN_MAX_SEQUENCE_LENGTH + _QWEN_PROMPT_TEMPLATE_START_IDX, + ), + None, + ] +@@ -208,7 +235,10 @@ class QwenImagePipelineConfig(QwenImageRolloutPipelineMixin, ImagePipelineConfig + if tok_kwargs.get("max_length") is not None: + tok_kwargs["padding"] = "max_length" + else: +- tok_kwargs.setdefault("max_length", 1024) ++ tok_kwargs.setdefault( ++ "max_length", ++ _QWEN_MAX_SEQUENCE_LENGTH + _QWEN_PROMPT_TEMPLATE_START_IDX, ++ ) + tok_kwargs["padding"] = True + return tokenizer(prompts, **tok_kwargs) + +diff --git a/python/sglang/multimodal_gen/configs/post_training/rl_rollout.py b/python/sglang/multimodal_gen/configs/post_training/rl_rollout.py +index 9103832..15737cc 100644 +--- a/python/sglang/multimodal_gen/configs/post_training/rl_rollout.py ++++ b/python/sglang/multimodal_gen/configs/post_training/rl_rollout.py +@@ -11,7 +11,7 @@ from typing import Any, Callable + + from sglang.multimodal_gen.utils import StoreBoolean + +-_VALID_ROLLOUT_SDE_TYPES = ("sde", "cps", "ode") ++_VALID_ROLLOUT_SDE_TYPES = ("sde", "cps", "ode", "dance") + + + @dataclass +diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py +index c5636c5..7ddbad7 100644 +--- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py ++++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py +@@ -241,6 +241,16 @@ class SamplingParams: + # 0-indexed denoising-loop step filters; None = all steps. + rollout_sde_step_indices: list[int] | None = None + rollout_return_step_indices: list[int] | None = None ++ # Driver-supplied rollout reproducibility data. These fields are consumed by ++ # the diffusion post-training path; prepare_request copies them into Req so ++ # the scheduler, denoising and latent-preparation stages see the same recipe ++ # the trainer will replay. ++ initial_noise_group_ids: list[str] | None = None ++ initial_noise_latent_shape: list[int] | None = None ++ initial_noise_seed: int | None = None ++ denoise_seeds: list[str] | None = None ++ sigmas: list[float] | None = None ++ timesteps: list[float] | None = None + # if True, disallow user params to override subclass-defined protected fields + no_override_protected_fields: bool = field( + default=False, metadata={"batch_sig_exclude": True} +diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py +index 45b3676..236e19b 100644 +--- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py ++++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/io_struct.py +@@ -37,6 +37,17 @@ class UpdateWeightFromTensorCheckerReqInput: + expected_named_tensors_sha256: dict[str, str] + + ++@dataclass ++class SetLoraFromTensorReqInput: ++ """Request to install a LoRA adapter from an in-memory tensor payload.""" ++ ++ serialized_tensors: str | bytes ++ lora_name: str = "default" ++ target: str = "transformer" ++ strength: float = 1.0 ++ merge_mode: str | None = None ++ ++ + @dataclass + class GetWeightsChecksumReqInput: + """Compute SHA-256 checksum of loaded module weights for verification.""" +diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py +index 554feea..251dde7 100644 +--- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py ++++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/weights_api.py +@@ -6,6 +6,7 @@ from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + GetWeightsChecksumReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, ++ SetLoraFromTensorReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, +@@ -138,6 +139,51 @@ async def update_weights_from_tensor_checker(request: Request): + ) + + ++@router.post("/set_lora_from_tensor") ++async def set_lora_from_tensor(request: Request): ++ """Install a LoRA adapter from base64(torch.save(dict[str, Tensor])).""" ++ import base64 ++ ++ body = await request.json() ++ serialized = body.get("serialized_tensors") ++ if serialized is None: ++ return orjson_response( ++ {"success": False, "message": "serialized_tensors is required"}, ++ status_code=400, ++ ) ++ ++ req = SetLoraFromTensorReqInput( ++ serialized_tensors=( ++ base64.b64decode(serialized) if isinstance(serialized, str) else serialized ++ ), ++ lora_name=body.get("lora_name", "default"), ++ target=body.get("target", "transformer"), ++ strength=float(body.get("strength", 1.0)), ++ merge_mode=body.get("merge_mode"), ++ ) ++ ++ try: ++ response = await async_scheduler_client.forward(req) ++ except Exception as e: ++ return orjson_response({"success": False, "message": str(e)}, status_code=500) ++ ++ if response.error is not None: ++ return orjson_response( ++ {"success": False, "message": str(response.error)}, ++ status_code=400, ++ ) ++ result = response.output or {} ++ success = bool(result.get("success", False)) ++ return orjson_response( ++ { ++ "success": success, ++ "message": result.get("message", "Unknown status"), ++ "adapted_layers": result.get("adapted_layers"), ++ }, ++ status_code=200 if success else 400, ++ ) ++ ++ + @router.post("/get_weights_checksum") + async def get_weights_checksum(request: Request): + """Return SHA-256 checksum of each requested module's weights.""" +diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +index 1494f58..d3cb820 100644 +--- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py ++++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +@@ -759,6 +759,19 @@ def prepare_request( + VSA_sparsity=server_args.attention_backend_config.VSA_sparsity, + ) + sampling_params.apply_request_extra(req) ++ for name in ( ++ "initial_noise_group_ids", ++ "initial_noise_latent_shape", ++ "initial_noise_seed", ++ "denoise_seeds", ++ "sigmas", ++ "timesteps", ++ ): ++ value = getattr(sampling_params, name, None) ++ if value is not None: ++ if name == "timesteps": ++ value = torch.as_tensor(value, dtype=torch.float32) ++ setattr(req, name, value) + if getattr(sampling_params, "max_sequence_length", None) is not None: + req.max_sequence_length = sampling_params.max_sequence_length + +diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py +index cf64593..d0efdf3 100644 +--- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py ++++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py +@@ -20,6 +20,7 @@ from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( + GetWeightsChecksumReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, ++ SetLoraFromTensorReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, +@@ -135,6 +136,7 @@ class Scheduler(SchedulerWarmupMixin, SchedulerPostTrainingMixin, SchedulerDisag + ShutdownReq: self._handle_shutdown, + ReleaseRealtimeSessionReq: self._handle_release_realtime_session, + GetDisaggStatsReq: self._handle_get_disagg_stats, ++ SetLoraFromTensorReqInput: self._handle_set_lora_from_tensor, + UpdateWeightFromDiskReqInput: self._handle_update_weights_from_disk, + UpdateWeightFromTensorReqInput: self._handle_update_weights_from_tensor, + UpdateWeightFromTensorCheckerReqInput: ( +diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +index 1f25398..2ecef7d 100644 +--- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py ++++ b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_flow_match_euler_discrete.py +@@ -350,21 +350,23 @@ class FlowMatchEulerDiscreteScheduler( + sigmas_array = np.array(sigmas).astype(np.float32) + num_inference_steps = len(sigmas_array) + +- # 2. Perform timestep shifting. Either no shifting is applied, or resolution-dependent shifting of +- # "exponential" or "linear" type is applied +- if self.config.use_dynamic_shifting: +- assert mu is not None, "mu cannot be None when use_dynamic_shifting is True" +- sigmas_array = self.time_shift(mu, 1.0, sigmas_array) +- else: +- sigmas_array = ( +- self.shift * sigmas_array / (1 + (self.shift - 1) * sigmas_array) +- ) ++ # 2. Perform timestep shifting. Driver-supplied sigmas are already in the ++ # exact sigma space the trainer will replay, so do not shift or ++ # stretch them again. ++ if sigmas is None: ++ if self.config.use_dynamic_shifting: ++ assert mu is not None, "mu cannot be None when use_dynamic_shifting is True" ++ sigmas_array = self.time_shift(mu, 1.0, sigmas_array) ++ else: ++ sigmas_array = ( ++ self.shift * sigmas_array / (1 + (self.shift - 1) * sigmas_array) ++ ) + +- # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value +- if self.config.shift_terminal: +- sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) +- sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor) +- sigmas_array = sigmas_tensor.numpy() ++ # 3. If required, stretch the sigmas schedule to terminate at the configured `shift_terminal` value ++ if self.config.shift_terminal: ++ sigmas_tensor = torch.from_numpy(sigmas_array).to(dtype=torch.float32) ++ sigmas_tensor = self.stretch_shift_to_terminal(sigmas_tensor) ++ sigmas_array = sigmas_tensor.numpy() + + # 4. If required, convert sigmas to one of karras, exponential, or beta sigma schedules + if self.config.use_karras_sigmas: +diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +index 8f3b26a..fbceed3 100644 +--- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py ++++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +@@ -722,6 +722,21 @@ class LoRAPipeline(ComposedPipelineBase): + if adapter_config.get("lora_alpha") is not None: + adapter_lora_alpha = int(adapter_config["lora_alpha"]) + ++ self._ingest_lora_state_dict( ++ lora_state_dict, ++ lora_nickname, ++ adapter_lora_alpha=adapter_lora_alpha, ++ ) ++ self.loaded_adapter_paths[lora_nickname] = lora_path ++ logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path) ++ ++ def _ingest_lora_state_dict( ++ self, ++ lora_state_dict: dict[str, torch.Tensor], ++ lora_nickname: str, ++ adapter_lora_alpha: int | None = None, ++ ) -> None: ++ """Map adapter keys to SGLang DiT names and stage them on device.""" + if lora_nickname in self.lora_adapters: + self.lora_adapters[lora_nickname].clear() + +@@ -764,9 +779,91 @@ class LoRAPipeline(ComposedPipelineBase): + f"Dit target weight name {target_name} already exists in lora_adapters[{lora_nickname}]" + ) + self.lora_adapters[lora_nickname][target_name] = weight.to(self.device) +- self.loaded_adapter_paths[lora_nickname] = lora_path + self.loaded_adapter_alphas[lora_nickname] = adapter_lora_alpha +- logger.info("Rank %d: loaded LoRA adapter %s", rank, lora_path) ++ ++ def set_lora_from_tensors( ++ self, ++ lora_nickname: str, ++ named_tensors: dict[str, torch.Tensor], ++ target: str = "transformer", ++ strength: float = 1.0, ++ merge_mode: str | None = None, ++ ) -> int: ++ """Apply a freshly supplied in-memory LoRA adapter.""" ++ if target not in self.VALID_TARGETS: ++ raise ValueError( ++ f"Invalid target: {target}. Valid targets: {self.VALID_TARGETS}" ++ ) ++ ++ if not self.lora_initialized: ++ with self._temporarily_disable_offload( ++ target="all", use_module_names_only=True ++ ): ++ self.convert_to_lora_layers() ++ ++ lora_state_dict = normalize_lora_state_dict( ++ dict(named_tensors), logger=logger ++ ) ++ self._ingest_lora_state_dict(lora_state_dict, lora_nickname) ++ tensor_path = f"" ++ self.loaded_adapter_paths[lora_nickname] = tensor_path ++ ++ target_modules, error = self._get_target_lora_layers(target) ++ if error: ++ raise ValueError(f"set_lora_from_tensors: {error}") ++ ++ merge_mode = self._resolve_lora_merge_mode(None, merge_mode) ++ merge_weights_by_module = { ++ module_name: self._should_merge_lora_for_layers( ++ module_name, lora_layers_dict, merge_mode ++ ) ++ for module_name, lora_layers_dict in target_modules ++ } ++ ++ if self._needs_lora_weight_update_context( ++ target_modules, merge_weights_by_module ++ ): ++ weight_update_context = self._temporarily_disable_offload( ++ target_modules=target_modules ++ ) ++ else: ++ weight_update_context = nullcontext() ++ ++ adapted_count = 0 ++ rank = dist.get_rank() ++ with weight_update_context: ++ for module_name, lora_layers_dict in target_modules: ++ effective_merge_weights = merge_weights_by_module[module_name] ++ count = self._apply_lora_to_layers( ++ lora_layers_dict, ++ [lora_nickname], ++ [tensor_path], ++ rank, ++ [strength], ++ clear_existing=True, ++ merge_weights=effective_merge_weights, ++ ) ++ adapted_count += count ++ self.cur_adapter_name[module_name] = lora_nickname ++ self.cur_adapter_path[module_name] = tensor_path ++ self.is_lora_merged[module_name] = effective_merge_weights ++ self.cur_adapter_strength[module_name] = strength ++ self.cur_adapter_config[module_name] = ( ++ [lora_nickname], ++ [strength], ++ ) ++ ++ logger.info( ++ "Rank %d: in-memory LoRA adapter %s applied to %d layers " ++ "(target=%s, strength=%.2f, merge_mode=%s)", ++ rank, ++ lora_nickname, ++ adapted_count, ++ target, ++ strength, ++ merge_mode, ++ ) ++ return adapted_count + + def set_lora( + self, +diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +index bc1cb59..4a74aa7 100644 +--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py ++++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +@@ -6,6 +6,7 @@ Denoising stage for diffusion pipelines. + """ + + import gc ++import hashlib + import inspect + import math + import time +@@ -135,6 +136,30 @@ from sglang.multimodal_gen.runtime.utils.torch_compile import ( + from sglang.multimodal_gen.utils import dict_to_3d_list + + logger = init_logger(__name__) ++_MAX_TORCH_SEED = (1 << 63) - 1 ++ ++ ++def _resolve_base_seed(batch) -> int | None: ++ seed = getattr(batch, "seed", None) ++ if seed is None: ++ seed = getattr(getattr(batch, "sampling_params", None), "seed", None) ++ return int(seed) if seed is not None else None ++ ++ ++def _make_step_generators( ++ base_seed: int, ++ step_index: int, ++ denoise_seeds: list[str], ++) -> list[torch.Generator]: ++ generators = [] ++ for seed_key in denoise_seeds: ++ payload = f"{int(base_seed)}::step::{int(step_index)}::sample::{str(seed_key)}".encode("utf-8") ++ digest = hashlib.blake2b(payload, digest_size=8).digest() ++ seed = int.from_bytes(digest, byteorder="big", signed=False) % _MAX_TORCH_SEED ++ generator = torch.Generator(device="cpu") ++ generator.manual_seed(seed) ++ generators.append(generator) ++ return generators + + + def _ensure_tensor_model_output(model_output): +@@ -1229,6 +1254,15 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin): + below; mirror them in the override if those markers are needed. + """ + use_nvtx = self.current_use_nvtx ++ denoise_seeds = getattr(batch, "denoise_seeds", None) ++ base_seed = _resolve_base_seed(batch) ++ if getattr(batch, "rollout", False) and denoise_seeds is not None and base_seed is not None: ++ ctx.extra_step_kwargs["generator"] = _make_step_generators( ++ base_seed, ++ int(step.step_index), ++ [str(seed) for seed in denoise_seeds], ++ ) ++ + # 1. Prepare latent inputs in the model's compute dtype. + latent_model_input = ctx.latents.to(ctx.target_dtype) + if batch.image_latent is not None: +diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +index f9a28a2..40c705f 100644 +--- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py ++++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +@@ -5,6 +5,7 @@ + Latent preparation stage for diffusion pipelines. + """ + ++import hashlib + from dataclasses import dataclass + from typing import Any + +@@ -26,6 +27,48 @@ from sglang.multimodal_gen.runtime.server_args import ServerArgs + from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + + logger = init_logger(__name__) ++_MAX_TORCH_SEED = (1 << 63) - 1 ++ ++ ++def _derive_group_seed(base_seed: int, group_id: str) -> int: ++ payload = f"{int(base_seed)}::{str(group_id)}".encode("utf-8") ++ digest = hashlib.blake2b(payload, digest_size=8).digest() ++ return int.from_bytes(digest, byteorder="big", signed=False) % (_MAX_TORCH_SEED + 1) ++ ++ ++def _regen_initial_noise( ++ noise_group_ids: list[str], ++ base_seed: int, ++ latent_shape: list[int] | tuple[int, ...], ++ device, ++ dtype, ++): ++ tensors = [] ++ cache = {} ++ for raw_group_id in noise_group_ids: ++ group_id = str(raw_group_id) ++ noise = cache.get(group_id) ++ if noise is None: ++ generator = torch.Generator(device=torch.device("cpu")) ++ generator.manual_seed(_derive_group_seed(base_seed, group_id)) ++ noise = torch.randn( ++ tuple(int(v) for v in latent_shape), ++ generator=generator, ++ device="cpu", ++ dtype=torch.float32, ++ ) ++ cache[group_id] = noise ++ tensors.append(noise) ++ return torch.stack(tensors, dim=0).to(device=device, dtype=dtype) ++ ++ ++def _driver_xt_recipe(batch): ++ group_ids = getattr(batch, "initial_noise_group_ids", None) ++ latent_shape = getattr(batch, "initial_noise_latent_shape", None) ++ base_seed = getattr(batch, "initial_noise_seed", None) ++ if group_ids and latent_shape is not None and base_seed is not None: ++ return [str(v) for v in group_ids], [int(v) for v in latent_shape], int(base_seed) ++ return None + + + @dataclass(frozen=True) +@@ -127,6 +170,7 @@ class LatentPreparationStage(PipelineStage): + latents = batch.latents + height = batch.height + width = batch.width ++ driver_xt_recipe = _driver_xt_recipe(batch) + + # TODO(will): remove this once we add input/output validation for stages + if self.requires_batch_height_width(batch, server_args) and ( +@@ -142,7 +186,34 @@ class LatentPreparationStage(PipelineStage): + ) + + # Generate or use provided latents +- if latents is None: ++ if driver_xt_recipe is not None: ++ group_ids, latent_shape, base_seed = driver_xt_recipe ++ spec = self.get_latent_preparation_spec( ++ batch, server_args, len(group_ids), latent_num_frames, device ++ ) ++ raw_latents = _regen_initial_noise( ++ group_ids, ++ base_seed, ++ latent_shape, ++ device=spec.device, ++ dtype=spec.dtype, ++ ) ++ ++ latent_ids = ( ++ server_args.pipeline_config.maybe_prepare_latent_ids(raw_latents) ++ if spec.prepare_latent_ids ++ else None ++ ) ++ ++ if latent_ids is not None: ++ batch.latent_ids = latent_ids.to(device=device) ++ ++ latents = raw_latents ++ if spec.pack_latents: ++ latents = server_args.pipeline_config.maybe_pack_latents( ++ raw_latents, len(group_ids), batch ++ ) ++ elif latents is None: + spec = self.get_latent_preparation_spec( + batch, server_args, batch_size, latent_num_frames, device + ) +@@ -202,7 +273,8 @@ class LatentPreparationStage(PipelineStage): + indexed_batches = group + group_batches = [batch for _, batch in indexed_batches] + if len(group_batches) == 1 or any( +- batch.latents is not None for batch in group_batches ++ batch.latents is not None or _driver_xt_recipe(batch) is not None ++ for batch in group_batches + ): + for index, batch in indexed_batches: + results[index] = self(batch, server_args) +diff --git a/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py +index f652059..d2e6eb7 100644 +--- a/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py ++++ b/python/sglang/multimodal_gen/runtime/post_training/gpu_worker_post_training_mixin.py +@@ -7,6 +7,7 @@ from sglang.multimodal_gen.runtime.loader.weight_utils import compute_weights_ch + from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload import ( + iter_materialized_weights, + ) ++from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch + from sglang.multimodal_gen.runtime.post_training.tensor_update_checker import ( + TensorUpdateChecker, + ) +@@ -19,12 +20,58 @@ from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions + + if TYPE_CHECKING: + from sglang.multimodal_gen.runtime.entrypoints.post_training.io_struct import ( ++ SetLoraFromTensorReqInput, + UpdateWeightFromTensorCheckerReqInput, + UpdateWeightFromTensorReqInput, + ) + + + class GPUWorkerPostTrainingMixin: ++ def set_lora_from_tensors( ++ self, ++ req: SetLoraFromTensorReqInput, ++ ) -> OutputBatch: ++ """Install a LoRA adapter from a torch.save blob of tensor weights.""" ++ import io ++ ++ import torch ++ ++ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline ++ ++ if not isinstance(self.pipeline, LoRAPipeline): ++ return OutputBatch(error="Lora is not enabled") ++ ++ try: ++ named_tensors = torch.load( ++ io.BytesIO(req.serialized_tensors), ++ map_location="cpu", ++ weights_only=True, ++ ) ++ adapted = self.pipeline.set_lora_from_tensors( ++ req.lora_name, ++ named_tensors, ++ target=req.target, ++ strength=req.strength, ++ merge_mode=req.merge_mode, ++ ) ++ if adapted <= 0: ++ return OutputBatch( ++ error=f"LoRA adapter {req.lora_name} did not match any layer" ++ ) ++ except Exception as e: ++ return OutputBatch(error=str(e)) ++ ++ return OutputBatch( ++ output={ ++ "success": True, ++ "message": ( ++ f"Applied in-memory LoRA {req.lora_name} to " ++ f"{adapted} layers of {req.target}" ++ ), ++ "adapted_layers": adapted, ++ } ++ ) ++ + def update_weights_from_disk( + self, + model_path: str, +@@ -135,7 +182,7 @@ class GPUWorkerPostTrainingMixin: + + def get_weights_checksum( + self, module_names: list[str] | None = None +- ) -> dict[str, str]: ++ ) -> dict[str, str | int]: + if not self.pipeline: + return {"error": "Pipeline is not initialized"} + +@@ -143,15 +190,16 @@ class GPUWorkerPostTrainingMixin: + names = module_names if module_names is not None else list(all_modules.keys()) + + checksums: dict[str, str] = {} ++ tensor_count = 0 + for name in names: + module = all_modules.get(name) + if module is None: + checksums[name] = "not_found" + continue +- checksums[name] = compute_weights_checksum( +- iter_materialized_weights(module) +- ) +- return checksums ++ materialized_weights = list(iter_materialized_weights(module)) ++ tensor_count += len(materialized_weights) ++ checksums[name] = compute_weights_checksum(materialized_weights) ++ return {**checksums, "tensor_count": tensor_count} + + def _select_rank_scoped_payload( + self, +diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py +index 025643b..8dfe997 100644 +--- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py ++++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_post_training_mixin.py +@@ -6,6 +6,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa + + + class SchedulerPostTrainingMixin: ++ def _handle_set_lora_from_tensor(self, reqs: List[Any]) -> OutputBatch: ++ req = reqs[0] ++ return self.worker.set_lora_from_tensors(req) ++ + def _handle_update_weights_from_disk(self, reqs: List[Any]) -> OutputBatch: + req = reqs[0] + success, message = self.worker.update_weights_from_disk( +diff --git a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py +index 3777ad4..c4e91ea 100644 +--- a/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py ++++ b/python/sglang/multimodal_gen/runtime/post_training/scheduler_rl_mixin.py +@@ -1,7 +1,9 @@ + # SPDX-License-Identifier: Apache-2.0 + """Flow-matching rollout step utilities for log-prob computation.""" + ++import hashlib + import math ++import os + from typing import Any, Union + + import torch +@@ -18,6 +20,25 @@ from sglang.multimodal_gen.runtime.post_training.scheduler_rl_debug_mixin import + ) + + _LOG_SQRT_2PI = math.log(math.sqrt(2 * math.pi)) ++_MAX_TORCH_SEED = (1 << 63) - 1 ++ ++ ++def _resolve_base_seed(batch) -> int | None: ++ seed = getattr(batch, "seed", None) ++ if seed is None: ++ seed = getattr(getattr(batch, "sampling_params", None), "seed", None) ++ return int(seed) if seed is not None else None ++ ++ ++def _resolve_fallback_seed(batch) -> int: ++ base_seed = _resolve_base_seed(batch) ++ denoise_seeds = getattr(batch, "denoise_seeds", None) ++ sample_key = str(denoise_seeds[0]) if denoise_seeds else None ++ if base_seed is not None and sample_key is not None: ++ payload = f"{int(base_seed)}::fallback::sample::{sample_key}".encode("utf-8") ++ digest = hashlib.blake2b(payload, digest_size=8).digest() ++ return int.from_bytes(digest, byteorder="big", signed=False) % _MAX_TORCH_SEED ++ return int.from_bytes(os.urandom(8), byteorder="big") % _MAX_TORCH_SEED + + + class SchedulerRLMixin(SchedulerRLDebugMixin): +@@ -86,7 +107,15 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): + B = local_shape[0] + if isinstance(generator, torch.Generator): + assert B == 1, "Generator must be a list if batch size is not 1" +- generator = [generator] ++ per_request_generator = getattr(batch, "_relax_noise_gen", None) ++ if per_request_generator is None: ++ per_request_generator = torch.Generator(device=device) ++ per_request_generator.manual_seed(_resolve_fallback_seed(batch)) ++ try: ++ batch._relax_noise_gen = per_request_generator ++ except AttributeError: ++ pass ++ generator = [per_request_generator] + else: + assert ( + len(generator) == B +@@ -96,11 +125,14 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): + rollout_session_data, rollout_session_data.latents_shape, device, dtype + ) + for i in range(B): +- torch.randn( +- rollout_session_data.latents_shape, +- out=buffer[i : i + 1], +- generator=generator[i], +- ) ++ gen = generator[i] ++ gen_device = getattr(gen, "device", None) ++ sample_shape = tuple(buffer[i].shape) ++ if gen is not None and gen_device is not None and gen_device.type != buffer.device.type: ++ tmp = torch.randn(sample_shape, generator=gen, dtype=dtype, device=gen_device) ++ buffer[i].copy_(tmp) ++ else: ++ torch.randn(sample_shape, out=buffer[i], generator=gen) + + sharded_noise, _ = rollout_session_data.pipeline_config.shard_latents_for_sp( + batch=batch, latents=buffer +@@ -128,6 +160,7 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): + 1. ``"sde"``: Standard stochastic differential equation transition (Gaussian). + 2. ``"cps"``: Coupled Particle Sampling. + 3. ``"ode"``: Deterministic ODE step (no diffusion noise). ++ 4. ``"dance"``: FlowGRPO dance-style transition used by Relax diffusion RL. + """ + rollout_session_data = self._get_rollout_session_data(batch) + sde_type = batch.rollout_sde_type +@@ -234,6 +267,34 @@ class SchedulerRLMixin(SchedulerRLDebugMixin): + log_prob_no_const + ), "p_ode is always 0, true log_prob is meaningless, set rollout_log_prob_no_const to True to enable log_prob computation" + ++ elif effective_sde_type == "dance": ++ model_output = model_output.float() ++ sample = sample.float() ++ variance_noise = self._rollout_variance_noise( ++ batch, model_output, generator ++ ) ++ full_variance_noise = rollout_session_data.noise_buffer ++ std_dev_t = current_sigma.new_tensor(noise_level) ++ noise_std_dev = std_dev_t * torch.sqrt(-1 * dt) ++ prev_sample_mean = ( ++ sample * (1 + std_dev_t**2 / (2 * current_sigma) * dt) ++ + model_output ++ * (1 + std_dev_t**2 * (1 - current_sigma) / (2 * current_sigma)) ++ * dt ++ ) ++ prev_sample = prev_sample_mean + variance_noise * noise_std_dev ++ if tuple(variance_noise.shape) != tuple(full_variance_noise.shape): ++ raise RuntimeError( ++ "rollout_sde_type='dance' does not support sequence-parallel rollout: " ++ f"the per-step log-prob is computed over the unsharded noise " ++ f"{tuple(full_variance_noise.shape)} while this rank only owns " ++ f"{tuple(variance_noise.shape)}, so the cross-rank reduction would " ++ "over-count it by the SP degree. Run the diffusion engine with " ++ "--rollout-num-gpus-per-engine 1 until the reduction convention " ++ "is confirmed against the live engine." ++ ) ++ log_prob_no_const_val = -((full_variance_noise * noise_std_dev) ** 2) ++ + else: + raise ValueError(f"Unsupported sde_type: {sde_type}") + +diff --git a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +index 7c470d0..2b50e0e 100644 +--- a/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py ++++ b/python/sglang/multimodal_gen/runtime/post_training/weights_updater.py +@@ -493,6 +493,13 @@ class WeightsUpdater: + lora_rank=lora_rank, + ) + ++ if getattr(self.pipeline, "lora_initialized", False): ++ return False, ( ++ "Refusing update_weights_from_tensor: this pipeline has LoRA " ++ "layers installed. Use adapter-only sync (/set_lora_from_tensor) " ++ "for a LoRA run." ++ ) ++ + if target_modules is None: + target_modules = [_DEFAULT_TENSOR_TARGET_MODULE] + try: diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 240e20de..eda3406c 100644 --- a/python/sglang/srt/disaggregation/decode.py diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts index c1afd6b95..664e780c1 100644 --- a/docs/.vitepress/config.mts +++ b/docs/.vitepress/config.mts @@ -277,7 +277,8 @@ export default defineConfig({ { text: 'Metrics Service', link: '/en/guide/metrics-service-detailed' }, { text: 'Notification System', link: '/en/guide/notification-system' }, { text: 'Update Weights Pipeline', link: '/en/guide/update-weights-pipeline' }, - { text: 'Low-Rank Adaptation (LoRA) Training', link: '/en/guide/low-rank-adaptation-training' } + { text: 'Low-Rank Adaptation (LoRA) Training', link: '/en/guide/low-rank-adaptation-training' }, + { text: 'Diffusion Generative RL', link: '/en/guide/diffusion-generative-rl' } ] }, { @@ -393,7 +394,8 @@ export default defineConfig({ { text: 'Metrics 服务', link: '/zh/guide/metrics-service-detailed' }, { text: '通知系统', link: '/zh/guide/notification-system' }, { text: '权重更新流水线优化', link: '/zh/guide/update-weights-pipeline' }, - { text: '低秩适配（LoRA）训练', link: '/zh/guide/low-rank-adaptation-training' } + { text: '低秩适配（LoRA）训练', link: '/zh/guide/low-rank-adaptation-training' }, + { text: '扩散生成式 RL', link: '/zh/guide/diffusion-generative-rl' } ] }, { diff --git a/docs/draft/multimodal-gen-rl-design.md b/docs/draft/multimodal-gen-rl-design.md new file mode 100644 index 000000000..9237d0fe6 --- /dev/null +++ b/docs/draft/multimodal-gen-rl-design.md @@ -0,0 +1,1167 @@ +# 多模生成 RL 适配设计 + +> **当前发布范围：Qwen-Image T2I。** FlowGRPO + FSDP2（全参数或 LoRA）+ SGLang Native Diffusion + 同步 colocate。已验证配方是 LoRA adapter-sync：100 轮完整跑通并对齐 Qwen-Image LoRA 参考曲线（最终 eval 0.8643，对照 0.8620）。 +> +> **设计上的多模型范围（Qwen-Image-Edit I2I、WAN 2.2 T2V/I2V/V2V、LTX-2.3 T2AV）已下线。** 完整代码只保留在**本地** git 分支 `backup/diffusion-generative-rl-full`（commit `f3e7e9e2`、`f7e30890`）。 +> +> ::: danger 备份分支尚未推送 +> 该分支和这两个 commit **只存在于本地仓库，`origin` 上没有**，所以对任何其他读者来说这个指针目前是无效的。若要让这条扩展参考真正可用，必须先 `git push origin backup/diffusion-generative-rl-full`；在那之前，请把本文中所有「已退役到备份分支」理解为「代码已从工作分支删除，只在某台机器的本地分支上」。 +> ::: +> +> 设计基线：Relax `e59cd7288995b13d6fac8cfbae2b638e45fb29b7`，SGLang `4ad418d2c3d43cb3c699bc9419d32673b1fca7d8`。初版日期 2026-07-16；本次按仓库实现回写，日期 2026-08-10。 +> +> 阅读约定：本文既是设计文档也是实现记录。**「已实现」**表示工作分支上有对应代码；**「已退役」**表示代码只在上述本地备份分支；**「设计保留、未实现」**表示只有设计、没有代码。 + +## 0. 现状总结 + +### 0.1 实际发布的东西 + +| 维度 | 现状 | +|---|---| +| 任务 | 只有 `t2i`（Qwen-Image）。`GENERATION_TASKS = ("t2i",)` —— 词表已随代码收敛，自带 adapter 的调用方必须同时把任务名加进来 | +| 算法 | FlowGRPO（组内中心化 advantage + clipped policy loss），`--advantage-estimator grpo` | +| 训练后端 | `--train-backend fsdp` → `FSDPTrainRayActor`，FSDP2，纯 DP（无 TP/PP/CP）；`--fsdp-trainable-mode` 支持 `full` 与 `lora` | +| Rollout 后端 | `SGLangNativeGenerationEngine`（SGLang 主线 native diffusion + Relax 静态 patch 契约检查） | +| Reward | PickScore（`relax/engine/rewards/pickscore.py`） | +| 部署模式 | 仅同步 colocate（`--colocate`），fully-async / hybrid 在预检直接拒绝 | +| 权重同步 | 每步同步。full/merge：完整 transformer，DTensor `full_tensor()` → 分桶 → CUDA-IPC → engine commit，带 count/bytes/checksum 校验与版本回滚；adapter 模式只发 LoRA 张量 | +| Checkpoint | DCP sharded（LoRA 只存 adapter）+ `COMMITTED` 协议 + 保留轮转 + 离线 HF/diffusers 与 HF-PEFT 导出 | +| 验证入口 | `scripts/training/diffusion/run-qwen-image-t2i-lora-8xgpu.sh`（默认 adapter sync，唯一复现参考曲线的配方）；`run-qwen-image-t2i-8xgpu.sh` 是全参参考 | + +### 0.2 已退役到 `backup/diffusion-generative-rl-full` 的东西 + +```text +relax/models/wan_video/{__init__.py,adapter.py} # WAN 2.2 T2V/I2V/V2V +relax/models/ltx_av/{__init__.py,adapter.py} # LTX-2.3 T2AV +relax/engine/rewards/editreward.py # EditReward (I2I) +relax/engine/rewards/video_reward.py # VideoPickScore / condition consistency / VideoCLIPDelta +relax/engine/rewards/t2av_reward.py # CLAP + 复合 reward +examples/diffusion/qwen_image_edit/i2i_full.yaml +examples/diffusion/wan22/{t2v,i2v,v2v}_full.yaml +examples/diffusion/ltx23/t2av_full.yaml +tests/models/{wan_video,ltx_av}/test_adapter.py +tests/engine/rewards/test_{editreward,video_reward,t2av_reward}.py +``` + +退役理由：这些路径没有 GPU 端到端验证，保留在主干上会让「配置得出来但跑不通」的表面积远大于实际可用面。 + +**这次退役不是纯删文件，抽象也一起收敛了**（与初版设计的最大差异，见 §4.1）：多 track（video + audio）支撑已经删除 —— `policy_tracks`、`track_weights`、`combine_track_logp`、`TrackLogp` 都不在了，`replay_transition` 直接返回 `torch.Tensor`。适配器协议同时去掉了 `encode_conditions` 与 `freeze_non_policy_modules`（冻结在 `load_train_model` 里一次做完）。因此从备份分支回接 T2AV 需要重新引入多 track 的 log-prob 合成，不再是「加一个 adapter 就行」。 + +`QwenImageAdapter.supported_tasks` 现在也只有 `("t2i",)`：i2i 分支曾经被声明为支持，但 `replay_transition` 会静默忽略源图，也就是在训练一个无条件速度场 —— 这是把词表和实际 replay 路径对齐的直接原因。 + +### 0.3 与初版设计相比的语义变化 + +1. **`--num-updates-per-batch` 的语义是切分，不是复用**（§6.4）。`_plan_micro_batch_updates` 把本 rank 的训练切片划成 N 份互不相交、样本数相等的更新，每个样本只参与一次 optimizer step。因此真正的 step 边界是 `global_batch_size / num_updates_per_batch`。初版设计里「每个 prompt group 一次 step」和中间那版「`--global-batch-size` 才是 step 边界」都已不成立，`--fsdp-optimizer-step-per-group` 也已删除。 +2. **新增 LoRA**（§10.1）：`--fsdp-trainable-mode lora` + 共享的 `--lora-*` 参数，merge / adapter 两条 rollout 同步路径。 +3. **新增按 optimizer step 计数的 LR schedule**（§10.1）：`--fsdp-lr-scheduler {constant,linear,cosine}`，默认 `constant`。 +4. **`sde_indices` 可从 `num_sde_steps` + `sde_timestep_fraction` 推导，且默认被每轮重抽样覆盖**（§5.5 / §13）。 +5. **评测、checkpoint 轮转、权重同步校验、artifact 保留、profiler 从「计划中」变成「已实现」**。 +6. **不支持的 flag 从静默无效变成预检报错**（§13.1）。 +7. **`sde_type="sde"` 下训练 SDE step 0 被硬拒绝**（§6.1 / §13.1）。 + +### 0.4 已知限制（诚实清单） + +| 限制 | 说明 | +|---|---| +| 只有同步 colocate | `fully_async` / `hybrid` / 非 colocate 在 `validate_generative_config` 被拒；`update_weights_fully_async` 抛 `NotImplementedError` | +| 训练侧纯 DP | `_get_parallel_config` 固定 `tp_size=1, pp_size=1`；rollout 侧 TP 由 `--rollout-num-gpus-per-engine` 支持 | +| replay 的 DP 切分不完整 | rank 0 读整个 TQ partition 再 broadcast 数值行（`_read_train_partition` 里的 `TODO(agent)`）。组内候选按 `[dp_rank::dp_world]` 切分（`hydrate_micro_batches`），但 `group_size % dp_world != 0` 时退化为每 rank 全量重放（结果正确，算力冗余） | +| 无训练数据 dump / debug replay | `--save-debug-train-data` 只告警；`--load-debug-rollout-data` 直接拒绝。TQ 行是数值索引 + 磁盘 sidecar，不是 token 序列 | +| 无 MFU / TFLOPs | 共享 `FlopsCounter` 是 LLM 架构专用的，DiT 没有解析模型；改用 transitions/samples per second（§14） | +| 不支持 CFG | `guidance_scale != 1.0` 在 `build_rollout_request` 直接报错：actor 只回放一次正向条件前向，双前向 CFG replay 未实现 | +| 依赖 SGLang 静态 patch | `docker/patch/sglang/v0.5.15.post1.patch` 必须应用（内存态权重同步、`/set_lora_from_tensor`、offload 端点，以及 rollout/replay 一致性改动）；engine 启动前做只读契约检查（§9） | +| KL-to-ref 未实现 | 无 reference model、无 KL 项；`kl_coef` / `use_kl_loss` 非零直接报错 | +| resume 需要显式 `--load` | `_maybe_resume` 读的是 `--load` 而不是 `--save`，两个启动脚本都没设置它 | +| reward 只有进程内或 remote | `post_process` 跑在 rollout worker 里，没有 controller 创建的 placement group。`colocate` / `cpu` 都是进程内评分（区别只在允不允许用 CUDA），`remote` 走 HTTP；旧的 `dedicated` 只是 `colocate` 的同义词并额外打一条假告警，已删除 | + +## 1. 目标方案 + +### 1.1 任务矩阵 + +| 任务 ID | 输入 | 输出 | 模型 | Rollout | 在线 reward | 状态 | +|---|---|---|---|---|---|---| +| `t2i` | text | image | `Qwen/Qwen-Image` | SGLang Native Diffusion | PickScore | **已实现并验证**（`GENERATION_TASKS` 中唯一项） | +| `i2i` / `t2v` / `i2v` / `v2v` / `t2av` | — | — | Qwen-Image-Edit、WAN 2.2、LTX-2.3 | — | EditReward / VideoPickScore / CLAP 等 | 已退役（见 §0.2 与 §4.3）。任务 ID 也已从 `GENERATION_TASKS` 移除 | + +每个训练任务独立启动，一个进程组只持有一个模型族和一个任务配置。不做跨模型族的混合 batch 或联合 optimizer。 + +### 1.2 统一训练口径 + +1. 训练完整 diffusion/flow transformer；可选用 LoRA 只训练注入的 adapter（`--fsdp-trainable-mode lora`）。 +2. VAE、text encoder、image encoder 保持冻结，且不在 FSDP actor 的模块里 —— 它们只存在于 rollout engine。 +3. 训练后端固定为 PyTorch FSDP2，参数与 optimizer 使用 DCP sharded checkpoint（LoRA 只存 adapter）。 +4. 算法固定为 FlowGRPO，rollout 保存被选中的 SDE transition，actor replay 产生 `old_logp`。 +5. actor 与 rollout 使用同步 colocate。fully async 只保留 full-weight transport 的接口形状，未实现。 +6. GPU 数量由任务 profile 决定，没有全局 8 卡上限。跨节点 FSDP 与多 GPU rollout engine 都是合法配置。 +7. SGLang 主仓库的 native diffusion 是统一推理后端；SGLang-Omni 不进入这些 diffusion 任务的运行路径。 + +### 1.3 架构边界 + +运行时代码按 Relax 已有 ownership 放置，不创建 `relax/diffusion/` 顶层领域包： + +- FSDP 状态、checkpoint 和完整权重导出进入 `relax/backends/fsdp/`。 +- SGLang DiffGenerator 生命周期进入 `relax/backends/sglang/`。 +- 通用生成模型协议和 FlowGRPO 数学进入 `relax/models/`。 +- 各模型族的 condition、trajectory geometry 和参数边界进入各自模型目录（现在只剩 `relax/models/qwen_image/`）。 +- rollout 编排进入 `relax/engine/rollout/`。 +- scorer 实现进入 `relax/engine/rewards/`，Ray 资源生命周期进入 `relax/distributed/ray/`。 + +控制面继续使用现有 `Controller`、`Actor`、`Rollout`、`RolloutManager` 和 `grpo` 注册键。数据面继续使用现有 `RolloutDataSource`、`Sample`、`TransferQueue` 与 `GRPOGroupNSampler`。 + +## 2. Relax 能力复用 + +| 现有能力 | 使用方式 | 落地情况 | +|---|---|---| +| `Sample.multimodal_inputs` | 承载条件 image/video | 已实现（t2i 不使用） | +| `MultimodalTypes` | 复用 image/video/audio 类型与 placeholder | 生成输出走 artifact manifest，未新增 Sample 字段 | +| `RolloutDataSource` | 读取统一 JSONL，按 `n_samples_per_prompt` 展开 group | 已实现；额外改动：`hf_checkpoint` 为空时不加载 tokenizer/processor | +| `_shallow_copy_sample()` | 同组候选共享只读条件媒体 | 复用现状 | +| `--multimodal-keys` | 映射数据列到 image/video/audio | 由任务 YAML 配置；t2i 为 `null` | +| `RolloutManager` | 生命周期、健康检查、offload/onload、engine recovery | 已实现，engine class 由 `_resolve_rollout_engine_class` dotpath 解析 | +| `--rollout-function-path` | 指向统一 native generation driver | 已实现（`generate_rollout`，`evaluation=True` 分流到 `evaluate_rollout`） | +| `custom_reward_post_process_path` | 完整 group 生成后批量评分 | 已实现（`relax.engine.rewards.generative.post_process`） | +| `custom_convert_samples_to_train_data_path` | 生成轻量 diffusion TQ row | 已实现（`relax/utils/utils.py` 执行该 hook） | +| `TransferQueue` | 每个候选一条数值 row，partition 仍是训练屏障 | 已实现；大 trajectory 走磁盘 sidecar | +| `GRPOGroupNSampler` | 保证同 prompt 候选组完整分配 | 复用现状 | +| 共享 `--lora-*` flag | rank / alpha / dropout / target modules / merge-adapter 模式 | 已实现；注入与同步是 diffusers 形态的独立实现（`backends/fsdp/lora.py`），只在导出时复用 `megatron_peft_utils.write_hf_peft_adapter` 一类与后端无关的 helper | +| `UpdateWeightFromTensor` | pause、分桶、Ray IPC、engine fan-out、版本校验 | FSDP actor 独立实现 `_run_weight_transaction`，mirror 其 per-rank→per-engine 映射（不复用 Megatron iterator） | +| DCS coordinator | async 分支的 topology 与版本协调 | **设计保留、未实现**：`weight_update.py` 没有 FSDP comm backend factory，`checkpoint_service/client/engine.py` 未改动 | +| `train_dump_utils` | rollout JSONL 与 debug dump | **设计保留、未实现**：该文件未改动，未透传 artifact/trajectory manifest 摘要 | +| DCP/轮转目录约定 | `iter_0000001`、latest marker 与保留策略 | 已实现；`rotate_ckpt` 新增 `save_dir` 参数以匹配 `//iter_*` 布局 | +| `TrainProfiler` | torch / memory profiler | 已实现（`--use-pytorch-profiler`、`--record-memory-history` 等对 FSDP actor 生效） | + +以下核心对象保持原样（未因本方案改动其逻辑）： + +```text +relax/core/registry.py +relax/core/controller.py +relax/core/service.py +relax/components/actor.py +relax/components/rollout.py +relax/utils/types.py +relax/utils/multimodal/* +``` + +同步 colocate 使用 `ROLES_COLOCATE`。FlowGRPO advantage 在 rollout 侧的 reward post-process 内计算、replay 与 loss 在 FSDP actor 内完成，不启动独立 Advantages Serve。 + +> **修正**：初版写「不新增或修改 `relax/components/advantages.py`」。实际做了一处**非语义**改动 —— 把 `from megatron.core import mpu` 改成函数内延迟导入，使控制面在无 megatron 的 diffusion 镜像里可以 import。同类改动还有 `relax/engine/sft/eval/runner.py`、`relax/utils/data/stream_dataloader.py`、`relax/utils/rocm_checkpoint_writer.py`、`relax/distributed/checkpoint_service/{utils.py,backends/device_direct.py}`、`relax/models/__init__.py`（Qwen-Omni import 容错）。这些都是「megatron 变可选依赖」的兼容改动，不改变 token RL 行为。 + +## 3. 总体架构 + +```mermaid +flowchart LR + C["Controller / grpo"] --> A["components.Actor"] + C --> R["components.Rollout"] + A --> F["FSDPTrainRayActor"] + R --> M["existing RolloutManager"] + M --> E["SGLangNativeGenerationEngine"] + M --> D["existing RolloutDataSource"] + D --> S["standard Sample groups"] + E --> Q["QwenImageAdapter"] + E --> T["trajectory sidecars"] + E --> O["image artifacts"] + O --> RM["generative reward post_process"] + RM --> P["PickScoreScorer"] + RM --> X["existing TransferQueue"] + X --> F + F --> U["full-weight bucket sync (CUDA IPC)"] + U --> E +``` + +> 备份分支上同一张图还有 `WAN adapter` / `LTX adapter` 两个 engine 下游分支和 `EditReward/Video/CLAP` 三类 scorer；现在只剩 Qwen 一条。 + +### 3.1 Colocate 资源模型 + +actor 与 rollout 共享同一 placement group。每个任务满足： + +```text +actor_world_size == rollout_total_gpus +rollout_total_gpus % rollout_num_gpus_per_engine == 0 +num_rollout_engines = rollout_total_gpus / rollout_num_gpus_per_engine +``` + +Reward 只有两种落点：`post_process` 所在的 rollout worker 进程内（`colocate` 允许用 CUDA，`cpu` 禁用 CUDA），或外部 HTTP 服务（`remote`）。没有独立的 reward placement group —— `GenerativeRewardManager` 现在就是 remote 的 HTTP 代理，`LocalGenerativeRewardManager` 是进程内实现。两个启动脚本都用 `--reward-runtime colocate`：PickScore 约 5.6 GB，而它运行时 actor 已经 offload。 + +任务 profile 只提供容量规划起点，不构成固定上限： + +| Profile | 建议起始 GPU | Rollout TP | Reward | 状态 | +|---|---:|---:|---|---| +| Qwen T2I LoRA | 8 x 96GB | 1 | colocate PickScore | **已验证**：`rollout_batch_size=32`、`n_samples_per_prompt=8`、384x384、12 步、训练 3 个 SDE step、rank 64 adapter | +| Qwen T2I 全参 | 8 x 96GB | 1 | colocate PickScore | 全参参考：`rollout_batch_size=8`、`n_samples_per_prompt=8`、384x384（曲线未复现） | +| Qwen T2I（YAML profile） | 8 x 96GB | 1 | colocate PickScore | `examples/diffusion/qwen_image/t2i_{full,lora}.yaml`：48 prompts x 16 / 24 candidates（未在本环境验证，作为参考片段保留） | +| Qwen Edit / WAN / LTX | — | — | — | 已退役 | + +> **Rollout TP 分片权重同步已实现。** engine 跨 `rollout_num_gpus_per_engine` 张物理卡，与之同卡的 `tp` 个 FSDP rank 组成一个 Gloo gather group，向 engine 发送 `tp` 个 CUDA-IPC blob（worker `i` open physical GPU `base+i` 上的 handle 后内部再分片）—— mirror megatron `UpdateWeightFromTensor` 的 colocate IPC 路径（`_build_ipc_gather_topology` / `_run_weight_transaction`）。拓扑按 CUDA device UUID 匹配，不按 GPU 序号，因为 placement group 会重排。唯一硬约束是 `actor_world % rollout_num_gpus_per_engine == 0`（`validate_generative_config` fail-fast，运行时再校验 `world == sum(engine_gpu_counts)`）。`tp == 1` 退化为每 rank 一个单卡 engine。 + +OOM 时增加 actor/rollout world size、提高 rollout TP、减少 `n_samples_per_prompt`、缩小训练的 SDE step 子集，或改用 LoRA。Validator 不执行模式降级，不把 full-FT 自动切成 LoRA。 + +### 3.2 Owner 状态机 + +```text +TRAIN + -> WEIGHT_SYNC + -> ROLLOUT + -> REWARD + -> TRANSFER_COMMIT + -> TRAIN +``` + +- `TRAIN`：只有 FSDP parameter/optimizer shard 在计算设备。 +- `WEIGHT_SYNC`：FSDP parameter shard 与 SGLang weight receiver 同驻，不运行 forward/backward/generate。此时 engine 只 onload transformer（`tags=[WEIGHTS]`）。 +- `ROLLOUT`：FSDP parameter 和 optimizer offload，SGLang encoder/DiT/VAE 全量 onload。 +- `REWARD`：全部 artifact 已提交；scorer 评分。 +- `TRANSFER_COMMIT`：完整 group 校验后写入 TQ partition，seal 后训练才可消费。 + +实现上 `REWARD` 与 `TRANSFER_COMMIT` 都发生在 `generate_rollout` 尾部的 `transfer_batch_to_data_system` 调用链里（reward post-process → converter → `async_put`），所以 `perf/reward_time` 覆盖的是「评分 + 转换 + TQ 落盘」三段之和。 + +## 4. 任务与模型适配 + +### 4.1 通用适配协议（已实现，`relax/models/generative.py`） + +```python +@runtime_checkable +class GenerativeModelAdapter(Protocol): + family: str + supported_tasks: tuple[str, ...] + + def load_train_model(self, config) -> torch.nn.Module: ... + def build_rollout_request(self, sample, sampling, seed) -> dict: ... + def validate_rollout_response(self, response) -> None: ... + def pack_trajectory(self, response) -> dict[str, torch.Tensor]: ... + def replay_transition(self, model, batch, step_index) -> torch.Tensor: ... + def artifact_tracks(self, response) -> list[ArtifactTrack]: ... + def weight_name_map(self, name: str) -> str: ... +``` + +另有两个**可选属性**，用 `getattr` 读取、刻意不声明为协议成员（这样加属性不会让已有 adapter 失效 —— 协议是 `runtime_checkable` 的，adapter 单测会 assert 这一点）：`lora_target_modules`（`--fsdp-trainable-mode lora` 的默认目标模块，模块名后缀）与 `lora_task_type`（PEFT task type，默认 `FEATURE_EXTRACTION`）。 + +相对初版设计，协议**收窄**了三处： + +- 去掉 `policy_tracks` 与多 track 支撑（见 §0.2 / §6.2）；`replay_transition` 直接返回 `torch.Tensor`。 +- 去掉 `freeze_non_policy_modules`：`load_train_model` 返回的模块会被直接交给 FSDP，所以「不该拿梯度的东西」必须在那里就设好 `requires_grad=False`，没有第二个冻结钩子。 +- 去掉 `encode_conditions`：条件由 engine 在 `denoising_env` 里回传，adapter 不再自己编码。 + +任务配置通过 `--model-adapter-path`（dotpath）选择实现。FSDP runtime、rollout driver、reward barrier 和 weight transport 只依赖该协议。 + +同模块还定义 `ArtifactTrack`、`FullWeightManifest`、`artifact_manifest_dict()`、`ordered_name_shape_hash()` 和 `resolve_sde_indices()`。 + +### 4.2 Qwen-Image T2I（已实现并验证） + +- policy：完整 `QwenImageTransformer2DModel`，`--fsdp-trainable-attr transformer`；LoRA 时只训练八个 attention 投影上的 adapter。 +- frozen：Qwen2.5-VL text encoder、VAE、scheduler —— 它们不在 FSDP actor 的模块里，只存在于 rollout engine。 +- condition：`prompt_embeds` 与 `prompt_embeds_mask`，由 engine 在 `denoising_env` 里回传。 +- output：PNG image track，从 `generated_output` 落盘到 `//rollout_*/group_*_sample_*.png`。 +- trajectory：packed latent（`image_x_t` / `image_x_next`）、`sigmas`、`sde_indices`、`image_grid`、seed hash。`image_grid` 由 engine 回显的 `height` / `width` 算出 —— 打包后的序列长度无法还原非正方形网格，所以这是重建 RoPE 的唯一来源。 +- reward：PickScore，组内中心化后除以该组自己的 std。 +- 硬约束：`guidance_scale` 必须为 `1.0`（`replay_transition` 只跑一次正向条件前向）。 + +### 4.3 已退役的模型族（WAN 2.2 / LTX-2.3 / Qwen-Image-Edit） + +这些设计的完整推理与代码都在备份分支上（见文首的推送提醒），本文不再复述，只留必要的回接提示： + +- **Qwen-Image-Edit I2I**：`QwenImageAdapter` 中的 i2i 分支已删除，`supported_tasks` 收敛为 `("t2i",)`。回接时要同时补上 source image 条件、`denoising_env.image_kwargs` 校验、EditReward scorer 与 aspect bucket 分组。 +- **WAN 2.2 T2V/I2V/V2V**：5D latent `[B,C,T,H,W]`、frame count 与 VAE temporal factor 的约束、MP4 track manifest、VideoPickScore / condition consistency / VideoCLIPDelta。 +- **LTX-2.3 T2AV**：video + audio 双 track、时长同步校验、CLAP 与 VideoPickScore 的加权合成。 + +::: warning 回接成本已经变高 +多 track 的通用支撑**已经删除**（`combine_track_logp`、`TrackLogp`、`track_weights` 都不在了），所以 T2AV 不再是「换个 adapter」就能回来的：需要重新引入按 track 加权合成 log-prob 的那一层。视频/音频任务的 `GENERATION_TASKS` 条目也要一并加回。 +::: + +## 5. 数据契约 + +### 5.1 统一 JSONL + +T2I（唯一在用的形态）： + +```json +{"prompt":"a tram passing through a rainy city at night","metadata":{"task":"t2i","prompt_id":"pickapic_0001"}} +``` + +条件类任务（设计保留）： + +```json +{"prompt":" replace the sky with a sunset","images":["/data/source/0001.png"],"metadata":{"task":"i2i","sample_id":"magicbrush_0001"}} +{"prompt":"