#!/usr/bin/env bash
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

# Step 2: Run SFT training with LlamaFactory.
#
# Required environment variables:
#   CONFIG_YAML - Name of the SFT config file in configs/sft/ (e.g. qwen3-4b-base-sft-qwen3-8b.yaml)
#   OUTPUT_DIR  - Directory for the SFT checkpoint output
#
# Optional:
#   NUM_NODES   - Number of nodes (default: 4)
#   NUM_GPUS    - GPUs per node (default: 8)
#   MASTER_ADDR - Master node address (default: localhost)
#
# Prerequisites:
#   - LlamaFactory installed (pip install llamafactory)
#   - SFT data generated by Step 1 and registered in LlamaFactory's dataset_info.json

set -euo pipefail

: "${CONFIG_YAML:?Set CONFIG_YAML (e.g. qwen3-4b-base-sft-qwen3-8b.yaml)}"
: "${OUTPUT_DIR:?Set OUTPUT_DIR for SFT checkpoint output}"

NUM_NODES="${NUM_NODES:-4}"
NUM_GPUS="${NUM_GPUS:-8}"
MASTER_ADDR="${MASTER_ADDR:-localhost}"
MASTER_PORT="${MASTER_PORT:-29500}"

# torchrun \
#     --nnodes "${NUM_NODES}" \
#     --nproc_per_node="${NUM_GPUS}" \
#     --rdzv_id $RANDOM \
#     --rdzv_backend c10d \
#     --rdzv_endpoint "${MASTER_ADDR}:29500" \
#     -m llamafactory.cli.train \
#     "configs/sft/${CONFIG_YAML}" \
#     "dataset_dir=configs/sft" \
#     "output_dir=${OUTPUT_DIR}"

FORCE_TORCHRUN=1 \
NNODES="${NUM_NODES}" \
NPROC_PER_NODE="${NUM_GPUS}" \
MASTER_ADDR="${MASTER_ADDR}" \
MASTER_PORT="${MASTER_PORT}" \
llamafactory-cli train \
    "configs/sft/${CONFIG_YAML}" \
    "dataset_dir=configs/sft" \
    "output_dir=${OUTPUT_DIR}"