wehub-resource-sync
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from deepspeed.moe.fused_expert_layout import classify_fused_gate_up_layout
|
|
|
|
|
|
def test_classifies_gate_up_first():
|
|
layout = classify_fused_gate_up_layout((8, 256, 64), (8, 64, 128))
|
|
|
|
assert layout is not None
|
|
assert layout.layout == "gate_up_first"
|
|
assert layout.hidden_size == 64
|
|
assert layout.ffn_hidden_size == 128
|
|
assert layout.needs_transpose is False
|
|
|
|
|
|
def test_classifies_hidden_first():
|
|
layout = classify_fused_gate_up_layout((8, 64, 256), (8, 128, 64))
|
|
|
|
assert layout is not None
|
|
assert layout.layout == "hidden_first"
|
|
assert layout.hidden_size == 64
|
|
assert layout.ffn_hidden_size == 128
|
|
assert layout.needs_transpose is True
|
|
|
|
|
|
def test_returns_none_for_unknown_layout():
|
|
assert classify_fused_gate_up_layout((8, 64, 64), (8, 64, 64)) is None
|
|
|
|
|
|
def test_returns_none_for_odd_inner_dim():
|
|
assert classify_fused_gate_up_layout((8, 255, 64), (8, 64, 127)) is None
|
|
|
|
|
|
def test_returns_none_for_non_3d():
|
|
assert classify_fused_gate_up_layout((64, 64), (64, 64)) is None
|