## See: docs/project-templates/lm-training-projects.md
-- extends "projects/lm_training_project.yaml"

[config_metadata]
    == super()
    ## Overrides
    -- set ns.config_name = "Tiny Llama v2"
    -- set ns.config_description = "Train Tiny Llama using the new lm_training_project.yaml template"
    -- set ns.model_name = model_name | default("v2")

    ## Train on all available GPUs via DDP
    -- set ns.trainer_type = trainer_type | default("ddp")

    ## Dataset
    -- set ns.dataset_proj = abspath(joinpath(ns.forgather_dir, "examples", "datasets", "roneneldan"))
    -- set ns.dataset_config = dataset_config | default("fast-iter-packed.yaml")
    -- set ns.dispatch_batches = dispatch_batches | default(False)

    ## 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("4M.yaml")

    ## Training budget (millions of tokens)
    ## Set to 10x Chinchilla optimal
    -- set ns.total_tokens = total_tokens | default(80)
    -- set ns.warmup_tokens = warmup_tokens | default(8)
    -- set ns.min_cooldown_tokens = min_cooldown_tokens | default(40)

    ## Batching
    -- set ns.seq_len = seq_len | default(2048)
    -- set ns.per_device_train_batch_size = batch_size | default(8)

    ## Step cadence
    ## Eval more frequently than the default.
    -- set ns.base_validation_tokens = 10

    ## LR scaling
    -- set ns.base_lr = lr | default(1.0e-3)
    
    ## Attention
    -- set ns.attn_implementation = attn_implementation | default("sdpa")

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

    ## Prompts
    -- set ns.eval_prompts_file = "prompts/tiny_stories.yaml"
    -- set ns.eval_max_new_tokens = 40
    -- set ns.generation_steps = 500

[trainer_args]
    ## See: docs/trainers/trainer_options.md
    == super()
    # ** Tiny Llama v2**
    # Checkpoint saving
    eval_on_save: True

    # Disable dynamic, as we are using fixed-sized sequence lengths; it's a little faster
    torch_compile_dynamic: False

    # Run dataset processing in another process and pipeline the batches
    dataloader_num_workers: 1
    dataloader_prefetch_factor: 16

    # This reduces the precision of float32 matrix-multiply for speed.
    float32_matmul_precision: {{ float32_matmul_precision | toyaml("high") }}

    # Enable bf16 mixed-precision training for speed.
    mixed_precision: {{ mixed_precision | toyaml("bf16") }}

    # Provide some insurance against gradient-spikes causing issues.
    max_grad_norm: 1.0

[datacollator]
    == super()
    max_length: {{ ns.seq_len }}
    ## Pad batches to 'max_length'; This minimizes recompiles, but
    ## 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 "callbacks/lm_callbacks.yaml"