# 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/llama_embed_nemotron_8b/llama_embed_nemotron_8b.yaml --nproc-per-node 8
# Adjust --nproc-per-node to the number of GPUs available on your machine.

recipe: TrainBiEncoderRecipe

seed: 125

temperature: 0.02

step_scheduler:
  global_batch_size: 128
  local_batch_size: 1
  ckpt_every_steps: 3000
  val_every_steps: 3000
  num_epochs: 1

model:
  _target_: nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.1-8B
  pooling: avg
  torch_dtype: bfloat16

tokenizer:
  _target_: transformers.AutoTokenizer.from_pretrained
  pretrained_model_name_or_path: meta-llama/Llama-3.1-8B

dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  dataset:
    _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
    model_type: bi_encoder
    data_dir_list:
      - ./embed_nemotron_dataset_v1/EmotionClassification/EmotionClassification.json
      - ./embed_nemotron_dataset_v1/FEVER/FEVER.json
      - ./embed_nemotron_dataset_v1/GooAQ/GooAQ.json
      - ./embed_nemotron_dataset_v1/HotpotQA/HotpotQA.json
      - ./embed_nemotron_dataset_v1/MAmmoTH2/MAmmoTH2.json
      - ./embed_nemotron_dataset_v1/MIRACL/MIRACL.json
      - ./embed_nemotron_dataset_v1/MSMARCO/MSMARCO.json
      - ./embed_nemotron_dataset_v1/NaturalQuestions/NaturalQuestions.json
      - ./embed_nemotron_dataset_v1/NFCorpus/NFCorpus.json
      - ./embed_nemotron_dataset_v1/PAQ/PAQ.json
      - ./embed_nemotron_dataset_v1/SciFact/SciFact.json
      - ./embed_nemotron_dataset_v1/SQuAD/SQuAD.json
      - ./embed_nemotron_dataset_v1/SyntheticClassificationData/SyntheticClassificationData.json
      - ./embed_nemotron_dataset_v1/TriviaQA/TriviaQA.json
    data_type: train
    n_passages: 5
    seed: 125
    do_shuffle: true
    use_dataset_instruction: true
  collate_fn:
    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
    q_max_len: 512
    p_max_len: 512
    query_prefix: ""
    passage_prefix: ""
    pad_to_multiple_of: 8
    use_dataset_instruction: true
  shuffle: true
  num_workers: 8

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

lr_scheduler:
  lr_warmup_steps: 2

checkpoint:
  enabled: true
  checkpoint_dir: ./output/llama_embed_nemotron_8b
  model_save_format: safetensors
  save_consolidated: true

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

  sequence_parallel: false

dist_env:
  backend: nccl
  timeout_minutes: 30
