[core] fix sleep mode and pytorch checkpoint compatibility (#13001)

Signed-off-by: youkaichao <youkaichao@gmail.com>
This commit is contained in:
youkaichao 2025-02-10 13:03:43 +08:00 committed by GitHub
parent 44607e07d3
commit b2496bb07f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 8 additions and 3 deletions

View File

@ -115,10 +115,16 @@ def test_cumem_with_cudagraph():
@fork_new_process_for_each_test
def test_end_to_end():
@pytest.mark.parametrize(
"model",
[
"meta-llama/Llama-3.2-1B", # sleep mode with safetensors
"facebook/opt-125m" # sleep mode with pytorch checkpoint
])
def test_end_to_end(model):
free, total = torch.cuda.mem_get_info()
used_bytes_baseline = total - free # in case other process is running
llm = LLM("meta-llama/Llama-3.2-1B", enable_sleep_mode=True)
llm = LLM(model, enable_sleep_mode=True)
prompt = "How are you?"
sampling_params = SamplingParams(temperature=0, max_tokens=10)
output = llm.generate(prompt, sampling_params)

View File

@ -462,7 +462,6 @@ def pt_weights_iterator(
state = torch.load(bin_file, map_location="cpu", weights_only=True)
yield from state.items()
del state
torch.cuda.empty_cache()
def get_gguf_extra_tensor_names(