From 679fc11a3ee52b4fb9c4f8014813c39b72fa9e63 Mon Sep 17 00:00:00 2001 From: Yu Zhang Date: Wed, 25 Jun 2025 05:31:25 +0800 Subject: [PATCH 1/5] Refactor modeling tests --- tests/models/test_modeling_abc.py | 52 +++++++++++++----- tests/models/test_modeling_bitnet.py | 52 +++++++++++++----- tests/models/test_modeling_comba.py | 52 +++++++++++++----- tests/models/test_modeling_deltanet.py | 52 +++++++++++++----- .../test_modeling_forgetting_transformer.py | 52 +++++++++++++----- tests/models/test_modeling_gated_deltanet.py | 51 ++++++++++++----- .../test_modeling_gated_deltaproduct.py | 48 +++++++++++----- tests/models/test_modeling_gla.py | 52 +++++++++++++----- tests/models/test_modeling_gsa.py | 52 +++++++++++++----- tests/models/test_modeling_hgrn.py | 52 +++++++++++++----- tests/models/test_modeling_hgrn2.py | 52 +++++++++++++----- tests/models/test_modeling_lightnet.py | 52 +++++++++++++----- tests/models/test_modeling_linear_attn.py | 52 +++++++++++++----- tests/models/test_modeling_mamba.py | 52 +++++++++++++----- tests/models/test_modeling_mamba2.py | 52 +++++++++++++----- tests/models/test_modeling_mesanet.py | 52 +++++++++++++----- tests/models/test_modeling_nsa.py | 52 +++++++++++++----- tests/models/test_modeling_path_attn.py | 52 +++++++++++++----- tests/models/test_modeling_retnet.py | 52 +++++++++++++----- tests/models/test_modeling_rodimus.py | 52 +++++++++++++----- tests/models/test_modeling_rwkv6.py | 52 +++++++++++++----- tests/models/test_modeling_rwkv7.py | 55 +++++++++++++------ tests/models/test_modeling_samba.py | 52 +++++++++++++----- tests/models/test_modeling_transformer.py | 52 +++++++++++++----- 24 files changed, 887 insertions(+), 359 deletions(-) diff --git a/tests/models/test_modeling_abc.py b/tests/models/test_modeling_abc.py index 9bfe28bdd6..8c84b697f0 100644 --- a/tests/models/test_modeling_abc.py +++ b/tests/models/test_modeling_abc.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype', 'use_l2warp'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, ABCConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, ABCConfig, dtype) diff --git a/tests/models/test_modeling_bitnet.py b/tests/models/test_modeling_bitnet.py index 60d9be71a3..c2b03e9d2a 100644 --- a/tests/models/test_modeling_bitnet.py +++ b/tests/models/test_modeling_bitnet.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, BitNetConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, BitNetConfig, dtype) diff --git a/tests/models/test_modeling_comba.py b/tests/models/test_modeling_comba.py index 0209905b84..2bf6137ce7 100644 --- a/tests/models/test_modeling_comba.py +++ b/tests/models/test_modeling_comba.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, CombaConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, CombaConfig, dtype) diff --git a/tests/models/test_modeling_deltanet.py b/tests/models/test_modeling_deltanet.py index 80cebb95d3..6ff3b1915b 100644 --- a/tests/models/test_modeling_deltanet.py +++ b/tests/models/test_modeling_deltanet.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, DeltaNetConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, DeltaNetConfig, dtype) diff --git a/tests/models/test_modeling_forgetting_transformer.py b/tests/models/test_modeling_forgetting_transformer.py index ff5d46aa27..1f0f916381 100644 --- a/tests/models/test_modeling_forgetting_transformer.py +++ b/tests/models/test_modeling_forgetting_transformer.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, ForgettingTransformerConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, ForgettingTransformerConfig, dtype) diff --git a/tests/models/test_modeling_gated_deltanet.py b/tests/models/test_modeling_gated_deltanet.py index 419b5db98a..6b294e4855 100644 --- a/tests/models/test_modeling_gated_deltanet.py +++ b/tests/models/test_modeling_gated_deltanet.py @@ -11,25 +11,46 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, GatedDeltaNetConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, GatedDeltaNetConfig, dtype) diff --git a/tests/models/test_modeling_gated_deltaproduct.py b/tests/models/test_modeling_gated_deltaproduct.py index 9e6665b68d..55b61d7f33 100644 --- a/tests/models/test_modeling_gated_deltaproduct.py +++ b/tests/models/test_modeling_gated_deltaproduct.py @@ -14,25 +14,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, GatedDeltaProductConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [5]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-{}".format(*test)) + for test in [ + (2, 4, 4000, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + dtype: torch.dtype, +): config = GatedDeltaProductConfig() config.num_hidden_layers = L config.use_forget_gate = False diff --git a/tests/models/test_modeling_gla.py b/tests/models/test_modeling_gla.py index 6d257930a5..f3f00c2d70 100644 --- a/tests/models/test_modeling_gla.py +++ b/tests/models/test_modeling_gla.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, GLAConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, GLAConfig, dtype) diff --git a/tests/models/test_modeling_gsa.py b/tests/models/test_modeling_gsa.py index 756c4c8e5f..a57c261c84 100644 --- a/tests/models/test_modeling_gsa.py +++ b/tests/models/test_modeling_gsa.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, GSAConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, GSAConfig, dtype) diff --git a/tests/models/test_modeling_hgrn.py b/tests/models/test_modeling_hgrn.py index 8f19473876..3dcf88936b 100644 --- a/tests/models/test_modeling_hgrn.py +++ b/tests/models/test_modeling_hgrn.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, HGRNConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, HGRNConfig, dtype) diff --git a/tests/models/test_modeling_hgrn2.py b/tests/models/test_modeling_hgrn2.py index e90ab1a980..1bce95a5c9 100644 --- a/tests/models/test_modeling_hgrn2.py +++ b/tests/models/test_modeling_hgrn2.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_hgrn2_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_hgrn2_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, HGRN2Config, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_hgrn2_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, HGRN2Config, dtype) diff --git a/tests/models/test_modeling_lightnet.py b/tests/models/test_modeling_lightnet.py index 63b521b3b3..a43870fd8a 100644 --- a/tests/models/test_modeling_lightnet.py +++ b/tests/models/test_modeling_lightnet.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, LightNetConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, LightNetConfig, dtype) diff --git a/tests/models/test_modeling_linear_attn.py b/tests/models/test_modeling_linear_attn.py index 391affea96..0aa90ec0c9 100644 --- a/tests/models/test_modeling_linear_attn.py +++ b/tests/models/test_modeling_linear_attn.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, LinearAttentionConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, LinearAttentionConfig, dtype) diff --git a/tests/models/test_modeling_mamba.py b/tests/models/test_modeling_mamba.py index 13ea546410..b6d57964be 100644 --- a/tests/models/test_modeling_mamba.py +++ b/tests/models/test_modeling_mamba.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, MambaConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, MambaConfig, dtype) diff --git a/tests/models/test_modeling_mamba2.py b/tests/models/test_modeling_mamba2.py index 2a1669be34..721ee8b690 100644 --- a/tests/models/test_modeling_mamba2.py +++ b/tests/models/test_modeling_mamba2.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_mamba2_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, Mamba2Config, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_mamba2_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, Mamba2Config, dtype) diff --git a/tests/models/test_modeling_mesanet.py b/tests/models/test_modeling_mesanet.py index 5e840a8fe1..a1c66920bf 100644 --- a/tests/models/test_modeling_mesanet.py +++ b/tests/models/test_modeling_mesanet.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, MesaNetConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, MesaNetConfig, dtype) diff --git a/tests/models/test_modeling_nsa.py b/tests/models/test_modeling_nsa.py index 55ca7870df..5dd3d918d7 100644 --- a/tests/models/test_modeling_nsa.py +++ b/tests/models/test_modeling_nsa.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, NSAConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, NSAConfig, dtype) diff --git a/tests/models/test_modeling_path_attn.py b/tests/models/test_modeling_path_attn.py index 817cc5840f..3d01914c1d 100644 --- a/tests/models/test_modeling_path_attn.py +++ b/tests/models/test_modeling_path_attn.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [32, 64]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, PaTHAttentionConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, PaTHAttentionConfig, dtype) diff --git a/tests/models/test_modeling_retnet.py b/tests/models/test_modeling_retnet.py index bb761888c5..4bdd9b7155 100644 --- a/tests/models/test_modeling_retnet.py +++ b/tests/models/test_modeling_retnet.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, RetNetConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, RetNetConfig, dtype) diff --git a/tests/models/test_modeling_rodimus.py b/tests/models/test_modeling_rodimus.py index a9155394bc..baea2c02fd 100644 --- a/tests/models/test_modeling_rodimus.py +++ b/tests/models/test_modeling_rodimus.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, RodimusConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, RodimusConfig, dtype) diff --git a/tests/models/test_modeling_rwkv6.py b/tests/models/test_modeling_rwkv6.py index 34d071495e..d90aff85d9 100644 --- a/tests/models/test_modeling_rwkv6.py +++ b/tests/models/test_modeling_rwkv6.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_rwkv6_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, RWKV6Config, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_rwkv6_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, RWKV6Config, dtype) diff --git a/tests/models/test_modeling_rwkv7.py b/tests/models/test_modeling_rwkv7.py index 4cb80eee32..170d9784e9 100644 --- a/tests/models/test_modeling_rwkv7.py +++ b/tests/models/test_modeling_rwkv7.py @@ -11,25 +11,48 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_rwkv7_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, RWKV7Config, dtype, use_l2warp) - # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_rwkv7_generation(L, B, T, H, D, dtype): + + +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, RWKV7Config, dtype) diff --git a/tests/models/test_modeling_samba.py b/tests/models/test_modeling_samba.py index db7f13a15b..4fb851f3fa 100644 --- a/tests/models/test_modeling_samba.py +++ b/tests/models/test_modeling_samba.py @@ -11,25 +11,47 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, SambaConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [4000]) -@pytest.mark.parametrize("H", [8]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): run_test_generation(L, B, T, H, D, SambaConfig, dtype) diff --git a/tests/models/test_modeling_transformer.py b/tests/models/test_modeling_transformer.py index ae15982b1f..2919bf5252 100644 --- a/tests/models/test_modeling_transformer.py +++ b/tests/models/test_modeling_transformer.py @@ -14,27 +14,49 @@ # =================================================================================== # Test for Modeling (Forward/Backward Pass) # =================================================================================== -@pytest.mark.parametrize("L", [4]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [1024]) -@pytest.mark.parametrize("H", [4]) -@pytest.mark.parametrize("D", [64, 128]) -@pytest.mark.parametrize("dtype", [torch.bfloat16]) -@pytest.mark.parametrize("use_l2warp", [True, False]) -def test_modeling(L, B, T, H, D, dtype, use_l2warp): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) + for test in [ + (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + (4, 4, 1024, 4, 128, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, + use_l2warp: bool, +): run_test_model_forward_backward(L, B, T, H, D, TransformerConfig, dtype, use_l2warp) # =================================================================================== # Test for Generation # =================================================================================== -@pytest.mark.parametrize("L", [2]) -@pytest.mark.parametrize("B", [4]) -@pytest.mark.parametrize("T", [2000]) -@pytest.mark.parametrize("H", [3]) -@pytest.mark.parametrize("D", [64]) -@pytest.mark.parametrize("dtype", [torch.float16]) -def test_generation(L, B, T, H, D, dtype): +@pytest.mark.parametrize( + ['L', 'B', 'T', 'H', 'D', 'dtype'], + [ + pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-{}".format(*test)) + for test in [ + (2, 4, 2000, 8, 64, torch.float16), + ] + ] +) +def test_generation( + L: int, + B: int, + T: int, + H: int, + D: int, + dtype: torch.dtype, +): config = TransformerConfig() config.num_hidden_layers = L config.num_heads = H From 8502d22e800957c3cd85ef0b902621f4ea9b02ed Mon Sep 17 00:00:00 2001 From: Yu Zhang Date: Wed, 25 Jun 2025 05:40:46 +0800 Subject: [PATCH 2/5] [CI] Enable verbose output --- .github/workflows/intel-a770.yml | 10 +++++----- .github/workflows/nvidia-4090.yml | 20 ++++++++++---------- .github/workflows/nvidia-a100.yml | 20 ++++++++++---------- .github/workflows/nvidia-h100.yml | 20 ++++++++++---------- 4 files changed, 35 insertions(+), 35 deletions(-) diff --git a/.github/workflows/intel-a770.yml b/.github/workflows/intel-a770.yml index c391d62e64..080b7f4e68 100644 --- a/.github/workflows/intel-a770.yml +++ b/.github/workflows/intel-a770.yml @@ -81,21 +81,21 @@ jobs: if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=1 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' continue-on-error: true run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on varlen test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' continue-on-error: true run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" - name: Test full compiling on all test files @@ -110,12 +110,12 @@ jobs: continue-on-error: true run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run full pytest on varlen test files if: false && github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' continue-on-error: true run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" diff --git a/.github/workflows/nvidia-4090.yml b/.github/workflows/nvidia-4090.yml index 82c93247c6..85553b565a 100644 --- a/.github/workflows/nvidia-4090.yml +++ b/.github/workflows/nvidia-4090.yml @@ -86,19 +86,19 @@ jobs: if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=1 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on varlen test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" - name: Test full compiling on all test files @@ -111,13 +111,13 @@ jobs: if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run full pytest on varlen test files if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" test-models: @@ -169,19 +169,19 @@ jobs: if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=1 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on varlen test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" - name: Test full compiling on all test files @@ -194,11 +194,11 @@ jobs: if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run full pytest on varlen test files if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" diff --git a/.github/workflows/nvidia-a100.yml b/.github/workflows/nvidia-a100.yml index f43c7dab33..aafb6ebeef 100644 --- a/.github/workflows/nvidia-a100.yml +++ b/.github/workflows/nvidia-a100.yml @@ -87,19 +87,19 @@ jobs: if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=1 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on varlen test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" - name: Test full compiling on all test files @@ -112,13 +112,13 @@ jobs: if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run full pytest on varlen test files if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" test-models: @@ -170,19 +170,19 @@ jobs: if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=1 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on varlen test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" - name: Test full compiling on all test files @@ -195,11 +195,11 @@ jobs: if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run full pytest on varlen test files if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" diff --git a/.github/workflows/nvidia-h100.yml b/.github/workflows/nvidia-h100.yml index 84c64d892a..baf65dfa55 100644 --- a/.github/workflows/nvidia-h100.yml +++ b/.github/workflows/nvidia-h100.yml @@ -87,19 +87,19 @@ jobs: if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=1 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on varlen test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" - name: Test full compiling on all test files @@ -112,13 +112,13 @@ jobs: if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run full pytest on varlen test files if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" test-models: @@ -170,19 +170,19 @@ jobs: if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=1 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run pytest on varlen test files if: steps.find-dependent-tests.outputs.test_files && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" - name: Test full compiling on all test files @@ -195,11 +195,11 @@ jobs: if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=1 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} - name: Run full pytest on varlen test files if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/') && steps.check_skip.outputs.skip_tests == 'false' run: | FLA_COMPILER_MODE=0 TRITON_PRINT_AUTOTUNING=0 SKIP_TEST_CHUNK_VARLEN=0 \ - pytest ${{ steps.find-dependent-tests.outputs.test_files }} || \ + pytest -s -v ${{ steps.find-dependent-tests.outputs.test_files }} || \ echo "Varlen tests failed (non-critical)" From 8238ded34205b170f76429ee3f784d3949d76b48 Mon Sep 17 00:00:00 2001 From: Yu Zhang Date: Wed, 25 Jun 2025 05:45:56 +0800 Subject: [PATCH 3/5] Include `dtype` in `run_test_model_forward_backward` calls --- tests/models/test_modeling_abc.py | 4 ++-- tests/models/test_modeling_base.py | 6 +++--- tests/models/test_modeling_bitnet.py | 4 ++-- tests/models/test_modeling_comba.py | 4 ++-- tests/models/test_modeling_deltanet.py | 4 ++-- tests/models/test_modeling_forgetting_transformer.py | 4 ++-- tests/models/test_modeling_gated_deltanet.py | 4 ++-- tests/models/test_modeling_gated_deltaproduct.py | 4 ++-- tests/models/test_modeling_gla.py | 4 ++-- tests/models/test_modeling_gsa.py | 4 ++-- tests/models/test_modeling_hgrn.py | 4 ++-- tests/models/test_modeling_hgrn2.py | 4 ++-- tests/models/test_modeling_lightnet.py | 4 ++-- tests/models/test_modeling_linear_attn.py | 4 ++-- tests/models/test_modeling_mamba.py | 4 ++-- tests/models/test_modeling_mamba2.py | 4 ++-- tests/models/test_modeling_mesanet.py | 4 ++-- tests/models/test_modeling_nsa.py | 4 ++-- tests/models/test_modeling_path_attn.py | 4 ++-- tests/models/test_modeling_retnet.py | 4 ++-- tests/models/test_modeling_rodimus.py | 4 ++-- tests/models/test_modeling_rwkv6.py | 4 ++-- tests/models/test_modeling_rwkv7.py | 4 ++-- tests/models/test_modeling_samba.py | 4 ++-- tests/models/test_modeling_transformer.py | 4 ++-- 25 files changed, 51 insertions(+), 51 deletions(-) diff --git a/tests/models/test_modeling_abc.py b/tests/models/test_modeling_abc.py index 8c84b697f0..551259a558 100644 --- a/tests/models/test_modeling_abc.py +++ b/tests/models/test_modeling_abc.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, ABCConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, ABCConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_base.py b/tests/models/test_modeling_base.py index 5b5e02433b..049b7b2331 100644 --- a/tests/models/test_modeling_base.py +++ b/tests/models/test_modeling_base.py @@ -31,8 +31,8 @@ def run_test_model_forward_backward( H: int, D: int, config_class: type, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): """ A foundational test for the forward and backward passes of a model. @@ -44,7 +44,7 @@ def run_test_model_forward_backward( if config_class.__name__ in NOT_READY_FOR_TESTING: pytest.skip(f"{config_class.__name__} is not yet ready for testing.") - model, config = create_model_and_config(config_class, L, H, D, dtype, use_l2warp=use_l2warp) + model, config = create_model_and_config(config_class, L, H, D, use_l2warp=use_l2warp, dtype=dtype) input_ids = torch.randint(low=0, high=config.vocab_size, size=(B, T), device=device) output_fixed = model(input_ids, output_hidden_states=True).hidden_states[-1] assert output_fixed.shape == (B, T, config.hidden_size) @@ -87,7 +87,7 @@ def run_test_generation( pytest.skip(f"{config_class.__name__} is not yet ready for testing.") if model is None: - model, config = create_model_and_config(config_class, L, H, D, dtype, use_l2warp=use_l2warp) + model, config = create_model_and_config(config_class, L, H, D, use_l2warp=use_l2warp, dtype=dtype) model.eval() model = model.to(dtype).to(device) diff --git a/tests/models/test_modeling_bitnet.py b/tests/models/test_modeling_bitnet.py index c2b03e9d2a..c354ae4fbd 100644 --- a/tests/models/test_modeling_bitnet.py +++ b/tests/models/test_modeling_bitnet.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, BitNetConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, BitNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_comba.py b/tests/models/test_modeling_comba.py index 2bf6137ce7..366fe6f0d9 100644 --- a/tests/models/test_modeling_comba.py +++ b/tests/models/test_modeling_comba.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, CombaConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, CombaConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_deltanet.py b/tests/models/test_modeling_deltanet.py index 6ff3b1915b..c8d8a76aea 100644 --- a/tests/models/test_modeling_deltanet.py +++ b/tests/models/test_modeling_deltanet.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, DeltaNetConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, DeltaNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_forgetting_transformer.py b/tests/models/test_modeling_forgetting_transformer.py index 1f0f916381..6ff0b4744b 100644 --- a/tests/models/test_modeling_forgetting_transformer.py +++ b/tests/models/test_modeling_forgetting_transformer.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, ForgettingTransformerConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, ForgettingTransformerConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_gated_deltanet.py b/tests/models/test_modeling_gated_deltanet.py index 6b294e4855..b91962cfa5 100644 --- a/tests/models/test_modeling_gated_deltanet.py +++ b/tests/models/test_modeling_gated_deltanet.py @@ -27,10 +27,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, GatedDeltaNetConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, GatedDeltaNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_gated_deltaproduct.py b/tests/models/test_modeling_gated_deltaproduct.py index 55b61d7f33..c369a98ee0 100644 --- a/tests/models/test_modeling_gated_deltaproduct.py +++ b/tests/models/test_modeling_gated_deltaproduct.py @@ -31,10 +31,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, GatedDeltaProductConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, GatedDeltaProductConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_gla.py b/tests/models/test_modeling_gla.py index f3f00c2d70..5fefcf3508 100644 --- a/tests/models/test_modeling_gla.py +++ b/tests/models/test_modeling_gla.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, GLAConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, GLAConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_gsa.py b/tests/models/test_modeling_gsa.py index a57c261c84..cc9bdb19ba 100644 --- a/tests/models/test_modeling_gsa.py +++ b/tests/models/test_modeling_gsa.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, GSAConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, GSAConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_hgrn.py b/tests/models/test_modeling_hgrn.py index 3dcf88936b..7a598586c4 100644 --- a/tests/models/test_modeling_hgrn.py +++ b/tests/models/test_modeling_hgrn.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, HGRNConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, HGRNConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_hgrn2.py b/tests/models/test_modeling_hgrn2.py index 1bce95a5c9..c9595e4c12 100644 --- a/tests/models/test_modeling_hgrn2.py +++ b/tests/models/test_modeling_hgrn2.py @@ -28,10 +28,10 @@ def test_hgrn2_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, HGRN2Config, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, HGRN2Config, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_lightnet.py b/tests/models/test_modeling_lightnet.py index a43870fd8a..9743f90557 100644 --- a/tests/models/test_modeling_lightnet.py +++ b/tests/models/test_modeling_lightnet.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, LightNetConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, LightNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_linear_attn.py b/tests/models/test_modeling_linear_attn.py index 0aa90ec0c9..a8a9ea4fc8 100644 --- a/tests/models/test_modeling_linear_attn.py +++ b/tests/models/test_modeling_linear_attn.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, LinearAttentionConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, LinearAttentionConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_mamba.py b/tests/models/test_modeling_mamba.py index b6d57964be..e80e8e715f 100644 --- a/tests/models/test_modeling_mamba.py +++ b/tests/models/test_modeling_mamba.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, MambaConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, MambaConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_mamba2.py b/tests/models/test_modeling_mamba2.py index 721ee8b690..b1e774e362 100644 --- a/tests/models/test_modeling_mamba2.py +++ b/tests/models/test_modeling_mamba2.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, Mamba2Config, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, Mamba2Config, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_mesanet.py b/tests/models/test_modeling_mesanet.py index a1c66920bf..a2b1d15163 100644 --- a/tests/models/test_modeling_mesanet.py +++ b/tests/models/test_modeling_mesanet.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, MesaNetConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, MesaNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_nsa.py b/tests/models/test_modeling_nsa.py index 5dd3d918d7..576878c696 100644 --- a/tests/models/test_modeling_nsa.py +++ b/tests/models/test_modeling_nsa.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, NSAConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, NSAConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_path_attn.py b/tests/models/test_modeling_path_attn.py index 3d01914c1d..7421630822 100644 --- a/tests/models/test_modeling_path_attn.py +++ b/tests/models/test_modeling_path_attn.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, PaTHAttentionConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, PaTHAttentionConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_retnet.py b/tests/models/test_modeling_retnet.py index 4bdd9b7155..6ca144876c 100644 --- a/tests/models/test_modeling_retnet.py +++ b/tests/models/test_modeling_retnet.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, RetNetConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, RetNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_rodimus.py b/tests/models/test_modeling_rodimus.py index baea2c02fd..d6990f2574 100644 --- a/tests/models/test_modeling_rodimus.py +++ b/tests/models/test_modeling_rodimus.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, RodimusConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, RodimusConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_rwkv6.py b/tests/models/test_modeling_rwkv6.py index d90aff85d9..1966377368 100644 --- a/tests/models/test_modeling_rwkv6.py +++ b/tests/models/test_modeling_rwkv6.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, RWKV6Config, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, RWKV6Config, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_rwkv7.py b/tests/models/test_modeling_rwkv7.py index 170d9784e9..44b95006a0 100644 --- a/tests/models/test_modeling_rwkv7.py +++ b/tests/models/test_modeling_rwkv7.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, RWKV7Config, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, RWKV7Config, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # Test for Generation diff --git a/tests/models/test_modeling_samba.py b/tests/models/test_modeling_samba.py index 4fb851f3fa..8ac3f27985 100644 --- a/tests/models/test_modeling_samba.py +++ b/tests/models/test_modeling_samba.py @@ -28,10 +28,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, SambaConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, SambaConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== diff --git a/tests/models/test_modeling_transformer.py b/tests/models/test_modeling_transformer.py index 2919bf5252..2772c4873d 100644 --- a/tests/models/test_modeling_transformer.py +++ b/tests/models/test_modeling_transformer.py @@ -31,10 +31,10 @@ def test_modeling( T: int, H: int, D: int, - dtype: torch.dtype, use_l2warp: bool, + dtype: torch.dtype, ): - run_test_model_forward_backward(L, B, T, H, D, TransformerConfig, dtype, use_l2warp) + run_test_model_forward_backward(L, B, T, H, D, TransformerConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== From 540fe57a4c4f1c4c80cbed8f5adf6fca698e9f49 Mon Sep 17 00:00:00 2001 From: Yu Zhang Date: Wed, 25 Jun 2025 06:00:18 +0800 Subject: [PATCH 4/5] Refactor variable length constant name and update test parameter order --- tests/models/test_modeling_abc.py | 2 +- tests/models/test_modeling_base.py | 4 ++-- tests/models/testing_utils.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/models/test_modeling_abc.py b/tests/models/test_modeling_abc.py index 551259a558..3656e39a57 100644 --- a/tests/models/test_modeling_abc.py +++ b/tests/models/test_modeling_abc.py @@ -12,7 +12,7 @@ # Test for Modeling (Forward/Backward Pass) # =================================================================================== @pytest.mark.parametrize( - ['L', 'B', 'T', 'H', 'D', 'dtype', 'use_l2warp'], + ['L', 'B', 'T', 'H', 'D', 'use_l2warp', 'dtype'], [ pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) for test in [ diff --git a/tests/models/test_modeling_base.py b/tests/models/test_modeling_base.py index 049b7b2331..12cbb31d62 100644 --- a/tests/models/test_modeling_base.py +++ b/tests/models/test_modeling_base.py @@ -11,7 +11,7 @@ from .testing_utils import ( GENERATION_UNSUPPORTED, HOPPER_EXCLUSIVE, - MODELING_UNSUPPORTED_VAR_LEN, + MODELING_UNSUPPORTED_VARLEN, NOT_READY_FOR_TESTING, create_model_and_config ) @@ -49,7 +49,7 @@ def run_test_model_forward_backward( output_fixed = model(input_ids, output_hidden_states=True).hidden_states[-1] assert output_fixed.shape == (B, T, config.hidden_size) - if config_class.__name__ in MODELING_UNSUPPORTED_VAR_LEN: + if config_class.__name__ in MODELING_UNSUPPORTED_VARLEN: pytest.skip(f"Variable length not supported for {config_class.__name__}.") cu_seqlens = torch.arange(0, B * T + 1, T, dtype=torch.int32, device=device) diff --git a/tests/models/testing_utils.py b/tests/models/testing_utils.py index 6bca1f5114..e65f1a70c3 100644 --- a/tests/models/testing_utils.py +++ b/tests/models/testing_utils.py @@ -9,7 +9,7 @@ from fla.utils import device # Models that do not yet support variable sequence lengths (for modeling tests) -MODELING_UNSUPPORTED_VAR_LEN = [ +MODELING_UNSUPPORTED_VARLEN = [ "ABCConfig", "ForgettingTransformerConfig", "LinearAttentionConfig", "LightNetConfig", "Mamba2Config", "MambaConfig", "MesaNetConfig", "SambaConfig", "RodimusConfig", From d5c7aa205e3e2cc0ca92cb5fdaf18dc97dbd41c1 Mon Sep 17 00:00:00 2001 From: Yu Zhang Date: Wed, 25 Jun 2025 06:39:07 +0800 Subject: [PATCH 5/5] only headdim 32/64 --- tests/models/test_modeling_path_attn.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/models/test_modeling_path_attn.py b/tests/models/test_modeling_path_attn.py index 7421630822..d51d877032 100644 --- a/tests/models/test_modeling_path_attn.py +++ b/tests/models/test_modeling_path_attn.py @@ -16,9 +16,9 @@ [ pytest.param(*test, id="L{}-B{}-T{}-H{}-D{}-use_l2warp{}-{}".format(*test)) for test in [ - (4, 4, 1024, 4, 64, True, torch.bfloat16), + (4, 4, 1024, 4, 32, True, torch.bfloat16), + (4, 4, 1024, 4, 32, False, torch.bfloat16), (4, 4, 1024, 4, 64, False, torch.bfloat16), - (4, 4, 1024, 4, 128, False, torch.bfloat16), ] ] )