p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] test.py (python)
def print_compiled_memory_stats(compiled_stats):
  """Prints a summary of the compiled memory statistics."""
  if compiled_stats is None:
    return

  def bytes_to_gb(num_bytes):
    return num_bytes / (1024**3)

  output_gb = bytes_to_gb(compiled_stats.output_size_in_bytes)
  temp_gb = bytes_to_gb(compiled_stats.temp_size_in_bytes)
  argument_gb = bytes_to_gb(compiled_stats.argument_size_in_bytes)
  alias_gb = bytes_to_gb(compiled_stats.alias_size_in_bytes)
  host_temp_gb = bytes_to_gb(compiled_stats.host_temp_size_in_bytes)
  total_gb = output_gb + temp_gb + argument_gb - alias_gb

  max_logging.log(
      f"Total memory size: {total_gb:.1f} GB, Output size: {output_gb:.1f} GB, Temp size: {temp_gb:.1f} GB, "
      f"Argument size: {argument_gb:.1f} GB, Host temp size: {host_temp_gb:.1f} GB."
  )
[content] train_tokenizer.py (python)
    # Pre-compile and print memory analysis for the jitted train_step
    if jax.process_index() == 0:
        try:
            first_videos = next(dataloader)
            sample_inputs = dict(videos=first_videos)
            compiled = train_step.lower(tokenizer, optimizer, sample_inputs).compile()
            print_compiled_memory_stats(compiled.memory_analysis())
            # Do not skip the first batch during training
            dataloader = itertools.chain([first_videos], dataloader)
        except Exception as e:
            print(f"Compiled memory analysis unavailable: {e}")
[content] train_tokenizer.py (python)
def print_compiled_memory_stats(compiled_stats):
    if compiled_stats is None:
        return
    def bytes_to_gb(num_bytes):
        return num_bytes / (1024**3)
    output_gb = bytes_to_gb(compiled_stats.output_size_in_bytes)
    temp_gb = bytes_to_gb(compiled_stats.temp_size_in_bytes)
    argument_gb = bytes_to_gb(compiled_stats.argument_size_in_bytes)
    alias_gb = bytes_to_gb(compiled_stats.alias_size_in_bytes)
    host_temp_gb = bytes_to_gb(compiled_stats.host_temp_size_in_bytes)
    total_gb = output_gb + temp_gb + argument_gb - alias_gb
    print(
        f"Total memory size: {total_gb:.1f} GB, Output size: {output_gb:.1f} GB, Temp size: {temp_gb:.1f} GB, "
        f"Argument size: {argument_gb:.1f} GB, Host temp size: {host_temp_gb:.1f} GB."
    )
[content] train_tokenizer.py (python)
def print_compiled_memory_stats(compiled_stats):
    if compiled_stats is None:
        return
    def bytes_to_gb(num_bytes):
        return num_bytes / (1024**3)
    output_gb = bytes_to_gb(compiled_stats.output_size_in_bytes)
    temp_gb = bytes_to_gb(compiled_stats.temp_size_in_bytes)
    argument_gb = bytes_to_gb(compiled_stats.argument_size_in_bytes)
    alias_gb = bytes_to_gb(compiled_stats.alias_size_in_bytes)
    host_temp_gb = bytes_to_gb(compiled_stats.host_temp_size_in_bytes)
    total_gb = output_gb + temp_gb + argument_gb - alias_gb
    print(
        f"Total memory size: {total_gb:.1f} GB, Output size: {output_gb:.1f} GB, Temp size: {temp_gb:.1f} GB, "
        f"Argument size: {argument_gb:.1f} GB, Host temp size: {host_temp_gb:.1f} GB."
    )


def print_compiled_cost_analysis(cost_stats):
    if cost_stats is None:
        return
    flops = float(cost_stats.get("flops", 0.0))
    bytes_accessed = float(cost_stats.get("bytes accessed", 0.0))
    transcendentals = cost_stats.get("transcendentals", None)
    optimal_seconds = cost_stats.get("optimal_seconds", None)
    gb = bytes_accessed / (1024 ** 3) if bytes_accessed else 0.0
    intensity = (flops / bytes_accessed) if bytes_accessed else float("nan")
    print(
        f"FLOPs: {flops:.3e}, Bytes: {bytes_accessed:.3e} ({gb:.1f} GB), "
        f"Intensity: {intensity:.1f} FLOPs/byte, Transcendentals: {transcendentals}, "
        f"Optimal seconds: {optimal_seconds}"
    )
[content] utils/utils.py (python)

import optax


def get_lr_schedule(
    lr_schedule: str,
    init_lr: float,
    max_lr: float,
    decay_end: float,
    total_steps: int,
    warmup_steps: int,
    wsd_decay_steps: int,
) -> optax.Schedule:
    supported_schedules = ["wsd", "cos"]
    if lr_schedule == "cos":
        assert (
            warmup_steps <= total_steps
        ), "Warmup steps can't be greater than total steps."
        return optax.warmup_cosine_decay_schedule(
            init_value=init_lr,
            peak_value=max_lr,
            warmup_steps=warmup_steps,
            decay_steps=total_steps,  # Note: decay_steps includes the warmup steps, so we need to pass total value
            end_value=decay_end,
        )
    elif lr_schedule == "wsd":
        assert (
            warmup_steps + wsd_decay_steps <= total_steps
        ), "Warmup and decay period is longer than total steps."
        schedules = [
            optax.linear_schedule(
                init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps
            ),
            optax.constant_schedule(value=max_lr),
            optax.linear_schedule(
                init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps
            ),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
    else:
        raise ValueError(
            f"Learning rate schedule not supported. Please use one of {supported_schedules}"
        )
[content] utils/utils.py (python)

from jax.tree_util import tree_map, tree_reduce


def _count_leaf(x):
    """Count parameters in a single leaf node."""
    if hasattr(x, "size"):
        return x.size
    return 0


def _count_component(component_params):
    """Count total parameters in a component."""
    return tree_reduce(
        lambda x, y: x + y, tree_map(_count_leaf, component_params), initializer=0
    )


def count_parameters_by_component(params):
    """Count parameters for each component of the model.

    Args:
        params: Model parameters from nnx.split(model, nnx.Param, ...)

    Returns:
        Dictionary with parameter counts for each component
    """
    component_names = list(params.keys())
    print(f"Counting all components: {component_names}")

    counts = {}
    total_params = 0

    for name in component_names:
        component_params = params[name]
        count = _count_component(component_params)
        counts[name] = count
        total_params += count

    counts["total"] = total_params
    return counts
[tab] utils/train_utils.py (python)
import optax
from jax.tree_util import tree_map, tree_reduce


def get_lr_schedule(
    lr_schedule: str,
    init_lr: float,
    max_lr: float,
    decay_end: float,
    total_steps: int,
    warmup_steps: int,
    wsd_decay_steps: int,
) -> optax.Schedule:
    supported_schedules = ["wsd", "cos"]
    if lr_schedule == "cos":
        assert (
            warmup_steps <= total_steps
        ), "Warmup steps can't be greater than total steps."
        return optax.warmup_cosine_decay_schedule(
            init_value=init_lr,
            peak_value=max_lr,
            warmup_steps=warmup_steps,
            decay_steps=total_steps,  # Note: decay_steps includes the warmup steps, so we need to pass total value
            end_value=decay_end,
        )
    elif lr_schedule == "wsd":
        assert (
            warmup_steps + wsd_decay_steps <= total_steps
        ), "Warmup and decay period is longer than total steps."
        schedules = [
            optax.linear_schedule(
                init_value=init_lr, end_value=max_lr, transition_steps=warmup_steps
            ),
            optax.constant_schedule(value=max_lr),
            optax.linear_schedule(
                init_value=max_lr, end_value=decay_end, transition_steps=wsd_decay_steps
            ),
        ]
        boundaries = [warmup_steps, total_steps - wsd_decay_steps]
        return optax.join_schedules(schedules, boundaries)
    else:
        raise ValueError(
            f"Learning rate schedule not supported. Please use one of {supported_schedules}"
        )


def _count_leaf(x):
    """Count parameters in a single leaf node."""
    if hasattr(x, "size"):
        return x.size
    return 0


def _count_component(component_params):
    """Count total parameters in a component."""
    return tree_reduce(
        lambda x, y: x + y, tree_map(_count_leaf, component_params), initializer=0
    )


def count_parameters_by_component(params):
    """Count parameters for each component of the model.

    Args:
        params: Model parameters from nnx.split(model, nnx.Param, ...)

    Returns:
        Dictionary with parameter counts for each component
    """
    component_names = list(params.keys())
    print(f"Counting all components: {component_names}")

    counts = {}
    total_params = 0

    for name in component_names:
        component_params = params[name]
        count = _count_component(component_params)
        counts[name] = count
        total_params += count

    counts["total"] = total_params
    return counts
[content] utils/train_utils.py (python)



def bytes_to_gb(num_bytes):
    return num_bytes / (1024**3)

def print_compiled_memory_stats(compiled_stats):
    output_gb = bytes_to_gb(compiled_stats.output_size_in_bytes)
    temp_gb = bytes_to_gb(compiled_stats.temp_size_in_bytes)
    argument_gb = bytes_to_gb(compiled_stats.argument_size_in_bytes)
    alias_gb = bytes_to_gb(compiled_stats.alias_size_in_bytes)
    host_temp_gb = bytes_to_gb(compiled_stats.host_temp_size_in_bytes)
    total_gb = output_gb + temp_gb + argument_gb - alias_gb
    print(
        f"Total memory size: {total_gb:.1f} GB, Output size: {output_gb:.1f} GB, Temp size: {temp_gb:.1f} GB, "
        f"Argument size: {argument_gb:.1f} GB, Host temp size: {host_temp_gb:.1f} GB."
    )

def print_compiled_cost_analysis(cost_stats):
    flops = float(cost_stats.get("flops", 0.0))
    bytes_accessed = float(cost_stats.get("bytes accessed", 0.0))
    gb = bytes_to_gb(bytes_accessed) if bytes_accessed else 0.0
    intensity = (flops / bytes_accessed) if bytes_accessed else float("nan")
    print(
        f"FLOPs: {flops:.3e}, Bytes: {bytes_accessed:.3e} ({gb:.1f} GB), "
        f"Intensity: {intensity:.1f} FLOPs/byte"
    )
[content] train_lam.py (python)

    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(tokenizer, optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
[tab] tests/data/generate_dummy_data.py (python)
import pickle
import numpy as np
from pathlib import Path
from array_record.python.array_record_module import ArrayRecordWriter


def generate_dummy_arrayrecord(
    output_path: Path,
    num_videos: int = 5,
    episode_length: int = 16,
    height: int = 90,
    width: int = 160,
    channels: int = 3,
    num_action_tokens: int = 11,
    vocab_size: int = 68,  # default vocab size from minerl
    seed: int = 42,
):
    """Generates a dummy ArrayRecord file with synthetic video data for testing."""
    print(f"Generating dummy ArrayRecord file at {output_path}")

    output_path.parent.mkdir(parents=True, exist_ok=True)

    writer = ArrayRecordWriter(str(output_path), "group_size:1")
    try:
        for i in range(num_videos):
            np.random.seed(seed + i)
            dummy_video = np.random.randint(
                0, 256, size=(episode_length, height, width, channels), dtype=np.uint8
            )

            actions = np.random.randint(
                0,
                vocab_size,
                size=(episode_length, num_action_tokens),
                dtype=np.uint8,
            )

            record = {
                "raw_video": dummy_video.tobytes(),
                "sequence_length": episode_length,
                "actions": actions,
            }

            writer.write(pickle.dumps(record))
    finally:
        writer.close()

    print("Dummy ArrayRecord generation complete.")


if __name__ == "__main__":
    test_dir = Path("data_arrayrecord/dummy")
    test_dir.mkdir(parents=True, exist_ok=True)
    dummy_file = test_dir / "dummy_test_shard.array_record"

    generate_dummy_arrayrecord(
        dummy_file, num_videos=1000
    )
    print(f"Generated dummy file: {dummy_file}")
[content] train_dynamics.py (python)

    if jax.process_index() == 0:
        first_videos = next(dataloader)
        sample_inputs = dict(videos=first_videos)
        compiled = train_step.lower(lam, optimizer, sample_inputs).compile()
        print_compiled_memory_stats(compiled.memory_analysis())
        print_compiled_cost_analysis(compiled.cost_analysis())
        # Do not skip the first batch during training
        dataloader = itertools.chain([first_videos], dataloader)
[tab] slurm/dev/alfred/horeka/jobs_old/coinrun/latent_action_ablation/train_dynamics_coinrun.sh (shellscript)
# Log the sbatch script
cat $0

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

job_name=$SLURM_JOB_NAME
slurm_job_id=$SLURM_JOB_ID

ws_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared'

CHECKPOINT_DIR=$ws_dir/data/checkpoints/$job_name_$slurm_job_id
mkdir -p $CHECKPOINT_DIR

data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/coinrun/coinrun_tfrecords'

lam_checkpoint='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/lam_ckpt_dir/lam_1748080132_50000'
tokenizer_checkpoint='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/checkpoints/tokenizer_ckpt_dir/tokenizer_1748080132_10000'


python train_dynamics.py \
    --ckpt_dir $CHECKPOINT_DIR \
    --lam_checkpoint $lam_checkpoint \
    --image_height 64 \
    --image_width 64 \
    --tokenizer_checkpoint $tokenizer_checkpoint \
    --batch_size=48 \
    --num_steps=10 \
    --warmup_steps=1 \
    --log_checkpoint_interval=500 \
    --log_image_interval=50 \
    --seed=0 \
    --name=$job_name \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --log \
    --entity instant-uv \
    --data_dir $data_dir \
    --project jafar