Yangzhang/codonfm 10b (#1305)

### Description

<!-- Provide a detailed description of the changes in this PR -->

add fsdp2 support for TE

#### Usage

<!--- How does a user interact with the changed code -->

```python
python -m src.runner pretrain \
    --exp_name "$exp_name" \
    --model_name encodon_80m \
    --data_path /data/ncbi/processed_unfiltered/ \
    --process_item mlm_memmap \
    --dataset_name CodonMemmapDataset \
    --lr $learning_rate \
    --num_gpus $num_gpus \
    --num_nodes $num_nodes \
    --train_batch_size $train_batch_size \
    --val_batch_size $val_batch_size \
    --num_workers $num_workers \
    --collate_fn thd \
    --attn_input_format thd \
    --use_transformer_engine \
    --bf16 \
    --split_name_prefix nopathogen \
    --use_transformer_engine \
    --checkpoints_dir results/${exp_name}/checkpoints/ \
    --out_dir results/${exp_name}/ \
   --enable_fsdp

```

### 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. By default, only
basic unit tests are run.

-
[ciflow:skip](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:skip)
- Skip all CI tests for this PR
-
[ciflow:notebooks](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:notebooks)
- Run Jupyter notebooks execution tests for bionemo2
-
[ciflow:slow](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:slow)
- Run slow single GPU integration tests marked as @pytest.mark.slow for
bionemo2
-
[ciflow:all](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:all)
- Run all tests (unit tests, slow tests, and notebooks) for bionemo2.
This label can be used to enforce running tests for all bionemo2.
-
[ciflow:all-recipes](https://github.com/NVIDIA/bionemo-framework/blob/main/docs/docs/main/contributing/contributing.md#ciflow:all-recipes)
- Run tests for all recipes (under bionemo-recipes). This label can be
used to enforce running tests for all recipes.

Unit tests marked as `@pytest.mark.multi_gpu` or
`@pytest.mark.distributed` are not run in the PR pipeline.

For more details, see [CONTRIBUTING](CONTRIBUTING.md)

> [!NOTE]
> By default, only basic unit tests are run. Add appropriate labels to
enable an additional test coverage.

#### 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.

### Pre-submit Checklist

<!--- Ensure all items are completed before submitting -->

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

---------

Signed-off-by: Yang Zhang <[email protected]>

Author: yzhang123

bionemo-recipes/recipes/codonfm_ptl_te/experiment_scripts/pretraining/encodon_filtered/mlm/encodon_1b.sh

@@ -11,11 +11,13 @@ val_batch_size=4
 effective_batch_size=$((train_batch_size * num_gpus * num_nodes))
 num_workers=12
 
+exp_name="encodon_1b_latest_${learning_rate}_${effective_batch_size}_nopathogen"
+
 # Note if you would like to use WandB please add --enable_wandb, --project_name and --entity.
 
 # - run
 python -m src.runner pretrain \
-    --exp_name encodon_1b_baseline_${learning_rate}_${effective_batch_size}_nopathogen \
+    --exp_name "$exp_name" \
     --model_name encodon_1b \
     --data_path /data/ncbi/processed_unfiltered/ \
     --process_item mlm_memmap \
@@ -31,4 +33,5 @@ python -m src.runner pretrain \
     --num_workers $num_workers \
     --bf16 \
     --split_name_prefix nopathogen \
-    --out_dir /workspace/codonfm/results \
+    --checkpoints_dir results/${exp_name}/checkpoints/ \
+    --out_dir results/${exp_name}/ \

bionemo-recipes/recipes/codonfm_ptl_te/experiment_scripts/pretraining/encodon_filtered/mlm/encodon_600m.sh

@@ -32,4 +32,5 @@ python -m src.runner pretrain \
     --use_transformer_engine \
     --bf16 \
     --split_name_prefix nopathogen \
-    --checkpoints_dir results/checkpoints/${exp_name} \
+    --checkpoints_dir results/${exp_name}/checkpoints/ \
+    --out_dir results/${exp_name}/ \

bionemo-recipes/recipes/codonfm_ptl_te/experiment_scripts/pretraining/encodon_filtered/mlm/encodon_80m.sh

@@ -31,4 +31,5 @@ python -m src.runner pretrain \
     --bf16 \
     --split_name_prefix nopathogen \
     --use_transformer_engine \
-    --checkpoints_dir results/checkpoints/${exp_name} \
+    --checkpoints_dir results/${exp_name}/checkpoints/ \
+    --out_dir results/${exp_name}/

bionemo-recipes/recipes/codonfm_ptl_te/requirements.txt

@@ -148,7 +148,6 @@ natsort==8.4.0
 nbclient==0.10.2
 nbconvert==7.16.6
 nbformat==5.10.4
-nemo_run @ git+https://github.com/NVIDIA/NeMo-Run.git@3ec63b951a3cf3733358f3ed2a55e87bf466d263
 nest-asyncio==1.6.0
 networkx==3.5
 ninja==1.11.1.4

bionemo-recipes/recipes/codonfm_ptl_te/src/config.py

@@ -33,14 +33,15 @@
 from src.models.encodon_pl import EncodonPL
 from src.models.encodon_te_pl import EncodonTEPL
 from src.tokenizer import Tokenizer
+from src.utils.fsdp_config import get_fsdp_strategy
 from src.utils.grad_norm_callback import GradientNormLogger
 from src.utils.pred_writer import PredWriter
 from src.utils.scheduler import linear_scheduler_with_warmup_lr_lambda
 from src.utils.timer import StepTimingCallback
 
 
 # Datasets
-def get_dataset_config(args: Any, process_item_cfg: fdl.Partial) -> fdl.Config:
+def get_dataset_config(args: Any, process_item_cfg: fdl.Partial) -> fdl.Config:  # noqa: C901
     """Builds the dataset configuration."""
     class_name = args.dataset_name
     if class_name == "CodonMemmapDataset":
@@ -49,6 +50,8 @@ def get_dataset_config(args: Any, process_item_cfg: fdl.Partial) -> fdl.Config:
         module_path = "src.data.mutation_dataset"
     elif class_name == "CodonBertDataset":
         module_path = "src.data.codon_bert_dataset"
+    elif class_name == "SimpleCodonDataset":
+        module_path = "src.data.simple_codon_dataset"
     else:
         raise ValueError(f"Unknown dataset name: {class_name}")
 
@@ -94,6 +97,12 @@ def get_dataset_config(args: Any, process_item_cfg: fdl.Partial) -> fdl.Config:
             tokenizer=tokenizer_cfg,
             process_item=process_item_cfg,
         )
+    elif class_name == "SimpleCodonDataset":
+        # SimpleCodonDataset doesn't need data_path, tokenizer, or most other args
+        dataset_cfg = fdl.Partial(
+            dataset_class,
+            process_item=process_item_cfg,
+        )
     else:
         print(f"Warning: Using generic config for dataset '{args.dataset_name}'.")
         dataset_cfg = fdl.Partial(dataset_class, **common_args)
@@ -114,6 +123,7 @@ def get_callbacks_config(args: Any) -> Dict[str, fdl.Config]:
             mode="min",
             save_top_k=1,
             auto_insert_metric_name=False,
+            enable_version_counter=False,
         ),
         "early_stopping": fdl.Config(
             EarlyStopping,
@@ -217,6 +227,12 @@ def get_logger_config(args: Any) -> fdl.Config:
 
 # Model
 MODEL_ARCHITECTURES: Dict[str, Dict[str, Any]] = {
+    "encodon_200k": {
+        "hidden_size": 128,
+        "intermediate_size": 512,
+        "num_attention_heads": 4,
+        "num_hidden_layers": 2,
+    },
     "encodon_80m": {
         "hidden_size": 1024,
         "intermediate_size": 4096,
@@ -235,6 +251,12 @@ def get_logger_config(args: Any) -> fdl.Config:
         "num_attention_heads": 16,
         "num_hidden_layers": 18,
     },
+    "encodon_10b": {
+        "hidden_size": 5120,
+        "intermediate_size": 20480,
+        "num_attention_heads": 40,
+        "num_hidden_layers": 34,
+    },
 }
 
 
@@ -304,12 +326,24 @@ def get_model_config(args: Any) -> fdl.Config:
 # Trainer
 def get_trainer_config(args: Any) -> Dict[str, Any]:
     """Builds the trainer configuration arguments."""
+    # Configure strategy based on args
+    if args.enable_fsdp:
+        # Use proper FSDP/FSDP2 strategy with auto-wrap policy
+        # This ensures FSDP uses LESS memory than DDP
+        strategy = get_fsdp_strategy(
+            cpu_offload=getattr(args, "fsdp_cpu_offload", False), activation_checkpointing=False, use_fsdp2=True
+        )
+    elif args.mode == "finetune":
+        strategy = "ddp_find_unused_parameters_true"
+    else:
+        strategy = "ddp"
+
     trainer_kwargs = dict(  # noqa: C408
         num_nodes=args.num_nodes,
         devices=args.num_gpus,
         max_steps=args.max_steps,
         default_root_dir=args.out_dir,
-        strategy="ddp" if args.mode != "finetune" else "ddp_find_unused_parameters_true",
+        strategy=strategy,
         precision="bf16-mixed" if getattr(args, "bf16", False) else "32-true",
         limit_val_batches=args.limit_val_batches,
         log_every_n_steps=args.log_every_n_steps,

bionemo-recipes/recipes/codonfm_ptl_te/src/data/simple_codon_dataset.py

@@ -0,0 +1,182 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-Apache2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""Simple synthetic codon dataset for testing and demo purposes.
+
+This dataset generates random sequences on-the-fly without requiring any data files.
+Useful for quick testing, debugging, and development without setting up real data.
+"""
+
+from typing import Callable, Optional
+
+import numpy as np
+from torch.utils.data import Dataset
+
+from src.data.metadata import MetadataFields
+
+
+class SimpleCodonDataset(Dataset):
+    """Simple synthetic dataset that generates random codon sequences.
+
+    This dataset is useful for:
+    - Quick testing without setting up data files
+    - Debugging model training loops
+    - Development and prototyping
+    - FSDP/distributed training tests
+
+    Args:
+        num_samples (int): Number of samples in the dataset. Defaults to 1000.
+        seq_length (int): Length of each sequence. Defaults to 2048.
+        vocab_size (int): Size of the vocabulary. Defaults to 69 (codon vocabulary size).
+        split_name (str): Split name ('train', 'val', 'test', or 'all'). Defaults to 'all'.
+        train_ratio (float): Ratio of training samples. Defaults to 0.8.
+        val_ratio (float): Ratio of validation samples. Defaults to 0.1.
+        process_item (Callable, optional): Function to process items. Not used in this dataset.
+        seed (int, optional): Random seed for reproducibility.
+    """
+
+    def __init__(
+        self,
+        num_samples: int = 10000,
+        seq_length: int = 2048,
+        vocab_size: int = 69,
+        split_name: str = "all",
+        train_ratio: float = 0.8,
+        val_ratio: float = 0.1,
+        process_item: Optional[Callable] = None,
+        seed: Optional[int] = None,
+        **kwargs,
+    ):
+        """Initialize the SimpleCodonDataset."""
+        self.num_samples = num_samples
+        self.seq_length = seq_length
+        self.vocab_size = vocab_size
+        self.split_name = split_name
+        self.train_ratio = train_ratio
+        self.val_ratio = val_ratio
+        self.test_ratio = 1.0 - train_ratio - val_ratio
+        self.process_item = process_item
+        self.seed = seed
+
+        # Calculate split boundaries
+        train_end = int(num_samples * train_ratio)
+        val_end = train_end + int(num_samples * val_ratio)
+
+        # Set the actual samples for this split
+        if split_name == "train":
+            self.start_idx = 0
+            self.end_idx = train_end
+        elif split_name == "val":
+            self.start_idx = train_end
+            self.end_idx = val_end
+        elif split_name == "test":
+            self.start_idx = val_end
+            self.end_idx = num_samples
+        else:  # 'all'
+            self.start_idx = 0
+            self.end_idx = num_samples
+
+        self.actual_num_samples = self.end_idx - self.start_idx
+
+    def __len__(self):
+        """Return the number of samples in this split."""
+        return self.actual_num_samples
+
+    def __getitem__(self, idx):
+        """Generate a random codon sequence sample.
+
+        Args:
+            idx: Index of the sample to retrieve.
+
+        Returns:
+            Dictionary containing:
+                - INPUT_IDS: Random token IDs (numpy array)
+                - LABELS: Random labels for MLM (numpy array)
+                - ATTENTION_MASK: All ones (no padding) (numpy array)
+                - INPUT_MASK: All ones (no masking) (numpy array)
+        """
+        # Use deterministic random generation based on seed and index
+        if self.seed is not None:
+            rng = np.random.default_rng(self.seed + self.start_idx + idx)
+        else:
+            rng = np.random.default_rng()
+
+        return {
+            MetadataFields.INPUT_IDS: rng.integers(0, self.vocab_size, size=self.seq_length, dtype=np.int64),
+            MetadataFields.LABELS: rng.integers(0, self.vocab_size, size=self.seq_length, dtype=np.int64),
+            MetadataFields.ATTENTION_MASK: np.ones(self.seq_length, dtype=bool),
+            MetadataFields.INPUT_MASK: np.ones(self.seq_length, dtype=bool),
+        }
+
+    def get_train(self, process_item: Optional[Callable] = None) -> "SimpleCodonDataset":
+        """Return the training split of the dataset.
+
+        Args:
+            process_item: Optional processing function (not used in this dataset).
+
+        Returns:
+            SimpleCodonDataset instance for the training split.
+        """
+        return SimpleCodonDataset(
+            num_samples=self.num_samples,
+            seq_length=self.seq_length,
+            vocab_size=self.vocab_size,
+            split_name="train",
+            train_ratio=self.train_ratio,
+            val_ratio=self.val_ratio,
+            process_item=process_item or self.process_item,
+            seed=self.seed,
+        )
+
+    def get_validation(self, process_item: Optional[Callable] = None) -> "SimpleCodonDataset":
+        """Return the validation split of the dataset.
+
+        Args:
+            process_item: Optional processing function (not used in this dataset).
+
+        Returns:
+            SimpleCodonDataset instance for the validation split.
+        """
+        return SimpleCodonDataset(
+            num_samples=self.num_samples,
+            seq_length=self.seq_length,
+            vocab_size=self.vocab_size,
+            split_name="val",
+            train_ratio=self.train_ratio,
+            val_ratio=self.val_ratio,
+            process_item=process_item or self.process_item,
+            seed=self.seed,
+        )
+
+    def get_test(self, process_item: Optional[Callable] = None) -> "SimpleCodonDataset":
+        """Return the test split of the dataset.
+
+        Args:
+            process_item: Optional processing function (not used in this dataset).
+
+        Returns:
+            SimpleCodonDataset instance for the test split.
+        """
+        return SimpleCodonDataset(
+            num_samples=self.num_samples,
+            seq_length=self.seq_length,
+            vocab_size=self.vocab_size,
+            split_name="test",
+            train_ratio=self.train_ratio,
+            val_ratio=self.val_ratio,
+            process_item=process_item or self.process_item,
+            seed=self.seed,
+        )