architecture-experiments

Repository to host architecture experiments and development using Paxml and Praxis

Example Run

export LOG_DIR=... # can be gcp bucket or local directory

python -m paxml.main \
    --exp=experiments.pile.BaseTransformerPile \
    --job_log_dir=${LOG_DIR}

ToDos

• ST-Moe Implementation
• muP Implementation