diff --git a/python/sglang/srt/debug_utils/dumper.py b/python/sglang/srt/debug_utils/dumper.py index 249106a94..364ebbf63 100644 --- a/python/sglang/srt/debug_utils/dumper.py +++ b/python/sglang/srt/debug_utils/dumper.py @@ -842,30 +842,36 @@ class _Grafter: role = _GraftRole(cfg.grafter_role) is_send = self._is_sender(role=role, direction=direction) - # 1+1 broadcast: sender side ships the tensor as a pickled object; - # recv side feeds it through the user transform (default: identity) - # and `value.copy_()` the result. - sender_rank = 0 if direction == _GraftDirection.B2T else 1 - obj_list: list = [None] + # all-gather over the graft world; sender ranks contribute `value`, + # recv ranks contribute None (their local target is private and + # shouldn't leak). all_gather_object is pickle-routed, so tensor + # shapes may differ across sender ranks. + total_world = cfg.grafter_baseline_world_size + cfg.grafter_target_world_size + my_contribution = value if is_send else None + gathered: list = [None] * total_world + dist.all_gather_object(gathered, my_contribution, group=self._pg) + if is_send: - obj_list = [value] _log(f"[Grafter] send role={role.value} dir={direction.value} tags={tags}") - dist.broadcast_object_list(obj_list, src=sender_rank, group=self._pg) - if is_send: return - received = obj_list[0] - if isinstance(received, torch.Tensor): - # Pickled CUDA tensors restore to their original-device name; - # that may not match this process's local device, so normalize. - received = received.to(value.device) + sender_contribs = self._sender_slice(direction=direction, gathered=gathered) + # Pickled CUDA tensors restore to their original-device name; that + # may not match this process's local device, so normalize. + sender_tensors = [ + (t.to(value.device) if isinstance(t, torch.Tensor) else t) + for t in sender_contribs + ] # Transform + copy_ are wrapped: a buggy user transform must NOT # crash the whole training/inference run. On error we log the full # traceback and skip this graft point; downstream sees the recv # side's original tensor unchanged. try: - value_to_override = self._apply_transform([received], target=value) - _log(f"[Grafter] recv role={role.value} dir={direction.value} tags={tags}") + value_to_override = self._apply_transform(sender_tensors, target=value) + _log( + f"[Grafter] recv role={role.value} dir={direction.value} " + f"tags={tags} n_senders={len(sender_tensors)}" + ) value.copy_(value_to_override) except Exception as e: _log( @@ -875,6 +881,12 @@ class _Grafter: f"{traceback.format_exc()}" ) + def _sender_slice(self, *, direction: "_GraftDirection", gathered: list) -> list: + cfg = self._config + if direction == _GraftDirection.B2T: + return gathered[: cfg.grafter_baseline_world_size] + return gathered[cfg.grafter_baseline_world_size :] + def _apply_transform( self, received_list: list, @@ -890,16 +902,27 @@ class _Grafter: def _default_transform( received_list: list, *, target: torch.Tensor ) -> torch.Tensor: - """Identity-by-rank fallback. For the 1+1 setup currently supported, - just returns the single received tensor; requires shape match.""" - candidate = received_list[0] + """Identity-by-rank fallback. Requires #senders == #recvs and + shape(received_list[my_recv_rank]) == shape(target). Otherwise raises + and asks the user for a transform.""" + my_recv_rank = dist.get_rank() + recv_world_size = dist.get_world_size() + if len(received_list) != recv_world_size: + raise RuntimeError( + f"[Grafter] no grafter_transform_path set; default " + f"identity-by-rank requires #senders == #recvs but got " + f"#senders={len(received_list)} vs #recvs={recv_world_size}. " + f"Provide a transform via " + f"DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol." + ) + candidate = received_list[my_recv_rank] if candidate.shape != target.shape: raise RuntimeError( f"[Grafter] no grafter_transform_path set; default " f"identity-by-rank requires matching shapes but " - f"received_list[0].shape={tuple(candidate.shape)} != " - f"target.shape={tuple(target.shape)}. Provide a transform via " - f"DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol." + f"received_list[{my_recv_rank}].shape={tuple(candidate.shape)} " + f"!= target.shape={tuple(target.shape)}. Provide a transform " + f"via DUMPER_GRAFTER_TRANSFORM_PATH=pkg.module.symbol." ) return candidate @@ -938,11 +961,26 @@ class _Grafter: dist.is_initialized() ), "[Grafter] default torch.distributed must be initialized" role = _GraftRole(cfg.grafter_role) - global_rank = 0 if role == _GraftRole.BASELINE else 1 + local_world = dist.get_world_size() + local_rank = dist.get_rank() + if role == _GraftRole.BASELINE: + assert local_world == cfg.grafter_baseline_world_size, ( + f"[Grafter] grafter_baseline_world_size={cfg.grafter_baseline_world_size} " + f"but dist.get_world_size()={local_world}; they must match on the baseline side" + ) + global_rank = local_rank + else: + assert local_world == cfg.grafter_target_world_size, ( + f"[Grafter] grafter_target_world_size={cfg.grafter_target_world_size} " + f"but dist.get_world_size()={local_world}; they must match on the target side" + ) + global_rank = cfg.grafter_baseline_world_size + local_rank total_world = cfg.grafter_baseline_world_size + cfg.grafter_target_world_size init_method = f"tcp://{cfg.grafter_master_address}:{cfg.grafter_master_port}" _log( f"[Grafter] init group: role={role.value} " + f"baseline_world={cfg.grafter_baseline_world_size} " + f"target_world={cfg.grafter_target_world_size} " f"rank={global_rank} init_method={init_method} " f"backend={cfg.grafter_backend} name={cfg.grafter_group_name}" ) diff --git a/test/registered/debug_utils/test_dumper.py b/test/registered/debug_utils/test_dumper.py index a7f5cec63..6da7ef7c6 100644 --- a/test/registered/debug_utils/test_dumper.py +++ b/test/registered/debug_utils/test_dumper.py @@ -3246,5 +3246,308 @@ class TestGrafterDistributed: dist.destroy_process_group(grafter._pg) +def _run_graft_test_cpu_multi( + worker_func, *, baseline_world: int, target_world: int, **kwargs +): + """Spawn (baseline_world + target_world) CPU-only processes (gloo backend). + + Used to exercise asymmetric multi-rank cases (e.g. 4 baseline ranks and + 2 target ranks) that we can't run on the 2-GPU CI fleet. Each role gets + its OWN default PG (gloo, world=role_world); the graft cross-system PG + spans all ranks. + + The worker function receives (role, local_rank, **kwargs). + """ + import torch.multiprocessing as mp + + role_ports = { + "baseline": find_available_port(29800), + "target": find_available_port(29900), + } + + ctx = mp.get_context("spawn") + result_queue = ctx.Queue() + processes = [] + total = baseline_world + target_world + for global_rank in range(total): + if global_rank < baseline_world: + role = "baseline" + local_rank = global_rank + local_world = baseline_world + else: + role = "target" + local_rank = global_rank - baseline_world + local_world = target_world + p = ctx.Process( + target=_graft_cpu_worker_entry, + args=( + role, + local_rank, + local_world, + role_ports[role], + worker_func, + result_queue, + kwargs, + ), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() + + errors = [result_queue.get() for _ in range(total)] + errors = [e for e in errors if e] + if errors: + raise AssertionError("\n".join(errors)) + + +def _graft_cpu_worker_entry( + role, local_rank, local_world, port, worker_func, result_queue, kwargs +): + dist.init_process_group( + backend="gloo", + init_method=f"tcp://127.0.0.1:{port}", + world_size=local_world, + rank=local_rank, + ) + try: + worker_func(role=role, local_rank=local_rank, **kwargs) + result_queue.put(None) + except Exception as e: + result_queue.put( + f"role={role} local_rank={local_rank}: {e}\n{traceback.format_exc()}" + ) + finally: + dist.destroy_process_group() + + +def _make_multi_rank_config( + *, + role: str, + graft_port: int, + group_name: str, + baseline_world: int, + target_world: int, + transform_path: Optional[str], + direction: str, +) -> DumperConfig: + return DumperConfig( + grafter_enable=True, + grafter_role=role, + grafter_b2t_filter="name == 'x'" if direction == "b2t" else None, + grafter_t2b_filter="name == 'x'" if direction == "t2b" else None, + grafter_master_address="127.0.0.1", + grafter_master_port=graft_port, + grafter_baseline_world_size=baseline_world, + grafter_target_world_size=target_world, + grafter_backend="gloo", + grafter_group_name=group_name, + grafter_timeout=30, + grafter_transform_path=transform_path, + ) + + +class TestGrafterMultiRankCpu: + """Coverage of asymmetric multi-rank cases via CPU/gloo (CI fleet has + only 2 GPUs, which is too few for these cases).""" + + def test_4_baseline_2_target_b2t_with_user_transform(self, tmp_path: Path): + """4 baseline senders -> 2 target receivers via b2t graft.""" + module_name = "_xform_assert_4_senders" + (tmp_path / f"{module_name}.py").write_text( + "import torch\n" + "def transform(received_list, target):\n" + " rl = received_list\n" + " assert len(rl) == 4, f'expected 4 senders, got {len(rl)}'\n" + " for i, t in enumerate(rl):\n" + " v = float(t.flatten()[0].item())\n" + " assert v == float(i), f'rl[{i}][0]={v}, want {float(i)}'\n" + " return torch.full_like(target, 999.0)\n" + ) + graft_port = find_available_port(29655) + _run_graft_test_cpu_multi( + self._test_4b_2t_func, + baseline_world=4, + target_world=2, + graft_port=graft_port, + group_name="grafter_4b_2t", + transform_dir=str(tmp_path), + transform_path=f"{module_name}.transform", + ) + + @staticmethod + def _test_4b_2t_func( + role, local_rank, graft_port, group_name, transform_dir, transform_path + ): + sys.path.insert(0, transform_dir) + cfg = _make_multi_rank_config( + role=role, + graft_port=graft_port, + group_name=group_name, + baseline_world=4, + target_world=2, + transform_path=transform_path, + direction="b2t", + ) + grafter = _Grafter(config=cfg) + try: + if role == "baseline": + tensor = torch.full((3,), float(local_rank)) + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.full((3,), 99.0) + grafter.maybe_intercept(value=target, tags={"name": "x"}) + assert target.tolist() == [999.0, 999.0, 999.0], target.tolist() + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + + def test_2_target_4_baseline_t2b_with_user_transform(self, tmp_path: Path): + """Mirror image: 2 target senders -> 4 baseline receivers.""" + module_name = "_xform_assert_2_senders_t2b" + (tmp_path / f"{module_name}.py").write_text( + "import torch\n" + "def transform(received_list, target):\n" + " rl = received_list\n" + " assert len(rl) == 2, f'expected 2 senders, got {len(rl)}'\n" + " for i, t in enumerate(rl):\n" + " v = float(t.flatten()[0].item())\n" + " assert v == float(i + 100), f'rl[{i}][0]={v}'\n" + " return torch.full_like(target, 7.0)\n" + ) + graft_port = find_available_port(29670) + _run_graft_test_cpu_multi( + self._test_2t_4b_func, + baseline_world=4, + target_world=2, + graft_port=graft_port, + group_name="grafter_2t_4b", + transform_dir=str(tmp_path), + transform_path=f"{module_name}.transform", + ) + + @staticmethod + def _test_2t_4b_func( + role, local_rank, graft_port, group_name, transform_dir, transform_path + ): + sys.path.insert(0, transform_dir) + cfg = _make_multi_rank_config( + role=role, + graft_port=graft_port, + group_name=group_name, + baseline_world=4, + target_world=2, + transform_path=transform_path, + direction="t2b", + ) + grafter = _Grafter(config=cfg) + try: + if role == "target": + tensor = torch.full((3,), float(local_rank + 100)) + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.full((3,), 99.0) + grafter.maybe_intercept(value=target, tags={"name": "x"}) + assert target.tolist() == [7.0, 7.0, 7.0], target.tolist() + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + + def test_default_transform_with_asymmetric_world_logs_and_skips(self): + """Default identity-by-rank requires #senders == #recvs; with 4 vs 2 + and no user transform, recv catches RuntimeError + leaves target + unchanged.""" + graft_port = find_available_port(29675) + _run_graft_test_cpu_multi( + self._test_default_asym_func, + baseline_world=4, + target_world=2, + graft_port=graft_port, + group_name="grafter_default_asym", + ) + + @staticmethod + def _test_default_asym_func(role, local_rank, graft_port, group_name): + cfg = _make_multi_rank_config( + role=role, + graft_port=graft_port, + group_name=group_name, + baseline_world=4, + target_world=2, + transform_path=None, + direction="b2t", + ) + grafter = _Grafter(config=cfg) + try: + if role == "baseline": + tensor = torch.full((3,), float(local_rank)) + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.full((3,), 42.0) + with _capture_stdout() as captured: + grafter.maybe_intercept(value=target, tags={"name": "x"}) + assert target.tolist() == [42.0, 42.0, 42.0], target.tolist() + output = captured.getvalue() + assert "transform/copy_ raised RuntimeError" in output, output + assert "#senders=4" in output and "#recvs=2" in output, output + assert "Traceback (most recent call last)" in output, output + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + + def test_mixed_shape_senders_via_user_transform(self, tmp_path: Path): + """`all_gather_object` is pickle-routed, so sender ranks may + contribute tensors with DIFFERENT shapes. The user transform handles + the dispatch.""" + module_name = "_xform_concat_mixed_shape" + (tmp_path / f"{module_name}.py").write_text( + "import torch\n" + "def transform(received_list, target):\n" + " expected_shapes = [(i + 1,) for i in range(len(received_list))]\n" + " actual_shapes = [tuple(t.shape) for t in received_list]\n" + " assert actual_shapes == expected_shapes, actual_shapes\n" + " return torch.cat(received_list)\n" + ) + graft_port = find_available_port(29680) + _run_graft_test_cpu_multi( + self._test_mixed_shape_func, + baseline_world=4, + target_world=2, + graft_port=graft_port, + group_name="grafter_mixed_shape", + transform_dir=str(tmp_path), + transform_path=f"{module_name}.transform", + ) + + @staticmethod + def _test_mixed_shape_func( + role, local_rank, graft_port, group_name, transform_dir, transform_path + ): + sys.path.insert(0, transform_dir) + cfg = _make_multi_rank_config( + role=role, + graft_port=graft_port, + group_name=group_name, + baseline_world=4, + target_world=2, + transform_path=transform_path, + direction="b2t", + ) + grafter = _Grafter(config=cfg) + try: + if role == "baseline": + tensor = torch.full((local_rank + 1,), float(local_rank)) + grafter.maybe_intercept(value=tensor, tags={"name": "x"}) + else: + target = torch.zeros(10) + grafter.maybe_intercept(value=target, tags={"name": "x"}) + expected = [0.0] + [1.0] * 2 + [2.0] * 3 + [3.0] * 4 + assert target.tolist() == expected, target.tolist() + finally: + if grafter._pg is not None: + dist.destroy_process_group(grafter._pg) + + if __name__ == "__main__": sys.exit(pytest.main([__file__]))