p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] experiments/dynamics_grain_full_restore.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_dynamics.py \
    --batch_size 48 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] experiments/dynamics_grain_full_restore.sh (shellscript)
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=48:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_franz/big-runs/dynamics-cotraining-batchsize-scaling-linear-lr/%x_%j.log
#SBATCH --job-name=train_dynamics_batch_size_scaling_linear_lr_1_node
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate


data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_dynamics.py \
    --batch_size 48 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] experiments/tokenizer_grain_checkpointing.sh (shellscript)
#!/usr/bin/env bash
source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun --ntasks=8 python train_tokenizer.py \
    --batch_size 96 \
    --restore_ckpt \
    --save_ckpt \
    --log_checkpoint_interval 5 \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[tab] input_pipeline/download/openai/download_actions_files.py (python)
import subprocess
import json
import tyro
from dataclasses import dataclass
import os
from multiprocessing import Pool, cpu_count
from tqdm import tqdm


@dataclass
class Args:
    index_file: str = "data/open_ai_index_files/all_6xx_Jun_29.json"
    output_dir: str = "data/open_ai_minecraft_actions_files"
    num_workers: int = -1  # -1 means use all available cores


def flatten_path(relpath):
    """Convert nested path to flattened filename with subdirectory as prefix
    e.g. data/6.10/filename.mp4 -> 6.10_filename.mp4
    """

    parts = relpath.split("/")

    if len(parts) >= 3:
        subdir = parts[1]
        filename = parts[2]
        return f"{subdir}_{filename}"
    else:
        return relpath.replace("/", "_")


def download_file(args):
    try:
        url, base_dir, output_dir = args
        jsonl_url = url.rsplit(".", 1)[0] + ".jsonl"
        filename = flatten_path(jsonl_url)
        output_file = os.path.join(output_dir, filename)
        subprocess.run(["wget", "-q", base_dir + jsonl_url, "-O", output_file], check=True)
        return {'file': jsonl_url, 'success': True}
    except subprocess.CalledProcessError as e:
        # delete file if it exists
        if os.path.exists(output_file):
            os.remove(output_file)
        return {'file': jsonl_url, 'success': False, 'error': str(e)}


def download_actions_files(index_file: str, output_dir: str, num_workers: int):
    # load json file
    with open(index_file, "r") as f:
        data = json.load(f)

    base_dir = data["basedir"]
    urls = data["relpaths"]

    # Prepare arguments for each process
    args_list = [(url, base_dir, output_dir) for url in urls]

    results = []
    with tqdm(total=len(args_list), desc="Downloading actions files") as pbar:
        with Pool(processes=num_workers) as pool:
            for result in pool.imap_unordered(download_file, args_list):
                results.append(result)
                pbar.update(1)

    # save results to json
    meta_data_file_name = index_file.split("/")[-1].split(".")[0] + "_metadata.json"
    with open(os.path.join(output_dir, meta_data_file_name), "w") as f:
        json.dump(results, f)

    # print number of failed downloads
    failed_downloads = [result for result in results if not result["success"]]
    print(f"Number of failed downloads: {len(failed_downloads)}")

    # print number of successful downloads
    successful_downloads = [result for result in results if result["success"]]
    print(f"Number of successful downloads: {len(successful_downloads)}")

if __name__ == "__main__":
    args = tyro.cli(Args)
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    if args.num_workers == -1:
        args.num_workers = cpu_count()

    print(f"Index file: {args.index_file}")
    print(f"Output directory: {args.output_dir}")
    print(f"Number of workers: {args.num_workers}")

    download_actions_files(args.index_file, args.output_dir, args.num_workers)
[tab] experiments/dynamics_grain_tok_restore.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --time=12:00:00
#SBATCH --cpus-per-task=1
#SBATCH --gres=gpu:8
#SBATCH --mem=100GB

source .venv/bin/activate

data_dir="$PWD/data_arrayrecord/dummy"
ckpt_dir="$PWD/checkpoints/dynamics_openai_grain_tok_lam_restore"
tokenizer_ckpt_dir="$PWD/checkpoints/tokenizer_openai_grain_checkpointing"

export XLA_FLAGS=--xla_gpu_autotune_level=0
srun python train_dynamics.py \
    --batch_size 48 \
    --tokenizer_checkpoint $tokenizer_ckpt_dir \
    --ckpt_dir $ckpt_dir \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --init_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] slurm-13723.out (plaintext)
srun: error: hai001: task 4: Exited with exit code 1
srun: error: hai002: task 12: Exited with exit code 1
srun: error: hai001: tasks 1-3,5-7: Exited with exit code 1
srun: error: hai002: tasks 8-10,13-15: Exited with exit code 1
srun: error: hai001: task 0: Exited with exit code 1
srun: error: hai002: task 11: Exited with exit code 1