llama3_750m.toml

# 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 = false
save_aim_folder = "aim"

[model]
name = "llama3"
flavor = "750M"
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 = 4e-4

[training]
batch_size = 2
gradient_accumulation_steps = 16
seq_len = 2048
max_norm = 1.0  # grad norm clipping
warmup_steps = 500  # lr scheduler warm up, normally 20% of the train steps
steps = 80000
decay_steps_perc = 0.1
num_decays = 1
data_parallel_degree = -1
tensor_parallel_degree = 1
compile = true
dataset = "chemlactica_train"
data_processing_style = "chemlactica_style"
representation_type = "SMILES"

[validation]
valid_freq = 2000
enable_valid = true
dataset = "chemlactica_valid"  # supported datasets: chemlactica_valid_mini

[dataloader]
num_workers = 2

[experimental]
pipeline_parallel_degree = 1
enable_async_tensor_parallel = false

[checkpoint]
enable_checkpoint = true
save_folder = "yerevann/Llama-3.2-750M"
# load_folder = "yerevann/Llama-3.2-750M"
# load_at_step = 0
interval_type = "steps"
interval = 2000
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