train_soc_243_gpu.toml

# DeepH-pack GPU training config for 243 OpenMX SOC structures.
# Run with:
#   sbatch /home/jinliangl/home2/songby/20260514/ab-stacking/slurm_train_deeph_soc_243_gpu.sh

[system]
note = "OpenMX-3.9 SOC, 243 structures, 8-GPU training"
device = "gpu*8:0"
float_type = "fp32"
random_seed = 137
log_level = "info"
jax_memory_preallocate = false
show_train_process_bar = true

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

[data.dft]
data_dir_depth = 0
validation_check = false

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

[data.model_save]
best = true
latest = true
latest_interval = 1
latest_num = 10
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 = "8x0e+8x1o+8x2e+8x3o+8x4e"
latent_irreps = "8x0e+8x1o+8x2e+8x3o+8x4e"
latent_edge_cutoff = 100.0
num_blocks = 2
num_heads = 2
enable_bs3b_layer = false
bs3b_orbital_types = ""
consider_parity = true
standardize_gauge = false
vr_focus_size = 256

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

[process.train.dataloader]
batch_size = 1
train_size = 195
validate_size = 24
test_size = 24
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 = 2E-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 = 80
rtol = 0.05
cooldown = 20
accum_size = -1
init_scale = 0.1
warmup_steps = 10
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