# Wan Fun InP robotics training on SageMaker (1.3B or 14B).
# Hyperparams mirror the any4d_rwm3_robot41d (Cosmos 2B) reference command so
# training behavior carries over directly; the only size-specific bits are the
# experiment file and the 14B-only activation checkpointing.
#
# Prerequisites:
#   1. Wan checkpoint in S3 (one-time, already done for both sizes):
#        s3://tri-ml-sandbox-16011-us-west-2-datasets/wan/Wan2.1-Fun-1.3B-InP/
#        s3://tri-ml-sandbox-16011-us-west-2-datasets/wan/Wan2.2-Fun-A14B-InP/
#
#   2. Set env vars: SM_USER, TRI_PROJECT, TRI_OWNER_EMAIL, WANDB_API_KEY.
#
# Usage:
#   bash custom/sagemaker/run_wan.sh [SIZE] [INSTANCE_COUNT] [NAME] [EXTRA_OVERRIDES...]
#     SIZE: 1_3b | 14b  (default: 14b)
#
# Example:
#   bash custom/sagemaker/run_wan.sh 1_3b 2 swati_rwm3_robot33_wan1_3b
#   bash custom/sagemaker/run_wan.sh 14b 2 swati_rwm3_robot33_wan14b dataloader_train.batch_size=2

SIZE=${1:-14b} # default to 14B model
INSTANCE_COUNT=${2:-2}

case "${SIZE}" in
    1_3b)
        EXPERIMENT=any4d_rwm3_robot_wan
        ACT_CKPT_OVERRIDES=()  # 1.3B fits without activation checkpointing
        ;;
    14b)
        EXPERIMENT=any4d_rwm3_robot41d_wan14b
        ACT_CKPT_OVERRIDES=(++model.config.wan_activation_checkpoint=True)
        ;;
    *)
        echo "Unknown SIZE '${SIZE}' (expected '1_3b' or '14b')" >&2
        exit 1
        ;;
esac

NAME=${3:-swati_rwm3_robot33_bs2_ga2_wan${SIZE}}
shift 3 2>/dev/null || shift $#
EXTRA_OVERRIDES=("$@")

QUEUE=cv-wfm
BUILD_TYPE=full
VERSION=271-2stage
PRIORITY=100

# User/environment-specific paths. Override via env var so the script is reusable without
# editing; defaults preserve the original behavior.
JOB_GROUP=${SM_JOB_GROUP:-sm_swati}
LOCAL_ROOT=${SM_LOCAL_ROOT:-/home/swatigupta/Any4D/logs/a4d2/}
S3_ROOT=${SM_S3_ROOT:-s3://tri-ml-sandbox-16011-us-west-2-datasets/any4d/swati}

# FSDP sharding value can be at max 8 (its defined as GPUs per node)
bash custom/sagemaker/run_sm.sh                                                          \
    ${INSTANCE_COUNT} ${EXPERIMENT} ${NAME} ${QUEUE} ${BUILD_TYPE} ${VERSION} ${PRIORITY} \
    "job.group=${JOB_GROUP}"                                                             \
    "job.local_root=${LOCAL_ROOT}"                                                       \
    "job.s3_root=${S3_ROOT}"                                                             \
    dataloader_train.batch_size=2                                                        \
    trainer.grad_accum_iter=2                                                            \
    dataloader_train.num_workers=2                                                       \
    model.config.fsdp_shard_size=8                                        \
    model.config.collate_quiet=True                                                      \
    "${ACT_CKPT_OVERRIDES[@]}"                                                            \
    trainer.validation_iter=1000                                                         \
    trainer.skip_first_validation=delay                                                  \
    "++model.config.data_train_overrides.resize=[-16,448]"                               \
    "++model.config.data_train_overrides.length=33"                                      \
    "++model.config.data_val_overrides.resize=[-16,448]"                                 \
    "++model.config.data_val_overrides.length=33"                                        \
    "++model.config.data_val_overrides.subsample=64"                                     \
    "${EXTRA_OVERRIDES[@]}"