[tab] /fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/causal/train_dyn_causal_coinrun_32932.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/causal/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/dyn/causal/%x_%j.log
#SBATCH --job-name=train_dyn_causal_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_causal.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=train_dyn_causal_coinrun_$slurm_job_id \
    --tags dyn causal 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/job32932/slurm_script: line 15: .venv/bin/activate: No such file or directory
Running on 1 devices.
Traceback (most recent call last):
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/baselines/maskgit/train_dynamics_causal.py", line 818, in <module>
    main(args)
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/baselines/maskgit/train_dynamics_causal.py", line 421, in main
    genie, rng = build_model(args, rng)
                 ^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine_runner_1/jasmine/baselines/maskgit/train_dynamics_causal.py", line 110, in build_model
    genie = GenieMaskGIT(
            ^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 333, in __call__
    return _graph_node_meta_call(cls, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 345, in _graph_node_meta_call
    cls._pytree_meta_construct(node, *args, **kwargs)
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 336, in _pytree_meta_construct
    self.__init__(*args, **kwargs)
  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/models/genie.py", line 136, in __init__
    self.dynamics = DynamicsCausal(
                    ^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 333, in __call__
    return _graph_node_meta_call(cls, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 345, in _graph_node_meta_call
    cls._pytree_meta_construct(node, *args, **kwargs)
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 336, in _pytree_meta_construct
    self.__init__(*args, **kwargs)
  File "/fast/home/mihir.mahajan/Projects/jasmine/jasmine/models/dynamics.py", line 142, in __init__
    self.transformer = Transformer(
                       ^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 333, in __call__
    return _graph_node_meta_call(cls, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 353, in _graph_node_meta_call
    check_pytree(node)
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 326, in check_pytree
    pytree._check_value(name, value, new_status=None)
  File "/fast/home/mihir.mahajan/Projects/jasmine/.venv/lib/python3.12/site-packages/flax/nnx/pytreelib.py", line 588, in _check_value
    raise ValueError(
ValueError: Found unexpected Arrays on value of type <class 'list'> in static attribute 'blocks' of Pytree type '<class 'jasmine.utils.nn.Transformer'>'. This is an error starting from Flax version 0.12.0.
Consider one of the following options:

1. If the attribute is meant to be static, either remove the Array value or wrap it in a static container.
2. Wrap the value with nnx.data on assignment:

  _.blocks = nnx.data(...)

3. Annotate the class attribute with nnx.Data:

  class Transformer(Module):
    blocks: nnx.Data[list]

4. If the container is a list or dict, try using nnx.List(...) or nnx.Dict(...) instead.
5. Disable pytree for this class:

  class Transformer(Module, pytree=False):


W1028 00:31:11.043243 3617989 pjrt_client.cc:1469] WatchJobStateAsync failed for task goo.gle/debugproto  job_name: "jax_worker": CANCELLED: CANCELLED
Additional GRPC error information from remote target coordination_service while calling /tensorflow.CoordinationService/WatchJobState:
:UNKNOWN:Error received from peer  {grpc_status:1, grpc_message:"CANCELLED"} [type.googleapis.com/tensorflow.DerivedStatus='']
srun: error: hai004: task 0: Exited with exit code 1