diff --git a/nemo_automodel/_transformers/mfu.py b/nemo_automodel/_transformers/mfu.py index 3eb42f967f..3b4d8dc6ff 100644 --- a/nemo_automodel/_transformers/mfu.py +++ b/nemo_automodel/_transformers/mfu.py @@ -57,6 +57,8 @@ "910B": 354e12, "Ascend910": 354e12, "RTX 3070 Ti": 21.75e12, + "RTX PRO 6000 Blackwell Max-Q Workstation Edition": 438.9e12, + "RTX PRO 6000 Blackwell Workstation Edition": 503.8e12, } _UNIT_TO_SCALE = { diff --git a/tests/unit_tests/utils/test_flops_utils.py b/tests/unit_tests/utils/test_flops_utils.py index e7c7d52ccf..0922fd7a03 100644 --- a/tests/unit_tests/utils/test_flops_utils.py +++ b/tests/unit_tests/utils/test_flops_utils.py @@ -318,6 +318,17 @@ def test_automfu_h100_reference_uses_dense_bf16_peak(): assert h100_tflops == 989.0 +@pytest.mark.parametrize( + "device_name, expected_tflops", + [ + ("NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition", 438.9), + ("NVIDIA RTX PRO 6000 Blackwell Workstation Edition", 503.8), + ], +) +def test_get_device_flops_distinguishes_rtx_pro_6000_variants(device_name, expected_tflops): + assert get_device_flops(unit="T", device_name=device_name) == expected_tflops + + def test_get_device_flops_detects_current_cuda_device(monkeypatch): monkeypatch.setattr(mfu_module.torch.cuda, "is_available", lambda: True) monkeypatch.setattr(mfu_module.torch.cuda, "current_device", lambda: 3)