#!/bin/bash
#SBATCH --job-name=imgnt-movq
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
#SBATCH --cpus-per-task=96
#SBATCH --gres=gpu:8
#SBATCH --exclusive
#SBATCH --partition=production-cluster
#SBATCH --output=/admin/home/suraj/logs/maskgit-imagenet/%x-%j.out

set -x -e

source /admin/home/suraj/.bashrc
source /fsx/suraj/miniconda3/etc/profile.d/conda.sh
conda activate muse

echo "START TIME: $(date)"

MUSE_REPO=/admin/home/suraj/code/muse
OUTPUT_DIR=/fsx/suraj/imgnt-movq
LOG_PATH=$OUTPUT_DIR/main_log.txt
ACCELERATE_CONFIG_FILE="$OUTPUT_DIR/${SLURM_JOB_ID}_accelerate_config.yaml.autogenerated"
DEEPSPEED_CONFIG_FILE="$OUTPUT_DIR/${SLURM_JOB_ID}_ds_config.json.autogenerated"

mkdir -p $OUTPUT_DIR
touch $LOG_PATH
pushd $MUSE_REPO


GPUS_PER_NODE=8
NNODES=$SLURM_NNODES
NUM_GPUS=$((NUM_GPUS_PER_NODE*SLURM_NNODES))

# so processes know who to talk to
MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
MASTER_PORT=6000


# Auto-generate the accelerate config
cat << EOT > $ACCELERATE_CONFIG_FILE
compute_environment: LOCAL_MACHINE
deepspeed_config:
  deepspeed_multinode_launcher: standard
  deepspeed_config_file: $DEEPSPEED_CONFIG_FILE
  zero3_init_flag: false
distributed_type: DEEPSPEED
fsdp_config: {}
machine_rank: 0
main_process_ip: $MASTER_ADDR
main_process_port: $MASTER_PORT
main_training_function: main
num_machines: $SLURM_NNODES
num_processes: $NUM_GPUS
use_cpu: false
EOT

# Auto-generate the deepspeed config
cat << EOT > $DEEPSPEED_CONFIG_FILE
{
    "communication_data_type": "fp32",
    "bf16": {
        "enabled": true
    },
    "fp16": {
        "enabled": false,
        "auto_cast": true,
        "loss_scale": 0.0,
        "initial_scale_power": 32,
        "loss_scale_window": 1000,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 2,
        "allgather_partitions": true,
        "allgather_bucket_size": 5e8,
        "overlap_comm": false,
        "reduce_scatter": true,
        "reduce_bucket_size": "auto",
        "contiguous_gradients": true,
        "offload_optimizer": {
            "device": "none"
        },
        "offload_param": {
            "device": "none"
        }
    },
    "gradient_accumulation_steps": 1,
    "gradient_clipping": "auto",
    "train_batch_size": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "steps_per_print": 2000
}
EOT


PROGRAM=" \
    training/train_maskgit_imagenet.py config=configs/imagenet_movq.yaml \
    wandb.entity=psuraj \
    experiment.name=$(basename $OUTPUT_DIR) \
    experiment.output_dir=$OUTPUT_DIR \
    training.seed=343676232 \
    optimizer.params.learning_rate=6e-4 \
    dataset.params.num_workers=2 \
    model.vq_model.pretrained=openMUSE/movq-lion-high-res-f8-16384 \
    training.batch_size=48  \
    "

# Note: it is important to escape `$SLURM_PROCID` since we want the srun on each node to evaluate this variable
export LAUNCHER="accelerate launch \
    --rdzv_conf "rdzv_backend=c10d,rdzv_endpoint=$MASTER_ADDR:$MASTER_PORT,max_restarts=0,tee=3" \
    --config_file $ACCELERATE_CONFIG_FILE \
    --main_process_ip $MASTER_ADDR \
    --main_process_port $MASTER_PORT \
    --machine_rank \$SLURM_PROCID \
    "


export CMD="$LAUNCHER $PROGRAM"
echo $CMD

# hide duplicated errors using this hack - will be properly fixed in pt-1.12
# export TORCHELASTIC_ERROR_FILE=/tmp/torch-elastic-error.json

# force crashing on nccl issues like hanging broadcast
export NCCL_ASYNC_ERROR_HANDLING=1
# export NCCL_DEBUG=INFO
# export NCCL_DEBUG_SUBSYS=COLL
# export NCCL_SOCKET_NTHREADS=1
# export NCCL_NSOCKS_PERTHREAD=1
# export CUDA_LAUNCH_BLOCKING=1

# AWS specific
export NCCL_PROTO=simple
export RDMAV_FORK_SAFE=1
export FI_EFA_FORK_SAFE=1
export FI_EFA_USE_DEVICE_RDMA=1
export FI_PROVIDER=efa
export FI_LOG_LEVEL=1
export NCCL_IB_DISABLE=1
export NCCL_SOCKET_IFNAME=ens


# srun error handling:
# --wait=60: wait 60 sec after the first task terminates before terminating all remaining tasks
# --kill-on-bad-exit=1: terminate a step if any task exits with a non-zero exit code
SRUN_ARGS=" \
    --wait=60 \
    --kill-on-bad-exit=1 \
    "

clear; srun $SRUN_ARGS --jobid $SLURM_JOB_ID bash -c "$CMD" 2>&1 | tee $LOG_PATH

echo "END TIME: $(date)"
