Energy-Based Transformers (EBT)

The execution flow follows this hierarchy:

1. Job Script: Passes hyperparameters to train_model.py.
2. train_model.py: The primary entry point. It handles distributed training setup, GPU determination, model configuration, and argument parsing.
3. base_model_trainer.py: A PyTorch Lightning trainer that manages the training/validation loops, dataset instantiation (via setup), logging, and optimizers.
4. Model Modules: PyTorch Lightning modules (e.g., in model/nlp) that implement the actual architectures and forward/loss calculations.

Developer Tip: If you want to use EBT architectures in your own custom training loop, refer directly to the model/ directory for the model implementations and utilities.

The project relies on the following key tools:

• PyTorch Lightning: Used as a lightweight wrapper over PyTorch to handle distributed training, automate the training loop (no need for manual backwards() or zero_grad()), manage CUDA device placement, and provide callbacks.
• Weights and Biases (WandB): Used as the primary logger. It is configured to log hyperparameters, console logs (including model architecture and parameters), checkpoints, GPU counts, and gradients.
  • If you choose not to use WandB via the --no_wandb flag, logs are stored in logs/console.log.

Download the UCF dataset and annotations from the UCF website.

Important Note: The UCF website does not have a trusted certificate. If using wget to download, you must use the --no-check-certificate flag.

Required Directory Structure:

• `UCF101/
  • ucfTrainTestlist/` (annotations)
  • *.avi (videos)

Configuration: Set the --dataset_dir=<path_to_dataset> argument or the $SSV2_DIR environment variable.

SSv2 is distributed by Qualcomm. After downloading the video files and labels, you must concatenate and extract the videos.

Processing steps:

1. Concatenate and unzip video files:

   cat 20bn-something-something-v2-* > ssv2_archive
   tar -xvf ssv2_archive

2. Unzip labels:

   unzip 20bn-something-something-download-package-labels.zip

Configuration: Set the --dataset_dir=<path_to_dataset> argument or the $SSV2_DIR environment variable.

To set up the environment, it is recommended to use Conda. Create a new environment with Python 3.12 and install the dependencies from requirements.txt.

If you encounter issues with PyTorch or other packages, you can use alternative requirement files:

• gh200_requirements.txt: For GH200 systems.
• loose_requirements.txt: For systems without NVIDIA, PyTorch, or Triton packages.
• environment.yml: To create a conda environment using the provided YAML file.

conda create -n ebt python=3.12
conda activate ebt
pip install -r requirements.txt

Downloads for ImageNet, RedPajama, and MNIST are managed via HuggingFace. You must provide a HuggingFace User Access Token and specify a storage directory using environment variables before running any download scripts.

Required environment variables:

• HF_TOKEN: Your HuggingFace User Access Token.
• HF_HOME: The directory where HuggingFace downloads should be stored.

To use Kinetics-400 and Kinetics-600, use the scripts provided by the CVDFoundation.

Configuration:

1. Provide the path to the downloaded dataset using the --dataset_dir=<path_to_dataset> argument.
2. Set the $K400_DIR environment variable.

Inference scripts are located in job_scripts/[modality]/inference/. To run inference using a pretrained checkpoint, use the following flags:

• --only_test_model_ckpt: Path to your .ckpt file.
• --only_test: Tells the script to skip training and only perform testing.
• --execution_mode "inference": Sets the execution mode to inference.

Note: Most hyperparameters in the inference script are inherited from the checkpoint via train_model.py. Ensure you provide a valid path to your .ckpt file.

# Example inference command structure
# (Actual command depends on the specific script in job_scripts/)
./job_scripts/nlp/inference/ebt.sh --only_test_model_ckpt /path/to/model.ckpt --only_test --execution_mode "inference"

Set the following environment variables to manage Hugging Face caches and authentication:

• HF_HOME: Path to your data/models cache directory.
• HF_TOKEN: Your Hugging Face authentication token.

To log training progress, use wandb login within your activated environment.

export HF_HOME=/path/to/cache                                                              
export HF_TOKEN=your_token_here
wandb login