Merge branch 'main' into mvle/evo2-fine-tuning

Author: jwilber

3rdparty/NeMo

@@ -74,3 +74,28 @@ script: |-
     --disable-checkpointing \
     --early-stop-on-step=${stop_steps} \
     --create-tflops-callback;
+tests:
+  - logic_type: static
+    logic_spec:
+      exit_codes:
+        - 0
+      baselines:
+        consumed_samples:
+          operator: eq
+          value: 3072000.0
+        val_loss:
+          operator: range
+          max: 0.6678
+          min: 0.5630
+        reduced_train_loss:
+          operator: range
+          max: 0.1372
+          min: 0.0993
+        TFLOPS_per_GPU:
+          operator: range
+          max: 719.5984
+          min: 694.8020
+        val_acc:
+          operator: range
+          max: 0.8462
+          min: 0.7833

ci/benchmarks/partial-conv/esm2_finetuning.yaml

@@ -16,7 +16,7 @@ script_args:
   workspace: /workspace/bionemo2
   data_path: /data/evo2
   artefacts_url: https://__token__:${JET_GITLAB_TOKEN}@gitlab-master.nvidia.com/api/v4/projects/180496/packages/pypi/simple
-  file_name_wheel: cuhyena==v0.2.3+cuda12.9
+  file_name_wheel: subquadratic-ops
   model: evo2
   variant: train
   config_name: 1b
@@ -55,7 +55,7 @@ script: |-
   --enable-preemption \
   --ckpt-async-save \
   --use-megatron-comm-overlap-llama3-8k \
-  --use-b2b-causal-conv1d \
+  --use-subquadratic_ops \
   --overlap-grad-reduce \
   --clip-grad=${clip_grad} \
   --eod-pad-in-loss-mask \
@@ -78,7 +78,7 @@ script: |-
   --limit-val-batches=20 \
   --log-every-n-steps=50 \
   --val-check-interval=500 \
-  --use-b2b-causal-conv1d \
+  --use-subquadratic_ops \
   --create-tflops-callback \
   --create-tensorboard-logger \
   --result-dir=${tensorboard_dir} \

ci/benchmarks/partial-conv/evo2_pretrain.yaml

@@ -76,7 +76,7 @@ script: |-
   --limit-val-batches=20 \
   --log-every-n-steps=50 \
   --val-check-interval=200 \
-  --use-b2b-causal-conv1d \
+  --use-subquadratic_ops \
   --create-tflops-callback \
   --create-tensorboard-logger \
   --result-dir=${tensorboard_dir} \

ci/benchmarks/perf/evo2_pretrain.yaml

@@ -462,9 +462,9 @@ def parse_args(args: Optional[List[str]] = None) -> argparse.Namespace:
         help="Dropout probability for the attention layers.",
     )
     parser.add_argument(
-        "--use-b2b-causal-conv1d",
+        "--use-subquadratic_ops",
         action="store_true",
-        help="Use back-to-back causal convolution CUDA kernel for hyena short conv layers for improved performance.",
+        help="Use subquadratic_ops for improved performance.",
     )
     parser.add_argument(
         "--save-top-k",
@@ -597,8 +597,8 @@ def train(args: argparse.Namespace) -> nl.Trainer:
         config_modifiers_init["ffn_hidden_size"] = args.ffn_hidden_size
     if args.use_targeted_variance_loss:
         config_modifiers_init["use_targeted_variance_loss"] = True
-    if args.use_b2b_causal_conv1d:
-        config_modifiers_init["use_b2b_causal_conv1d"] = True
+    if args.use_subquadratic_ops:
+        config_modifiers_init["use_subquadratic_ops"] = True
     if args.hybrid_override_pattern:
         config_modifiers_init["hybrid_override_pattern"] = args.hybrid_override_pattern
     if args.num_layers: