[tab] /fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/train_dyn_dit_32889.log (log)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=24:00:00
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --mem=100GB
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/diffusion/%x_%j.log
#SBATCH --job-name=train_dyn_dit

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_breakout/breakout_episodes_80m/val

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

tokenizer_checkpoint=/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/interactive/32627

export PYTHONUNBUFFERED=1

srun python jasmine/train_dynamics_diffusion.py \
    --save_ckpt \
    --image_height=10 \
    --image_width=10 \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=128 \
    --patch_size=4 \
    --log_interval 1 \
    --val_steps 1 \
    --val_interval 10 \
    --log_image_interval 10 \
    --diffusion_denoise_steps=32 \
    --name=train_dyn_dit_breakout_dev \
    --tags dyn diffusion breakout \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
    --tokenizer_checkpoint=$tokenizer_checkpoint \

Running on 1 devices.
Counting all components: ['dynamics', 'lam', 'tokenizer']
wandb: Currently logged in as: mihir-mahajan2002 (instant-uv) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
wandb: Tracking run with wandb version 0.22.1
wandb: Run data is saved locally in /fast/home/mihir.mahajan/Projects/jasmine/wandb/run-20251027_134953-v1qx7iaj
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run train_dyn_dit_breakout_dev
wandb: ⭐️ View project at https://wandb.ai/instant-uv/jafar
wandb: 🚀 View run at https://wandb.ai/instant-uv/jafar/runs/v1qx7iaj
Parameter counts:
{'dynamics': 40510560, 'lam': 17640416, 'tokenizer': 33717536, 'total': 91868512}
WARNING:absl:Path /fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/interactive/32627/005260.orbax-checkpoint-tmp could not be identified as a finalized checkpoint path using <class 'orbax.checkpoint._src.path.atomicity.AtomicRenameTemporaryPath'>. Got error: Expected final AtomicRenameTemporaryPath (/fast/project/HFMI_ProFOUND/mihir/checkpoints/tokenizer/mae/interactive/32627/005260.orbax-checkpoint-tmp) not to end with ".orbax-checkpoint-tmp".
/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/orbax/checkpoint/_src/serialization/type_handlers.py:1269: UserWarning: Sharding info not provided when restoring. Populating sharding info from sharding file. Please note restoration time will be slightly increased due to reading from file. Note also that this option is unsafe when restoring on a different topology than the checkpoint was saved with.
  warnings.warn(
Total memory size: 2.0 GB, Output size: 1.0 GB, Temp size: 1.0 GB, Argument size: 1.0 GB, Host temp size: 0.0 GB.
FLOPs: 2.139e+11, Bytes: 3.797e+10 (35.4 GB), Intensity: 5.6 FLOPs/byte
Starting training from step 0...

Memstats: After params initialized:
	Using (GB) 1.2 / 71.27 (1.683738%) on cuda:0
Calculating validation metrics...
Step 10, validation loss: 0.8816443681716919
Calculating validation metrics...
Step 20, validation loss: 0.850609302520752
Calculating validation metrics...
Step 30, validation loss: 0.8010892868041992
Calculating validation metrics...
Step 40, validation loss: 0.7192520499229431
Calculating validation metrics...
Step 50, validation loss: 0.6315332651138306
Calculating validation metrics...
Step 60, validation loss: 0.5423277020454407
Calculating validation metrics...
Step 70, validation loss: 0.45385074615478516
Calculating validation metrics...
Step 80, validation loss: 0.3750932812690735
Calculating validation metrics...
Step 90, validation loss: 0.31556975841522217
Calculating validation metrics...
Step 100, validation loss: 0.2790672779083252
Calculating validation metrics...
Step 110, validation loss: 0.25860437750816345
Calculating validation metrics...
Step 120, validation loss: 0.24097195267677307
Calculating validation metrics...
Step 130, validation loss: 0.21976527571678162
Calculating validation metrics...
Step 140, validation loss: 0.20201927423477173
Calculating validation metrics...
Step 150, validation loss: 0.18454957008361816
Calculating validation metrics...
Step 160, validation loss: 0.1674945056438446
Calculating validation metrics...
Step 170, validation loss: 0.15205618739128113
Calculating validation metrics...
Step 180, validation loss: 0.13686075806617737
Calculating validation metrics...
Step 190, validation loss: 0.12516246736049652
Calculating validation metrics...
Step 200, validation loss: 0.11379821598529816
Calculating validation metrics...
Step 210, validation loss: 0.10418741405010223
Calculating validation metrics...
Step 220, validation loss: 0.09652290493249893
Calculating validation metrics...
Step 230, validation loss: 0.0894232839345932
Calculating validation metrics...
Step 240, validation loss: 0.08362925797700882
Calculating validation metrics...
Step 250, validation loss: 0.08002109825611115
Calculating validation metrics...
Step 260, validation loss: 0.0752774029970169
Calculating validation metrics...
srun: Job step aborted: Waiting up to 32 seconds for job step to finish.
[2025-10-27T13:55:46.056] error: *** JOB 32889 ON hai003 CANCELLED AT 2025-10-27T13:55:46 DUE to SIGNAL Terminated ***
[2025-10-27T13:55:46.057] error: *** STEP 32889.0 ON hai003 CANCELLED AT 2025-10-27T13:55:46 DUE to SIGNAL Terminated ***