# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# To run this recipe:
#   automodel examples/retrieval/bi_encoder/llama3_2_1b.yaml --nproc-per-node 8
# Adjust --nproc-per-node to the number of GPUs available on your machine.

recipe: TrainBiEncoderRecipe

seed: 42

temperature: 0.02

step_scheduler:
  global_batch_size: 128
  local_batch_size: 4
  ckpt_every_steps: 500
  val_every_steps: 500
  num_epochs: 1

dist_env:
  backend: nccl
  timeout_minutes: 1

model:
  _target_: nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
  pooling: avg
  l2_normalize: true
  use_liger_kernel: true
  use_sdpa_patching: true
  torch_dtype: bfloat16

tokenizer:
  _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
  add_eos_token: false

dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  dataset:
    _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
    model_type: bi_encoder
    data_dir_list:
      - hf://nvidia/embed-nemotron-dataset-v1/FEVER
      - hf://nvidia/embed-nemotron-dataset-v1/SyntheticClassificationData
    data_type: train
    n_passages: 5
    seed: 42
    do_shuffle: true
  collate_fn:
    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
    q_max_len: 512
    p_max_len: 512
    query_prefix: "query:"
    passage_prefix: "passage:"
    pad_to_multiple_of: 8
  shuffle: true
  num_workers: 0

# Optional: Uncomment to enable validation
# validation_dataloader:
#   _target_: torchdata.stateful_dataloader.StatefulDataLoader
#   dataset:
#     _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
#     model_type: bi_encoder
#     data_dir_list: training_datasets/validation.json
#     data_type: eval
#     n_passages: 5
#     seed: 42
#     do_shuffle: false
#     max_train_samples: 1000
#     train_data_select_offset: 0
#   collate_fn:
#     _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
#     q_max_len: 512
#     p_max_len: 512
#     query_prefix: "query:"
#     passage_prefix: "passage:"
#     padding: longest
#     pad_to_multiple_of: 8
#   batch_size: 2
#   shuffle: false
#   num_workers: 0

# LoRA/PEFT configuration
# Uncomment and configure for LoRA/PEFT
# peft:
#   _target_: nemo_automodel.components._peft.lora.PeftConfig
#   target_modules:
#     - q_proj
#     - k_proj
#     - v_proj
#     - o_proj
#     - gate_proj
#     - up_proj
#     - down_proj
#   exclude_modules: []
#   match_all_linear: True
#   dim: 16
#   alpha: 32
#   use_triton: true

optimizer:
  _target_: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam
  lr: 5.0e-6
  weight_decay: 0.01
  adam_w_mode: true
  bias_correction: true
  master_weights: true

lr_scheduler:
  lr_warmup_steps: 100
  lr_decay_style: linear

checkpoint:
  enabled: true
  checkpoint_dir: ./output/llama3_2_1b_encoder/checkpoints
  model_save_format: safetensors
  save_consolidated: true

distributed:
  strategy: fsdp2
  dp_size: none
  tp_size: 1
  cp_size: 1

  sequence_parallel: false

# Uncomment and configure for W&B logging
# wandb:
#   project: <your_wandb_project>
#   entity: <your_wandb_entity>
#   name: <your_wandb_exp_name>
#   save_dir: <your_wandb_save_dir>

ci:
  recipe_owner: oliverholworthy
  time: "01:30:00"
  nodes: 1
  embed_eval: true
