[CI] Add per-job uv venv isolation and upgrade CI version to Cuda 13 (#23119)
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
co-authored by
Kangyan Zhou
Claude Opus 4.7
Alison Shao
Mick
Kangyan Zhou
Claude Opus 4.7
Alison Shao
Mick
parent
03828f4205
commit
6ecd6f84db
@@ -0,0 +1,279 @@
|
||||
"""
|
||||
# TODO: Fails on cu13 venv migration. Ref: https://github.com/sgl-project/sglang/actions/runs/24616960626/job/71980705675?pr=23119
|
||||
# Should move back to registered test after it's fixed
|
||||
"""
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user