p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/single_gpu_dev.sh (shellscript)
#!/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 --output=/fast/project/HFMI_SynergyUnit/mihir/logs/nemo/%x_%j.log
#SBATCH --error=/fast/project/HFMI_SynergyUnit/mihir/logs/nemo/%x_%j.log
#SBATCH --job-name=qwen3_600m_finetune
#SBATCH --mem=400GB
#SBATCH --qos=low

# Log the sbatch script
cat $0

source .uv_venv/bin/activate

python slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/test_full_finetune_4k.py
[tab] slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/test_full_finetune_4k copy.py (python)
import os

import nemo_run as run
from nemo.collections import llm
from lightning.pytorch.loggers import WandbLogger

os.environ["NEMORUN_HOME"] = "/fast/project/HFMI_SynergyUnit/tab_model/franz/nemorun_home"
os.environ["NEMO_CACHE_DIR"] = "/fast/project/HFMI_SynergyUnit/tab_model/franz/nemo_cache"
slurm_job_id = os.environ.get("SLURM_JOB_ID", "local")

NAME = f"qwen3_600m_full_4k_tokens_debug_{slurm_job_id}"
DATASET_DIR = "/fast/project/HFMI_SynergyUnit/tab_model/data/nemo_hf_part_jsonl_4k_tokens/" 

def recipe_to_dict(recipe, depth=4):
    if depth == 0:
        return dict()
    cfg_dict = dict()
    for k, v in vars(recipe)['__arguments__'].items():
        if isinstance(v, run.Config) or isinstance(v, dict):
            cfg_dict[k] = recipe_to_dict(v, depth-1)
        elif isinstance(v, (int, float, str, bool)):
            cfg_dict[k] = v
        else:
            cfg_dict[k] = str(v)
    return cfg_dict
                    
recipe = llm.qwen3_600m.finetune_recipe(
    name=NAME,
    dir=f"/fast/project/HFMI_SynergyUnit/mihir/logs",
    num_nodes=1,
    num_gpus_per_node=1,
    peft_scheme=None,
    # FIXME(f.srambical): check whether this is supported by ChatDataModule
    packed_sequence=False,
)

def local_executor_torchrun(nodes: int, devices: int) -> run.LocalExecutor:
    # Env vars for jobs are configured here
    env_vars = {
        "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
        "NCCL_NVLS_ENABLE": "0",
        "NVTE_DP_AMAX_REDUCE_INTERVAL": "0",
        "NVTE_ASYNC_AMAX_REDUCTION": "1",
    }

    executor = run.LocalExecutor(ntasks_per_node=devices, launcher="torchrun", env_vars=env_vars)

    return executor

def run_finetuning():
    recipe.resume.restore_config.path = "/fast/project/HFMI_SynergyUnit/tab_model/data/checkpoints/nemo_converted_weights_qwen3-600m/"
    recipe.data = run.Config(
        llm.ChatDataModule,
        dataset_root=DATASET_DIR,
        seq_length=4096,
        global_batch_size=32,
        micro_batch_size=4,
        num_workers=0,
    )

    MAX_STEPS = 11
    VAL_INTERVAL = 2
    NUM_CHECKPOINTS = 10
    recipe.trainer.max_steps = MAX_STEPS
    recipe.log.ckpt.every_n_train_steps = MAX_STEPS // NUM_CHECKPOINTS
    recipe.log.ckpt.save_top_k = NUM_CHECKPOINTS
    recipe.trainer.val_check_interval = VAL_INTERVAL

    cfg_dict = recipe_to_dict(recipe)
    recipe.log.wandb = run.Config(
        WandbLogger,
        name=NAME,
        project="crowd-pilot-nemo",
        save_dir="/fast/project/HFMI_SynergyUnit/mihir/logs/wandb",
        entity="instant-uv",
        config=cfg_dict,
    )

    executor = local_executor_torchrun(nodes=recipe.trainer.num_nodes, devices=recipe.trainer.devices)

    run.run(recipe, executor=executor, name=NAME)

# This condition is necessary for the script to be compatible with Python's multiprocessing module.
if __name__ == "__main__":
    run_finetuning()
[tab] slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/test_full_finetune_4k_requeue.py (python)
import os

import nemo_run as run
from nemo.collections import llm
from lightning.pytorch.loggers import WandbLogger

os.environ["NEMORUN_HOME"] = "/fast/project/HFMI_SynergyUnit/tab_model/franz/nemorun_home"
os.environ["NEMO_CACHE_DIR"] = "/fast/project/HFMI_SynergyUnit/tab_model/franz/nemo_cache"
slurm_job_id = os.environ.get("SLURM_JOB_ID", "local")

NAME = f"qwen3_600m_full_4k_tokens_debug_{slurm_job_id}"
DATASET_DIR = "/fast/project/HFMI_SynergyUnit/tab_model/data/nemo_hf_part_jsonl_4k_tokens/" 

def recipe_to_dict(recipe, depth=4):
    if depth == 0:
        return dict()
    cfg_dict = dict()
    for k, v in vars(recipe)['__arguments__'].items():
        if isinstance(v, run.Config) or isinstance(v, dict):
            cfg_dict[k] = recipe_to_dict(v, depth-1)
        elif isinstance(v, (int, float, str, bool)):
            cfg_dict[k] = v
        else:
            cfg_dict[k] = str(v)
    return cfg_dict
                    
recipe = llm.qwen3_600m.finetune_recipe(
    name=NAME,
    dir=f"/fast/project/HFMI_SynergyUnit/mihir/logs",
    num_nodes=1,
    num_gpus_per_node=1,
    peft_scheme=None,
    # FIXME(f.srambical): check whether this is supported by ChatDataModule
    packed_sequence=False,
)

def local_executor_torchrun(nodes: int, devices: int) -> run.LocalExecutor:
    # Env vars for jobs are configured here
    env_vars = {
        "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
        "NCCL_NVLS_ENABLE": "0",
        "NVTE_DP_AMAX_REDUCE_INTERVAL": "0",
        "NVTE_ASYNC_AMAX_REDUCTION": "1",
    }

    executor = run.LocalExecutor(ntasks_per_node=devices, launcher="torchrun", env_vars=env_vars)

    return executor

def run_finetuning():
    recipe.resume.restore_config.path = "/fast/project/HFMI_SynergyUnit/tab_model/data/checkpoints/nemo_converted_weights_qwen3-600m/"
    recipe.data = run.Config(
        llm.ChatDataModule,
        dataset_root=DATASET_DIR,
        seq_length=4096,
        global_batch_size=32,
        micro_batch_size=4,
        num_workers=0,
    )

    MAX_STEPS = 11
    VAL_INTERVAL = 2
    NUM_CHECKPOINTS = 10
    recipe.trainer.max_steps = MAX_STEPS
    recipe.log.ckpt.every_n_train_steps = MAX_STEPS // NUM_CHECKPOINTS
    recipe.log.ckpt.save_top_k = NUM_CHECKPOINTS
    recipe.trainer.val_check_interval = VAL_INTERVAL

    cfg_dict = recipe_to_dict(recipe)
    recipe.log.wandb = run.Config(
        WandbLogger,
        name=NAME,
        project="crowd-pilot-nemo",
        save_dir="/fast/project/HFMI_SynergyUnit/mihir/logs/wandb",
        entity="instant-uv",
        config=cfg_dict,
    )

    executor = local_executor_torchrun(nodes=recipe.trainer.num_nodes, devices=recipe.trainer.devices)

    run.run(recipe, executor=executor, name=NAME)

# This condition is necessary for the script to be compatible with Python's multiprocessing module.
if __name__ == "__main__":
    run_finetuning()
[content] slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/test_full_finetune_4k_requeue.py (python)

        recipe.resume = run.Config(
            nl.AutoResume,
            restore_config=run.Config(nl.RestoreConfig, path=CKPT_PATH),
            resume_if_exists=True,
        )
    else:
        recipe.resume.restore_config.path = BASE_CKPT_PATH
[content] slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/test_full_finetune_4k_requeue.py (python)
    if os.path.exists(CKPT_PATH):
        # find first directory under CKPT_PATH
        first_dir = next(os.path.dirname(CKPT_PATH) for _ in os.listdir(CKPT_PATH) if os.path.isdir(os.path.join(CKPT_PATH, _)))
        # use first directory as CKPT_PATH
        breakpoint()
        path = first_dir
        recipe.resume = run.Config(
            nl.AutoResume,
            restore_config=run.Config(nl.RestoreConfig, path=path),
            resume_if_exists=True,
        )
    else:
        recipe.resume.restore_config.path = BASE_CKPT_PATH
[content] slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/single_gpu_dev.sh (shellscript)

# --- signal trap to requeue job before timeout ---
requeue_job() {
    echo "[$(date)] caught sigusr1 (timeout warning), requeueing slurm job $SLURM_JOB_ID..."
    scontrol requeue $SLURM_JOB_ID
    exit 0
}

trap requeue_job sigusr1
[tab] slurm/jobs/mihir/berlin/diffusion/coinrun/dyn_dit_coinrun.sh (shellscript)
#!/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_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/mae/coinrun_train_tokenizer_mae/32925


srun python jasmine/baselines/diffusion/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=25 \
    --name=train_dyn_diffusion_coinrun_$slurm_job_id \
    --tags dyn diffusion coinrun \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[tab] slurm/jobs/mihir/berlin/maskgit/train_dynamics_causal.sh (shellscript)
#!/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 \
[tab] slurm/jobs/mihir/berlin/maskgit/train_tokenizer_vqvae.sh (shellscript)
#!/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/tokenizer/mae/%x_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/train/tokenizer/mae/%x_%j.log
#SBATCH --job-name=coinrun_train_tokenizer_vqvae

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/tokenizer/vqvae/$job_name/$slurm_job_id
mkdir -p $CHECKPOINT_DIR

srun python jasmine/baselines/maskgit/train_tokenizer_vqvae.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --name=train_tokenizer_vqvae_coinrun_$slurm_job_id \
    --tags tokenizer vqvae coinrun \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
[tab] slurm/jobs/mihir/berlin/maskgit/train_dynamics_maskgit.sh (shellscript)
#!/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 \
[tab] slurm/jobs/mihir/berlin/maskgit/sample_coinrun.sh (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=1
#SBATCH --ntasks-per-node=4
#SBATCH --time=01:00:00
#SBATCH --partition=accelerated
#SBATCH --cpus-per-task=5
#SBATCH --gres=gpu:4
#SBATCH --output=/fast/project/HFMI_ProFOUND/mihir/logs/sample/coinrun/diffusion/sample_coinrun_%j.log
#SBATCH --error=/fast/project/HFMI_ProFOUND/mihir/logs/sample/coinrun/diffusion/sample_coinrun_%j.log
#SBATCH --job-name=sample_coinrun

# Activate virtual environment
source .venv/bin/activate

array_records_dir=/fast/project/HFMI_ProFOUND/mihir/data/data_coinrun/coinrun_episodes_50m/test

CHECKPOINT_PATH=/fast/project/HFMI_ProFOUND/mihir/checkpoints/dyn/train_dyn_maskgit_coinrun/33019

echo "Sampling from checkpoint: $CHECKPOINT_PATH"

srun python jasmine/baselines/maskgit/sample_maskgit.py \
    --checkpoint "$CHECKPOINT_PATH" \
    --data_dir=$array_records_dir \
    --seq_len=16 \
    --batch_size=5 \
    --seed=3 \
    --start_frame=4 \
    --image_height=64 \
    --image_width=64 \
    --print-action-indices \
    --output_dir "gifs/maskgit/coinrun/submission" \
[tab] slurm/jobs/mihir/berlin/jasmine/diffusion/coinrun/dyn_dit_coinrun_lr_3e-5.sh (shellscript)
#!/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_diffusion_coinrun_lr_3e-5

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/mae/coinrun_train_tokenizer_mae/32632

srun python jasmine/baselines/diffusion/train_dynamics_diffusion.py \
    --save_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --diffusion_denoise_steps=25 \
    --max_lr=3e-5 \
    --name=train_dyn_diffusion_coinrun_lr_3e-5_$slurm_job_id \
    --tags dyn diffusion coinrun lr_3e-5 \
    --entity instant-uv \
    --project jafar \
    --data_dir $array_records_dir_train \
    --val_data_dir $array_records_dir_val \
    --tokenizer_checkpoint=$tokenizer_checkpoint \
[tab] slurm/jobs/mihir/horeka/action_space_scaling/co_training/train_dynamics_cotraining_8_actions.sbatch (shellscript)
#!/usr/bin/env bash

#SBATCH --nodes=2
#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_mihir/big-runs/dynamics-cotraining-actionspace-scaling/%x_%j.log
#SBATCH --error=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/logs/logs_mihir/big-runs/dynamics-cotraining-actionspace-scaling/%x_%j.log
#SBATCH --job-name=dynamics_cotraining_action_space_scaling_8

# 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=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/big-runs/dynamics-cotraining-actionspace-scaling/$job_name
mkdir -p $CHECKPOINT_DIR

env | grep SLURM

tokenizer_ckpt_dir=/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/train_tokenizer_batch_size_scaling_16_node/3321526/tokenizer_22000/

env | grep SLURM

srun python jasmine/train_dynamics.py \
    --save_ckpt \
    --restore_ckpt \
    --ckpt_dir $CHECKPOINT_DIR \
    --batch_size=96 \
    --min_lr=0 \
    --max_lr=1.5e-5 \
    --log_image_interval=1000 \
    --log \
    --num_latent_actions=8 \
    --log_checkpoint_interval=1000 \
    --name=dynamics-cotraining-action-space-scaling-8-$slurm_job_id \
    --tags dynamics cotraining action-space-scaling 8 \
    --entity instant-uv \
    --project jafar \
    --tokenizer_checkpoint=$tokenizer_ckpt_dir \
    --data_dir $array_records_dir
[content] slurm/dev/mihir/berlin/crowd_pilot/nemo/iterating_on_stuff/single_gpu_dev.sh (shellscript)
# set checkpoint flag based on restart count
restart_count=$(scontrol show job $SLURM_JOB_ID | grep -o 'Restarts=[0-9]*' | cut -d'=' -f2)

if [ $restart_count -eq 0 ]; then
    restore_ckpt_flag="--no-restore-ckpt"
else
    restore_ckpt_flag="--restore-ckpt"
fi