minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CI] Add per-job uv venv isolation and upgrade CI version to Cuda 13 (#23119)

Browse Source 
Co-authored-by: Kangyan Zhou <zky314343421@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: Alison Shao <a.shao@wustl.edu>
Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Baizhou Zhang
2026-04-19 05:32:36 -07:00
committed by GitHub
co-authored by Kangyan Zhou Claude Opus 4.7 Alison Shao Mick
parent 03828f4205
commit 6ecd6f84db
39 changed files with 892 additions and 239 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/rl/test_multi_instance_release_memory_occupation.py
-274
View File
@@ -1,274 +0,0 @@
import gc
import multiprocessing
import os
import traceback
import unittest
from multiprocessing import Process
from typing import Iterable, Tuple
import torch
import torch.distributed as dist
from torch.distributed.device_mesh import init_device_mesh
from transformers import AutoModelForCausalLM
from sglang.srt.entrypoints.engine import Engine as SglangEngine
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
CustomTestCase,
find_available_port,
)
register_cuda_ci(est_time=57, suite="stage-c-test-4-gpu-h100")
register_amd_ci(
est_time=64,
suite="stage-c-test-4-gpu-amd",
disabled="torch_memory_saver incompatible with ROCm (libcuda.so.1 not found)",
)
TEST_SUITE = dict(
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
mem_fraction_static=0.83,
dp_size=2,
tp_size=2,
)
# Minimum expected memory change in MB for each operation.
# Llama-3.2-1B bf16 is ~2GB total, ~1GB per TP rank.
# KV cache with mem_fraction_static=0.83 is much larger.
MIN_DELTA_MB = 200
class EngineWrapper:
"""
A wrapper around Sglang engine to mock multi instance cases such as RL training.
"""
def __init__(
self, model_path, random_seed, mem_fraction_static, device_mesh_cpu, base_gpu_id
):
self._device_mesh_cpu = device_mesh_cpu
self._tp_rank = device_mesh_cpu.get_local_rank()
self._rank = device_mesh_cpu.get_rank()
self._tp_size = device_mesh_cpu.size()
tp_size_per_node = self._tp_size
node_rank = self._tp_rank // tp_size_per_node
first_rank_in_node = self._tp_rank % tp_size_per_node == 0
engine_kwargs = dict(
model_path=model_path,
random_seed=random_seed,
mem_fraction_static=mem_fraction_static,
base_gpu_id=base_gpu_id,
enable_memory_saver=True,
tp_size=self._tp_size,
node_rank=node_rank,
nnodes=1,
)
self._engine = None
if first_rank_in_node:
os.environ["SGLANG_BLOCK_NONZERO_RANK_CHILDREN"] = "0"
self._engine = SglangEngine(**engine_kwargs)
dist.barrier(group=self._device_mesh_cpu.get_group())
def update_weights_from_tensor(
self, named_tensors: Iterable[Tuple[str, torch.Tensor]]
):
if self._tp_rank == 0:
self._engine.update_weights_from_tensor(list(named_tensors))
self._engine.flush_cache()
dist.barrier(group=self._device_mesh_cpu.get_group())
def release_memory_occupation(self, tags):
if self._tp_rank == 0:
self._engine.release_memory_occupation(tags)
dist.barrier(group=self._device_mesh_cpu.get_group())
def resume_memory_occupation(self, tags):
if self._tp_rank == 0:
self._engine.resume_memory_occupation(tags)
dist.barrier(group=self._device_mesh_cpu.get_group())
def shutdown(self):
if self._tp_rank == 0:
self._engine.shutdown()
dist.barrier(group=self._device_mesh_cpu.get_group())
def get_gpu_memory_mb(device_id: int) -> float:
"""Return device-level GPU memory used in MB."""
free, total = torch.cuda.mem_get_info(device_id)
return (total - free) / (1024**2)
def assert_memory_decreased(before_mb, after_mb, step_name):
delta = before_mb - after_mb
assert delta > MIN_DELTA_MB, (
f"[{step_name}] Expected memory decrease > {MIN_DELTA_MB} MB, "
f"got delta={delta:.0f} MB (before={before_mb:.0f}, after={after_mb:.0f})"
)
def assert_memory_increased(before_mb, after_mb, step_name):
delta = after_mb - before_mb
assert delta > MIN_DELTA_MB, (
f"[{step_name}] Expected memory increase > {MIN_DELTA_MB} MB, "
f"got delta={delta:.0f} MB (before={before_mb:.0f}, after={after_mb:.0f})"
)
class TestMultiInstanceReleaseMemoryOccupation(CustomTestCase):
@classmethod
def setUpClass(cls):
multiprocessing.set_start_method("spawn")
def test_multi_instance_release_memory_occupation(self):
master_port = find_available_port(23456)
dp_size = TEST_SUITE["dp_size"]
tp_size = TEST_SUITE["tp_size"]
world_size = dp_size * tp_size
processes = []
output_reader, output_writer = multiprocessing.Pipe(duplex=False)
for rank in range(world_size):
p = Process(
target=_run_sglang_subprocess,
kwargs=dict(
rank=rank,
dp_size=dp_size,
tp_size=tp_size,
model_path=TEST_SUITE["model_path"],
master_port=master_port,
output_writer=output_writer,
mem_fraction_static=TEST_SUITE["mem_fraction_static"],
),
)
p.start()
processes.append(p)
for _ in range(world_size):
self.assertTrue(
output_reader.recv(), f"Subprocess fail. Check the logs above."
)
for p in processes:
p.join()
def _run_sglang_subprocess(
rank: int,
dp_size: int,
tp_size: int,
model_path: str,
master_port: int,
output_writer,
mem_fraction_static: float,
):
engine = None
try:
os.environ["MASTER_ADDR"] = "localhost"
os.environ["MASTER_PORT"] = str(master_port)
dist.init_process_group(
rank=rank,
device_id=torch.device(f"cuda:{rank}"),
world_size=dp_size * tp_size,
)
torch.cuda.set_device(rank)
base_gpu_id = rank // tp_size * tp_size
mesh_kwargs = dict(
mesh_shape=(dp_size, tp_size, 1), mesh_dim_names=["dp", "tp", "pp"]
)
inference_device_mesh_cpu = init_device_mesh("cpu", **mesh_kwargs)
# Only TP master ranks (rank % tp_size == 0) create the Engine and
# measure memory. Non-master ranks share the same GPU and would see
# device-level memory from the master's Engine workers, causing
# unpredictable assertion results.
is_tp_master = rank % tp_size == 0
engine = EngineWrapper(
model_path=model_path,
random_seed=42,
mem_fraction_static=mem_fraction_static,
device_mesh_cpu=inference_device_mesh_cpu["tp"],
base_gpu_id=base_gpu_id,
)
print(f"subprocess[{rank=}] engine created, {is_tp_master=}", flush=True)
# 1 - release kv cache
if is_tp_master:
mem_before = get_gpu_memory_mb(rank)
print(f"GPU{rank} before releasing KV cache: {mem_before:.0f} MB")
engine.release_memory_occupation(tags=["kv_cache"])
if is_tp_master:
mem_after = get_gpu_memory_mb(rank)
assert_memory_decreased(mem_before, mem_after, "release KV cache")
# 2 - release sglang weights
if is_tp_master:
mem_before = get_gpu_memory_mb(rank)
print(f"GPU{rank} before releasing weights: {mem_before:.0f} MB")
engine.release_memory_occupation(tags=["weights"])
if is_tp_master:
mem_after = get_gpu_memory_mb(rank)
assert_memory_decreased(mem_before, mem_after, "release weights")
# 3 - load hf model (TP master only)
hf_model = None
if is_tp_master:
mem_before = get_gpu_memory_mb(rank)
print(f"GPU{rank} before loading HF model: {mem_before:.0f} MB")
# Avoid device_map= which triggers accelerate dispatch hooks in
# transformers v5, preventing clean memory release on del.
hf_model = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
).to(f"cuda:{rank}")
mem_after = get_gpu_memory_mb(rank)
assert_memory_increased(mem_before, mem_after, "load HF model")
dist.barrier(group=inference_device_mesh_cpu["tp"].get_group())
# 4 - resume sglang weights and update from hf model
engine.resume_memory_occupation(tags=["weights"])
engine.update_weights_from_tensor(
named_tensors=list(hf_model.named_parameters()) if hf_model else []
)
# 5 - release hf model (TP master only)
if is_tp_master:
mem_before = get_gpu_memory_mb(rank)
print(f"GPU{rank} before releasing HF model: {mem_before:.0f} MB")
del hf_model
gc.collect()
torch.cuda.empty_cache()
mem_after = get_gpu_memory_mb(rank)
assert_memory_decreased(mem_before, mem_after, "release HF model")
dist.barrier(group=inference_device_mesh_cpu["tp"].get_group())
# 6 - resume kv cache
if is_tp_master:
mem_before = get_gpu_memory_mb(rank)
print(f"GPU{rank} before resuming KV cache: {mem_before:.0f} MB")
engine.resume_memory_occupation(tags=["kv_cache"])
if is_tp_master:
mem_after = get_gpu_memory_mb(rank)
assert_memory_increased(mem_before, mem_after, "resume KV cache")
print(f"GPU{rank} final memory: {mem_after:.0f} MB")
execution_ok = True
except Exception as e:
print(f"subprocess[{rank=}] has error: {e}", flush=True)
traceback.print_exc()
execution_ok = False
output_writer.send(execution_ok)
output_writer.close()
if engine:
engine.shutdown()
if __name__ == "__main__":
unittest.main()