---
id: metaflow-torchrun
version: "0.2.2"
license: unclear
license_treatment: unclear
maintenance: active
---

# metaflow-torchrun — A torchrun decorator for Metaflow

License: unclear · Maintenance: active · Downloads: 82.5K/mo

## What it is and what it does

metaflow-torchrun is a plugin that bridges Metaflow workflows and PyTorch's torchrun distributed training framework. It provides a @torchrun decorator that lets you run existing PyTorch distributed programs (like DDP training) as parallel steps in a Metaflow DAG without modifying your training code. The decorator automatically handles torchrun argument selection based on Metaflow compute resource requests (CPU, GPU, memory), network discovery, and subprocess orchestration.

You use it by stacking the @torchrun decorator on a Metaflow step, then calling current.torch.run() with your training entrypoint and arguments. This pattern is useful for multi-node training jobs submitted to AWS Batch or Kubernetes, where you want Metaflow to manage task orchestration and torchrun to manage the distributed torch processes within each task.

Use it for:

- Run multi-node PyTorch DDP training as a step in a Metaflow workflow without subprocess boilerplate
- Orchestrate distributed GPT or transformer training across parallel Metaflow tasks on Kubernetes or AWS Batch
- Execute existing torchrun scripts inside Metaflow without rewriting training code for the workflow framework
- Automate network discovery and torchrun argument configuration based on Metaflow resource decorators

## Worth the install?

AI-flagged interpretation of the facts above — verify before relying.

Provides a Metaflow decorator that integrates PyTorch distributed training via torchrun, allowing you to run multi-node torch programs as parallel Metaflow tasks without modifying your training code.

Yes, if you use Metaflow and PyTorch distributed training together. The package has no runtime dependencies, low install friction, active maintenance, and solves a real integration gap. License treatment is unclear—verify Apache License compliance before production use. No known vulnerabilities.

## Install

pip install metaflow-torchrun
uv add metaflow-torchrun
poetry add metaflow-torchrun

## Installing metaflow-torchrun

Before you install: Installation is straightforward with no runtime dependencies. The package is actively maintained with a recent release, though Python version support is unspecified.

License in practice: License treatment is unclear; the description mentions Apache License but SPDX metadata is not present in the package record, so verify the actual license terms before use.

Quickstart:

pip install metaflow-torchrun

from metaflow import FlowSpec, step, torchrun

class MyFlow(FlowSpec):
    @kubernetes(gpu=1)
    @torchrun
    @step
    def train(self):
        current.torch.run(
            entrypoint="main.py",
            entrypoint_args={"main-arg-1": "123"},
            nproc_per_node=1
        )

Requires Metaflow to be installed and configured; GPU/compute environment must support the resource decorators used (e.g., @kubernetes or @batch).

Verify before relying:

- Exact Python version requirements and compatibility range
- Whether Metaflow itself is a required dependency or assumed to be pre-installed
- GPU/CUDA version requirements for torchrun integration

## Package facts

- License: not declared (unclear)
- Python support: unspecified
- Install friction: low
- Maintenance: active
- Downloads: 82.5K/month (top 15,000 on PyPI)
- Known vulnerabilities: none known

## Tags

metaflow pytorch distributed training, torchrun metaflow decorator, pytorch ddp metaflow integration, multi-node torch metaflow, distributed training workflow orchestration, pytorch elastic run metaflow, torch parallel tasks metaflow, pytorch-distributed, workflow-orchestration, multi-node-training

[View on SkillFed](https://skillfed.io/packages/metaflow-torchrun) · [View on PyPI](https://pypi.org/project/metaflow-torchrun/)