-- extends 'samantha.yaml'

[config_metadata]
    == super()
    -- set ns.config_name = "Samantha Llama7B 16 GB"
    -- set ns.config_description = "Train on 16 GB GPU"
    -- set ns.log_name = "16gb"

    ## Dataset
    -- set ns.dataset_config = dataset_config | default("samantha.yaml")

    ## Batching
    -- set ns.seq_len = seq_len | default(1280)
    -- set ns.per_device_train_batch_size = batch_size | default(1)

    ## Attention
    -- set ns.attn_implementation = attn_implementation | default("sdpa")

    ## Compile settings
    -- set ns.torch_compile = compile | default(False)

[trainer_args]
    == super()
    gradient_checkpointing: {{ gradient_checkpointing | default(True) }}
    fuse_optim_with_backward: {{ fuse_optim_with_backward | default(True) }}
    ## This appears to be incompatible with flex-attention, so we have fallen back to SDPA
    enable_activation_offloading: {{ activation_offloading | default(True) }}

[datacollator]
data_collator: &data_collator !singleton:forgather.ml.data_collator:DataCollatorForCausalLM@DataCollatorForCausalLM
    tokenizer: *tokenizer
    return_tensors: pt