updated configs for benchmarks (#833)

### Description
Updating configs to adjust them to the recent changes in the CI

### Type of changes
<!-- Mark the relevant option with an [x] -->

- [ ]  Bug fix (non-breaking change which fixes an issue)
- [ ]  New feature (non-breaking change which adds functionality)
- [ ]  Refactor
- [ ]  Documentation update
- [ ]  Other (please describe):

### CI Pipeline Configuration
Configure CI behavior by applying the relevant labels:

-
[SKIP_CI](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/user-guide/contributing/contributing.md#skip_ci)
- Skip all continuous integration tests
-
[INCLUDE_NOTEBOOKS_TESTS](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/user-guide/contributing/contributing.md#include_notebooks_tests)
- Execute notebook validation tests in pytest
-
[INCLUDE_SLOW_TESTS](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/user-guide/contributing/contributing.md#include_slow_tests)
- Execute tests labelled as slow in pytest for extensive testing

> [!NOTE]
> By default, the notebooks validation tests are skipped unless
explicitly enabled.

#### Authorizing CI Runs

We use
[copy-pr-bot](https://docs.gha-runners.nvidia.com/apps/copy-pr-bot/#automation)
to manage authorization of CI
runs on NVIDIA's compute resources.

* If a pull request is opened by a trusted user and contains only
trusted changes, the pull request's code will
automatically be copied to a pull-request/ prefixed branch in the source
repository (e.g. pull-request/123)
* If a pull request is opened by an untrusted user or contains untrusted
changes, an NVIDIA org member must leave an
`/ok to test` comment on the pull request to trigger CI. This will need
to be done for each new commit.

### Usage
<!--- How does a user interact with the changed code -->
```python
TODO: Add code snippet
```

### Pre-submit Checklist
<!--- Ensure all items are completed before submitting -->

 - [ ] I have tested these changes locally
 - [ ] I have updated the documentation accordingly
 - [ ] I have added/updated tests as needed
 - [ ] All existing tests pass successfully

---------

Signed-off-by: dorotat <dorotat@nvidia.com>

Author: dorotat-nv

ci/benchmarks/partial-conv/esm2_pretrain.yaml

@@ -1,31 +1,22 @@
 scope: partial-conv
 time_limit: 14400
+key_segments:
+  # Modify keys to be renamed (str) or excluded (False) from run identifier. By default, all args under script_args are included.
+  data_path: False
 script_args:
   # All arguments referenced in the script string must be specified here.
   # Arguments not referenced in the script string must have the 'arg' field specified.
   # See jet/core/configs.py for the specification of the configuration class
-  workspace:
-    value: /workspace/bionemo2
-    key_segment: False
-  data_path:
-    value: /data/20240809_uniref_2024_03/data
-    key_segment: False
-  model:
-    value: esm2
-  variant:
-    value: train
-  config_name:
-    value: 650M
-  precision:
-    value: [bf16-mixed]
-  nodes:
-    value: [4]
-  gpus:
-    value: 8
-  batch_size:
-    value: 16
-  max_steps:
-    value: 500000
+  workspace: /workspace/bionemo2
+  data_path: /data/20240809_uniref_2024_03/data
+  model: esm2
+  variant: train
+  config_name: 650M
+  precision: [bf16-mixed]
+  nodes: [4]
+  gpus: 8
+  batch_size: 16
+  max_steps: 500000
   stop_steps: 26500
 script: |-
   WANDB_API_KEY=$BIONEMO_WANDB_API_KEY ${variant}_${model} \

ci/benchmarks/partial-conv/evo2_pretrain.yaml

@@ -1,19 +1,20 @@
 scope: partial-conv
 time_limit: 14400
+key_segments:
+  # Modify keys to be renamed (str) or excluded (False) from run identifier. By default, all args under script_args are included.
+  data_path: False
+  clip_grad: False
+  lr: False
+  min_lr: False
+  wu_steps: False
 script_args:
   # All arguments referenced in the script string must be specified here.
   # Arguments not referenced in the script string must have the 'arg' field specified.
   # See jet/core/configs.py for the specification of the configuration class
-  workspace:
-    value: /workspace/bionemo2
-    key_segment: False
-  data_path:
-    value: /data/evo2
-    key_segment: False
-  model:
-    value: evo2
-  variant:
-    value: train
+  workspace: /workspace/bionemo2
+  data_path: /data/evo2
+  model: evo2
+  variant: train
   config_name: 1b
   precision: fp8
   gpus: 8
@@ -26,22 +27,12 @@ script_args:
   tp: 1
   seq_len: 8192
   acc_grad: 1
-  clip_grad:
-    value: 250
-    key_segment: False
+  clip_grad: 250
   seed: 3735928559
-  lr:
-    value: 0.00015
-    key_segment: False
-  min_lr:
-    value: 0.000015
-    key_segment: False
-  wu_steps:
-    value: 5000
-    key_segment: False
-  wd:
-    value: 0.1
-    key_segment: False
+  lr: 0.00015
+  min_lr: 0.000015
+  wu_steps: 5000
+  wd: 0.1
 script: |-
   WANDB_API_KEY=$BIONEMO_WANDB_API_KEY ${variant}_${model} \
   -d /workspace/bionemo2/sub-packages/bionemo-evo2/examples/configs/full_pretrain_shortphase_config.yaml \

ci/benchmarks/partial-conv/geneformer_pretrain.yaml

@@ -1,38 +1,27 @@
 scope: partial-conv
 time_limit: 14400
+key_segments:
+  # Modify keys to be renamed (str) or excluded (False) from run identifier. By default, all args under script_args are included.
+  data_path: False
+  val_check_interval: False
+  lr: False
 script_args:
   # All arguments referenced in the script string must be specified here.
   # Arguments not referenced in the script string must have the 'arg' field specified.
   # See jet/core/configs.py for the specification of the configuration class
-  workspace:
-    value: /workspace/bionemo2
-    key_segment: False
-  data_path:
-    value: /data/cellxgene_scdl
-    key_segment: False
-  model:
-    value: geneformer
-  variant:
-    value: train
-  config_name:
-    value: geneformer_config
-  precision:
-    value: [bf16-mixed]
-  nodes:
-    value: [2]
-  gpus:
-    value: 8
-  batch_size:
-    value: 32
-  max_steps:
-    value: 37000
-  lr:
-    value: 0.001
-  val_check_interval:
-    value: 500
-  acc:
-    value: 1
-
+  workspace: /workspace/bionemo2
+  data_path: /data/cellxgene_scdl
+  model: geneformer
+  variant: train
+  config_name: geneformer_config
+  precision: [bf16-mixed]
+  nodes: [2]
+  gpus: 8
+  batch_size: 32
+  max_steps: 37000
+  lr: 0.001
+  val_check_interval: 500
+  acc_grad: 1
 script: |-
    WANDB_API_KEY=$BIONEMO_WANDB_API_KEY ${variant}_${model} \
     --data-dir ${data_path} \
@@ -55,6 +44,6 @@ script: |-
     --wandb-job-type=${pipeline_label} \
     --cosine-rampup-frac 0.004331629559040111 \
     --cosine-hold-frac 0.021658147795200554 \
-    --accumulate-grad-batches ${acc} \
+    --accumulate-grad-batches ${acc_grad} \
     --precision ${precision} \
     --disable-checkpointing;

ci/benchmarks/perf/esm2_pretrain.yaml

@@ -1,15 +1,14 @@
 scope: perf
 time_limit: 1800
+key_segments:
+  # Modify keys to be renamed (str) or excluded (False) from run identifier. By default, all args under script_args are included.
+  data_path: False
 script_args:
   # All arguments referenced in the script string must be specified here.
   # Arguments not referenced in the script string must have the 'arg' field specified.
   # See jet/core/configs.py for the specification of the configuration class
-  workspace:
-    value: /workspace/bionemo2
-    key_segment: False
-  data_path:
-    value: /data/20240809_uniref_2024_03/data
-    key_segment: False
+  workspace: /workspace/bionemo2
+  data_path: /data/20240809_uniref_2024_03/data
   model: esm2
   variant: train
   config_name: 650M

ci/benchmarks/perf/evo2_pretrain.yaml

@@ -1,19 +1,20 @@
 scope: perf
 time_limit: 1800
+key_segments:
+  # Modify keys to be renamed (str) or excluded (False) from run identifier. By default, all args under script_args are included.
+  data_path: False
+  clip_grad: False
+  lr: False
+  min_lr: False
+  wu_steps: False
 script_args:
   # All arguments referenced in the script string must be specified here.
   # Arguments not referenced in the script string must have the 'arg' field specified.
   # See jet/core/configs.py for the specification of the configuration class
-  workspace:
-    value: /workspace/bionemo2
-    key_segment: False
-  data_path:
-    value: /data/evo2
-    key_segment: False
-  model:
-    value: evo2
-  variant:
-    value: train
+  workspace: /workspace/bionemo2
+  data_path: /data/evo2
+  model: evo2
+  variant: train
   config_name: 1b
   precision: fp8
   gpus: 8
@@ -23,22 +24,12 @@ script_args:
   cp: 1
   seq_len: 8192
   acc_grad: 1
-  clip_grad:
-    value: 250
-    key_segment: False
+  clip_grad: 250
   seed: 3735928559
-  lr:
-    value: 0.00015
-    key_segment: False
-  min_lr:
-    value: 0.000015
-    key_segment: False
-  wu_steps:
-    value: 5000
-    key_segment: False
-  wd:
-    value: 0.1
-    key_segment: False
+  lr: 0.00015
+  min_lr: 0.000015
+  wu_steps: 5000
+  wd: 0.1
   products:
     - nodes: 2
       pp: 1

ci/benchmarks/perf/geneformer_pretrain.yaml

@@ -1,33 +1,25 @@
 scope: perf
 time_limit: 3600
+key_segments:
+  # Modify keys to be renamed (str) or excluded (False) from run identifier. By default, all args under script_args are included.
+  data_path: False
+  val_check_interval: False
+  lr: False
 script_args:
   # All arguments referenced in the script string must be specified here.
   # Arguments not referenced in the script string must have the 'arg' field specified.
   # See jet/core/configs.py for the specification of the configuration class
-  workspace:
-    value: /workspace/bionemo2
-    key_segment: False
-  data_path:
-    value: /data/cellxgene_scdl
-    key_segment: False
-  model:
-    value: geneformer
-  variant:
-    value: train
-  config_name:
-    value: geneformer_config
-  precision:
-    value: [bf16-mixed]
-  gpus:
-    value: 8
-  max_steps:
-    value: 1000
-  lr:
-    value: 0.001
-  val_check_interval:
-    value: 500
-  acc:
-    value: 1
+  workspace: /workspace/bionemo2
+  data_path: /data/cellxgene_scdl
+  model: geneformer
+  variant: train
+  config_name: geneformer_config
+  precision: [bf16-mixed]
+  gpus:  8
+  max_steps: 1000
+  lr: 0.001
+  val_check_interval: 500
+  acc_grad: 1
   products:
     - nodes: 1
       batch_size: 32
@@ -52,10 +44,11 @@ script: |-
     --lr ${lr} \
     --create-tensorboard-logger \
     --result-dir=${tensorboard_dir} \
+    --wandb-group=${model}_${variant}_${config_name}__${target} \
     --wandb-project ${wandb_project_name} \
     --wandb-job-type=${pipeline_label} \
     --cosine-rampup-frac 0.004331629559040111 \
     --cosine-hold-frac 0.021658147795200554 \
-    --accumulate-grad-batches ${acc} \
+    --accumulate-grad-batches ${acc_grad} \
     --precision ${precision} \
     --disable-checkpointing;