train_soc_abstacking_243_Bi_s3p2d2f1_gpu_lmax6_batch2.toml

# DeepH-pack GPU training config for the ab-stacking Bi s3p2d2f1 OpenMX SOC dataset.
# Short comparison run with batch_size=2. It writes to a separate output
# directory from the batch_size=1 production run.
#
# Run with:
#   sbatch /home/jinliangl/home2/songby/20260514/ab-stacking/slurm_train_deeph_soc_abstacking_243_Bi_s3p2d2f1_gpu_lmax6_batch2.sh

[system]
note = "OpenMX-3.9 SOC, ab-stacking 243 structures, Bi8.0-s3p2d2f1, lmax=6 8-GPU batch_size=2 comparison training"
device = "gpu*8:0-7"
float_type = "fp32"
random_seed = 137
log_level = "info"
jax_memory_preallocate = false
show_train_process_bar = false

[data]
inputs_dir = "/home/jinliangl/home2/songby/20260514/ab-stacking/deeph_inputs_soc_abstacking_243_Bi_s3p2d2f1"
outputs_dir = "/home/jinliangl/home2/songby/20260514/ab-stacking/outputs_train_soc_abstacking_243_Bi_s3p2d2f1_gpu_lmax6_batch2"

[data.dft]
data_dir_depth = 0
validation_check = false

[data.graph]
dataset_name = "OPENMX39-SOC-ABSTACKING-243-Bi-s3p2d2f1-GPU"
graph_type = "H"
storage_type = "memory"
disk_shards_num = 1
disk_shards_indices = []
disk_mem_buffer_size = 4096
common_orbital_types = ""
parallel_num = 64
only_save_graph = false

[data.model_save]
best = true
latest = true
latest_interval = 5
latest_num = 5
latest_cache_mid_step_interval = -1

[model]
net_type = "sparrow"
target_type = "H"
loss_type = "mse"

[model.advanced]
gaussian_basis_rmax = 7.5
net_irreps = "16x0e+16x1o+16x2e+16x3o+16x4e+16x5o+16x6e"
latent_irreps = "16x0e+16x1o+16x2e+16x3o+16x4e+16x5o+16x6e"
latent_edge_cutoff = 100.0
num_blocks = 4
num_heads = 4
enable_bs3b_layer = false
bs3b_orbital_types = ""
consider_parity = true
standardize_gauge = false
vr_focus_size = 256

[process.train]
max_epoch = 1500
multi_way_jit_num = 1
ahead_of_time_compile = false
do_remat = false

[process.train.dataloader]
batch_size = 2
train_size = 195
validate_size = 24
test_size = 24
dataset_split_json = "/home/jinliangl/home2/songby/20260514/ab-stacking/outputs_train_soc_abstacking_243_gpu_large/2026-05-20_14-40-23/dataset_split.json"
only_use_train_loss = false

[process.train.drop]
dropout_rate = 0.05
stochastic_depth = 0.0
proj_rate = 0.0

[process.train.optimizer]
type = "adamw"
init_learning_rate = 1E-3
clip_norm_factor = -1.0
momentum = 0.8
betas = [0.9, 0.999]
weight = 0.001
eps = 1E-8

[process.train.scheduler]
min_learning_rate_scale = 1E-4
type = "reduce_on_plateau"
factor = 0.5
patience = 100
rtol = 0.03
cooldown = 20
accum_size = -1
init_scale = 0.1
warmup_steps = 20
decay_steps = -1
end_scale = -1.0

[process.train.continued]
enable = false
new_training_data = false
new_optimizer = false
previous_output_dir = ""
load_model_type = "latest"
load_model_epoch = -1

[process.train.with_plugin]
enable = false
plugin = "null"
env_path = "./user/need/to/set/this"
script_path = "./user/need/to/set/this"
backend = "null"
dump_intermediate = false