CrystaLLM-pi_alex_mp_20_base / training_config.jsonc
c-bone's picture
Use the standard CIF tokenizer
6b7077d verified
Raw History Blame Contribute Delete
1.72 kB
{
// Data Arguments
//################
"dataset_HF": "c-bone/alex_mp_20",
"pretrained_tokenizer_dir": "_utils/HF-cif-tokenizer",
"context_length": 1024,
// Filters
"remove_CIFs_above_context": false,
"remove_CIFs_with_unk": true,
// "activate_conditionality": "None",
// Model Arguments
//#################
// Model Depth
// n_positions has been tied to context_length
"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,
// Learning Rate and Optimizer
"learning_rate": 0.001,
"lr_scheduler_type": "cosine_with_min_lr",
"lr_scheduler_kwargs": {
"min_lr_rate": 0.01
},
"warmup_ratio": 0.0,
"adam_beta1": 0.9,
"adam_beta2": 0.98,
"grad_clip": 1.0,
"weight_decay": 0.01,
"optimizer": "muon",
"muon_lr": 0.01,
"muon_momentum": 0.95,
// Logging
"output_dir": "model_ckpts/alex_mp_20_text/",
"save_total_limit": 2,
"report_to": "wandb",
"wandb_project_folder": "CrystaLLM-pi",
"eval_strategy": "steps",
"eval_steps": 10000,
"logging_steps": 50,
"save_strategy": "steps",
"max_steps": 375000,
"early_stopping_patience": 15,
"early_stopping_threshold": 0.000005,
// 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,
"deepspeed_config": "_config_files/deepspeed_default.json",
// CodeCarbon Arguments
//######################
"codecarbon": true,
"tracker_project": "CrystaLLM-pi"
}