Add update_weights_from_disk endpoint to Engine (#4102)

Co-authored-by: zhaochenyang20 <zhaochen20@outlook.com>
This commit is contained in:
Jhin
2025-03-05 14:25:18 -06:00
committed by GitHub
parent ef9d3b3c2c
commit 70b3c6eeb1
3 changed files with 239 additions and 31 deletions

View File

@@ -15,6 +15,7 @@ distributed setup.
import gc
import os
import random
import time
import unittest
@@ -529,8 +530,9 @@ class TestUpdateWeightsFromDistributed(unittest.TestCase):
assert torch.cuda.device_count() >= 2, "At least 2 GPUs are required"
# test_suits : tp, dp, model_name, backend
if is_in_ci():
mode = random.choice(["Engine", "Server"])
test_suits = [
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, "Engine"),
(1, 1, DEFAULT_SMALL_MODEL_NAME_FOR_TEST, mode),
]
else:
test_suits = [