[ci][cu13] Bump torch_memory_saver to 0.0.9.post1; restore manual tests (#23182)
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Kangyan-Zhou
co-authored by
Claude Opus 4.7
Claude Opus 4.7
parent
674a80d5d8
commit
cd27baaffd
@@ -1,178 +0,0 @@
|
||||
import asyncio
|
||||
import os
|
||||
import unittest
|
||||
|
||||
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
|
||||
from sglang.srt.weight_sync.utils import update_weights
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
register_cuda_ci(est_time=32, suite="stage-b-test-1-gpu-large")
|
||||
|
||||
"""
|
||||
# TODO: torch_memory_saver wheel is built against libcudart.so.12, fails to LD_PRELOAD in Cu13 venv. Ref: https://github.com/sgl-project/sglang/actions/runs/24604424372/job/71968573867
|
||||
# Should move back to registered test after it's fixed
|
||||
"""
|
||||
|
||||
|
||||
class AsyncEngine(Engine):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
async def update_weights_from_tensor(self, update_weights_request):
|
||||
return await self.tokenizer_manager.update_weights_from_tensor(
|
||||
update_weights_request, None
|
||||
)
|
||||
|
||||
|
||||
def is_distributed_available():
|
||||
"""Check if distributed training environment is available"""
|
||||
required_vars = ["RANK", "WORLD_SIZE", "MASTER_ADDR", "MASTER_PORT"]
|
||||
return all(var in os.environ for var in required_vars)
|
||||
|
||||
|
||||
def setup_single_process_distributed():
|
||||
"""Setup distributed environment for single process testing"""
|
||||
if not is_distributed_available():
|
||||
os.environ["RANK"] = "0"
|
||||
os.environ["WORLD_SIZE"] = "1"
|
||||
os.environ["MASTER_ADDR"] = "localhost"
|
||||
os.environ["MASTER_PORT"] = "12356"
|
||||
os.environ["LOCAL_RANK"] = "0"
|
||||
|
||||
|
||||
class TestUtilsUpdateWeights(unittest.TestCase):
|
||||
"""Test class for utils.update_weights function"""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Setup distributed environment and test fixtures for the entire test class"""
|
||||
cls.setup_distributed()
|
||||
cls.setup_test_engine()
|
||||
cls.setup_test_model()
|
||||
cls.setup_device_mesh()
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
"""Cleanup after all tests"""
|
||||
if hasattr(cls, "engine") and cls.engine:
|
||||
cls.engine.shutdown()
|
||||
|
||||
# Cleanup distributed
|
||||
if dist.is_initialized():
|
||||
dist.destroy_process_group()
|
||||
|
||||
@classmethod
|
||||
def setup_distributed(cls):
|
||||
"""Setup distributed environment for testing"""
|
||||
setup_single_process_distributed()
|
||||
|
||||
if not dist.is_initialized():
|
||||
try:
|
||||
dist.init_process_group(
|
||||
backend="nccl" if torch.cuda.is_available() else "gloo"
|
||||
)
|
||||
except Exception as e:
|
||||
raise unittest.SkipTest(
|
||||
f"Could not initialize distributed backend: {e}"
|
||||
)
|
||||
|
||||
cls.rank = dist.get_rank()
|
||||
cls.world_size = dist.get_world_size()
|
||||
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.set_device(cls.rank % torch.cuda.device_count())
|
||||
|
||||
# Set up environment variables
|
||||
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
|
||||
os.environ["NCCL_CUMEM_ENABLE"] = "0"
|
||||
os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "4"
|
||||
os.environ["CUDA_MODULE_LOADING"] = "AUTO"
|
||||
|
||||
@classmethod
|
||||
def setup_test_engine(cls):
|
||||
"""Setup test engine"""
|
||||
if cls.rank == 0:
|
||||
cls.engine = AsyncEngine(
|
||||
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
dtype="bfloat16",
|
||||
mem_fraction_static=0.3,
|
||||
enable_memory_saver=True,
|
||||
tp_size=cls.world_size,
|
||||
disable_cuda_graph=False,
|
||||
)
|
||||
else:
|
||||
cls.engine = None
|
||||
|
||||
@classmethod
|
||||
def setup_test_model(cls):
|
||||
"""Load test model"""
|
||||
try:
|
||||
cls.model = AutoModelForCausalLM.from_pretrained(
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
device_map="cpu",
|
||||
trust_remote_code=True,
|
||||
low_cpu_mem_usage=True,
|
||||
torch_dtype=(
|
||||
torch.float16 if torch.cuda.is_available() else torch.float32
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
raise unittest.SkipTest(f"Could not load test model: {e}")
|
||||
|
||||
@classmethod
|
||||
def setup_device_mesh(cls):
|
||||
"""Create device mesh for testing"""
|
||||
if not torch.cuda.is_available():
|
||||
raise unittest.SkipTest("CUDA not available for device mesh")
|
||||
|
||||
cls.device_mesh_key = "tp"
|
||||
cls.mesh = init_device_mesh(
|
||||
"cuda", (cls.world_size,), mesh_dim_names=(cls.device_mesh_key,)
|
||||
)
|
||||
|
||||
def create_test_params_batch(self, model, num_params=64):
|
||||
"""Create a batch of test parameters from the model"""
|
||||
param_names = []
|
||||
test_tensors = []
|
||||
|
||||
# Get first few parameters from the model for testing
|
||||
for i, (name, tensor) in enumerate(model.named_parameters()):
|
||||
if i >= num_params:
|
||||
break
|
||||
param_names.append(name)
|
||||
# Create test tensor with known values, matching original shape and dtype
|
||||
test_tensor = torch.full_like(tensor, 1.5, dtype=tensor.dtype).cuda()
|
||||
test_tensors.append(test_tensor)
|
||||
|
||||
return list(zip(param_names, test_tensors))
|
||||
|
||||
def test_utils_update_weights(self):
|
||||
"""Test basic functionality of utils.update_weights"""
|
||||
|
||||
async def async_test():
|
||||
# Create test parameters batch
|
||||
params_batch = self.create_test_params_batch(self.model, num_params=2)
|
||||
|
||||
# Test the utils.update_weights function
|
||||
result = await update_weights(
|
||||
engine=self.engine,
|
||||
params_batch=params_batch,
|
||||
device_mesh_key=self.device_mesh_key,
|
||||
device_mesh=self.mesh,
|
||||
load_format=None,
|
||||
)
|
||||
|
||||
self.assertIn("Success", result)
|
||||
|
||||
# Run the async test
|
||||
asyncio.run(async_test())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,279 +0,0 @@
|
||||
"""
|
||||
# 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