diff --git a/src/molgen3D/config/paths.yaml b/src/molgen3D/config/paths.yaml index 8dc0c3a..9920324 100644 --- a/src/molgen3D/config/paths.yaml +++ b/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,6 +37,7 @@ 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" @@ -36,10 +45,12 @@ base_paths: 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: diff --git a/src/molgen3D/config/pretrain/qwen3_06b.toml b/src/molgen3D/config/pretrain/qwen3_06b.toml index c41f27c..8609358 100644 --- a/src/molgen3D/config/pretrain/qwen3_06b.toml +++ b/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 diff --git a/src/molgen3D/config/pretrain/qwen3_06b_revisited_cartesian_isomeric_4e.toml b/src/molgen3D/config/pretrain/qwen3_06b_revisited_cartesian_isomeric_4e.toml new file mode 100644 index 0000000..c4600bf --- /dev/null +++ b/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" diff --git a/src/molgen3D/config/pretrain/qwen3_06b_revisited_quantile_binned_isomeric_4e.toml b/src/molgen3D/config/pretrain/qwen3_06b_revisited_quantile_binned_isomeric_4e.toml new file mode 100644 index 0000000..8ed4e3c --- /dev/null +++ b/src/molgen3D/config/pretrain/qwen3_06b_revisited_quantile_binned_isomeric_4e.toml @@ -0,0 +1,96 @@ +[job] +description = "qwen3_06b_pre_4e_revisited_quantile_binned_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_binned_258" +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_binned_258" + +[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 = 27200 +max_norm = 1.0 +dtype = "bfloat16" +dataset = "molgen_jsonl" +dataset_path = "data:revisited_quantile_binned_isomeric_train" + +[validation] +enable = false +dataset = "molgen_jsonl" +dataset_path = "data:revisited_quantile_binned_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_quantile_binned_isomeric_train" +tokenizer_key = "qwen3_0.6b_binned_258" +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 = [6800, 13600, 20400, 27200] +lr_min = 0 +decay_frac = 0.1 + +[checkpoint] +enable = true +interval = 6800 +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" diff --git a/src/molgen3D/config/pretrain/qwen3_06b_revisited_uniform_binned_isomeric_4e.toml b/src/molgen3D/config/pretrain/qwen3_06b_revisited_uniform_binned_isomeric_4e.toml new file mode 100644 index 0000000..fef9516 --- /dev/null +++ b/src/molgen3D/config/pretrain/qwen3_06b_revisited_uniform_binned_isomeric_4e.toml @@ -0,0 +1,96 @@ +[job] +description = "qwen3_06b_pre_4e_revisited_uniform_binned_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_binned_258" +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_binned_258" + +[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 = 27200 +max_norm = 1.0 +dtype = "bfloat16" +dataset = "molgen_jsonl" +dataset_path = "data:revisited_uniform_binned_isomeric_train" + +[validation] +enable = false +dataset = "molgen_jsonl" +dataset_path = "data:revisited_uniform_binned_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_uniform_binned_isomeric_train" +tokenizer_key = "qwen3_0.6b_binned_258" +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 = [6800, 13600, 20400, 27200] +lr_min = 0 +decay_frac = 0.1 + +[checkpoint] +enable = true +interval = 6800 +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"