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)" diff --git a/tests/models/test_modeling_abc.py b/tests/models/test_modeling_abc.py index 9bfe28bdd6..3656e39a57 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): - run_test_model_forward_backward(L, B, T, H, D, ABCConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, ABCConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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_base.py b/tests/models/test_modeling_base.py index 5b5e02433b..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 ) @@ -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,12 +44,12 @@ 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) - 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) @@ -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 60d9be71a3..c354ae4fbd 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): - run_test_model_forward_backward(L, B, T, H, D, BitNetConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, BitNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..366fe6f0d9 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): - run_test_model_forward_backward(L, B, T, H, D, CombaConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, CombaConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..c8d8a76aea 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): - run_test_model_forward_backward(L, B, T, H, D, DeltaNetConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, DeltaNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..6ff0b4744b 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): - run_test_model_forward_backward(L, B, T, H, D, ForgettingTransformerConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, ForgettingTransformerConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..b91962cfa5 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): - run_test_model_forward_backward(L, B, T, H, D, GatedDeltaNetConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, GatedDeltaNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..c369a98ee0 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): - run_test_model_forward_backward(L, B, T, H, D, GatedDeltaProductConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, GatedDeltaProductConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..5fefcf3508 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): - run_test_model_forward_backward(L, B, T, H, D, GLAConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, GLAConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..cc9bdb19ba 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): - run_test_model_forward_backward(L, B, T, H, D, GSAConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, GSAConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..7a598586c4 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): - run_test_model_forward_backward(L, B, T, H, D, HGRNConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, HGRNConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..c9595e4c12 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): - run_test_model_forward_backward(L, B, T, H, D, HGRN2Config, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, HGRN2Config, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..9743f90557 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): - run_test_model_forward_backward(L, B, T, H, D, LightNetConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, LightNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..a8a9ea4fc8 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): - run_test_model_forward_backward(L, B, T, H, D, LinearAttentionConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, LinearAttentionConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..e80e8e715f 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): - run_test_model_forward_backward(L, B, T, H, D, MambaConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, MambaConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..b1e774e362 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): - run_test_model_forward_backward(L, B, T, H, D, Mamba2Config, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, Mamba2Config, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..a2b1d15163 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): - run_test_model_forward_backward(L, B, T, H, D, MesaNetConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, MesaNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..576878c696 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): - run_test_model_forward_backward(L, B, T, H, D, NSAConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, NSAConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..d51d877032 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): - run_test_model_forward_backward(L, B, T, H, D, PaTHAttentionConfig, 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, 32, True, torch.bfloat16), + (4, 4, 1024, 4, 32, False, torch.bfloat16), + (4, 4, 1024, 4, 64, False, torch.bfloat16), + ] + ] +) +def test_modeling( + L: int, + B: int, + T: int, + H: int, + D: int, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, PaTHAttentionConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..6ca144876c 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): - run_test_model_forward_backward(L, B, T, H, D, RetNetConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, RetNetConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..d6990f2574 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): - run_test_model_forward_backward(L, B, T, H, D, RodimusConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, RodimusConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..1966377368 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): - run_test_model_forward_backward(L, B, T, H, D, RWKV6Config, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, RWKV6Config, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..44b95006a0 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): - run_test_model_forward_backward(L, B, T, H, D, RWKV7Config, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, RWKV7Config, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..8ac3f27985 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): - run_test_model_forward_backward(L, B, T, H, D, SambaConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, SambaConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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..2772c4873d 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): - run_test_model_forward_backward(L, B, T, H, D, TransformerConfig, 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, + use_l2warp: bool, + dtype: torch.dtype, +): + run_test_model_forward_backward(L, B, T, H, D, TransformerConfig, use_l2warp=use_l2warp, dtype=dtype) # =================================================================================== # 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 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",