Add new pretraining configurations and update paths for Qwen3 model

src/molgen3D/config/paths.yaml

@@ -1,26 +1,34 @@
 base_paths:
   project_root: "."
   data_root:
+    - "/data/vtarasov"
     - "/data/molgen"
     - "data"
   geom_data_root:
+    - "/data/vtarasov"
     - "/data/molgen"
     - "/data/molgen/GEOM_data"
   geom_dataset_root:
-    - "/data/molgen/GEOM_data/torsional_diff_gdrive/extracted"
+    - "/mnt/weka/vtarasov/torsional_gdrive/extracted"
+    - "/data/molgen/GEOM_data/torsional_gdrive/extracted"
   ckpts_root:
+    - "/mnt/weka/vtarasov/checkpoints"
     - "/data/chem-project/checkpoints"
   hf_yerevann_root:
     - "/nfs/h100/raid/chem/checkpoints/hf/yerevann"
     - "/data/chem-project/checkpoints/hf/yerevann"
   qwen_yerevann_root:
     - "/data/chem-project/checkpoints"
+    - "/mnt/weka/data/chem-project/checkpoints"
+    - "/home/chem-project/checkpoints"
     - "/data/chem-project/checkpoints/yerevann"
     - "/nfs/h100/raid/chem/checkpoints/yerevann"
   grpo_root:
+    - "/mnt/weka/vtarasov/conf_grpo"
     - "/data/chem-project/checkpoints/conf_grpo"
     - "/nfs/h100/raid/chem/checkpoints/conf_grpo"
   qwen3_grpo_root:
+    - "/mnt/weka/vtarasov/conf_grpo"
     - "/data/chem-project/checkpoints/conf_grpo"
     - "/nfs/h100/raid/chem/checkpoints/conf_grpo"
   gen_results_root: "outputs/gen_results"
@@ -29,17 +37,20 @@ base_paths:
   pretrain_logs_root: "outputs/pretrain_logs"
   wandb_root: "wandb_runs"
   qwen3_0.6b_base_model:
+    - "/mnt/weka/vtarasov/Qwen3-0.6B-Base"
     - "/data/chem-project/checkpoints/Qwen3-0.6B-Base"
     - "/nfs/h100/raid/chem/checkpoints/yerevann/qwen3_06b/Qwen3-0.6B-Base"
     - "/data/chem-project/checkpoints/yerevann/qwen3_06b/Qwen3-0.6B-Base"
   tokenizers_root: "src/molgen3D/training/tokenizers"
 
 ckpts:
   qwen3_06b:
+    - "/mnt/weka/vtarasov/qwen3_06b"
     - "/data/chem-project/checkpoints/qwen3_06b"
     - "/nfs/h100/raid/chem/checkpoints/yerevann/qwen3_06b"
     - "/data/chem-project/checkpoints/yerevann/qwen3_06b"
   qwen3_grpo:
+    - "/mnt/weka/vtarasov/conf_grpo"
     - "/data/chem-project/checkpoints/conf_grpo"
     - "/nfs/h100/raid/chem/checkpoints/conf_grpo"
 
@@ -48,6 +59,7 @@ data:
   distinct_smi: "tmp/distinct_smi_filtered_13.pickle"
   xl_smi: "xl_smi.pickle"
   qm9_smi: "qm9_smi.pickle"
+  revisited_smi: "revisited_smi.pickle"
   icl_5_smi: "icl_5_smi.pickle"
   icl_10_smi: "icl_10_smi.pickle"
   icl_20_smi: "icl_20_smi.pickle"
@@ -116,12 +128,75 @@ data:
   filtered_conformers_test:
     - "3DMolGen_data/binned_filtered/processed_strings/test"
     - "binned_filtered/processed_strings/test"
-
+  revisited_cartesian_train:
+    - "geom_revisited_cartesian/processed_strings/train"
+  revisited_cartesian_valid:
+    - "geom_revisited_cartesian/processed_strings/valid"
+  revisited_cartesian_test:
+    - "geom_revisited_cartesian/processed_strings/test"
+
+  revisited_quantile_binned_train:
+    - "geom_revisited_quantile_binned/processed_strings/train"
+  revisited_quantile_binned_valid:
+    - "geom_revisited_quantile_binned/processed_strings/valid"
+  revisited_quantile_binned_test:
+    - "geom_revisited_quantile_binned/processed_strings/test"
+
+  revisited_uniform_binned_train:
+    - "geom_revisited_uniform_binned/processed_strings/train"
+  revisited_uniform_binned_valid:
+    - "geom_revisited_uniform_binned/processed_strings/valid"
+  revisited_uniform_binned_test:
+    - "geom_revisited_uniform_binned/processed_strings/test"
+
+  revisited_cartesian_isomeric_grouped_train:
+    - "geom_revisited_cartesian_isomeric_grouped/processed_strings/train"
+  revisited_cartesian_isomeric_grouped_valid:
+    - "geom_revisited_cartesian_isomeric_grouped/processed_strings/valid"
+  revisited_cartesian_isomeric_grouped_test:
+    - "geom_revisited_cartesian_isomeric_grouped/processed_strings/test"
+
+  revisited_quantile_binned_isomeric_grouped_train:
+    - "geom_revisited_quantile_binned_isomeric_grouped/processed_strings/train"
+  revisited_quantile_binned_isomeric_grouped_valid:
+    - "geom_revisited_quantile_binned_isomeric_grouped/processed_strings/valid"
+  revisited_quantile_binned_isomeric_grouped_test:
+    - "geom_revisited_quantile_binned_isomeric_grouped/processed_strings/test"
+
+  revisited_uniform_binned_isomeric_grouped_train:
+    - "geom_revisited_uniform_binned_isomeric_grouped/processed_strings/train"
+  revisited_uniform_binned_isomeric_grouped_valid:
+    - "geom_revisited_uniform_binned_isomeric_grouped/processed_strings/valid"
+  revisited_uniform_binned_isomeric_grouped_test:
+    - "geom_revisited_uniform_binned_isomeric_grouped/processed_strings/test"
+
+
+  revisited_uniform_binned_isomeric_train:
+    - "geom_revisited_uniform_binned_isomeric/processed_strings/train"
+  revisited_uniform_binned_isomeric_valid:
+    - "geom_revisited_uniform_binned_isomeric/processed_strings/valid"
+  revisited_uniform_binned_isomeric_test:
+    - "geom_revisited_uniform_binned_isomeric/processed_strings/test"
+
+  revisited_quantile_binned_isomeric_train:
+    - "geom_revisited_quantile_binned_isomeric/processed_strings/train"
+  revisited_quantile_binned_isomeric_valid:
+    - "geom_revisited_quantile_binned_isomeric/processed_strings/valid"
+  revisited_quantile_binned_isomeric_test:
+    - "geom_revisited_quantile_binned_isomeric/processed_strings/test"
+
+  revisited_cartesian_isomeric_grouped_train:
+    - "geom_revisited_cartesian_isomeric_grouped/processed_strings/train"
+  revisited_cartesian_isomeric_grouped_valid:
+    - "geom_revisited_cartesian_isomeric_grouped/processed_strings/valid"
+  revisited_cartesian_isomeric_grouped_test:
+    - "geom_revisited_cartesian_isomeric_grouped/processed_strings/test"
 tokenizers:
   llama3_chem_v1: "src/molgen3D/training/tokenizers/Llama-3.2-chem-1B-v1"
   qwen3_0.6b_custom: "src/molgen3D/training/tokenizers/Qwen3_tokenizer_custom"
   qwen3_0.6b_origin: "src/molgen3D/training/tokenizers/Qwen3_tokenizer_original"
   qwen3_0.6b_binned: "src/molgen3D/training/tokenizers/Qwen3_tokenizer_binned"
+  qwen3_0.6b_binned_258: "src/molgen3D/training/tokenizers/Qwen3_tokenizer_binned_258"
 
 models:
   # ---------- BASE MODELS ----------
@@ -285,6 +360,30 @@ models:
       "4e": "step-22000-hf"
       "5e": "step-27500-hf"
 
+  qw600_pre_binned_revisited_uniform_isomeric:
+    root: "qwen3_06b/260308-0317-2013-qwen3_06b_pre_4e_revisited_uniform_binned_isomeric/"
+    steps:
+      "1e": "step-6800-hf"
+      "2e": "step-13600-hf"
+      "3e": "step-20400-hf"
+      "4e": "step-27200-hf"
+
+  qw600_pre_binned_revisited_quantile_isomeric:
+    root: "qwen3_06b/260308-0316-a956-qwen3_06b_pre_4e_revisited_quantile_binned_isomeric/"
+    steps:
+      "1e": "step-6800-hf"
+      "2e": "step-13600-hf"
+      "3e": "step-20400-hf"
+      "4e": "step-27200-hf"
+
+  qw600_pre_binned_revisited_cartesian_isomeric:
+    root: "qwen3_06b/260308-0316-05c6-qwen3_06b_pre_4e_revisited_cartesian_isomeric/"
+    steps:
+      "1e": "step-20500-hf"
+      "2e": "step-41000-hf"
+      "3e": "step-61500-hf"
+      "4e": "step-82000-hf"
+
   # ---------- GRPO / FINETUNED MODELS (same structure) ----------
 
   m380_1e_1xgrpo_100e_100s:

src/molgen3D/config/pretrain/qwen3_06b.toml

@@ -6,14 +6,14 @@ custom_config_module = "molgen3D.training.pretraining.config.custom_job_config"
 
 [molgen_run]
 init_mode = "hf_pretrain" # scratch | hf_pretrain | resume
-tokenizer_tag = "tokenizers:qwen3_0.6b_custom"
+tokenizer_tag = "tokenizers:qwen3_0.6b_binned_258"
 base_model_tag = "base_paths:qwen3_0.6b_base_model"
 # resume_run_path_tag = "ckpts:qwen3_06b/251230-1352-da5b-qwen3_06b_pre_5e_48effb_4kseq/step-90000"
 
 [model]
 name = "molgen_qwen3"
 flavor = "0.6B"
-hf_assets_path = "src/molgen3D/training/tokenizers/Qwen3_tokenizer_binned"
+hf_assets_path = "src/molgen3D/training/tokenizers/Qwen3_tokenizer_binned_258"
 
 [experimental]
 custom_import = "molgen3D.training.pretraining.torchtitan_model.qwen3_custom"
@@ -23,18 +23,18 @@ gc_freq = 1000
 seq_len = 4096
 local_batch_size = 6
 global_batch_size = 48
-steps = 34000
+steps = 28000
 max_norm = 1.0
 dtype = "bfloat16"
 dataset = "molgen_jsonl"
 # dataset_path = "data:geom_conformers_train"
-dataset_path = "data:conformers_train"
+dataset_path = "data:revisited_uniform_binned_train"
 
 [validation]
-enable = true
+enable = false
 dataset = "molgen_jsonl"
 # dataset_path = "data:geom_conformers_valid"
-dataset_path = "data:conformers_valid"
+dataset_path = "data:revisited_uniform_binned_valid"
 local_batch_size = 4
 seq_len = 4096
 freq = 500
@@ -45,8 +45,8 @@ numerical_validation_freq = 10000
 
 [molgen_data]
 # train_path_key = "conformers_train"
-train_path = "data:conformers_train"
-tokenizer_key = "qwen3_0.6b_custom"
+train_path = "data:revisited_q_binned_train"
+tokenizer_key = "qwen3_0.6b_binned_258"
 min_emb_len = 0
 shuffle_lines = true
 infinite = true
@@ -59,7 +59,7 @@ prefetch_factor = 2
 lookahead_limit = 100
 shuffle_buffer_size = 128
 #serialization_mode = "pairs"
-serialization_mode = "isomer_units"
+serialization_mode = "pairs"
 
 [optimizer]
 name = "AdamW"
@@ -76,13 +76,13 @@ weight_decay = 0.1
 [wsds_scheduler]
 enable = true
 warmup_steps = 500
-checkpoints = [8500, 17000, 25500, 34000]
+checkpoints = [7000, 14000, 21000, 28000]
 lr_min = 0
 decay_frac = 0.1
 
 [checkpoint]
 enable = true
-interval = 20000
+interval = 7000
 keep_latest_k = 4
 folder = "checkpoint"
 create_seed_checkpoint = false

src/molgen3D/config/pretrain/qwen3_06b_revisited_cartesian_isomeric_4e.toml

@@ -0,0 +1,96 @@
+[job]
+description = "qwen3_06b_pre_4e_revisited_cartesian_isomeric"
+dump_folder = "pretrain_runs"
+print_config = true
+custom_config_module = "molgen3D.training.pretraining.config.custom_job_config"
+
+[molgen_run]
+init_mode = "hf_pretrain" # scratch | hf_pretrain | resume
+tokenizer_tag = "tokenizers:qwen3_0.6b_custom"
+base_model_tag = "base_paths:qwen3_0.6b_base_model"
+
+[model]
+name = "molgen_qwen3"
+flavor = "0.6B"
+hf_assets_path = "src/molgen3D/training/tokenizers/Qwen3_tokenizer_custom"
+
+[experimental]
+custom_import = "molgen3D.training.pretraining.torchtitan_model.qwen3_custom"
+
+[training]
+gc_freq = 1000
+seq_len = 4096
+local_batch_size = 6
+global_batch_size = 48
+steps = 82000
+max_norm = 1.0
+dtype = "bfloat16"
+dataset = "molgen_jsonl"
+dataset_path = "data:revisited_cartesian_isomeric_train"
+
+[validation]
+enable = false
+dataset = "molgen_jsonl"
+dataset_path = "data:revisited_cartesian_isomeric_valid"
+local_batch_size = 4
+seq_len = 4096
+freq = 500
+steps = 200
+numerical_validation = false
+num_val_molecules = 10
+numerical_validation_freq = 10000
+
+[molgen_data]
+train_path = "data:revisited_cartesian_isomeric_train"
+tokenizer_key = "qwen3_0.6b_custom"
+min_emb_len = 0
+shuffle_lines = true
+infinite = true
+seed = 2025
+num_workers = 8
+pin_memory = true
+drop_last = true
+persistent_workers = false
+prefetch_factor = 2
+lookahead_limit = 100
+shuffle_buffer_size = 128
+serialization_mode = "pairs"
+
+[optimizer]
+name = "AdamW"
+lr = 8e-4
+beta1 = 0.9
+beta2 = 0.95
+eps = 1e-8
+weight_decay = 0.1
+
+[wsds_scheduler]
+enable = true
+warmup_steps = 500
+checkpoints = [20500, 41000, 61500, 82000]
+lr_min = 0
+decay_frac = 0.1
+
+[checkpoint]
+enable = true
+interval = 20500
+keep_latest_k = 4
+folder = "checkpoint"
+create_seed_checkpoint = false
+initial_load_model_only = true
+initial_load_in_hf = true
+initial_load_path = "base_paths:qwen3_0.6b_base_model"
+last_save_in_hf = true
+async_mode = "async"
+
+[metrics]
+log_freq = 20
+enable_wandb = false
+save_for_all_ranks = false
+save_tb_folder = "tb"
+
+[parallelism]
+data_parallel_replicate_degree = 1
+data_parallel_shard_degree = -1
+tensor_parallel_degree = 1
+fsdp_reshard_after_forward = "default"