🐍 Python Library

👀 Overview

Launch and train across all your accelerators, using your favorite framework + backend combo.

ezpz simplifies the process of:

• Setting up + launching distributed training:

  • import ezpz as ez

    • RANK = ez.setup_torch(backend=backend) [for backend \(\in\) {DDP, deepspeed, horovod}]{.dim-text}

    • RANK = ez.get_rank()

    • LOCAL_RANK = ez.get_local_rank()

    • WORLD_SIZE = ez.get_world_size()

    [(see ezpz/dist.py for more details).]{.dim-text}

• Using your favorite framework:

  • framework=pytorch + backend={DDP, deepspeed, horovod}

  • framework=tensorflow + backend=horovod

  • ez.get_torch_device(): {cuda, xpu, mps, cpu}

  • ez.get_torch_backend(): {nccl, ccl, gloo}

  2ez 😎.

• Writing device agnostic code:

  • ezpz.get_torch_device()

    ```pycon
    >>> import ezpz as ez
    >>> DEVICE = ez.get_torch_device()
    >>> model = torch.nn.Linear(10, 10)
    >>> model.to(DEVICE)
    >>> x = torch.randn((10, 10), device=DEVICE)
    >>> y = model(x)
    >>> y.device
    device(type='mps', index=0)
    ```
    

• Using wandb:

  • ez.setup_wandb(project_name='ezpz')

• Full support for any {device + framework + backend}:
• device: {GPU, XPU, MPS, CPU}
• framework: {torch, deepspeed, horovod, tensorflow}
• backend: {DDP, deepspeed, horovod}

Install

To install¹:

python3 -m pip install -e "git+https://github.com/saforem2/ezpz#egg=ezpz" --require-virtualenv

📂 Directory Structure

• 📂 ezpz / src / ezpz/
• 📂 bin/:
  • utils.sh: Shell utilities for ezpz
• 📂 conf/:
  • ⚙️ config.yaml: Default TrainConfig object
  • ⚙️ ds_config.json: DeepSpeed configuration
• 📂 log/: Logging configuration.
• 🐍 __about__.py: Version information
• 🐍 __init__.py: Main module
• 🐍 __main__.py: Entry point
• 🐍 configs.py: Configuration module
• 🐍cria.py: Baby Llama
• 🐍dist.py: Distributed training module
• 🐍history.py: History module
• 🐍jobs.py: Jobs module
• 🐍model.py: Model module
• 🐍plot.py: Plot modul
• 🐍profile.py: Profile module
• 🐍runtime.py: Runtime module
• 🐍test.py: Test module
• 🐍test_dist.py: Distributed training test module
• 🐍train.py: train module
• 🐍trainer.py: trainer module
• 🐍utils.py: utility module

📂 /ezpz/src/ezpz/
┣━━ 📂 bin/
┃   ┣━━ 📄 affinity.sh
┃   ┣━━ 📄 getjobenv
┃   ┣━━ 📄 savejobenv
┃   ┣━━ 📄 saveslurmenv
┃   ┣━━ 📄 setup.sh
┃   ┣━━ 📄 train.sh
┃   ┗━━ 📄 utils.sh
┣━━ 📂 conf/
┃   ┣━━ 📂 hydra/
┃   ┃   ┗━━ 📂 job_logging/
┃   ┃       ┣━━ ⚙️ colorlog1.yaml
┃   ┃       ┣━━ ⚙️ custom.yaml
┃   ┃       ┗━━ ⚙️ enrich.yaml
┃   ┣━━ 📂 logdir/
┃   ┃   ┗━━ ⚙️ default.yaml
┃   ┣━━ ⚙️ config.yaml
┃   ┣━━ 📄 ds_config.json
┃   ┗━━ ⚙️ ds_config.yaml
┣━━ 📂 log/
┃   ┣━━ 📂 conf/
┃   ┃   ┗━━ 📂 hydra/
┃   ┃       ┗━━ 📂 job_logging/
┃   ┃           ┗━━ ⚙️ enrich.yaml
┃   ┣━━ 🐍 __init__.py
┃   ┣━━ 🐍 __main__.py
┃   ┣━━ 🐍 config.py
┃   ┣━━ 🐍 console.py
┃   ┣━━ 🐍 handler.py
┃   ┣━━ 🐍 style.py
┃   ┣━━ 🐍 test.py
┃   ┗━━ 🐍 test_log.py
┣━━ 🐍 __about__.py
┣━━ 🐍 __init__.py
┣━━ 🐍 __main__.py
┣━━ 🐍 configs.py
┣━━ 🐍 cria.py
┣━━ 🐍 dist.py
┣━━ 🐍 history.py
┣━━ 🐍 jobs.py
┣━━ 🐍 loadjobenv.py
┣━━ 🐍 model.py
┣━━ 🐍 plot.py
┣━━ 🐍 profile.py
┣━━ 🐍 runtime.py
┣━━ 🐍 savejobenv.py
┣━━ 🐍 test.py
┣━━ 🐍 test_dist.py
┣━━ 🐍 train.py
┣━━ 🐍 trainer.py
┗━━ 🐍 utils.py

---

1. Note the --require-virtualenv isn’t strictly required, but I highly recommend to always try and work within a virtual environment, when possible. ↩