LoRA support for qwen3.5 and nemotron3 (#23594)
Co-authored-by: Yanbin Jiang <jybsuper@gmail.com>
This commit is contained in:
Opher Lieber
co-authored by
Yanbin Jiang
Yanbin Jiang
parent
0b1fbdba15
commit
c8c1c9261d
@@ -614,6 +614,110 @@ class TestChunkedSGMV(unittest.TestCase):
|
||||
f"QKV missing k_proj batch_size={batch_size}",
|
||||
)
|
||||
|
||||
def test_4_slice_gdn_qkvz(self):
|
||||
"""Test 4-slice shrink+expand operations (GDN in_proj_qkvz)."""
|
||||
num_slices = 4
|
||||
# GDN-style: 4 slices with different sizes [2048, 2048, 4096, 4096]
|
||||
slice_offsets = torch.tensor(
|
||||
[0, 2048, 4096, 8192, 12288], dtype=torch.int32, device=self.device
|
||||
)
|
||||
total_out = 12288
|
||||
max_slice_size = 4096
|
||||
|
||||
for batch_size in [1, 2, 16]:
|
||||
with self.subTest(batch_size=batch_size):
|
||||
# Build batch (reuse the 3-slice helper just for x, batch_info, etc.)
|
||||
x, _, batch_info, seq_lengths, lora_assignments = (
|
||||
self.create_test_batch(BatchComposition.MIXED, batch_size)
|
||||
)
|
||||
|
||||
# Build 4-slice LoRA weights and stack them manually
|
||||
lora_names = [n for n in self.lora_configs if n != "_NO_LORA_"]
|
||||
max_rank = max(self.lora_configs[n][0] for n in lora_names)
|
||||
stacked_a = torch.zeros(
|
||||
len(lora_names),
|
||||
num_slices * max_rank,
|
||||
self.input_dim,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
stacked_b = torch.zeros(
|
||||
len(lora_names),
|
||||
total_out,
|
||||
max_rank,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
for i, name in enumerate(lora_names):
|
||||
rank = self.lora_configs[name][0]
|
||||
if rank > 0:
|
||||
stacked_a[i, : num_slices * rank, :] = torch.randn(
|
||||
num_slices * rank,
|
||||
self.input_dim,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
stacked_b[i, :, :rank] = torch.randn(
|
||||
total_out, rank, dtype=self.dtype, device=self.device
|
||||
)
|
||||
|
||||
lora_assignments_tensor = torch.tensor(
|
||||
lora_assignments, dtype=torch.int32, device="cpu"
|
||||
)
|
||||
seq_lengths_tensor = torch.tensor(
|
||||
seq_lengths, dtype=torch.int32, device="cpu"
|
||||
)
|
||||
lora_ranks_tensor = batch_info.lora_ranks.detach().cpu()
|
||||
scalings_tensor = batch_info.scalings.detach().cpu()
|
||||
|
||||
# Shrink
|
||||
chunked_shrink = chunked_sgmv_lora_shrink_forward(
|
||||
x, stacked_a, batch_info, num_slices=num_slices
|
||||
)
|
||||
reference_shrink = reference_sgmv_shrink(
|
||||
x,
|
||||
stacked_a,
|
||||
lora_assignments_tensor,
|
||||
seq_lengths_tensor,
|
||||
lora_ranks_tensor,
|
||||
scalings_tensor,
|
||||
num_slices=num_slices,
|
||||
)
|
||||
self._compare_shrink_outputs(
|
||||
chunked_shrink,
|
||||
reference_shrink,
|
||||
seq_lengths,
|
||||
lora_assignments,
|
||||
batch_info,
|
||||
num_slices=num_slices,
|
||||
test_name=f"4-slice shrink bs={batch_size}",
|
||||
)
|
||||
|
||||
# Expand
|
||||
chunked_expand = chunked_sgmv_lora_expand_forward(
|
||||
reference_shrink,
|
||||
stacked_b,
|
||||
batch_info,
|
||||
slice_offsets,
|
||||
max_slice_size,
|
||||
base_output=None,
|
||||
)
|
||||
reference_expand = reference_sgmv_expand(
|
||||
reference_shrink,
|
||||
stacked_b,
|
||||
lora_assignments_tensor,
|
||||
seq_lengths_tensor,
|
||||
lora_ranks_tensor,
|
||||
slice_offsets,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
chunked_expand,
|
||||
reference_expand,
|
||||
rtol=self.RTOL,
|
||||
atol=self.ATOL,
|
||||
msg=f"4-slice expand failed bs={batch_size}",
|
||||
)
|
||||
|
||||
# === Batch Composition Tests ===
|
||||
|
||||
def test_uniform_lora_batch(self):
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
# Copyright 2023-2025 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
"""
|
||||
Regression test for NVIDIA-Nemotron-3-Super-120B-A12B-BF16 LoRA logprob accuracy.
|
||||
|
||||
Compares SGLang LoRA logprobs against reference training logprobs from a
|
||||
pre-computed dataset. The LoRA adapter and reference data are downloaded from:
|
||||
https://huggingface.co/datasets/opherlie/lora-test-case-NVIDIA-Nemotron-3-Super-120B-A12B-BF16
|
||||
|
||||
Usage:
|
||||
python -m unittest test_lora_nemotron_3_super_120b_a12b_logprob_diff
|
||||
"""
|
||||
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=300,
|
||||
suite="stage-c-test-4-gpu-b200-small",
|
||||
)
|
||||
|
||||
BASE_MODEL = "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-BF16"
|
||||
LORA_HF_REPO = "opherlie/lora-test-case-NVIDIA-Nemotron-3-Super-120B-A12B-BF16"
|
||||
LORA_BACKEND = "triton"
|
||||
MAX_LORA_RANK = 64
|
||||
TP_SIZE = 4
|
||||
MOE_RUNNER_BACKEND = "triton"
|
||||
EXPERTS_SHARED_OUTER_LORAS = True
|
||||
LORA_USE_VIRTUAL_EXPERTS = True
|
||||
DISABLE_SHARED_EXPERTS_FUSION = True
|
||||
|
||||
KL_THRESHOLD = 2.5e-3
|
||||
|
||||
|
||||
def kl_v2(a, b):
|
||||
a = torch.tensor(a) if not torch.is_tensor(a) else a
|
||||
b = torch.tensor(b) if not torch.is_tensor(b) else b
|
||||
return (((a - b) ** 2) * 0.5).mean().item()
|
||||
|
||||
|
||||
def get_prompt_logprobs(engine, input_ids, lora_path):
|
||||
if isinstance(input_ids, torch.Tensor):
|
||||
input_ids = [input_ids.tolist()]
|
||||
elif not isinstance(input_ids[0], list):
|
||||
input_ids = [input_ids]
|
||||
out = engine.generate(
|
||||
input_ids=input_ids,
|
||||
sampling_params={"max_new_tokens": 0, "temperature": 0.0},
|
||||
return_logprob=True,
|
||||
logprob_start_len=0,
|
||||
lora_path=lora_path,
|
||||
)
|
||||
if isinstance(out, list):
|
||||
out = out[0]
|
||||
return [logprob for logprob, _, _ in out["meta_info"]["input_token_logprobs"]][1:]
|
||||
|
||||
|
||||
class TestLoRANemotron3Super120B_A12B_LogprobDiff(CustomTestCase):
|
||||
|
||||
def test_lora_nemotron_3_super_120b_a12b_logprob_accuracy(self):
|
||||
adapter_path = snapshot_download(
|
||||
LORA_HF_REPO,
|
||||
repo_type="dataset",
|
||||
)
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=BASE_MODEL,
|
||||
tp_size=TP_SIZE,
|
||||
enable_lora=True,
|
||||
max_lora_rank=MAX_LORA_RANK,
|
||||
lora_paths={"my_lora": adapter_path},
|
||||
lora_backend=LORA_BACKEND,
|
||||
moe_runner_backend=MOE_RUNNER_BACKEND,
|
||||
experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS,
|
||||
lora_use_virtual_experts=LORA_USE_VIRTUAL_EXPERTS,
|
||||
disable_shared_experts_fusion=DISABLE_SHARED_EXPERTS_FUSION,
|
||||
)
|
||||
|
||||
try:
|
||||
cdata = torch.load(
|
||||
os.path.join(adapter_path, "compare_sample_train_data.pt"),
|
||||
weights_only=False,
|
||||
)
|
||||
|
||||
base_logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path=None)
|
||||
logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path="my_lora")
|
||||
|
||||
base_t = torch.tensor(base_logprobs)
|
||||
lora_t = torch.tensor(logprobs)
|
||||
diff = (base_t - lora_t).abs()
|
||||
print(
|
||||
f"[VERIFY] base vs lora: mean_diff={diff.mean().item():.6f}, "
|
||||
f"max_diff={diff.max().item():.6f}, "
|
||||
f"identical={torch.equal(base_t, lora_t)}"
|
||||
)
|
||||
|
||||
self.assertFalse(
|
||||
torch.equal(base_t, lora_t),
|
||||
"LoRA logprobs should differ from base model logprobs",
|
||||
)
|
||||
|
||||
kl_sglang_trainer = kl_v2(cdata["training_logprobs"], logprobs)
|
||||
kl_orig_trainer = kl_v2(
|
||||
cdata["training_logprobs"], cdata["sampling_logprobs"]
|
||||
)
|
||||
kl_sglang_orig = kl_v2(logprobs, cdata["sampling_logprobs"])
|
||||
|
||||
print(f"KL(orig_sampler, trainer) = {kl_orig_trainer:.6e}")
|
||||
print(f"KL(sglang, trainer) = {kl_sglang_trainer:.6e}")
|
||||
print(f"KL(sglang, orig_sampler) = {kl_sglang_orig:.6e}")
|
||||
|
||||
self.assertLessEqual(
|
||||
kl_sglang_trainer,
|
||||
KL_THRESHOLD,
|
||||
f"KL(sglang, trainer) = {kl_sglang_trainer:.6e} exceeds "
|
||||
f"threshold {KL_THRESHOLD}",
|
||||
)
|
||||
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
mp.set_start_method("spawn")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
unittest.main(warnings="ignore", verbosity=2)
|
||||
finally:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
@@ -0,0 +1,157 @@
|
||||
# Copyright 2023-2025 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
"""
|
||||
Regression test for Qwen3.5-35B-A3B LoRA logprob accuracy.
|
||||
|
||||
Compares SGLang LoRA logprobs against reference training logprobs from a
|
||||
pre-computed dataset. The LoRA adapter and reference data are downloaded from:
|
||||
https://huggingface.co/datasets/opherlie/lora-test-case-Qwen3.5-35B-A3B
|
||||
|
||||
Usage:
|
||||
python -m unittest test_lora_qwen3_5_35b_a3b_logprob_diff
|
||||
"""
|
||||
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=160,
|
||||
suite="stage-c-test-4-gpu-b200-small",
|
||||
)
|
||||
|
||||
BASE_MODEL = "Qwen/Qwen3.5-35B-A3B"
|
||||
LORA_HF_REPO = "opherlie/lora-test-case-Qwen3.5-35B-A3B"
|
||||
LORA_BACKEND = "triton"
|
||||
MAX_LORA_RANK = 64
|
||||
TP_SIZE = 4
|
||||
MOE_RUNNER_BACKEND = "triton"
|
||||
EXPERTS_SHARED_OUTER_LORAS = True
|
||||
LORA_USE_VIRTUAL_EXPERTS = True
|
||||
DISABLE_SHARED_EXPERTS_FUSION = True
|
||||
|
||||
KL_THRESHOLD = 1e-3
|
||||
|
||||
|
||||
def kl_v2(a, b):
|
||||
a = torch.tensor(a) if not torch.is_tensor(a) else a
|
||||
b = torch.tensor(b) if not torch.is_tensor(b) else b
|
||||
return (((a - b) ** 2) * 0.5).mean().item()
|
||||
|
||||
|
||||
def get_prompt_logprobs(engine, input_ids, lora_path):
|
||||
if isinstance(input_ids, torch.Tensor):
|
||||
input_ids = [input_ids.tolist()]
|
||||
elif not isinstance(input_ids[0], list):
|
||||
input_ids = [input_ids]
|
||||
out = engine.generate(
|
||||
input_ids=input_ids,
|
||||
sampling_params={"max_new_tokens": 0, "temperature": 0.0},
|
||||
return_logprob=True,
|
||||
logprob_start_len=0,
|
||||
lora_path=lora_path,
|
||||
)
|
||||
if isinstance(out, list):
|
||||
out = out[0]
|
||||
return [logprob for logprob, _, _ in out["meta_info"]["input_token_logprobs"]][1:]
|
||||
|
||||
|
||||
class TestLoRAQwen3_5_35B_A3B_LogprobDiff(CustomTestCase):
|
||||
|
||||
def test_lora_qwen3_5_35b_a3b_logprob_accuracy(self):
|
||||
adapter_path = snapshot_download(
|
||||
LORA_HF_REPO,
|
||||
repo_type="dataset",
|
||||
)
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=BASE_MODEL,
|
||||
tp_size=TP_SIZE,
|
||||
enable_lora=True,
|
||||
max_lora_rank=MAX_LORA_RANK,
|
||||
lora_paths={"my_lora": adapter_path},
|
||||
lora_backend=LORA_BACKEND,
|
||||
moe_runner_backend=MOE_RUNNER_BACKEND,
|
||||
experts_shared_outer_loras=EXPERTS_SHARED_OUTER_LORAS,
|
||||
lora_use_virtual_experts=LORA_USE_VIRTUAL_EXPERTS,
|
||||
disable_shared_experts_fusion=DISABLE_SHARED_EXPERTS_FUSION,
|
||||
# OOM's on logits with the defaults here, as this test-case is longer context and qwen3.5 has larger vocab
|
||||
chunked_prefill_size=8192,
|
||||
mem_fraction_static=0.8,
|
||||
)
|
||||
|
||||
try:
|
||||
cdata = torch.load(
|
||||
os.path.join(adapter_path, "compare_sample_train_data.pt"),
|
||||
weights_only=False,
|
||||
)
|
||||
|
||||
base_logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path=None)
|
||||
logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path="my_lora")
|
||||
|
||||
base_t = torch.tensor(base_logprobs)
|
||||
lora_t = torch.tensor(logprobs)
|
||||
diff = (base_t - lora_t).abs()
|
||||
print(
|
||||
f"[VERIFY] base vs lora: mean_diff={diff.mean().item():.6f}, "
|
||||
f"max_diff={diff.max().item():.6f}, "
|
||||
f"identical={torch.equal(base_t, lora_t)}"
|
||||
)
|
||||
|
||||
self.assertFalse(
|
||||
torch.equal(base_t, lora_t),
|
||||
"LoRA logprobs should differ from base model logprobs",
|
||||
)
|
||||
|
||||
kl_sglang_trainer = kl_v2(cdata["training_logprobs"], logprobs)
|
||||
kl_orig_trainer = kl_v2(
|
||||
cdata["training_logprobs"], cdata["sampling_logprobs"]
|
||||
)
|
||||
kl_sglang_orig = kl_v2(logprobs, cdata["sampling_logprobs"])
|
||||
|
||||
print(f"KL(orig_sampler, trainer) = {kl_orig_trainer:.6e}")
|
||||
print(f"KL(sglang, trainer) = {kl_sglang_trainer:.6e}")
|
||||
print(f"KL(sglang, orig_sampler) = {kl_sglang_orig:.6e}")
|
||||
|
||||
self.assertLessEqual(
|
||||
kl_sglang_trainer,
|
||||
KL_THRESHOLD,
|
||||
f"KL(sglang, trainer) = {kl_sglang_trainer:.6e} exceeds "
|
||||
f"threshold {KL_THRESHOLD}",
|
||||
)
|
||||
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
mp.set_start_method("spawn")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
unittest.main(warnings="ignore", verbosity=2)
|
||||
finally:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
@@ -0,0 +1,146 @@
|
||||
# Copyright 2023-2025 SGLang Team
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
|
||||
"""
|
||||
Regression test for Qwen3.5-4B LoRA logprob accuracy.
|
||||
|
||||
Compares SGLang LoRA logprobs against reference training logprobs from a
|
||||
pre-computed dataset. The LoRA adapter and reference data are downloaded from:
|
||||
https://huggingface.co/datasets/opherlie/lora-test-case-Qwen3.5-4B
|
||||
|
||||
Usage:
|
||||
python -m unittest test_lora_qwen3_5_4b_logprob_diff
|
||||
"""
|
||||
|
||||
import multiprocessing as mp
|
||||
import os
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=90,
|
||||
suite="stage-b-test-1-gpu-large",
|
||||
)
|
||||
|
||||
BASE_MODEL = "Qwen/Qwen3.5-4B"
|
||||
LORA_HF_REPO = "opherlie/lora-test-case-Qwen3.5-4B"
|
||||
LORA_BACKEND = "triton"
|
||||
MAX_LORA_RANK = 64
|
||||
TP_SIZE = 1
|
||||
|
||||
KL_THRESHOLD = 4e-3
|
||||
|
||||
|
||||
def kl_v2(a, b):
|
||||
a = torch.tensor(a) if not torch.is_tensor(a) else a
|
||||
b = torch.tensor(b) if not torch.is_tensor(b) else b
|
||||
return (((a - b) ** 2) * 0.5).mean().item()
|
||||
|
||||
|
||||
def get_prompt_logprobs(engine, input_ids, lora_path):
|
||||
if isinstance(input_ids, torch.Tensor):
|
||||
input_ids = [input_ids.tolist()]
|
||||
elif not isinstance(input_ids[0], list):
|
||||
input_ids = [input_ids]
|
||||
out = engine.generate(
|
||||
input_ids=input_ids,
|
||||
sampling_params={"max_new_tokens": 0, "temperature": 0.0},
|
||||
return_logprob=True,
|
||||
logprob_start_len=0,
|
||||
lora_path=lora_path,
|
||||
)
|
||||
if isinstance(out, list):
|
||||
out = out[0]
|
||||
return [logprob for logprob, _, _ in out["meta_info"]["input_token_logprobs"]][1:]
|
||||
|
||||
|
||||
class TestLoRAQwen3_5_4BLogprobDiff(CustomTestCase):
|
||||
|
||||
def test_lora_qwen3_5_4b_logprob_accuracy(self):
|
||||
adapter_path = snapshot_download(
|
||||
LORA_HF_REPO,
|
||||
repo_type="dataset",
|
||||
)
|
||||
|
||||
engine = sgl.Engine(
|
||||
model_path=BASE_MODEL,
|
||||
tp_size=TP_SIZE,
|
||||
enable_lora=True,
|
||||
max_lora_rank=MAX_LORA_RANK,
|
||||
lora_paths={"my_lora": adapter_path},
|
||||
lora_backend=LORA_BACKEND,
|
||||
)
|
||||
|
||||
try:
|
||||
cdata = torch.load(
|
||||
os.path.join(adapter_path, "compare_sample_train_data.pt"),
|
||||
weights_only=False,
|
||||
)
|
||||
|
||||
base_logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path=None)
|
||||
logprobs = get_prompt_logprobs(engine, cdata["tokens"], lora_path="my_lora")
|
||||
|
||||
base_t = torch.tensor(base_logprobs)
|
||||
lora_t = torch.tensor(logprobs)
|
||||
diff = (base_t - lora_t).abs()
|
||||
print(
|
||||
f"[VERIFY] base vs lora: mean_diff={diff.mean().item():.6f}, "
|
||||
f"max_diff={diff.max().item():.6f}, "
|
||||
f"identical={torch.equal(base_t, lora_t)}"
|
||||
)
|
||||
|
||||
self.assertFalse(
|
||||
torch.equal(base_t, lora_t),
|
||||
"LoRA logprobs should differ from base model logprobs",
|
||||
)
|
||||
|
||||
kl_sglang_trainer = kl_v2(cdata["training_logprobs"], logprobs)
|
||||
kl_orig_trainer = kl_v2(
|
||||
cdata["training_logprobs"], cdata["sampling_logprobs"]
|
||||
)
|
||||
kl_sglang_orig = kl_v2(logprobs, cdata["sampling_logprobs"])
|
||||
|
||||
print(f"KL(orig_sampler, trainer) = {kl_orig_trainer:.6e}")
|
||||
print(f"KL(sglang, trainer) = {kl_sglang_trainer:.6e}")
|
||||
print(f"KL(sglang, orig_sampler) = {kl_sglang_orig:.6e}")
|
||||
|
||||
self.assertLessEqual(
|
||||
kl_sglang_trainer,
|
||||
KL_THRESHOLD,
|
||||
f"KL(sglang, trainer) = {kl_sglang_trainer:.6e} exceeds "
|
||||
f"threshold {KL_THRESHOLD}",
|
||||
)
|
||||
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
mp.set_start_method("spawn")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
unittest.main(warnings="ignore", verbosity=2)
|
||||
finally:
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
Reference in New Issue
Block a user