Skip to content

Commit 7dfd17f

Browse files
authored
add fp8 decorator to esm2 tests (#1183)
Forgot these in my fp8+thd PR, adding to fix tests on devices without FP8 Signed-off-by: Peter St. John <pstjohn@nvidia.com>
1 parent c374faf commit 7dfd17f

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

bionemo-recipes/recipes/esm2_native_te/tests/test_train.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -281,6 +281,7 @@ def test_sanity_mfsdp_thd(distributed_cleanup, tmp_path, monkeypatch, recipe_pat
281281
main_mfsdp(sanity_config)
282282

283283

284+
@requires_fp8
284285
def test_sanity_ddp_thd_fp8(distributed_cleanup, tmp_path, monkeypatch, recipe_path):
285286
if torch.cuda.get_device_capability() == (12, 0):
286287
# TODO(BIONEMO-2840): On sm120, we need to set NVTE_FUSED_ATTN to 0 since TE will choose fused attn by default,
@@ -303,6 +304,7 @@ def test_sanity_ddp_thd_fp8(distributed_cleanup, tmp_path, monkeypatch, recipe_p
303304
main_ddp(sanity_config)
304305

305306

307+
@requires_fp8
306308
def test_sanity_mfsdp_thd_fp8(distributed_cleanup, tmp_path, monkeypatch, recipe_path):
307309
if torch.cuda.get_device_capability() == (12, 0):
308310
# TODO(BIONEMO-2840): On sm120, we need to set NVTE_FUSED_ATTN to 0 since TE will choose fused attn by default,

0 commit comments

Comments
 (0)