p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/dev/mihir/horeka/overfit_sample/causal/dynamics_overfit_sample.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:1
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/yoloruns/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/yoloruns/%x_%j.log
#SBATCH --job-name=train_dynamics_overfit_sample_causal

# Log the sbatch script
cat $0

module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
# source .venv/bin/activate

array_records_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data_new/open_ai_minecraft_arrayrecords_chunked

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

CHECKPOINT_DIR=$ws_dir/checkpoints/causal-fr/overfit-seed69-1-no-noise-4/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/tokenizer/train_tokenizer_1e-4/3388151


env | grep SLURM

srun python train_dynamics.py \
    --save_ckpt \
    --seed=69 \
    --num_steps=5000 \
    --warmup_steps=0 \
    --wsd_decay_steps=0 \
    --ckpt_dir $CHECKPOINT_DIR \
    --dyna_type=causal \
    --batch_size=8 \
    --init_lr=1e-4 \
    --max_lr=1e-4 \
    --log_image_interval=100 \
    --log \
    --log_checkpoint_interval=1000 \
    --name=dynamics-causal-overfit-$slurm_job_id \
    --tags dynamics causal overfit \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir \
    --no-use-flash-attention \
    --dyna_dim=128 \
    --dyna_num_blocks=2 \
    --dyna_num_heads=4
[content] utils/nn.py (python)
def _create_flash_attention_fn(use_flash_attention: bool, is_causal: bool) -> Callable:
    """
    Create an attention function that uses flash attention if enabled.

    Flax MultiHeadAttention provides tensors with shape (batch..., length, num_heads, head_dim)
    jax.nn.dot_product_attention expects (batch, length, num_heads, head_dim).

    We need to reshape to ensure compatibility. cuDNN's flash attention additionally
    requires a sequence length that is a multiple of 4. We pad the sequence length to the nearest
    multiple of 4 and mask accordingly.
    """

    def attention_fn(query, key, value, bias=None, mask=None, **kwargs):
        implementation = "cudnn" if use_flash_attention else None

        def _rearrange(x):
            return einops.rearrange(x, "... l h d -> (...) l h d")

        def _pad(x):
            return jnp.pad(x, ((0, 0), (0, pad_size), (0, 0), (0, 0)))

        def _fuse_masks(mask: jax.Array, attention_mask: jax.Array) -> jax.Array:
            mask_bool = mask.astype(jnp.bool_)
            expanded_mask = jnp.pad(
                mask_bool, ((0, pad_size), (0, pad_size)), constant_values=False
            )
            return jnp.logical_and(attention_mask, expanded_mask)

        original_shape = query.shape
        original_seq_len = query.shape[-3]

        # Pad to nearest multiple of 4
        target_seq_len = ((original_seq_len + 3) // 4) * 4
        pad_size = target_seq_len - original_seq_len

        query_4d = _pad(_rearrange(query))
        key_4d = _pad(_rearrange(key))
        value_4d = _pad(_rearrange(value))

        attention_mask = jnp.ones((target_seq_len, target_seq_len), dtype=jnp.bool_)
        attention_mask = attention_mask.at[original_seq_len:, :].set(False)
        attention_mask = attention_mask.at[:, original_seq_len:].set(False)

        mask_4d = (
            _fuse_masks(mask, attention_mask) if mask is not None else attention_mask
        )
        mask_4d = mask_4d[jnp.newaxis, jnp.newaxis, :, :]  # (1, 1, seq_len, seq_len)

        bias_4d = _pad(_rearrange(bias)) if bias is not None else None

        # NOTE: jax.nn.dot_product_attention does not support dropout
        output_4d = jax.nn.dot_product_attention(
            query=query_4d,
            key=key_4d,
            value=value_4d,
            bias=bias_4d,
            mask=mask_4d,
            implementation=implementation,
            is_causal=is_causal,
        )
        return output_4d[..., :original_seq_len, :, :].reshape(original_shape)

    return attention_fn