Train conformers

submitit_train.py

@@ -9,8 +9,9 @@
 
 if __name__ == "__main__":
     executor = submitit.AutoExecutor(folder="~/slurm_jobs/titan/job_%j")
-    n_gpus = 6
-    node = "h100"
+    # executor = submitit.local.local.LocalExecutor(folder="~/slurm_jobs/titan/job_%j")
+    n_gpus = 4
+    node = "a100"
     executor.update_parameters(
         name="titan",
         timeout_min=24 * 24 * 60,
@@ -26,7 +27,8 @@
         for _ in range(1):
             # train_config = './train_configs/chemlactica_125m.toml'
             # train_config = './train_configs/chemlactica_1.3b.toml'
-            train_config = "./train_configs/llama3.2_1b.toml"
+            train_config = "./train_configs/llama3.2_1b_conformers.toml"
+            # train_config = "./train_configs/llama3.2_1b_conversion.toml"
             # train_config = './train_configs/debug_model.toml'
             function = submitit.helpers.CommandFunction(
                 [

torchtitan/datasets/hf_datasets.py

@@ -25,7 +25,7 @@
 
 from torchtitan.logging import logger
 from torchtitan.tokenizers.tokenizer import Tokenizer
-from torchtitan.utils.dataset_utils import chemlactica_style_data_processing
+from torchtitan.utils.dataset_utils import chemlactica_style_data_processing, conformer_data_processing
 
 from datasets import load_dataset
 from datasets.distributed import split_dataset_by_node
@@ -38,13 +38,16 @@
     "c4": "allenai/c4",
     "chemlactica_train_mini": "test/assets/chemlactica_train_mini",
     "chemlactica_train": "/nfs/dgx/raid/chem/data/rdkit_computed_rel+form/train_rdkit_computed_rel+form",
+    "conformers_train": "/auto/home/menuab/code/3DMolGen/data/pcqm/train",
+    "conformers_valid": "/auto/home/menuab/code/3DMolGen/data/pcqm/valid",
     # valid
     "chemlactica_valid": "/nfs/dgx/raid/chem/data/rdkit_computed_rel+form",
     "chemlactica_valid_mini": "test/assets/chemlactica_valid_mini",
 }
 
 _supported_data_processing_styles = {
-    "chemlactica_style": chemlactica_style_data_processing
+    "chemlactica_style": chemlactica_style_data_processing,
+    "conformer_style": conformer_data_processing
 }
 
 

torchtitan/utils/dataset_utils.py

@@ -3,7 +3,7 @@
 
 import orjson
 import json
-from .text_format_utils import generate_formatted_string, delete_empty_tags
+from .text_format_utils import generate_formatted_string, delete_empty_tags, generate_formatted_conformer_string
 import torch
 import os
 from pathlib import Path
@@ -33,6 +33,18 @@ def chemlactica_style_data_processing(sample_json, rng, representation_type):
     return sample_json
 
 
+def conformer_data_processing(sample_json, rng, representation_type):
+    try:
+        compound = json.loads(sample_json)
+        sample_json = generate_formatted_conformer_string(
+            compound, rng, representation_type
+        )
+    except Exception as e:
+        print(e)
+        sample_json = ""
+    return sample_json
+
+
 def sft_formatting_prompts_func(example):
     output_texts = []
     for i in range(len(example["smiles"])):

torchtitan/utils/text_format_utils.py

@@ -53,6 +53,30 @@ def convert_representation(smiles, representation_type):
         # logger.info(f"{e}. Could not encode molecule {smiles} with representation {representation_type}")
         return smiles
 
+def generate_formatted_conformer_string(compound_json, rng, representation_type):
+    canonical_smiles = compound_json.get("canonical_smiles", "")
+    canonical_smiles = f"[SMILES]{canonical_smiles}[/SMILES]"
+
+    conformers = compound_json.get("conformers", "")
+    label = compound_json.get("pcqm4v2_label", "")
+
+    key_value_pairs = []
+    compound_formatted_string = ""
+
+    if rng.integers(2) == 0:
+        compound_formatted_string += canonical_smiles
+    else:
+        key_value_pairs.append(canonical_smiles)
+
+    key_value_pairs.append(f"[CONFORMER]{conformers['embedded_smiles']}[/CONFORMER]")
+    if label != "nan":
+        key_value_pairs.append(f"[PROPERTY]{label:.2f}[/PROPERTY]" )
+
+    rng.shuffle(key_value_pairs)
+    for kv in key_value_pairs:
+        compound_formatted_string += kv
+
+    return compound_formatted_string
 
 def generate_formatted_string(compound_json, rng, representation_type):
     key_value_pairs = []

train_configs/llama3.2_1b_conformers.toml

@@ -0,0 +1,76 @@
+# torchtitan Config.toml
+
+[job]
+dump_folder = "/nfs/h100/raid/chem/checkpoints"
+description = "Llama 3.2 training"
+use_for_integration_test = false
+
+[profiling]
+enable_profiling = false
+save_traces_folder = "profile_trace"
+profile_freq = 10
+enable_memory_snapshot = false
+save_memory_snapshot_folder = "memory_snapshot"
+
+[metrics]
+log_freq = 1
+enable_color_printing = true
+enable_aim = true
+save_aim_folder = "aim_conformers"
+
+[model]
+name = "llama3"
+flavor = "1B"
+norm_type = "rmsnorm"  # layernorm / np_layernorm / rmsnorm / compiled_rmsnorm / fused_rmsnorm
+tokenizer_path = "torchtitan/tokenizers/Llama-3.2-chem-1B-v1/"
+
+[optimizer]
+name = "AdamW"
+lr = 6e-4
+
+[training]
+batch_size = 6
+gradient_accumulation_steps = 1
+seq_len = 2048
+warmup_steps = 1  # lr scheduler warm up, normally 20% of the train steps
+max_norm = 1.0  # grad norm clipping
+steps = 5200
+data_parallel_degree = -1
+tensor_parallel_degree = 1
+compile = true
+# dataset = "c4"  # supported datasets: c4_test (2K), c4 (177M)
+# dataset = "chemlactica_train_mini"  # supported datasets: c4_test (2K), c4 (177M), chemlactica_train_mini (4K)
+dataset = "conformers_train"
+data_processing_style="conformer_style"
+representation_type = "SMILES"
+
+[validation]
+valid_freq = 1000
+enable_valid = true
+dataset = "conformers_valid"  # supported datasets: chemlactica_valid_mini
+
+[dataloader]
+num_workers = 1
+
+[experimental]
+pipeline_parallel_degree = 1
+enable_async_tensor_parallel = false
+
+[checkpoint]
+enable_checkpoint = true
+save_folder = "yerevann/Llama-3.2-1B_conformers"
+load_folder = "meta-llama/Llama-3.2-1B"
+# load_folder = "yerevann/Llama-3.2-1B/7b98d06b463e45ea8db87d05"
+load_at_step = 0
+interval_type = "steps"
+interval = 5200
+model_weights_only = false
+export_dtype = "float32"
+async_mode = "async_with_pinned_mem"  # ["disabled", "async", "async_with_pinned_mem"]
+
+[activation_checkpoint]
+mode = 'none'  # ['none', 'selective', 'full']
+selective_ac_option = '2'  # 'int' = ac every positive int layer or 'op', ac based on ops policy
+
+[float8]
+enable_float8_linear = false