[Test] Consolidate test cleanup and CI taxonomy (net -11.4K lines) (#37436)
Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
This commit is contained in:
Xiaoyu Zhang
co-authored by
Mick Qian
Mick Qian
parent
6a1ff90f2d
commit
4d23a4fa6d
@@ -802,269 +802,6 @@ class TestReduceSum:
|
||||
without_dim_names(unsharder_result.tensors[0]), without_dim_names(expected)
|
||||
)
|
||||
|
||||
def test_recompute_pseudo_mismatch(self) -> None:
|
||||
"""_verify_replicated_group returns failed check for RECOMPUTE_PSEUDO axis mismatch."""
|
||||
tensor_a = torch.ones(4)
|
||||
tensor_b = torch.ones(4) + 0.1
|
||||
|
||||
checks: list[ReplicatedCheckResult] = _verify_replicated_group(
|
||||
[tensor_a, tensor_b],
|
||||
axis=ParallelAxis.RECOMPUTE_PSEUDO,
|
||||
group_index=0,
|
||||
)
|
||||
assert len(checks) == 1
|
||||
assert checks[0].axis == "recompute_pseudo"
|
||||
assert checks[0].group_index == 0
|
||||
assert checks[0].compared_index == 1
|
||||
assert checks[0].baseline_index == 0
|
||||
assert not checks[0].passed
|
||||
assert checks[0].diff.max_abs_diff == pytest.approx(0.1, abs=1e-5)
|
||||
|
||||
|
||||
class TestThdCpConcat:
|
||||
def test_single_seq(self) -> None:
|
||||
"""Single seq THD unshard: 2 ranks → per-seq concat."""
|
||||
rank0 = apply_dim_names(torch.tensor([1, 2, 3]), ["t"])
|
||||
rank1 = apply_dim_names(torch.tensor([4, 5, 6]), ["t"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
expected = torch.tensor([1, 2, 3, 4, 5, 6])
|
||||
assert torch.equal(without_dim_names(unsharder_result.tensors[0]), expected)
|
||||
|
||||
def test_multi_seq(self) -> None:
|
||||
"""Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
|
||||
# rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
|
||||
# rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
|
||||
seq_a_r0 = torch.arange(0, 50)
|
||||
seq_b_r0 = torch.arange(100, 132)
|
||||
pad_r0 = torch.full((46,), -1)
|
||||
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0, pad_r0]), ["t"])
|
||||
|
||||
seq_a_r1 = torch.arange(50, 100)
|
||||
seq_b_r1 = torch.arange(132, 164)
|
||||
pad_r1 = torch.full((46,), -2)
|
||||
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1, pad_r1]), ["t"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
|
||||
|
||||
# seqA: r0(50) + r1(50) = 100 tokens, values 0..99
|
||||
assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
|
||||
# seqB: r0(32) + r1(32) = 64 tokens
|
||||
assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
|
||||
# pad: r0(46) + r1(46) = 92 tokens
|
||||
assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
|
||||
|
||||
def test_with_hidden_dim(self) -> None:
|
||||
"""THD unshard with trailing hidden dim: shape [T, H]."""
|
||||
torch.manual_seed(42)
|
||||
hidden: int = 4
|
||||
# rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
|
||||
# rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
|
||||
seq_a_r0 = torch.randn(3, hidden)
|
||||
seq_b_r0 = torch.randn(2, hidden)
|
||||
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0]), ["t", "h"])
|
||||
|
||||
seq_a_r1 = torch.randn(3, hidden)
|
||||
seq_b_r1 = torch.randn(2, hidden)
|
||||
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1]), ["t", "h"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
|
||||
|
||||
assert unsharded.shape == (10, hidden)
|
||||
assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
|
||||
assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
|
||||
|
||||
def test_with_leading_batch_dim(self) -> None:
|
||||
"""THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
|
||||
torch.manual_seed(42)
|
||||
batch: int = 2
|
||||
hidden: int = 4
|
||||
# rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
|
||||
# rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
|
||||
seq_a_r0 = torch.randn(batch, 3, hidden)
|
||||
seq_b_r0 = torch.randn(batch, 2, hidden)
|
||||
rank0 = apply_dim_names(torch.cat([seq_a_r0, seq_b_r0], dim=1), ["b", "t", "h"])
|
||||
|
||||
seq_a_r1 = torch.randn(batch, 3, hidden)
|
||||
seq_b_r1 = torch.randn(batch, 2, hidden)
|
||||
rank1 = apply_dim_names(torch.cat([seq_a_r1, seq_b_r1], dim=1), ["b", "t", "h"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.CP,
|
||||
params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, [rank0, rank1])
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
unsharded: torch.Tensor = without_dim_names(unsharder_result.tensors[0])
|
||||
|
||||
assert unsharded.shape == (batch, 10, hidden)
|
||||
# seqA: r0(3) + r1(3) = 6 tokens per batch
|
||||
assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
|
||||
# seqB: r0(2) + r1(2) = 4 tokens per batch
|
||||
assert torch.equal(
|
||||
unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
|
||||
)
|
||||
|
||||
|
||||
class TestReduceSum:
|
||||
def test_basic_tp2_reduce(self) -> None:
|
||||
"""2 partial tensors sum to full tensor."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
part_a = full_tensor * 0.6
|
||||
part_b = full_tensor * 0.4
|
||||
|
||||
dim_specs = parse_dims("h[tp:partial] d").dims
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
assert isinstance(plans[0].params, ReduceSumParams)
|
||||
|
||||
named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(
|
||||
plans[0], named_parts
|
||||
)
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
assert torch.allclose(
|
||||
without_dim_names(unsharder_result.tensors[0]), full_tensor
|
||||
)
|
||||
|
||||
def test_tp4_reduce(self) -> None:
|
||||
"""4 partial tensors sum to full tensor."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)]
|
||||
|
||||
dim_specs = parse_dims("h[tp:partial] d").dims
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
|
||||
]
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 1
|
||||
|
||||
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(
|
||||
plans[0], named_parts
|
||||
)
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
assert torch.allclose(
|
||||
without_dim_names(unsharder_result.tensors[0]), full_tensor
|
||||
)
|
||||
|
||||
def test_multi_axis_concat_then_reduce(self) -> None:
|
||||
"""CP concat + TP reduce end-to-end."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8, 16)
|
||||
|
||||
cp_chunks = list(full_tensor.chunk(2, dim=1))
|
||||
# Each CP chunk is held as partial sums across TP ranks
|
||||
tensors: list[torch.Tensor] = []
|
||||
parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
|
||||
for cp_rank in range(2):
|
||||
for tp_rank in range(2):
|
||||
tensors.append(cp_chunks[cp_rank] * 0.5)
|
||||
parallel_infos.append(
|
||||
{
|
||||
ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
|
||||
ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
|
||||
}
|
||||
)
|
||||
|
||||
dim_specs = parse_dims("b s[cp] h[tp:partial]").dims
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
assert len(plans) == 2
|
||||
|
||||
current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
|
||||
for plan in plans:
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(plan, current)
|
||||
current = unsharder_result.tensors
|
||||
|
||||
assert len(current) == 1
|
||||
assert torch.allclose(without_dim_names(current[0]), full_tensor)
|
||||
|
||||
def test_reduce_scrambled_ranks(self) -> None:
|
||||
"""Scrambled rank order — sum is commutative so result is the same."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
parts: list[torch.Tensor] = [
|
||||
full_tensor * 0.1,
|
||||
full_tensor * 0.2,
|
||||
full_tensor * 0.3,
|
||||
full_tensor * 0.4,
|
||||
]
|
||||
|
||||
parallel_infos = [
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
|
||||
{ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
|
||||
]
|
||||
dim_specs = parse_dims("h[tp:partial] d").dims
|
||||
plans = compute_unsharder_plan(dim_specs, parallel_infos)
|
||||
|
||||
named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(
|
||||
plans[0], named_parts
|
||||
)
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
assert torch.allclose(
|
||||
without_dim_names(unsharder_result.tensors[0]), full_tensor
|
||||
)
|
||||
|
||||
def test_reduce_preserves_named_dims(self) -> None:
|
||||
"""Named tensor dimensions are preserved through reduce_sum."""
|
||||
dim_specs = parse_dims("h[tp:partial] d").dims
|
||||
part_a = apply_dim_names(torch.randn(4, 8), ["h", "d"])
|
||||
part_b = apply_dim_names(torch.randn(4, 8), ["h", "d"])
|
||||
|
||||
plan = UnsharderPlan(
|
||||
axis=ParallelAxis.TP,
|
||||
params=ReduceSumParams(),
|
||||
groups=[[0, 1]],
|
||||
)
|
||||
unsharder_result: UnsharderResult = execute_unsharder_plan(
|
||||
plan, [part_a, part_b]
|
||||
)
|
||||
|
||||
assert len(unsharder_result.tensors) == 1
|
||||
assert get_dim_names(unsharder_result.tensors[0]) == ("h", "d")
|
||||
expected = apply_dim_names(
|
||||
without_dim_names(part_a) + without_dim_names(part_b), ["h", "d"]
|
||||
)
|
||||
assert torch.allclose(
|
||||
without_dim_names(unsharder_result.tensors[0]), without_dim_names(expected)
|
||||
)
|
||||
|
||||
|
||||
class TestFusedDimExecutor:
|
||||
def test_fused_tp2_concat(self) -> None:
|
||||
|
||||
@@ -19,79 +19,6 @@ from sglang.test.ci.ci_register import register_cpu_ci
|
||||
register_cpu_ci(est_time=20, stage="weekly", runner_config="cpu")
|
||||
|
||||
|
||||
class TestComputeTensorInfo:
|
||||
def test_basic_tensor_returns_correct_shape_and_dtype(self) -> None:
|
||||
tensor = torch.randn(2, 3)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.shape == [2, 3]
|
||||
assert info.dtype == "torch.float32"
|
||||
assert info.stats.mean == pytest.approx(tensor.float().mean().item(), abs=1e-4)
|
||||
|
||||
def test_include_sample_false_returns_none_sample(self) -> None:
|
||||
tensor = torch.randn(2, 3)
|
||||
info = compute_tensor_info(tensor, include_sample=False)
|
||||
assert info.sample is None
|
||||
|
||||
def test_include_sample_true_returns_string_sample(self) -> None:
|
||||
tensor = torch.randn(2, 3)
|
||||
info = compute_tensor_info(tensor, include_sample=True)
|
||||
assert info.sample is not None
|
||||
assert isinstance(info.sample, str)
|
||||
|
||||
def test_empty_tensor_stats_are_zero(self) -> None:
|
||||
tensor = torch.tensor([])
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.stats.mean == 0.0
|
||||
assert info.stats.std == 0.0
|
||||
assert info.shape == [0]
|
||||
|
||||
def test_integer_tensor_converted_to_float_for_stats(self) -> None:
|
||||
"""Integer tensors should be cast to float internally for stats computation."""
|
||||
tensor = torch.tensor([1, 2, 3, 4], dtype=torch.int32)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.dtype == "torch.int32"
|
||||
assert info.stats.mean == pytest.approx(2.5, abs=1e-4)
|
||||
assert info.stats.min == pytest.approx(1.0, abs=1e-4)
|
||||
assert info.stats.max == pytest.approx(4.0, abs=1e-4)
|
||||
|
||||
def test_bfloat16_tensor_shape_and_stats(self) -> None:
|
||||
"""bfloat16 tensors produce correct shape and dtype string."""
|
||||
tensor = torch.ones(3, 4, dtype=torch.bfloat16)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.shape == [3, 4]
|
||||
assert info.dtype == "torch.bfloat16"
|
||||
assert info.stats.mean == pytest.approx(1.0, abs=1e-2)
|
||||
|
||||
def test_multidimensional_shape(self) -> None:
|
||||
"""Shape is preserved for high-rank tensors."""
|
||||
tensor = torch.randn(2, 3, 4, 5)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.shape == [2, 3, 4, 5]
|
||||
|
||||
def test_scalar_tensor(self) -> None:
|
||||
"""Scalar (0-dim) tensor produces empty shape list."""
|
||||
tensor = torch.tensor(3.14)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert info.shape == []
|
||||
assert info.stats.mean == pytest.approx(3.14, abs=1e-4)
|
||||
assert info.stats.min == pytest.approx(3.14, abs=1e-4)
|
||||
assert info.stats.max == pytest.approx(3.14, abs=1e-4)
|
||||
|
||||
def test_include_sample_true_contains_tensor_representation(self) -> None:
|
||||
"""Sample string should contain some recognizable tensor content."""
|
||||
tensor = torch.tensor([1.0, 2.0])
|
||||
info = compute_tensor_info(tensor, include_sample=True)
|
||||
assert info.sample is not None
|
||||
assert "1." in info.sample or "2." in info.sample
|
||||
|
||||
def test_percentiles_present_for_small_tensor(self) -> None:
|
||||
"""Small tensors (< threshold) should have percentile data."""
|
||||
tensor = torch.randn(100)
|
||||
info = compute_tensor_info(tensor)
|
||||
assert len(info.stats.percentiles) > 0
|
||||
assert 50 in info.stats.percentiles
|
||||
|
||||
|
||||
class TestComputeTensorInfo:
|
||||
def test_basic_tensor_returns_correct_shape_and_dtype(self) -> None:
|
||||
tensor = torch.randn(2, 3)
|
||||
|
||||
@@ -389,47 +389,6 @@ class TestTorchSave:
|
||||
assert "skip the tensor" in captured.out
|
||||
|
||||
|
||||
class TestLog:
|
||||
def test_log_format(self):
|
||||
with _capture_stdout() as captured:
|
||||
_log("hello")
|
||||
out = captured.getvalue()
|
||||
assert "hello" in out, out
|
||||
assert "[Dumper, rank=" in out, out
|
||||
assert ", t=" in out, out
|
||||
|
||||
|
||||
class TestCompareTensorsQuick:
|
||||
def test_identical(self):
|
||||
a = torch.tensor([1.0, 2.0, 3.0])
|
||||
s = _compare_tensors_quick(a, a.clone())
|
||||
assert "rel_diff=0" in s, s
|
||||
assert "max_abs=0" in s, s
|
||||
|
||||
def test_diverged(self):
|
||||
a = torch.tensor([1.0, 2.0, 3.0])
|
||||
b = torch.tensor([1.0, 2.0, 4.0]) # last element differs by 1
|
||||
s = _compare_tensors_quick(a, b)
|
||||
assert "max_abs=1" in s, s
|
||||
assert "rel_diff=" in s, s
|
||||
|
||||
def test_shape_mismatch(self):
|
||||
s = _compare_tensors_quick(torch.zeros(3), torch.zeros(4))
|
||||
assert "shape mismatch" in s, s
|
||||
|
||||
def test_dtype_unified(self):
|
||||
s = _compare_tensors_quick(
|
||||
torch.zeros(3, dtype=torch.float32),
|
||||
torch.zeros(3, dtype=torch.float64),
|
||||
)
|
||||
assert "rel_diff=" in s, s
|
||||
assert "max_abs=" in s, s
|
||||
|
||||
def test_empty(self):
|
||||
s = _compare_tensors_quick(torch.zeros(0), torch.zeros(0))
|
||||
assert s == "empty"
|
||||
|
||||
|
||||
class TestCollectiveTimeout:
|
||||
def test_watchdog_fires_on_timeout(self):
|
||||
block_event = threading.Event()
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
@@ -12,20 +14,11 @@ from sglang.srt.layers.linear import LinearBase
|
||||
from sglang.srt.models.qwen2 import Qwen2MLP
|
||||
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
|
||||
from sglang.srt.utils import add_prefix, get_device
|
||||
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.layer_ut_utils import init_single_process_dist
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=9,
|
||||
stage="base-b",
|
||||
runner_config="1-gpu-small",
|
||||
disabled="Test uses pytest-style function without TestCase class - see #17145",
|
||||
)
|
||||
register_amd_ci(
|
||||
est_time=15,
|
||||
suite="stage-b-test-1-gpu-small-amd",
|
||||
disabled="Test uses pytest-style function without TestCase class - see #17145",
|
||||
)
|
||||
register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
TEST_HIDDEN_SIZE = 32
|
||||
|
||||
@@ -73,26 +66,29 @@ def init_weights(module):
|
||||
torch.nn.init.ones_(module.weight)
|
||||
|
||||
|
||||
def test_model_forward_dump(tmp_path):
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
device = get_device()
|
||||
init_single_process_dist(backend=get_default_distributed_backend(device))
|
||||
model = MockCausalLM()
|
||||
model.apply(init_weights)
|
||||
model = model.to(device=device, dtype=torch.bfloat16)
|
||||
dumper = register_forward_hook_for_model(
|
||||
model, tmp_path / "sglang_dump", [0], 0, 0, 0
|
||||
)
|
||||
class TestTensorDumpForwardHook(CustomTestCase):
|
||||
def test_model_forward_dump(self):
|
||||
set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
|
||||
device = get_device()
|
||||
init_single_process_dist(backend=get_default_distributed_backend(device))
|
||||
model = MockCausalLM()
|
||||
model.apply(init_weights)
|
||||
model = model.to(device=device, dtype=torch.bfloat16)
|
||||
|
||||
dir_path = dumper.get_dump_dir()
|
||||
inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01
|
||||
result = model(inp.to(device))
|
||||
data = torch.load(f"{dir_path}/Pass00000.pt")
|
||||
assert "model.layernorm" in data
|
||||
assert "model.mlp.down_proj" in data
|
||||
assert torch.allclose(
|
||||
data["model.mlp.down_proj"], result.cpu(), rtol=1e-5, atol=1e-5
|
||||
)
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
dumper = register_forward_hook_for_model(
|
||||
model, Path(temp_dir) / "sglang_dump", [0], 0, 0, 0
|
||||
)
|
||||
dir_path = dumper.get_dump_dir()
|
||||
inp = torch.randn(4, TEST_HIDDEN_SIZE, dtype=torch.bfloat16) * 0.01
|
||||
result = model(inp.to(device))
|
||||
data = torch.load(f"{dir_path}/Pass00000.pt")
|
||||
|
||||
self.assertIn("model.layernorm", data)
|
||||
self.assertIn("model.mlp.down_proj", data)
|
||||
torch.testing.assert_close(
|
||||
data["model.mlp.down_proj"], result.cpu(), rtol=1e-5, atol=1e-5
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user