CrystaLLM-pi_base / training_config.jsonc
c-bone's picture
Training config as in the reproduction repo
ff52b5e verified
Raw History Blame Contribute Delete
1.71 kB
{
// Data Arguments
//################
"dataset_HF": "c-bone/lematerial_clean",
"pretrained_tokenizer_dir": "HF-cif-tokenizer",
"context_length": 1024,
"dataset_streaming": false,
// Filters
"remove_CIFs_above_context": false,
"remove_CIFs_with_unk": true,
// Conditional Arguments
//#######################
"activate_conditionality": "None",
// Model Arguments
//#################
// Model Depth
"n_embd": 512,
"n_layer": 8,
"n_head": 8,
// Dropout
"residual_dropout": 0.1,
"embedding_dropout": 0.1,
"attention_dropout": 0.1,
// Trainer Arguments
//###################
// Batching
"train_batch_size": 32,
"eval_batch_size": 32,
"gradient_accumulation_steps": 1,
"auto_find_batch_size": false,
// Learning Rate and Optimizer
"learning_rate": 0.001,
"lr_scheduler_type": "cosine_with_min_lr",
"lr_scheduler_kwargs": {
"min_lr_rate": 0.001
},
"warmup_steps": 0,
"adam_beta1": 0.85,
"adam_beta2": 0.98,
"grad_clip": 1.0,
"weight_decay": 0.1,
// Logging
"output_dir": "model_ckpts/mpdb-small-base-lematerial",
"save_total_limit": 2,
"report_to": "wandb",
"eval_strategy": "steps",
"eval_steps": 5000,
"logging_steps": 500,
"save_strategy": "steps",
"max_steps": 1250000,
"early_stopping_patience": 10,
"early_stopping_threshold": 0.00001,
// Utils
"seed": 1,
"data_seed": 1,
"load_best_model_at_end": true,
"metric_for_best_model": "eval_loss",
"greater_is_better": false,
"torch_compile": true,
"fp16": true,
"deepspeed_config": "_config_files/deepspeed_default.json",
// CodeCarbon Arguments
//######################
"codecarbon": true,
"tracker_project": "CrystaLLM-pi"
}