-- extends "training_script/causal_lm/causal_lm.yaml"
## ---------------------------------------------------------------------------
## LM Training Project
##
## A complete training project template that computes the number of training
## steps from a target token budget. The budget is typically set relative to
## the "Chinchilla-Optimal" ratio (see background below).
##
## Includes a cosine-decay LR scheduler with warmup and an AdamW optimizer
## with automatic LR scaling based on global batch size. This allows you to
## change the batch size, number of GPUs, or gradient accumulation steps
## without manually retuning the learning rate. The scaling uses a power-law
## rule:
##
##   lr = base_lr * (global_batch_size / base_batch_size) ** alpha
##
## ---------------------------------------------------------------------------
## Configuration Parameters
##
## All token counts are specified in MILLIONS unless noted otherwise.
##
## Training Budget
##   total_tokens           (int)   Total training tokens.          Default: 560
##   warmup_tokens          (int)   LR warmup tokens.               Default: 56
##   min_cooldown_tokens    (int)   Minimum cosine-decay window;    Default: 300
##                                  prevents decay to 0 on short runs.
##
## Batching
##   seq_len                (int)   Maximum sequence length.        Default: 512
##   batch_size             (int)   Per-device training batch size. Default: 32
##   gradient_accumulation_steps
##                          (int)   Gradient accumulation steps.    Default: 1
##   batch_density          (float) Estimated fraction of non-pad   Default: 0.90
##                                  tokens per batch; used to
##                                  correct token-count estimates.
##
## Data
##   dataset_project        (path)  Path to dataset project.        Default: examples/datasets/HuggingFaceTB
##   dataset_config         (str)   Dataset project configuration.  Default: smollm-corpus/fineweb-edu-packed.yaml
##   dispatch_batches       (bool)  When True, rank-0 loads and     Default: False
##                                  dispatches all batches (DDP).
##                                  Prefer False for throughput:
##                                  centralised loading adds per-step
##                                  dispatch overhead and serialises
##                                  data loading on rank-0. Enable only
##                                  when (a) the dataset doesn't support
##                                  sharding, or (b) you need the DDP run
##                                  to see the same example sequence as
##                                  a single-GPU / pipeline run so that
##                                  training curves are directly
##                                  comparable.
##
## Model
##   model_project          (path)  Path to model project.          Default: examples/models/llama
##   model_config           (str)   Model project configuration.    Default: small.yaml (28M parameters)
##   attn_implementation    (str)   Attention backend.              Default: sdpa
##
## Optimizer / Scheduler / LR Scaling
##   lr                     (float) Base learning rate at           Default: 3e-4
##                                  base_batch_size.
##   base_batch_size        (int)   Reference batch size (tokens)   Default: 16384
##                                  for LR scaling calculation.
##   lr_alpha               (float) Scaling exponent.               Default: 0.5
##                                    0.0 = no scaling
##                                    0.5 = sqrt (noise-dominated)
##                                    1.0 = linear (signal-dominated)
##
## Trainer Selection
##   trainer_type           (str)   Trainer backend.                Choices: basic, ddp, fsdp2, pipeline
##                                                                  Default: basic
##   nproc_per_node         (str|int) Processes per node.           Default: auto (1 for basic,
##                                                                           "gpu" for ddp/fsdp2/pipeline)
##   pipeline_schedule      (str)   PP schedule class.              Default: ScheduleInterleaved1F1B
##   microbatch_scale       (int)   PP microbatch scale factor.     Default: 1
##
## Step Cadence (logging / eval / save intervals)
##   step_cadence           (float) Multiplier applied to the base  Default: 1.0
##                                  token intervals below.
##   base_logging_tokens    (int)   Tokens (M) between log steps.   Default: 1
##   base_validation_tokens (int)   Tokens (M) between eval steps.  Default: 25
##   base_save_tokens       (int)   Tokens (M) between save steps.  Default: 500
##
## Hardware / Performance
##   peak_hardware_flops    (float) Peak device FLOPS for MFU       Default: auto-detected
##                                  computation. Auto-detected from
##                                  the current GPU, or read from
##                                  ~/.config/forgather/hardware.yaml.
##                                  See docs/trainers/training-performance-metrics.md
##
## Precision / Compilation
##   default_dtype          (str)   Torch dtype for model construction. Choices: float32, bfloat16, float16
##   float32_matmul_precision
##                          (str)   Approximate float32 matmul with bf16. Choices: highest, high, medium
##   mixed_precision        (str)   AMP dtype.                      Choices: bf16, fp16, no
##   fp8_recipe             (str)   FP8 recipe for linear layers.   Choices: tensorwise, rowwise, rowwise_with_gw_hp
##   qat_recipe             (str)   QAT recipe for linear layers.   Choices: see forgather.ml.qat_recipes.QAT_RECIPES (rendered into --qat-recipe --help)
##   compile                (bool)  Enable torch.compile.           Default: False
##   torch_compile_mode     (str)   torch.compile mode.             Default: default
##                                  Choices: default, max-autotune, max-autotune-no-cudagraphs.
##                                  max-autotune is incompatible with gradient accumulation
##                                  and pipeline parallel; use max-autotune-no-cudagraphs
##                                  instead. Zero-bubble pipeline schedules are incompatible
##                                  with torch_compile entirely.
##
## Memory / Performance
##   gradient_checkpointing (bool)  Enable activation checkpointing Default: False
##                                  on models that implement the HF
##                                  gradient-checkpointing API.
##                                  Trades extra recompute for lower
##                                  peak activation memory.
##   activation_offloading  (bool)  Offload saved activations to    Default: False
##                                  CPU memory during backward. Pair
##                                  with gradient_checkpointing for
##                                  maximum memory savings.
##   activation_memory_budget
##                          (float) Fraction passed to torch._functorch
##                                  activation_memory_budget. Requires
##                                  compile=True. See
##                                  https://pytorch.org/blog/activation-checkpointing-techniques
##   fuse_optim_with_backward
##                          (bool)  Apply optimizer updates inside  Default: False
##                                  the backward grad hook so each
##                                  gradient is consumed and freed
##                                  immediately. Supported by the
##                                  basic and pipeline trainers.
##                                  Incompatible with DDP/Accelerate
##                                  (they need gradients intact for
##                                  the all-reduce step), with
##                                  max_grad_norm gradient clipping
##                                  (clipping needs the full gradient
##                                  vector before any update),
##                                  gradient_accumulation_steps > 1,
##                                  and fp16 AMP.
##   construct_model_on     (str)   Where to build the model.       Default: default
##                                  Choices: default, device, meta.
##                                    default: build on CPU, move to device.
##                                    device:  build directly on device.
##                                    meta:    build on meta then materialize empty
##                                             on device. If no checkpoint is
##                                             available the basic trainer falls
##                                             back to CPU construction
##                                             automatically. Ignored by the
##                                             pipeline trainer, which always
##                                             constructs on meta (with rank-0
##                                             doing a one-time CPU
##                                             initialisation + dispatch when no
##                                             checkpoint is present).
##   gc_threshold           (float) Ratio of allocated-to-total GPU Default: 0.9
##                                  memory at which to run empty_cache
##                                  and (conditionally) gc.collect.
##                                  Real trade-off: too low and frequent
##                                  cache flushes slow training (each
##                                  flush synchronises the CUDA stream);
##                                  too high and you OOM from allocator
##                                  fragmentation. Tune against the
##                                  OOM edge.
##
## Misc
##   seed                   (int)   Random seed.                    Default: 42
##   resume                 (bool)  Resume from latest checkpoint.  Default: True
##   save_strategy          (str)   Checkpoint save strategy.       Default: steps
##
## See docs/trainers/trainer_options.md for the complete list of options
## supported by each trainer and full descriptions of their semantics.
##
## ---------------------------------------------------------------------------
## Batch-Size / Learning-Rate Scaling: Background
##
## The transition between scaling regimes is governed by the critical batch
## size B_crit, which can be estimated from the gradient noise scale:
##
##   B_simple ~ tr(Sigma) / ||G||^2
##
## where G is the mean gradient and Sigma is the per-example gradient
## covariance. Practically, B_simple can be estimated cheaply by computing
## gradient norms at two batch sizes (e.g., B=1 and B=64) and taking their
## ratio -- see McCandlish et al. (2018) Appendix A for the exact estimator.
##
## When batch_size << B_crit (signal-dominated), each sample contributes
## fresh gradient signal, so linear scaling (alpha=1.0) is appropriate.
## When batch_size >> B_crit (noise-dominated), additional samples mainly
## reduce gradient noise, so sqrt scaling (alpha=0.5) is appropriate.
##
## Caveats:
##   - B_crit increases over training as "easy" gradient signal is exhausted,
##     suggesting batch size warmup as a practical strategy.
##   - The noise scale proxy relies on SGD + well-conditioned gradient
##     assumptions; recent work (Mayberry et al., 2025) finds it can
##     underestimate B_crit by orders of magnitude for Adam on LLMs, so
##     treat it as a rough lower bound.
##
## References:
##   McCandlish et al. (2018) "An Empirical Model of Large-Batch Training"
##     https://arxiv.org/abs/1812.06162
##   Mayberry et al. (2025) "Critical Batch Size Revisited"
##     https://arxiv.org/abs/2505.23971
##
## ---------------------------------------------------------------------------
## Chinchilla-Optimal Training: Background
##
## The "Chinchilla scaling law" (Hoffmann et al., 2022) empirically determined
## that prior large language models were significantly *undertrained* relative
## to their parameter count. The key finding: for a compute-optimal model,
## training tokens should scale *linearly* with model parameters -- roughly
## 20 tokens per parameter.
##
##   Optimal tokens ~ 20 x N
##   where N = number of model parameters
##
## This overturned the earlier Kaplan et al. (2020) wisdom, which favoured
## scaling parameters over data given a fixed compute budget (FLOPs).
## Hoffmann et al. showed that halving model size while doubling tokens yields
## equal or better loss at the same compute cost -- as demonstrated by
## Chinchilla (70B params, 1.4T tokens) outperforming Gopher (280B params,
## 300B tokens) despite using ~4x fewer parameters.
##
## Practical derivation:
##   C ~ 6 x N x D          (approximate FLOPs for a transformer fwd+bwd pass)
##   Minimise L(N, D) s.t. C = 6ND  =>  N_opt ~ C^0.5,  D_opt ~ C^0.5
##   Fitted constant from Hoffmann et al.: D_opt ~ 20 x N_opt
##
## Caveats:
##   - The "20x" ratio is a fitted constant; Epoch AI (2024) and others
##     suggest the true optimum may be higher (~40-100x) once inference costs
##     are factored in.
##   - Assumes a fixed compute budget; if data is the bottleneck, the ratio
##     shifts.
##   - Applies to autoregressive transformer LLMs; may not generalise to other
##     architectures or modalities.
##
## References:
##   Hoffmann et al. (2022) "Training Compute-Optimal Large Language Models"
##     https://arxiv.org/abs/2203.15556
##   Kaplan et al. (2020) "Scaling Laws for Neural Language Models"
##     https://arxiv.org/abs/2001.08361
##   Besiroglu et al. / Epoch AI (2024) "Chinchilla Scaling: A Replication Attempt"
##     https://arxiv.org/abs/2404.10102
## ---------------------------------------------------------------------------
##
## ===== Metadata and parameter defaults ======================================
##
[config_metadata]
    == super()
    -- set ns.config_name = "LM Training Project"
    -- set ns.config_description = "Default LM Training Project"
    -- set ns.nproc_per_node = nproc_per_node | default("gpu")

    ## Dataset
    -- set ns.dataset_proj = abspath(joinpath(ns.forgather_dir, "examples", "datasets", "HuggingFaceTB"))
    -- set ns.dataset_config = dataset_config | default("smollm-corpus/fineweb-edu-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("small.yaml")

    ## Training budget (millions of tokens)
    -- set ns.total_tokens = total_tokens | default(560)
    -- set ns.warmup_tokens = warmup_tokens | default(56)
    -- set ns.min_cooldown_tokens = min_cooldown_tokens | default(300)

    ## Batching
    -- set ns.seq_len = seq_len | default(512)
    -- set ns.per_device_train_batch_size = batch_size | default(32)
    -- set ns.gradient_accumulation_steps = gradient_accumulation_steps | default(1)
    -- set ns.batch_density = batch_density | default(0.95)

    ## Step cadence
    -- set ns.step_cadence = step_cadence | default(1.0)
    -- set ns.base_logging_tokens = base_logging_tokens | default(1)
    -- set ns.base_validation_tokens = base_validation_tokens | default(25)
    -- set ns.base_save_tokens = base_save_tokens | default(500)

    ## Hardware
    -- set ns.peak_hardware_flops = peak_hardware_flops | default(get_peak_hardware_flops())

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

    ## LR scaling
    -- set ns.base_lr = lr | default(3e-4)
    -- set ns.min_lr_ratio = min_lr_ratio | default(0.1)
    -- set ns.base_batch_size = base_batch_size | default(16384)
    -- set ns.lr_alpha = lr_alpha | default(0.5)

    ## Trainer selection
    -- set ns.trainer_type = trainer_type | default("basic")

    ## Pipeline Parallel inputs (only used when trainer_type == "pipeline")
    -- set ns.pipe_schedule_factory = pipeline_schedule | default("ScheduleInterleaved1F1B")
    -- set ns.microbatch_scale = microbatch_scale | default(1)

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

## ===== Derived globals ======================================================

[globals]
    == super()
    ## Map trainer_type to trainer template and launch config
    -- if ns.trainer_type == "ddp"
        -- set ns.trainer_class = 'trainers/ddp_trainer.yaml'
    -- elif ns.trainer_type == "fsdp2"
        -- set ns.trainer_class = 'trainers/fsdp2_trainer.yaml'
    -- elif ns.trainer_type == "pipeline"
        -- set ns.trainer_class = 'trainers/pipeline_trainer.yaml'
        -- set ns.dispatch_batches = True
    -- else
        -- set ns.trainer_class = 'trainers/trainer.yaml'
        -- set ns.nproc_per_node = 1
    -- endif

    ## Convert M-token inputs to raw token counts
    -- set ns.total_tokens = ns.total_tokens * 1000 * 1000
    -- set ns.warmup_tokens = ns.warmup_tokens * 1000 * 1000
    -- set ns.min_cooldown_tokens = ns.min_cooldown_tokens * 1000 * 1000

    ## Token intervals for logging / eval / save
    -- set ns.logging_tokens = ns.base_logging_tokens * 1000 * 1000 * ns.step_cadence
    -- set ns.validation_tokens = ns.base_validation_tokens * 1000 * 1000 * ns.step_cadence
    -- set ns.save_tokens = ns.base_save_tokens * 1000 * 1000 * ns.step_cadence

    ## Batch size and step calculations
    -- set ns.world_size = (getenv("WORLD_SIZE", "1") | int)
    -- set ns.effective_per_device_batch_size = ns.gradient_accumulation_steps * ns.per_device_train_batch_size
    -- set ns.global_batch_size = ns.effective_per_device_batch_size * ns.world_size
    -- set ns.tokens_per_step = (ns.seq_len * ns.global_batch_size * ns.batch_density) | int
    -- set ns.min_cooldown_steps = ns.min_cooldown_tokens // ns.tokens_per_step
    -- set ns.warmup_steps = ns.warmup_tokens // ns.tokens_per_step
    -- set ns.total_steps = ns.total_tokens // ns.tokens_per_step
    -- set ns.total_peak_hardware_flops = ns.world_size * ns.peak_hardware_flops

    ## LR Scaling
    -- set ns.global_lr = ns.base_lr * (ns.tokens_per_step / ns.base_batch_size) ** ns.lr_alpha
    -- set ns.min_lr = ns.global_lr * ns.min_lr_ratio

    ## Compile
    -- if ns.torch_compile == True and ns.gradient_accumulation_steps != 1 and ns.torch_compile_mode == "max-autotune"
        == raise("max-autotune is incompatible with gradient accumulation; try 'max-autotune-no-cudagraphs'")
    -- endif

    ## Pipeline Parallel derived values
    -- if ns.trainer_type == "pipeline"
        -- if ns.torch_compile == True
            -- if ns.pipe_schedule_factory in ["ScheduleInterleavedZeroBubble", "ScheduleZBVZeroBubble"]
                == raise("Torch compile is incompatible with zero-bubble schedulers'")
            -- elif ns.torch_compile_mode == "max-autotune"
                == raise("max-autotune is incompatible with pipeline training; try 'max-autotune-no-cudagraphs'")
            -- endif
        -- endif
        -- if ns.pipe_schedule_factory in ["ScheduleGPipe", "Schedule1F1B"]
            -- set ns.stages_per_rank = 1
        -- else
            -- set ns.stages_per_rank = 2
        -- endif
        -- if (ns.per_device_train_batch_size % (ns.stages_per_rank * ns.microbatch_scale)) != 0
            == raise("per_device_train_batch_size must be evenly divisible by stages_per_rank * microbatch_scale")
        -- endif
        -- set ns.per_stage_batch_size = ns.per_device_train_batch_size // (ns.stages_per_rank * ns.microbatch_scale)
        -- set ns.n_microbatches = ns.world_size * ns.stages_per_rank * ns.microbatch_scale
        -- set ns.pp_batch_size = ns.n_microbatches * ns.per_stage_batch_size
        -- if ns.pipe_schedule_factory == "ScheduleZBVZeroBubble"
            -- set ns.pp_stage_type = "v"
        -- else
            -- set ns.pp_stage_type = "loop"
        -- endif
    -- endif

[variable_listing]
    == super()
# **LM Training Project**
# ns.per_device_train_batch_size: {{ ns.per_device_train_batch_size }}
# ns.gradient_accumulation_steps: {{ ns.gradient_accumulation_steps }}
# ns.effective_per_device_batch_size: {{ ns.effective_per_device_batch_size }}
# ns.global_batch_size: {{ ns.global_batch_size }}
# ns.seq_len: {{ ns.seq_len }} tokens
# ns.total_steps: {{ ns.total_steps }} steps
# ns.warmup_steps: {{ ns.warmup_steps }}
# ns.min_cooldown_steps: {{ ns.min_cooldown_steps }}
# ns.total_tokens: {{ ns.total_tokens // (1000 * 1000)}}M
# ns.tokens_per_step: {{ ns.tokens_per_step }} tokens
# ns.total_peak_hardware_flops: {{ ns.total_peak_hardware_flops / 1e12 }} TFLOPS
# ns.base_lr: {{ ns.base_lr }}
# ns.min_lr: {{ ns.min_lr }}
# ns.base_batch_size: {{ ns.base_batch_size }}
# ns.lr_alpha: {{ ns.lr_alpha }}
# ns.global_lr: {{ ns.global_lr }}
# ns.trainer_type: {{ ns.trainer_type }}
# ns.trainer_class: {{ ns.trainer_class }}
-- if ns.trainer_type == "pipeline"
# Pipeline Parallel:
# ns.stages_per_rank: {{ ns.stages_per_rank }}
# ns.per_stage_batch_size: {{ ns.per_stage_batch_size }}
# ns.n_microbatches: {{ ns.n_microbatches }}
# ns.pp_batch_size: {{ ns.pp_batch_size }}
# ns.pp_stage_type: {{ ns.pp_stage_type }}
-- endif

## ===== Dataset ==============================================================
##
[datasets_preprocessor_args]
.define: &tokenizer_args !dict
    truncation: True
    max_length: {{ ns.seq_len }}

[datasets_definition]
    [dataset_project_pp_args]
.define: &dataset_project_pp_args !dict

    [dataset_project]
.define: &dataset_dict !call:forgather:from_project
    project_dir: "{{ ns.dataset_proj }}"
    config_template: "{{ ns.dataset_config }}"
    targets: [ "train_dataset", "eval_dataset" ]
    pp_kwargs: *dataset_project_pp_args
    preprocess_args: *tokenizer_args
    tokenizer: *tokenizer
    shard_dataset: {{ ns.dispatch_batches == False }}

    [dataset_splits]
train_dataset: &train_dataset !call:getitem [ *dataset_dict, 'train_dataset' ]
eval_dataset: &eval_dataset !call:getitem [ *dataset_dict, 'eval_dataset' ]

## ===== Model ================================================================
##
[model_constructor_args]
.define: &model_constructor_args
    attn_implementation: "{{ ns.attn_implementation }}"

[construct_new_model]
    [model_project_pp_args]
.define: &model_project_pp_args
    output_dir: "{{ ns.output_dir }}"

    [model_project]
.define: &model_dict !call:forgather:from_project
    project_dir: "{{ ns.model_project_dir }}"
    config_template: "{{ ns.model_project_config }}"
    targets: [ "pretrained_tokenizer", "model" ]
    pp_kwargs: *model_project_pp_args
    model_constructor_args: *model_constructor_args

    [model_assets]
tokenizer: &tokenizer !call:getitem [ *model_dict, 'pretrained_tokenizer' ]
model: &model !call:getitem [ *model_dict, 'model' ]

## ===== Trainer ==============================================================
##
[trainer_args]
    == super()
    seed: {{ seed | toyaml(42) }}
    eval_strategy: steps
    save_strategy: {{ save_strategy | toyaml('steps') }}

    logging_steps: {{ (ns.logging_tokens // ns.tokens_per_step) | int }}
    eval_steps: {{ (ns.validation_tokens // ns.tokens_per_step) | int }}
    save_steps: {{ (ns.save_tokens // ns.tokens_per_step) | int }}
    max_steps: {{ max_steps | toyaml(ns.total_steps) }}

    per_device_train_batch_size: {{ ns.per_device_train_batch_size }}
    per_device_eval_batch_size: {{ ns.per_device_train_batch_size * 2 }}  # 2x train batch size
    gradient_accumulation_steps: {{ ns.gradient_accumulation_steps }}

    resume_from_checkpoint: {{ resume | default(True) }}
    torch_compile: {{ ns.torch_compile }}
    torch_compile_mode: {{ ns.torch_compile_mode }}
    dispatch_batches: {{ ns.dispatch_batches }}
    float32_matmul_precision: {{ float32_matmul_precision | toyaml(None) }}
    default_dtype: {{ default_dtype | toyaml(None) }}
    mixed_precision: {{ mixed_precision | toyaml(None) }}
    fp8_recipe: {{ fp8_recipe | toyaml(None) }}
    qat_recipe: {{ qat_recipe | toyaml(None) }}
    gradient_checkpointing: {{ gradient_checkpointing | default(False) }}
    fuse_optim_with_backward: {{ fuse_optim_with_backward | default(False) }}
    enable_activation_offloading: {{ activation_offloading | default(False) }}
    dataloader_drop_last: True
    construct_model_on: {{ construct_model_on | toyaml("default") }}
    activation_memory_budget: {{ activation_memory_budget | toyaml(None) }}
    gc_threshold: {{ gc_threshold | default(0.9) }}
    save_safetensors: {{ save_safetensors | default(False) }}
    debug_optimizer_groups: {{ debug_optimizer_groups | default(False) }}
    max_grad_norm: {{ max_grad_norm | toyaml(None) }}

-- if ns.trainer_type == "pipeline"
    # Pipeline Parallel
    n_microbatches: {{ ns.n_microbatches }}
    is_multistage: {{ (ns.stages_per_rank > 1) | toyaml }}
    stages_per_rank: {{ ns.stages_per_rank }}
    pp_stage_type: {{ ns.pp_stage_type | toyaml }}
    per_device_train_batch_size: {{ ns.pp_batch_size }}
    per_device_eval_batch_size: {{ ns.pp_batch_size }}
    dispatch_batches: null
-- endif

[trainer_definition]
    -- include ns.trainer_class

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

[lr_scheduler]
# LR Scheduler (cosine decay with warmup)
# total_steps is clamped to min_cooldown_steps ({{ ns.min_cooldown_steps }})
# so that short runs do not decay the LR all the way to zero.
lr_scheduler: &lr_scheduler !partial:forgather.ml.optim:CosineLRScheduler@lr_scheduler
    warmup_steps: {{ ns.warmup_steps }}
    total_steps: {{ [ns.min_cooldown_steps, ns.total_steps] | max | int }}
    min_lr: {{ ns.min_lr | toyaml }}

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

[optimizer_factories]
    ## Block for extending
    [optimizer_factories_]
    [no_decay_optim]
.define: &no_decay_optim ~

[optimizer_groups]
optimizer_groups: &optimizer_groups
    no_decay:
        regex: 'norm|bias|embed|lm_head|alibi_slopes'
        config:
            weight_decay: 0.0
        factory: *no_decay_optim

## ===== CLI arguments ========================================================
##
[dynamic_args]
    == super()
    model_project:
        names: "--model-project"
        type: path
        group: "Model"
        help: "Path to model project for model initialization"
    model_config:
        names: "--model-config"
        group: "Model"
        help: "Model project configuration for model init"
    dataset_proj:
        names: "--dataset-project"
        type: path
        group: "Data"
        help: "Path to dataset project"
    dataset_config:
        names: "--dataset-config"
        group: "Data"
        help: "Dataset project configuration"
    seq_len:
        names: "--seq-len"
        type: int
        group: "Data"
        min: 1
        help: "Set maximum sequence length"
    batch_density:
        names: "--batch-density"
        type: float
        group: "Data"
        min: 0
        max: 1
        help: "Estimated fraction of non-pad tokens per batch; used to correct token-count estimates"
    resume:
        names: "--resume"
        type: bool
        group: "Trainer:Checkpointing"
        help: "Resume from checkpoint. Defaults to 'auto'"
    save_safetensors:
        names: "--save-safetensors"
        type: bool
        group: "Trainer:Checkpointing"
        help: "Save checkpoint using safe-tensors format"
    batch_size:
        names: "--batch-size"
        type: int
        group: "Trainer:Batch"
        min: 1
        help: "Set the per-device training batch size"
    gradient_accumulation_steps:
        names: "--gradient-accumulation-steps"
        type: int
        group: "Trainer:Batch"
        min: 1
        help: "Set the number of gradient accumulation steps"
    total_tokens:
        names: "--total-tokens"
        type: int
        group: "Trainer:Schedule"
        min: 1
        help: "Total training tokens (millions)"
    warmup_tokens:
        names: "--warmup-tokens"
        type: int
        group: "Trainer:Schedule"
        min: 0
        help: "Total warmup tokens (millions)"
    min_cooldown_tokens:
        names: "--min-cooldown-tokens"
        type: int
        group: "Trainer:Schedule"
        min: 0
        help: "Minimum tokens (millions) for LR decay phase"
    lr:
        names: "--lr"
        type: float
        group: "Trainer:LR-scaling"
        min: 0
        help: "Learning rate"
    min_lr_ratio:
        names: "--min-lr-ratio"
        type: float
        group: "Trainer:LR-scaling"
        min: 0
        max: 1
        help: "The minimum LR ratio, relative to the base-lr"
    base_batch_size:
        names: "--base-batch-size"
        type: int
        group: "Trainer:LR-scaling"
        min: 1
        help: "Reference batch size (tokens) for LR scaling calculation"
    lr_alpha:
        names: "--lr_alpha"
        type: float
        group: "Trainer:LR-scaling"
        min: 0
        max: 1
        help: "LR scaling exponent: 0.0 = no scaling, 0.5 = sqrt, 1.0 = linear"
    max_grad_norm:
        names: "--max-grad-norm"
        type: float
        group: "Trainer"
        min: 0
        help: "Maximum gradient norm for gradient clipping"
    seed:
        names: "--seed"
        type: int
        group: "Trainer"
        min: 0
        help: "Random seed"
    trainer_type:
        names: "--trainer-type"
        choices: [ "basic", "ddp", "fsdp2", "pipeline" ]
        default: "basic"
        group: "Trainer"
        help: "Trainer: basic (single GPU), ddp (multi-GPU), fsdp2 (sharded DP), pipeline (pipeline parallel)"
    dispatch_batches:
        names: "--dispatch-batches"
        type: bool
        group: "Trainer:Distributed"
        help: "Dispatch all batches from rank-0 instead of sharding (DDP)"
    nproc_per_node:
        names: "--nproc-per-node"
        type: int
        group: "Trainer:Distributed"
        min: 1
        help: "Processes per node: 'gpu' for all GPUs, or integer."
    pipeline_schedule:
        names: "--pipeline-schedule"
        default: "ScheduleInterleaved1F1B"
        group: "Trainer:Distributed"
        choices:
            - "ScheduleGPipe"
            - "Schedule1F1B"
            - "ScheduleInterleaved1F1B"
            - "ScheduleLoopedBFS"
            - "ScheduleInterleavedZeroBubble"
            - "ScheduleZBVZeroBubble"
        help: "Pipeline schedule class (pipeline trainer only)"
    microbatch_scale:
        names: "--microbatch-scale"
        default: 1
        type: int
        group: "Trainer:Distributed"
        min: 1
        help: "Microbatch scale factor (pipeline trainer only)"
    default_dtype:
        names: "--default-dtype"
        choices: [ "float32", "bfloat16", "float16" ]
        group: "Precision"
        help: "Default torch dtype for model construction"
    float32_matmul_precision:
        names: "--float32-matmul-precision"
        choices: [ "highest", "high", "medium" ]
        group: "Precision"
        help: "Approximate float32 matmul with bf16; highest = full float32"
    mixed_precision:
        names: "--mixed-precision"
        choices: [ "bf16", "fp16", "no" ]
        group: "Precision"
        help: "Automatic Mixed Precision (AMP) dtype"
    fp8_recipe:
        names: "--fp8-recipe"
        choices: [ "tensorwise", "rowwise", "rowwise_with_gw_hp" ]
        group: "Precision"
        help: "FP8 training recipe (torchao Float8Linear)"
    qat_recipe:
        names: "--qat-recipe"
        choices: {{ qat_recipes() | toyaml }}
        group: "Precision"
        help: "Quantization-aware training recipe (torchao QAT prepare; run `forgather finalize --quantize <recipe>` after training)"
    compile:
        names: "--compile"
        type: bool
        group: "Performance"
        help: "Enable torch.compile"
    torch_compile_mode:
        names: "--torch-compile-mode"
        choices: [ "default", "max-autotune", "max-autotune-no-cudagraphs" ]
        group: "Performance"
        help: "Torch compile mode argument."
    fuse_optim_with_backward:
        names: "--fuse-optim-with-backward"
        type: bool
        group: "Performance"
        help: "Combine gradient computation with optimizer step"
    construct_model_on:
        names: "--construct-model-on"
        choices: [ "default", "meta", "device" ]
        group: "Performance"
        help: "Where to construct the model?"
    gradient_checkpointing:
        names: [ "--gradient-checkpointing", "-G" ]
        type: bool
        group: "Memory"
        help: "Enable gradient (activation) checkpointing, when supported"
    activation_offloading:
        names: "--activation-offloading"
        type: bool
        group: "Memory"
        help: "Save activations to CPU memory"
    activation_memory_budget:
        names: "--activation-memory-budget"
        type: float
        group: "Memory"
        min: 0
        max: 1
        help: "https://pytorch.org/blog/activation-checkpointing-techniques"
    gc_threshold:
        names: "--gc-threshold"
        type: float
        group: "Memory"
        min: 0
        max: 1
        help: "Ratio of reserved to total GPU memory to trigger GC"
    step_cadence:
        names: "--step-cadence"
        type: float
        group: "Logging"
        min: 0
        help: "Scale log/eval/save intervals by this factor"
    base_logging_tokens:
        names: "--base-logging-tokens"
        type: float
        group: "Logging"
        min: 0
        help: "Tokens (M) between log steps"
    base_validation_tokens:
        names: "--base-validation-tokens"
        type: float
        group: "Logging"
        min: 0
        help: "Tokens (M) between eval steps"
    base_save_tokens:
        names: "--base-save-tokens"
        type: float
        group: "Logging"
        min: 0
        help: "Tokens (M) between save steps"
    peak_hardware_flops:
        names: "--peak-hardware-flops"
        type: float
        group: "Logging"
        min: 0
        help: "Peak device FLOPS for MFU computation"
    debug_optimizer_groups:
        names: "--debug-optimizer-groups"
        action: "store_true"
        group: "Logging"
        help: "Show how parameters are mapped to optimizer param groups"
