p-doom/crowd-code-dataset-0.2 huggingface dataset:
[tab] input_pipeline/download/huggingface/download_openai_array_records.sh (shellscript)
#!/bin/bash

# Download and extract array records from Hugging Face
# 
# This script performs a two-step process:
# 1. Downloads compressed array records from a Hugging Face dataset repository
# 2. Extracts the compressed tar files in parallel for better performance
#
# Usage:
#   ./download_openai_array_records.sh [hf_download_dir] [final_dataset_dir]
#
# Arguments:
#   hf_download_dir    - Directory to store compressed downloads (default: data/minecraft_arrayrecords_compressed)
#   final_dataset_dir  - Directory for extracted array records (default: data/minecraft_arrayrecords)

# Set default directories if not provided as arguments
hf_download_dir="${1:-data/minecraft_arrayrecords_compressed}" 
final_dataset_dir="${2:-data/minecraft_arrayrecords}"          

mkdir -p $hf_download_dir
mkdir -p $final_dataset_dir

# Step 1: Download compressed dataset from Hugging Face
echo "Starting download from Hugging Face..."
repo_id=p-doom/open_ai_minecraft_arrayrecords_chunked
start_time_hf_download=$(date +%s)

HF_HUB_ENABLE_HF_TRANSFER=1 HF_HUB_DISABLE_SYMLINKS=1 \
huggingface-cli download --repo-type dataset $repo_id --local-dir $hf_download_dir

end_time_hf_download=$(date +%s)
echo "Download completed. Time taken: $((end_time_hf_download - start_time_hf_download)) seconds"

# Step 2: Extract compressed array records in parallel
echo "Starting parallel extraction of tar files..."
num_workers=64  # Number of parallel extraction processes
start_time_uncompress=$(date +%s)

# Find all shard tar files and extract them in parallel:
xargs -0 -P $num_workers -I {} bash -c 'echo "Extracting {}"; tar -xf "{}" -C "'$final_dataset_dir'"'

end_time_uncompress=$(date +%s)

# Display timing summary
echo "================================"
echo "Extraction completed successfully!"
echo "Uncompress time: $((end_time_uncompress - start_time_uncompress)) seconds"
echo "Download time: $((end_time_hf_download - start_time_hf_download)) seconds"
echo "Total time: $((end_time_uncompress - start_time_hf_download)) seconds"
echo "Final dataset location: $final_dataset_dir"
[tab] test/test_torch.py (python)
import datetime
import torch

torch.backends.cuda.matmul.allow_tf32 = True

MATRIX_DIM = 32768
STEPS = 10

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

A = torch.ones((MATRIX_DIM, MATRIX_DIM), device=device, dtype=torch.float32)
B = torch.ones((MATRIX_DIM, MATRIX_DIM), device=device, dtype=torch.float32)

num_bytes = A.numel() * A.element_size()
total_num_bytes_crossing_to_hbm = num_bytes * 3

total_num_flops = 2 * MATRIX_DIM * MATRIX_DIM**2

def matmul(X, Y):
  return X @ Y

# warmup
with torch.no_grad():
  _ = matmul(A, B)
  if device.type == 'cuda':
    torch.cuda.synchronize()

start_time = datetime.datetime.now()
with torch.no_grad():
  for i in range(STEPS):
    C = A @ B
    if device.type == 'cuda':
      torch.cuda.synchronize()
end_time = datetime.datetime.now()

average_time_per_step = (end_time - start_time).total_seconds() / STEPS

print(f"{average_time_per_step}, teraflops per second: {total_num_flops / average_time_per_step / 1e12}, gigabytes per second: {total_num_bytes_crossing_to_hbm / average_time_per_step / 1e9}")
[tab] test.py (python)
import jax
import jax.numpy as jnp
import numpy as np
import functools

@jax.jit
def matmul(a, b):
    return jnp.matmul(a, b)

MATRIX_SIZE = 1024*12
a = jnp.ones((MATRIX_SIZE, MATRIX_SIZE))
b = jnp.ones((MATRIX_SIZE, MATRIX_SIZE))

compiled_matmul = matmul.lower(a, b).compile()

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

  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."
  )

print_compiled_memory_stats(compiled_matmul.memory_analysis())
[tab] salient_restore_failing_minimal_example.py (python)
import os
import shutil

import jax.numpy as jnp
from flax import nnx
import orbax.checkpoint as ocp

# The hunch is that parameters like biases in Linear layers are only created if
# they are used (e.g. `use_bias=True`). When checkpointing a model created with
# `use_bias=False`, the bias parameter is not saved.
#
# Attempting to restore this checkpoint into a model instance created with
# `use_bias=True` is expected to fail, as the checkpointed state has a
# different structure (missing the 'bias' leaf in the PyTree) from the target model.

ckpt_dir = '/tmp/flax_nnx_bug_ckpt'
if os.path.exists(ckpt_dir):
    shutil.rmtree(ckpt_dir)
os.makedirs(ckpt_dir)


class MyModel(nnx.Module):
  """A simple model containing a Linear layer."""

  def __init__(self, use_bias: bool, *, rngs: nnx.Rngs):
    self.linear = nnx.Linear(10, 20, use_bias=use_bias, rngs=rngs)

  def __call__(self, x):
    return self.linear(x)


rngs = nnx.Rngs(0)
model_no_bias = MyModel(use_bias=False, rngs=rngs)

dummy_input = jnp.ones((1, 10))
model_no_bias(dummy_input)

state_no_bias, _ = nnx.split(model_no_bias)
print('State of model_no_bias:', state_no_bias)

# Save the checkpoint.
checkpointer = ocp.StandardCheckpointer()
checkpointer.save(ckpt_dir, model_no_bias)
print(f"Checkpoint saved to '{ckpt_dir}'.
")


# 3. Instantiate a new model with use_bias=True.
print('Step 2: Creating a new model with use_bias=True.')
model_with_bias = MyModel(use_bias=True, rngs=rngs)

# Initialize parameters. This will create the 'bias' parameter.
model_with_bias(dummy_input)

# Display the state. Note the presence of 'bias'.
state_with_bias, _ = nnx.split(model_with_bias)
print('State of model_with_bias (before restore):', state_with_bias)
print('')


# 4. Attempt to restore the checkpoint into the new model.
# This is expected to fail because the checkpoint is missing the 'bias' parameter
# that exists in 'model_with_bias'.
print('Step 3: Attempting to restore the checkpoint.')
try:
  checkpoints.restore_checkpoint(model_with_bias, ckpt_dir)
  print('Restore successful. (This was not expected)')
except Exception as e:
  print('Restore failed as expected.')
  print(f'Error: {e}')
finally:
  # Clean up the checkpoint directory.
  shutil.rmtree(ckpt_dir)