dataset refactor, add 350M config

Signed-off-by: Peter St. John <pstjohn@nvidia.com>

Author: pstjohn

recipes/esm2_accelerate/dataset.py

@@ -16,59 +16,70 @@
 # Create the dataset -- here, we just use a simple parquet file with some raw protein sequences
 # stored in the repo itself to avoid external dependencies.
 
-from pathlib import Path
-
-from datasets import load_dataset
+from datasets import IterableDataset, load_dataset
 from transformers import AutoTokenizer
 from transformers.data.data_collator import DataCollatorForLanguageModeling
 
 
-def infinite_dataloader(dataloader, sampler):
-    """Create an infinite iterator that automatically restarts at the end of each epoch."""
-    epoch = 0
-    while True:
-        sampler.set_epoch(epoch)  # Update epoch for proper shuffling
-        for batch in dataloader:
-            yield batch
-        epoch += 1  # Increment epoch counter after completing one full pass
-
-
-def create_datasets_and_collator(tokenizer_name: str, max_length: int = 1024):
-    """Create a dataloader for the dataset.
+def create_datasets_and_collator(
+    tokenizer_name: str,
+    train_load_dataset_kwargs: dict,
+    eval_load_dataset_kwargs: dict,
+    max_seq_length: int = 1024,
+    truncate_eval_dataset: int | None = None,
+):
+    """Create datasets and a data collator to pass to the huggingface trainer.
 
     Args:
         tokenizer_name: The name of the tokenizer to pull from the HuggingFace Hub.
-        max_length: The maximum length of the protein sequences.
+        train_load_dataset_kwargs: Keyword arguments to pass to `load_dataset` for the train dataset.
+        eval_load_dataset_kwargs: Keyword arguments to pass to `load_dataset` for the eval dataset.
+        max_seq_length: The maximum length of the protein sequences.
+        truncate_eval_dataset: If not `None`, the eval dataset will be truncated to this number of examples.
+
+    This assumes that the dataset has a "sequence" column that will be tokenized.
 
     Returns:
         Tuple of (train_dataset, eval_dataset, data_collator).
     """
-    # We copy this parquet file to the container to avoid external dependencies, modify if you're
-    # using a local dataset. If you're reading this and scaling up the dataset to a larger size,
-    # look into `set_transform` and other streaming options from the `datasets` library.
-    data_path = Path(__file__).parent / "train.parquet"
-    train_dataset = load_dataset("parquet", data_files=data_path.as_posix(), split="train")
-    eval_dataset = train_dataset.select(range(10))
+    train_dataset = load_dataset(**train_load_dataset_kwargs)
+    eval_dataset = load_dataset(**eval_load_dataset_kwargs)
+    if truncate_eval_dataset is not None:
+        if isinstance(eval_dataset, IterableDataset):
+            raise ValueError(
+                "Cannot truncate an IterableDataset, don't use streaming datasets for eval if you want to truncate."
+            )
+        eval_dataset = eval_dataset.select(range(truncate_eval_dataset))
 
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
 
-    def tokenize_function(examples):
+    def tokenize_function(sequence):
         """Tokenize the protein sequences."""
         return tokenizer(
-            examples["sequence"],
+            sequence,
             truncation=True,
             padding="max_length",
-            max_length=max_length,
+            max_length=max_seq_length,
             return_tensors="pt",
         )
 
-    for dataset in [train_dataset, eval_dataset]:
-        dataset.set_transform(tokenize_function)
+    train_dataset = train_dataset.map(
+        tokenize_function,
+        batched=True,
+        input_columns=["sequence"],
+        remove_columns=train_dataset.column_names,
+    )
+    eval_dataset = eval_dataset.map(
+        tokenize_function,
+        batched=True,
+        input_columns=["sequence"],
+        remove_columns=eval_dataset.column_names,
+    )
 
     data_collator = DataCollatorForLanguageModeling(
         tokenizer=tokenizer,
         mlm_probability=0.15,
-        pad_to_multiple_of=max_length,
+        pad_to_multiple_of=max_seq_length,
     )
 
     return train_dataset, eval_dataset, data_collator

recipes/esm2_accelerate/hydra_config/L0_sanity.yaml

@@ -5,6 +5,20 @@ defaults:
 model_tag: "nvidia/esm2_t6_8M_UR50D"
 stop_after_n_steps: 250
 
+dataset:
+  tokenizer_name: ${model_tag}
+  max_seq_length: 1024
+  train_load_dataset_kwargs:
+    path: "parquet"
+    split: "train"
+    data_files: "train.parquet"
+    streaming: True
+  eval_load_dataset_kwargs:
+    path: "parquet"
+    split: "train"
+    data_files: "train.parquet"
+  truncate_eval_dataset: 10
+
 trainer:
   run_name: "esm2_t6_8M_UR50D_sanity"
   per_device_train_batch_size: 2
@@ -13,5 +27,5 @@ trainer:
   eval_steps: 1000
   logging_steps: 10
   report_to: "none"
-  dataloader_num_workers: 4
+  dataloader_num_workers: 1
   warmup_steps: 0

recipes/esm2_accelerate/hydra_config/L0_sanity_amplify.yaml

@@ -5,6 +5,20 @@ defaults:
 model_tag: "nvidia/AMPLIFY_120M"
 stop_after_n_steps: 250
 
+dataset:
+  tokenizer_name: ${model_tag}
+  max_seq_length: 1024
+  train_load_dataset_kwargs:
+    path: "parquet"
+    split: "train"
+    data_files: "train.parquet"
+    streaming: True
+  eval_load_dataset_kwargs:
+    path: "parquet"
+    split: "train"
+    data_files: "train.parquet"
+  truncate_eval_dataset: 10
+
 trainer:
   run_name: "amplify_120M_sanity"
   per_device_train_batch_size: 2

recipes/esm2_accelerate/hydra_config/L1_amplify_350M_partial_conv.yaml

@@ -0,0 +1,13 @@
+defaults:
+  - defaults_amplify
+  - _self_
+
+stop_after_n_steps: 20_000
+trainer:
+  run_name: "L1-350M-partial-conv"
+  eval_steps: 1_000
+  save_steps: 1_000
+  logging_steps: 10
+  report_to: "wandb"
+  per_device_train_batch_size: 128
+  per_device_eval_batch_size: 256

…erate/hydra_config/L1_15B_perf_test.yaml …/hydra_config/L1_esm2_15B_perf_test.yamlrecipes/esm2_accelerate/hydra_config/L1_15B_perf_test.yaml renamed to recipes/esm2_accelerate/hydra_config/L1_esm2_15B_perf_test.yaml recipes/esm2_accelerate/hydra_config/L1_15B_perf_test.yaml renamed to recipes/esm2_accelerate/hydra_config/L1_esm2_15B_perf_test.yaml

@@ -1,6 +1,23 @@
 model_tag: "nvidia/esm2_t6_8M_UR50D"
 stop_after_n_steps: 500_000
-max_seq_length: 1024
+
+dataset:
+  tokenizer_name: ${model_tag}
+  max_seq_length: 1024
+  # TODO(BIONEMO-2783): Replace this with our ESM-2 parquet dataset when it's ready.
+  train_load_dataset_kwargs:
+    path: "chandar-lab/UR100P"
+    split: "train"
+    revision: "refs/convert/parquet"
+    streaming: True
+  eval_load_dataset_kwargs:
+    path: "chandar-lab/UR100P"
+    split: "test"
+    revision: "refs/convert/parquet"
+  # Whether to truncate the eval dataset; HF Trainer will run the full eval dataset each eval step.
+  # If set to an integer, the eval dataset will be truncated to that number of examples.
+  truncate_eval_dataset: null
+
 trainer:
   output_dir: "results"
   run_name: ???

recipes/esm2_accelerate/hydra_config/defaults.yaml

@@ -48,10 +48,7 @@ def main(args: DictConfig):
     config = AutoConfig.from_pretrained(args.model_tag, trust_remote_code=True)
     model = AutoModelForMaskedLM.from_config(config, trust_remote_code=True, dtype=torch.bfloat16)
 
-    train_dataset, eval_dataset, data_collator = create_datasets_and_collator(
-        tokenizer_name=args.model_tag,
-        max_length=args.max_seq_length,
-    )
+    train_dataset, eval_dataset, data_collator = create_datasets_and_collator(**args.dataset)
 
     training_args = TrainingArguments(**args.trainer)