Hugging Face's logo

Anbinh93
/
MAVT

Model card Files
xet
 Community
MAVT
File size: 4,940 Bytes
251713e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
#!/bin/bash
#SBATCH --job-name=mavt-s3-threed
#SBATCH --partition=defq
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=2
#SBATCH --gpus-per-node=2
#SBATCH --cpus-per-task=24
#SBATCH --mem=120G
#SBATCH --time=48:00:00
#SBATCH --output=logs/stage3_dgx_%j.log
#SBATCH --error=logs/stage3_dgx_%j.err

# ============================================================================
# MAVT Stage 3 — Image + Video + 3D triplane
#   - Init from stage1_5 step=120k (best converged Stage 1 ckpt;
#     Stage 2 has no ckpt >15k due to save_top_k=1 bug in old config).
#   - 2x A100 DDP, bf16-mixed
#   - 50k steps, batch=6/GPU × 2 GPU × accum=2 = effective batch 24
#   - LR 2e-5, warmup 500
#   - SigLIP2 fully unfrozen (stage 3)
#   - 3D from dataset/universal_3d/ (30 519 triplane objects, captions/3d.json)
#
# Prerequisites (already in place):
#   - dataset/universal_3d/3d_objects/renders → symlink to dataset/tripplane
#   - dataset/universal_3d/captions/3d.json (30 519 LVIS captions)
#   - dataset/image10k/train (WDS image shards)
#   - dataset/dataset_10m (video shards; 51% known corrupt, loader handles)
#
# Submit:
#   sbatch train_stage3_dgx.sh
# ============================================================================

set -euo pipefail

PROJECT_DIR="/home/user02/linhdang/Antokenizer"
cd "$PROJECT_DIR"

# --- Data paths -----------------------------------------------------------
IMAGE_SHARDS_DIR="$PROJECT_DIR/dataset/image10k/train"
VIDEO_SHARDS_DIR="$PROJECT_DIR/dataset/dataset_10m"
UNIVERSAL_ROOT="$PROJECT_DIR/dataset/universal_3d"

# --- Init checkpoint ------------------------------------------------------
INIT_CKPT="checkpoints/stage1_5/balanced/mavt-stage1_5-balanced-step=0120000-val/loss=0.1238.ckpt"

mkdir -p logs checkpoints/stage3_dgx

# --- Environment ----------------------------------------------------------
if [ -d "$PROJECT_DIR/.venv" ]; then
    source "$PROJECT_DIR/.venv/bin/activate"
elif [ -d "$HOME/miniconda3" ]; then
    source "$HOME/miniconda3/etc/profile.d/conda.sh"
    conda activate base
fi

export PYTHONPATH="$PROJECT_DIR/src:${PYTHONPATH:-}"
export TORCH_NCCL_BLOCKING_WAIT=1
export OMP_NUM_THREADS=8
export TOKENIZERS_PARALLELISM=false

# --- 3D data sanity -------------------------------------------------------
THREED_RENDERS="$UNIVERSAL_ROOT/3d_objects/renders"
N_3D=$(find -L "$THREED_RENDERS" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | wc -l)
echo "========================================"
echo "  MAVT Stage 3 — Image + Video + 3D"
echo "  GPUs requested:  2"
echo "  Image shards:    $IMAGE_SHARDS_DIR"
echo "  Video shards:    $VIDEO_SHARDS_DIR"
echo "  Universal root:  $UNIVERSAL_ROOT"
echo "  3D objects:      $N_3D renders"
echo "  Init ckpt:       $INIT_CKPT"
echo "========================================"

if [ "$N_3D" -lt 1000 ]; then
    echo "[FATAL] Too few 3D objects ($N_3D). Expected ≥ 1000."
    exit 1
fi

if [ ! -f "$INIT_CKPT" ]; then
    echo "[FATAL] Init checkpoint not found: $INIT_CKPT"
    exit 1
fi

# Use init_from_ckpt (soft restart, keeps model weights but resets optimizer).
# For Stage 3 init from Stage 1, this is the intended behavior — new threed
# poolers are created via prepare_poolers() during setup('fit').
CKPT_ARG="--model.init_from_ckpt $INIT_CKPT"

# --- Launch ---------------------------------------------------------------
srun --mpi=none python train.py fit \
    --config configs/model/mavt_base.yaml \
    --config configs/train/universal_data/stage3_universal.yaml \
    --config configs/train/universal_data/stage3_paths.yaml \
    --data.image_shards_dir "$IMAGE_SHARDS_DIR" \
    --data.video_shards_dir "$VIDEO_SHARDS_DIR" \
    --data.video_max_shards 100 \
    --data.universal_data_root "$UNIVERSAL_ROOT" \
    --data.active_modalities '["image", "video", "threed"]' \
    --model.active_modalities '["image", "video", "threed"]' \
    --data.image_resolution 256 \
    --data.video_frames 16 \
    --data.video_resolution 256 \
    --data.triplane_res 256 \
    --data.batch_size 6 \
    --data.num_workers 6 \
    --data.pin_memory true \
    --data.persistent_workers true \
    --data.prefetch_factor 3 \
    --model.training_stage 3 \
    --model.init_siglip2 true \
    --model.use_lpips true \
    --model.use_clip false \
    --model.w_l1 1.0 \
    --model.w_lpips 0.2 \
    --model.w_sem 0.3 \
    --model.w_temp 0.05 \
    --model.warmup_steps 500 \
    --model.total_steps 50000 \
    --model.weight_decay 0.01 \
    --model.grad_clip 1.0 \
    $CKPT_ARG \
    --trainer.devices 2 \
    --trainer.strategy ddp_find_unused_parameters_true \
    --trainer.precision bf16-mixed \
    --trainer.max_steps 50000 \
    --trainer.accumulate_grad_batches 2 \
    --trainer.log_every_n_steps 50 \
    --trainer.val_check_interval 1000 \
    --trainer.logger.class_path lightning.pytorch.loggers.WandbLogger \
    --trainer.logger.init_args.project mavt \
    --trainer.logger.init_args.name stage3_threed