-- extends "training_script/training_script_type.yaml"
-- from 'inc/formatting.jinja' import h2, sep

[config_metadata]
    == super()
    -- set ns.config_class = ns.config_class + ".causal_lm"

[datasets_preprocessor_args]
tokenizer_args: &tokenizer_args !dict
    truncation: True

[model_constructor_args]
## These are-runtime model construction arguments which
## typically control things like Flash Attention or what
## data-type to convert the model to.
# https://huggingface.co/docs/transformers/en/model_doc/auto
.define: &model_constructor_args
    # See: https://huggingface.co/docs/transformers/en/attention_interface
    attn_implementation: "{{ attn_implementation | default('sdpa') }}"

[model_definition]
## Depending upon the value 'create_new_model,' either load an existing
## model or construct a new one.
## The details for constructing a new model need to be filled in by
## a child template, while an existing model should work 'as-is'
    [construct_new_model]
## Undefined model constructor

[trainer_callbacks]
    -- include 'callbacks/loggers.yaml'

[datacollator]
# Data collator for causal model
# Batches are dynamically padded to longest sequence
# labels are set to input_ids, with pad tokens set to -100
data_collator: &data_collator !singleton:forgather.ml.data_collator:DataCollatorForCausalLM@DataCollatorForCausalLM
    tokenizer: *tokenizer
    return_tensors: pt

[trainer_args]
    == super()
    save_strategy: "{{ save_strategy | default('no') }}"
    max_steps: {{ max_steps | default(-1) }}
    output_dir: "{{ ns.output_dir }}"
    logging_dir: "{{ ns.logging_dir }}"

[trainer_dependencies]
    [fused_loss_factory]
.define: &fused_loss_factory null

[trainer_definition]
    ## Set default trainer
    -- include 'trainers/trainer.yaml'

[dynamic_args]
    == super()
    attn_implementation:
        names: "--attn-implementation"
        type: "str"
        choices: [ "eager", "sdpa", "flash_attention_2", "flex_attention" ]
        group: "Model"
        help: "Attention implementation"
    verbose_info:
        names: "--verbose-info"
        action: "store_true"
        group: "Logging"
        help: "Display verbose training info on startup"
