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

Migrate 4-GPU/8-GPU workflow jobs to stage-c and add CI registry decorators (#17299)

Browse Source
This commit is contained in:
Alison Shao
2026-01-31 22:37:22 -08:00
committed by GitHub
parent 95180484e9
commit a0bae4c343
39 changed files with 189 additions and 284 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/rl/test_multi_instance_release_memory_occupation.py
+256
View File
@@ -0,0 +1,256 @@
import multiprocessing
import os
import time
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_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=64, suite="stage-c-test-4-gpu-h100")
TEST_SUITE = dict(
model_path=DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
mem_fraction_static=0.83,
dp_size=2,
tp_size=2,
)
class EngineWrapper:
"""
A wrapper around Sglang engine to mock multi instance cases such as RL traing.
"""
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_gb(gpu_id=0):
return torch.cuda.device_memory_used() / 1024**3
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_device = init_device_mesh("cuda", **mesh_kwargs)
inference_device_mesh_cpu = init_device_mesh("cpu", **mesh_kwargs)
print(
f"subprocess[{rank=},{base_gpu_id=},{rank=},{tp_size=}] {inference_device_mesh_device=} {inference_device_mesh_cpu=}"
)
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage before starting Engine: {_mem_usage}")
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=}", flush=True)
# 1 - release kv cache
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage before releasing Sgl KV cache: {_mem_usage}")
engine.release_memory_occupation(tags=["kv_cache"])
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage < _mem_usage
), f"Memory usage after releasing KV cache must be reduced! before: {_mem_usage} vs after: {_curr_usage}"
# 2 - release sglang weights
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage before releasing Sgl weights: {_mem_usage}")
engine.release_memory_occupation(tags=["weights"])
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage < _mem_usage
), f"Memory usage after releasing weights must be reduced! before: {_mem_usage} vs after: {_curr_usage}"
# 3 - load hf model
_mem_usage = get_gpu_memory_gb(rank)
print(
f"GPU{rank} Memory usage after releasing Sgl weights and kv cache: {_mem_usage}"
)
hf_model = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map=f"cuda:{rank}",
trust_remote_code=True,
).cuda()
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage > _mem_usage
), f"Memory usage after loading hf model must be increased! before: {_mem_usage} vs after: {_curr_usage}"
# 4 - resume sglang weights and update the weights
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage after loading hf model: {_mem_usage}")
engine.resume_memory_occupation(tags=["weights"])
engine.update_weights_from_tensor(
named_tensors=list(hf_model.named_parameters())
)
# 5 - release hf model
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage after resuming Sgl weights: {_mem_usage}")
del hf_model
hf_model = None
torch.cuda.empty_cache()
time.sleep(3)
torch.cuda.empty_cache()
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage < _mem_usage
), f"Memory usage after releasing hf model must be reduced! before: {_mem_usage} vs after: {_curr_usage}"
# 6 - resume slgang kv cache
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage after releasing hf model: {_mem_usage}")
engine.resume_memory_occupation(tags=["kv_cache"])
_curr_usage = get_gpu_memory_gb(rank)
assert (
_curr_usage > _mem_usage
), f"Memory usage after resuming kv cache must be increased! before: {_mem_usage} vs after: {_curr_usage}"
# 7 - Final checking!
_mem_usage = get_gpu_memory_gb(rank)
print(f"GPU{rank} Memory usage after resuming Sgl KV cache: {_mem_usage}")
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()
test/registered/rl/test_release_memory_occupation.py
+479
View File
@@ -0,0 +1,479 @@
"""Test memory release and resume operations for SGLang engine in hybrid RL training.
This test suite evaluates the SGLang engine's memory management capabilities, focusing
on releasing and resuming memory occupation for KV cache and model weights. It simulates
an RL workflow where the SGLang engine acts as a rollout engine for experience collection.
The process involves initializing the engine, sending a small number of requests to simulate
rollout, releasing memory to mimic offloading during RL training, resuming memory occupation,
updating weights with a trained HuggingFace model, and verifying the updated weights.
Detailed in our proposal (https://github.com/sgl-project/sglang/pull/7099), two test cases
are included:
1. Basic Release and Resume: Uses a lower mem_fraction_static (0.6) to control memory allocation
and avoid OOM errors carefully. This test simulates a scenario without multi-stage memory management,
ensuring the engine can release and resume memory occupation while maintaining functionality after
weight updates.
2. Multi-Stage Release and Resume: Employs a higher mem_fraction_static (0.85) to simulate higher
memory pressure, leveraging multi-stage memory management. It sequentially releases and resumes
KV cache and model weights, verifying memory deallocation and reallocation at each stage, and
ensuring correct weight updates and text generation.
3. Tensor Parallel Tests: Tests memory release and resume operations with different tensor parallel
configurations (tp=1, tp=2) to ensure proper memory management in distributed settings. For different
data parallel size, we test it in verl.
NOTE: This test is temporarily disabled.
"""
import os
import time
import unittest
import torch
from transformers import AutoModelForCausalLM
import sglang as sgl
from sglang.srt.constants import (
GPU_MEMORY_TYPE_CUDA_GRAPH,
GPU_MEMORY_TYPE_KV_CACHE,
GPU_MEMORY_TYPE_WEIGHTS,
)
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import (
DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT,
CustomTestCase,
)
register_cuda_ci(
est_time=200,
suite="stage-c-test-4-gpu-h100",
disabled="Temporarily disabled - needs investigation",
)
# (temporarily) set to true to observe memory usage in nvidia-smi more clearly
_DEBUG_EXTRA = False
def get_gpu_memory_gb():
return torch.cuda.device_memory_used() / 1024**3
class TestReleaseMemoryOccupation(CustomTestCase):
def _setup_engine(
self,
model_name,
mem_fraction_static=0.8,
tp_size=1,
ep_size=1,
enable_weights_cpu_backup=False,
):
"""Common setup for engine and HF model."""
os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1"
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=mem_fraction_static,
tp_size=tp_size,
ep_size=ep_size,
enable_weights_cpu_backup=enable_weights_cpu_backup,
# disable_cuda_graph=True, # for debugging only
)
return engine
def _common_test_params(self):
"""Common test parameters."""
return {
"prompt": "Today is a sunny day and I like",
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
"expect_output_before_update_weights": " to spend it outdoors. I decided to",
"expect_output_after_update_weights": " to go for a walk. I like",
"prompt_moe": "The weather is nice today, and I want to",
"sampling_params_moe": {"temperature": 0, "max_new_tokens": 16},
"expect_output_before_update_weights_moe": " go to the park. I have a picnic basket, a book, and a",
"expect_output_after_update_weights_moe": " go to the park. I have a lot of things to do, but I",
"prompt_hybrid_mamba": "The weather is nice today, and I want to",
"sampling_params_hybrid_mamba": {"temperature": 0, "max_new_tokens": 16},
"expect_output_before_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
"expect_output_after_update_weights_hybrid_mamba": " go out for a walk. But I don't know what to wear. Can",
}
def _test_initial_generation(
self, engine, prompt, sampling_params, expect_output_before_update_weights
):
"""Test initial generation and memory allocation."""
print("generate (#1)")
outputs = engine.generate(prompt, sampling_params)["text"]
self.assertEqual(outputs, expect_output_before_update_weights)
if _DEBUG_EXTRA:
time.sleep(3)
def test_release_and_resume_occupation(self):
# Without multi-stage release and resume, we need to carefully control the memory fraction to avoid OOM
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
assert (
torch.cuda.device_count() >= 2
), "Need at least 2 GPUs for tensor parallel tests"
for tp_size in [1, 2]:
print(f"Testing tp_size={tp_size} for test_release_and_resume_occupation")
engine = self._setup_engine(
model_name=model_name, mem_fraction_static=0.6, tp_size=tp_size
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
)
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
print("generate (#2)")
outputs = engine.generate(params["prompt"], params["sampling_params"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights"])
engine.shutdown()
def test_release_and_resume_occupation_with_weights_cpu_backup(self):
# Test release and resume occupation with weights CPU backup
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
print("Testing test_release_and_resume_occupation_with_weights_cpu_backup")
engine = self._setup_engine(
model_name=model_name,
mem_fraction_static=0.6,
enable_weights_cpu_backup=True,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
print("generate post resume")
outputs = engine.generate(params["prompt"], params["sampling_params"])["text"]
self.assertEqual(outputs, params["expect_output_before_update_weights"])
engine.shutdown()
def test_multi_stage_release_and_resume(self):
# With multi-stage release and resume, we can set the memory fraction to 0.85 without concern of OOM
model_name = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
for tp_size in [1, 2]:
if tp_size == 2 and torch.cuda.device_count() < 2:
continue
print(f"Testing tp_size={tp_size} for test_multi_stage_release_and_resume")
os.environ["SGLANG_MEMORY_SAVER_CUDA_GRAPH"] = "1"
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=0.85, # Higher memory pressure
tp_size=tp_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt"],
params["sampling_params"],
params["expect_output_before_update_weights"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
gpu_memory_usage_after_release_kv_cache = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_kv_cache,
gpu_memory_usage_before_release,
)
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS])
gpu_memory_usage_after_release_weights = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_weights,
gpu_memory_usage_after_release_kv_cache,
)
engine.release_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])
gpu_memory_usage_after_release_cuda_graph = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release_cuda_graph,
gpu_memory_usage_after_release_weights,
)
print(f"Release took {time.perf_counter() - t:.2f}s")
print(
f"Memory: {gpu_memory_usage_before_release:.1f} → {gpu_memory_usage_after_release_kv_cache:.1f} → {gpu_memory_usage_after_release_weights:.1f} → {gpu_memory_usage_after_release_cuda_graph:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
gpu_memory_usage_before_resume = get_gpu_memory_gb()
# gpu_memory_usage_after_release_weights and gpu_memory_usage_before_resume should be close
self.assertAlmostEqual(
gpu_memory_usage_after_release_weights,
gpu_memory_usage_before_resume,
delta=3.0,
)
print(f"Resume weights took {time.perf_counter() - t:.2f}s")
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_CUDA_GRAPH])
gpu_memory_usage_after_resume_cuda_graph = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_cuda_graph,
gpu_memory_usage_before_resume,
)
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_WEIGHTS])
gpu_memory_usage_after_resume_weights = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_weights,
gpu_memory_usage_after_resume_cuda_graph,
)
# Update weights from a trained model to serving engine, and then destroy the trained model
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
)
gpu_memory_usage_after_loaded_hf_model = get_gpu_memory_gb()
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
engine.resume_memory_occupation(tags=[GPU_MEMORY_TYPE_KV_CACHE])
gpu_memory_usage_after_resume_kv_cache = get_gpu_memory_gb()
self.assertGreater(
gpu_memory_usage_after_resume_kv_cache,
gpu_memory_usage_after_resume_weights,
)
print(f"Resume + update took {time.perf_counter() - t:.2f}s")
print(
f"Memory: {gpu_memory_usage_before_resume:.1f} → {gpu_memory_usage_after_resume_cuda_graph:.1f} → {gpu_memory_usage_after_resume_weights:.1f} → {gpu_memory_usage_after_loaded_hf_model:.1f} → {gpu_memory_usage_after_resume_kv_cache:.1f} GB"
)
print("generate (#2)")
outputs = engine.generate(params["prompt"], params["sampling_params"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights"])
engine.shutdown()
def test_moe_model_release_and_resume(self):
# Test with MoE model
model_name = DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_CHAT
tp_size = ep_size = 2
print(
f"Testing tp_size={tp_size} and ep_size={ep_size} for test_moe_model_release_and_resume"
)
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
mem_fraction_static=0.5,
tp_size=tp_size,
ep_size=ep_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt_moe"],
params["sampling_params_moe"],
params["expect_output_before_update_weights_moe"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
hf_model_new = AutoModelForCausalLM.from_pretrained(
DEFAULT_SMALL_MOE_MODEL_NAME_FOR_TEST_BASE,
torch_dtype="bfloat16",
device_map="cuda",
)
engine.update_weights_from_tensor(list(hf_model_new.named_parameters()))
# destroy the hf model
del hf_model_new
torch.cuda.empty_cache()
print("generate (#2)")
outputs = engine.generate(params["prompt_moe"], params["sampling_params_moe"])[
"text"
]
self.assertEqual(outputs, params["expect_output_after_update_weights_moe"])
engine.shutdown()
def test_hybrid_mamba_model_release_and_resume(self):
# Test with Hybrid Mamba model
model_name = DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST
tp_size = 4
print(
f"Testing tp_size={tp_size} for test_hybrid_mamba_model_release_and_resume"
)
engine = sgl.Engine(
model_path=model_name,
random_seed=42,
enable_memory_saver=True,
tp_size=tp_size,
)
params = self._common_test_params()
self._test_initial_generation(
engine,
params["prompt_hybrid_mamba"],
params["sampling_params_hybrid_mamba"],
params["expect_output_before_update_weights_hybrid_mamba"],
)
t = time.perf_counter()
gpu_memory_usage_before_release = get_gpu_memory_gb()
engine.release_memory_occupation()
gpu_memory_usage_after_release = get_gpu_memory_gb()
self.assertLess(
gpu_memory_usage_after_release,
gpu_memory_usage_before_release,
)
print(
f"Release took {time.perf_counter() - t:.2f}s, memory: {gpu_memory_usage_before_release:.1f} GB → {gpu_memory_usage_after_release:.1f} GB"
)
if _DEBUG_EXTRA:
time.sleep(3)
t = time.perf_counter()
engine.resume_memory_occupation()
print(
f"Resume took {time.perf_counter() - t:.2f}s, memory: {get_gpu_memory_gb():.1f} GB"
)
engine.update_weights_from_disk(model_name)
# destroy the hf model
torch.cuda.empty_cache()
print("generate (#2)")
outputs = engine.generate(
params["prompt_hybrid_mamba"], params["sampling_params_hybrid_mamba"]
)["text"]
self.assertEqual(
outputs, params["expect_output_after_update_weights_hybrid_mamba"]
)
engine.shutdown()
if __name__ == "__main__":
unittest.main()