run_ddp_train.sh

#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "${ROOT_DIR}"

CUDA_DEVICES="${CUDA_DEVICES:-2}"
MASTER_PORT="${MASTER_PORT:-17777}"
SEED="${SEED:-143}"
PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
NCCL_ASYNC_ERROR_HANDLING="${NCCL_ASYNC_ERROR_HANDLING:-1}"
NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-1}"
NCCL_P2P_DISABLE="${NCCL_P2P_DISABLE:-0}"
VERIFY_MERGED_INIT_CKPT="${VERIFY_MERGED_INIT_CKPT:-1}"

CONFIG="${CONFIG:-configs/deimv2/deimv2_dinov3_s_crowd_human_mot_distilled.yml}"
OUTPUT_DIR="${OUTPUT_DIR:-outputs_dinov3_small_distilled_dinov3_tiny}"
COCO_INIT_CKPT="${COCO_INIT_CKPT:-${INIT_CKPT:-ckpts/deimv2_dinov3_s_coco.pth}}"
DISTILLED_BACKBONE_CKPT="${DISTILLED_BACKBONE_CKPT:-${BACKBONE_CKPT:-outputs/lightly_distill/deimv2_vitt_crowd_human_mot_full/exported_models/exported_last.pt}}"
MERGED_INIT_CKPT="${MERGED_INIT_CKPT:-${OUTPUT_DIR}/init_coco_with_distilled_backbone.pth}"
TRAIN_ANN="${TRAIN_ANN:-data/crowd_human_mot_coco/annotations/instances_train_id01_crowdhuman_nega_minus1000neg_le100.json}"
VAL_ANN="${VAL_ANN:-data/crowd_human_mot_coco/annotations/instances_val_id01_crowdhuman_nega_plus1000neg.json}"

IFS=',' read -r -a GPU_ARRAY <<< "${CUDA_DEVICES}"
NPROC_PER_NODE="${NPROC_PER_NODE:-${#GPU_ARRAY[@]}}"

if [[ "${NPROC_PER_NODE}" -ne "${#GPU_ARRAY[@]}" ]]; then
  echo "Warning: expected ${#GPU_ARRAY[@]} processes for GPUs ${CUDA_DEVICES}, got NPROC_PER_NODE=${NPROC_PER_NODE}" >&2
fi

for required in "${CONFIG}" "${TRAIN_ANN}" "${VAL_ANN}"; do
  if [[ ! -e "${required}" ]]; then
    echo "Missing required file: ${required}" >&2
    exit 1
  fi
done

mkdir -p "${OUTPUT_DIR}"
timestamp="$(date +%Y%m%d_%H%M%S)"
log_file="${OUTPUT_DIR}/train_${timestamp}.log"

build_merged_init_ckpt() {
  local coco_ckpt="$1"
  local distilled_backbone_ckpt="$2"
  local merged_ckpt="$3"

  mkdir -p "$(dirname "${merged_ckpt}")"
  python - "${coco_ckpt}" "${distilled_backbone_ckpt}" "${merged_ckpt}" <<'PY'
import sys
import torch

coco_ckpt, distilled_ckpt, merged_ckpt = sys.argv[1:4]

base_state = torch.load(coco_ckpt, map_location="cpu")
if isinstance(base_state, dict) and "model" in base_state and isinstance(base_state["model"], dict):
    base_model_state = base_state["model"]
    merged_state = dict(base_state)
elif isinstance(base_state, dict):
    base_model_state = base_state
    merged_state = {"model": base_model_state}
else:
    raise RuntimeError(f"Unsupported COCO checkpoint format: {type(base_state)}")

distilled_state = torch.load(distilled_ckpt, map_location="cpu")
if isinstance(distilled_state, dict) and "model" in distilled_state and isinstance(distilled_state["model"], dict):
    distilled_state = distilled_state["model"]
elif isinstance(distilled_state, dict) and "state_dict" in distilled_state and isinstance(distilled_state["state_dict"], dict):
    distilled_state = distilled_state["state_dict"]
elif not isinstance(distilled_state, dict):
    raise RuntimeError(f"Unsupported distilled checkpoint format: {type(distilled_state)}")

def normalize_source_key(key: str) -> str:
    if key.startswith("module."):
        key = key[len("module."):]
    for prefix in ("backbone.dinov3._model.", "dinov3._model.", "_model."):
        if key.startswith(prefix):
            return key[len(prefix):]
    return key

replaced = 0
missing = []
shape_mismatch = []
for src_key, src_tensor in distilled_state.items():
    if not torch.is_tensor(src_tensor):
        continue
    suffix = normalize_source_key(src_key)
    target_key = f"backbone.dinov3._model.{suffix}"
    if target_key not in base_model_state:
        missing.append((src_key, target_key))
        continue
    if base_model_state[target_key].shape != src_tensor.shape:
        shape_mismatch.append(
            (src_key, tuple(src_tensor.shape), target_key, tuple(base_model_state[target_key].shape))
        )
        continue
    base_model_state[target_key] = src_tensor
    replaced += 1

if replaced == 0:
    raise RuntimeError("No distilled backbone parameters were merged into the COCO checkpoint.")

merged_state["model"] = base_model_state
torch.save(merged_state, merged_ckpt)

print(f"[merge] wrote merged checkpoint: {merged_ckpt}")
print(f"[merge] replaced backbone params: {replaced}")
if missing:
    print(f"[merge] missing target keys: {len(missing)} (first: {missing[0][0]} -> {missing[0][1]})")
if shape_mismatch:
    first = shape_mismatch[0]
    print(f"[merge] shape mismatches: {len(shape_mismatch)} (first: {first[0]} {first[1]} -> {first[2]} {first[3]})")
PY
}

verify_merged_init_ckpt() {
  local coco_ckpt="$1"
  local distilled_backbone_ckpt="$2"
  local merged_ckpt="$3"

  python - "${coco_ckpt}" "${distilled_backbone_ckpt}" "${merged_ckpt}" <<'PY'
import sys
import torch

coco_ckpt, distilled_ckpt, merged_ckpt = sys.argv[1:4]

def extract_model_state(state, name):
    if isinstance(state, dict) and "model" in state and isinstance(state["model"], dict):
        return state["model"]
    if isinstance(state, dict):
        return state
    raise RuntimeError(f"Unsupported {name} checkpoint format: {type(state)}")

def normalize_source_key(key: str) -> str:
    if key.startswith("module."):
        key = key[len("module."):]
    for prefix in ("backbone.dinov3._model.", "dinov3._model.", "_model."):
        if key.startswith(prefix):
            return key[len(prefix):]
    return key

base = extract_model_state(torch.load(coco_ckpt, map_location="cpu"), "COCO")
merged = extract_model_state(torch.load(merged_ckpt, map_location="cpu"), "merged")
distilled = torch.load(distilled_ckpt, map_location="cpu")
if isinstance(distilled, dict) and "model" in distilled and isinstance(distilled["model"], dict):
    distilled = distilled["model"]
elif isinstance(distilled, dict) and "state_dict" in distilled and isinstance(distilled["state_dict"], dict):
    distilled = distilled["state_dict"]
elif not isinstance(distilled, dict):
    raise RuntimeError(f"Unsupported distilled checkpoint format: {type(distilled)}")

distilled_items = [(k, v) for k, v in distilled.items() if torch.is_tensor(v)]

bb_ok = 0
bb_missing = 0
bb_shape = 0
bb_mismatch = 0
for src_key, src_tensor in distilled_items:
    target_key = f"backbone.dinov3._model.{normalize_source_key(src_key)}"
    if target_key not in merged:
        bb_missing += 1
        continue
    if merged[target_key].shape != src_tensor.shape:
        bb_shape += 1
        continue
    if torch.equal(merged[target_key], src_tensor):
        bb_ok += 1
    else:
        bb_mismatch += 1

non_backbone_diff = 0
for key, tensor in base.items():
    if key.startswith("backbone.dinov3._model."):
        continue
    if key not in merged:
        non_backbone_diff += 1
        continue
    if merged[key].shape != tensor.shape:
        non_backbone_diff += 1
        continue
    if not torch.equal(merged[key], tensor):
        non_backbone_diff += 1

print(f"[verify] distilled backbone tensors: {len(distilled_items)}")
print(f"[verify] matched distilled tensors in merged backbone: {bb_ok}")
print(f"[verify] backbone missing={bb_missing}, shape_mismatch={bb_shape}, value_mismatch={bb_mismatch}")
print(f"[verify] non-backbone tensors changed vs COCO init: {non_backbone_diff}")

if bb_ok != len(distilled_items) or bb_missing or bb_shape or bb_mismatch or non_backbone_diff:
    raise SystemExit(1)
PY
}

RESUME_CKPT="${RESUME_CKPT:-${OUTPUT_DIR}/last.pth}"
START_CKPT="${START_CKPT:-${OUTPUT_DIR}/last.pth}"
mode_args=()
if [[ -f "${RESUME_CKPT}" ]]; then
  echo "Resuming interrupted run from ${RESUME_CKPT}"
  mode_args=(-r "${RESUME_CKPT}")
elif [[ -f "${START_CKPT}" ]]; then
  echo "Resume checkpoint not found; starting fresh run from tuning checkpoint ${START_CKPT}"
  mode_args=(-t "${START_CKPT}")
else
  for required in "${COCO_INIT_CKPT}" "${DISTILLED_BACKBONE_CKPT}"; do
    if [[ ! -e "${required}" ]]; then
      echo "Missing required initialization file: ${required}" >&2
      exit 1
    fi
  done

  if [[ "${REBUILD_MERGED_INIT_CKPT:-0}" == "1" || ! -f "${MERGED_INIT_CKPT}" ]]; then
    echo "Preparing merged init checkpoint:"
    echo "  COCO init: ${COCO_INIT_CKPT}"
    echo "  Distilled backbone: ${DISTILLED_BACKBONE_CKPT}"
    echo "  Merged output: ${MERGED_INIT_CKPT}"
    build_merged_init_ckpt "${COCO_INIT_CKPT}" "${DISTILLED_BACKBONE_CKPT}" "${MERGED_INIT_CKPT}"
  else
    echo "Using existing merged init checkpoint ${MERGED_INIT_CKPT}"
  fi

  if [[ "${VERIFY_MERGED_INIT_CKPT}" == "1" ]]; then
    echo "Verifying merged init checkpoint integrity..."
    verify_merged_init_ckpt "${COCO_INIT_CKPT}" "${DISTILLED_BACKBONE_CKPT}" "${MERGED_INIT_CKPT}"
  fi

  echo "No resume/tuning checkpoint found; starting finetune from merged checkpoint ${MERGED_INIT_CKPT}"
  mode_args=(-t "${MERGED_INIT_CKPT}")
fi

cmd=(
  torchrun
  --master_port="${MASTER_PORT}"
  --nproc_per_node="${NPROC_PER_NODE}"
  train.py
  -c "${CONFIG}"
  --use-amp
  --seed "${SEED}"
  --output-dir "${OUTPUT_DIR}"
  -u "DINOv3STAs.weights_path=null"
  -u "train_dataloader.dataset.ann_file=${TRAIN_ANN}"
  -u "val_dataloader.dataset.ann_file=${VAL_ANN}"
  "${mode_args[@]}"
)

if [[ "$#" -gt 0 ]]; then
  cmd+=("$@")
fi

echo "Launching command:"
echo "PYTORCH_CUDA_ALLOC_CONF=${PYTORCH_CUDA_ALLOC_CONF} NCCL_ASYNC_ERROR_HANDLING=${NCCL_ASYNC_ERROR_HANDLING} NCCL_IB_DISABLE=${NCCL_IB_DISABLE} NCCL_P2P_DISABLE=${NCCL_P2P_DISABLE} CUDA_VISIBLE_DEVICES=${CUDA_DEVICES} ${cmd[*]}"

if [[ "${DRY_RUN:-0}" == "1" ]]; then
  echo "DRY_RUN=1 set; exiting before training launch."
  exit 0
fi

PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF}" \
NCCL_ASYNC_ERROR_HANDLING="${NCCL_ASYNC_ERROR_HANDLING}" \
NCCL_IB_DISABLE="${NCCL_IB_DISABLE}" \
NCCL_P2P_DISABLE="${NCCL_P2P_DISABLE}" \
CUDA_VISIBLE_DEVICES="${CUDA_DEVICES}" \
"${cmd[@]}" 2>&1 | tee "${log_file}"