br: adjust sequence length cap for fixed memory requirements"

Signed-off-by: Brian Roland <broland@nvidia.com>

Author: broland-hat

sub-packages/bionemo-evo2/tests/bionemo/evo2/test_evo2.py

@@ -16,6 +16,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from calendar import c
 import logging
 import os
 import time
@@ -48,9 +49,44 @@
 logger.setLevel(logging.DEBUG)  # Capture all levels in the logger itself
 
 
-MEM_REQUIREMENT_1B_GB = 18  # add 0.6 GB to max mem reserved, and round up
-MEM_REQUIREMENT_7B_GB = 48
+def determine_memory_requirement_and_skip_if_not_met(
+    ckpt_name: str, flash_decode: bool | None = None
+) -> int:
+    """Determine the memory requirement for a given checkpoint and flash decode condition.
+    ckpt_name : str
+        the name of the checkpoint to test
+    flash_decode: bool | None
+        whether to test with flash decode
+    Returns:
+        The input sequence length cap, for the model sin the checkpoint, given certain memory requirements.
+        If the memory requirement is not met, the test is skipped.
+    """
+    
+    if "1b" in ckpt_name:
+        model_size = "1b"
+        seq_len_cap = 6000
+        memory_needed_by_test = 17  # max reserved rounded up, for stand-alone test
+    
+    elif "7b" in ckpt_name:
+        model_size = "7b"
+        seq_len_cap = 4000
+        memory_needed_by_test = 32  # max reserved rounded up, for stand-alone test
+    else:
+        raise ValueError(f"{ckpt_name=} is not supported for testing")
+    
+    skip_condition_flash = (flash_decode is None or flash_decode)
+    gb_available = torch.cuda.mem_get_info()[0] / 1024**3
+    skip_condition = gb_available < memory_needed_by_test and skip_condition_flash
 
+    if skip_condition:
+        pytest.skip(
+            ", ".join([
+                f"Inference API requires at least {memory_needed_by_test}GB of available memory for {model_size} models",
+                f"{gb_available=}"
+            ])
+        )
+    
+    return seq_len_cap
 
 
 def load_weights_sharded_inplace_nemo2_to_mcore(
@@ -369,11 +405,8 @@ def check_matchrate(*, ckpt_name, matchrate, assert_matchrate=True):
 )
 def test_forward(sequences: list[str], ckpt_name: str, expected_matchpercents: list[float]):
     assert len(sequences) > 0
-    gb_available = torch.cuda.mem_get_info()[0] / 1024**3
-    if (gb_available < MEM_REQUIREMENT_1B_GB and "1b" in ckpt_name) or (gb_available < MEM_REQUIREMENT_7B_GB and "7b" in ckpt_name):
-        pytest.skip(
-            f"Inference API requires more than 38GB of memory for 1b models, or 50GB for 7b models. {gb_available=}"
-        )
+    seq_len_cap = determine_memory_requirement_and_skip_if_not_met(ckpt_name)
+    
     is_fp8_supported, compute_capability, device_info = check_fp8_support(torch.cuda.current_device())
     skip = "evo2/1b-8k:" in ckpt_name and not is_fp8_supported
     if skip:
@@ -385,7 +418,7 @@ def test_forward(sequences: list[str], ckpt_name: str, expected_matchpercents: l
     )
     matchrates = []
     for seq in sequences:
-        seq = seq[:6000]  # TODO: artificial limit, megatron uses more memory. Vortex can process full sequences
+        seq = seq[:seq_len_cap]  # TODO: artificial limit, megatron uses more memory. Vortex can process full sequences
         with torch.no_grad():
             device = torch.cuda.current_device()
             tokens = torch.tensor([mcore_tokenizer.tokenize(seq)], device=device)
@@ -431,13 +464,12 @@ def test_forward(sequences: list[str], ckpt_name: str, expected_matchpercents: l
 )
 def test_forward_manual(sequences: list[str], ckpt_name: str, expected_matchpercents: list[float], flash_decode: bool):
     assert len(sequences) > 0
+    seq_len_cap = determine_memory_requirement_and_skip_if_not_met(ckpt_name, flash_decode)
+    
     is_fp8_supported, compute_capability, device_info = check_fp8_support(torch.cuda.current_device())
     skip = "evo2/1b-8k:" in ckpt_name and not is_fp8_supported
-    gb_available = torch.cuda.mem_get_info()[0] / 1024**3
-    if (gb_available < MEM_REQUIREMENT_1B_GB and flash_decode) or (gb_available < MEM_REQUIREMENT_7B_GB and flash_decode and "7b" in ckpt_name):
-        pytest.skip(
-            f"Inference API requires more than 38GB of memory for 1b models, or 50GB for 7b models. {gb_available=}"
-        )
+    
+    
     vortex_style_fp8 = is_fp8_supported and "bf16" not in ckpt_name
     if skip:
         # This checkpoint is sensitive to FP8, so we skip it if it is not supported on the current device.
@@ -484,7 +516,7 @@ def test_forward_manual(sequences: list[str], ckpt_name: str, expected_matchperc
             forward_kwargs = {}
         matchrates = []
         for seq in sequences:
-            seq = seq[:6000]  # TODO: artificial limit, megatron uses more memory. Vortex can process full sequences
+            seq = seq[:seq_len_cap]  # TODO: artificial limit, megatron uses more memory. Vortex can process full sequences
             with torch.no_grad():
                 device = torch.cuda.current_device()
                 # tokens = torch.tensor([tokenizer.tokenize(seq)], device=device)
@@ -547,12 +579,9 @@ def test_batch_generate(
     sequences: list[str], ckpt_name: str, model_tokenizer_provider: Callable, expected_matchpercents: list[float]
 ):
     assert len(sequences) > 0
+    determine_memory_requirement_and_skip_if_not_met(ckpt_name)
+    
     is_fp8_supported, compute_capability, device_info = check_fp8_support(torch.cuda.current_device())
-    gb_available = torch.cuda.mem_get_info()[0] / 1024**3
-    if (gb_available < MEM_REQUIREMENT_1B_GB and "1b" in ckpt_name) or (gb_available < MEM_REQUIREMENT_7B_GB and "7b" in ckpt_name):
-        pytest.skip(
-            f"Inference API requires more than 38GB of memory for 1b models, or 50GB for 7b models. {gb_available=}"
-        )
     skip = "evo2/1b-8k:" in ckpt_name and not is_fp8_supported
     if skip:
         # This checkpoint is sensitive to FP8, so we skip it if it is not supported on the current device.
@@ -619,11 +648,8 @@ def test_batch_generate_coding_sequences(
     expected_matchpercents: list[float],
 ):
     assert len(coding_sequences) > 0
-    gb_available = torch.cuda.mem_get_info()[0] / 1024**3
-    if (gb_available < MEM_REQUIREMENT_1B_GB and "1b" in ckpt_name) or (gb_available < MEM_REQUIREMENT_7B_GB and "7b" in ckpt_name):
-        pytest.skip(
-            f"Inference API requires more than 38GB of memory for 1b models, or 50GB for 7b models. {gb_available=}"
-        )
+    determine_memory_requirement_and_skip_if_not_met(ckpt_name)
+    
     is_fp8_supported, compute_capability, device_info = check_fp8_support(torch.cuda.current_device())
     skip = "evo2/1b-8k:" in ckpt_name and not is_fp8_supported
     if skip:
@@ -728,11 +754,8 @@ def test_generate_speed(
     expected_tokens_sec: float,
 ):
     is_fp8_supported, compute_capability, device_info = check_fp8_support(torch.cuda.current_device())
-    gb_available = torch.cuda.mem_get_info()[0] / 1024**3
-    if (gb_available < MEM_REQUIREMENT_1B_GB and "1b" in ckpt_name) or (gb_available < MEM_REQUIREMENT_7B_GB and "7b" in ckpt_name):
-        pytest.skip(
-            f"Inference API requires more than 38GB of memory for 1b models, or 50GB for 7b models. {gb_available=}"
-        )
+    determine_memory_requirement_and_skip_if_not_met(ckpt_name)
+    
     skip = "evo2/1b-8k:" in ckpt_name and not is_fp8_supported
     if skip:
         # This checkpoint is sensitive to FP8, so we skip it if it is not supported on the current device.