-- extends "projects/lm_training_project.yaml"

[config_metadata]
    == super()
    ## Overrides
    -- set ns.config_name = "Pretrain Small-LLM"
    -- set ns.config_description = "Pretrain a small LLM"
    -- set ns.model_name = model_name | default("sllm")
    -- set ns.trainer_type = trainer_type | default("ddp")

    ## Dataset
    -- set ns.dataset_proj = abspath(joinpath(ns.forgather_dir, "examples", "datasets", "HuggingFaceTB"))
    -- set ns.dataset_config = dataset_config | default("smollm-corpus/interleaved-packed.yaml")

    ## Model
    -- set ns.model_project_dir = abspath(model_project | default((joinpath(ns.forgather_dir, "examples", "models", "llama"))))
    -- set ns.model_project_config = model_config | default("medium.yaml")

    ## Training budget (millions of tokens)
    -- set ns.total_tokens = total_tokens | default(2270)
    -- set ns.warmup_tokens = warmup_tokens | default(114)
    -- set ns.min_cooldown_tokens = min_cooldown_tokens | default(6800)

    ## Batching
    -- set ns.seq_len = seq_len | default(4096)
    -- set ns.per_device_train_batch_size = batch_size | default(4)

    ## LR scaling
    -- set ns.base_lr = lr | default(1.5e-4)

    ## LR Scheduler
    -- set ns.constant_lr_ratio = 1.0 / 3
    -- set ns.cooldown_p = 0.3
    -- set ns.annealing_tokens = annealing_tokens | default(227)
    
    ## Attention
    -- set ns.attn_implementation = attn_implementation | default("flex_attention")

    ## Compile
    -- set ns.torch_compile = compile | default("True")
    -- set ns.torch_compile_mode = torch_compile_mode | default("max-autotune")

[globals]
    == super()
    -- set ns.constant_lr = ns.global_lr * ns.constant_lr_ratio
    -- set ns.cooldown_steps = [ns.min_cooldown_steps, ns.total_steps * ns.cooldown_p] | max | int
    -- set ns.annealing_tokens = ns.annealing_tokens * 1000 * 1000
    -- set ns.annealing_steps = ns.annealing_tokens // ns.tokens_per_step
    -- set ns.decay_start_step = (ns.total_steps - ns.annealing_steps) | int

[variable_listing]
    == super()
# ns.constant_lr_ratio: {{ ns.constant_lr_ratio }}
# ns.constant_lr: {{ ns.constant_lr }}
# ns.cooldown_p: {{ ns.cooldown_p }}
# ns.cooldown_steps: {{ ns.cooldown_steps }}
# ns.annealing_steps: {{ ns.annealing_steps }}

[trainer_args]
    == super()

    # Checkpoint saving
    save_strategy: {{ save_strategy | toyaml('steps') }}
    save_total_limit: 4
    preserve_best_model: True
    best_model_metric: "loss"
    preserve_n_best: 1
    eval_on_save: True

    # Compile
    torch_compile_dynamic: False

    # Dataloader
    dataloader_num_workers: 1
    dataloader_prefetch_factor: 4

    # Misc.
    float32_matmul_precision: {{ float32_matmul_precision | toyaml("high") }}
    default_dtype: {{ default_dtype | toyaml("float32") }}
    mixed_precision: {{ mixed_precision | toyaml("bf16") }}
    max_grad_norm: 1.0

[optimizer]
optimizer: &optimizer !partial:torch:optim.AdamW
    lr: {{ ns.global_lr | toyaml }}
    foreach: True

[lr_scheduler]
lr_scheduler: &lr_scheduler !partial:forgather.ml.optim:InfiniteLRScheduler@lr_scheduler
    warmup_steps: {{ ns.warmup_steps }}
    cooldown_steps: {{ ns.cooldown_steps }}
    constant_lr: {{ ns.constant_lr | toyaml }}
    min_lr: {{ ns.min_lr | toyaml }}
    annealing_type: "rsqrt"
    start_annealing: {{ start_annealing | toyaml(False) }}
    annealing_steps: {{ ns.annealing_steps }}

[datacollator]
    == super()
    max_length: {{ ns.seq_len }}
    ## Pad batches to 'max_length'; This minimizes recompiles, but should
    ## may be inefficient if not using a 'packed' dataset configuration.
    padding: "max_length"

[fused_loss_factory]
### Fused output-linear x cross-entropy-loss kernel
.define: &fused_loss_factory !partial:forgather.ml.loss:LinearCrossEntropyLoss

[trainer_callbacks]
    -- include 'project.logger_config'

[dynamic_args]
    == super()
    annealing_tokens:
        names: "--annealing-tokens"
        type: int
        group: "Trainer:Schedule"
        min: 1
        help: "Number of tokens in annealing phase."
    start_annealing:
        names: "--start-annealing"
        action: "store_true"
        group: "Trainer:Schedule"
        help: "Manually trigger annealing phase"

#-------------------- project.logger_config --------------------
-- extends "callbacks/lm_callbacks.yaml"