[tab] /fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/maskgit/train_dyn_maskgit_coinrun_33018.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/maskgit/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/maskgit/%x_%j.log
#SBATCH --job-name=train_dyn_maskgit_coinrun

cat $0

source .venv/bin/activate

array_records_dir_train=/fast/project/HFMI_ProFOUND/mihir/data/data_coinrun/coinrun_episodes_50m/train
array_records_dir_val=/fast/project/HFMI_ProFOUND/mihir/data/data_coinrun/coinrun_episodes_50m/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/vqvae/coinrun_train_tokenizer_vqvae/32874


srun python jasmine/baselines/maskgit/train_dynamics_maskgit.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=train_dyn_maskgit_coinrun_$slurm_job_id \
    --tags dyn maskgit coinrun \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
/var/spool/slurm/job33018/slurm_script: line 15: .venv/bin/activate: No such file or directory
wandb: Currently logged in as: mihir-mahajan2002 (instant-uv) to https://api.wandb.ai. Use `wandb login --relogin` to force relogin
wandb: setting up run 007ax80i
wandb: Tracking run with wandb version 0.22.1
wandb: Run data is saved locally in /fast/home/mihir.mahajan/Projects/jasmine_runner_1/wandb/run-20251028_234514-007ax80i
wandb: Run `wandb offline` to turn off syncing.
wandb: Syncing run train_dyn_maskgit_coinrun_33018
wandb: ⭐️ View project at https://wandb.ai/instant-uv/jafar
wandb: 🚀 View run at https://wandb.ai/instant-uv/jafar/runs/007ax80i
/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(
Running on 1 devices.
Counting all components: ['dynamics', 'lam', 'tokenizer']
Parameter counts:
{'dynamics': 26555904, 'lam': 17640416, 'tokenizer': 34489696, 'total': 78686016}
Traceback (most recent call last):
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/baselines/maskgit/train_dynamics_maskgit.py", line 822, in <module>
    main(args)
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/baselines/maskgit/train_dynamics_maskgit.py", line 652, in main
    compiled = train_step.lower(optimizer, first_batch).compile()
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/compilation.py", line 476, in lower
    lowered = self.jitted_fn.lower(*pure_args, **pure_kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/compilation.py", line 126, in __call__
    out = self.f(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/baselines/maskgit/train_dynamics_maskgit.py", line 509, in train_step
    (loss, (recon, metrics)), grads = nnx.value_and_grad(loss_fn, has_aux=True)(
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/graph.py", line 2045, in update_context_manager_wrapper
    return f(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/autodiff.py", line 163, in grad_wrapper
    fn_out = gradded_fn(*pure_args)
             ^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/transforms/autodiff.py", line 88, in __call__
    out = self.f(*args)
          ^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/baselines/maskgit/train_dynamics_maskgit.py", line 507, in loss_fn
    return dynamics_loss_fn(model, inputs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/baselines/maskgit/train_dynamics_maskgit.py", line 492, in dynamics_loss_fn
    outputs = model(inputs)
              ^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/models/genie.py", line 185, in __call__
    assert type(latent_actions) == jax.Array
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
Filtering out episode with length 8, which is shorter than the requested sequence length 16.
Filtering out episode with length 10, which is shorter than the requested sequence length 16.
Filtering out episode with length 14, which is shorter than the requested sequence length 16.
Filtering out episode with length 9, which is shorter than the requested sequence length 16.
Filtering out episode with length 10, which is shorter than the requested sequence length 16.
Filtering out episode with length 8, which is shorter than the requested sequence length 16.
Filtering out episode with length 15, which is shorter than the requested sequence length 16.
wandb: 
wandb: 🚀 View run train_dyn_maskgit_coinrun_33018 at: https://wandb.ai/instant-uv/jafar/runs/007ax80i
wandb: Find logs at: ../../../../fast/home/mihir.mahajan/Projects/jasmine_runner_1/wandb/run-20251028_234514-007ax80i/logs
W1028 23:45:40.596809 3994001 pjrt_client.cc:1469] WatchJobStateAsync failed for task goo.gle/debugproto    job_name: "jax_worker": UNAVAILABLE: Cancelling all calls
Additional GRPC error information from remote target coordination_service while calling /tensorflow.CoordinationService/WatchJobState:
:UNKNOWN:Error received from peer  {grpc_message:"Cancelling all calls", grpc_status:14}
/opt/miniforge3/lib/python3.12/multiprocessing/resource_tracker.py:255: UserWarning: resource_tracker: There appear to be 20 leaked shared_memory objects to clean up at shutdown
  warnings.warn('resource_tracker: There appear to be %d '
srun: error: hai003: task 0: Exited with exit code 1