15750f56b2
- Update SLURM scripts to use correct CUDA modules (devel/cuda/12.4, intel compiler) - Add JAX downgrade to 0.4.35 for CuDNN 9.5.1 compatibility - Fix JAX_PLATFORMS environment variable (cuda vs gpu,cpu) - Update README with cluster-specific JAX installation steps - Tested successfully: Both PyTorch and JAX working on GPU with full training
46 lines
1.2 KiB
Bash
46 lines
1.2 KiB
Bash
#!/bin/bash
|
|
#SBATCH --job-name=fasttd3_test
|
|
#SBATCH --account=hk-project-p0022232
|
|
#SBATCH --partition=accelerated
|
|
#SBATCH --time=02:00:00
|
|
#SBATCH --gres=gpu:1
|
|
#SBATCH --ntasks=1
|
|
#SBATCH --cpus-per-task=8
|
|
#SBATCH --mem=32G
|
|
#SBATCH --output=fasttd3_%j.out
|
|
#SBATCH --error=fasttd3_%j.err
|
|
|
|
# Load necessary modules
|
|
module purge
|
|
module load devel/cuda/12.4
|
|
module load compiler/intel/2025.1_llvm
|
|
|
|
# Navigate to the project directory
|
|
cd $SLURM_SUBMIT_DIR
|
|
|
|
# Activate the virtual environment
|
|
source .venv/bin/activate
|
|
|
|
# Set environment variables for proper GPU usage
|
|
export CUDA_VISIBLE_DEVICES=$SLURM_LOCALID
|
|
export JAX_PLATFORMS="cuda"
|
|
export JAX_ENABLE_X64=True
|
|
|
|
# Ensure wandb is logged in (set WANDB_API_KEY environment variable)
|
|
# export WANDB_API_KEY=your_api_key_here
|
|
# For testing, use offline mode
|
|
export WANDB_MODE=offline
|
|
|
|
# Run FastTD3 training with MuJoCo Playground environment
|
|
python fast_td3/train.py \
|
|
--env_name T1JoystickFlatTerrain \
|
|
--exp_name FastTD3_HoReKa_Test \
|
|
--seed 42 \
|
|
--total_timesteps 25000 \
|
|
--num_envs 1024 \
|
|
--batch_size 4096 \
|
|
--eval_interval 5000 \
|
|
--render_interval 0 \
|
|
--project FastTD3_HoReKa
|
|
|
|
echo "Job completed at $(date)" |