Environment Setup and Running Instructions

Environment Setup

1. Create a Conda Environment: Create a new Conda environment named moe with Python 3.10:
   conda create -n moe python=3.10 -y
2. Activate the Conda Environment: Activate the newly created environment:
   conda activate moe
3. Run Environment Setup Script: Execute the setup script to configure dependencies:
   bash moe.bash
4. Build DGL with CUDA Support: If using CUDA 12.8, build DGL with the provided script:
   bash build_dgl_cuda128.sh
5. Optional Runtime Overrides (recommended for stable runs):
   source runtime_env_overrides.sh

Debugging with Small Sample Sizes

1. Overfit on 1 Sample:
   bash

   1python moe.py \ 2 --stage overfit \ 3 --overfit-docs 1 \ 4 --epochs 10 \ 5 --device auto
2. Custom Sample Limit:
   bash

   1python moe.py \ 2 --stage train \ 3 --train-limit 10 \ 4 --eval-limit 5 \ 5 --epochs 2 \ 6 --device auto

Notes

• Ensure Conda is installed on your system before running the commands.
• The moe.bash script should contain all necessary dependency installations.
• The build_dgl_cuda128.sh script is specific to CUDA 12.8. Adjust if using a different CUDA version.

Running the Pipeline

Full Data Run

1. Distant Supervision Cleaning:
   bash

   1python moe.py \ 2 --stage ds \ 3 --distant-file train_distant.json \ 4 --eval-file dev.json \ 5 --distant-clean-file train_distant_clean.json \ 6 --build-distant-clean \ 7 --distant-topk 2 \ 8 --device auto
2. Supervised Training:
   bash

   1python -u moe.py \ 2--stage train \ 3--epochs 20 \ 4--full-train \ 5--full-eval \ 6--full-test \ 7--pretrained-gnn-checkpoint "<wandb_artifact_url_or_ref>" \ 8--no-pair-markers \ 9--result-dir inference_results \ 10--result-file result.json

Smoke Test

1. Distant Supervision Cleaning (Smoke Test):
   bash

   1python moe.py \ 2 --stage ds \ 3 --distant-file train_distant.json \ 4 --eval-file dev.json \ 5 --distant-clean-file train_distant_clean.json \ 6 --build-distant-clean \ 7 --distant-topk 2 \ 8 --device auto \ 9 --debug
2. Supervised Training (Smoke Test):
   bash

   1python moe.py \ 2 --stage train \ 3 --train-file train_annotated.json \ 4 --distant-clean-file train_distant_clean.json \ 5 --eval-file dev.json \ 6 --test-file dev.json \ 7 --epochs 2 \ 8 --num-experts 2 \ 9 --capacity-factor 1.0 \ 10 --lambda-moe 0.1 \ 11 --lambda-scl 0.05 \ 12 --scl-temp 0.1 \ 13 --candidate-keep-ratio 0.3 \ 14 --adaptive-threshold-scale 1.0 \ 15 --max-pairs-per-doc 25 \ 16 --max-seq-length 1024 \ 17 --result-dir inference_results \ 18 --result-file result.json \ 19 --device auto \ 20 --debug
3. Overfit Smoke Test (fastest end-to-end check):
   bash

   1python moe.py \ 2 --stage overfit \ 3 --overfit-docs 1 \ 4 --epochs 2 \ 5 --model-id sshleifer/tiny-gpt2 \ 6 --result-dir inference_results \ 7 --result-file result.json \ 8 --device auto \ 9 --debug

Command Validity Notes

• The commands above are aligned with current moe.py arguments (including --max-seq-length and multi-label pipeline settings).
• If you want logs on Weights & Biases, keep default behavior (do not pass --no-wandb).
• W&B now defaults to --wandb-mode online. If your shell/directory was previously set to offline, the run still forces online unless you explicitly pass --wandb-mode offline.
• To avoid silent fallback to offline when cloud init fails, add --wandb-no-offline-fallback.
• If you hit 403 Forbidden, avoid stale hardcoded credentials and run with your own account context, for example: wandb login --relogin then add --wandb-project <your_project> (optionally --wandb-entity <your_team>).
• For CPU-only debugging, use --model-id sshleifer/tiny-gpt2 to avoid loading a very large backbone.
• With W&B enabled, the run uploads artifacts for:
  • best checkpoint (best_model.pt)
  • inference result JSON and files from --result-dir (for example inference_results/*.json)
  • full workspace snapshot at run start (all folders and child files under workspace)
• Server-side checkpoint retention keeps only the latest run folder under checkpoints/.
• To warm-start from a previous run, pass --pretrained-gnn-checkpoint with one of:
  • local file path (for example checkpoints/<run_name>/best_model.pt)
  • W&B artifact reference (for example <entity>/<project>/<artifact_name>:v0)
  • W&B artifact URL (for example https://wandb.ai/<entity>/<project>/artifacts/model/<artifact_name>/v0)
• Add --pretrained-gnn-strict if you want strict key matching when loading the checkpoint.

Command full data