Installation¤

Requirements¤

Python >= 3.10
torch >= 2.5
torchvision >= 0.20
torchaudio >= 2.5
transformers
peft
datasets

Install from PyPI¤

pip install unitorch

Install with Optional Extras¤

# All extras (recommended)
pip install "unitorch[all]"

# DeepSpeed support
pip install "unitorch[deepspeed]"

# Image generation (FLUX)
pip install "unitorch[diffusers]"

# Megatron-LM support
pip install "unitorch[megatron]"

# FastAPI / ONNX / W&B
pip install "unitorch[others]"

Install from Source¤

pip install "git+https://github.com/fuliucansheng/unitorch"

# With extras
pip install "git+https://github.com/fuliucansheng/unitorch#egg=unitorch[all]"

Install with CUDA Extensions¤

# Build the ngram CUDA kernel
UNITORCH_EXTENSIONS=NGRAM pip install .

Environment Variables¤

Variable	Default	Purpose
`UNITORCH_CACHE`	`~/.cache/unitorch`	Model/data cache (also sets `HF_HOME`, `TRANSFORMERS_CACHE`)
`UNITORCH_TEMP`	`/tmp/unitorch`	Temporary files
`UNITORCH_HOME`	`~/.unitorch`	Home directory
`UNITORCH_DEBUG`	`INFO`	Log level: `OFF`, `INFO`, `DETAIL`, `CPU`, `ALL`
`UNITORCH_EXTENSIONS`	(unset)	Set to `NGRAM` to build CUDA extensions