diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..7410066 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,30 @@ +name: Tests + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + test: + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.10"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e ".[dev]" + + - name: Run CPU tests + run: python -m pytest -m "not gpu" -v test diff --git a/.gitignore b/.gitignore index 9235f7a..eac13f2 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,7 @@ __pycache__/ .venv piper.egg-info/ -out/ \ No newline at end of file +out/ +.vscode/ +.codex +ec2-trust-policy.json diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..c2195fe --- /dev/null +++ b/Dockerfile @@ -0,0 +1,25 @@ +FROM rayproject/ray:2.44.1-py310-cu128 + +USER root + +RUN apt-get update && \ + apt-get install -y --no-install-recommends ca-certificates wget && \ + . /etc/os-release && \ + arch="$(dpkg --print-architecture)" && \ + case "$arch" in \ + amd64) cuda_arch="x86_64" ;; \ + arm64) cuda_arch="sbsa" ;; \ + *) echo "Unsupported architecture: $arch" >&2; exit 1 ;; \ + esac && \ + cuda_repo="https://developer.download.nvidia.com/compute/cuda/repos/${ID}${VERSION_ID//./}/${cuda_arch}" && \ + wget -q "${cuda_repo}/cuda-keyring_1.1-1_all.deb" -O /tmp/cuda-keyring.deb && \ + dpkg -i /tmp/cuda-keyring.deb && \ + apt-get update && \ + apt-get install -y --no-install-recommends cuda-nsight-systems-12-8 && \ + rm -rf /var/lib/apt/lists/* /tmp/cuda-keyring.deb + +USER ray + +COPY requirements.txt /tmp/requirements.txt +RUN pip install --no-cache-dir -r /tmp/requirements.txt \ + --extra-index-url https://download.pytorch.org/whl/cu128 diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..bd82f15 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Megan Frisella + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/README.md b/README.md index be7fe34..8460243 100644 --- a/README.md +++ b/README.md @@ -1,53 +1,203 @@ # Piper -Piper is a PyTorch library for training large models with flexible pipeline parallel schedules. - -## Environment setup: conda -We assume a Linux-based environment - -1. Create a conda environment with `python==3.10` -2. Install the requirements in `requirements.txt` -3. Modify PyTorch and Ray dependencies according to the instructions below - -## Modifying Ray dependency - -**Ray** - -Tensor transport backends currently only support 1 return value per task. -- WIP: Upstream this into Ray. -- Modifications (2): Comment out the [assertion in ActorMethod._remote()](https://github.com/ray-project/ray/blob/b70d990db786a1f2259dec0504acccf2590353f3/python/ray/actor.py#L824-L828) and add logic for [handling multiple return values with a GPU object manager](https://github.com/ray-project/ray/blob/b70d990db786a1f2259dec0504acccf2590353f3/python/ray/actor.py#L880-L887). -``` -####### PIPER MODIFICATION START ####### -# if num_returns != 1: -# raise ValueError( -# f"Currently, methods with tensor_transport={tensor_transport.name} only support 1 return value. " -# "Please make sure the actor method is decorated with `@ray.method(num_returns=1)` (the default)." -# ) -####### PIPER MODIFICATION END ####### -``` -``` -####### PIPER MODIFICATION START ####### -gpu_object_manager = ray._private.worker.global_worker.gpu_object_manager -if isinstance(object_refs, ObjectRef): - object_ref = object_refs - gpu_object_manager.add_gpu_object_ref( - object_ref, self._actor, tensor_transport - ) -else: - for object_ref in object_refs: - assert isinstance(object_ref, ObjectRef) - gpu_object_manager.add_gpu_object_ref( - object_ref, self._actor, tensor_transport - ) -####### PIPER MODIFICATION END ####### -``` - -## Training Llama in Piper -The llama test program `test/test_llama.py` supports GPipe, 1F1B and interleaved 1F1B schedules for 2 or 4 devices. -The `test/models/llama.py` file has example `forward` methods for one stage, two stage, and four stage partitions. -Ensure that the correct `forward` method is uncommented for the desired schedule (e.g. two stage for 1F1B on 2 devices, four stage for interleaved 1F1B on two devices). -DP training can also be turned on with the `dp_degree` flag. -Run the Llama test program for the 1F1B schedule for two devices: -``` -python3 -m test.test_llama --model LLAMA_DEBUG --schedule 1f1b --num_stages 2 --pp_degree 2 --dp_degree 1 +[![arXiv](https://img.shields.io/badge/arXiv-TODO-b31b1b.svg)](TODO) + +New distributed training strategies should not require new distributed runtimes; Piper gives PyTorch users direct control over model placement and GPU scheduling with lightweight annotations and a small scheduling language. + +## Updates + +* 2026-06 - Blog post: [User-Controlled Distributed Training for PyTorch](TODO). +* 2025-10 - Paper released on arXiv: [Piper: Towards Flexible Pipeline Parallelism for PyTorch](TODO). + +## Introduction + +Large training jobs increasingly combine multiple parallelism strategies such as pipeline, data, and expert parallelism with ZeRO-style sharding, creating placement and GPU scheduling choices that current frameworks cannot express cleanly. +Today, ML researchers and practitioners choose between one-off specialized systems that perform well but are hard to adapt, and general-purpose frameworks that are easier to use but expose limited control. + +Piper is a user-controllable distributed training system for PyTorch that separates model placement and GPU scheduling from model code and runtime implementation. +With lightweight model annotations and a small scheduling language, Piper lets users express, visualize, profile, and run high-performance training schedules such as DualPipe-style pipeline- and expert-parallel overlap. + +## Architecture + +![Piper architecture](figs/architecture.jpg) + +Piper has two user-facing inputs: + +* An annotated PyTorch model: standard model code with lightweight tags for schedulable regions such as pipeline stages and MoE experts. +* A schedule-directive program: instructions that tell the Piper compiler how to shard, replicate, order, and overlap those schedulable regions. + +The compiler traces the model with TorchDynamo, splits the graph by Piper annotations, builds a distributed training DAG IR, and applies DAG rewrites according to the schedule directives. +The directive rewrites insert point-to-point pipeline communication, DP collectives, ZeRO gather/scatter collectives, EP all-to-all communication, temporal edges, device assignments, and logical stream assignments. + +The runtime decomposes the global DAG into per-device execution plans and runs them on Ray actors. +Each actor manages local CUDA streams, process groups/communicators, model-state buffers, and intermediate tensors. + +## Installation + +Requires Python 3.10+ on Linux with CUDA GPUs. +The current setup has been tested with Python 3.10 and CUDA 12.x. + +```bash +python3.10 -m venv .venv +source .venv/bin/activate +python -m pip install --upgrade pip +python -m pip install -e . +``` + +If you do not already have Python 3.10 locally: + +```bash +conda create -n piper python=3.10 -y +conda activate piper +python -m pip install --upgrade pip +python -m pip install -e . +``` + +## Quickstart + +From the repository root, run the Qwen example with a DualPipeV schedule: + +```bash +python examples/test_harness.py \ + --test-file examples/test_qwen.py \ + --base-schedule examples/base-schedules/pp4_dp2_ep2_v_placement.json \ + --schedule dualpipev \ + --ranks 2 \ + --mbs 4 +``` + +The base schedule `pp4_dp2_ep2_v_placement.json` describes a PP x DP x EP placement with four annotated pipeline regions mapped onto two physical pipeline ranks in a V layout. +Regions `PP=0` and `PP=3` run on device group `[0, 2]`, while `PP=1` and `PP=2` run on device group `[1, 3]`. +Within each device group, Piper replicates non-expert regions for DP and shards expert regions for EP. +This example expects four visible CUDA devices. +The harness appends DualPipeV `split` and `order` directives for two physical pipeline ranks and four microbatches. + +Run the Qwen example with PP x ZeRO-3 1F1B schedule: + +```bash +python examples/test_harness.py \ + --test-file examples/test_qwen.py \ + --base-schedule examples/base-schedules/pp2_dp2_ep2_zero3.json \ + --schedule 1f1b \ + --ranks 2 \ + --mbs 4 +``` + +The base schedule `pp2_dp2_ep2_zero3.json` describes a PP x DP x EP placement with two pipeline stages, DP degree two, EP sharding, and ZeRO-3-style gradient and parameter sharding. +Stage `PP=0` runs on device group `[0, 2]`, while `PP=1` runs on device group `[1, 3]`. +This example also expects four visible CUDA devices. +The harness appends 1F1B `split` and `order` directives for two pipeline ranks and four microbatches. +Each run creates `out//` with the complete generated schedule and metrics. + +## Overview of Inputs + +### Annotations + +`piper.annotate(tag)` is a context manager that attaches Piper metadata to all PyTorch operations traced inside the scope. +The `tag` is a non-empty string naming a schedulable dimension of the model, such as `PP` for pipeline regions or `EP` for expert regions. +For each tag name, Piper assigns integer indices in trace order, so repeated `piper.annotate("PP")` scopes become `PP=0`, `PP=1`, and so on. + +The schedule directives program uses these tag names and indices in filters to select regions of the traced model. +For example, a filter can select one concrete region by its tag index, all regions with a tag, or regions that match a combination of tags. +See the blog walkthrough section on [annotating a Qwen3 MoE model](TODO) for an example and further details. + +### Schedule Directives Program + +A schedule directive program is a JSON array of directive objects. +Each directive has an `op` string naming the rewrite to apply to the distributed training DAG IR. +Most directives use a `filter` object to select the model region to apply the directive to. +The `order` directive uses `filters` to describe a sequence of filter groups. +A `filter` is a JSON object whose keys are tag names and whose values are the tag indices to match. + +Filter keys can refer to annotation tags, new tags added by previous directives, or the compiler-provided `PASS` tag which supports `F` (forward), `B` (backward), `BI` (backward for inputs), and `BW` (backward for weights) for different training step stages. +Filter values can be tag indices or the special value `"*"` to match any concrete index for the key. + +The empty filter `{}` matches the entire DAG IR. +All key/value pairs in a filter are conjunctive, so `{"PP": 1, "EP": "*"}` matches nodes that are in pipeline region `PP=1` and have any `EP` index. + +Supported directives: + +* `place` assigns matched compute regions to device groups and names the stream used for inserted PP send/recv communication. + * `op`: `"place"` + * `filter`: Selects the model regions to place. + * `devices`: Non-empty list of CUDA device IDs for the placement group. + * `stream`: Optional logical stream name for inserted point-to-point communication. +* `replicate` replicates matched regions across devices and inserts DP synchronization. + * `op`: `"replicate"`. + * `filter`: Selects the model regions to replicate. + * `devices`: Non-empty list of CUDA device IDs across which the region is replicated. + * `reduce_stream`: Optional logical stream name for gradient reduction collective communication. + * `gather_stream`: Optional logical stream name for ZeRO-3 parameter materialization collective communication. + * `bucket_size`: Optional parameter bucket size in MB for finer-grained synchronization. + * `shard_grads`: Optional boolean enabling ZeRO-2-style gradient sharding. + * `shard_params`: Optional boolean enabling ZeRO-3-style parameter sharding. +* `shard` shards matched regions across devices and inserts all-to-all communication, typically for expert regions. + * `op`: `"shard"` + * `filter`: Selects the model regions to shard. + * `devices`: Non-empty list of CUDA device IDs across which the region is sharded. + * `stream`: Optional logical stream name for inserted all-to-all collective communication. +* `split` duplicates the matched DAG by a named microbatch dimension. + * `op`: `"split"` + * `filter`: Selects the sub-DAG to duplicate. + * `dim_name`: Non-empty string naming the new split dimension. + * `num_microbatches`: Positive integer number of copies to create. +* `order` adds temporal dependencies between filter groups. + * `op`: `"order"` + * `filters`: List of at least two non-empty filter groups. + * Each filter group is a list of filter objects that occupy the same ordering slot. + * Consecutive filter groups create temporal dependencies from one slot to the next. + * Multiple filters in the same group permit Piper to interleave those sub-DAGs. + +In practice, base schedules under `examples/base-schedules/` describe model placement and composed parallelism choices, while `examples/test_harness.py` appends generated `split` and `order` directives for schedule families such as `1f1b`, `interleaved_1f1b`, `zerobubble`, and `dualpipev`. +For more detail, see the blog walkthrough sections on [PP x DP x EP placement](https://github.com/uw-syfi/uw-syfi.github.io/blob/piper-blog/_posts/2026-06-05-piper.md#scheduling-dualpipe-like-pp-x-dp-x-ep-model-placement), [DualPipe-like pipeline scheduling](https://github.com/uw-syfi/uw-syfi.github.io/blob/piper-blog/_posts/2026-06-05-piper.md#scheduling-a-dualpipe-like-pipeline-schedule), and [schedule builders](https://github.com/uw-syfi/uw-syfi.github.io/blob/piper-blog/_posts/2026-06-05-piper.md#generating-directives-with-schedule-builders). + +## Overview of Outputs + +Every harness run creates a timestamped directory under `out/`: + +```text +out// +|-- _pp_mbs.json +`-- results.csv +``` + +`results.csv` contains a row for each SPMD rank (e.g., each DP rank) reporting the mean/std iteration time, training throughput in tokens/sec, and per-PP-rank peak memory in GB. + +Optional artifact generation flags: + +* `--viz` renders the generated pipeline schedule and per-PP-rank DAG IRs under the run directory. +* `--pytorch-profiler --pytorch-profiler-iters ` runs extra profiled iterations and writes combined Chrome trace files per SPMD rank; Piper annotates GPU events with DAG IR node labels. + +## Development + +The GitHub Actions workflow runs the CPU-only pytest suite on pushes and pull requests: + +```bash +python -m pytest -m "not gpu" -v test +``` + +There are currently no pytest tests marked `gpu` under `test/`. +To run the full pytest suite, including any future GPU-marked tests, use: + +```bash +python -m pytest -v test +``` + +For GPU end-to-end validation, run the quickstart examples above on a machine with four visible CUDA devices. +To run all Qwen example schedules, use `examples/run_qwen_examples.py`. + +## Citation + +If you use Piper in your research, please cite: + +```bibtex +@inproceedings{frisella2025piper, + title = {Piper: Towards Flexible Pipeline Parallelism for PyTorch}, + author = {Frisella, Megan and Oentoro, Arvin and Gao, Xiangyu and Bernstein, Gilbert and Wang, Stephanie}, + booktitle = {Proceedings of the 4th Workshop on Practical Adoption Challenges of ML for Systems}, + year = {2025}, + publisher = {Association for Computing Machinery}, + doi = {10.1145/3766882.3767187}, + url = {https://doi.org/10.1145/3766882.3767187} +} ``` diff --git a/artifact/Dockerfile b/artifact/Dockerfile new file mode 100644 index 0000000..9fd13ba --- /dev/null +++ b/artifact/Dockerfile @@ -0,0 +1,104 @@ +FROM rayproject/ray:2.44.1-py310-cu128 + +SHELL ["/bin/bash", "-lc"] +USER root +WORKDIR /workspace + +ARG TORCHTITAN_COMMIT=b01adfb544b4331ecab090ebdb50b2296cd8eb6a +ARG MEGATRON_REF=26.04-alpha.rc2 +ARG TORCHTITAN_TORCH_INDEX_URL=https://download.pytorch.org/whl/nightly/cu128 +ARG APEX_REF=master + +ENV DEBIAN_FRONTEND=noninteractive +ENV TORCHTITAN_CONDA_ENV=torchtitan +ENV MEGATRON_CONDA_ENV=megatron +ENV DEEPSPEED_CONDA_ENV=deepspeed +ENV PIPER_CONDA_ENV=piper +ENV MEGATRON_WORKSPACE=/workspace/Megatron-LM +ENV TORCHTITAN_WORKSPACE=/workspace/torchtitan +ENV PIPER_WORKSPACE=/workspace/piper +ENV ARTIFACT_WORKSPACE=/workspace/artifact + +RUN apt-get update && \ + apt-get install -y --no-install-recommends ca-certificates git wget build-essential ninja-build && \ + . /etc/os-release && \ + arch="$(dpkg --print-architecture)" && \ + case "$arch" in \ + amd64) cuda_arch="x86_64" ;; \ + arm64) cuda_arch="sbsa" ;; \ + *) echo "Unsupported architecture: $arch" >&2; exit 1 ;; \ + esac && \ + cuda_repo="https://developer.download.nvidia.com/compute/cuda/repos/${ID}${VERSION_ID//./}/${cuda_arch}" && \ + wget -q "${cuda_repo}/cuda-keyring_1.1-1_all.deb" -O /tmp/cuda-keyring.deb && \ + dpkg -i /tmp/cuda-keyring.deb && \ + apt-get update && \ + apt-get install -y --no-install-recommends \ + cuda-nsight-systems-12-8 \ + cuda-nvcc-12-8 \ + cuda-cudart-dev-12-8 && \ + rm -rf /var/lib/apt/lists/* /tmp/cuda-keyring.deb + +RUN mkdir -p /workspace/artifact/patches +COPY artifact/patches /workspace/artifact/patches + +RUN git clone https://github.com/pytorch/torchtitan.git /workspace/torchtitan && \ + cd /workspace/torchtitan && \ + git checkout "${TORCHTITAN_COMMIT}" && \ + git apply /workspace/artifact/patches/torchtitan-qwen-e2e.patch + +RUN git clone https://github.com/NVIDIA/Megatron-LM.git /workspace/Megatron-LM && \ + cd /workspace/Megatron-LM && \ + git checkout "${MEGATRON_REF}" + +RUN git clone https://github.com/NVIDIA/apex.git /workspace/apex && \ + cd /workspace/apex && \ + git checkout "${APEX_REF}" + +RUN conda create -y -n torchtitan python=3.10 && \ + conda run -n torchtitan pip install --no-cache-dir --pre torch --index-url "${TORCHTITAN_TORCH_INDEX_URL}" && \ + conda run -n torchtitan pip install --no-cache-dir smart_open requests && \ + conda run -n torchtitan pip install --no-cache-dir --pre --no-deps torchdata --index-url https://download.pytorch.org/whl/nightly/cpu && \ + conda run -n torchtitan pip install --no-cache-dir -e /workspace/torchtitan && \ + conda run -n torchtitan pip install --no-cache-dir --upgrade fsspec && \ + conda clean -afy + +RUN conda create -y -n megatron python=3.12 && \ + conda run -n megatron pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu128 && \ + conda run -n megatron pip install --no-cache-dir \ + numpy packaging pandas pydantic pyyaml einops importlib-metadata \ + nvdlfw-inspect onnx onnxscript tensorboard && \ + conda run -n megatron pip install --no-cache-dir --no-deps \ + transformer-engine==2.13.0 \ + transformer-engine-torch==2.13.0 \ + transformer-engine-cu12==2.13.0 && \ + cd /workspace/apex && conda run -n megatron pip install --no-cache-dir --no-build-isolation \ + --config-settings "--build-option=--cpp_ext" \ + --config-settings "--build-option=--cuda_ext" \ + . && \ + conda run -n megatron pip install --no-cache-dir -e /workspace/Megatron-LM --no-deps && \ + conda clean -afy + +RUN conda create -y -n deepspeed python=3.10 && \ + conda run -n deepspeed pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu128 && \ + conda run -n deepspeed pip install --no-cache-dir deepspeed && \ + conda clean -afy + +COPY pyproject.toml requirements.txt README.md /workspace/piper/ +COPY src /workspace/piper/src +COPY examples /workspace/piper/examples +COPY test /workspace/piper/test + +RUN conda create -y -n piper python=3.10 && \ + conda run -n piper pip install --no-cache-dir torch==2.10.0 --index-url https://download.pytorch.org/whl/cu128 && \ + cd /workspace/piper && conda run -n piper pip install --no-cache-dir -e . \ + --extra-index-url https://download.pytorch.org/whl/cu128 && \ + conda clean -afy + +COPY artifact/e2e_eval.py artifact/run_deepspeed.py artifact/run_megatron.py artifact/sitecustomize.py /workspace/artifact/ +COPY artifact/scripts /workspace/artifact/scripts +COPY artifact/README.md /workspace/artifact/README.md + +RUN chmod +x /workspace/artifact/scripts/*.sh && \ + mkdir -p /workspace/eval-out /tmp/piper/ray_tmp + +WORKDIR /workspace/piper diff --git a/artifact/README.md b/artifact/README.md new file mode 100644 index 0000000..4fe4d4d --- /dev/null +++ b/artifact/README.md @@ -0,0 +1,103 @@ +# Piper Qwen E2E Artifact + +This directory is the reproducible entrypoint for the paper e2e evals across +Piper, TorchTitan, Megatron, and DeepSpeed. + +## Build + +From the piper repository root: + +```bash +docker build -f artifact/Dockerfile -t piper-e2e-artifact:latest . +``` + +The image pins TorchTitan to upstream commit +`b01adfb544b4331ecab090ebdb50b2296cd8eb6a` and applies +`artifact/patches/torchtitan-qwen-e2e.patch`. Megatron defaults to the stable +release tag `26.04-alpha.rc2`; pass `--build-arg MEGATRON_REF=` +to override it. + +## Local Single-Node Runs + +Run the smallest smoke eval on local GPUs: + +```bash +python artifact/e2e_eval.py \ + --backend local \ + --systems torchtitan piper \ + --sweeps local \ + --image piper-e2e-artifact:latest +``` + +Run the full default matrix locally: + +```bash +python artifact/e2e_eval.py --backend local --image piper-e2e-artifact:latest +``` + +Local mode uses one Docker container per experiment and maps +`artifact/out/e2e-eval/` into `/workspace/eval-out` inside the +container. For local mode, `nnode=1` and `ngpu=pp*dp`. If +`/m-coriander/coriander/mfris/torchtitan/assets/hf` exists, it is mounted +read-only into `/workspace/torchtitan/assets/hf` so TorchTitan can find the +downloaded Qwen tokenizer assets. Override this with +`--local-torchtitan-assets-path`. + +## AWS Existing-Node Runs + +AWS mode targets an existing EC2 cluster over SSH. Required environment: + +```bash +export SSH_KEY=/path/to/key.pem +export HEAD_PUBLIC_IP= +export HEAD_PRIVATE_IP= +export WORKER1_PRIVATE_IP= +export WORKER2_PRIVATE_IP= +export WORKER3_PRIVATE_IP= +``` + +Start containers automatically if the image already exists on every node: + +```bash +python artifact/e2e_eval.py \ + --backend aws \ + --aws-start-containers \ + --image piper-e2e-artifact:latest \ + --container piper_artifact +``` + +If containers are already running, omit `--aws-start-containers`. AWS mode uses +`nnode=dp` and `ngpu=pp`, matching the original EC2 experiment layout. + +## Useful Options + +```bash +python artifact/e2e_eval.py --dry-run +python artifact/e2e_eval.py --backend local --sweeps local --local-cuda-visible-devices 0 +python artifact/e2e_eval.py --backend local --sweeps local --local-torchtitan-assets-path /path/to/assets/hf +python artifact/e2e_eval.py --backend local --sweeps local --no-torchtitan-use-bmm-experts +python artifact/e2e_eval.py --backend local --sweeps local --piper-use-inductor --torchtitan-compile +python artifact/e2e_eval.py --systems piper torchtitan --sweeps schedule +python artifact/e2e_eval.py --sweeps zero --bucket-size-mb 25 +python artifact/e2e_eval.py --backend aws --nsight --systems piper +``` + +Outputs are written under `artifact/out/e2e-eval//`: + +- `results.csv` +- `.png` for combined plots with successful rows +- `/*.log` + +## Notes + +- Local Docker runs bind-mount this `artifact/` directory plus Piper `src/`, + `examples/`, and `test/` into the container, so runner/script/source edits do + not require rebuilding the image. Rebuild only when dependencies or the + Dockerfile image contents change. +- The TorchTitan patch is intentionally limited to the Qwen e2e runtime needs: + Qwen3 MoE eval configs, mesh ordering, timing/memory logs, dataset retry + behavior, and small runtime compatibility fixes. +- The piper runner invokes the current tracked entrypoint: + `examples/test_harness.py --test-file examples/test_qwen.py`. +- Piper base schedules are generated per experiment, so the artifact is not + limited to the checked-in `examples/base-schedules/*.json` files. diff --git a/artifact/e2e_eval.py b/artifact/e2e_eval.py new file mode 100755 index 0000000..25dda2f --- /dev/null +++ b/artifact/e2e_eval.py @@ -0,0 +1,1122 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import csv +import json +import math +import os +import re +import shlex +import shutil +import signal +import subprocess +import sys +import tempfile +import time +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Iterable, Sequence + +ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") +TT_ITER_RE = re.compile(r"Final \d+ iter times.*?avg:\s*([\d.]+)\s*s,\s*std:\s*([\d.]+)\s*s") +TT_FINAL_BY_RANK_RE = re.compile( + r"\[rank(\d+)\].*?Final \d+ iter times.*?avg:\s*([\d.]+)\s*s,\s*std:\s*([\d.]+)\s*s" +) +TT_MEM_RE = re.compile(r"\[rank(\d+)\].*?memory:\s*([\d.]+)GiB") +DS_STEP_RE = re.compile(r"\[Step\s+\d+/\d+\].*?step_time=([\d.]+)s") +DS_MEM_RE = re.compile(r"\[rank(\d+)\]\s+peak_memory_allocated_gb=([\d.]+)\s+peak_memory_reserved_gb=([\d.]+)") +MG_ITER_RE = re.compile(r"elapsed time per iteration \(ms\):\s*([\d.]+)") +MG_MEM_RE = re.compile(r"\[Rank\s+(\d+)\].*?max allocated:\s*([\d.]+)\s*\|") +OOM_RE = re.compile(r"(?:cuda\s+out\s+of\s+memory|out\s+of\s+memory|std::bad_alloc|\boom\b)", re.IGNORECASE) + +DEFAULT_SYSTEMS = ("torchtitan", "megatron", "deepspeed", "piper") +DEFAULT_SCHEDULE_SWEEP = ("1f1b", "interleaved1f1b", "zerobubble", "dualpipe") +DEFAULT_IMAGE = "piper-e2e-artifact:latest" +DEFAULT_CONTAINER = "piper_artifact" +DEFAULT_LOCAL_TORCHTITAN_ASSETS = Path("/m-coriander/coriander/mfris/torchtitan/assets/hf") +CONTAINER_OUT = Path("/workspace/eval-out") +CONTAINER_ARTIFACT = Path("/workspace/artifact") +CONTAINER_TORCHTITAN_HF_ASSETS = Path("/workspace/torchtitan/assets/hf") + +RUNTIME_SCHEDULES = { + "1f1b": "1F1B", + "interleaved1f1b": "Interleaved1F1B", + "zerobubble": "InterleavedZeroBubble", + "dualpipe": "DualPipeV", +} +PIPER_SCHEDULES = { + "1f1b": "1f1b", + "interleaved1f1b": "interleaved_1f1b", + "zerobubble": "zerobubble", + "dualpipe": "dualpipev", +} +PIPER_ZERO_STAGES = {"zero0": 0, "zero1": 1, "zero2": 2, "zero3": 3} +SYSTEM_COLORS = { + "torchtitan": "#4C72B0", + "megatron": "#55A868", + "deepspeed": "#C44E52", + "piper": "#8172B3", +} + + +def remapped_cuda_visible_devices(host_devices: str) -> str: + devices = [device.strip() for device in host_devices.split(",") if device.strip()] + return ",".join(str(index) for index, _device in enumerate(devices)) or "0" + + +def docker_gpus_request(cuda_visible_devices: str | None) -> str: + if cuda_visible_devices is None: + return "all" + return f'"device={cuda_visible_devices}"' + + +@dataclass(frozen=True) +class Experiment: + system: str + sweep: str + config: str + pp: int + dp: int + ep: int + zero_level: str + schedule: str + mb_size: int + seq_len: int + gradient_accumulation: bool = True + bucket_size_mb: float | None = None + ar_a2a_same_stream: bool = False + overlap_chunks: bool = False + + +@dataclass +class Result: + system: str + sweep: str + config: str + pp: int + dp: int + ep: int + zero_level: str + schedule: str + mb_size: int + seq_len: int + local_batch_size: int + global_batch_size: int + nnode: int + ngpu: int + log_path: str + metrics_path: str + returncode: int + status: str + iter_time_mean: float | None + iter_time_stddev: float | None + peak_memory_gb_by_rank: str + failure_reason: str + + +def normalize_schedule(name: str) -> str: + mapping = { + "1f1b": "1f1b", + "interleaved1f1b": "interleaved1f1b", + "interleaved-1f1b": "interleaved1f1b", + "interleaved_1f1b": "interleaved1f1b", + "zerobubble": "zerobubble", + "zero-bubble": "zerobubble", + "interleavedzerobubble": "zerobubble", + "interleaved-zero-bubble": "zerobubble", + "dualpipe": "dualpipe", + "dualpipev": "dualpipe", + } + key = name.strip().lower() + if key not in mapping: + raise ValueError(f"Unsupported schedule: {name}") + return mapping[key] + + +def build_experiments( + *, + systems: Iterable[str], + schedule_sweep: Iterable[str], + seq_len: int, + enabled_sweeps: set[str], + gradient_accumulation: bool, + ar_a2a_same_stream: bool, + overlap_chunks: bool, + bucket_size_mb: float | None, +) -> list[Experiment]: + experiments: list[Experiment] = [] + scalability_defaults = { + "config": "qwen3_9b", + "ep": 1, + "zero_level": "zero0", + "schedule": "1f1b", + "mb_size": 4, + } + zero_defaults = { + "config": "qwen3_1b", + "pp": 8, + "dp": 2, + "ep": 1, + "schedule": "1f1b", + } + zero_levels_by_system = { + "torchtitan": ("zero1", "zero2", "zero3"), + "megatron": ("zero1",), + "deepspeed": ("zero1",), + "piper": ("zero1", "zero2", "zero3"), + } + zero_sweep_values = (16, 24, 32, 34, 36, 38, 40) + schedule_defaults = { + "config": "qwen3_9b", + "pp": 2, + "dp": 2, + "ep": 2, + "zero_level": "zero1", + "mb_size": 4, + } + local_defaults = { + "config": "qwen3_1b", + "pp": 1, + "dp": 1, + "ep": 1, + "zero_level": "zero1", + "schedule": "1f1b", + "mb_size": 8, + } + supported_schedules_by_system = { + "torchtitan": {"1f1b", "interleaved1f1b", "dualpipe"}, + "megatron": {"1f1b", "interleaved1f1b"}, + "deepspeed": {"1f1b"}, + "piper": {"1f1b", "interleaved1f1b", "dualpipe"}, + } + + for system in systems: + if "scalability" in enabled_sweeps: + for pp in (4, 8): + for dp in (1, 2, 4): + experiments.append( + Experiment( + system=system, + sweep="scalability", + config=str(scalability_defaults["config"]), + pp=pp, + dp=dp, + ep=int(scalability_defaults["ep"]), + zero_level=str(scalability_defaults["zero_level"]), + schedule=str(scalability_defaults["schedule"]), + mb_size=int(scalability_defaults["mb_size"]), + seq_len=seq_len, + gradient_accumulation=gradient_accumulation, + bucket_size_mb=bucket_size_mb, + ar_a2a_same_stream=ar_a2a_same_stream, + overlap_chunks=overlap_chunks, + ) + ) + + if "zero" in enabled_sweeps: + for zero_level in zero_levels_by_system[system]: + for sweep_value in zero_sweep_values: + experiments.append( + Experiment( + system=system, + sweep="zero", + config=str(zero_defaults["config"]), + pp=int(zero_defaults["pp"]), + dp=int(zero_defaults["dp"]), + ep=int(zero_defaults["ep"]), + zero_level=zero_level, + schedule=str(zero_defaults["schedule"]), + mb_size=sweep_value, + seq_len=seq_len, + gradient_accumulation=gradient_accumulation, + bucket_size_mb=bucket_size_mb, + ar_a2a_same_stream=ar_a2a_same_stream, + overlap_chunks=overlap_chunks, + ) + ) + + if "schedule" in enabled_sweeps: + for schedule in schedule_sweep: + if schedule not in supported_schedules_by_system[system]: + continue + experiments.append( + Experiment( + system=system, + sweep="schedule", + config=str(schedule_defaults["config"]), + pp=int(schedule_defaults["pp"]), + dp=int(schedule_defaults["dp"]), + ep=int(schedule_defaults["ep"]), + zero_level=str(schedule_defaults["zero_level"]), + schedule=schedule, + mb_size=int(schedule_defaults["mb_size"]), + seq_len=seq_len, + gradient_accumulation=gradient_accumulation, + bucket_size_mb=bucket_size_mb, + ar_a2a_same_stream=ar_a2a_same_stream, + overlap_chunks=overlap_chunks, + ) + ) + + if "local" in enabled_sweeps: + experiments.append( + Experiment( + system=system, + sweep="local", + config=str(local_defaults["config"]), + pp=int(local_defaults["pp"]), + dp=int(local_defaults["dp"]), + ep=int(local_defaults["ep"]), + zero_level=str(local_defaults["zero_level"]), + schedule=str(local_defaults["schedule"]), + mb_size=int(local_defaults["mb_size"]), + seq_len=seq_len, + gradient_accumulation=gradient_accumulation, + bucket_size_mb=bucket_size_mb, + ar_a2a_same_stream=ar_a2a_same_stream, + overlap_chunks=overlap_chunks, + ) + ) + return experiments + + +def local_batch_size(exp: Experiment) -> int: + if exp.sweep == "local": + return exp.mb_size + return exp.pp * 2 * exp.mb_size + + +def global_batch_size(exp: Experiment) -> int: + return local_batch_size(exp) * exp.dp + + +def experiment_label(exp: Experiment) -> str: + if exp.sweep == "scalability": + return f"pp={exp.pp}, dp={exp.dp}" + if exp.sweep == "zero": + return f"zero={exp.zero_level}, mb={exp.mb_size}" + if exp.sweep == "local": + return "local" + return exp.schedule + + +def piper_model_name(config: str) -> str: + return {"qwen3_1b": "1B", "qwen3_9b": "9B"}[config] + + +def piper_num_mbs(exp: Experiment) -> int: + if exp.sweep == "local": + return 1 + return exp.pp * 2 + + +def slug(exp: Experiment, *, nsight: bool = False) -> str: + bucket_part = f"-bucket{exp.bucket_size_mb:g}" if exp.bucket_size_mb is not None else "" + nsight_part = "-nsight1" if nsight else "" + return ( + f"{exp.system}-{exp.sweep}-qwen{exp.config.removeprefix('qwen3_')}-" + f"sched_{exp.schedule}-pp{exp.pp}-dp{exp.dp}-ep{exp.ep}-" + f"{exp.zero_level}{bucket_part}{nsight_part}-bs{exp.mb_size}-" + f"sl{exp.seq_len}-mbs{piper_num_mbs(exp)}" + ) + + +def torchtitan_hf_assets_path(config: str) -> str | None: + return { + "qwen3_1b": "/workspace/torchtitan/assets/hf/Qwen3-0.6B", + "qwen3_9b": "/workspace/torchtitan/assets/hf/Qwen3-8B", + }.get(config) + + +def backend_layout(exp: Experiment, backend_name: str) -> tuple[int, int]: + if backend_name == "local": + return 1, exp.pp * exp.dp + return exp.dp, exp.pp + + +def in_container_command(exp: Experiment, args: argparse.Namespace, metrics_container_path: str | None = None) -> list[str]: + nnode = "{nnode}" + ngpu = "{ngpu}" + node_rank = "{node_rank}" + master_addr = "{master_addr}" + + if exp.system == "torchtitan": + nnode_value, ngpu_value = backend_layout(exp, args.backend) + if exp.ep > 1: + dp_replicate_degree = 1 + dp_shard_degree = exp.ep + else: + dp_replicate_degree = exp.dp if exp.zero_level == "zero1" else 1 + dp_shard_degree = 1 if exp.zero_level == "zero1" else exp.dp + tt_args = [ + "--parallelism.pipeline_parallel_degree", str(exp.pp), + "--parallelism.expert_parallel_degree", str(exp.ep), + "--parallelism.pipeline_parallel_schedule", RUNTIME_SCHEDULES[exp.schedule], + "--parallelism.pipeline_parallel_microbatch_size", str(exp.mb_size), + "--parallelism.data_parallel_replicate_degree", str(dp_replicate_degree), + "--parallelism.data_parallel_shard_degree", str(dp_shard_degree), + "--training.seq_len", str(exp.seq_len), + "--training.local_batch_size", str(local_batch_size(exp)), + "--training.global_batch_size", str(global_batch_size(exp)), + ] + hf_assets_path = torchtitan_hf_assets_path(exp.config) + if hf_assets_path is not None: + tt_args.extend(["--hf_assets_path", hf_assets_path]) + if exp.zero_level == "zero2": + tt_args.extend(["--parallelism.fsdp_reshard_after_forward", "never"]) + elif exp.zero_level == "zero3": + tt_args.extend(["--parallelism.fsdp_reshard_after_forward", "always"]) + if exp.schedule == "dualpipe" or not args.torchtitan_compile: + tt_args.append("--compile.no-enable") + command = [ + "/workspace/artifact/scripts/run_torchtitan.sh", + "--nnode", nnode, + "--ngpu", ngpu, + "--node-rank", node_rank, + "--master-addr", master_addr, + "--master-port", "29500", + "--module", "qwen3", + "--config", exp.config, + "--log-rank", ",".join(str(i) for i in range(min(nnode_value * ngpu_value, int(args.max_log_ranks)))), + ] + if args.nsight: + command.append("--nsight") + if args.torchtitan_use_bmm_experts: + command.append("--use-bmm-experts") + command.extend(["--", *tt_args]) + return command + + if exp.system == "megatron": + dp_megatron = exp.dp // exp.ep if exp.ep > 1 else exp.dp + command = [ + "/workspace/artifact/scripts/run_megatron.sh", + "--nnode", nnode, + "--ngpu", ngpu, + "--node-rank", node_rank, + "--master-addr", master_addr, + "--master-port", "29500", + "--model", exp.config, + "--pp", str(exp.pp), + "--dp", str(dp_megatron), + "--ep", str(exp.ep), + ] + if args.nsight: + command.append("--nsight") + command.extend([ + "--", + "--micro-bs", str(exp.mb_size), + "--global-bs", str(global_batch_size(exp)), + "--seq-length", str(exp.seq_len), + "--schedule", exp.schedule, + "--zero-level", exp.zero_level, + "--train-iters", str(args.baseline_steps), + ]) + return command + + if exp.system == "deepspeed": + return [ + "/workspace/artifact/scripts/run_deepspeed.sh", + "--nnode", nnode, + "--ngpu", ngpu, + "--node-rank", node_rank, + "--master-addr", master_addr, + "--master-port", "29501", + "--model", exp.config, + "--", + "--pp", str(exp.pp), + "--dp", str(exp.dp), + "--ep", str(exp.ep), + "--micro-bs", str(exp.mb_size), + "--global-bs", str(global_batch_size(exp)), + "--seq-len", str(exp.seq_len), + "--schedule", exp.schedule, + "--zero-stage", {"zero1": "1", "zero2": "2", "zero3": "3"}[exp.zero_level], + "--steps", str(args.baseline_steps), + ] + + command = [ + "/workspace/artifact/scripts/run_piper.sh", + "--model", piper_model_name(exp.config), + "--schedule", PIPER_SCHEDULES[exp.schedule], + "--pp", str(exp.pp), + "--dp", str(exp.dp), + "--zero-stage", str(PIPER_ZERO_STAGES[exp.zero_level]), + "--batch-size", str(exp.mb_size), + "--seq-len", str(exp.seq_len), + "--mbs", str(piper_num_mbs(exp)), + "--warmup", str(args.piper_warmup), + "--iters", str(args.piper_iters), + "--iteration-sleep", f"{args.piper_iteration_sleep:g}", + "--port", str(args.piper_ray_port), + "--temp-dir", "/tmp/piper/ray_tmp", + "--use-inductor" if args.piper_use_inductor and exp.sweep != "schedule" else "--no-use-inductor", + ] + if exp.ep > 1: + command.append("--ep") + if exp.bucket_size_mb is not None: + command.extend(["--bucket-size", f"{exp.bucket_size_mb:g}"]) + if args.nsight: + command.append("--nsight") + if metrics_container_path: + command.extend(["--metrics-out", metrics_container_path]) + if args.backend == "aws": + command.extend(["--address", "{master_addr}"]) + return command + + +def render_command(command: Sequence[str], *, nnode: int, ngpu: int, node_rank: int, master_addr: str) -> list[str]: + values = { + "nnode": str(nnode), + "ngpu": str(ngpu), + "node_rank": str(node_rank), + "master_addr": master_addr, + } + return [part.format(**values) for part in command] + + +class Backend: + name: str + + def layout(self, exp: Experiment) -> tuple[int, int]: + return backend_layout(exp, self.name) + + def run(self, exp: Experiment, command: Sequence[str], log_path: Path) -> int: + raise NotImplementedError + + def prepare_piper(self, args: argparse.Namespace) -> None: + return None + + def cleanup_after_experiment(self, exp: Experiment) -> None: + return None + + def fetch_file(self, container_path: str, destination: Path) -> bool: + return Path(container_path).is_file() + + def dry_run_lines(self, exp: Experiment, command: Sequence[str]) -> list[str]: + raise NotImplementedError + + +class LocalBackend(Backend): + name = "local" + + def __init__( + self, + image: str, + out_dir: Path, + extra_docker_args: Sequence[str], + cuda_visible_devices: str | None, + torchtitan_assets_path: Path | None, + ): + self.image = image + self.out_dir = out_dir + self.extra_docker_args = list(extra_docker_args) + self.cuda_visible_devices = cuda_visible_devices + self.container_cuda_visible_devices = ( + remapped_cuda_visible_devices(cuda_visible_devices) + if cuda_visible_devices is not None + else None + ) + self.artifact_src_dir = Path(__file__).resolve().parent + self.workspace_dir = self.artifact_src_dir.parent + self.torchtitan_assets_path = torchtitan_assets_path if torchtitan_assets_path and torchtitan_assets_path.is_dir() else None + + def run(self, exp: Experiment, command: Sequence[str], log_path: Path) -> int: + nnode, ngpu = self.layout(exp) + rendered = render_command(command, nnode=nnode, ngpu=ngpu, node_rank=0, master_addr="127.0.0.1") + gpus = docker_gpus_request(self.cuda_visible_devices) + docker_cmd = [ + "docker", "run", "--rm", + "--gpus", gpus, + "--network", "host", + "--ipc", "host", + "--shm-size", "32g", + "-v", f"{self.out_dir}:{CONTAINER_OUT}", + "-v", f"{self.artifact_src_dir}:{CONTAINER_ARTIFACT}:ro", + "-v", f"{self.workspace_dir / 'src'}:/workspace/piper/src:ro", + "-v", f"{self.workspace_dir / 'examples'}:/workspace/piper/examples:ro", + "-v", f"{self.workspace_dir / 'test'}:/workspace/piper/test:ro", + ] + if self.torchtitan_assets_path is not None: + docker_cmd.extend(["-v", f"{self.torchtitan_assets_path}:{CONTAINER_TORCHTITAN_HF_ASSETS}:ro"]) + if self.cuda_visible_devices is not None: + docker_cmd.extend([ + "-e", f"NVIDIA_VISIBLE_DEVICES={self.cuda_visible_devices}", + "-e", f"CUDA_VISIBLE_DEVICES={self.container_cuda_visible_devices}", + ]) + docker_cmd.extend([ + *self.extra_docker_args, + self.image, + "bash", "-lc", shlex.join(rendered), + ]) + log_path.parent.mkdir(parents=True, exist_ok=True) + with log_path.open("w", encoding="utf-8") as log_file: + proc = subprocess.run(docker_cmd, stdout=log_file, stderr=subprocess.STDOUT, text=True) + return proc.returncode + + def dry_run_lines(self, exp: Experiment, command: Sequence[str]) -> list[str]: + nnode, ngpu = self.layout(exp) + rendered = render_command(command, nnode=nnode, ngpu=ngpu, node_rank=0, master_addr="127.0.0.1") + cuda_env = ( + f"-e CUDA_VISIBLE_DEVICES={shlex.quote(str(self.container_cuda_visible_devices))} " + if self.cuda_visible_devices is not None + else "" + ) + extra_args = " ".join(shlex.quote(arg) for arg in self.extra_docker_args) + extra_args = f"{extra_args} " if extra_args else "" + gpus = docker_gpus_request(self.cuda_visible_devices) + nvidia_env = ( + f"-e NVIDIA_VISIBLE_DEVICES={shlex.quote(self.cuda_visible_devices)} " + if self.cuda_visible_devices is not None + else "" + ) + return [ + f"docker run --rm --gpus {shlex.quote(gpus)} --network host --ipc host " + f"{nvidia_env}" + f"{cuda_env}" + f"-v {shlex.quote(str(self.out_dir))}:{CONTAINER_OUT} " + f"-v {shlex.quote(str(self.artifact_src_dir))}:{CONTAINER_ARTIFACT}:ro " + f"-v {shlex.quote(str(self.workspace_dir / 'src'))}:/workspace/piper/src:ro " + f"-v {shlex.quote(str(self.workspace_dir / 'examples'))}:/workspace/piper/examples:ro " + f"-v {shlex.quote(str(self.workspace_dir / 'test'))}:/workspace/piper/test:ro " + f"{self._torchtitan_assets_mount_for_dry_run()}" + f"{extra_args}{shlex.quote(self.image)} bash -lc {shlex.quote(shlex.join(rendered))}" + ] + + def _torchtitan_assets_mount_for_dry_run(self) -> str: + if self.torchtitan_assets_path is None: + return "" + return f"-v {shlex.quote(str(self.torchtitan_assets_path))}:{CONTAINER_TORCHTITAN_HF_ASSETS}:ro " + + +class AwsSshBackend(Backend): + name = "aws" + + def __init__( + self, + *, + image: str, + container: str, + start_containers: bool, + ray_port: int, + ): + self.image = image + self.container = container + self.start_containers = start_containers + self.ray_port = ray_port + self.ssh_key = require_env("SSH_KEY") + self.head_public_ip = require_env("HEAD_PUBLIC_IP") + self.head_private_ip = require_env("HEAD_PRIVATE_IP") + self.workers = workers_from_env() + if self.start_containers: + self._start_containers() + + def _ssh(self, target: str, remote_command: str, *, worker: bool = False) -> list[str]: + cmd = [ + "ssh", + "-i", self.ssh_key, + "-o", "StrictHostKeyChecking=no", + ] + if worker: + proxy = ( + "ProxyCommand=" + f"ssh -i {self.ssh_key} -o StrictHostKeyChecking=no " + f"-W %h:%p ubuntu@{self.head_public_ip}" + ) + cmd.extend(["-o", proxy]) + cmd.extend([f"ubuntu@{target}", remote_command]) + return cmd + + def _node_specs(self, nnode: int) -> list[tuple[str, str, bool]]: + if nnode > 1 and len(self.workers) < nnode - 1: + raise RuntimeError(f"Need {nnode - 1} WORKER*_PRIVATE_IP values, found {len(self.workers)}") + nodes = [("head", self.head_public_ip, False)] + nodes.extend((f"worker{i}", ip, True) for i, ip in enumerate(self.workers[: nnode - 1], start=1)) + return nodes + + def _docker_exec_command(self, rendered: Sequence[str], *, node_rank: int) -> str: + inner = shlex.join(rendered) + docker = [ + "docker", "exec", + "-e", f"NODE_RANK={node_rank}", + "-e", "NCCL_SOCKET_IFNAME=ens32", + "-e", "GLOO_SOCKET_IFNAME=ens32", + self.container, + "bash", "-lc", inner, + ] + return shlex.join(docker) + + def _start_containers(self) -> None: + remote = ( + f"docker rm -f {shlex.quote(self.container)} >/dev/null 2>&1 || true; " + f"docker run -d --name {shlex.quote(self.container)} --gpus all " + "--network host --ipc host --shm-size 32g " + f"{shlex.quote(self.image)} sleep infinity" + ) + for label, target, worker in self._node_specs(len(self.workers) + 1): + print(f"[aws] starting container on {label}", file=sys.stderr) + subprocess.run(self._ssh(target, remote, worker=worker), check=True) + + def prepare_piper(self, args: argparse.Namespace) -> None: + nnode = max(1, len(self.workers) + 1) + stop = f"docker exec {shlex.quote(self.container)} bash -lc 'conda run -n piper ray stop -f >/dev/null 2>&1 || true'" + for _label, target, worker in self._node_specs(nnode): + subprocess.run(self._ssh(target, stop, worker=worker), check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + head_cmd = ( + f"docker exec {shlex.quote(self.container)} bash -lc " + + shlex.quote( + "conda run -n piper ray start --head " + f"--node-ip-address={self.head_private_ip} " + f"--port={self.ray_port} " + "--disable-usage-stats " + "--temp-dir=/tmp/piper/ray_tmp" + ) + ) + subprocess.run(self._ssh(self.head_public_ip, head_cmd), check=True) + for _label, target, worker in self._node_specs(nnode)[1:]: + worker_cmd = ( + f"docker exec {shlex.quote(self.container)} bash -lc " + + shlex.quote( + "conda run -n piper ray start " + f"--address={self.head_private_ip}:{self.ray_port} " + "--disable-usage-stats " + "--temp-dir=/tmp/piper/ray_tmp" + ) + ) + subprocess.run(self._ssh(target, worker_cmd, worker=worker), check=True) + + def run(self, exp: Experiment, command: Sequence[str], log_path: Path) -> int: + nnode, ngpu = self.layout(exp) + nodes = self._node_specs(nnode) + log_path.parent.mkdir(parents=True, exist_ok=True) + + if exp.system == "piper": + rendered = render_command( + command, + nnode=nnode, + ngpu=ngpu, + node_rank=0, + master_addr=self.head_private_ip, + ) + remote = self._docker_exec_command(rendered, node_rank=0) + with log_path.open("w", encoding="utf-8") as log_file: + proc = subprocess.run(self._ssh(self.head_public_ip, remote), stdout=log_file, stderr=subprocess.STDOUT, text=True) + return proc.returncode + + with tempfile.TemporaryDirectory(prefix="piper-e2e-node-logs.") as tmp_dir: + tmp_path = Path(tmp_dir) + procs: list[tuple[str, subprocess.Popen[str], Path]] = [] + for node_rank, (label, target, worker) in enumerate(nodes): + rendered = render_command( + command, + nnode=nnode, + ngpu=ngpu, + node_rank=node_rank, + master_addr=self.head_private_ip, + ) + remote = self._docker_exec_command(rendered, node_rank=node_rank) + node_log = tmp_path / f"{log_path.stem}.{label}.log" + handle = node_log.open("w", encoding="utf-8") + proc = subprocess.Popen(self._ssh(target, remote, worker=worker), stdout=handle, stderr=subprocess.STDOUT, text=True) + handle.close() + procs.append((label, proc, node_log)) + + returncodes = [] + for _label, proc, _node_log in procs: + returncodes.append(proc.wait()) + with log_path.open("w", encoding="utf-8") as combined: + for label, _proc, node_log in procs: + combined.write(f"===== {label} =====\n") + if node_log.exists(): + combined.write(node_log.read_text(encoding="utf-8", errors="replace")) + combined.write("\n") + return 0 if all(code == 0 for code in returncodes) else next(code for code in returncodes if code != 0) + + def fetch_file(self, container_path: str, destination: Path) -> bool: + remote = f"docker exec {shlex.quote(self.container)} bash -lc {shlex.quote('test -f ' + shlex.quote(container_path) + ' && cat ' + shlex.quote(container_path))}" + result = subprocess.run(self._ssh(self.head_public_ip, remote), check=False, capture_output=True, text=True) + if result.returncode != 0 or not result.stdout: + return False + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_text(result.stdout, encoding="utf-8") + return True + + def cleanup_after_experiment(self, exp: Experiment) -> None: + if exp.system != "piper": + return + stop = f"docker exec {shlex.quote(self.container)} bash -lc 'conda run -n piper ray stop -f >/dev/null 2>&1 || true'" + for _label, target, worker in self._node_specs(max(1, len(self.workers) + 1)): + subprocess.run(self._ssh(target, stop, worker=worker), check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) + + def dry_run_lines(self, exp: Experiment, command: Sequence[str]) -> list[str]: + nnode, ngpu = self.layout(exp) + nodes = self._node_specs(nnode) + lines = [] + if exp.system == "piper": + rendered = render_command(command, nnode=nnode, ngpu=ngpu, node_rank=0, master_addr=self.head_private_ip) + lines.append(self._docker_exec_command(rendered, node_rank=0)) + return lines + for node_rank, (label, _target, _worker) in enumerate(nodes): + rendered = render_command(command, nnode=nnode, ngpu=ngpu, node_rank=node_rank, master_addr=self.head_private_ip) + lines.append(f"{label}: {self._docker_exec_command(rendered, node_rank=node_rank)}") + return lines + + +def require_env(name: str) -> str: + value = os.environ.get(name) + if not value: + raise SystemExit(f"Missing required environment variable: {name}") + return value + + +def workers_from_env() -> list[str]: + workers: list[str] = [] + index = 1 + while True: + value = os.environ.get(f"WORKER{index}_PRIVATE_IP") + if not value: + break + workers.append(value) + index += 1 + return workers + + +def read_text_lossy(path: Path) -> str: + if not path.is_file(): + return "" + return path.read_text(encoding="utf-8", errors="replace").replace("\0", "") + + +def strip_ansi(text: str) -> str: + return ANSI_RE.sub("", text) + + +def parse_log(system: str, log_path: Path, metrics_path: Path | None = None) -> tuple[float | None, float | None, str, bool, str]: + text = strip_ansi(read_text_lossy(log_path)) + is_oom = OOM_RE.search(text) is not None + reason = "oom" if is_oom else "" + + if system == "torchtitan": + peak: dict[int, float] = {} + iter_by_rank: dict[int, tuple[float, float]] = {} + last_iter: tuple[float, float] | None = None + for line in text.splitlines(): + m = TT_FINAL_BY_RANK_RE.search(line) + if m: + iter_by_rank[int(m.group(1))] = (float(m.group(2)), float(m.group(3))) + last_iter = (float(m.group(2)), float(m.group(3))) + m = TT_ITER_RE.search(line) + if m: + last_iter = (float(m.group(1)), float(m.group(2))) + m = TT_MEM_RE.search(line) + if m: + peak[int(m.group(1))] = max(peak.get(int(m.group(1)), 0.0), float(m.group(2))) + iter_mean, iter_std = iter_by_rank.get(0, last_iter or (None, None)) + peak_str = "/".join(f"{peak[rank]:.2f}" for rank in sorted(peak)) if peak else "" + return iter_mean, iter_std, peak_str, is_oom, reason + + if system == "deepspeed": + step_times = [float(m.group(1)) for m in DS_STEP_RE.finditer(text)][-5:] + peak = {int(m.group(1)): float(m.group(2)) for m in DS_MEM_RE.finditer(text)} + peak_str = "/".join(f"{peak[rank]:.2f}" for rank in sorted(peak)) if peak else "" + if not step_times: + return None, None, peak_str, is_oom, reason + mean = sum(step_times) / len(step_times) + return mean, pstdev(step_times, mean), peak_str, is_oom, reason + + if system == "megatron": + iter_times_s = [float(m.group(1)) / 1000.0 for m in MG_ITER_RE.finditer(text)][-5:] + peak = {int(m.group(1)): float(m.group(2)) / 1024.0 for m in MG_MEM_RE.finditer(text)} + peak_str = "/".join(f"{peak[rank]:.2f}" for rank in sorted(peak)) if peak else "" + if not iter_times_s: + return None, None, peak_str, is_oom, reason + mean = sum(iter_times_s) / len(iter_times_s) + return mean, pstdev(iter_times_s, mean), peak_str, is_oom, reason + + if system == "piper" and metrics_path is not None and metrics_path.is_file(): + rows = list(csv.DictReader(metrics_path.open("r", encoding="utf-8"))) + times: list[tuple[float, float, int]] = [] + peak_values: list[str] = [] + for row in rows: + try: + mean = float(row.get("iter_time_mean_s") or "") + except ValueError: + continue + try: + std = float(row.get("iter_time_std_s") or 0.0) + except ValueError: + std = 0.0 + try: + samples = int(float(row.get("samples") or 1)) + except ValueError: + samples = 1 + times.append((mean, std, samples)) + for key, value in row.items(): + if key.startswith("peak_memory_pp") and value: + try: + peak_values.append(f"{key.removeprefix('peak_memory_pp').removesuffix('_gb')}:{float(value):.3f}") + except ValueError: + pass + total_samples = sum(samples for _mean, _std, samples in times) + if total_samples: + mean = sum(mean * samples for mean, _std, samples in times) / total_samples + std = sum(std * samples for _mean, std, samples in times) / total_samples + return mean, std, ";".join(peak_values), is_oom, reason + return None, None, "", is_oom, reason + + +def pstdev(values: Sequence[float], mean: float) -> float: + if not values: + return 0.0 + return math.sqrt(sum((value - mean) ** 2 for value in values) / len(values)) + + +def write_csv(path: Path, rows: list[Result]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + fieldnames = list(Result.__annotations__.keys()) + with path.open("w", newline="", encoding="utf-8") as handle: + writer = csv.DictWriter(handle, fieldnames=fieldnames) + writer.writeheader() + for row in rows: + writer.writerow(asdict(row)) + + +def row_tps(row: Result) -> float | None: + if row.iter_time_mean is None or row.iter_time_mean <= 0: + return None + return row.global_batch_size * row.seq_len / row.iter_time_mean + + +def result_label(row: Result) -> str: + if row.sweep == "scalability": + return f"pp={row.pp},dp={row.dp}" + if row.sweep == "zero": + return f"{row.zero_level},mb={row.mb_size}" + if row.sweep == "local": + return "local" + return row.schedule + + +def write_outputs(base_out: Path, rows: list[Result]) -> None: + write_csv(base_out / "results.csv", rows) + for sweep in ("scalability", "zero", "schedule", "local"): + sweep_rows = [row for row in rows if row.sweep == sweep] + if sweep_rows: + save_plot(sweep_rows, base_out / f"{sweep}.png", title=f"Qwen3 {sweep} throughput") + + +def save_plot(rows: list[Result], output_path: Path, *, title: str) -> None: + ok_rows = [row for row in rows if row.status == "ok" and row_tps(row) is not None] + if not ok_rows: + return + try: + import matplotlib + + matplotlib.use("Agg") + import matplotlib.pyplot as plt + except Exception as exc: # pragma: no cover - plotting is optional for dry environments + print(f"Skipping plot {output_path}: {exc}", file=sys.stderr) + return + + output_path.parent.mkdir(parents=True, exist_ok=True) + labels = [f"{row.system}\n{result_label(row)}" for row in ok_rows] + values = [float(row_tps(row) or 0.0) for row in ok_rows] + colors = [SYSTEM_COLORS.get(row.system, "#4C72B0") for row in ok_rows] + fig, ax = plt.subplots(figsize=(max(10, len(ok_rows) * 0.7), 5)) + ax.bar(range(len(ok_rows)), values, color=colors) + ax.set_title(title) + ax.set_ylabel("Tokens / second") + ax.set_xticks(range(len(ok_rows))) + ax.set_xticklabels(labels, rotation=45, ha="right") + ax.grid(axis="y", linestyle="--", alpha=0.35) + fig.tight_layout() + fig.savefig(output_path, dpi=200) + plt.close(fig) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run the Piper paper Qwen e2e artifact.") + parser.add_argument("--backend", choices=("local", "aws"), default="local") + parser.add_argument("--image", default=DEFAULT_IMAGE) + parser.add_argument("--container", default=DEFAULT_CONTAINER) + parser.add_argument("--aws-start-containers", action="store_true") + parser.add_argument("--systems", nargs="+", choices=list(DEFAULT_SYSTEMS), default=list(DEFAULT_SYSTEMS)) + parser.add_argument("--sweeps", nargs="+", choices=("scalability", "zero", "schedule", "local"), default=["scalability", "zero", "schedule"]) + parser.add_argument("--schedule-sweep", nargs="+", default=list(DEFAULT_SCHEDULE_SWEEP)) + parser.add_argument("--seq-len", type=int, default=512) + parser.add_argument("--out-dir", default=None) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--nsight", action="store_true") + parser.add_argument("--bucket-size-mb", type=float, default=None) + parser.add_argument("--gradient-accumulation", dest="gradient_accumulation", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--ar-a2a-same-stream", dest="ar_a2a_same_stream", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--overlap-chunks", dest="overlap_chunks", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--piper-warmup", type=int, default=3) + parser.add_argument("--piper-iters", type=int, default=10) + parser.add_argument("--piper-iteration-sleep", type=float, default=0.0) + parser.add_argument("--piper-use-inductor", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--torchtitan-compile", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument("--torchtitan-use-bmm-experts", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--piper-ray-port", type=int, default=6379) + parser.add_argument("--baseline-steps", type=int, default=8) + parser.add_argument("--max-log-ranks", type=int, default=8) + parser.add_argument( + "--local-cuda-visible-devices", + "--cuda-visible-devices", + dest="local_cuda_visible_devices", + default=None, + help=( + "Restrict local Docker runs to these host GPU IDs. Inside the filtered " + "container, CUDA_VISIBLE_DEVICES is remapped to logical IDs." + ), + ) + default_torchtitan_assets = os.environ.get("TORCHTITAN_HF_ASSETS") + if default_torchtitan_assets is None and DEFAULT_LOCAL_TORCHTITAN_ASSETS.is_dir(): + default_torchtitan_assets = str(DEFAULT_LOCAL_TORCHTITAN_ASSETS) + parser.add_argument( + "--local-torchtitan-assets-path", + default=default_torchtitan_assets, + help="Host path mounted read-only to /workspace/torchtitan/assets/hf for local TorchTitan runs.", + ) + parser.add_argument("--docker-arg", action="append", default=[], help="Extra argument passed to docker run in local mode.") + return parser.parse_args() + + +def make_backend(args: argparse.Namespace, base_out: Path) -> Backend: + if args.backend == "local": + torchtitan_assets_path = Path(args.local_torchtitan_assets_path).resolve() if args.local_torchtitan_assets_path else None + return LocalBackend( + args.image, + base_out, + args.docker_arg, + args.local_cuda_visible_devices, + torchtitan_assets_path, + ) + return AwsSshBackend( + image=args.image, + container=args.container, + start_containers=args.aws_start_containers, + ray_port=args.piper_ray_port, + ) + + +def main() -> int: + args = parse_args() + schedule_sweep = [normalize_schedule(item) for item in args.schedule_sweep] + base_out = Path(args.out_dir).resolve() if args.out_dir else (Path("artifact") / "out" / "e2e-eval" / time.strftime("%Y%m%d_%H%M%S")).resolve() + base_out.mkdir(parents=True, exist_ok=True) + backend = make_backend(args, base_out) + + experiments = build_experiments( + systems=args.systems, + schedule_sweep=schedule_sweep, + seq_len=args.seq_len, + enabled_sweeps=set(args.sweeps), + gradient_accumulation=args.gradient_accumulation, + ar_a2a_same_stream=args.ar_a2a_same_stream, + overlap_chunks=args.overlap_chunks, + bucket_size_mb=args.bucket_size_mb, + ) + grouped: dict[str, list[Experiment]] = {system: [] for system in args.systems} + for exp in experiments: + grouped[exp.system].append(exp) + + print(f"output directory: {base_out}") + print(f"planned experiments: {len(experiments)}") + overall_failed = False + + def _terminate(signum: int, _frame) -> None: + raise SystemExit(128 + signum) + + signal.signal(signal.SIGINT, _terminate) + signal.signal(signal.SIGTERM, _terminate) + + results: list[Result] = [] + for system in args.systems: + system_out = base_out / system + system_out.mkdir(parents=True, exist_ok=True) + system_experiments = grouped[system] + + for index, exp in enumerate(system_experiments, start=1): + exp_slug = slug(exp, nsight=args.nsight) + log_path = system_out / f"{index:02d}_{exp_slug}.log" + metrics_name = f".{exp_slug}.metrics.csv" + host_metrics_path = base_out / metrics_name + container_metrics_path = str(CONTAINER_OUT / metrics_name) + command = in_container_command(exp, args, container_metrics_path if system == "piper" else None) + nnode, ngpu = backend.layout(exp) + + print(f"[{system} {index}/{len(system_experiments)}] {exp.sweep}: {experiment_label(exp)} config={exp.config}") + for line in backend.dry_run_lines(exp, command): + print(" " + line) + if args.dry_run: + continue + + if system == "piper": + backend.prepare_piper(args) + returncode = 1 + try: + returncode = backend.run(exp, command, log_path) + if system == "piper" and not host_metrics_path.is_file(): + backend.fetch_file(container_metrics_path, host_metrics_path) + if system == "piper": + for leaked_json in base_out.glob("*.json"): + leaked_json.unlink() + iter_mean, iter_std, peak_str, is_oom, reason = parse_log( + system, + log_path, + metrics_path=host_metrics_path if system == "piper" else None, + ) + if returncode == 0 and iter_mean is not None: + status = "ok" + elif is_oom: + status = "oom" + reason = reason or "oom" + else: + status = "failed" + reason = reason or f"exit_code={returncode}" + if status != "ok": + overall_failed = True + results.append( + Result( + system=system, + sweep=exp.sweep, + config=exp.config, + pp=exp.pp, + dp=exp.dp, + ep=exp.ep, + zero_level=exp.zero_level, + schedule=exp.schedule, + mb_size=exp.mb_size, + seq_len=exp.seq_len, + local_batch_size=local_batch_size(exp), + global_batch_size=global_batch_size(exp), + nnode=nnode, + ngpu=ngpu, + log_path=str(log_path), + metrics_path="", + returncode=returncode, + status=status, + iter_time_mean=iter_mean, + iter_time_stddev=iter_std, + peak_memory_gb_by_rank=peak_str, + failure_reason=reason, + ) + ) + if host_metrics_path.exists(): + host_metrics_path.unlink() + write_csv(base_out / "results.csv", results) + finally: + if host_metrics_path.exists(): + host_metrics_path.unlink() + backend.cleanup_after_experiment(exp) + + if not args.dry_run: + write_outputs(base_out, results) + print(f"wrote results under {base_out}") + return 1 if overall_failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/artifact/patches/torchtitan-qwen-e2e.patch b/artifact/patches/torchtitan-qwen-e2e.patch new file mode 100644 index 0000000..042a6a7 --- /dev/null +++ b/artifact/patches/torchtitan-qwen-e2e.patch @@ -0,0 +1,718 @@ +diff --git a/torchtitan/config/manager.py b/torchtitan/config/manager.py +index 9c9ab6cf..335b388b 100644 +--- a/torchtitan/config/manager.py ++++ b/torchtitan/config/manager.py +@@ -102,18 +102,21 @@ class ConfigManager: + # Import config_registry from module based on module specification + if module_name in all_supported: + # short module from supported module list (search models first, then experiments) ++ first_exc = None + for prefix in ("torchtitan.models", "torchtitan.experiments"): + module_path = f"{prefix}.{module_name}.config_registry" + try: + module = importlib.import_module(module_path) + break +- except ImportError: ++ except ImportError as e: ++ if first_exc is None: ++ first_exc = e + continue + if module is None: + raise ImportError( + f"Cannot import config_registry for module '{module_name}' " + f"from torchtitan.models or torchtitan.experiments" +- ) ++ ) from first_exc + else: + # Fully qualified module path: try appending .config_registry first, + # then fall back to importing directly (e.g., torchtitan.models.llama3 +diff --git a/torchtitan/components/checkpoint.py b/torchtitan/components/checkpoint.py +index d31dcda1..1c6becc1 100644 +--- a/torchtitan/components/checkpoint.py ++++ b/torchtitan/components/checkpoint.py +@@ -22,9 +22,13 @@ import torch.distributed as dist + import torch.distributed.checkpoint as dcp + import torch.nn as nn +-from torch.distributed.checkpoint import HuggingFaceStorageWriter +-from torch.distributed.checkpoint._consolidate_hf_safetensors import ( +- consolidate_safetensors_files_on_every_rank, +-) ++try: ++ from torch.distributed.checkpoint import HuggingFaceStorageWriter ++ from torch.distributed.checkpoint._consolidate_hf_safetensors import ( ++ consolidate_safetensors_files_on_every_rank, ++ ) ++except ImportError: ++ HuggingFaceStorageWriter = None ++ consolidate_safetensors_files_on_every_rank = None + from torch.distributed.checkpoint.staging import DefaultStager, StagingOptions + from torch.distributed.checkpoint.state_dict import ( + get_model_state_dict, +@@ -531,6 +534,10 @@ class CheckpointManager(Configurable): + checkpoint_save_id: str | None = None + fqn_to_index_mapping: dict[Any, int] | None = None + if to_hf: ++ if HuggingFaceStorageWriter is None: ++ raise RuntimeError( ++ "HuggingFaceStorageWriter is unavailable in this torch build." ++ ) + assert ( + self.sd_adapter is not None + ), "trying to save checkpoint in HF safetensors format, but sd_adapter is not provided." +@@ -956,6 +963,10 @@ class CheckpointManager(Configurable): + states = self._flattened_model_states_sd() + + if self.last_save_in_hf: ++ if consolidate_safetensors_files_on_every_rank is None: ++ raise RuntimeError( ++ "HF safetensors consolidation is unavailable in this torch build." ++ ) + assert ( + self.last_save_model_only + ), "Only model can be saved when saving in HF safetensors format." +diff --git a/torchtitan/distributed/parallel_dims.py b/torchtitan/distributed/parallel_dims.py +index 5fcf646a..a267cd3f 100644 +--- a/torchtitan/distributed/parallel_dims.py ++++ b/torchtitan/distributed/parallel_dims.py +@@ -93,9 +93,9 @@ class ParallelDims: + which is created by flattening the batch and cp dimensions. + This API performs the following unflatten operations from the world mesh: + +- ["pp", "batch", "cp", "tp"] # dataloading_mesh +- ["pp", "dp_replicate", "fsdp", "tp"] # dense_mesh +- ["pp", "dp_replicate", "efsdp", "ep", "etp"] # sparse_mesh ++ ["batch", "pp", "cp", "tp"] # dataloading_mesh ++ ["dp_replicate", "fsdp", "pp", "tp"] # dense_mesh ++ ["dp_replicate", "efsdp", "ep", "pp", "etp"] # sparse_mesh + + Note: DeviceMesh currently recreates the process group for each dimension. + It should share the process group for the same dim group to avoid unnecessary +@@ -143,19 +143,19 @@ class ParallelDims: + ) + dataloading_mesh = unflatten_mesh( + self._world_mesh, +- ("pp", "batch", "cp", "tp"), +- (self.pp, batch, self.cp, self.tp), ++ ("batch", "pp", "cp", "tp"), ++ (batch, self.pp, self.cp, self.tp), + ) + loss_mesh = dataloading_mesh["batch", "cp"]._flatten("loss_mesh") + dense_mesh = unflatten_mesh( + self._world_mesh, +- ("pp", "dp_replicate", "fsdp", "tp"), +- (self.pp, self.dp_replicate, fsdp, self.tp), ++ ("dp_replicate", "fsdp", "pp", "tp"), ++ (self.dp_replicate, fsdp, self.pp, self.tp), + ) + sparse_mesh = unflatten_mesh( + self._world_mesh, +- ("pp", "dp_replicate", "efsdp", "ep", "etp"), +- (self.pp, self.dp_replicate, efsdp, self.ep, self.etp), ++ ("dp_replicate", "efsdp", "ep", "pp", "etp"), ++ (self.dp_replicate, efsdp, self.ep, self.pp, self.etp), + ) + + self._global_meshes = { +diff --git a/torchtitan/models/llama4/parallelize.py b/torchtitan/models/llama4/parallelize.py +index 03a0cc21..de795c44 100644 +--- a/torchtitan/models/llama4/parallelize.py ++++ b/torchtitan/models/llama4/parallelize.py +@@ -1,9 +1,10 @@ + # Copyright (c) Meta Platforms, Inc. and affiliates. + # All rights reserved. + # + # This source code is licensed under the BSD-style license found in the + # LICENSE file in the root directory of this source tree. + ++from typing import NamedTuple + from typing import Any + + import torch +@@ -11,9 +12,16 @@ import torch.nn as nn + from torch.distributed.device_mesh import DeviceMesh + from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard, MixedPrecisionPolicy +-from torch.distributed.fsdp._fully_shard._fsdp_common import ( +- FSDPMeshInfo, +- ShardPlacementResult, +-) ++try: ++ from torch.distributed.fsdp._fully_shard._fsdp_common import ( ++ FSDPMeshInfo, ++ ShardPlacementResult, ++ ) ++except ImportError: ++ from torch.distributed.fsdp._fully_shard._fsdp_common import FSDPMeshInfo ++ ++ class ShardPlacementResult(NamedTuple): ++ placement: object ++ mesh_info: FSDPMeshInfo + from torch.distributed.tensor import Partial, Replicate, Shard + from torch.distributed.tensor.parallel import ( + ColwiseParallel, +diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py +index 7cc6e890..f04bc298 100644 +--- a/torchtitan/distributed/pipeline_parallel.py ++++ b/torchtitan/distributed/pipeline_parallel.py +@@ -240,6 +240,13 @@ def build_pipeline_schedule( + f"with {n_microbatches} microbatches and {num_total_stages} stages." + ) + ++ if hasattr(schedule, "pipeline_order"): ++ pp_rank = torch.distributed.get_rank(stages[0].group) ++ rank_ops = schedule.pipeline_order.get(pp_rank, []) ++ logger.info( ++ f"PP rank {pp_rank} pipeline operations ({len(rank_ops)} steps): {rank_ops}" ++ ) ++ + if parallelism.pipeline_parallel_expert_parallel_overlap and isinstance( + schedule, ScheduleDualPipeV + ): +diff --git a/torchtitan/hf_datasets/text_datasets.py b/torchtitan/hf_datasets/text_datasets.py +index 51a5126a..f41038a4 100644 +--- a/torchtitan/hf_datasets/text_datasets.py ++++ b/torchtitan/hf_datasets/text_datasets.py +@@ -25,7 +25,23 @@ from torchtitan.tools.logging import logger + + def _load_c4_dataset(dataset_path: str, split: str): + """Load C4 dataset with default configuration.""" +- return load_dataset(dataset_path, name="en", split=split, streaming=True) ++ import time ++ ++ from huggingface_hub.errors import HfHubHTTPError ++ ++ max_retries = 5 ++ for attempt in range(max_retries): ++ try: ++ return load_dataset(dataset_path, name="en", split=split, streaming=True) ++ except HfHubHTTPError as e: ++ if attempt == max_retries - 1: ++ raise ++ wait = 30 * (2**attempt) ++ logger.warning( ++ f"HF API error loading C4 (attempt {attempt + 1}/{max_retries}): {e}. " ++ f"Retrying in {wait}s..." ++ ) ++ time.sleep(wait) + + + def _process_c4_text(sample: dict[str, Any]) -> str: +diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py +index fe40d387..70c1a64a 100644 +--- a/torchtitan/models/common/attention.py ++++ b/torchtitan/models/common/attention.py +@@ -23,12 +23,10 @@ from torch.nn.attention import sdpa_kernel, SDPBackend + from torch.nn.attention.flex_attention import ( + _mask_mod_signature, + _score_mod_signature, +- AuxRequest, + BlockMask, + create_block_mask, + flex_attention, + ) +-from torch.nn.attention.varlen import varlen_attn + from torch.types import Number + + from torchtitan.models.common.linear import Linear +@@ -187,6 +185,8 @@ class VarlenAttentionWrapper(LocalMapAttention): + xk_packed = xk_packed.to(torch.bfloat16) + xv_packed = xv_packed.to(torch.bfloat16) + ++ from torch.nn.attention.varlen import varlen_attn ++ + return varlen_attn( + xq_packed, + xk_packed, +@@ -260,6 +260,8 @@ class FlexAttentionWrapper(LocalMapAttention): + # 2. `self._compiled_flex_attn` is not correct, `self` will be passed in + # as the first argument, which will cause an error. + # `FlexAttentionWrapper._compiled_flex_attn` is correct. ++ from torch.nn.attention.flex_attention import AuxRequest ++ + out, aux = FlexAttentionWrapper._compiled_flex_attn( + q, + k, +@@ -317,7 +319,7 @@ class ScaledDotProductAttentionWrapper(LocalMapAttention): + super().__init__() + if not self.sdpa_backends: + self.sdpa_backends = [ +- SDPBackend.CUDNN_ATTENTION, ++ # SDPBackend.CUDNN_ATTENTION, + SDPBackend.FLASH_ATTENTION, + SDPBackend.MATH, + ] +diff --git a/torchtitan/models/common/moe/moe.py b/torchtitan/models/common/moe/moe.py +index acd2de5f..0de207ab 100644 +--- a/torchtitan/models/common/moe/moe.py ++++ b/torchtitan/models/common/moe/moe.py +@@ -28,6 +28,12 @@ def _run_experts_for_loop( + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + ) -> torch.Tensor: ++ ++ # Avoid synchronization with the host (may race as implemented right now) ++ # num_tokens_per_expert_host = torch.empty(len(num_tokens_per_expert), dtype=torch.int64, device="cpu", pin_memory=True) ++ # num_tokens_per_expert_host.copy_(num_tokens_per_expert.to(torch.int64), non_blocking=True) ++ # num_tokens_per_expert_list = num_tokens_per_expert_host.tolist() ++ + # NOTE: this would incur a synchronization between device and host + num_tokens_per_expert_list = num_tokens_per_expert.tolist() + +@@ -78,6 +84,24 @@ def _run_experts_grouped_mm( + + return out + ++ ++def _run_experts_bmm( ++ w1: torch.Tensor, ++ w2: torch.Tensor, ++ w3: torch.Tensor, ++ x: torch.Tensor, ++ num_tokens_per_expert: torch.Tensor, ++) -> torch.Tensor: ++ dim = x.shape[-1] ++ num_experts = w1.shape[0] ++ counts = num_tokens_per_expert.to(device=x.device, dtype=torch.long) ++ capacity = int(counts.max().item()) if counts.numel() else 0 ++ starts = torch.cumsum(counts, dim=0) - counts ++ positions = torch.arange(capacity, device=x.device, dtype=torch.long) ++ zero = x.new_zeros((capacity, dim)) ++ ++ packed_rows = [] ++ for expert_idx in range(num_experts): ++ valid = positions < counts[expert_idx] ++ source_idx = starts[expert_idx] + positions ++ safe_idx = torch.where(valid, source_idx, positions.new_zeros(())) ++ packed_rows.append(torch.where(valid[:, None], x[safe_idx], zero)) ++ packed = torch.stack(packed_rows, dim=0) ++ ++ h = F.silu(torch.bmm(packed, w1.transpose(-2, -1))) ++ h = h * torch.bmm(packed, w3.transpose(-2, -1)) ++ expert_out = torch.bmm(h, w2.transpose(-2, -1)) ++ ++ out = x.new_zeros(x.shape) ++ for expert_idx in range(num_experts): ++ valid = positions < counts[expert_idx] ++ dest_idx = starts[expert_idx] + positions ++ safe_idx = torch.where(valid, dest_idx, positions.new_zeros(())) ++ src = torch.where(valid[:, None], expert_out[expert_idx], zero) ++ out = out.scatter_add(0, safe_idx[:, None].expand(-1, dim), src) ++ return out ++ ++ + class GroupedExperts(Module): + @dataclass(kw_only=True, slots=True) + class Config(Module.Config): +@@ -85,6 +109,7 @@ class GroupedExperts(Module): + hidden_dim: int = field(init=False) + num_experts: int = field(init=False) + use_grouped_mm: bool = True ++ use_bmm: bool = False + + def __init__(self, config: Config): + super().__init__() +@@ -103,6 +128,7 @@ class GroupedExperts(Module): + torch.empty(config.num_experts, config.hidden_dim, config.dim) + ) + self.use_grouped_mm = config.use_grouped_mm ++ self.use_bmm = config.use_bmm + + def forward( + self, +@@ -122,6 +148,9 @@ class GroupedExperts(Module): + w2 = self.w2 + w3 = self.w3 + ++ if self.use_bmm: ++ return _run_experts_bmm(w1, w2, w3, x, num_tokens_per_expert) ++ + if self.use_grouped_mm: + # NOTE: If EP is not used, we need to pad the indices + # to prepare for grouped_mm; +diff --git a/torchtitan/models/qwen3/__init__.py b/torchtitan/models/qwen3/__init__.py +index d4bace1c..846d9e59 100644 +--- a/torchtitan/models/qwen3/__init__.py ++++ b/torchtitan/models/qwen3/__init__.py +@@ -272,6 +272,88 @@ qwen3_configs = { + ), + ), + # Qwen3-MoE models ++ "1B-A0.7B": Qwen3Model.Config( ++ vocab_size=151936, ++ dim=1024, ++ n_layers=16, ++ tok_embeddings=Embedding.Config(), ++ output=Linear.Config(), ++ norm=RMSNorm.Config(eps=1e-6), ++ enable_weight_tying=False, ++ layer=Qwen3TransformerBlock.Config( ++ attention_norm=RMSNorm.Config(eps=1e-6), ++ ffn_norm=RMSNorm.Config(eps=1e-6), ++ moe_enabled=True, ++ moe=MoE.Config( ++ hidden_dim=3584, ++ num_experts=4, ++ num_shared_experts=0, ++ score_before_experts=False, ++ load_balance_coeff=None, ++ router=TokenChoiceTopKRouter.Config( ++ top_k=2, ++ score_func="softmax", ++ ), ++ ), ++ feed_forward=FeedForward.Config(hidden_dim=3584), ++ attention=GQAttention.Config( ++ n_heads=16, ++ n_kv_heads=8, ++ head_dim=64, ++ q_norm=RMSNorm.Config(eps=1e-6), ++ k_norm=RMSNorm.Config(eps=1e-6), ++ attn_backend="sdpa", ++ rope_backend="cos_sin", ++ ), ++ ), ++ rope=RoPE.Config( ++ dim=64, ++ max_seq_len=2048, ++ theta=1000000.0, ++ backend="cos_sin", ++ ), ++ ), ++ "9B-A3B": Qwen3Model.Config( ++ vocab_size=151936, ++ dim=2048, ++ n_layers=24, ++ tok_embeddings=Embedding.Config(), ++ output=Linear.Config(), ++ norm=RMSNorm.Config(eps=1e-6), ++ enable_weight_tying=False, ++ layer=Qwen3TransformerBlock.Config( ++ attention_norm=RMSNorm.Config(eps=1e-6), ++ ffn_norm=RMSNorm.Config(eps=1e-6), ++ moe_enabled=True, ++ moe=MoE.Config( ++ hidden_dim=7168, ++ num_experts=8, ++ num_shared_experts=0, ++ score_before_experts=False, ++ load_balance_coeff=None, ++ router=TokenChoiceTopKRouter.Config( ++ top_k=2, ++ score_func="softmax", ++ ), ++ ), ++ feed_forward=FeedForward.Config(hidden_dim=7168), ++ attention=GQAttention.Config( ++ n_heads=32, ++ n_kv_heads=8, ++ head_dim=64, ++ q_norm=RMSNorm.Config(eps=1e-6), ++ k_norm=RMSNorm.Config(eps=1e-6), ++ attn_backend="sdpa", ++ rope_backend="cos_sin", ++ ), ++ ), ++ rope=RoPE.Config( ++ dim=64, ++ max_seq_len=2048, ++ theta=1000000.0, ++ backend="cos_sin", ++ ), ++ ), + "debugmodel_moe": Qwen3Model.Config( + vocab_size=2048, + dim=256, +diff --git a/torchtitan/models/qwen3/config_registry.py b/torchtitan/models/qwen3/config_registry.py +index 80aed1f6..8e9094a8 100644 +--- a/torchtitan/models/qwen3/config_registry.py ++++ b/torchtitan/models/qwen3/config_registry.py +@@ -11,6 +11,7 @@ from torchtitan.components.optimizer import OptimizersContainer + from torchtitan.config import ( + ActivationCheckpointConfig, + ParallelismConfig, ++ CompileConfig, + TrainingConfig, + ) + from torchtitan.hf_datasets.text_datasets import ( +@@ -202,13 +203,18 @@ def qwen3_moe_debug() -> Trainer.Config: + optimizer=OptimizersContainer.Config(lr=3e-4), + lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), + training=TrainingConfig( +- local_batch_size=4, +- seq_len=4096, +- steps=10, ++ global_batch_size=32, # Should be divisible by pipeline_parallel_microbatch_size ++ local_batch_size=16, # Should be divisible by pipeline_parallel_microbatch_size ++ seq_len=512, ++ steps=8, # Few iterations for profiling ++ dtype="bfloat16", + ), + parallelism=ParallelismConfig( +- expert_parallel_degree=1, ++ expert_parallel_degree=2, + expert_tensor_parallel_degree=1, ++ pipeline_parallel_degree=2, ++ pipeline_parallel_microbatch_size=4, # Configurable micro batch size ++ data_parallel_shard_degree=-1, + ), + checkpoint=CheckpointManager.Config( + interval=10, +@@ -221,6 +227,202 @@ def qwen3_moe_debug() -> Trainer.Config: + ) + + ++def qwen3_1b_single() -> Trainer.Config: ++ return Trainer.Config( ++ hf_assets_path="./assets/hf/Qwen3-0.6B", ++ model_spec=model_registry("1B-A0.7B"), ++ dataloader=HuggingFaceTextDataLoader.Config( ++ dataset="c4", ++ ), ++ optimizer=OptimizersContainer.Config(lr=8e-4), ++ lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), ++ training=TrainingConfig( ++ local_batch_size=8, # Should be divisible by pipeline_parallel_microbatch_size ++ seq_len=32, ++ steps=8, # Few iterations for profiling ++ dtype="bfloat16", ++ ), ++ parallelism=ParallelismConfig( ++ expert_parallel_degree=1, ++ expert_tensor_parallel_degree=1, ++ pipeline_parallel_degree=1, ++ pipeline_parallel_microbatch_size=1, # Configurable micro batch size ++ data_parallel_shard_degree=1, ++ ), ++ checkpoint=CheckpointManager.Config( ++ enable=False, ++ ), ++ activation_checkpoint=ActivationCheckpointConfig( ++ mode="none", ++ ), ++ # compile=CompileConfig(enable=True), ++ ) ++ ++ ++def qwen3_1b_single_no_grouped_mm() -> Trainer.Config: ++ cfg = qwen3_1b_single() ++ assert cfg.model_spec is not None ++ assert cfg.model_spec.model.layer.moe is not None ++ cfg.model_spec.model.layer.moe.experts.use_grouped_mm = False ++ return cfg ++ ++ ++def qwen3_1b() -> Trainer.Config: ++ return Trainer.Config( ++ hf_assets_path="./assets/hf/Qwen3-0.6B", ++ model_spec=model_registry("1B-A0.7B"), ++ dataloader=HuggingFaceTextDataLoader.Config( ++ dataset="c4", ++ ), ++ optimizer=OptimizersContainer.Config(lr=8e-4), ++ lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), ++ training=TrainingConfig( ++ # global_batch_size=128, # Should be divisible by pipeline_parallel_microbatch_size ++ local_batch_size=32, # Should be divisible by pipeline_parallel_microbatch_size ++ seq_len=32, ++ steps=8, # Few iterations for profiling ++ dtype="bfloat16", ++ ), ++ parallelism=ParallelismConfig( ++ expert_parallel_degree=2, ++ expert_tensor_parallel_degree=1, ++ pipeline_parallel_degree=4, ++ pipeline_parallel_microbatch_size=8, # Configurable micro batch size ++ data_parallel_shard_degree=-1, ++ ), ++ checkpoint=CheckpointManager.Config( ++ enable=False, ++ ), ++ activation_checkpoint=ActivationCheckpointConfig( ++ mode="none", ++ ), ++ compile=CompileConfig(enable=True), ++ ) ++ ++ ++def qwen3_1b_no_grouped_mm() -> Trainer.Config: ++ cfg = qwen3_1b() ++ assert cfg.model_spec is not None ++ assert cfg.model_spec.model.layer.moe is not None ++ cfg.model_spec.model.layer.moe.experts.use_grouped_mm = False ++ return cfg ++ ++ ++def qwen3_9b_single() -> Trainer.Config: ++ return Trainer.Config( ++ hf_assets_path="./assets/hf/Qwen3-8B", ++ model_spec=model_registry("9B-A3B"), ++ dataloader=HuggingFaceTextDataLoader.Config( ++ dataset="c4", ++ ), ++ optimizer=OptimizersContainer.Config(lr=8e-4), ++ lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), ++ training=TrainingConfig( ++ local_batch_size=128, # Should be divisible by pipeline_parallel_microbatch_size ++ seq_len=512, ++ steps=8, # Few iterations for profiling ++ dtype="bfloat16", ++ ), ++ parallelism=ParallelismConfig( ++ pipeline_parallel_degree=8, ++ pipeline_parallel_microbatch_size=8, # Configurable micro batch size ++ ), ++ checkpoint=CheckpointManager.Config( ++ enable=False, ++ ), ++ activation_checkpoint=ActivationCheckpointConfig( ++ mode="none", ++ ), ++ compile=CompileConfig(enable=True), ++ ) ++ ++def qwen3_9b() -> Trainer.Config: ++ """Generic Qwen3 9B config for runtime PP/DP/EP/schedule/batch/ZeRO overrides.""" ++ return Trainer.Config( ++ hf_assets_path="./assets/hf/Qwen3-8B", ++ model_spec=model_registry("9B-A3B"), ++ dataloader=HuggingFaceTextDataLoader.Config(dataset="c4"), ++ optimizer=OptimizersContainer.Config(lr=8e-4), ++ lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), ++ training=TrainingConfig( ++ global_batch_size=64, ++ local_batch_size=64, ++ seq_len=512, ++ steps=8, ++ dtype="bfloat16", ++ ), ++ parallelism=ParallelismConfig( ++ expert_parallel_degree=1, ++ pipeline_parallel_degree=8, ++ pipeline_parallel_schedule="1F1B", ++ pipeline_parallel_microbatch_size=4, ++ data_parallel_replicate_degree=1, ++ data_parallel_shard_degree=1, ++ fsdp_reshard_after_forward="default", ++ ), ++ checkpoint=CheckpointManager.Config(enable=False), ++ activation_checkpoint=ActivationCheckpointConfig(mode="none"), ++ compile=CompileConfig(enable=True), ++ ) ++ ++ ++def qwen3_30b() -> Trainer.Config: ++ return Trainer.Config( ++ hf_assets_path="./assets/hf/Qwen3-1.7B", ++ model_spec=model_registry("30B-A3B"), ++ dataloader=HuggingFaceTextDataLoader.Config( ++ dataset="c4", ++ ), ++ optimizer=OptimizersContainer.Config(lr=8e-4), ++ lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), ++ training=TrainingConfig( ++ global_batch_size=256, # Should be divisible by pipeline_parallel_microbatch_size ++ local_batch_size=128, # Should be divisible by pipeline_parallel_microbatch_size ++ seq_len=1024, ++ steps=8, # Few iterations for profiling ++ dtype="bfloat16", ++ ), ++ parallelism=ParallelismConfig( ++ expert_parallel_degree=2, ++ expert_tensor_parallel_degree=1, ++ pipeline_parallel_degree=8, ++ pipeline_parallel_microbatch_size=8, # Configurable micro batch size ++ data_parallel_shard_degree=-1, ++ ), ++ checkpoint=CheckpointManager.Config( ++ enable=False, ++ ), ++ activation_checkpoint=ActivationCheckpointConfig( ++ mode="none", ++ ), ++ compile=CompileConfig(enable=True), ++ ) ++ ++ ++def qwen3_30b_single() -> Trainer.Config: ++ return Trainer.Config( ++ hf_assets_path="./assets/hf/Qwen3-1.7B", ++ model_spec=model_registry("30B-A3B"), ++ dataloader=HuggingFaceTextDataLoader.Config( ++ dataset="c4", ++ ), ++ optimizer=OptimizersContainer.Config(lr=8e-4), ++ lr_scheduler=LRSchedulersContainer.Config(warmup_steps=2), ++ training=TrainingConfig( ++ local_batch_size=16, # Should be divisible by pipeline_parallel_microbatch_size ++ seq_len=256, ++ steps=8, # Few iterations for profiling ++ dtype="bfloat16", ++ ), ++ checkpoint=CheckpointManager.Config( ++ enable=False, ++ ), ++ activation_checkpoint=ActivationCheckpointConfig( ++ mode="none", ++ ), ++ # compile=CompileConfig(enable=True), ++ ) ++ + def sft_qwen3_8b_math() -> Trainer.Config: + """Qwen3-8B SFT on GSM8K math dataset.""" + +diff --git a/torchtitan/trainer.py b/torchtitan/trainer.py +index 6137578f..b65973d4 100644 +--- a/torchtitan/trainer.py ++++ b/torchtitan/trainer.py +@@ -234,6 +234,12 @@ class Trainer(torch.distributed.checkpoint.stateful.Stateful, Configurable): + self.tokenizer = config.tokenizer.build(tokenizer_path=config.hf_assets_path) + + # build dataloader ++ # Non-zero local ranks wait for local rank 0 to prime the HuggingFace ++ # dataset cache before making their own requests, avoiding HTTP timeouts ++ # when multiple processes hit the HF servers simultaneously. ++ local_rank = int(os.environ.get("LOCAL_RANK", 0)) ++ if local_rank != 0: ++ torch.distributed.barrier() + self.dataloader = config.dataloader.build( + dp_world_size=batch_degree, + dp_rank=batch_rank, +@@ -241,6 +247,8 @@ class Trainer(torch.distributed.checkpoint.stateful.Stateful, Configurable): + seq_len=config.training.seq_len, + local_batch_size=config.training.local_batch_size, + ) ++ if local_rank == 0: ++ torch.distributed.barrier() + + # build model (using meta init) + model_config = model_spec.model +@@ -819,8 +827,10 @@ class Trainer(torch.distributed.checkpoint.stateful.Stateful, Configurable): + base_folder=config.dump_folder, + ) as memory_profiler, + ): ++ iter_times = [] + data_iterator = self.batch_generator(self.dataloader) + while self.should_continue_training(): ++ start = time.perf_counter() + self.step += 1 + self.gc_handler.run(self.step) + try: +@@ -852,6 +862,19 @@ class Trainer(torch.distributed.checkpoint.stateful.Stateful, Configurable): + timeout=timedelta(seconds=config.comm.train_timeout_seconds), + parallel_dims=self.parallel_dims, + ) ++ end = time.perf_counter() ++ iter_time = end - start ++ iter_times.append(iter_time) ++ logger.info(f"Training iter {self.step} took {iter_time:.5f} s") ++ ++ if len(iter_times) > 0: ++ last_n = iter_times[-5:] ++ mean = sum(last_n) / len(last_n) ++ variance = sum((t - mean) ** 2 for t in last_n) / len(last_n) ++ std = variance**0.5 ++ logger.info( ++ f"Final {len(last_n)} iter times — avg: {mean:.5f} s, std: {std:.5f} s" ++ ) + + if torch.distributed.get_rank() == 0: + logger.info("Sleeping 2 seconds for other ranks to complete") diff --git a/artifact/run_deepspeed.py b/artifact/run_deepspeed.py new file mode 100644 index 0000000..68116fd --- /dev/null +++ b/artifact/run_deepspeed.py @@ -0,0 +1,621 @@ + +#!/usr/bin/env python3 +""" +DeepSpeed training script for Qwen3 MoE models. + +Parallelism layout (8 GPUs, 2 nodes x 4 GPUs/node): + - PP = 4 (intra-node pipeline parallel) + - EP = 2 (inter-node expert parallel, splits DP dim) + - DP = 2 (data parallel, consumed by EP) + +Config: + - qwen3_1b: global_bs=128, micro_bs=8, seq_len=2048, grad_accum=8 + - qwen3_9b: global_bs=256, micro_bs=8, seq_len=2048, grad_accum=16 + +Topology design (see setup_custom_topology for details): + We use ProcessTopology(axes=['data', 'pipe'], dims=[2, 4]) so that the + 'pipe' axis varies fastest → PP stages map to adjacent ranks (intra-node). + Then we pre-register EP process groups with the correct cross-node rank + pairs BEFORE deepspeed.initialize(), because DeepSpeed's default EP group + creation assumes pipe-slowest layout and would produce wrong groups. +""" + +import argparse +import os +import time + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import Dataset + +import deepspeed +from deepspeed.pipe import PipelineModule, LayerSpec +from deepspeed.runtime.pipe.topology import ProcessTopology +from deepspeed.utils import groups as ds_groups + +# --------------------------------------------------------------------------- +# Qwen3 architecture constants +# --------------------------------------------------------------------------- +ROPE_THETA = 1_000_000.0 +MODEL_CONFIGS = { + "qwen3_1b": { + "vocab_size": 151936, + "dim": 1024, + "n_layers": 16, + "n_heads": 16, + "n_kv_heads": 8, + "head_dim": 64, + "moe_inter_dim": 3584, + "num_experts": 4, + "top_k": 2, + "seq_len": 2048, + "global_batch_size": 128, + "micro_batch_size": 8, + }, + "qwen3_9b": { + "vocab_size": 151936, + "dim": 2048, + "n_layers": 24, + "n_heads": 32, + "n_kv_heads": 8, + "head_dim": 64, + "moe_inter_dim": 7168, + "num_experts": 8, + "top_k": 2, + "seq_len": 2048, + "global_batch_size": 256, + "micro_batch_size": 8, + }, +} + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() + return (x.float() * norm).type_as(x) * self.weight + + +def precompute_rope(dim: int, seq_len: int, theta: float = 1e6): + """Return (cos, sin) each of shape [seq_len, dim//2].""" + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + t = torch.arange(seq_len, dtype=torch.float32) + freqs = torch.outer(t, freqs) # [seq_len, dim//2] + return torch.cos(freqs), torch.sin(freqs) + + +def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor): + """x: [B, n_heads, S, head_dim]. cos/sin: [S, head_dim//2].""" + half = x.shape[-1] // 2 + x1, x2 = x[..., :half], x[..., half:] + cos = cos[:x.shape[2]].unsqueeze(0).unsqueeze(0) # [1, 1, S, half] + sin = sin[:x.shape[2]].unsqueeze(0).unsqueeze(0) + out1 = x1 * cos - x2 * sin + out2 = x2 * cos + x1 * sin + return torch.cat([out1, out2], dim=-1) + + +def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: + """Repeat KV heads to match query head count.""" + if n_rep == 1: + return x + B, n_kv, S, D = x.shape + return x[:, :, None, :, :].expand(B, n_kv, n_rep, S, D).reshape(B, n_kv * n_rep, S, D) + + +# --------------------------------------------------------------------------- +# Model layers (each takes and returns a single tensor for PipelineModule) +# +# Convention: the pipeline tensor is [B, S, D] for hidden states. +# The first layer converts input_ids → hidden, the last converts hidden → logits. +# --------------------------------------------------------------------------- + +class EmbeddingLayer(nn.Module): + """Token embedding: input_ids [B, S] → hidden [B, S, D].""" + def __init__(self, vocab_size: int, dim: int): + super().__init__() + self.tok_emb = nn.Embedding(vocab_size, dim) + + def forward(self, input_ids): + return self.tok_emb(input_ids) + + +class TransformerBlock(nn.Module): + """ + Single Qwen3 MoE transformer block. + GQA attention + MoE FFN, with pre-norm (RMSNorm) and residual connections. + + Takes hidden [B, S, D], returns hidden [B, S, D]. + """ + def __init__(self, model_config: dict[str, int], layer_id: int = 0): + super().__init__() + self.layer_id = layer_id + self.dim = int(model_config["dim"]) + self.n_heads = int(model_config["n_heads"]) + self.n_kv_heads = int(model_config["n_kv_heads"]) + self.head_dim = int(model_config["head_dim"]) + + # --- Attention --- + self.attn_norm = RMSNorm(self.dim) + self.q_proj = nn.Linear(self.dim, self.n_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.n_heads * self.head_dim, self.dim, bias=False) + # qk_norm + self.q_norm = RMSNorm(self.head_dim) + self.k_norm = RMSNorm(self.head_dim) + # RoPE cache (registered as buffer so it moves with .to(device)) + cos, sin = precompute_rope(self.head_dim, int(model_config["seq_len"]), ROPE_THETA) + self.register_buffer("rope_cos", cos, persistent=False) + self.register_buffer("rope_sin", sin, persistent=False) + + # --- MoE FFN --- + self.ffn_norm = RMSNorm(self.dim) + # Build a single expert template, then wrap with DeepSpeed MoE + expert = SwiGLUExpert(self.dim, int(model_config["moe_inter_dim"])) + from deepspeed.moe.layer import MoE + self.moe = MoE( + hidden_size=self.dim, + expert=expert, + num_experts=int(model_config["num_experts"]), + ep_size=int(model_config["ep_size"]), + k=int(model_config["top_k"]), + capacity_factor=1.0, + eval_capacity_factor=1.0, + min_capacity=4, + noisy_gate_policy=None, + drop_tokens=False, # no token dropping for training stability + use_residual=False, + use_tutel=False, + ) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + B, S, D = hidden.shape + + # ---------- Self-Attention ---------- + residual = hidden + h = self.attn_norm(hidden) + + q = self.q_proj(h).view(B, S, self.n_heads, self.head_dim) + k = self.k_proj(h).view(B, S, self.n_kv_heads, self.head_dim) + v = self.v_proj(h).view(B, S, self.n_kv_heads, self.head_dim) + + # qk_norm + q = self.q_norm(q) + k = self.k_norm(k) + + # [B, heads, S, D] + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + # RoPE + q = apply_rope(q, self.rope_cos, self.rope_sin) + k = apply_rope(k, self.rope_cos, self.rope_sin) + + # GQA: repeat KV heads + n_rep = self.n_heads // self.n_kv_heads + k = repeat_kv(k, n_rep) + v = repeat_kv(v, n_rep) + + # Scaled dot-product attention (causal) + attn_out = F.scaled_dot_product_attention( + q, k, v, is_causal=True, dropout_p=0.0 + ) # [B, heads, S, head_dim] + attn_out = attn_out.transpose(1, 2).contiguous().view(B, S, -1) + hidden = residual + self.o_proj(attn_out) + + # ---------- MoE FFN ---------- + residual = hidden + h = self.ffn_norm(hidden) + # DeepSpeed MoE expects [B*S, D] or [B, S, D]; it handles both + moe_out, aux_loss, _ = self.moe(h) + hidden = residual + moe_out + + return hidden + + +class SwiGLUExpert(nn.Module): + """Single SwiGLU expert: gate_proj + up_proj → SiLU → down_proj.""" + def __init__(self, dim: int, inter_dim: int): + super().__init__() + self.gate_proj = nn.Linear(dim, inter_dim, bias=False) + self.up_proj = nn.Linear(dim, inter_dim, bias=False) + self.down_proj = nn.Linear(inter_dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) + + +class FinalNorm(nn.Module): + """Final RMSNorm before LM head.""" + def __init__(self, dim: int): + super().__init__() + self.norm = RMSNorm(dim) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return self.norm(hidden) + + +class LMHead(nn.Module): + """Linear projection to vocab logits: [B, S, D] → [B*S, V].""" + def __init__(self, dim: int, vocab_size: int): + super().__init__() + self.vocab_size = vocab_size + self.head = nn.Linear(dim, vocab_size, bias=False) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + # Flatten for cross-entropy: [B*S, V] + return self.head(hidden).view(-1, self.vocab_size) + + +# --------------------------------------------------------------------------- +# Loss function for PipelineModule +# --------------------------------------------------------------------------- + +def loss_fn(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + """Cross-entropy loss. logits: [B*S, V], labels: [B, S] or [B*S].""" + labels = labels.view(-1) + return F.cross_entropy(logits, labels) + + +# --------------------------------------------------------------------------- +# Mock dataset +# --------------------------------------------------------------------------- + +class MockTokenDataset(Dataset): + """Random token dataset for benchmarking. No real data needed.""" + + def __init__(self, vocab_size: int, seq_len: int, num_samples: int, seed: int = 42): + self.vocab_size = vocab_size + self.seq_len = seq_len + self.num_samples = num_samples + self.seed = seed + + def __len__(self): + return self.num_samples + + def __getitem__(self, idx): + rng = torch.Generator() + rng.manual_seed(self.seed + idx) + tokens = torch.randint(0, self.vocab_size, (self.seq_len + 1,), generator=rng) + return tokens[:-1], tokens[1:] # (input_ids, labels) + + +# --------------------------------------------------------------------------- +# Pre-register EP groups for custom topology +# --------------------------------------------------------------------------- + +def pre_register_ep_groups(ep_size: int, pp_size: int, dp_size: int): + """ + Pre-register expert-parallel (EP) and expert-data-parallel process groups + in DeepSpeed's global groups module BEFORE deepspeed.initialize() is called. + + Why this is needed: + ------------------ + We use ProcessTopology(axes=['data', 'pipe'], dims=[dp, pp]) so that the + 'pipe' axis is fastest-varying → PP stages are adjacent ranks → intra-node. + + Rank layout with 2 nodes × 4 GPUs: + rank 0: data=0, pipe=0 (node 0) + rank 1: data=0, pipe=1 (node 0) + rank 2: data=0, pipe=2 (node 0) + rank 3: data=0, pipe=3 (node 0) + rank 4: data=1, pipe=0 (node 1) + rank 5: data=1, pipe=1 (node 1) + rank 6: data=1, pipe=2 (node 1) + rank 7: data=1, pipe=3 (node 1) + + PP groups (varying pipe, same data): [0,1,2,3] and [4,5,6,7] → intra-node ✓ + DP groups (varying data, same pipe): [0,4], [1,5], [2,6], [3,7] → cross-node ✓ + + EP should align with DP (same PP stage, different data ranks): + EP groups: [0,4], [1,5], [2,6], [3,7] → cross-node ✓ + + However, DeepSpeed's _create_expert_and_data_parallel() assumes pipe is the + SLOWEST axis and creates EP groups from consecutive ranks: [0,1], [2,3], ... + which would be WRONG for our layout (those pair different PP stages on the + same node). + + Solution: we pre-populate the _EXPERT_PARALLEL_GROUP and + _EXPERT_DATA_PARALLEL_GROUP dicts. When DeepSpeed's MoE._create_process_groups + runs, it checks `if group_name not in groups._get_expert_parallel_group_dict()` + and skips creation if the groups already exist. + """ + import torch.distributed as dist + + rank = dist.get_rank() + world_size = dist.get_world_size() + assert world_size == pp_size * dp_size + + group_name = f"ep_size_{ep_size}" + + # EP groups: ranks sharing the same pipe stage but different data index + # With axes=['data', 'pipe'], rank = data_idx * pp_size + pipe_idx + # So for a given pipe stage p, EP peers are: [p, p + pp_size, p + 2*pp_size, ...] + # With dp_size=2, pp_size=4: EP groups = [0,4], [1,5], [2,6], [3,7] + for pipe_stage in range(pp_size): + ep_ranks = list(range(pipe_stage, world_size, pp_size)) + # ep_ranks has dp_size elements; chunk into groups of ep_size + for start in range(0, len(ep_ranks), ep_size): + ranks = ep_ranks[start : start + ep_size] + group = dist.new_group(ranks) + if rank in ranks: + ds_groups._EXPERT_PARALLEL_GROUP[group_name] = group + ds_groups._EXPERT_PARALLEL_GROUP_RANKS[group_name] = ranks + if rank == 0: + print(f" EP group ({group_name}): {ranks}") + + # Expert-data-parallel groups: ranks with the same EP role but different + # data-parallel shards. With ep_size == dp_size (our case), each EP group + # spans the entire DP dimension, so expert-data-parallel is trivially each + # rank alone (no all-reduce needed for MoE params). + # More generally: for each pipe stage, chunk DP ranks into EP groups, + # then expert-data-parallel = ranks at the same position across EP groups. + for pipe_stage in range(pp_size): + dp_ranks = list(range(pipe_stage, world_size, pp_size)) + # With ep_size == dp_size, there's only 1 EP group per pipe stage, + # so expert-data-parallel groups are just individual ranks. + # With ep_size < dp_size, we'd interleave. + num_ep_groups = dp_size // ep_size + for pos_in_ep in range(ep_size): + edp_ranks = [dp_ranks[g * ep_size + pos_in_ep] for g in range(num_ep_groups)] + group = dist.new_group(edp_ranks) + if rank in edp_ranks: + ds_groups._EXPERT_DATA_PARALLEL_GROUP[group_name] = group + ds_groups._EXPERT_DATA_PARALLEL_GROUP_RANKS[group_name] = edp_ranks + if rank == 0: + print(f" Expert-DP group ({group_name}): {edp_ranks}") + + +# --------------------------------------------------------------------------- +# Build PipelineModule +# --------------------------------------------------------------------------- + +def build_pipeline_model(model_config: dict[str, int], args): + """Construct PipelineModule for the selected Qwen3 model.""" + layers = [] + + # Embedding + layers.append(LayerSpec(EmbeddingLayer, int(model_config["vocab_size"]), int(model_config["dim"]))) + + for i in range(int(model_config["n_layers"])): + layers.append(LayerSpec(TransformerBlock, model_config, layer_id=i)) + + # Final norm + LM head + layers.append(LayerSpec(FinalNorm, int(model_config["dim"]))) + layers.append(LayerSpec(LMHead, int(model_config["dim"]), int(model_config["vocab_size"]))) + + # ----------------------------------------------------------------------- + # Topology: PP=4 intra-node, DP=2 cross-node (consumed by EP=2) + # 8 GPUs total: 4 PP stages × 2 DP ranks + # + # We use ProcessTopology with axes=['data', 'pipe'] so that 'pipe' is the + # fastest-varying axis. This ensures PP stages are adjacent ranks and + # thus intra-node (torchrun assigns rank = node_rank * nproc + local). + # + # Rank layout: + # rank 0: data=0, pipe=0 (node 0, GPU 0) + # rank 1: data=0, pipe=1 (node 0, GPU 1) + # rank 2: data=0, pipe=2 (node 0, GPU 2) + # rank 3: data=0, pipe=3 (node 0, GPU 3) + # rank 4: data=1, pipe=0 (node 1, GPU 0) + # rank 5: data=1, pipe=1 (node 1, GPU 1) + # rank 6: data=1, pipe=2 (node 1, GPU 2) + # rank 7: data=1, pipe=3 (node 1, GPU 3) + # + # PP groups (intra-node): [0,1,2,3] and [4,5,6,7] + # DP groups (cross-node): [0,4], [1,5], [2,6], [3,7] + # EP groups (cross-node): [0,4], [1,5], [2,6], [3,7] (pre-registered) + # ----------------------------------------------------------------------- + topo = ProcessTopology(axes=['data', 'pipe'], dims=[args.dp, args.pp]) + + model = PipelineModule( + layers=layers, + topology=topo, + loss_fn=loss_fn, + partition_method="parameters", + activation_checkpoint_interval=0, + ) + return model + + +def parameter_counts(model_config: dict[str, int]) -> tuple[int, int]: + """Return total and active parameter counts for the local Qwen3 MoE model.""" + vocab_size = int(model_config["vocab_size"]) + dim = int(model_config["dim"]) + n_heads = int(model_config["n_heads"]) + n_kv_heads = int(model_config["n_kv_heads"]) + head_dim = int(model_config["head_dim"]) + n_layers = int(model_config["n_layers"]) + hidden_dim = int(model_config["moe_inter_dim"]) + num_experts = int(model_config["num_experts"]) + top_k = int(model_config["top_k"]) + + embeddings_and_head = 2 * vocab_size * dim + attention = ( + dim * (n_heads * head_dim) + + dim * (n_kv_heads * head_dim) + + dim * (n_kv_heads * head_dim) + + (n_heads * head_dim) * dim + ) + norms = dim + head_dim + head_dim + dim + router = dim * num_experts + expert = 3 * dim * hidden_dim + total_block = attention + norms + router + num_experts * expert + active_block = attention + norms + router + top_k * expert + return embeddings_and_head + n_layers * total_block, embeddings_and_head + n_layers * active_block + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + parser = argparse.ArgumentParser(description="DeepSpeed Qwen3 MoE PP4+EP2") + parser.add_argument("--model", choices=tuple(MODEL_CONFIGS), default="qwen3_1b") + parser.add_argument("--pp", type=int, default=4) + parser.add_argument("--dp", type=int, default=2) + parser.add_argument("--ep", type=int, default=1) + parser.add_argument("--zero-stage", type=int, default=1, choices=(0, 1, 2, 3)) + parser.add_argument("--schedule", choices=("1f1b",), default="1f1b") + parser.add_argument("--micro-bs", type=int, default=None) + parser.add_argument("--global-bs", type=int, default=None) + parser.add_argument("--seq-len", type=int, default=None) + parser.add_argument("--local_rank", type=int, default=-1) + parser.add_argument("--steps", type=int, default=8) + parser.add_argument("--seed", type=int, default=42) + parser = deepspeed.add_config_arguments(parser) + args = parser.parse_args() + + # Initialize distributed + deepspeed.init_distributed() + torch.manual_seed(args.seed) + + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + if args.schedule != "1f1b": + raise ValueError("DeepSpeed runner currently supports only 1f1b schedule") + if args.pp > 1 and args.zero_stage > 1: + raise ValueError("DeepSpeed pipeline parallelism is not compatible with ZeRO-2/3") + if args.dp < args.ep or args.dp % args.ep != 0: + raise ValueError("--dp must be divisible by --ep and at least as large as --ep") + assert world_size == args.pp * args.dp, \ + f"Expected {args.pp * args.dp} GPUs (PP={args.pp} x DP={args.dp}), got {world_size}" + model_config = dict(MODEL_CONFIGS[args.model]) + model_config["ep_size"] = args.ep + if args.seq_len is not None: + model_config["seq_len"] = args.seq_len + if args.micro_bs is not None: + model_config["micro_batch_size"] = args.micro_bs + if args.global_bs is not None: + model_config["global_batch_size"] = args.global_bs + grad_accum_steps = int(model_config["global_batch_size"]) // ( + int(model_config["micro_batch_size"]) * args.dp + ) + if grad_accum_steps < 1: + raise ValueError("Derived gradient_accum_steps must be >= 1") + + if rank == 0: + print(f"{'='*60}") + print(f"{args.model} — DeepSpeed") + print(f" World size: {world_size}") + print(f" PP stages: {args.pp}") + print(f" EP size: {args.ep}") + print(f" DP size: {args.dp}") + print(f" ZeRO stage: {args.zero_stage}") + print(f" Seq len: {model_config['seq_len']}") + print(f" Global BS: {model_config['global_batch_size']}") + print(f" Micro BS: {model_config['micro_batch_size']}") + print(f" Grad accum: {grad_accum_steps}") + print(f" Num layers: {model_config['n_layers']}") + print(f" Num experts: {model_config['num_experts']}") + print(f" Top-K: {model_config['top_k']}") + total_params, active_params = parameter_counts(model_config) + print(f" Total params: {total_params:,}") + print(f" Active params:{active_params:,}") + print(f"{'='*60}") + + # Pre-register EP groups with correct cross-node topology BEFORE + # deepspeed.initialize() so that the automatic (wrong) creation is skipped. + if rank == 0: + print("\nPre-registering EP process groups (cross-node):") + pre_register_ep_groups(ep_size=args.ep, pp_size=args.pp, dp_size=args.dp) + + # Build model + model = build_pipeline_model(model_config, args) + + # Dataset: enough samples for all steps + # PipelineEngine handles batching internally + num_samples = int(model_config["global_batch_size"]) * (args.steps + 2) + dataset = MockTokenDataset( + vocab_size=int(model_config["vocab_size"]), + seq_len=int(model_config["seq_len"]), + num_samples=num_samples, + seed=args.seed, + ) + + ds_config = { + "train_micro_batch_size_per_gpu": int(model_config["micro_batch_size"]), + "train_batch_size": int(model_config["global_batch_size"]), + "gradient_accumulation_steps": grad_accum_steps, + "steps_per_print": 1, + "bf16": {"enabled": True}, + "optimizer": { + "type": "Adam", + "params": { + "lr": 3.0e-4, + "betas": [0.9, 0.95], + "eps": 1.0e-8, + "weight_decay": 0.1, + }, + }, + "zero_optimization": { + "stage": args.zero_stage, + }, + } + + # Initialize DeepSpeed (auto-detects PipelineModule → PipelineEngine) + engine, _, _, _ = deepspeed.initialize( + args=args, + config=ds_config, + model=model, + model_parameters=[p for p in model.parameters() if p.requires_grad], + training_data=dataset, + ) + torch.cuda.reset_peak_memory_stats() + + if rank == 0: + print(f"\nPipeline stage: {engine.stage_id}, " + f"num micro-batches: {engine.micro_batches}") + print(f"Starting training for {args.steps} steps...\n") + + # Training loop + start_time = time.time() + for step in range(1, args.steps + 1): + step_start = time.time() + loss = engine.train_batch() + step_time = time.time() - step_start + + if rank == 0: + loss_val = loss.item() if torch.is_tensor(loss) else loss + print(f"[Step {step:3d}/{args.steps}] loss={loss_val:.4f} " + f"step_time={step_time:.2f}s") + + total_time = time.time() - start_time + torch.distributed.barrier() + peak_allocated_gb = torch.cuda.max_memory_allocated() / (1024 ** 3) + peak_reserved_gb = torch.cuda.max_memory_reserved() / (1024 ** 3) + local_peak = { + "rank": rank, + "peak_allocated_gb": peak_allocated_gb, + "peak_reserved_gb": peak_reserved_gb, + } + gathered_peaks = [None for _ in range(world_size)] + torch.distributed.all_gather_object(gathered_peaks, local_peak) + if rank == 0: + for peak in sorted(gathered_peaks, key=lambda item: int(item["rank"])): + print( + f"[rank{peak['rank']}] " + f"peak_memory_allocated_gb={float(peak['peak_allocated_gb']):.3f} " + f"peak_memory_reserved_gb={float(peak['peak_reserved_gb']):.3f}", + flush=True, + ) + if rank == 0: + print(f"\nTraining complete. Total time: {total_time:.2f}s " + f"({total_time / args.steps:.2f}s/step)") + + +if __name__ == "__main__": + main() diff --git a/artifact/run_megatron.py b/artifact/run_megatron.py new file mode 100644 index 0000000..f5866bb --- /dev/null +++ b/artifact/run_megatron.py @@ -0,0 +1,349 @@ +################ Below is training script for Megatron ################ + + + +""" +Train a Qwen3-30B-A3B style MoE model using mock (synthetic) data. +No real dataset required — useful for debugging and benchmarking. + +Architecture: Qwen3-30B-A3B MoE + - 48 layers, hidden=2048, 32 attention heads (GQA: 4 query groups), kv_channels=128 + - FFN: dense ffn_hidden_size=6144 (unused for MoE layers), moe_ffn_hidden_size=768 + - 128 experts, top-8 routing, all layers are MoE (moe_layer_freq=1) + - RoPE: base=1000000, position_embedding_type=rope + - QK LayerNorm enabled + +Parallelism notes: + - EP (expert parallelism) is orthogonal to TP/PP/DP/CP. + - Total GPUs = tp * pp * dp * cp * ep. + - With alltoall dispatcher, EP tokens are exchanged across ep ranks within each dp group. + - Sequence parallelism (--sp) requires --tp > 1. + +Usage: + python examples/train_moe_30b_mock.py --nnodes N --nproc-per-node N + [--tensorboard-dir DIR] + [--tp TP] [--pp PP] [--dp DP] [--cp CP] [--ep EP] [--sp] + [--seq-length N] [--micro-bs N] [--global-bs N] + [--use-tp-pp-dp-mapping] + [--master-addr ADDR] [--master-port PORT] + [--disable-background-mode] + +Must be run from the Megatron-LM root directory. +""" + +import argparse +import os +import subprocess +import sys + +MODEL_CONFIGS = { + "qwen3_1b": { + "num_layers": 16, + "hidden_size": 1024, + "num_attention_heads": 16, + "num_query_groups": 8, + "kv_channels": 64, + "ffn_hidden_size": 3584, + "num_experts": 4, + "moe_router_topk": 2, + "moe_ffn_hidden_size": 3584, + "seq_length": 512, + "max_position_embeddings": 2048, + "vocab_size": 151936, + }, + "qwen3_9b": { + "num_layers": 24, + "hidden_size": 2048, + "num_attention_heads": 32, + "num_query_groups": 8, + "kv_channels": 64, + "ffn_hidden_size": 7168, + "num_experts": 8, + "moe_router_topk": 2, + "moe_ffn_hidden_size": 7168, + "seq_length": 2048, + "max_position_embeddings": 2048, + "vocab_size": 151936, + }, +} + +MEGATRON_OPTIONAL_FLAGS = { + "--tensorboard-dir", + "--profile", + "--use-pytorch-profiler", + "--profile-step-start", + "--profile-step-end", + "--pytorch-profiler-collect-shapes", +} + + +def _detect_supported_megatron_flags() -> set[str]: + """Return optional CLI flags supported by this checkout's pretrain_gpt.py.""" + try: + proc = subprocess.run( + [sys.executable, "pretrain_gpt.py", "--help"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=False, + ) + except OSError: + return set() + help_text = f"{proc.stdout}\n{proc.stderr}" + return {flag for flag in MEGATRON_OPTIONAL_FLAGS if flag in help_text} + +def _setup_distributed(nnodes: int, master_addr: str, background_mode: bool): + env_node_rank = os.environ.get("NODE_RANK") or os.environ.get("GROUP_RANK") or "0" + + if background_mode: + resolved_master = master_addr or os.environ.get("MASTER_ADDR") or "127.0.0.1" + return int(env_node_rank), resolved_master + + if not master_addr: + raise ValueError("--master-addr is required when --disable-background-mode is set") + if nnodes == 1: + return 0, master_addr + return int(env_node_rank), master_addr + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--tensorboard-dir", default="tensorboard/moe_30b_mock", + help="Directory for TensorBoard logs and profiler traces.") + parser.add_argument( + "--model", + choices=tuple(MODEL_CONFIGS), + default="qwen3_1b", + help="Qwen model preset to run.", + ) + parser.add_argument("--micro-bs", type=int, default=16, + help="Micro batch size.") + parser.add_argument("--global-bs", type=int, default=64, + help="Global batch size.") + parser.add_argument("--train-iters", type=int, default=8, + help="Training iterations.") + parser.add_argument("--tp", type=int, default=1, + help="Tensor model parallel size.") + parser.add_argument("--pp", type=int, default=1, + help="Pipeline model parallel size.") + parser.add_argument("--dp", type=int, default=1, + help="Data parallel size.") + parser.add_argument("--cp", type=int, default=1, + help="Context parallel size.") + parser.add_argument("--ep", type=int, default=1, + help="Expert model parallel size.") + parser.add_argument("--seq-length", type=int, default=512, + help="Sequence length (default: 512).") + parser.add_argument("--sp", action="store_true", default=False, + help="Enable sequence parallelism (requires --tp > 1).") + parser.add_argument( + "--schedule", + choices=("1f1b", "interleaved1f1b"), + default="1f1b", + help="Pipeline schedule.", + ) + parser.add_argument( + "--zero-level", + choices=("zero1", "zero2", "zero3"), + default="zero1", + help="Megatron data-parallel sharding mode.", + ) + parser.add_argument("--use-tp-pp-dp-mapping", action="store_true", default=False, + help="Use tp-cp-ep-pp-dp rank ordering (PP intra-node) " + "instead of default tp-cp-ep-dp-pp (PP cross-node).") + parser.add_argument("--nnodes", type=int, required=True, + help="Number of nodes.") + parser.add_argument("--nproc-per-node", type=int, required=True, + help="GPUs per node.") + parser.add_argument("--master-addr", default=None, + help="Master node address (required when --disable-background-mode).") + parser.add_argument("--master-port", default="6000", + help="Master node port.") + parser.add_argument("--disable-background-mode", action="store_true", default=False, + help="Disable hostfile-based rank discovery; requires --master-addr.") + parser.add_argument("--nsight", action="store_true", default=False, + help="Enable Megatron/PyTorch profiler output under --tensorboard-dir.") + parser.add_argument("--rerun-mode", choices=("disabled", "validate_results", "report_stats"), default="disabled", + help="Megatron rerun engine mode. Defaults to disabled for benchmark timing.") + args = parser.parse_args() + if "--seq-length" not in os.sys.argv: + args.seq_length = MODEL_CONFIGS[args.model]["seq_length"] + return args + + +def build_training_args( + tensorboard_dir: str, + micro_bs: int, + global_bs: int, + seq_length: int, + tp: int, + pp: int, + dp: int, + cp: int, + ep: int, + sp: bool, + use_tp_pp_dp_mapping: bool, + model: str, + train_iters: int, + schedule: str, + zero_level: str, + enable_profiler: bool, + rerun_mode: str, + supported_flags: set[str] | None = None, +) -> list: + if sp and tp == 1: + raise ValueError("--sp requires --tp > 1") + + model_config = MODEL_CONFIGS[model] + supported_flags = supported_flags or set() + args = [ + "--num-layers", str(model_config["num_layers"]), + "--hidden-size", str(model_config["hidden_size"]), + "--num-attention-heads", str(model_config["num_attention_heads"]), + "--kv-channels", str(model_config["kv_channels"]), + "--ffn-hidden-size", str(model_config["ffn_hidden_size"]), + "--seq-length", str(seq_length), + "--max-position-embeddings", str(model_config["max_position_embeddings"]), + "--normalization", "RMSNorm", + "--norm-epsilon", "1e-6", + "--position-embedding-type", "rope", + "--rotary-percent", "1.0", + "--rotary-base", "1000000", + "--use-rotary-position-embeddings", + "--swiglu", + "--disable-bias-linear", + "--untie-embeddings-and-output-weights", + "--group-query-attention", + "--num-query-groups", str(model_config["num_query_groups"]), + "--qk-layernorm", + "--attention-dropout", "0.0", + "--hidden-dropout", "0.0", + "--use-mcore-models", + "--transformer-impl", "transformer_engine", + "--use-flash-attn", + "--num-experts", str(model_config["num_experts"]), + "--moe-router-topk", str(model_config["moe_router_topk"]), + "--moe-ffn-hidden-size", str(model_config["moe_ffn_hidden_size"]), + "--moe-layer-freq", "1", + "--moe-router-load-balancing-type", "aux_loss", + "--moe-aux-loss-coeff", "0.001", + "--moe-token-dispatcher-type", "alltoall", + "--moe-grouped-gemm", + "--tensor-model-parallel-size", str(tp), + "--pipeline-model-parallel-size", str(pp), + "--context-parallel-size", str(cp), + "--expert-model-parallel-size", str(ep), + "--micro-batch-size", str(micro_bs), + "--global-batch-size", str(global_bs), + "--train-iters", str(train_iters), + "--weight-decay", "0.1", + "--adam-beta1", "0.9", + "--adam-beta2", "0.95", + "--init-method-std", "0.01", + "--clip-grad", "1.0", + "--bf16", + "--lr", "3.0e-4", + "--lr-decay-style", "cosine", + "--min-lr", "3.0e-5", + "--lr-warmup-iters", "0", + "--lr-decay-iters", "2", + "--mock-data", + "--vocab-size", str(model_config["vocab_size"]), + "--tokenizer-type", "NullTokenizer", + "--log-interval", "1", + "--eval-interval", "100000", + "--eval-iters", "0", + "--save-interval", "100000", + "--no-gradient-accumulation-fusion", + "--rerun-mode", rerun_mode, + ] + if "--tensorboard-dir" in supported_flags: + args.extend(["--tensorboard-dir", tensorboard_dir]) + else: + print("[megatron] pretrain_gpt.py does not support --tensorboard-dir; skipping") + if enable_profiler: + profiler_args: list[str] = [] + if "--profile" in supported_flags: + profiler_args.append("--profile") + if "--use-pytorch-profiler" in supported_flags: + profiler_args.append("--use-pytorch-profiler") + if "--profile-step-start" in supported_flags: + profiler_args.extend(["--profile-step-start", "5"]) + if "--profile-step-end" in supported_flags: + profiler_args.extend(["--profile-step-end", "8"]) + if "--pytorch-profiler-collect-shapes" in supported_flags: + profiler_args.append("--pytorch-profiler-collect-shapes") + if profiler_args: + args.extend(profiler_args) + else: + print("[megatron] profiler flags unsupported by pretrain_gpt.py; running without profiler flags") + if dp > 1: + args.extend([ + "--use-distributed-optimizer", + "--overlap-grad-reduce", + "--overlap-param-gather", + "--data-parallel-sharding-strategy", + { + "zero1": "optim", + "zero2": "optim_grads", + "zero3": "optim_grads_params", + }[zero_level], + ]) + if zero_level in ("zero2", "zero3"): + args.append("--use-megatron-fsdp") + if schedule == "interleaved1f1b": + args.extend(["--num-layers-per-virtual-pipeline-stage", "1"]) + if sp: + args.append("--sequence-parallel") + if use_tp_pp_dp_mapping: + args.append("--use-tp-pp-dp-mapping") + return args + + +def main(): + # TODO: + # os.environ["CUDA_DEVICE_MAX_CONNECTIONS"] = "1" + + args = parse_args() + os.makedirs(args.tensorboard_dir, exist_ok=True) + + background_mode = not args.disable_background_mode + node_rank, master_addr = _setup_distributed( + args.nnodes, args.master_addr, background_mode, + ) + + print("=" * 60) + print(f"Node rank: {node_rank} / {args.nnodes}") + print(f"Master: {master_addr}:{args.master_port}") + print(f"GPUs/node: {args.nproc_per_node}") + print("=" * 60) + + training_args = build_training_args( + args.tensorboard_dir, args.micro_bs, args.global_bs, args.seq_length, + args.tp, args.pp, args.dp, args.cp, args.ep, args.sp, + args.use_tp_pp_dp_mapping, args.model, args.train_iters, args.schedule, args.zero_level, + args.nsight, args.rerun_mode, + _detect_supported_megatron_flags(), + ) + + from torch.distributed.run import main as torchrun_main + + # Use torchrun's static rendezvous path for multi-node runs so every node + # consistently connects to the fixed master_addr/master_port pair. + torchrun_args = [ + f"--nproc_per_node={args.nproc_per_node}", + f"--nnodes={args.nnodes}", + f"--node_rank={node_rank}", + "--rdzv_backend=static", # TODO: + f"--master_addr={master_addr}", + f"--master_port={args.master_port}", + ] + torchrun_main([ + *torchrun_args, + "pretrain_gpt.py", + *training_args, + ]) + + +if __name__ == "__main__": + main() diff --git a/artifact/scripts/make_piper_base_schedule.py b/artifact/scripts/make_piper_base_schedule.py new file mode 100755 index 0000000..5f0b8f3 --- /dev/null +++ b/artifact/scripts/make_piper_base_schedule.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate a Piper base schedule for Qwen e2e evals.") + parser.add_argument("--pp", type=int, required=True, help="Physical pipeline ranks.") + parser.add_argument("--dp", type=int, required=True, help="Data-parallel replicas.") + parser.add_argument("--virtual-stages", type=int, default=1) + parser.add_argument("--layout", choices=("linear", "v"), default="linear") + parser.add_argument("--zero-stage", type=int, default=1, choices=(0, 1, 2, 3)) + parser.add_argument("--ep", action="store_true") + parser.add_argument("--bucket-size", type=float, default=None) + parser.add_argument("--output", type=Path, required=True) + return parser.parse_args() + + +def devices_for_stage(stage: int, physical_pp: int, dp: int, layout: str) -> list[int]: + if layout == "v" and stage >= physical_pp: + physical_rank = (2 * physical_pp - 1 - stage) % physical_pp + else: + physical_rank = stage % physical_pp + return [physical_rank + replica * physical_pp for replica in range(dp)] + + +def main() -> int: + args = parse_args() + if args.pp <= 0 or args.dp <= 0 or args.virtual_stages <= 0: + raise SystemExit("--pp, --dp, and --virtual-stages must be positive") + + directives: list[dict] = [] + stage_count = args.pp * args.virtual_stages + for stage in range(stage_count): + devices = devices_for_stage(stage, args.pp, args.dp, args.layout) + directives.append( + { + "op": "place", + "filter": {"PP": stage}, + "devices": devices, + "stream": "pp_stream", + } + ) + + if args.dp > 1: + for stage in range(stage_count): + devices = devices_for_stage(stage, args.pp, args.dp, args.layout) + directive = { + "op": "replicate", + "filter": {"PP": stage}, + "devices": devices, + "reduce_stream": "reduce_stream", + } + if args.zero_stage == 2: + directive["shard_grads"] = True + elif args.zero_stage == 3: + directive["gather_stream"] = "gather_stream" + directive["shard_grads"] = True + directive["shard_params"] = True + directive["bucket_size"] = int(args.bucket_size or 1000) + elif args.bucket_size is not None: + directive["bucket_size"] = int(args.bucket_size) + directives.append(directive) + + if args.ep: + for stage in range(stage_count): + directives.append( + { + "op": "shard", + "filter": {"PP": stage, "EP": "*"}, + "devices": devices_for_stage(stage, args.pp, args.dp, args.layout), + "stream": "ep_stream", + } + ) + + args.output.parent.mkdir(parents=True, exist_ok=True) + with args.output.open("w", encoding="utf-8") as handle: + json.dump(directives, handle, indent=2) + handle.write("\n") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/artifact/scripts/run_deepspeed.sh b/artifact/scripts/run_deepspeed.sh new file mode 100755 index 0000000..d26590d --- /dev/null +++ b/artifact/scripts/run_deepspeed.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +set -euo pipefail + +NNODE=1 +NGPU=8 +NODE_RANK=0 +MASTER_ADDR=127.0.0.1 +MASTER_PORT=29501 +MODEL=qwen3_1b +TRAIN_ARGS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + --nnode) NNODE="$2"; shift 2 ;; + --ngpu) NGPU="$2"; shift 2 ;; + --node-rank) NODE_RANK="$2"; shift 2 ;; + --master-addr) MASTER_ADDR="$2"; shift 2 ;; + --master-port) MASTER_PORT="$2"; shift 2 ;; + --model) MODEL="$2"; shift 2 ;; + --) shift; TRAIN_ARGS=("$@"); break ;; + *) echo "Unknown argument: $1" >&2; exit 2 ;; + esac +done + +cd /workspace/piper +exec conda run -n deepspeed env \ + NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-}" \ + GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-}" \ + torchrun \ + "--nnodes=${NNODE}" \ + "--nproc_per_node=${NGPU}" \ + "--node_rank=${NODE_RANK}" \ + "--master_addr=${MASTER_ADDR}" \ + "--master_port=${MASTER_PORT}" \ + /workspace/artifact/run_deepspeed.py \ + --model "${MODEL}" \ + "${TRAIN_ARGS[@]}" diff --git a/artifact/scripts/run_megatron.sh b/artifact/scripts/run_megatron.sh new file mode 100755 index 0000000..143ac4b --- /dev/null +++ b/artifact/scripts/run_megatron.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +set -euo pipefail + +NNODE=1 +NGPU=8 +NODE_RANK=0 +MASTER_ADDR=127.0.0.1 +MASTER_PORT=29500 +MODEL=qwen3_1b +TP=1 +PP=1 +DP=1 +CP=1 +EP=1 +NSIGHT=false +TRAIN_ARGS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + --nnode) NNODE="$2"; shift 2 ;; + --ngpu) NGPU="$2"; shift 2 ;; + --node-rank) NODE_RANK="$2"; shift 2 ;; + --master-addr) MASTER_ADDR="$2"; shift 2 ;; + --master-port) MASTER_PORT="$2"; shift 2 ;; + --model) MODEL="$2"; shift 2 ;; + --tp) TP="$2"; shift 2 ;; + --pp) PP="$2"; shift 2 ;; + --dp) DP="$2"; shift 2 ;; + --cp) CP="$2"; shift 2 ;; + --ep) EP="$2"; shift 2 ;; + --nsight) NSIGHT=true; shift ;; + --) shift; TRAIN_ARGS=("$@"); break ;; + *) echo "Unknown argument: $1" >&2; exit 2 ;; + esac +done + +NSIGHT_ARGS=() +if $NSIGHT; then + NSIGHT_ARGS=(--nsight --tensorboard-dir "/tmp/megatron_profiler/${MODEL}_pp${PP}_dp${DP}_ep${EP}") +fi + +cd /workspace/Megatron-LM +exec conda run -n megatron env \ + NODE_RANK="${NODE_RANK}" \ + NCCL_P2P_DISABLE=1 \ + NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-}" \ + GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-}" \ + python /workspace/artifact/run_megatron.py \ + --model "${MODEL}" \ + --nnodes "${NNODE}" \ + --nproc-per-node "${NGPU}" \ + --master-addr "${MASTER_ADDR}" \ + --master-port "${MASTER_PORT}" \ + --disable-background-mode \ + --tp "${TP}" --pp "${PP}" --dp "${DP}" --cp "${CP}" --ep "${EP}" \ + "${NSIGHT_ARGS[@]}" \ + "${TRAIN_ARGS[@]}" diff --git a/artifact/scripts/run_piper.sh b/artifact/scripts/run_piper.sh new file mode 100755 index 0000000..dbed516 --- /dev/null +++ b/artifact/scripts/run_piper.sh @@ -0,0 +1,127 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL=1B +SCHEDULE=1f1b +PP=1 +DP=1 +EP=false +ZERO_STAGE=1 +BUCKET_SIZE="" +BATCH_SIZE=4 +SEQ_LEN=512 +MBS=1 +WARMUP=3 +ITERS=10 +ITERATION_SLEEP=0 +RAY_ADDRESS="" +RAY_PORT=6379 +TEMP_DIR=/tmp/piper/ray_tmp +METRICS_OUT="" +NSIGHT=false +USE_INDUCTOR_ARG=--use-inductor +PP_OUTER_ARG=--no-pp-outer + +while [[ $# -gt 0 ]]; do + case "$1" in + --model) MODEL="$2"; shift 2 ;; + --schedule) SCHEDULE="$2"; shift 2 ;; + --pp) PP="$2"; shift 2 ;; + --dp) DP="$2"; shift 2 ;; + --ep) EP=true; shift ;; + --zero-stage) ZERO_STAGE="$2"; shift 2 ;; + --bucket-size) BUCKET_SIZE="$2"; shift 2 ;; + --batch-size) BATCH_SIZE="$2"; shift 2 ;; + --seq-len) SEQ_LEN="$2"; shift 2 ;; + --mbs) MBS="$2"; shift 2 ;; + --warmup) WARMUP="$2"; shift 2 ;; + --iters) ITERS="$2"; shift 2 ;; + --iteration-sleep) ITERATION_SLEEP="$2"; shift 2 ;; + --address) RAY_ADDRESS="$2"; shift 2 ;; + --port) RAY_PORT="$2"; shift 2 ;; + --temp-dir) TEMP_DIR="$2"; shift 2 ;; + --metrics-out) METRICS_OUT="$2"; shift 2 ;; + --nsight) NSIGHT=true; shift ;; + --use-inductor) USE_INDUCTOR_ARG=--use-inductor; shift ;; + --no-use-inductor) USE_INDUCTOR_ARG=--no-use-inductor; shift ;; + --gradient-accumulation|--no-gradient-accumulation) shift ;; + --ar-a2a-same-stream|--no-ar-a2a-same-stream) shift ;; + --overlap-chunks|--no-overlap-chunks) shift ;; + --pp-outer) PP_OUTER_ARG=--pp-outer; shift ;; + --no-pp-outer) PP_OUTER_ARG=--no-pp-outer; shift ;; + *) echo "Unknown argument: $1" >&2; exit 2 ;; + esac +done + +case "$SCHEDULE" in + 1f1b) HARNESS_SCHEDULE=1f1b; VIRTUAL_STAGES=1 ;; + interleaved1f1b|interleaved-1f1b|interleaved_1f1b) HARNESS_SCHEDULE=interleaved_1f1b; VIRTUAL_STAGES=2 ;; + zerobubble|interleaved-zerobubble|interleavedzerobubble) HARNESS_SCHEDULE=zerobubble; VIRTUAL_STAGES=1 ;; + dualpipe|dualpipev) HARNESS_SCHEDULE=dualpipev; VIRTUAL_STAGES=2 ;; + *) echo "Unsupported Piper schedule: $SCHEDULE" >&2; exit 2 ;; +esac + +BASE_SCHEDULE="$(mktemp /tmp/piper-base-schedule.XXXXXX.json)" +BASE_ARGS=( + --pp "$PP" + --dp "$DP" + --virtual-stages "$VIRTUAL_STAGES" + --layout "$([[ "$HARNESS_SCHEDULE" == "dualpipev" ]] && echo v || echo linear)" + --zero-stage "$ZERO_STAGE" + --output "$BASE_SCHEDULE" +) +if $EP; then + BASE_ARGS+=(--ep) +fi +if [[ -n "$BUCKET_SIZE" ]]; then + BASE_ARGS+=(--bucket-size "$BUCKET_SIZE") +fi +python /workspace/artifact/scripts/make_piper_base_schedule.py "${BASE_ARGS[@]}" + +HARNESS_ARGS=( + examples/test_harness.py + --test-file examples/test_qwen.py + --base-schedule "$BASE_SCHEDULE" + --schedule "$HARNESS_SCHEDULE" + --ranks "$PP" + --mbs "$MBS" + --temp-dir "$TEMP_DIR" +) +if [[ -n "$RAY_ADDRESS" ]]; then + HARNESS_ARGS+=(--address "$RAY_ADDRESS" --port "$RAY_PORT") +fi + +TEST_ARGS=( + --model "$MODEL" + --batch-size "$BATCH_SIZE" + --seq-len "$SEQ_LEN" + --warmup "$WARMUP" + --iters "$ITERS" + --iteration-sleep "$ITERATION_SLEEP" + "$USE_INDUCTOR_ARG" + "$PP_OUTER_ARG" +) +if $NSIGHT; then + TEST_ARGS+=(--nsight) +fi + +cd /workspace/piper +export PYTHONPATH="/workspace/piper:/workspace/piper/examples:${PYTHONPATH:-}" +set +e +conda run -n piper python "${HARNESS_ARGS[@]}" "${TEST_ARGS[@]}" +STATUS=$? +set -e + +if [[ -n "$METRICS_OUT" ]]; then + mkdir -p "$(dirname "$METRICS_OUT")" + LATEST_RUN="$(ls -td out/* 2>/dev/null | head -1 || true)" + if [[ -n "$LATEST_RUN" && -f "$LATEST_RUN/results.csv" ]]; then + cp "$LATEST_RUN/results.csv" "$METRICS_OUT" + echo "Piper metrics copied to $METRICS_OUT" + else + echo "Piper metrics were not found under /workspace/piper/out" >&2 + fi +fi + +rm -f "$BASE_SCHEDULE" +exit "$STATUS" diff --git a/artifact/scripts/run_torchtitan.sh b/artifact/scripts/run_torchtitan.sh new file mode 100755 index 0000000..6b17058 --- /dev/null +++ b/artifact/scripts/run_torchtitan.sh @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +set -euo pipefail + +NNODE=1 +NGPU=8 +NODE_RANK=0 +MASTER_ADDR=127.0.0.1 +MASTER_PORT=29500 +MODULE=qwen3 +CONFIG=qwen3_9b +LOG_RANK=0 +NSIGHT=false +USE_BMM_EXPERTS=false +TRAIN_ARGS=() + +while [[ $# -gt 0 ]]; do + case "$1" in + --nnode) NNODE="$2"; shift 2 ;; + --ngpu) NGPU="$2"; shift 2 ;; + --node-rank) NODE_RANK="$2"; shift 2 ;; + --master-addr) MASTER_ADDR="$2"; shift 2 ;; + --master-port) MASTER_PORT="$2"; shift 2 ;; + --module) MODULE="$2"; shift 2 ;; + --config) CONFIG="$2"; shift 2 ;; + --log-rank) LOG_RANK="$2"; shift 2 ;; + --nsight) NSIGHT=true; shift ;; + --use-bmm-experts) USE_BMM_EXPERTS=true; shift ;; + --) shift; TRAIN_ARGS=("$@"); break ;; + *) echo "Unknown argument: $1" >&2; exit 2 ;; + esac +done + +RUNNER=( + torchrun + "--nnodes=${NNODE}" + "--nproc_per_node=${NGPU}" + "--node_rank=${NODE_RANK}" + "--master_addr=${MASTER_ADDR}" + "--master_port=${MASTER_PORT}" + "--local-ranks-filter" "${LOG_RANK}" + "--role" "rank" + "--tee" "3" + -m torchtitan.train + --module "${MODULE}" + --config "${CONFIG}" + "${TRAIN_ARGS[@]}" +) + +if $NSIGHT; then + mkdir -p /workspace/eval-out/nsight + RUNNER=(nsys profile --trace=cuda,nvtx,osrt --force-overwrite=true --output "/workspace/eval-out/nsight/torchtitan_node${NODE_RANK}" "${RUNNER[@]}") +fi + +cd /workspace/torchtitan +exec conda run -n torchtitan env \ + PYTHONPATH="/workspace/artifact:${PYTHONPATH:-}" \ + TORCHTITAN_USE_BMM_EXPERTS="$($USE_BMM_EXPERTS && echo 1 || echo 0)" \ + PYTORCH_ALLOC_CONF=expandable_segments:True \ + TORCHFT_LIGHTHOUSE=http://localhost:29510 \ + NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-}" \ + GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-}" \ + "${RUNNER[@]}" diff --git a/artifact/sitecustomize.py b/artifact/sitecustomize.py new file mode 100644 index 0000000..1d6d302 --- /dev/null +++ b/artifact/sitecustomize.py @@ -0,0 +1,64 @@ +"""Runtime hooks for artifact-only TorchTitan experiments.""" + +import os + + +if os.environ.get("TORCHTITAN_USE_BMM_EXPERTS") == "1": + try: + import torch + import torch.nn.functional as F + from torch.distributed.tensor import DTensor + + from torchtitan.models.common.moe import moe as moe_mod + + def _run_experts_bmm( + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, + ) -> torch.Tensor: + dim = x.shape[-1] + num_experts = w1.shape[0] + counts = num_tokens_per_expert.to(device=x.device, dtype=torch.long) + capacity = int(counts.max().item()) if counts.numel() else 0 + starts = torch.cumsum(counts, dim=0) - counts + positions = torch.arange(capacity, device=x.device, dtype=torch.long) + zero = x.new_zeros((capacity, dim)) + + packed_rows = [] + for expert_idx in range(num_experts): + valid = positions < counts[expert_idx] + source_idx = starts[expert_idx] + positions + safe_idx = torch.where(valid, source_idx, positions.new_zeros(())) + packed_rows.append(torch.where(valid[:, None], x[safe_idx], zero)) + packed = torch.stack(packed_rows, dim=0) + + h = F.silu(torch.bmm(packed, w1.transpose(-2, -1))) + h = h * torch.bmm(packed, w3.transpose(-2, -1)) + expert_out = torch.bmm(h, w2.transpose(-2, -1)) + + out = x.new_zeros(x.shape) + for expert_idx in range(num_experts): + valid = positions < counts[expert_idx] + dest_idx = starts[expert_idx] + positions + safe_idx = torch.where(valid, dest_idx, positions.new_zeros(())) + src = torch.where(valid[:, None], expert_out[expert_idx], zero) + out = out.scatter_add(0, safe_idx[:, None].expand(-1, dim), src) + return out + + def _bmm_experts_forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: + if isinstance(self.w1, DTensor): + w1 = self.w1.to_local() + w2 = self.w2.to_local() + w3 = self.w3.to_local() + else: + w1 = self.w1 + w2 = self.w2 + w3 = self.w3 + return _run_experts_bmm(w1, w2, w3, x, num_tokens_per_expert) + + moe_mod.GroupedExperts.forward = _bmm_experts_forward + print("[artifact] TorchTitan GroupedExperts.forward patched to BMM experts", flush=True) + except Exception as exc: + raise RuntimeError("Failed to install TorchTitan BMM experts artifact hook") from exc diff --git a/examples/__init__.py b/examples/__init__.py new file mode 100644 index 0000000..11e4542 --- /dev/null +++ b/examples/__init__.py @@ -0,0 +1 @@ +"""Runnable Piper examples and benchmark harnesses.""" diff --git a/examples/base-schedules/pp2_dp2_ep2.json b/examples/base-schedules/pp2_dp2_ep2.json new file mode 100644 index 0000000..6baf5f8 --- /dev/null +++ b/examples/base-schedules/pp2_dp2_ep2.json @@ -0,0 +1,70 @@ +[ + { + "op": "place", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "shard", + "filter": { + "PP": 0, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 1, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "ep_stream" + } +] diff --git a/examples/base-schedules/pp2_dp2_ep2_bucket100.json b/examples/base-schedules/pp2_dp2_ep2_bucket100.json new file mode 100644 index 0000000..3956e3d --- /dev/null +++ b/examples/base-schedules/pp2_dp2_ep2_bucket100.json @@ -0,0 +1,72 @@ +[ + { + "op": "place", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "reduce_stream": "dp_stream", + "bucket_size": 100 + }, + { + "op": "replicate", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "reduce_stream": "dp_stream", + "bucket_size": 100 + }, + { + "op": "shard", + "filter": { + "PP": 0, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 1, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "ep_stream" + } +] diff --git a/examples/base-schedules/pp2_dp2_ep2_custom_order.json b/examples/base-schedules/pp2_dp2_ep2_custom_order.json new file mode 100644 index 0000000..51cdb0e --- /dev/null +++ b/examples/base-schedules/pp2_dp2_ep2_custom_order.json @@ -0,0 +1,140 @@ +[ + { + "op": "place", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "shard", + "filter": { + "PP": 0, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 1, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "ep_stream" + }, + { + "op": "split", + "filter": {}, + "dim_name": "MB", + "num_microbatches": 2 + }, + { + "op": "order", + "filters": [ + [ + { + "PP": 0, + "MB": 0, + "PASS": "F" + } + ], + [ + { + "PP": 0, + "MB": 1, + "PASS": "F" + } + ], + [ + { + "PP": 0, + "MB": 0, + "PASS": "B" + } + ], + [ + { + "PP": 0, + "MB": 1, + "PASS": "B" + } + ] + ] + }, + { + "op": "order", + "filters": [ + [ + { + "PP": 1, + "MB": 0, + "PASS": "F" + } + ], + [ + { + "PP": 1, + "MB": 0, + "PASS": "B" + }, + { + "PP": 1, + "MB": 1, + "PASS": "F" + } + ], + [ + { + "PP": 1, + "MB": 1, + "PASS": "B" + } + ] + ] + } +] diff --git a/examples/base-schedules/pp2_dp2_ep2_zero2.json b/examples/base-schedules/pp2_dp2_ep2_zero2.json new file mode 100644 index 0000000..be2379d --- /dev/null +++ b/examples/base-schedules/pp2_dp2_ep2_zero2.json @@ -0,0 +1,72 @@ +[ + { + "op": "place", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "reduce_stream": "dp_stream", + "shard_grads": true + }, + { + "op": "replicate", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "reduce_stream": "dp_stream", + "shard_grads": true + }, + { + "op": "shard", + "filter": { + "PP": 0, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "dp_stream" + }, + { + "op": "shard", + "filter": { + "PP": 1, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "dp_stream" + } +] diff --git a/examples/base-schedules/pp2_dp2_ep2_zero3.json b/examples/base-schedules/pp2_dp2_ep2_zero3.json new file mode 100644 index 0000000..1b0939c --- /dev/null +++ b/examples/base-schedules/pp2_dp2_ep2_zero3.json @@ -0,0 +1,76 @@ +[ + { + "op": "place", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "gather_stream": "dp_stream", + "reduce_stream": "dp_stream", + "shard_grads": true, + "shard_params": true + }, + { + "op": "replicate", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "gather_stream": "dp_stream", + "reduce_stream": "dp_stream", + "shard_grads": true, + "shard_params": true + }, + { + "op": "shard", + "filter": { + "PP": 0, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "dp_stream" + }, + { + "op": "shard", + "filter": { + "PP": 1, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "dp_stream" + } +] diff --git a/examples/base-schedules/pp4_dp2_ep2_interleaved.json b/examples/base-schedules/pp4_dp2_ep2_interleaved.json new file mode 100644 index 0000000..11e335c --- /dev/null +++ b/examples/base-schedules/pp4_dp2_ep2_interleaved.json @@ -0,0 +1,138 @@ +[ + { + "op": "place", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 2 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 3 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 2 + }, + "devices": [ + 0, + 2 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 3 + }, + "devices": [ + 1, + 3 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "shard", + "filter": { + "PP": 0, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 1, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 2, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 3, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "ep_stream" + } +] diff --git a/examples/base-schedules/pp4_dp2_ep2_v_placement.json b/examples/base-schedules/pp4_dp2_ep2_v_placement.json new file mode 100644 index 0000000..7e2fae2 --- /dev/null +++ b/examples/base-schedules/pp4_dp2_ep2_v_placement.json @@ -0,0 +1,138 @@ +[ + { + "op": "place", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 2 + }, + "devices": [ + 1, + 3 + ], + "stream": "pp_stream" + }, + { + "op": "place", + "filter": { + "PP": 3 + }, + "devices": [ + 0, + 2 + ], + "stream": "pp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 0 + }, + "devices": [ + 0, + 2 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 1 + }, + "devices": [ + 1, + 3 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 2 + }, + "devices": [ + 1, + 3 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "replicate", + "filter": { + "PP": 3 + }, + "devices": [ + 0, + 2 + ], + "reduce_stream": "dp_stream" + }, + { + "op": "shard", + "filter": { + "PP": 0, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 1, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 2, + "EP": "*" + }, + "devices": [ + 1, + 3 + ], + "stream": "ep_stream" + }, + { + "op": "shard", + "filter": { + "PP": 3, + "EP": "*" + }, + "devices": [ + 0, + 2 + ], + "stream": "ep_stream" + } +] diff --git a/examples/build_schedule.py b/examples/build_schedule.py new file mode 100644 index 0000000..b25439c --- /dev/null +++ b/examples/build_schedule.py @@ -0,0 +1,334 @@ +"""Build JSON order directives for common pipeline schedules.""" + +from __future__ import annotations + +from collections import defaultdict +from dataclasses import dataclass +from typing import Literal + +Pass = Literal["F", "B", "BI", "BW"] + + +@dataclass(frozen=True) +class _Op: + pp: int + mb: int + pass_: Pass + + +# A "slot" is one timestep in a rank's row. A slot holds one op normally; when +# multiple ops are nested in the same order-directive filter group they share a +# slot and visualize as a vertically-stacked cell, indicating that they run +# interleaved within the same scheduling block. +_Slot = list[_Op] + + +def _filter(pp: int, mb: int, pass_: Pass) -> dict[str, int | str]: + return {"PP": pp, "MB": mb, "PASS": pass_} + + +def _order_directive(ops: list[_Op]) -> dict: + return _order_directive_from_slots([[op] for op in ops]) + + +def _order_directive_from_slots(slots: list[_Slot]) -> dict: + return { + "op": "order", + "filters": [ + [_filter(op.pp, op.mb, op.pass_) for op in slot] + for slot in slots + ], + } + + +def build_1f1b_schedule(n_ranks: int, n_mbs: int) -> list[dict]: + """Return one 1F1B order directive per rank. + + This preserves the old PipelineSchedule builder's non-None per-rank op order: + rank r warms up with ``n_ranks - 1 - r`` forwards, then alternates forward + and fused backward until all microbatches have completed. + """ + _validate_positive("n_ranks", n_ranks) + _validate_positive("n_mbs", n_mbs) + + rows: list[list[_Op]] = [[] for _ in range(n_ranks)] + fwd_mb = [0] * n_ranks + bwd_mb = [0] * n_ranks + + for rank in range(n_ranks): + warmup = n_ranks - 1 - rank + for _ in range(min(warmup, n_mbs)): + rows[rank].append(_Op(pp=rank, mb=fwd_mb[rank], pass_="F")) + fwd_mb[rank] += 1 + + for rank in range(n_ranks): + while bwd_mb[rank] < n_mbs: + if fwd_mb[rank] < n_mbs: + rows[rank].append(_Op(pp=rank, mb=fwd_mb[rank], pass_="F")) + fwd_mb[rank] += 1 + rows[rank].append(_Op(pp=rank, mb=bwd_mb[rank], pass_="B")) + bwd_mb[rank] += 1 + + return [_order_directive(row) for row in rows] + + +def build_zerobubble_schedule(n_ranks: int, n_mbs: int) -> list[dict]: + """Return one ZeroBubble order directive per rank. + + This is the old ZB-1 schedule builder expressed as order directives. Each + rank owns one pipeline stage, fused backward is split into BI/BW, and BW is + deferred by ``rank`` backward-input steps. + """ + _validate_positive("n_ranks", n_ranks) + _validate_positive("n_mbs", n_mbs) + + rows: list[list[_Op]] = [] + for rank in range(n_ranks): + warmup = min(n_ranks - 1 - rank, n_mbs) + fwd_bwd_ops = n_mbs - warmup + cooldown_ops = n_mbs - fwd_bwd_ops + deferred_bw_count = rank + + fwd_mb = 0 + bwd_mb = 0 + bw_mb = 0 + bwdi_count = 0 + emitted_bw_count = 0 + ops: list[_Op] = [] + + for op_idx in range(warmup + fwd_bwd_ops + cooldown_ops): + if op_idx < warmup: + ops.append(_Op(pp=rank, mb=fwd_mb, pass_="F")) + fwd_mb += 1 + continue + + if op_idx < warmup + fwd_bwd_ops: + ops.append(_Op(pp=rank, mb=fwd_mb, pass_="F")) + fwd_mb += 1 + + ops.append(_Op(pp=rank, mb=bwd_mb, pass_="BI")) + bwd_mb += 1 + bwdi_count += 1 + + if bwdi_count > deferred_bw_count: + ops.append(_Op(pp=rank, mb=bw_mb, pass_="BW")) + bw_mb += 1 + emitted_bw_count += 1 + + while emitted_bw_count < bwdi_count: + ops.append(_Op(pp=rank, mb=bw_mb, pass_="BW")) + bw_mb += 1 + emitted_bw_count += 1 + + rows.append(ops) + + return [_order_directive(row) for row in rows] + + +def build_interleaved_1f1b_schedule( + n_ranks: int, + n_mbs: int, + n_virtual_stages: int, +) -> list[dict]: + """Return one interleaved 1F1B order directive per physical rank. + + Rank ``r`` owns virtual pipeline stages ``r + k * n_ranks``. The op order is + the old interleaved PipelineSchedule order with the placeholder ``None`` slots + omitted. + """ + _validate_positive("n_ranks", n_ranks) + _validate_positive("n_mbs", n_mbs) + _validate_positive("n_virtual_stages", n_virtual_stages) + + microbatches_per_round = n_ranks + + def warmup_ops(rank: int) -> int: + return min( + (n_virtual_stages - 1) * n_ranks + 2 * (n_ranks - 1 - rank), + n_mbs * n_virtual_stages, + ) + + def forward_stage_index(step: int, rank: int) -> int: + local_index = (step // microbatches_per_round) % n_virtual_stages + return local_index * n_ranks + rank + + def backward_stage_index(step: int, rank: int, warmup: int) -> int: + local_index = ( + n_virtual_stages + - 1 + - ((step - warmup) // microbatches_per_round) % n_virtual_stages + ) + return local_index * n_ranks + rank + + rows: list[list[_Op]] = [] + for rank in range(n_ranks): + warmup = warmup_ops(rank) + microbatch_ops = n_virtual_stages * n_mbs + fwd_bwd_ops = microbatch_ops - warmup + cooldown_ops = microbatch_ops - fwd_bwd_ops + + fwd_mb: dict[int, int] = defaultdict(int) + bwd_mb: dict[int, int] = defaultdict(int) + ops: list[_Op] = [] + + for op_idx in range(warmup + fwd_bwd_ops + cooldown_ops): + if op_idx < warmup: + stage = forward_stage_index(op_idx, rank) + mb = fwd_mb[stage] + fwd_mb[stage] += 1 + ops.append(_Op(pp=stage, mb=mb, pass_="F")) + elif op_idx < warmup + fwd_bwd_ops: + stage = forward_stage_index(op_idx, rank) + mb = fwd_mb[stage] + fwd_mb[stage] += 1 + ops.append(_Op(pp=stage, mb=mb, pass_="F")) + + stage = backward_stage_index(op_idx, rank, warmup) + mb = bwd_mb[stage] + bwd_mb[stage] += 1 + ops.append(_Op(pp=stage, mb=mb, pass_="B")) + else: + stage = backward_stage_index(op_idx, rank, warmup) + mb = bwd_mb[stage] + bwd_mb[stage] += 1 + ops.append(_Op(pp=stage, mb=mb, pass_="B")) + + rows.append(ops) + + return [_order_directive(row) for row in rows] + + +def build_dualpipev_schedule(n_ranks: int, n_mbs: int) -> list[dict]: + """Return DualPipeV order directives with overlapped FWD/BWD slots. + + Rank ``r`` owns V-layout stages ``r`` and ``2 * n_ranks - 1 - r``. The + old ``FWD_BWD`` chunk is represented as one nested order filter group + containing the corresponding FWD filter and fused BWD filter. + """ + _validate_positive("n_ranks", n_ranks) + _validate_positive("n_mbs", n_mbs) + if n_mbs < 2 * n_ranks: + raise ValueError( + f"dualpipev requires n_mbs >= 2 * n_ranks, got n_mbs={n_mbs}, " + f"n_ranks={n_ranks}" + ) + + rows: list[list[_Slot]] = [[] for _ in range(n_ranks)] + + for rank in range(n_ranks): + s0 = rank + s1 = 2 * n_ranks - 1 - rank + slots = rows[rank] + counts: dict[tuple[int, str], int] = {} + weight_queue: list[tuple[int, int]] = [] + + def count(stage: int, key: str) -> int: + return counts.get((stage, key), 0) + + def inc(stage: int, key: str) -> None: + counts[(stage, key)] = count(stage, key) + 1 + + def inc_bwd(stage: int) -> None: + inc(stage, "i") + inc(stage, "w") + + def append_op(stage: int, mb: int, pass_: Pass) -> None: + slots.append([_Op(pp=stage, mb=mb, pass_=pass_)]) + + def fwd(stage: int) -> None: + mb = count(stage, "f") + append_op(stage, mb, "F") + inc(stage, "f") + + def bwdi(stage: int) -> None: + mb = count(stage, "i") + append_op(stage, mb, "BI") + weight_queue.append((stage, mb)) + inc(stage, "i") + + def drain_w() -> None: + if not weight_queue: + return + stage, mb = weight_queue.pop(0) + append_op(stage, mb, "BW") + inc(stage, "w") + + def full_bwd(stage: int) -> None: + mb = count(stage, "i") + append_op(stage, mb, "B") + inc_bwd(stage) + + def overlap_fb(fwd_stage: int, bwd_stage: int) -> None: + fwd_mb = count(fwd_stage, "f") + bwd_mb = count(bwd_stage, "i") + slots.append([ + _Op(pp=fwd_stage, mb=fwd_mb, pass_="F"), + _Op(pp=bwd_stage, mb=bwd_mb, pass_="B"), + ]) + inc(fwd_stage, "f") + inc_bwd(bwd_stage) + + # Phase 1: F0 warmup. + for _ in range((n_ranks - rank - 1) * 2): + fwd(s0) + + # Phase 2: F0F1. + for _ in range(rank + 1): + fwd(s0) + fwd(s1) + + # Phase 3: I1 W1 F1. + for _ in range(n_ranks - rank - 1): + bwdi(s1) + drain_w() + fwd(s1) + + # Phase 4: Main overlapped F0B1 + F1B0. + for i in range(n_mbs - n_ranks * 2 + rank + 1): + if i == 0 and rank == n_ranks - 1: + fwd(s0) + full_bwd(s1) + else: + overlap_fb(s0, s1) + overlap_fb(s1, s0) + + # Phase 5: B1 + F1B0. + for _ in range(n_ranks - rank - 1): + full_bwd(s1) + overlap_fb(s1, s0) + + # Phase 6: B1B0, switching the second half to split backward. + enable_zb = False + for i in range(rank + 1): + if i == (rank + 1) // 2 and rank % 2 == 1: + enable_zb = True + (bwdi if enable_zb else full_bwd)(s1) + if i == (rank + 1) // 2 and rank % 2 == 0: + enable_zb = True + (bwdi if enable_zb else full_bwd)(s0) + + # Phase 7: W0 B0. + for _ in range(n_ranks - rank - 1): + drain_w() + (bwdi if enable_zb else full_bwd)(s0) + + # Phase 8: W0. + for _ in range(rank + 1): + drain_w() + + return [_order_directive_from_slots(row) for row in rows] + + +def _validate_positive(name: str, value: int) -> None: + if not isinstance(value, int) or value <= 0: + raise ValueError(f"{name} must be a positive integer, got {value!r}") + + +def main() -> None: + from test_harness import main as harness_main + + harness_main() + + +if __name__ == "__main__": + main() diff --git a/examples/models/__init__.py b/examples/models/__init__.py new file mode 100644 index 0000000..88ecaf4 --- /dev/null +++ b/examples/models/__init__.py @@ -0,0 +1 @@ +"""Example model definitions used by Piper harness programs.""" diff --git a/test/models/llama.py b/examples/models/llama3.py similarity index 82% rename from test/models/llama.py rename to examples/models/llama3.py index 04ce362..78c3062 100644 --- a/test/models/llama.py +++ b/examples/models/llama3.py @@ -13,8 +13,7 @@ import torch.nn.functional as F from torch import Tensor, nn from torch.nn.utils import parameters_to_vector - -from src.piper import distributed_stage +from src.piper import annotate # import fairscale.nn.model_parallel.initialize as fs_init # from fairscale.nn.model_parallel.layers import ( @@ -25,6 +24,7 @@ logger = logging.getLogger(__name__) +PP_TAG = "PP" @dataclass @@ -101,6 +101,20 @@ class ModelArgs: max_seq_len=2048, ) +LLAMA_70B = ModelArgs( + dim=8192, + n_layers=80, + n_heads=64, + n_kv_heads=8, + vocab_size=128256, + multiple_of=4096, + ffn_dim_multiplier=1.3, + norm_eps=1e-5, + rope_theta=500000, + max_batch_size=32, + max_seq_len=8192, +) + class RMSNorm(torch.nn.Module): def __init__(self, dim: int, eps: float = 1e-6): @@ -345,17 +359,19 @@ def forward( return out -def partition(*args): - pass - - class Transformer(nn.Module): - def __init__(self, params: ModelArgs, seq_len: int): + def __init__(self, params: ModelArgs, seq_len: int, num_stages: int | None = None): super().__init__() self.params = params self.vocab_size = params.vocab_size self.n_layers = params.n_layers self.seq_len = seq_len + if num_stages is None: + from src.state import piper_metadata + + schedule_info = piper_metadata.schedule_info or {} + num_stages = int(schedule_info.get("num_stages", schedule_info.get("pp_degree", 2))) + self.num_stages = num_stages def log_size(layer, indent=0): num_params = sum(p.numel() for p in layer.parameters()) @@ -389,6 +405,7 @@ def log_size(layer, indent=0): ) log_size(self.output) + # Register freq_cis and mask as buffers so they are moved with the model self.freqs_cis = precompute_freqs_cis( params.dim // params.n_heads, self.seq_len, @@ -399,87 +416,17 @@ def log_size(layer, indent=0): mask = torch.triu(mask, diagonal=1) self.mask = torch.hstack([torch.zeros((self.seq_len, 0)), mask]) - # """ - # forward method for interleaved-1f1b schedule - # requires: - # - 2 devices - # - 4 stages - # - n_layers is divisible by 4 - # """ - # def forward(self, tokens: torch.Tensor): - - # distributed_stage(0, actor_id=0, optim=torch.optim.Adam) - - # h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens - # start_pos = 0 - - # for layer in self.layers[:self.n_layers//4]: - # h = layer(h, start_pos, self.freqs_cis, self.mask) - - # distributed_stage(1, actor_id=1, optim=torch.optim.Adam) - - # for layer in self.layers[self.n_layers//4:self.n_layers//2]: - # h = layer(h, start_pos, self.freqs_cis, self.mask) - - # distributed_stage(2, actor_id=0, optim=torch.optim.Adam) - - # for layer in self.layers[self.n_layers//2:3*self.n_layers//4]: - # h = layer(h, start_pos, self.freqs_cis, self.mask) - - # distributed_stage(3, actor_id=1, optim=torch.optim.Adam) - - # for layer in self.layers[3*self.n_layers//4:]: - # h = layer(h, start_pos, self.freqs_cis, self.mask) - - # h = self.norm(h) if self.norm else h - # output = self.output(h).float() if self.output else h - - # return output - - """ - forward method for 1f1b schedule - requires: - - 2 devices - - 2 stages - - n_layers is divisible by 2 - """ def forward(self, tokens: torch.Tensor): - - distributed_stage(0, actor_id=0) - - h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens start_pos = 0 - - for layer in self.layers[:self.n_layers//2]: - h = layer(h, start_pos, self.freqs_cis, self.mask) - - distributed_stage(1, actor_id=1) - - for layer in self.layers[self.n_layers//2:]: - h = layer(h, start_pos, self.freqs_cis, self.mask) - - h = self.norm(h) if self.norm else h - output = self.output(h).float() if self.output else h - + for stage_id in range(self.num_stages): + layer_start = stage_id * self.n_layers // self.num_stages + layer_end = (stage_id + 1) * self.n_layers // self.num_stages + with annotate(PP_TAG): + if stage_id == 0: + h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens + for layer in self.layers[layer_start:layer_end]: + h = layer(h, start_pos, self.freqs_cis, self.mask) + if stage_id == self.num_stages - 1: + h = self.norm(h) if self.norm else h + output = self.output(h).float() if self.output else h return output - - # """ - # forward method for no pp - # requires: - # - 1 device - # - 1 stages - # """ - # def forward(self, tokens: torch.Tensor): - - # distributed_stage(0, actor_id=0, optim=torch.optim.Adam) - - # h = self.tok_embeddings(tokens) if self.tok_embeddings else tokens - # start_pos = 0 - - # for layer in self.layers: - # h = layer(h, start_pos, self.freqs_cis, self.mask) - - # h = self.norm(h) if self.norm else h - # output = self.output(h).float() if self.output else h - - # return output \ No newline at end of file diff --git a/examples/models/qwen3.py b/examples/models/qwen3.py new file mode 100644 index 0000000..8ef8903 --- /dev/null +++ b/examples/models/qwen3.py @@ -0,0 +1,466 @@ +from typing import Optional +import torch +import torch.nn.functional as F +from torch import nn +from torch.nn.attention import sdpa_kernel, SDPBackend + +from torchtitan.models.qwen3 import Qwen3Model, Qwen3ModelArgs +from torchtitan.models.qwen3.model.model import TransformerBlock +from torchtitan.models.moe import MoE, MoEArgs +from torchtitan.protocols.model import AttentionMasksType +from src.piper import annotate + +PP_TAG = "PP" +EP_TAG = "EP" + + +def _run_bmm_experts( + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + if x.is_meta: + return torch.empty_like(x) + + dim = x.shape[-1] + num_experts = w1.shape[0] + counts = num_tokens_per_expert.to(device=x.device, dtype=torch.long) + # Keep capacity shape-derived so torch.compile never has to materialize a + # routed-token count as a Python scalar. + capacity = x.shape[0] + starts = torch.cumsum(counts, dim=0) - counts + positions = torch.arange(capacity, device=x.device, dtype=torch.long) + zero = x.new_zeros((capacity, dim)) + + packed_rows = [] + for expert_idx in range(num_experts): + valid = positions < counts[expert_idx] + source_idx = starts[expert_idx] + positions + safe_idx = torch.where(valid, source_idx, positions.new_zeros(())) + packed_rows.append(torch.where(valid[:, None], x[safe_idx], zero)) + packed = torch.stack(packed_rows, dim=0) + + h = F.silu(torch.bmm(packed, w1.transpose(-2, -1))) + h = h * torch.bmm(packed, w3.transpose(-2, -1)) + expert_out = torch.bmm(h, w2.transpose(-2, -1)) + + out = x.new_zeros(x.shape) + for expert_idx in range(num_experts): + valid = positions < counts[expert_idx] + dest_idx = starts[expert_idx] + positions + safe_idx = torch.where(valid, dest_idx, positions.new_zeros(())) + src = torch.where(valid[:, None], expert_out[expert_idx], zero) + out = out.scatter_add(0, safe_idx[:, None].expand(-1, dim), src) + return out + + +_BMM_EXPERTS_LIBS = [] + + +def _register_bmm_experts_op() -> None: + if ( + hasattr(torch.ops, "piper_artifact") + and hasattr(torch.ops.piper_artifact, "bmm_experts") + ): + return + + def_lib = torch.library.Library("piper_artifact", "DEF") + def_lib.define( + "bmm_experts(Tensor w1, Tensor w2, Tensor w3, Tensor x, Tensor num_tokens_per_expert) -> Tensor" + ) + impl_lib = torch.library.Library("piper_artifact", "IMPL") + impl_lib.impl("bmm_experts", _run_bmm_experts, "CPU") + impl_lib.impl("bmm_experts", _run_bmm_experts, "CUDA") + impl_lib.impl("bmm_experts", _run_bmm_experts, "Autograd") + impl_lib.impl("bmm_experts", _run_bmm_experts, "Meta") + _BMM_EXPERTS_LIBS.extend([def_lib, impl_lib]) + + +_register_bmm_experts_op() + + +def bmm_experts( + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + x: torch.Tensor, + num_tokens_per_expert: torch.Tensor, +) -> torch.Tensor: + return torch.ops.piper_artifact.bmm_experts( + w1, w2, w3, x, num_tokens_per_expert + ) + + +def create_qwen3_config(name: str) -> Qwen3ModelArgs: + """Create Qwen3 model config based on name.""" + match name: + case '9M': + return Qwen3ModelArgs( + vocab_size=2048, + dim=256, + n_layers=4, + n_heads=8, + n_kv_heads=4, + head_dim=32, + hidden_dim=512, + norm_eps=1e-6, + qk_norm=True, + max_seq_len=2048, + rope_theta=1000000.0, + depth_init=True, + enable_weight_tying=False, + moe_enabled=True, + moe_inter_dim=128, + moe_args=MoEArgs( + num_experts=8, + top_k=2, + use_grouped_mm=True, + num_expert_groups=None, + num_limited_groups=None, + score_func="softmax", + route_norm=False, + route_scale=1.0, + gate_bias=False, + score_before_experts=False, + num_shared_experts=0, + load_balance_coeff=None, + _debug_force_load_balance=False, + ), + ) + case '1B': + return Qwen3ModelArgs( + vocab_size=151936, + dim=1024, + n_layers=16, + n_heads=16, + n_kv_heads=8, + head_dim=64, + hidden_dim=3584, + norm_eps=1e-6, + qk_norm=True, + max_seq_len=2048, + rope_theta=1000000.0, + depth_init=True, + enable_weight_tying=False, + moe_enabled=True, + moe_inter_dim=3584, + moe_args=MoEArgs( + num_experts=4, + top_k=2, + use_grouped_mm=True, + num_expert_groups=None, + num_limited_groups=None, + score_func="softmax", + route_norm=False, + route_scale=1.0, + gate_bias=False, + score_before_experts=False, + num_shared_experts=0, + load_balance_coeff=None, + _debug_force_load_balance=False, + ), + ) + case '9B': + return Qwen3ModelArgs( + vocab_size=151936, + dim=2048, + n_layers=4, + n_heads=32, + n_kv_heads=8, + head_dim=64, + hidden_dim=7168, + norm_eps=1e-6, + qk_norm=True, + max_seq_len=2048, + rope_theta=1000000.0, + depth_init=True, + enable_weight_tying=False, + moe_enabled=True, + moe_inter_dim=7168, + moe_args=MoEArgs( + num_experts=8, + top_k=2, + use_grouped_mm=True, + num_expert_groups=None, + num_limited_groups=None, + score_func="softmax", + route_norm=False, + route_scale=1.0, + gate_bias=False, + score_before_experts=False, + num_shared_experts=0, + load_balance_coeff=None, + _debug_force_load_balance=False, + ), + ) + case '48B': + return Qwen3ModelArgs( + vocab_size=151936, + dim=4096, + n_layers=4, #32, + n_heads=32, + n_kv_heads=8, + head_dim=128, + hidden_dim=14336, + norm_eps=1e-6, + qk_norm=True, + max_seq_len=2048, + rope_theta=1000000.0, + depth_init=True, + enable_weight_tying=False, + moe_enabled=True, + moe_inter_dim=14336, + moe_args=MoEArgs( + num_experts=8, + top_k=2, + use_grouped_mm=True, + num_expert_groups=None, + num_limited_groups=None, + score_func="softmax", + route_norm=False, + route_scale=1.0, + gate_bias=False, + score_before_experts=False, + num_shared_experts=0, + load_balance_coeff=None, + _debug_force_load_balance=False, + ), + ) + case '30B-A3B': + return Qwen3ModelArgs( + vocab_size=151936, + dim=2048, + n_layers=48, + n_heads=32, + n_kv_heads=4, + head_dim=128, + hidden_dim=6144, + norm_eps=1e-6, + qk_norm=True, + max_seq_len=262144, + rope_theta=1000000.0, + enable_weight_tying=False, + moe_enabled=True, + moe_inter_dim=768, + moe_args=MoEArgs( + num_experts=128, + num_shared_experts=0, + top_k=8, + score_func="softmax", + route_norm=True, + route_scale=1.0, + score_before_experts=False, + load_balance_coeff=None, + ), + ) + case '30B-A3B-half': + return Qwen3ModelArgs( + vocab_size=151936, + dim=2048, + n_layers=24, + n_heads=32, + n_kv_heads=4, + head_dim=128, + hidden_dim=6144, + norm_eps=1e-6, + qk_norm=True, + max_seq_len=262144, + rope_theta=1000000.0, + enable_weight_tying=False, + moe_enabled=True, + moe_inter_dim=768, + moe_args=MoEArgs( + num_experts=64, + num_shared_experts=0, + top_k=8, + score_func="softmax", + route_norm=True, + route_scale=1.0, + score_before_experts=False, + load_balance_coeff=None, + ), + ) + case '72B': + return Qwen3ModelArgs( + vocab_size=152064, + dim=8192, + n_layers=4, #80, + n_heads=64, + n_kv_heads=8, + head_dim=128, + hidden_dim=29568, + norm_eps=1e-5, + qk_norm=True, + max_seq_len=131072, + rope_theta=1000000.0, + depth_init=True, + enable_weight_tying=False, + moe_enabled=False, + ) + case _: + raise ValueError(f"Unknown model config: {name}") + + +class BmmExperts(nn.Module): + """Triton-free expert GEMM with GroupedExperts-compatible routing. + + Mirrors torchtitan ``GroupedExperts`` parameter shapes/names so the param + registry is unchanged, but runs the per-expert SwiGLU via batched matmul + instead of ``torch._grouped_mm``. The input must be sorted by expert, and + ``num_tokens_per_expert`` gives each expert's dynamic token count. + """ + + def __init__(self, dim: int, hidden_dim: int, num_experts: int): + super().__init__() + self.num_experts = num_experts + self.w1 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim)) + self.w2 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim)) + self.w3 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim)) + + def forward(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor: + counts = num_tokens_per_expert.to(device=x.device, dtype=torch.long) + return bmm_experts(self.w1, self.w2, self.w3, x, counts) + + +class AnnotatedMoE(MoE): + def __init__(self, moe_args: MoEArgs, dim: int, hidden_dim: int): + super().__init__(moe_args, dim=dim, hidden_dim=hidden_dim) + # Replace the grouped-mm (Triton) experts with a batched-matmul variant. + self.experts = BmmExperts( + dim=dim, hidden_dim=hidden_dim, num_experts=moe_args.num_experts + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + bs, slen, dim = x.shape + x = x.view(-1, dim) + + ( + top_scores, + selected_experts_indices, + num_tokens_per_expert, + ) = self.router(x, self.expert_bias) + + ( + top_scores_experts_sorted, + token_indices_experts_sorted, + num_tokens_per_expert, + ) = self.reorderer(top_scores, selected_experts_indices) + + # shape (bs*slen*top_k, dim) + routed_input = x[token_indices_experts_sorted // self.router.top_k] + + if self.score_before_experts: + routed_input = ( + routed_input.to(torch.float32) + * top_scores_experts_sorted.reshape(-1, 1) + ).to(x.dtype) + + # expert component + with annotate(EP_TAG): + gathered_output = self.experts(routed_input, num_tokens_per_expert) + + gathered_output = gathered_output.reshape(-1, dim) + out = self.shared_experts(x) if self.shared_experts is not None else None + + routed_output_unsorted = torch.zeros( + (bs * slen * self.router.top_k, dim), + dtype=gathered_output.dtype, + device=gathered_output.device, + ) + routed_output_unsorted[token_indices_experts_sorted] = gathered_output + routed_output_unsorted = routed_output_unsorted.reshape( + -1, self.router.top_k, dim + ) + + if not self.score_before_experts: + out_experts = ( + torch.bmm( + top_scores.reshape(-1, 1, self.router.top_k), + routed_output_unsorted.float(), + ) + .to(x.dtype) + .squeeze(1) + ) + else: + out_experts = routed_output_unsorted.sum(dim=1) + + if out is None: + return out_experts.reshape(bs, slen, dim) + return (out + out_experts).reshape(bs, slen, dim) + + +class AnnotatedQwen3TransformerBlock(TransformerBlock): + def __init__(self, layer_id: int, model_args: Qwen3ModelArgs): + super().__init__(layer_id, model_args) + + self.layer_id = layer_id + if self.moe_enabled: + self.moe = AnnotatedMoE( + model_args.moe_args, dim=model_args.dim, hidden_dim=model_args.moe_inter_dim + ) + + def forward( + self, + x: torch.Tensor, + freqs_cis: torch.Tensor, + attention_masks: AttentionMasksType | None, + positions: torch.Tensor | None = None, + ): + x = x + self.attention( + self.attention_norm(x), freqs_cis, attention_masks, positions + ) + + if self.moe_enabled: + x = x + self.moe(self.ffn_norm(x)) + else: + x = x + self.feed_forward(self.ffn_norm(x)) + return x + +class PiperQwen3Model(Qwen3Model): + """ + Qwen3Model with pipeline stage annotations. + + Wraps torchtitan's Qwen3Model and adds stage boundary annotations + in the forward pass for pipeline parallelism, and EP dispatch/combine + annotations in the MoE layers. + """ + + def __init__(self, config: Qwen3ModelArgs, num_stages: int): + super().__init__(config) + if num_stages is None: + from src.state import piper_metadata + + schedule_info = piper_metadata.schedule_info or {} + num_stages = int(schedule_info.get("num_stages", schedule_info.get("pp_degree", 2))) + self.num_stages = num_stages + + # Replace TransformerBlock layers with AnnotatedQwen3TransformerBlock + for layer_id in range(config.n_layers): + self.layers[str(layer_id)] = AnnotatedQwen3TransformerBlock(layer_id, config) + + def forward( + self, + tokens: torch.Tensor, + attention_masks: Optional[AttentionMasksType] = None, + positions: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + num_layers = len(self.layers) + + for stage_id in range(self.num_stages): + layer_start = stage_id * num_layers // self.num_stages + layer_end = (stage_id + 1) * num_layers // self.num_stages + with annotate(PP_TAG): + if stage_id == 0: + h = self.tok_embeddings(tokens) if self.tok_embeddings is not None else tokens + + for i in range(layer_start, layer_end): + layer = self.layers[str(i)] + h = layer(h, self.rope_cache, attention_masks, positions) + + if stage_id == self.num_stages - 1: + h = self.norm(h) if self.norm is not None else h + output = self.output(h) if self.output is not None else h + + return output diff --git a/examples/run_qwen_examples.py b/examples/run_qwen_examples.py new file mode 100755 index 0000000..15bd1ac --- /dev/null +++ b/examples/run_qwen_examples.py @@ -0,0 +1,172 @@ +#!/usr/bin/env python3 +"""Run the Qwen example schedule set with the normal execution path.""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +from pathlib import Path + + +QWEN_MODELS = ("9M", "1B", "9B", "48B", "30B-A3B", "30-A3B-half", "72B") + +# name, base schedule, generated schedule, pp ranks, microbatches +RUNS = ( + ( + "dualpipev", + "examples/base-schedules/pp4_dp2_ep2_v_placement.json", + "dualpipev", + 2, + 4, + ), + ("zero3_1f1b", "examples/base-schedules/pp2_dp2_ep2_zero3.json", "1f1b", 2, 4), + ( + "bucket100_1f1b", + "examples/base-schedules/pp2_dp2_ep2_bucket100.json", + "1f1b", + 2, + 4, + ), + ( + "interleaved_1f1b", + "examples/base-schedules/pp4_dp2_ep2_interleaved.json", + "interleaved_1f1b", + 2, + 4, + ), + ("zero2_1f1b", "examples/base-schedules/pp2_dp2_ep2_zero2.json", "1f1b", 2, 4), + ( + "custom_order", + "examples/base-schedules/pp2_dp2_ep2_custom_order.json", + "custom", + 2, + 2, + ), + ("zerobubble", "examples/base-schedules/pp2_dp2_ep2.json", "zerobubble", 2, 4), +) + + +def main(argv: list[str]) -> int: + args = _parse_args(argv) + repo_root = Path(__file__).resolve().parent.parent + python_bin = _python_bin(repo_root) + forwarded_args = _forwarded_args(args) + + print(f"Python: {python_bin}", flush=True) + if forwarded_args: + print(f"Forwarded args: {' '.join(forwarded_args)}", flush=True) + + status = 0 + for name, base_schedule, schedule, ranks, mbs in RUNS: + print(flush=True) + print( + f"==> {name}: {schedule} {base_schedule} --ranks {ranks} --mbs {mbs}", + flush=True, + ) + cmd = [ + python_bin, + "examples/test_harness.py", + "--test-file", + "examples/test_qwen.py", + "--base-schedule", + base_schedule, + "--schedule", + schedule, + "--ranks", + str(ranks), + "--mbs", + str(mbs), + *forwarded_args, + ] + result = subprocess.run(cmd, cwd=repo_root, check=False) + if result.returncode: + print(f"FAILED: {name}", file=sys.stderr, flush=True) + status = 1 + else: + print(f"PASSED: {name}", flush=True) + + return status + + +def _parse_args(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run all Qwen example schedules with the normal execution path." + ) + parser.add_argument("--model", choices=QWEN_MODELS) + parser.add_argument("--batch-size", type=int) + parser.add_argument("--seq-len", type=int) + parser.add_argument("--warmup", type=int) + parser.add_argument("--iters", type=int) + parser.add_argument("--iteration-sleep", type=float) + parser.add_argument("--activation-checkpointing", action="store_true") + parser.add_argument("--nsight", action="store_true") + parser.add_argument("--viz", action="store_true") + parser.add_argument( + "--use-inductor", + action=argparse.BooleanOptionalAction, + default=None, + help="Compile actor stage modules with torch.compile.", + ) + parser.add_argument( + "--pp-outer", + action=argparse.BooleanOptionalAction, + default=None, + help="Use PP as the outer placement dimension.", + ) + parser.add_argument("--pytorch-profiler", action="store_true") + parser.add_argument("--pytorch-profiler-iters", type=int) + parser.add_argument("--address") + parser.add_argument("--port", type=int) + parser.add_argument("--temp-dir", help="Ray temp directory.") + parser.add_argument("--ray-namespace") + return parser.parse_args(argv) + + +def _forwarded_args(args: argparse.Namespace) -> list[str]: + forwarded: list[str] = [] + _append_value(forwarded, "--model", args.model) + _append_value(forwarded, "--batch-size", args.batch_size) + _append_value(forwarded, "--seq-len", args.seq_len) + _append_value(forwarded, "--warmup", args.warmup) + _append_value(forwarded, "--iters", args.iters) + _append_value(forwarded, "--iteration-sleep", args.iteration_sleep) + if args.activation_checkpointing: + forwarded.append("--activation-checkpointing") + if args.nsight: + forwarded.append("--nsight") + if args.viz: + forwarded.append("--viz") + if args.use_inductor is not None: + forwarded.append("--use-inductor" if args.use_inductor else "--no-use-inductor") + if args.pp_outer is not None: + forwarded.append("--pp-outer" if args.pp_outer else "--no-pp-outer") + if args.pytorch_profiler: + forwarded.append("--pytorch-profiler") + _append_value(forwarded, "--pytorch-profiler-iters", args.pytorch_profiler_iters) + _append_value(forwarded, "--address", args.address) + _append_value(forwarded, "--port", args.port) + _append_value(forwarded, "--temp-dir", args.temp_dir) + _append_value(forwarded, "--ray-namespace", args.ray_namespace) + return forwarded + + +def _append_value(out: list[str], flag: str, value: object | None) -> None: + if value is not None: + out.extend([flag, str(value)]) + + +def _python_bin(repo_root: Path) -> str: + if os.environ.get("PYTHON"): + return os.environ["PYTHON"] + + venv_python = repo_root / ".venv" / "bin" / "python" + if venv_python.exists(): + return str(venv_python) + + return sys.executable + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/examples/test_harness.py b/examples/test_harness.py new file mode 100644 index 0000000..962906b --- /dev/null +++ b/examples/test_harness.py @@ -0,0 +1,835 @@ +"""Build a schedule JSON and run a model test with it. + +This is a small harness around model example modules such as +``examples.test_qwen`` and ``examples.test_llama``. It consumes schedule/runtime +arguments, parses the selected test module's arguments, starts Ray, creates the +Piper placement group, and runs the model through ``PiperProgramCoordinator``. +""" + +from __future__ import annotations + +import argparse +import csv +import importlib.util +import json +import logging +import re +import statistics +import sys +from datetime import datetime +from pathlib import Path + +import ray +from ray.util.placement_group import placement_group_table + +from src.coordinator import PiperProgramCoordinator, create_piper_placement_group +from src.schedule import load_schedule_directives +from src.visualization import visualize_order_directives + +from build_schedule import ( + build_1f1b_schedule, + build_dualpipev_schedule, + build_interleaved_1f1b_schedule, + build_zerobubble_schedule, +) + + +def main() -> None: + parser = argparse.ArgumentParser( + description=( + "Build a full schedule JSON from a base schedule and generated order " + "directives, then run a model test with it." + ) + ) + parser.add_argument( + "--test-file", + "--test", + required=True, + help=( + "Example file to run directly, e.g. examples/test_qwen.py." + ), + ) + parser.add_argument( + "--base-schedule", + required=True, + type=Path, + help="Input JSON schedule without order directives.", + ) + parser.add_argument( + "--schedule", + "--schedule-name", + dest="schedule", + required=True, + choices=("1f1b", "interleaved_1f1b", "zerobubble", "dualpipev", "custom"), + help=( + "Order directive schedule variant to append. Use 'custom' to keep the " + "split/order directives already in --base-schedule untouched." + ), + ) + parser.add_argument( + "--ranks", + type=int, + default=None, + help="Number of PP ranks. Required unless --schedule custom.", + ) + parser.add_argument( + "--mbs", + type=int, + default=None, + help="Number of microbatches. Required unless --schedule custom.", + ) + parser.add_argument( + "--virtual-stages", + type=int, + default=None, + help=argparse.SUPPRESS, + ) + parser.add_argument( + "--keep-existing-order", + action="store_true", + help="Keep any order directives already present in the base schedule.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Write the generated schedule and print the test command without running it.", + ) + parser.add_argument( + "--viz", + action="store_true", + help="Render the generated schedule and per-rank TrainingDAG visualizations.", + ) + parser.add_argument( + "--pytorch-profiler", + action="store_true", + help="Run extra iterations under torch.profiler and combine the per-actor " + "chrome traces into one trace per dp-rank under out/.", + ) + parser.add_argument( + "--address", + default="", + help="Ray head address to connect to. Omit to start a local Ray session.", + ) + parser.add_argument( + "--port", + type=int, + default=4567, + help="Ray head port to connect to when --address is set.", + ) + parser.add_argument( + "--temp-dir", + default="/tmp/piper/ray_tmp", + help="Ray temp directory.", + ) + parser.add_argument( + "--ray-namespace", + default=None, + help="Ray namespace. Defaults to the selected test module name.", + ) + args, test_args = parser.parse_known_args() + + if args.schedule != "custom": + missing = [ + flag for flag, value in (("--ranks", args.ranks), ("--mbs", args.mbs)) + if value is None + ] + if missing: + parser.error( + f"the following arguments are required unless --schedule custom: {', '.join(missing)}" + ) + if args.virtual_stages is not None: + parser.error( + "--virtual-stages is internal; it is 2 for interleaved_1f1b/dualpipev " + "and 1 otherwise" + ) + args.virtual_stages = _virtual_stages_for_schedule(args.schedule) + + # Each run gets its own timestamped output directory under out/, holding its + # results.csv (one row per DP rank), schedule viz, and combined profiles. + run_dir = Path("out") / datetime.now().strftime("%Y%m%d_%H%M%S") + run_dir.mkdir(parents=True, exist_ok=True) + + generated_schedule, viz_path = build_schedule_file(args, run_dir) + forwarded_args = _replace_schedule_arg(test_args, generated_schedule) + if args.viz: + forwarded_args.append("--viz") + + profile_dir: Path | None = None + if args.pytorch_profiler and not args.dry_run: + # Absolute path on the shared filesystem so remote actors (possibly on + # other nodes) write to the same place the harness later reads. + profile_dir = (run_dir / "pytorch_profiles").resolve() + if profile_dir.exists(): + for stale in profile_dir.glob("dp*_pp*.json"): + stale.unlink() + profile_dir.mkdir(parents=True, exist_ok=True) + + if args.dry_run: + dry_run_args = list(forwarded_args) + if args.pytorch_profiler: + dry_run_args.append("--pytorch-profiler") + _validate_test_args(args.test_file, dry_run_args) + cmd = [ + sys.executable, + str(Path(__file__).resolve()), + "--test-file", + args.test_file, + "--base-schedule", + str(args.base_schedule), + "--schedule", + args.schedule, + ] + if args.ranks is not None: + cmd.extend(["--ranks", str(args.ranks)]) + if args.mbs is not None: + cmd.extend(["--mbs", str(args.mbs)]) + cmd.extend(dry_run_args) + print(" ".join(cmd)) + print(f"generated_schedule={generated_schedule}") + if viz_path is not None: + print(f"schedule_viz={viz_path}") + return + + if args.pytorch_profiler: + forwarded_args.append("--pytorch-profiler") + dp_metrics = _run_test_module(args, forwarded_args, profile_dir=profile_dir) + if dp_metrics: + rows = _metrics_rows(dp_metrics) + results_csv = run_dir / "results.csv" + _write_results_csv(results_csv, rows) + print(f"metrics written to {results_csv}") + if profile_dir is not None: + _combine_pytorch_profiles( + profile_dir, + run_dir, + schedule=args.schedule, + pp=args.ranks if args.ranks is not None else "X", + mbs=args.mbs if args.mbs is not None else "X", + ) + + +def _validate_test_args(test_file: str, test_args: list[str]) -> None: + _, test_module = _load_test_file(test_file) + test_module.parse_args(test_args) + + +def _run_test_module( + args: argparse.Namespace, + test_args: list[str], + *, + profile_dir: Path | None = None, +) -> list[dict]: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + ) + module_name, test_module = _load_test_file(args.test_file) + test_ns = test_module.parse_args(test_args) + if hasattr(test_ns, "temp_dir"): + test_ns.temp_dir = args.temp_dir + test_ns.profile_dir = str(profile_dir) if profile_dir is not None else "" + if hasattr(test_module, "_derive_num_stages"): + test_ns.num_stages = test_module._derive_num_stages(test_ns.schedule_directives_file) + + namespace = args.ray_namespace or module_name.rsplit(".", 1)[-1].replace("test_", "") + ray.init( + address=f"{args.address}:{args.port}" if args.address else None, + namespace=namespace, + log_to_driver=True, + include_dashboard=False, + _temp_dir=args.temp_dir, + ) + try: + pp_outer = getattr(test_ns, "pp_outer", False) + pg = create_piper_placement_group(test_ns.schedule_directives_file, pp_outer=pp_outer) + ray.get(pg.ready(), timeout=600) + logging.getLogger(module_name).info(placement_group_table(pg)) + coordinator = PiperProgramCoordinator.remote( + schedule_directives_file=test_ns.schedule_directives_file, + pp_outer=pp_outer, + ) + handles = coordinator.run_program.remote(test_module.main, pg, test_ns, pg) + dp_metrics = ray.get(handles) + return dp_metrics + finally: + ray.shutdown() + + +def _combine_pytorch_profiles( + profile_dir: Path, + out_dir: Path, + schedule: str, + pp: int | str, + mbs: int | str, +) -> None: + """Group per-actor chrome traces (dp{dp}_pp{pp}.json) by dp-rank and merge + each group into one trace per dp-rank under out_dir. The combined filename + encodes the run configuration: schedule, pp degree, dp degree, mbs, dp-rank. + + Each pp-rank's events keep their own pid namespace via a collision-free + remap, and process_name metadata is prefixed with the pp rank so the merged + timeline groups work by stage. The intermediate per-actor traces and the + (now-empty) profile_dir are removed once combining succeeds. + """ + name_re = re.compile(r"^dp(\d+)_pp(\d+)\.json$") + groups: dict[int, dict[int, Path]] = {} + for path in sorted(profile_dir.glob("dp*_pp*.json")): + m = name_re.match(path.name) + if not m: + continue + dp_rank, pp_rank = int(m.group(1)), int(m.group(2)) + groups.setdefault(dp_rank, {})[pp_rank] = path + + if not groups: + print(f"pytorch profiler: no trace files found in {profile_dir}") + return + + dp_degree = len(groups) + out_dir.mkdir(parents=True, exist_ok=True) + consumed: list[Path] = [] + for dp_rank, pp_paths in sorted(groups.items()): + combined_events: list = [] + next_pid = 0 + for pp_rank, path in sorted(pp_paths.items()): + with path.open(encoding="utf-8") as f: + trace = json.load(f) + events = trace.get("traceEvents", []) + annotated = _annotate_cuda_events_with_dag_labels(events) + pid_map: dict = {} + for ev in events: + opid = ev.get("pid") + if opid is not None and opid not in pid_map: + pid_map[opid] = next_pid + next_pid += 1 + for ev in events: + if "pid" in ev and ev["pid"] in pid_map: + ev["pid"] = pid_map[ev["pid"]] + if ev.get("name") == "process_name": + pname = ev.get("args", {}).get("name", "") + ev["args"]["name"] = f"dp{dp_rank} pp{pp_rank}: {pname}" + combined_events.append(ev) + consumed.append(path) + if annotated: + print( + f"pytorch profiler: annotated {annotated} CUDA event(s) " + f"with DAG labels for dp{dp_rank} pp{pp_rank}" + ) + out_name = ( + f"pytorch_profile_{schedule}_pp{pp}_dp{dp_degree}" + f"_mbs{mbs}_dprank{dp_rank}.json" + ) + out_path = out_dir / out_name + with out_path.open("w", encoding="utf-8") as f: + json.dump({"traceEvents": combined_events}, f) + print( + f"pytorch profiler: combined {len(pp_paths)} pp-rank trace(s) " + f"-> {out_path}" + ) + + for path in consumed: + path.unlink(missing_ok=True) + try: + profile_dir.rmdir() + except OSError: + pass + + +def _is_dag_node_label(name: object) -> bool: + return isinstance(name, str) and ":uid" in name and not name.startswith("##") + + +def _annotate_cuda_events_with_dag_labels(events: list[dict]) -> int: + """Prefix GPU-side profiler events with their enclosing Piper DAG node. + + ``torch.profiler`` emits Piper ``record_function`` ranges on the CPU thread, + while CUDA kernels and copies appear on GPU lanes with PyTorch/Inductor + names such as ``## Call CompiledFxGraph ...``. CPU launch events and GPU + events share ``args["External id"]``; use that to carry the enclosing DAG + node label onto the GPU events in the combined Chrome trace. Then rewrite + nested non-DAG GPU annotations so the visible annotation stack stays rooted + in Piper's TrainingDAG labels instead of Inductor's opaque graph hashes. + """ + ranges_by_thread: dict[tuple[object, object], list[tuple[float, float, str]]] = {} + for ev in events: + if ev.get("ph") != "X" or ev.get("cat") != "user_annotation": + continue + label = ev.get("name") + if not _is_dag_node_label(label): + continue + ts = ev.get("ts") + dur = ev.get("dur") + if not isinstance(ts, (int, float)) or not isinstance(dur, (int, float)): + continue + ranges_by_thread.setdefault((ev.get("pid"), ev.get("tid")), []).append( + (float(ts), float(ts) + float(dur), label) + ) + for ranges in ranges_by_thread.values(): + ranges.sort(key=lambda item: (item[0], -item[1])) + + external_id_to_label: dict[object, str] = {} + active: dict[tuple[object, object], list[tuple[float, str]]] = {} + cpu_events = sorted( + ( + ev for ev in events + if ev.get("ph") == "X" + and ev.get("cat") != "kernel" + and str(ev.get("cat", "")).startswith(("cpu_op", "user_annotation")) + ), + key=lambda ev: float(ev.get("ts", 0.0) or 0.0), + ) + + range_idx = {key: 0 for key in ranges_by_thread} + for ev in cpu_events: + ts = ev.get("ts") + if not isinstance(ts, (int, float)): + continue + key = (ev.get("pid"), ev.get("tid")) + ranges = ranges_by_thread.get(key) + if not ranges: + continue + stack = active.setdefault(key, []) + while stack and stack[-1][0] < float(ts): + stack.pop() + idx = range_idx[key] + while idx < len(ranges) and ranges[idx][0] <= float(ts): + _start, end, label = ranges[idx] + if end >= float(ts): + stack.append((end, label)) + idx += 1 + range_idx[key] = idx + if not stack: + continue + ext_id = (ev.get("args") or {}).get("External id") + if ext_id is not None: + external_id_to_label[ext_id] = stack[-1][1] + + annotated = 0 + cuda_event_cats = {"kernel", "gpu_memcpy", "gpu_memset"} + for ev in events: + if ev.get("ph") != "X" or ev.get("cat") not in cuda_event_cats: + continue + args = ev.setdefault("args", {}) + ext_id = args.get("External id") + label = external_id_to_label.get(ext_id) + if not label: + continue + name = ev.get("name", "") + prefix = f"{label}::" + if isinstance(name, str) and name.startswith(prefix): + continue + args.setdefault("piper_dag_node", label) + args.setdefault("original_cuda_name", name) + ev["name"] = f"{prefix}{name}" + annotated += 1 + annotated += _rewrite_nested_gpu_annotations(events) + return annotated + + +def _rewrite_nested_gpu_annotations(events: list[dict]) -> int: + gpu_annotations = [ + ev for ev in events + if ev.get("ph") == "X" and ev.get("cat") == "gpu_user_annotation" + ] + dag_ranges_by_lane: dict[tuple[object, object], list[tuple[float, float, str]]] = {} + for ev in gpu_annotations: + label = ev.get("name") + if not _is_dag_node_label(label): + continue + ts = ev.get("ts") + dur = ev.get("dur") + if not isinstance(ts, (int, float)) or not isinstance(dur, (int, float)): + continue + dag_ranges_by_lane.setdefault((ev.get("pid"), ev.get("tid")), []).append( + (float(ts), float(ts) + float(dur), label) + ) + for ranges in dag_ranges_by_lane.values(): + ranges.sort(key=lambda item: (item[0], item[1] - item[0])) + + cuda_events_by_lane: dict[tuple[object, object], list[dict]] = {} + for ev in events: + if ev.get("ph") == "X" and ev.get("cat") in {"kernel", "gpu_memcpy", "gpu_memset"}: + cuda_events_by_lane.setdefault((ev.get("pid"), ev.get("tid")), []).append(ev) + for lane_events in cuda_events_by_lane.values(): + lane_events.sort(key=lambda ev: float(ev.get("ts", 0.0) or 0.0)) + + rewritten = 0 + for ev in gpu_annotations: + name = ev.get("name") + if _is_dag_node_label(name): + continue + ts = ev.get("ts") + dur = ev.get("dur") + if not isinstance(ts, (int, float)) or not isinstance(dur, (int, float)): + continue + start = float(ts) + end = start + float(dur) + lane = (ev.get("pid"), ev.get("tid")) + label = _enclosing_dag_gpu_label(dag_ranges_by_lane.get(lane, []), start, end) + if label is None: + label = _contained_cuda_dag_label( + cuda_events_by_lane.get(lane, []), start, end + ) + if label is None: + continue + args = ev.setdefault("args", {}) + args.setdefault("original_gpu_annotation", name) + args.setdefault("piper_dag_node", label) + ev["name"] = label + rewritten += 1 + return rewritten + + +def _enclosing_dag_gpu_label( + dag_ranges: list[tuple[float, float, str]], start: float, end: float +) -> str | None: + label = None + best_duration = None + for dag_start, dag_end, dag_label in dag_ranges: + if dag_start > start: + break + if dag_end < end: + continue + duration = dag_end - dag_start + if best_duration is None or duration < best_duration: + label = dag_label + best_duration = duration + return label + + +def _contained_cuda_dag_label( + cuda_events: list[dict], start: float, end: float +) -> str | None: + labels: dict[str, float] = {} + for ev in cuda_events: + ev_start = ev.get("ts") + ev_dur = ev.get("dur") + if not isinstance(ev_start, (int, float)) or not isinstance(ev_dur, (int, float)): + continue + ev_start = float(ev_start) + if ev_start > end: + break + ev_end = ev_start + float(ev_dur) + if ev_end < start: + continue + label = (ev.get("args") or {}).get("piper_dag_node") + if not label: + continue + overlap = max(0.0, min(end, ev_end) - max(start, ev_start)) + labels[label] = labels.get(label, 0.0) + overlap + if not labels: + return None + return max(labels.items(), key=lambda item: item[1])[0] + + +def _metrics_rows(dp_metrics: list[dict]) -> list[dict]: + """Summarize the raw metrics into one CSV row per DP rank (no averaging + across DP ranks). + + Each row's timing scalars come from that DP rank's own iter times; peak + memory is taken from that DP rank's per-(global=pp)-rank values and emitted + as one peak_memory_pp{i}_gb column per pp rank. + """ + rows: list[dict] = [] + for m in sorted(dp_metrics, key=lambda d: int(d.get("dp_rank", 0))): + iter_times = m.get("iter_times_s") or [] + mean_iter = statistics.fmean(iter_times) if iter_times else float("nan") + std_iter = statistics.pstdev(iter_times) if len(iter_times) > 1 else 0.0 + tokens = ( + (m.get("batch_size") or 0) + * (m.get("num_microbatches") or 1) + * (m.get("seq_len") or 0) + ) + throughput = tokens / mean_iter if iter_times and mean_iter else float("nan") + + row: dict = { + "dp_rank": m.get("dp_rank"), + "model": m.get("model"), + "schedule": m.get("schedule"), + "pp": m.get("pp"), + "dp": m.get("dp"), + "batch_size": m.get("batch_size"), + "num_microbatches": m.get("num_microbatches"), + "seq_len": m.get("seq_len"), + "samples": len(iter_times), + "iter_time_mean_s": mean_iter, + "iter_time_std_s": std_iter, + "throughput_tokens_per_s": throughput, + } + # peak_memory_by_rank is keyed by global rank; within a DP replica the + # global ranks ascend with pp stage (both pp-inner and pp-outer layouts), + # so enumerate them as local stage indices for columns that align across + # DP-rank rows. + peak = m.get("peak_memory_by_rank") or {} + for stage, rank in enumerate(sorted(peak, key=lambda r: int(r))): + row[f"peak_memory_pp{stage}_gb"] = float(peak[rank]) / (1024 ** 3) + rows.append(row) + return rows + + +def _write_results_csv(csv_path: Path, rows: list[dict]) -> None: + """Write ``rows`` (one per DP rank) to a fresh per-run results CSV.""" + csv_path.parent.mkdir(parents=True, exist_ok=True) + fieldnames: list[str] = [] + for row in rows: + for key in row: + if key not in fieldnames: + fieldnames.append(key) + with csv_path.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames, restval="") + writer.writeheader() + writer.writerows(rows) + + +def build_schedule_file(args: argparse.Namespace, run_dir: Path) -> tuple[Path, str | None]: + base = _load_schedule(args.base_schedule) + if args.schedule == "custom": + # Pass the base schedule through untouched: trust its existing split and + # order directives. Only used for visualization below. + schedule = list(base) + order_directives = [ + d for d in schedule + if isinstance(d, dict) and d.get("op") == "order" + ] + else: + if args.keep_existing_order: + schedule = list(base) + else: + schedule = [ + directive for directive in base + if not (isinstance(directive, dict) and directive.get("op") == "order") + ] + _set_num_microbatches(schedule, args.mbs) + + if args.schedule == "1f1b": + order_directives = build_1f1b_schedule(args.ranks, args.mbs) + elif args.schedule == "zerobubble": + order_directives = build_zerobubble_schedule(args.ranks, args.mbs) + elif args.schedule == "dualpipev": + order_directives = build_dualpipev_schedule(args.ranks, args.mbs) + else: + order_directives = build_interleaved_1f1b_schedule( + args.ranks, + args.mbs, + args.virtual_stages, + ) + + _validate_generated_order_placement( + schedule, + order_directives, + schedule_name=args.schedule, + ranks=args.ranks, + ) + schedule.extend(order_directives) + schedule_stem = _schedule_attr_stem( + args.schedule, + args.ranks, + args.mbs, + args.virtual_stages, + base_schedule=args.base_schedule, + ) + run_schedule = run_dir / f"{schedule_stem}.json" + with run_schedule.open("w", encoding="utf-8") as f: + json.dump(schedule, f, indent=2) + f.write("\n") + + viz_path = None + if args.viz: + viz_output = run_dir / f"{schedule_stem}.png" + viz_path = visualize_order_directives(order_directives, viz_output) + return run_schedule, viz_path + + +def _schedule_attr_stem( + schedule_name: str, + ranks: int | None, + mbs: int | None, + virtual_stages: int | None, + base_schedule: Path | None = None, +) -> str: + if schedule_name == "custom": + # ranks/mbs/virtual_stages aren't required for custom; derive a stable + # name from the base schedule filename instead. + if base_schedule is not None: + return f"custom_{Path(base_schedule).stem}" + return "custom" + suffix = f"{schedule_name}_pp{ranks}_mbs{mbs}" + if schedule_name == "interleaved_1f1b": + suffix += f"_v{virtual_stages}" + return suffix + + +def _virtual_stages_for_schedule(schedule_name: str) -> int: + if schedule_name in ("interleaved_1f1b", "dualpipev"): + return 2 + return 1 + + +def _load_schedule(path: Path) -> list[dict]: + return load_schedule_directives(str(path)) + + +def _set_num_microbatches(schedule: list[dict], n_mbs: int) -> None: + """Ensure the schedule has exactly one split directive with num_microbatches=n_mbs. + + If a split directive already exists, its num_microbatches is overwritten. + Otherwise a default MB-split directive is appended. + """ + found = False + for directive in schedule: + if directive.get("op") == "split": + directive["num_microbatches"] = n_mbs + found = True + if not found: + schedule.append({ + "op": "split", + "filter": {}, + "dim_name": "MB", + "num_microbatches": n_mbs, + }) + + +def _validate_generated_order_placement( + base_schedule: list[dict], + order_directives: list[dict], + *, + schedule_name: str, + ranks: int | None, +) -> None: + """Generated order rows are per physical rank and must be device-local.""" + stage_devices = _stage_devices_from_place_directives(base_schedule) + if not stage_devices: + return + + errors: list[str] = [] + for row_idx, directive in enumerate(order_directives): + row_devices: set[tuple[int, ...]] = set() + for slot_idx, slot in enumerate(_order_filter_slots(directive.get("filters", []))): + slot_devices: set[tuple[int, ...]] = set() + slot_pps: list[int] = [] + for flt in slot: + spec = _filter_to_dict(flt) + pp = spec.get("PP") + if pp is None or pp == "*": + continue + pp = int(pp) + slot_pps.append(pp) + dev = stage_devices.get(pp) + if dev is None: + errors.append( + f"row {row_idx} slot {slot_idx} references PP{pp}, " + "but no matching place directive exists" + ) + continue + slot_devices.add(dev) + if len(slot_devices) > 1: + errors.append( + f"row {row_idx} slot {slot_idx} groups stages {slot_pps} " + f"across devices {sorted(slot_devices)}" + ) + row_devices.update(slot_devices) + if len(row_devices) > 1: + errors.append( + f"row {row_idx} orders stages across devices {sorted(row_devices)}" + ) + + if not errors: + return + + hint = "" + if schedule_name == "dualpipev" and ranks is not None: + expected = [ + (rank, 2 * ranks - 1 - rank) + for rank in range(ranks) + ] + hint = ( + " DualPipeV uses V-layout virtual stages; each pair " + f"{expected} must be placed on the same physical device set." + ) + raise ValueError( + "generated order directives are not compatible with the base placement: " + + "; ".join(errors[:4]) + + hint + ) + + +def _stage_devices_from_place_directives(schedule: list[dict]) -> dict[int, tuple[int, ...]]: + stage_devices: dict[int, tuple[int, ...]] = {} + for directive in schedule: + if directive.get("op") != "place": + continue + devices = directive.get("devices", directive.get("device")) + if not isinstance(devices, list): + continue + spec = _filter_to_dict(directive.get("filter", [])) + pp = spec.get("PP") + if pp is None or pp == "*": + continue + stage_devices[int(pp)] = tuple(sorted(int(d) for d in devices)) + return stage_devices + + +def _order_filter_slots(filters: object) -> list[list[object]]: + if not isinstance(filters, list): + raise ValueError(f"order directive requires filters list, got {type(filters)}") + slots: list[list[object]] = [] + for item in filters: + if isinstance(item, list) and item: + if not all(isinstance(flt, dict) for flt in item): + raise ValueError(f"invalid order filter group: {item}") + slots.append(list(item)) + else: + raise ValueError(f"invalid order filter group: {item}") + return slots + + +def _filter_to_dict(flt: object) -> dict: + if isinstance(flt, dict): + return dict(flt) + raise ValueError(f"invalid filter: {flt}") + + +def _replace_schedule_arg(test_args: list[str], schedule_path: Path) -> list[str]: + return _replace_arg(test_args, "--schedule-directives-file", str(schedule_path)) + + +def _replace_arg(test_args: list[str], flag: str, value: str) -> list[str]: + """Drop any existing occurrence of ``flag`` (and its value) from forwarded + test args, then append ``flag value`` so the harness setting wins.""" + out: list[str] = [] + skip_next = False + for arg in test_args: + if skip_next: + skip_next = False + continue + if arg == flag: + skip_next = True + continue + if arg.startswith(f"{flag}="): + continue + out.append(arg) + out.extend([flag, value]) + return out + + +def _load_test_file(test_file: str): + path = Path(test_file) + if path.suffix != ".py": + raise ValueError(f"--test-file must be a Python file path, got: {test_file}") + if not path.is_absolute(): + path = Path.cwd() / path + if not path.exists(): + raise FileNotFoundError(f"test file does not exist: {path}") + + sys.path.insert(0, str(path.parent)) + module_name = path.stem + spec = importlib.util.spec_from_file_location(module_name, path) + if spec is None or spec.loader is None: + raise ImportError(f"could not load Python file: {path}") + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module_name, module + + +if __name__ == "__main__": + main() diff --git a/examples/test_llama.py b/examples/test_llama.py new file mode 100644 index 0000000..d024c63 --- /dev/null +++ b/examples/test_llama.py @@ -0,0 +1,192 @@ +"""End-to-end LLaMA test for the JSON-driven TrainingDAG backend.""" +import argparse +import os +import time + +import ray +import torch + +from src.piper import piper_exec_dag +from src.compile import piper_setup +from src.state import ( + LOG_LEVEL, + create_logger, + piper_metadata, +) + +from models.llama3 import ( + LLAMA_1B, + LLAMA_3B, + LLAMA_8B, + LLAMA_70B, + LLAMA_DEBUG, + Transformer, + precompute_freqs_cis, +) + +logger = create_logger("test_llama", LOG_LEVEL) + + +def _model_config(name: str): + match name: + case "debug": + return LLAMA_DEBUG + case "1b": + return LLAMA_1B + case "3b": + return LLAMA_3B + case "8b": + return LLAMA_8B + case "70b": + return LLAMA_70B + raise ValueError(f"Unknown LLaMA model: {name}") + + +def _raw_metrics(args, iter_times, peak_memory_stats): + """Assemble raw per-dp-rank measurements + run config for the harness. + + No derived statistics are computed here; the harness summarizes the metrics + and writes the CSV. ``peak_memory_by_rank`` maps global rank -> peak bytes. + """ + info = dict(getattr(piper_metadata, "schedule_info", {}) or {}) + return { + "dp_rank": int(os.environ["PIPER_DP_RANK"]), + "model": args.model, + "schedule": info.get( + "name", os.path.splitext(os.path.basename(args.schedule_directives_file))[0] + ), + "schedule_directives_file": args.schedule_directives_file, + "pp": info.get("pp_degree"), + "dp": info.get("dp_degree"), + "batch_size": args.batch_size, + "num_microbatches": int(info.get("num_microbatches", 1)), + "seq_len": args.seq_len, + "iter_times_s": [float(t) for t in iter_times], + "peak_memory_by_rank": { + int(rank): int(max_alloc) for rank, max_alloc in peak_memory_stats + }, + } + + +def main(args, pg): + llama_config = _model_config(args.model) + logger.info(args) + + loss_mod = torch.nn.CrossEntropyLoss() + loss_fn = lambda output, labels: loss_mod(output.view(-1, output.size(-1)), labels.view(-1)) + + x = torch.randint(0, llama_config.vocab_size, (args.batch_size, args.seq_len)) + y = torch.randint(0, llama_config.vocab_size, (args.batch_size, args.seq_len)) + + freqs_cis = precompute_freqs_cis( + llama_config.dim // llama_config.n_heads, + args.seq_len, + llama_config.rope_theta, + ) + causal = torch.full((args.seq_len, args.seq_len), float("-inf")) + causal = torch.triu(causal, diagonal=1) + mask = torch.hstack([torch.zeros((args.seq_len, 0)), causal]) + + piper_setup( + Transformer, + model_args=(llama_config, args.seq_len), + optim_fn=torch.optim.Adam, + example_inputs=[x], + example_outputs=y, + activation_checkpointing=args.activation_checkpointing, + num_checkpoints=args.num_checkpoints, + model_dtype=torch.bfloat16, + pg=pg, + nsight=args.nsight, + temp_dir=args.temp_dir, + visualize_dag=args.viz, + const_attrs={"freqs_cis": freqs_cis, "mask": mask}, + use_inductor=args.use_inductor, + pp_outer=args.pp_outer, + schedule_directives_file=args.schedule_directives_file, + ) + + actors = piper_metadata.actors + logger.info(f"Running {args.warmup} warmup iterations") + for _ in range(args.warmup): + piper_exec_dag(loss_fn, log_stats=True) + time.sleep(1) + + ray.get([actor.reset_peak_memory.remote() for actor in actors.values()]) + logger.info(f"Running {args.iters} timed iterations") + iter_times = [] + for _ in range(args.iters): + start = time.perf_counter() + piper_exec_dag(loss_fn, log_stats=True) + end = time.perf_counter() + iter_times.append(end - start) + time.sleep(1) + + peak_memory_stats = ray.get( + [actor.get_and_reset_peak_memory_stats.remote() for actor in actors.values()] + ) + + metrics = _raw_metrics(args, iter_times, peak_memory_stats) + + if args.pytorch_profiler: + profile_dir = getattr(args, "profile_dir", "") or os.path.join( + "out", "pytorch_profiles" + ) + logger.info(f"Running {args.pytorch_profiler_iters} PyTorch-profiled iterations") + ray.get([actor.start_pytorch_profiler.remote() for actor in actors.values()]) + for _ in range(args.pytorch_profiler_iters): + piper_exec_dag(loss_fn) + time.sleep(1) + ray.get([ + actor.stop_pytorch_profiler.remote(profile_dir) + for actor in actors.values() + ]) + + if args.nsight: + logger.info("Stopping Piper actors so Nsight Systems reports are flushed") + try: + ray.get([actor.__ray_terminate__.remote() for actor in actors.values()]) + except ray.exceptions.ActorDiedError as exc: + logger.info(f"Piper actors stopped for Nsight flush: {exc}") + return metrics + + +def parse_args(argv=None): + parser = argparse.ArgumentParser(description="Run LLaMA with JSON-driven TrainingDAG execution") + parser.add_argument("--model", choices=["debug", "1b", "3b", "8b", "70b"], default="debug") + parser.add_argument("--batch-size", type=int, default=16) + parser.add_argument("--seq-len", type=int, default=256) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iters", type=int, default=5) + parser.add_argument("--activation-checkpointing", action="store_true", default=False) + parser.add_argument( + "--num-checkpoints", + type=int, + default=1, + help="Number of sequential activation-checkpoint regions per Piper bucket", + ) + parser.add_argument("--nsight", action="store_true", default=False) + parser.add_argument("--viz", action="store_true", default=False) + parser.add_argument("--temp-dir", default="/tmp/piper/ray_tmp") + parser.add_argument("--use-inductor", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--pp-outer", action=argparse.BooleanOptionalAction, default=False) + parser.add_argument( + "--schedule-directives-file", + type=str, + default="examples/base-schedules/pp2.json", + help="JSON file containing TrainingDAG schedule directives", + ) + parser.add_argument( + "--pytorch-profiler", + action="store_true", + default=False, + help="Run extra iterations under torch.profiler on every actor and write " + "per-actor chrome traces (combined per dp-rank by test_harness).", + ) + parser.add_argument( + "--pytorch-profiler-iters", + type=int, + default=3, + help="Number of iterations to run under the PyTorch profiler.", + ) + return parser.parse_args(argv) diff --git a/examples/test_qwen.py b/examples/test_qwen.py new file mode 100644 index 0000000..acae165 --- /dev/null +++ b/examples/test_qwen.py @@ -0,0 +1,203 @@ +"""End-to-end Qwen test for the JSON-driven TrainingDAG backend.""" +import ray +import torch +import argparse +import time +import os + +from src.compile import piper_setup +from src.piper import piper_exec_dag +from src.schedule import load_schedule_directives +from src.state import piper_metadata, create_logger, LOG_LEVEL + +from models.qwen3 import PiperQwen3Model, create_qwen3_config +from torchtitan.models.qwen3.model.model import precompute_rope_cache + +logger = create_logger("test_qwen", LOG_LEVEL) + + +def _raw_metrics(args, iter_times, peak_memory_stats): + """Assemble raw per-dp-rank measurements + run config for the harness. + + No derived statistics are computed here; the harness summarizes the metrics + and writes the CSV. ``peak_memory_by_rank`` maps global rank -> peak bytes. + """ + info = dict(getattr(piper_metadata, "schedule_info", {}) or {}) + return { + "dp_rank": int(os.environ["PIPER_DP_RANK"]), + "model": args.model, + "schedule": info.get( + "name", os.path.splitext(os.path.basename(args.schedule_directives_file))[0] + ), + "schedule_directives_file": args.schedule_directives_file, + "pp": info.get("pp_degree"), + "dp": info.get("dp_degree"), + "batch_size": args.batch_size, + "num_microbatches": int(info.get("num_microbatches", 1)), + "seq_len": args.seq_len, + "iter_times_s": [float(t) for t in iter_times], + "peak_memory_by_rank": { + int(rank): int(max_alloc) for rank, max_alloc in peak_memory_stats + }, + } + + +def main(args, pg): + batch_size = args.batch_size + + config = create_qwen3_config(args.model) + num_stages = int( + getattr(args, "num_stages", 0) + or _derive_num_stages(args.schedule_directives_file) + ) + + x = torch.randint(0, config.vocab_size, (batch_size, args.seq_len)) + y = torch.randint(0, config.vocab_size, (batch_size, args.seq_len)) + + _ce = torch.nn.CrossEntropyLoss() + loss_fn = lambda output, labels: _ce(output.view(-1, output.size(-1)), labels.view(-1)) + + rope_cache = precompute_rope_cache( + config.head_dim, + config.max_seq_len, + config.rope_theta, + ) + piper_setup( + PiperQwen3Model, + model_args=(config, num_stages), + optim_fn=torch.optim.Adam, + example_inputs=[x], + example_outputs=y, + activation_checkpointing=args.activation_checkpointing, + model_dtype=torch.bfloat16, + pg=pg, + nsight=args.nsight, + temp_dir=args.temp_dir, + visualize_dag=args.viz, + const_attrs={"rope_cache": rope_cache}, + use_inductor=args.use_inductor, + pp_outer=args.pp_outer, + schedule_directives_file=args.schedule_directives_file, + ) + + del x, y + + actors = piper_metadata.actors + + logger.info(f"Running {args.warmup} warmup iterations") + for _ in range(args.warmup): + piper_exec_dag(loss_fn) + if args.iteration_sleep > 0: + time.sleep(args.iteration_sleep) + + logger.info(f"Running {args.iters} timed iterations") + ray.get([actor.reset_peak_memory.remote() for actor in actors.values()]) + iter_times = [] + for _ in range(args.iters): + start = time.perf_counter() + piper_exec_dag(loss_fn, log_stats=True) + end = time.perf_counter() + iter_times.append(end - start) + if args.iteration_sleep > 0: + time.sleep(args.iteration_sleep) + + peak_memory_stats = ray.get( + [actor.get_and_reset_peak_memory_stats.remote() for actor in actors.values()] + ) + + metrics = _raw_metrics(args, iter_times, peak_memory_stats) + + if args.pytorch_profiler: + profile_dir = getattr(args, "profile_dir", "") or os.path.join( + "out", "pytorch_profiles" + ) + logger.info(f"Running {args.pytorch_profiler_iters} PyTorch-profiled iterations") + ray.get([actor.start_pytorch_profiler.remote() for actor in actors.values()]) + for _ in range(args.pytorch_profiler_iters): + piper_exec_dag(loss_fn) + if args.iteration_sleep > 0: + time.sleep(args.iteration_sleep) + ray.get([ + actor.stop_pytorch_profiler.remote(profile_dir) + for actor in actors.values() + ]) + + if args.nsight: + logger.info("Stopping Piper actors so Nsight Systems reports are flushed") + try: + ray.get([actor.__ray_terminate__.remote() for actor in actors.values()]) + except ray.exceptions.ActorDiedError as exc: + logger.info(f"Piper actors stopped for Nsight flush: {exc}") + return metrics + + +def _derive_num_stages(schedule_directives_file: str) -> int: + schedule_directives = load_schedule_directives(schedule_directives_file) + num_stages = sum( + 1 + for directive in schedule_directives + if isinstance(directive, dict) and directive.get("op") == "place" + ) + if num_stages <= 0: + raise ValueError( + f"schedule directives file must contain at least one place directive: " + f"{schedule_directives_file}" + ) + return num_stages + + +def parse_args(argv=None): + parser = argparse.ArgumentParser( + description="Test JSON-driven Qwen TrainingDAG execution" + ) + parser.add_argument('--model', choices=['9M', '1B', '9B', '48B', '30B-A3B', '30-A3B-half', '72B'], default='9M', + help='Model configuration: 9M, 1B, 9B, 48B, 30B-A3B, 30-A3B-half, or 72B (default: 9M)') + parser.add_argument("--batch-size", type=int, default=4) + parser.add_argument("--seq-len", type=int, default=1024) + parser.add_argument("--warmup", type=int, default=3) + parser.add_argument("--iters", type=int, default=3) + parser.add_argument("--iteration-sleep", type=float, default=0.0) + parser.add_argument('--activation-checkpointing', action='store_true', default=False) + parser.add_argument("--nsight", action="store_true", default=False, + help="Whether to use Nsight Systems for tracing") + parser.add_argument("--viz", action="store_true", default=False, + help="Save schedule and per-rank DAG visualizations") + parser.add_argument("--temp-dir", default="/tmp/piper/ray_tmp", + help="Ray temp directory (default: /tmp/piper/ray_tmp)") + parser.add_argument( + "--use-inductor", + action=argparse.BooleanOptionalAction, + default=True, + help="Whether actors torch.compile stage GraphModules in _load_stage (default: true)", + ) + parser.add_argument( + "--pp-outer", + action=argparse.BooleanOptionalAction, + default=False, + help=( + "Use PP as the outer placement dim (one pipeline stage per node, " + "all DP replicas for that stage colocated). Makes per-stage EP/DP " + "collectives intra-node at the cost of inter-node PP P2P. " + "Default: false (one DP replica per node, PP inner)." + ), + ) + parser.add_argument( + "--schedule-directives-file", + type=str, + default="examples/base-schedules/pp2.json", + help="JSON file containing schedule directives for the piper backend", + ) + parser.add_argument( + "--pytorch-profiler", + action="store_true", + default=False, + help="Run extra iterations under torch.profiler on every actor and write " + "per-actor chrome traces (combined per dp-rank by test_harness).", + ) + parser.add_argument( + "--pytorch-profiler-iters", + type=int, + default=3, + help="Number of iterations to run under the PyTorch profiler.", + ) + return parser.parse_args(argv) diff --git a/figs/architecture.jpg b/figs/architecture.jpg new file mode 100644 index 0000000..068cfd5 Binary files /dev/null and b/figs/architecture.jpg differ diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..62d361d --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,43 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "piper" +version = "0.1.0" +description = "User-controllable distributed training schedules for large-scale ML models." +readme = "README.md" +requires-python = ">=3.10" +license = "MIT" +dependencies = [ + "torchtitan", + "numpy", + "torch==2.10.0", + "ray==2.44.1", + "click<8.1", + "cupy-cuda12x", + "bitsandbytes", + "graphviz", + "matplotlib", +] + +[project.optional-dependencies] +dev = [ + "pytest", +] + +[project.urls] +Repository = "https://github.com/uw-syfi/piper" +Issues = "https://github.com/uw-syfi/piper/issues" +Paper = "TODO" +Blog = "TODO" + +[tool.setuptools.packages.find] +where = ["."] +include = ["src", "src.*", "test", "test.*", "examples", "examples.*"] + +[tool.pytest.ini_options] +testpaths = ["test"] +markers = [ + "gpu: tests requiring CUDA GPUs", +] diff --git a/requirements.txt b/requirements.txt index abc398c..376ad5c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,10 @@ -python==3.10 +torchtitan numpy -ray==2.53.0 -torch==2.9.1+cu128 +torch==2.10.0 +ray==2.44.1 +click<8.1 cupy-cuda12x +bitsandbytes +graphviz +matplotlib +pytest diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/actor.py b/src/actor.py new file mode 100644 index 0000000..4619e86 --- /dev/null +++ b/src/actor.py @@ -0,0 +1,747 @@ +import ray +import torch +import os +import re +from typing import Any +import gc +from torch.autograd.graph import set_warn_on_accumulate_grad_stream_mismatch +import torch.distributed as dist +from collections import defaultdict + +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from .state import ( + create_logger, + LOG_LEVEL, +) +from .fx import _deserialize_graphmodule, _serialize_graphmodule +from .executors import CommunicationExecutor, ComputeExecutor, DagExecutor +from .ordering import _serial_topological_order +from .runtime import BufferStore, EventStore, ParamStorage, RuntimeState, StageStore +from .tasks import training_dag_task_type as _training_dag_task_type + +CLEANUP_MEMORY = False + +logger = create_logger("actor", LOG_LEVEL) + +def _disable_functorch_donated_buffers() -> None: + import importlib + + config = importlib.import_module("torch._functorch.config") + config.donated_buffer = False + + +def _get_rank(pp_rank, dp_rank, pp_degree): + return pp_rank + dp_rank * pp_degree + + +def _create_actors( + num_actors, + optim_class, + profile=False, + no_nvtx: bool = False, + pg=None, + temp_dir: str = None, + use_inductor: bool = False, + pp_outer: bool = False, +): + dp_rank = int(os.environ["PIPER_DP_RANK"]) + world_size = int(os.environ["PIPER_WORLD_SIZE"]) + dp_degree = int(os.environ["PIPER_DP_DEGREE"]) + pp_degree = int(os.environ["PIPER_PP_DEGREE"]) + + from .state import piper_metadata + + for pp_rank in range(num_actors): + global_rank = _get_rank(pp_rank, dp_rank, pp_degree) + nsight_env = {"nsight": { + "t": "cuda,cudnn,cublas,nvtx", + "sample": "process-tree", + "backtrace": "dwarf", + "cudabacktrace": "sync:0,memory:0", + "python-backtrace": "cuda", + "stop-on-exit": "true", + }} if profile else {} + nccl_env = { + "env_vars": { + # "NCCL_SOCKET_IFNAME": "ens32", + # "GLOO_SOCKET_IFNAME": "ens32", + # "NCCL_P2P_DISABLE": "1", + # "NCCL_DEBUG": "INFO", + **({"TMPDIR": temp_dir} if (profile and temp_dir) else {}), + } + } + # When pp_outer=True, one bundle corresponds to one pipeline stage and + # holds all DP replicas for that stage (placement group shape is + # [{"GPU": dp}] * pp). Otherwise one bundle is one DP replica holding + # all PP ranks (shape is [{"GPU": pp}] * dp). + bundle_index = pp_rank if pp_outer else dp_rank + actor = PiperActor.options( + num_gpus=0.6, + runtime_env={**nsight_env, **nccl_env}, + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=bundle_index, + ), + ).remote( + pp_rank, + optim_class, + world_size, + dp_rank=dp_rank, + dp_degree=dp_degree, + pp_degree=pp_degree, + no_nvtx=no_nvtx, + use_inductor=use_inductor, + ) + piper_metadata.actors[pp_rank] = actor + + +@ray.remote +class PiperActor: + def __init__( + self, + pp_rank, + optim_class, + world_size, + dp_rank=0, + dp_degree=1, + pp_degree=1, + no_nvtx: bool = False, + use_inductor: bool = False, + ): + self.logger = create_logger("actor", LOG_LEVEL) + + # BWD_I uses torch.autograd.grad (not .backward), so AccumulateGrad nodes are + # traversed for stream-sync bookkeeping but never accumulate to p.grad. + # Suppress the spurious stream-mismatch warning. + set_warn_on_accumulate_grad_stream_mismatch(False) + + self.optim_class = optim_class + self.use_inductor = bool(use_inductor) + if self.use_inductor: + _disable_functorch_donated_buffers() + self.runtime = RuntimeState( + pp_rank=pp_rank, + dp_rank=dp_rank, + dp_degree=dp_degree, + pp_degree=pp_degree, + world_size=world_size, + no_nvtx=no_nvtx, + ) + + self.logger.debug( + f"Initializing Ray actor {self.runtime.global_rank} GPU {os.environ['CUDA_VISIBLE_DEVICES']}" + ) + + self.inputs = None + self.labels = None + + self.stages = StageStore() + # accumuate loss for each microbatch + self.loss = [] + + # DAG execution state + self.dag = None + self.sorted_dag_nodes = None + + self.buffers = BufferStore() + self.events = EventStore() + self.communication = CommunicationExecutor(self.runtime, self.stages, self.logger) + self.compute = ComputeExecutor(self.runtime, self.stages, self.logger) + self.params = ParamStorage(self.runtime, self.stages, self.logger) + self.dag_executor = DagExecutor( + self.runtime, + self.stages, + self.buffers, + self.events, + self.params, + self.communication, + self.compute, + self.logger, + ) + + # Non-trainable constant tensor attributes (e.g. freqs_cis, mask) pushed + # from the coordinator before compilation so _load_stage can fill them in + # instead of zero-initializing. Keyed by bare attribute name (e.g. "freqs_cis"). + self.model_const_attrs: dict = {} + + def get_and_reset_peak_memory_stats(self) -> tuple: + """Return (global_rank, max_memory_allocated_bytes) and reset peak stats.""" + max_alloc = torch.cuda.max_memory_allocated() + torch.cuda.reset_peak_memory_stats() + return self.runtime.global_rank, max_alloc + + def reset_peak_memory(self): + torch.cuda.reset_peak_memory_stats() + + def _nvtx_push(self, label: str) -> None: + self.runtime.nvtx_push(label) + + def _nvtx_pop(self) -> None: + self.runtime.nvtx_pop() + + def start_pytorch_profiler(self) -> None: + """Begin a torch.profiler session spanning the next run_dag iterations. + + Each node's execution in DagExecutor is wrapped in a record_function + labelled the same as its NVTX range, so the resulting trace identifies + per-node work. + """ + self.runtime.torch_profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + ) + self.runtime.torch_profiler.__enter__() + self.runtime.pytorch_profiler_enabled = True + self.logger.info(f"Actor {self.runtime.global_rank}: PyTorch profiler started") + + def stop_pytorch_profiler(self, profile_dir: str) -> str: + """End the profiler session and export a chrome trace. + + The file is named ``dp{dp_rank}_pp{pp_rank}.json`` inside *profile_dir* + so the harness can group same-dp-rank actors. Returns the written path. + """ + self.runtime.pytorch_profiler_enabled = False + self.runtime.torch_profiler.__exit__(None, None, None) + os.makedirs(profile_dir, exist_ok=True) + filepath = os.path.join( + profile_dir, f"dp{self.runtime.dp_rank}_pp{self.runtime.pp_rank}.json" + ) + self.runtime.torch_profiler.export_chrome_trace(filepath) + self.runtime.torch_profiler = None + self.logger.info( + f"Actor {self.runtime.global_rank}: PyTorch profiler trace written to {filepath}" + ) + return filepath + + def load_input(self, inputs): + self.inputs = [inp.to(self.runtime.device) for inp in inputs] + self.logger.debug(f"Actor {self.runtime.global_rank} loaded inputs {len(self.inputs)}") + + def load_labels(self, labels): + self.labels = labels.to(self.runtime.device) + self.logger.debug(f"Actor {self.runtime.global_rank} loaded labels {self.labels.shape}") + + def load_const_attrs(self, const_attrs: dict) -> None: + """Store non-trainable constant tensor attributes (e.g. freqs_cis, mask). + + *const_attrs* maps bare attribute name → CPU tensor. Values are moved to + the actor's device so ``_load_stage`` can copy them directly. + """ + self.model_const_attrs = {k: v.to(self.runtime.device) for k, v in const_attrs.items()} + + def get_node_ip_and_free_port(self): + import socket + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("0.0.0.0", 0)) + port = s.getsockname()[1] + return ray.util.get_node_ip_address(), port + + def _join_process_groups(self, master_addr, master_port): + + self.logger.info(f"Actor {self.runtime.global_rank} using GPU {os.environ['CUDA_VISIBLE_DEVICES']}, master addr {master_addr}:{master_port}") + + init_method = f"tcp://{master_addr}:{master_port}" + + self.runtime.device = f"cuda:{self.runtime.global_rank % torch.cuda.device_count()}" + torch.cuda.set_device(self.runtime.device) + + if self.runtime.pp_degree > 1 or self.runtime.dp_degree > 1: + dist.init_process_group( + "nccl", + init_method=init_method, + rank=self.runtime.global_rank, + world_size=self.runtime.world_size, + ) + + if self.runtime.dp_degree > 1: + self._join_dp_process_group() + if self.runtime.pp_degree > 1: + self._join_pp_process_group() + + self.logger.info(f"Actor {self.runtime.global_rank} joined process groups") + + def _join_dp_process_group(self): + num_dp_groups = self.runtime.world_size // self.runtime.dp_degree + for dp_group_id in range(num_dp_groups): + group_ranks = [ + (dp_group_id + num_dp_groups * i) for i in range(self.runtime.dp_degree) + ] + # Two separate NCCL communicators over the same ranks: one for allreduce, + # one for all2all. Sharing a communicator causes both op types to run on + # the same internal NCCL proxy stream, which prevents true overlap. + process_group = dist.new_group(ranks=group_ranks, backend="nccl") + ep_process_group = dist.new_group(ranks=group_ranks, backend="nccl") + if self.runtime.global_rank % num_dp_groups == dp_group_id: + self.runtime.dp_group = process_group + self.runtime.ep_group = ep_process_group + + def _join_pp_process_group(self): + num_pp_groups = self.runtime.world_size // self.runtime.pp_degree + + for pp_group_id in range(num_pp_groups): + group_ranks = [ + (pp_group_id * self.runtime.pp_degree + i) for i in range(self.runtime.pp_degree) + ] + lo_hi_group = dist.new_group(ranks=group_ranks, backend="nccl") + hi_lo_group = dist.new_group(ranks=group_ranks, backend="nccl") + + if self.runtime.global_rank in group_ranks: + self.runtime.pp_lo_hi = lo_hi_group + self.runtime.pp_hi_lo = hi_lo_group + + def _derive_dag_bucket_modes(self, training_dag: Any) -> None: + self.stages.param_sharded_ubids = set() + self.stages.grad_sharded_ubids = set() + for node in training_dag.nodes.values(): + meta = getattr(node, "node_meta", {}) or {} + ubid = meta.get("bucket_key") + if ubid is None: + continue + if node.node_kind == "ALL_GATHER_COMM": + self.stages.param_sharded_ubids.add(ubid) + elif node.node_kind == "REDUCE_SCATTER_COMM": + self.stages.grad_sharded_ubids.add(ubid) + self.stages.zero_managed_ubids = self.stages.param_sharded_ubids | self.stages.grad_sharded_ubids + + def _relocate_meta_devices(self, gm) -> None: + """Rewrite baked ``torch.device('meta')`` literals to the actor device. + + Meta-device compilation captures whatever device was current at trace + time (inside ``type_as`` / ``.to(device=x.device)`` / device-aware + factory ops like ``torch.zeros(..., device=...)``). That literal + round-trips through serialization as ``meta`` and would otherwise make + those ops emit meta tensors at runtime, mismatching the CUDA inputs. + Lifted params/buffers/inputs are unaffected (placed on device directly), + so this only touches device literals embedded in node args/kwargs. + """ + import torch.fx as fx + + device = self.runtime.device + replaced = 0 + + def _fix(value): + nonlocal replaced + if isinstance(value, torch.device): + if value.type == "meta": + replaced += 1 + return torch.device(device) + return value + if isinstance(value, tuple): + return tuple(_fix(v) for v in value) + if isinstance(value, list): + return [_fix(v) for v in value] + if isinstance(value, dict): + return {k: _fix(v) for k, v in value.items()} + return value + + for module in gm.modules(): + if not isinstance(module, fx.GraphModule): + continue + before = replaced + for node in module.graph.nodes: + node.args = _fix(node.args) + node.kwargs = _fix(node.kwargs) + if replaced != before: + module.recompile() + + def _load_stage( + self, + stage_id: int, + modules_data: list, + a2a_boundaries: dict = None, + use_activation_checkpointing: bool = False, + ) -> None: + """Load a (possibly bucketed) stage. + + *modules_data* is a list of dicts, one per module/bucket, each with keys: + ``gm_data``, ``graphargs``, ``input_idxs``, ``param_idxs``, ``bucket_key``. + + A non-bucketed stage is represented as a single-element list. + All per-bucket data structures are keyed by ``bucket_key`` so run_dag + can look up bucket data without knowing which stage owns a bucket. + """ + self.logger.debug( + f"Loading stage {stage_id} ({len(modules_data)} module(s)) on actor {self.runtime.global_rank}" + ) + + g = torch.Generator(device=self.runtime.device) + g.manual_seed(1000 * self.runtime.global_rank + stage_id) + + first_gm = None + + for b_idx, bd in enumerate(modules_data): + ubid: Any = bd["bucket_key"] + bucket = self.stages.ensure_bucket(ubid) + ac_num_subgraphs = int(bd.get("ac_num_subgraphs", 1)) + ac_requested_subgraphs = int(bd.get("ac_requested_subgraphs", ac_num_subgraphs)) + gms = [_deserialize_graphmodule(gm_data) for gm_data in bd["gm_data_list"]] if "gm_data_list" in bd else [_deserialize_graphmodule(bd["gm_data"])] + for gm in gms: + self._relocate_meta_devices(gm) + if self.use_inductor: + compiled_gms = [] + for subgraph_idx, gm in enumerate(gms): + compiled_gm = torch.compile(gm) + compiled_gms.append(compiled_gm) + self.logger.debug( + f"[load_stage_compile] rank={self.runtime.global_rank} stage={stage_id} " + f"ubid={ubid} subgraph={subgraph_idx} compiled=True" + ) + gms = compiled_gms + + if b_idx == 0: + first_gm = gms[0] + + forward_args = list(bd["graphargs"]) + b_input_idxs = list(bd["input_idxs"]) + b_param_idxs = list(bd["param_idxs"]) + apply_zero = bool(bd.get("apply_zero", True)) + shared_placeholder_names = list( + bd.get("shared_placeholder_names", bd.get("placeholder_names", [])) + ) + # Extract FX placeholder names for each param index. + bucket.param_names = [ + shared_placeholder_names[i] if i < len(shared_placeholder_names) else f"ubid{ubid}_p{i}" + for i in b_param_idxs + ] + + # Save input tensor metadata for pre-allocating FWD recv buffers. + # Stored as a list of (shape, dtype, requires_grad) in input-slot order. + recv_meta = [] + for i in b_input_idxs: + meta = forward_args[i] + if meta is not None: + recv_meta.append((tuple(meta.shape), meta.dtype, + bool(getattr(meta, "requires_grad", False)))) + forward_args[i] = None # clear slot; run_dag will fill it at execution time + bucket.forward_input_meta = recv_meta + + bucket.input_idxs = b_input_idxs + + # Realize parameter tensors. + realized = [None] * len(forward_args) + for i, arg in enumerate(forward_args): + if arg is None: + continue + t = torch.empty(arg.shape, dtype=arg.dtype, device=self.runtime.device) + if arg.requires_grad: + t.requires_grad_(True) + torch.nn.init.normal_(t, mean=0.0, std=0.02, generator=g) + else: + # Non-trainable slot: try to fill from const attrs (freqs_cis, mask, …) + # before falling back to zeros. Dynamo names direct model attrs as + # "l_self_", so strip that prefix to get the bare name. + ph_name = shared_placeholder_names[i] if i < len(shared_placeholder_names) else "" + attr_name = re.sub(r'^l_self_', '', ph_name) + const_val = self.model_const_attrs.get(attr_name) + if ( + const_val is not None + and tuple(const_val.shape) == tuple(arg.shape) + and const_val.dtype == arg.dtype + ): + t.copy_(const_val) + else: + t.zero_() + realized[i] = t + + shared_name_to_idx = { + name: i for i, name in enumerate(shared_placeholder_names) + } + subgraph_specs = [] + for gm in gms: + sub_placeholder_names = [ + n.name for n in gm.graph.nodes if n.op == "placeholder" + ] + dynamic_names = [ + name for name in sub_placeholder_names + if name not in shared_name_to_idx + ] + subgraph_specs.append((gm.forward, sub_placeholder_names, dynamic_names)) + + def _bucket_forward_runner( + shared_args, + _specs=subgraph_specs, + _name_to_idx=shared_name_to_idx, + _use_ac=use_activation_checkpointing, + _stage_id=stage_id, + _ubid=ubid, + _shared_placeholder_names=tuple(shared_placeholder_names), + ): + out = None + for subgraph_idx, (forward_impl, placeholder_names, dynamic_names) in enumerate(_specs): + if len(dynamic_names) == 0: + dynamic_values = [] + elif len(dynamic_names) == 1: + dyn_arg = out + if isinstance(dyn_arg, (tuple, list)) and len(dyn_arg) == 1: + dyn_arg = dyn_arg[0] + dynamic_values = [dyn_arg] + else: + if not isinstance(out, (tuple, list)): + raise RuntimeError( + f"Bucket forward expected {len(dynamic_names)} dynamic inputs " + f"but previous subgraph produced {type(out).__name__}" + ) + dynamic_values = list(out) + if len(dynamic_values) != len(dynamic_names): + raise RuntimeError( + f"Bucket forward expected {len(dynamic_names)} dynamic inputs " + f"but previous subgraph produced {len(dynamic_values)} values" + ) + dynamic_name_to_value = { + name: value for name, value in zip(dynamic_names, dynamic_values) + } + call_args = [] + for name in placeholder_names: + if name in dynamic_name_to_value: + call_args.append(dynamic_name_to_value[name]) + else: + call_args.append(shared_args[_name_to_idx[name]]) + if _use_ac: + out = torch.utils.checkpoint.checkpoint( + forward_impl, + *call_args, + use_reentrant=False, + ) + else: + out = forward_impl(*call_args) + return out + + bucket.forward_fn = _bucket_forward_runner + bucket.forward_args = realized + bucket.param_idxs = b_param_idxs + bucket.activation_checkpoint_subgraph_count = ac_num_subgraphs + trainable_idxs = [ + i for i in b_param_idxs + if realized[i] is not None and realized[i].requires_grad + ] + bucket.trainable_param_idxs = trainable_idxs + + zero_managed = ( + self.runtime.dp_degree > 1 + and apply_zero + and bool(trainable_idxs) + and ubid in self.stages.zero_managed_ubids + ) + params_sharded = ubid in self.stages.param_sharded_ubids + grads_sharded = ubid in self.stages.grad_sharded_ubids + + if zero_managed: + trainable = [realized[i] for i in trainable_idxs] + flat_params = torch.cat([p.detach().view(-1) for p in trainable]).contiguous() + flat_params.requires_grad_(False) + orig_numel = flat_params.numel() + dp = self.runtime.dp_degree + shard_size = (orig_numel + dp - 1) // dp + padded_numel = shard_size * dp + if padded_numel > orig_numel: + padded = flat_params.new_zeros(padded_numel) + padded[:orig_numel].copy_(flat_params) + flat_params = padded + + offset = 0 + view_specs = [] + for idx, p in zip(trainable_idxs, trainable): + numel = p.numel() + realized[idx] = realized[idx].detach() + realized[idx].data = flat_params[offset:offset + numel].view(p.shape) + realized[idx].requires_grad_(True) + realized[idx].grad_dtype = self.params.grad_buffer_dtype + view_specs.append((realized[idx], offset, numel, tuple(p.shape))) + offset += numel + + shard_start = self.runtime.dp_rank * shard_size + if params_sharded: + shard_param = flat_params[shard_start:shard_start + shard_size].detach().clone() + else: + shard_param = flat_params[shard_start:shard_start + shard_size] + shard_param.requires_grad_(True) + # Keep ZeRO comm/storage buffers in fp32, but let the optimizer consume + # grads in the shard-param dtype to match Adam's bf16 foreach path. + shard_param.grad_dtype = None + + bucket.flat_params = flat_params + bucket.flat_grads = ( + None + if grads_sharded + else torch.zeros(padded_numel, dtype=self.params.grad_buffer_dtype, device=self.runtime.device) + ) + bucket.shard_param = shard_param + bucket.shard_optimizer = self.optim_class([shard_param]) + bucket.reduce_scatter_grads = ( + torch.zeros(shard_size, dtype=self.params.grad_buffer_dtype, device=self.runtime.device) + if grads_sharded + else None + ) + bucket.param_shard_info = (shard_start, shard_size, orig_numel) + bucket.param_view_specs = view_specs + bucket.full_params_fresh = False + + if params_sharded: + storage = flat_params.untyped_storage() + if storage.size() != 0: + storage.resize_(0) + + optim = None + else: + # Keep trainable params as separate tensors for the non-ZeRO path. + bucket.param_view_specs = [] + bucket.flat_params = None + bucket.flat_grads = None + bucket.shard_param = None + bucket.shard_optimizer = None + bucket.reduce_scatter_grads = None + bucket.param_shard_info = None + bucket.full_params_fresh = False + trainable_for_optim = [realized[i] for i in trainable_idxs] + optim = self.optim_class(trainable_for_optim, fused=True) if trainable_for_optim else None + bucket.optimizer = optim + + # Keep first GraphModule for compatibility with external inspection tools. + self.stages.graph_modules[stage_id] = first_gm + + def load_training_dag(self, training_dag: Any) -> None: + """Load per-PP-rank TrainingDAG compute nodes using existing _load_stage logic. + + This is an adapter for the new TrainingDAG representation where each + COMPUTE/FWD node with ``node_meta['gm']`` is equivalent to one bucket/module. + The runtime executes TrainingDAG nodes via ``run_dag``; this method + only populates actor-side bucket/module state. + """ + if training_dag is None: + raise ValueError("load_training_dag requires non-None training_dag") + if not hasattr(training_dag, "nodes"): + raise TypeError("load_training_dag expected object with 'nodes' field") + + self._derive_dag_bucket_modes(training_dag) + self.runtime.initialize_streams_for_training_dag(training_dag) + + # Clear any prior module state before loading a new training DAG. + self.stages.clear_loaded_modules() + + stage_to_modules: dict[int, list[dict]] = defaultdict(list) + + # Deterministic order by stage/segment then uid. + compute_nodes = [ + n for n in training_dag.nodes.values() + if getattr(n, "node_kind", None) == "COMPUTE" + and getattr(n, "compute_subkind", None) == "FWD" + ] + compute_nodes.sort( + key=lambda n: ( + int(getattr(n, "node_meta", {}).get("stage_id", 10**9)), + int(getattr(n, "node_meta", {}).get("segment_id", 10**9)), + str(getattr(n, "uid", "")), + ) + ) + + seen_bucket_keys: set[Any] = set() + skipped_dupe_fwd_nodes = 0 + for node in compute_nodes: + meta = getattr(node, "node_meta", {}) or {} + bucket_key = meta.get("bucket_key", getattr(node, "uid", None)) + if bucket_key in seen_bucket_keys: + skipped_dupe_fwd_nodes += 1 + continue + seen_bucket_keys.add(bucket_key) + gm = meta.get("gm") + gm_data = meta.get("gm_data") + stage_id = meta.get("stage_id") + input_idxs = meta.get("input_idxs") + param_idxs = meta.get("param_idxs") + graphargs = meta.get("graphargs") + input_names = meta.get("input_names", []) + output_names = meta.get("output_names", []) + + if gm_data is None and gm is None: + raise ValueError( + f"load_training_dag: compute node {getattr(node, 'uid', '')} " + f"is missing required metadata field 'gm_data' (or fallback 'gm')" + ) + if stage_id is None or input_idxs is None or param_idxs is None or graphargs is None: + raise ValueError( + f"load_training_dag: compute node {getattr(node, 'uid', '')} " + f"is missing required metadata fields" + ) + + module_data = { + "gm_data": gm_data if gm_data is not None else _serialize_graphmodule(gm), + "graphargs": list(graphargs), + "input_idxs": list(input_idxs), + "param_idxs": list(param_idxs), + "placeholder_names": list(input_names), + "output_names": list(output_names), + "shared_placeholder_names": list(input_names), + "bucket_key": bucket_key, + "ac_num_subgraphs": 1, + "ac_requested_subgraphs": 1, + "apply_zero": bool(meta.get("apply_zero", True)), + "training_dag_uid": getattr(node, "uid", None), + "triton_constant_args": dict(meta.get("triton_constant_args", {})), + } + stage_to_modules[int(stage_id)].append(module_data) + + assert stage_to_modules, ( + "load_training_dag: no FWD compute nodes with gm metadata found" + ) + + for stage_id in sorted(stage_to_modules.keys()): + self._load_stage( + stage_id=stage_id, + modules_data=stage_to_modules[stage_id], + a2a_boundaries={}, + use_activation_checkpointing=False, + ) + + # Build runtime adjacency on TrainingDAG nodes so run_dag can execute directly. + for n in training_dag.nodes.values(): + n.data_preds = [] + n.data_succs = [] + n.temporal_preds = [] + n.temporal_succs = [] + if n.node_kind == "SEND_COMM": + n.peer_pp_rank = n.node_meta.get("peer_pp_rank") + elif n.node_kind == "RECV_COMM": + n.peer_pp_rank = n.node_meta.get("peer_pp_rank") + else: + n.peer_pp_rank = None + for e in training_dag.edges: + if e.src_uid not in training_dag.nodes or e.dst_uid not in training_dag.nodes: + continue + src = training_dag.nodes[e.src_uid] + dst = training_dag.nodes[e.dst_uid] + if e.dep_kind == "temporal": + src.temporal_succs.append(dst) + dst.temporal_preds.append(src) + else: + src.data_succs.append(dst) + dst.data_preds.append(src) + + for n in training_dag.nodes.values(): + n.task_type = _training_dag_task_type(n) + mb = n.tag.get("MB", 0) + st = n.tag.get("PP", 0) + n.batches = [type("RuntimeBatch", (), {"stage_id": st, "mb_idx": mb})()] + + self.dag = training_dag + self.sorted_dag_nodes = [ + training_dag.nodes[uid] for uid in _serial_topological_order(training_dag) + ] + + def run_dag(self, loss_fn=None): + # Mark the entire iteration boundary for the NVTX timeline. + iter_idx = getattr(self, "_iter_counter", 0) + self._iter_counter = iter_idx + 1 + self._nvtx_push(f"iter_{iter_idx}_rank_{self.runtime.global_rank}") + self.dag_executor.run( + self.dag, + self.sorted_dag_nodes, + self.inputs, + self.labels, + self.loss, + loss_fn=loss_fn, + ) + self._nvtx_pop() diff --git a/src/backward.py b/src/backward.py new file mode 100644 index 0000000..fb5bf56 --- /dev/null +++ b/src/backward.py @@ -0,0 +1,169 @@ +# =========================================================================================== +# This file contains helper functions to implement the split backward pass. +# +# These helper functions were copied from the PyTorch pipelining backward implementation. +# +# PyTorch Team, +# "torch.distributed.pipelining._backward", +# https://github.com/pytorch/pytorch/blob/main/torch/distributed/pipelining/_backward.py +# +# There are some minor implementation differences, but these functions are nearly identical. +# +# =========================================================================================== + + +from __future__ import annotations + +import collections +from typing import Any, Dict, Iterator, List, Optional, Set, Tuple + +import torch +from torch.autograd.graph import GradientEdge, Node +from torch.nn import Parameter + + +def _get_grad_fn_or_grad_acc(t: torch.Tensor) -> Optional[Node]: + """ + Return the autograd Node for a tensor: + - non-leaf tensors: t.grad_fn + - leaf tensors requiring grad (Parameters): AccumulateGrad node (created lazily) + """ + if not t.requires_grad: + return None + if t.grad_fn is not None: + return t.grad_fn + + # If t is a leaf tensor, we create a view to generate its AccumulateGrad node + viewed = t.view_as(t) + grad_fn = viewed.grad_fn + if grad_fn is None: + raise RuntimeError( + "Tried to get grad accumulator but got None. " + "Are you in a no-grad context?" + ) + return grad_fn.next_functions[0][0] + + +def reverse_closure( + roots: List[Node], + target_nodes: Set[Node], + reverse_edges: Dict[Node, List[Node]], +) -> Tuple[Set[Node], Set[Node]]: + """ + From roots, follow reverse_edges to collect reachable nodes, but stop traversal + when hitting any node in target_nodes. Returns the closure (set of reachable nodes) and + the set of target nodes that were able to be visited. + """ + closure: Set[Node] = set() + visited_targets: Set[Node] = set() + q: collections.deque[Node] = collections.deque() + + for node in roots: + if node is None or node in closure: + continue + closure.add(node) + q.append(node) + + while q: + currNode = q.popleft() + for node in reverse_edges.get(currNode, []): + if node is None or node in closure: + continue + if node in target_nodes: + visited_targets.add(node) + continue + closure.add(node) + q.append(node) + + return closure, visited_targets + + +def construct_reverse_graph(roots: List[Node]) -> Dict[Node, List[Node]]: + """ + Builds a reverse graph. In a backward graph, for each node X, + reverse_edges[X] is the list of nodes that X sends gradients to. + """ + q: collections.deque[Node] = collections.deque() + seen: Set[Node] = set() + reverse_edges: Dict[Node, List[Node]] = collections.defaultdict(list) + + for node in roots: + if node is not None and node not in seen: + seen.add(node) + q.append(node) + + while q: + node = q.popleft() + for fn, _ in node.next_functions: + if fn is None: + continue + + # fn -> node in reverse graph + reverse_edges[fn].append(node) + if fn not in seen: + seen.add(fn) + q.append(fn) + + return reverse_edges + + +def get_param_groups( + inputs: List[Node], + params: List[Node], + reverse_edges: Dict[Node, List[Node]], +) -> List[Dict[str, Any]]: + """ + Returns a list of parameter groups. Parameters are in the same group if they share the same + intermediate boundary nodes where the 'param-side' graph intersects the input closure. + """ + inputs_closure, _ = reverse_closure(inputs, set(), reverse_edges) + + intermediate_to_param_group: Dict[Node, Dict[str, Any]] = {} + solo_groups: List[Dict[str, Any]] = [] + + for p in params: + if p is None: + continue + + _, intersected_nodes = reverse_closure([p], inputs_closure, reverse_edges) + + if not intersected_nodes: + # if the parameter doesn't intersect inputs (unused or disconnected), then it + # we place it in its own group + solo_groups.append({"params": {p}, "intermediates": []}) + continue + + # Merge into an existing group if any intersected intermediate node already has one + group: Optional[Dict[str, Any]] = None + for inter_node in intersected_nodes: + if inter_node in intermediate_to_param_group: + group = intermediate_to_param_group[inter_node] + break + + if group is None: + group = {"params": set(), "intermediates": []} + + group["params"].add(p) + + # remap all intersected intermediates to this group + for inter_node in intersected_nodes: + intermediate_to_param_group[inter_node] = group + + # add unique groups to param_groups + param_groups: List[Dict[str, Any]] = [] + seen_ids: Set[int] = set() + for params in intermediate_to_param_group.values(): + if id(params) in seen_ids: + continue + seen_ids.add(id(params)) + param_groups.append(params) + + for params in param_groups: + inters = [k for k, v in intermediate_to_param_group.items() if v is params] + # sort nodes to maintain deterministic order when iterating over them in backward_weight + inters.sort(key=id) + params["intermediates"] = inters + + param_groups.extend(solo_groups) + + return param_groups diff --git a/src/bucket.py b/src/bucket.py new file mode 100644 index 0000000..6490f63 --- /dev/null +++ b/src/bucket.py @@ -0,0 +1,450 @@ +import operator +from collections import defaultdict + +import torch +import torch.fx as fx + +from .state import LOG_LEVEL, create_logger + +logger = create_logger("bucket", LOG_LEVEL) + + +# --------------------------------------------------------------------------- +# bucket_stage - stage-level parameter bucketing (placeholder params) +# --------------------------------------------------------------------------- +def bucket_stage( + stage_gm: fx.GraphModule, + graphargs: list, + input_idxs: list[int], + param_idxs: list[int], + bucket_size_bytes: int = 25 * 1024 * 1024, + debug_name: str | None = None, +) -> list[tuple[fx.GraphModule, list[int], list[int], list]]: + """Split a stage GraphModule into per-parameter-bucket sub-modules. + + This function handles annotation segments where trainable parameters are + passed as **placeholder inputs** identified by *param_idxs*. + + Args: + stage_gm: The stage ``fx.GraphModule`` to split. + graphargs: Flat arg list for ``stage_gm.forward``; entries at + *param_idxs* are meta tensors, entries at *input_idxs* are None. + input_idxs: Positions of activation-input placeholders. + param_idxs: Positions of parameter placeholders. + bucket_size_bytes: Target bucket size (default 25 MB). + debug_name: Optional label included in bucket planning logs. + + Returns: + A list of ``(bucket_gm, bucket_input_idxs, bucket_param_idxs, + bucket_graphargs)`` tuples in execution order. Returns a + single-element list containing the original stage when no split is + needed. + """ + nodes = list(stage_gm.graph.nodes) + node_idx: dict[fx.Node, int] = {nd: i for i, nd in enumerate(nodes)} + ph_nodes = [nd for nd in nodes if nd.op == "placeholder"] + + param_ph_set = {ph_nodes[i] for i in param_idxs if i < len(ph_nodes)} + param_ph_list = [ph_nodes[i] for i in param_idxs if i < len(ph_nodes)] + input_ph_set = {ph_nodes[i] for i in input_idxs if i < len(ph_nodes)} + + if not param_ph_list: + return [(stage_gm, list(input_idxs), list(param_idxs), list(graphargs))] + + compute_set: frozenset[int] = frozenset( + i for i, nd in enumerate(nodes) + if nd.op not in ("placeholder", "get_attr", "output") + ) + + def _use_range(pnd: fx.Node) -> tuple[int, int]: + idxs = [node_idx[u] for u in pnd.users if node_idx[u] in compute_set] + return (min(idxs), max(idxs)) if idxs else (node_idx[pnd], node_idx[pnd]) + + def _size(pnd: fx.Node) -> int: + i = ph_nodes.index(pnd) + if i < len(graphargs) and graphargs[i] is not None and hasattr(graphargs[i], "numel"): + return int(graphargs[i].numel() * graphargs[i].element_size()) + ev = pnd.meta.get("example_value") + if ev is not None and hasattr(ev, "numel"): + return int(ev.numel() * ev.element_size()) + return 0 + + debug_name = debug_name or getattr(stage_gm, "__class__", type(stage_gm)).__name__ + param_ranges = {pn: _use_range(pn) for pn in param_ph_list} + param_sizes = {pn: _size(pn) for pn in param_ph_list} + + # Greedy bucket assignment + bucket_id: dict[fx.Node, int] = {} + cur_b, cur_sz = 0, 0 + for pn in param_ph_list: + sz = param_sizes[pn] + if cur_sz + sz > bucket_size_bytes and cur_sz > 0: + cur_b += 1 + cur_sz = 0 + bucket_id[pn] = cur_b + cur_sz += sz + n_init = cur_b + 1 + + # Singleton promotion + def _init_range(b: int) -> tuple[int, int]: + ms = [pn for pn in param_ph_list if bucket_id[pn] == b] + if not ms: + return (0, 0) + return (min(param_ranges[pn][0] for pn in ms), max(param_ranges[pn][1] for pn in ms)) + + init_ranges = [_init_range(b) for b in range(n_init)] + + def _bucket_at(idx: int) -> int: + for b, (lo, hi) in enumerate(init_ranges): + if lo <= idx <= hi: + return b + return 0 + + sn = n_init + for pn in param_ph_list: + f, l = param_ranges[pn] + if _bucket_at(f) != _bucket_at(l): + bucket_id[pn] = sn + sn += 1 + + # Build bucket list, sort, merge overlapping ranges + bm: dict[int, list[fx.Node]] = defaultdict(list) + for pn in param_ph_list: + bm[bucket_id[pn]].append(pn) + + blist: list[tuple[int, int, list[fx.Node]]] = [] + for members in bm.values(): + f = min(param_ranges[pn][0] for pn in members) + l = max(param_ranges[pn][1] for pn in members) + blist.append((f, l, members)) + blist.sort(key=lambda t: t[0]) + + merged: list[tuple[int, int, list[fx.Node]]] = [] + for f, l, ms in blist: + if merged and f <= merged[-1][1]: + pf, pl, pm = merged[-1] + merged[-1] = (pf, max(pl, l), pm + ms) + else: + merged.append((f, l, ms)) + + def _bucket_bytes(members: list[fx.Node]) -> int: + return sum(param_sizes[pn] for pn in members) + + def _format_members(members: list[fx.Node], *, limit: int = 8) -> str: + def _one(pn: fx.Node) -> str: + return f"{pn.name}:{param_sizes[pn]}" + + if len(members) <= limit: + return "[" + ", ".join(_one(pn) for pn in members) + "]" + head = ", ".join(_one(pn) for pn in members[:4]) + tail = ", ".join(_one(pn) for pn in members[-2:]) + return f"[{head}, ... ({len(members) - 6} more), {tail}]" + + def _log_bucket_plan(phase: str, seg_ranges: list[tuple[int, int, list[fx.Node]]]) -> None: + ownership: dict[fx.Node, list[int]] = defaultdict(list) + for seg_idx, (_f, _l, members) in enumerate(seg_ranges): + for pn in members: + ownership[pn].append(seg_idx) + missing = [pn.name for pn in param_ph_list if pn not in ownership] + duplicates = [ + (pn.name, buckets) + for pn, buckets in ownership.items() + if len(buckets) != 1 + ] + total_bytes = sum(param_sizes[pn] for pn in param_ph_list) + max_bucket_bytes = max((_bucket_bytes(ms) for _f, _l, ms in seg_ranges), default=0) + logger.info( + "bucket_stage plan phase=%s node=%s bucket_size_bytes=%d bucket_size_mib=%.6f " + "params=%d total_param_bytes=%d initial_buckets=%d final_buckets=%d " + "owned_params=%d unique_owned_params=%d missing_params=%d duplicate_params=%d max_bucket_bytes=%d", + phase, + debug_name, + bucket_size_bytes, + bucket_size_bytes / (1024 * 1024), + len(param_ph_list), + total_bytes, + n_init, + len(seg_ranges), + sum(len(ms) for _f, _l, ms in seg_ranges), + len(ownership), + len(missing), + len(duplicates), + max_bucket_bytes, + ) + if missing or duplicates: + logger.warning( + "bucket_stage ownership-invalid phase=%s node=%s missing=%s duplicates=%s", + phase, + debug_name, + missing[:16], + duplicates[:16], + ) + + merge_candidates: list[tuple[int, int, int]] = [] + for left_idx in range(len(seg_ranges) - 1): + left_bytes = _bucket_bytes(seg_ranges[left_idx][2]) + right_bytes = _bucket_bytes(seg_ranges[left_idx + 1][2]) + combined = left_bytes + right_bytes + if combined <= bucket_size_bytes: + merge_candidates.append((left_idx, left_idx + 1, combined)) + + for seg_idx, (first, last, members) in enumerate(seg_ranges): + bytes_ = _bucket_bytes(members) + if bytes_ <= bucket_size_bytes: + status = "within_limit" + elif len(members) == 1: + status = "single_param_over_limit" + else: + status = "forced_multi_param_over_limit" + logger.info( + "bucket_stage bucket-plan phase=%s node=%s bucket=%d/%d param_bytes=%d " + "params=%d range=%d:%d status=%s members=%s", + phase, + debug_name, + seg_idx, + len(seg_ranges), + bytes_, + len(members), + first, + last, + status, + _format_members(members), + ) + if merge_candidates: + logger.warning( + "bucket_stage fullness-warning phase=%s node=%s mergeable_adjacent_pairs=%s " + "bucket_size_bytes=%d", + phase, + debug_name, + merge_candidates[:16], + bucket_size_bytes, + ) + else: + logger.info( + "bucket_stage fullness-ok phase=%s node=%s adjacent_pairs=%d bucket_size_bytes=%d", + phase, + debug_name, + max(0, len(seg_ranges) - 1), + bucket_size_bytes, + ) + + _log_bucket_plan("range_merge", merged) + + alias_methods = { + "view", "_unsafe_view", "reshape", "transpose", "permute", "t", + "movedim", "moveaxis", "swapdims", "swapaxes", + "select", "narrow", "slice", "split", "chunk", "unbind", + "unsqueeze", "squeeze", "flatten", "expand", "diagonal", + "detach", "alias", "as_strided", + } + multi_output_alias_methods = {"split", "chunk", "unbind"} + + alias_function_names = { + "view", "_unsafe_view", "reshape", "transpose", "permute", "t", + "movedim", "moveaxis", "swapdims", "swapaxes", + "select", "narrow", "slice", "split", "chunk", "unbind", + "unsqueeze", "squeeze", "flatten", "expand", "diagonal", + "detach", "alias", "as_strided", + } + multi_output_alias_function_names = {"split", "chunk", "unbind"} + hard_forbid_cross_bucket_prefixes = ("bfloat16_",) + + def _target_name(target: object) -> str | None: + if isinstance(target, str): + return target + return getattr(target, "__name__", None) + + def _alias_passthrough_sources(nd: fx.Node) -> set[fx.Node]: + if nd.op == "call_method": + if nd.target in alias_methods: + base = next(iter(nd.all_input_nodes), None) + return set(alias_sources.get(base, set())) if base is not None else set() + return set() + + if nd.op != "call_function": + return set() + + if nd.target == operator.getitem: + base = nd.args[0] if nd.args else None + if isinstance(base, fx.Node): + return set(alias_sources.get(base, set())) + return set() + + name = _target_name(nd.target) + if name in alias_function_names: + base = next(iter(nd.all_input_nodes), None) + return set(alias_sources.get(base, set())) if base is not None else set() + return set() + + alias_sources: dict[fx.Node, set[fx.Node]] = {} + for nd in nodes: + if nd.op == "placeholder": + alias_sources[nd] = {nd} if nd in param_ph_set else set() + continue + if nd.op in ("get_attr", "output"): + alias_sources[nd] = set() + continue + alias_sources[nd] = _alias_passthrough_sources(nd) + + def _compute_seg_metadata( + seg_ranges: list[tuple[int, int, list[fx.Node]]] + ) -> tuple[dict[fx.Node, int], dict[fx.Node, int], list[list[fx.Node]]]: + """Return (node_seg, node_max_user_seg, seg_cross_in) for segment ranges.""" + cut_after = [seg_ranges[i][1] for i in range(len(seg_ranges) - 1)] + + def _seg_of(idx: int) -> int: + return sum(1 for c in cut_after if c < idx) + + node_seg: dict[fx.Node, int] = {} + for nd in nodes: + if nd.op == "output": + node_seg[nd] = len(seg_ranges) - 1 + elif nd.op == "placeholder": + if nd in param_ph_set: + user_idxs = [node_idx[u] for u in nd.users if node_idx[u] in compute_set] + node_seg[nd] = _seg_of(min(user_idxs)) if user_idxs else 0 + else: + node_seg[nd] = 0 + elif nd.op == "get_attr": + user_idxs = [node_idx[u] for u in nd.users if node_idx[u] in compute_set] + node_seg[nd] = _seg_of(min(user_idxs)) if user_idxs else 0 + else: + node_seg[nd] = _seg_of(node_idx[nd]) + + node_max_user_seg: dict[fx.Node, int] = {} + for nd in nodes: + if nd.op == "output": + continue + node_max_user_seg[nd] = ( + max(node_seg[u] for u in nd.users) if nd.users else node_seg[nd] + ) + + seg_cross_in: list[list[fx.Node]] = [[] for _ in range(len(seg_ranges))] + for nd in nodes: + if nd.op == "output": + continue + s = node_seg[nd] + mu = node_max_user_seg[nd] + # Only values that cross a segment boundary must be forwarded. + if mu <= s: + continue + for seg in range(s + 1, mu + 1): + seg_cross_in[seg].append(nd) + + return node_seg, node_max_user_seg, seg_cross_in + + n_segs = len(merged) + if n_segs == 1: + return [(stage_gm, list(input_idxs), list(param_idxs), list(graphargs))] + + while True: + node_seg, node_max_user_seg, seg_cross_in = _compute_seg_metadata(merged) + merged_boundary = False + for seg in range(1, len(merged)): + alias_crossers = [ + nd for nd in seg_cross_in[seg] + if alias_sources.get(nd) or any( + nd.name.startswith(prefix) for prefix in hard_forbid_cross_bucket_prefixes + ) + ] + if not alias_crossers: + continue + pf, _pl, pms = merged[seg - 1] + _cf, cl, cms = merged[seg] + merged[seg - 1] = (pf, cl, pms + cms) + del merged[seg] + merged_boundary = True + break + if not merged_boundary: + break + + _log_bucket_plan("final", merged) + + n_segs = len(merged) + if n_segs == 1: + return [(stage_gm, list(input_idxs), list(param_idxs), list(graphargs))] + + # Build sub-graphs + results: list[tuple[fx.GraphModule, list[int], list[int], list]] = [] + + for seg in range(n_segs): + sub_g = fx.Graph() + remap: dict[fx.Node, fx.Node] = {} + new_input_idxs: list[int] = [] + new_param_idxs: list[int] = [] + new_graphargs: list = [] + pos = 0 + + def _meta_tensor_for(nd: fx.Node): + """Return a meta tensor matching nd's shape/dtype, or None if unavailable.""" + ev = nd.meta.get("example_value") + if ev is None: + ev = nd.meta.get("val") + if ev is not None and hasattr(ev, "shape"): + return torch.empty(ev.shape, dtype=ev.dtype, device="meta", requires_grad=ev.requires_grad) + return None + + def _add_ph(nd: fx.Node, name: str, is_input: bool) -> None: + nonlocal pos + new_ph = sub_g.placeholder(name) + new_ph.type = nd.type + remap[nd] = new_ph + if is_input: + new_input_idxs.append(pos) + meta = _meta_tensor_for(nd) + if meta is None: + # Fallback: use the corresponding entry from the incoming + # graphargs for freshly-created placeholder nodes that lack + # example_value metadata. + i_orig = ph_nodes.index(nd) if nd in ph_nodes else -1 + if 0 <= i_orig < len(graphargs): + meta = graphargs[i_orig] + new_graphargs.append(meta) + else: + new_param_idxs.append(pos) + i_orig = ph_nodes.index(nd) if nd in ph_nodes else -1 + new_graphargs.append(graphargs[i_orig] if 0 <= i_orig < len(graphargs) else None) + pos += 1 + + if seg == 0: + # Non-param placeholders (activation inputs + other non-param inputs). + for nd in nodes: + if nd.op == "placeholder" and nd not in param_ph_set: + _add_ph(nd, nd.name, is_input=(nd in input_ph_set)) + else: + # Cross-segment activation inputs from the previous segment. + for orig in seg_cross_in[seg]: + _add_ph(orig, f"_xseg_{orig.name}", is_input=True) + + # Parameter placeholders belonging to this segment. + for nd in nodes: + if nd.op == "placeholder" and nd in param_ph_set and node_seg[nd] == seg: + _add_ph(nd, nd.name, is_input=False) + + # get_attr nodes for this segment. + for nd in nodes: + if nd.op == "get_attr" and node_seg[nd] == seg: + new_ga = sub_g.get_attr(nd.target) + new_ga.type = nd.type + remap[nd] = new_ga + + # Compute nodes for this segment (graph order preserved). + for nd in nodes: + if nd.op in ("placeholder", "get_attr", "output"): + continue + if node_seg[nd] != seg: + continue + remap[nd] = sub_g.node_copy(nd, arg_transform=lambda x, r=remap: r[x]) + + # Output node. + if seg == n_segs - 1: + orig_out = next(nd for nd in nodes if nd.op == "output") + sub_g.output(fx.map_arg(orig_out.args[0], lambda x: remap[x])) + else: + out_nodes = [remap[orig] for orig in seg_cross_in[seg + 1]] + sub_g.output(tuple(out_nodes) if len(out_nodes) != 1 else out_nodes[0]) + + sub_g.lint() + results.append((fx.GraphModule(stage_gm, sub_g), new_input_idxs, new_param_idxs, new_graphargs)) + + return results diff --git a/src/compile.py b/src/compile.py new file mode 100644 index 0000000..8280fdd --- /dev/null +++ b/src/compile.py @@ -0,0 +1,345 @@ +import ray +import torch +import copy +import os +import time +from pathlib import Path + +from .actor import _create_actors +from .state import ( + piper_metadata, + create_logger, + LOG_LEVEL, +) +from .fx import _serialize_graphmodule +from .piper import _reset_annotation_state, piper +from .schedule import ( + derive_schedule_info, + load_schedule_directives, +) +logger = create_logger("compile", LOG_LEVEL) + +_RANK0_ADDR_ACTOR = "piper_rank0_addr" +_COMPILED_DATA_ACTOR = "piper_compiled_data" + +@ray.remote +class _Rank0AddrStore: + """Named actor used to share global rank-0's IP+port across independent dp_rank workers.""" + def __init__(self, addr: str, port: int): + self._addr = addr + self._port = port + def get(self): + return self._addr, self._port + + +@ray.remote +class _CompiledDataStore: + """Named actor used to share compiled training-DAG data from dp_rank=0 to dp_rank>0.""" + def __init__(self): + self._training_dag_data = None + self._training_dag_ready = False + + def publish_training_dag_data(self, data: dict): + self._training_dag_data = data + self._training_dag_ready = True + + def is_ready(self) -> bool: + return self._training_dag_ready + + def is_training_dag_ready(self) -> bool: + return self._training_dag_ready + + def get(self) -> dict: + if not self.is_ready(): + return None + return dict(self._training_dag_data) + + def get_training_dag_data(self) -> dict: + return self._training_dag_data + + + + + +torch._dynamo.config.capture_scalar_outputs = True + + +def _prepare_training_dags_for_rpc(per_pp_training_dags: list) -> list: + """Return a Ray-safe copy of per-PP training DAGs. + + Raw fx.GraphModule objects in node_meta are replaced by serialized strings + under ``gm_data`` to avoid Ray deserialization invoking FX tracing on HOPs. + """ + rpc_dags = copy.deepcopy(per_pp_training_dags) + for dag in rpc_dags: + for node in dag.nodes.values(): + if getattr(node, "node_kind", None) != "COMPUTE": + continue + if getattr(node, "compute_subkind", None) != "FWD": + continue + meta = getattr(node, "node_meta", None) + if not isinstance(meta, dict): + raise ValueError(f"Malformed COMPUTE node metadata for uid={getattr(node, 'uid', '')}") + gm = meta.get("gm") + if gm is None: + raise ValueError( + f"Malformed COMPUTE/FWD node {getattr(node, 'uid', '')}: missing node_meta['gm']" + ) + meta["gm_data"] = _serialize_graphmodule(gm) + del meta["gm"] + return rpc_dags + + +def _compute_loss_pp_ranks(per_pp_training_dags: list) -> list[int]: + ranks: list[int] = [] + for pp_rank, dag in enumerate(per_pp_training_dags): + if any( + isinstance(getattr(node, "node_meta", None), dict) + and bool(node.node_meta.get("compute_loss", False)) + for node in dag.nodes.values() + ): + ranks.append(pp_rank) + return ranks + + +def _resolve_model_constructor_value(value, schedule_info: dict): + if callable(value): + return value(schedule_info) + return value + + +def _artifact_dir_for_schedule(schedule_directives_file: str) -> str: + parent = Path(schedule_directives_file).parent + return str(parent if str(parent) else Path("out")) + + +def piper_setup( + model_class, + model_args=(), + model_kwargs={}, + model_dtype=torch.bfloat16, + optim_fn=None, + example_inputs=None, + example_outputs=None, + activation_checkpointing=False, + num_checkpoints=1, + use_inductor: bool = False, + no_nvtx: bool = False, + pg=None, + nsight=False, + temp_dir: str = None, + visualize_dag: bool = False, + const_attrs: dict = None, + pp_outer: bool = False, + schedule_directives_file: str | None = None, +): + """ + Compile a model with the TrainingDAG backend. + + The model and example inputs are always traced on the meta device so the + full model is never materialized on the driver GPU at compile time. + + Args: + model: The model to compile. + optim_fn: Callable ``(params) -> Optimizer`` used to create optimizers. + example_inputs: Example inputs for tracing. + example_outputs: Example outputs (labels) for tracing. + use_inductor: When true, actors torch.compile stage GraphModules + during _load_stage with the default inductor backend. + schedule_directives_file: JSON schedule description consumed by the + TrainingDAG backend. When provided, Piper derives pp/dp/mbs + schedule metadata internally. + """ + + # Clear Dynamo's global compilation cache so that a previous piper_setup + # call (e.g. with a different model size or schedule) in the same Python + # worker process cannot contaminate this compilation via a stale cache hit. + torch._dynamo.reset() + + schedule_info = {} + if schedule_directives_file is not None: + schedule_directives = load_schedule_directives(schedule_directives_file) + schedule_info = derive_schedule_info(schedule_directives, schedule_directives_file) + else: + raise ValueError("piper_setup requires schedule_directives_file") + + piper_metadata.visualize_dag = visualize_dag + piper_metadata.artifact_dir = _artifact_dir_for_schedule(schedule_directives_file) + piper_metadata.schedule_directives = list(schedule_directives or []) + piper_metadata.schedule_directives_file = schedule_directives_file + piper_metadata.schedule_info = dict(schedule_info) + + # Reset DAG/compile fields so stale data from a prior run never leaks into + # this run if the backend is somehow not re-invoked. + piper_metadata.training_dag = None + piper_metadata.per_pp_training_dags = None + piper_metadata.compiled_data_store = None + + pp_degree = int(schedule_info["pp_degree"]) + + # All dp_ranks must agree on a single master_addr: the IP of the actor with + # global_rank=0 (pp_rank=0, dp_rank=0). dp_rank=0 publishes it via a named + # Ray actor; other dp_ranks wait until it's available. + dp_rank = int(os.environ["PIPER_DP_RANK"]) + _create_actors( + pp_degree, optim_fn, + profile=nsight, + no_nvtx=no_nvtx, pg=pg, temp_dir=temp_dir, + use_inductor=use_inductor, + pp_outer=pp_outer, + ) + + if dp_rank == 0: + # Create the compiled-data store early so dp_rank>0 can poll for it. + piper_metadata.compiled_data_store = _CompiledDataStore.options( + name=_COMPILED_DATA_ACTOR, + lifetime="detached", + num_cpus=0, + get_if_exists=True, + ).remote() + ray.get(piper_metadata.compiled_data_store.is_ready.remote()) + + # Get IP and a free port from actor 0's node — it will be the TCPStore server. + master_addr, master_port = ray.get(piper_metadata.actors[0].get_node_ip_and_free_port.remote()) + addr_store = _Rank0AddrStore.options( + name=_RANK0_ADDR_ACTOR, + lifetime="detached", + num_cpus=0, + ).remote(master_addr, master_port) + ray.get(addr_store.get.remote()) + else: + master_addr = master_port = None + deadline = time.monotonic() + 600 + while master_addr is None: + try: + store = ray.get_actor(_RANK0_ADDR_ACTOR) + master_addr, master_port = ray.get(store.get.remote()) + except ValueError: + if time.monotonic() > deadline: + raise TimeoutError( + f"Timed out waiting for Ray actor {_RANK0_ADDR_ACTOR}" + ) + time.sleep(0.05) + except Exception: + logger.exception("Failed while waiting for Ray actor %s", _RANK0_ADDR_ACTOR) + raise + + logger.debug(f"DP rank {dp_rank}: Master address for process groups: {master_addr}:{master_port}") + + ray.get([actor._join_process_groups.remote(master_addr, master_port) for actor in piper_metadata.actors.values()]) + + if dp_rank == 0: + try: + ray.kill(ray.get_actor(_RANK0_ADDR_ACTOR)) + except ValueError: + logger.debug("Ray actor %s already absent during cleanup", _RANK0_ADDR_ACTOR) + except Exception: + logger.exception("Failed to clean up Ray actor %s", _RANK0_ADDR_ACTOR) + raise + + # Push non-trainable constant tensor attributes (e.g. freqs_cis, rope_cache, mask) + # to all actors so _load_stage can initialize them correctly instead of zero-filling. + # Must happen on every dp_rank since each dp_rank owns its own actors. + _const_attrs = { + k: v.detach().cpu() if isinstance(v, torch.Tensor) else v + for k, v in (const_attrs or {}).items() + } + if _const_attrs: + ray.get([ + actor.load_const_attrs.remote(_const_attrs) + for actor in piper_metadata.actors.values() + ]) + + if dp_rank == 0: + # --- dp_rank=0: run torch.compile, build DAGs, publish compiled data --- + + # Build the model directly on meta device. + resolved_model_args = _resolve_model_constructor_value(model_args, piper_metadata.schedule_info) + resolved_model_kwargs = _resolve_model_constructor_value(model_kwargs, piper_metadata.schedule_info) + with torch.device("meta"): + model = model_class(*tuple(resolved_model_args or ()), **dict(resolved_model_kwargs or {})) + if model_dtype is not None: + model = model.to(model_dtype) + + num_params = sum(p.numel() for p in model.parameters()) + if model_dtype == torch.bfloat16: + param_size_gb = num_params * 2 / (1024**3) + elif model_dtype == torch.float32: + param_size_gb = num_params * 4 / (1024**3) + else: + raise ValueError(f"Unsupported model dtype: {model_dtype}") + logger.info(f"Model size: {num_params/(1e6):.0f} M parameters ({param_size_gb:.2f} GB), dtype: {model_dtype}") + + compiled = torch.compile(model, backend=piper, fullgraph=True) + compile_inputs = [x.to(device="meta") for x in example_inputs] + _reset_annotation_state() + _ = compiled(*compile_inputs) + + per_pp_training_dags = getattr(piper_metadata, "per_pp_training_dags", None) + if not per_pp_training_dags: + raise RuntimeError( + "piper backend did not produce per_pp_training_dags; backend transition incomplete" + ) + if len(per_pp_training_dags) != pp_degree: + raise RuntimeError( + f"Expected {pp_degree} per-PP DAGs, got {len(per_pp_training_dags)}" + ) + per_pp_training_dags_rpc = _prepare_training_dags_for_rpc(per_pp_training_dags) + piper_metadata.per_pp_training_dags = per_pp_training_dags_rpc + + ray.get([ + piper_metadata.actors[pp_rank].load_training_dag.remote(pp_dag) + for pp_rank, pp_dag in enumerate(per_pp_training_dags_rpc) + ]) + ray.get( + piper_metadata.compiled_data_store.publish_training_dag_data.remote( + {"per_pp_training_dags": per_pp_training_dags_rpc} + ) + ) + + else: + logger.debug(f"DP rank {dp_rank} waiting for dp_rank=0 training-DAG data...") + training_dag_data = None + deadline = time.monotonic() + 600 + while training_dag_data is None: + try: + store = ray.get_actor(_COMPILED_DATA_ACTOR) + if ray.get(store.is_training_dag_ready.remote()): + training_dag_data = ray.get(store.get_training_dag_data.remote()) + except ValueError: + if time.monotonic() > deadline: + raise TimeoutError( + f"Timed out waiting for Ray actor {_COMPILED_DATA_ACTOR}" + ) + except Exception: + logger.exception("Failed while waiting for Ray actor %s", _COMPILED_DATA_ACTOR) + raise + if training_dag_data is None: + time.sleep(0.2) + + per_pp_training_dags = training_dag_data["per_pp_training_dags"] + if len(per_pp_training_dags) != pp_degree: + raise RuntimeError( + f"Expected {pp_degree} per-PP DAGs from dp_rank=0, got {len(per_pp_training_dags)}" + ) + piper_metadata.per_pp_training_dags = per_pp_training_dags + + ray.get([ + piper_metadata.actors[pp_rank].load_training_dag.remote(pp_dag) + for pp_rank, pp_dag in enumerate(per_pp_training_dags) + ]) + + loss_pp_ranks = _compute_loss_pp_ranks(per_pp_training_dags) + if not loss_pp_ranks: + loss_pp_ranks = [pp_degree - 1] + logger.warning( + "No compute_loss node found in per-PP DAGs; falling back to labels on pp_rank=%s", + loss_pp_ranks[0], + ) + ray.get(piper_metadata.actors[0].load_input.remote(example_inputs)) + ray.get([ + piper_metadata.actors[pp_rank].load_labels.remote(example_outputs) + for pp_rank in loss_pp_ranks + ]) + + logger.info(f"DP rank {dp_rank} done.") diff --git a/src/coordinator.py b/src/coordinator.py new file mode 100644 index 0000000..ec226ed --- /dev/null +++ b/src/coordinator.py @@ -0,0 +1,106 @@ +import ray +from typing import Callable +import os + +from ray.util.placement_group import placement_group +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from .state import create_logger, LOG_LEVEL +from .schedule import load_schedule_info + + +# Coordinator needs GPUs when using profiling to infer stage boundaries +# @ray.remote(num_gpus=0.1) + +# Use manual stage annotations- more stable +@ray.remote(num_gpus=0.1) +def run_dp_rank(dp_rank, dp_degree, pp_degree, world_size, training_func: Callable, *args, **kwargs): + logger = create_logger("coordinator", LOG_LEVEL) + logger.debug(f"Running DP rank {dp_rank+1} of {dp_degree}") + + os.environ["PIPER_DP_RANK"] = str(dp_rank) + os.environ["PIPER_DP_DEGREE"] = str(dp_degree) + os.environ["PIPER_PP_DEGREE"] = str(pp_degree) + os.environ["PIPER_WORLD_SIZE"] = str(world_size) + os.environ["TORCH_LOGS"] = "+graph_breaks" + return training_func(*args, **kwargs) + + +@ray.remote +class PiperProgramCoordinator: + """Central Actor that Coordinates all the DP replicas of a single pipeline""" + + def __init__( + self, + pp_outer: bool = False, + schedule_directives_file: str | None = None, + ): + if schedule_directives_file is None: + raise ValueError("PiperProgramCoordinator requires schedule_directives_file") + info = load_schedule_info(schedule_directives_file) + self.dp_degree = info["dp_degree"] + self.pp_degree = info["pp_degree"] + self.world_size = self.dp_degree * self.pp_degree + # pp_outer=True means one PP stage per node (placement bundles keyed by + # pp_rank). In that mode DP drivers are spread across the pp bundles. + self.pp_outer = pp_outer + + def run_program(self, training_func: Callable, pg, *args, **kwargs): + from .compile import _RANK0_ADDR_ACTOR, _COMPILED_DATA_ACTOR + logger = create_logger("coordinator", LOG_LEVEL) + try: + ray.kill(ray.get_actor(_RANK0_ADDR_ACTOR)) + except ValueError: + logger.debug("No stale Ray actor named %s to kill", _RANK0_ADDR_ACTOR) + except Exception: + logger.exception("Failed to kill stale Ray actor named %s", _RANK0_ADDR_ACTOR) + raise + # Kill any stale compiled-data store from a previous run so that + # dp_rank>0 workers cannot read outdated (e.g. wrong-model) stage data. + try: + ray.kill(ray.get_actor(_COMPILED_DATA_ACTOR)) + except ValueError: + logger.debug("No stale Ray actor named %s to kill", _COMPILED_DATA_ACTOR) + except Exception: + logger.exception("Failed to kill stale Ray actor named %s", _COMPILED_DATA_ACTOR) + raise + + return ray.get( + [ + run_dp_rank.options( + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=pg, + placement_group_bundle_index=( + dp_rank % self.pp_degree if self.pp_outer else dp_rank + ), + ) + ).remote( + dp_rank, + self.dp_degree, + self.pp_degree, + self.world_size, + training_func, + *args, + **kwargs, + ) + for dp_rank in range(self.dp_degree) + ] + ) + + +def create_piper_placement_group(schedule_directives_file: str, pp_outer: bool = False): + info = load_schedule_info(schedule_directives_file) + pp_degree = info["pp_degree"] + dp_degree = info["dp_degree"] + + if pp_outer: + drivers_per_bundle = (dp_degree + pp_degree - 1) // pp_degree + return placement_group( + [{"CPU": dp_degree + drivers_per_bundle, "GPU": dp_degree}] * pp_degree, + strategy="STRICT_SPREAD", + ) + + return placement_group( + [{"CPU": pp_degree, "GPU": pp_degree}] * dp_degree, + strategy="SPREAD", + ) diff --git a/src/dag.py b/src/dag.py new file mode 100644 index 0000000..37e88ac --- /dev/null +++ b/src/dag.py @@ -0,0 +1,340 @@ +from dataclasses import dataclass, field +from typing import Any + +import torch.fx as fx + +from .fx import AnnotationSegment + +def _collect_triton_constant_args(gm: fx.GraphModule) -> dict[int, Any]: + from torch._higher_order_ops.triton_kernel_wrap import kernel_side_table + out: dict[int, Any] = {} + for node in gm.graph.nodes: + if node.op != "call_function": + continue + target_str = getattr(node.target, "__name__", "") or str(node.target) + if "triton_kernel_wrapper" not in target_str: + continue + idx = node.kwargs.get("constant_args_idx") + if idx is not None and idx in kernel_side_table.constant_args: + out[int(idx)] = kernel_side_table.constant_args[idx] + return out + + +@dataclass(frozen=True) +class TrainingDAGEdge: + src_uid: str + dst_uid: str + dep_kind: str # "data" | "temporal" + tensor_name: str | None = None + + +@dataclass +class TrainingDAGNode: + uid: str + node_kind: str # "COMPUTE" | "SEND_COMM" | "RECV_COMM" | "REDUCE_COMM" | "ALL_GATHER_COMM" | "REDUCE_SCATTER_COMM" | "A2A_COMM" | "ORDER_DUMMY" + compute_subkind: str | None # "FWD" | "BWD" | "BWD_I" | "BWD_W" when node_kind == "COMPUTE" + tag: dict[str, int | None] + device: list[int] | None + stream: str + node_meta: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class TrainingDAG: + nodes: dict[str, TrainingDAGNode] = field(default_factory=dict) + edges: list[TrainingDAGEdge] = field(default_factory=list) + succs: dict[str, set[str]] = field(default_factory=dict) + preds: dict[str, set[str]] = field(default_factory=dict) + edge_keys: set[tuple[str, str, str, str | None]] = field(default_factory=set) + + def add_node(self, node: TrainingDAGNode) -> None: + if node.uid in self.nodes: + raise ValueError(f"Duplicate node uid: {node.uid}") + self.nodes[node.uid] = node + self.succs[node.uid] = set() + self.preds[node.uid] = set() + + def add_edge(self, edge: TrainingDAGEdge) -> None: + if edge.src_uid not in self.nodes or edge.dst_uid not in self.nodes: + raise ValueError(f"Edge references unknown node: {edge}") + edge_key = (edge.src_uid, edge.dst_uid, edge.dep_kind, edge.tensor_name) + if edge_key in self.edge_keys: + return + self.edge_keys.add(edge_key) + self.edges.append(edge) + self.succs[edge.src_uid].add(edge.dst_uid) + self.preds[edge.dst_uid].add(edge.src_uid) + + + +def _flatten_output_nodes(out_arg: Any) -> list[fx.Node]: + out_nodes: list[fx.Node] = [] + + def _walk(v: Any) -> None: + if isinstance(v, fx.Node): + out_nodes.append(v) + elif isinstance(v, (tuple, list)): + for x in v: + _walk(x) + elif isinstance(v, dict): + for x in v.values(): + _walk(x) + + _walk(out_arg) + return out_nodes + + + +def _node_io_names(seg_gm: fx.GraphModule) -> tuple[list[str], list[str]]: + inputs = [n.name for n in seg_gm.graph.nodes if n.op == "placeholder"] + out_node = next((n for n in seg_gm.graph.nodes if n.op == "output"), None) + if out_node is None: + return inputs, [] + outputs = [n.name for n in _flatten_output_nodes(out_node.args[0])] + return inputs, outputs + +def _with_pass_tag(tag: dict[str, Any], pass_value: str | None) -> dict[str, Any]: + t = dict(tag) + t["PASS"] = pass_value + return t + +def build_training_dag( + annotation_segments: list[AnnotationSegment], +) -> TrainingDAG: + """Build a TrainingDAG where each node is one annotated GraphModule segment.""" + dag = TrainingDAG() + + global_prev_uid: str | None = None + + for segment in annotation_segments: + uid = f"s{segment.stage_id}.seg{segment.segment_id}" + tag = _with_pass_tag(segment.tag, "F") + in_names, out_names = _node_io_names(segment.gm) + node = TrainingDAGNode( + uid=uid, + node_kind="COMPUTE", + compute_subkind="FWD", + tag=tag, + device=None, + stream="default_stream", + node_meta={ + "bucket_key": uid, + "stage_id": segment.stage_id, + "segment_id": segment.segment_id, + "gm": segment.gm, + "input_idxs": list(segment.input_idxs), + "param_idxs": list(segment.param_idxs), + "graphargs": list(segment.graphargs), + "input_names": in_names, + "output_names": out_names, + "a2a_boundary_after": segment.a2a_boundary_after, + "triton_constant_args": _collect_triton_constant_args(segment.gm), + }, + ) + dag.add_node(node) + + # Explicit activation-flow ordering in annotation segment order. + if global_prev_uid is not None: + dag.add_edge(TrainingDAGEdge(src_uid=global_prev_uid, dst_uid=uid, dep_kind="data")) + global_prev_uid = uid + + # Build backward-pass nodes and dependencies by reversing forward data edges. + fwd_compute_uids = [ + uid for uid, node in dag.nodes.items() + if node.node_kind == "COMPUTE" and node.compute_subkind == "FWD" + ] + fwd_uid_to_bwd_uid: dict[str, str] = {} + for fwd_uid in fwd_compute_uids: + fwd_node = dag.nodes[fwd_uid] + bwd_uid = f"{fwd_uid}.bwd" + fwd_uid_to_bwd_uid[fwd_uid] = bwd_uid + dag.add_node( + TrainingDAGNode( + uid=bwd_uid, + node_kind="COMPUTE", + compute_subkind="BWD", + tag=_with_pass_tag(fwd_node.tag, "B"), + device=fwd_node.device, + stream="default_stream", + node_meta={ + "fwd_uid": fwd_uid, + "bucket_key": fwd_node.node_meta.get("bucket_key", fwd_uid), + "compute_loss": False, + }, + ) + ) + + fwd_edges = [ + e for e in dag.edges + if e.dep_kind == "data" + and dag.nodes[e.src_uid].node_kind == "COMPUTE" + and dag.nodes[e.dst_uid].node_kind == "COMPUTE" + and dag.nodes[e.src_uid].compute_subkind == "FWD" + and dag.nodes[e.dst_uid].compute_subkind == "FWD" + ] + + # Reverse every forward dependency for the backward pass: + # FWD: u -> v ==> BWD: v' -> u' + for e in fwd_edges: + bwd_src = fwd_uid_to_bwd_uid[e.dst_uid] + bwd_dst = fwd_uid_to_bwd_uid[e.src_uid] + dag.add_edge( + TrainingDAGEdge( + src_uid=bwd_src, + dst_uid=bwd_dst, + dep_kind="data", + tensor_name=e.tensor_name, + ) + ) + + # Bridge from forward phase into backward phase: + # last_fwd -> first_bwd (where first_bwd is the BWD node paired with last_fwd). + if fwd_compute_uids: + def _fwd_sort_key(uid: str) -> tuple[int, int]: + n = dag.nodes[uid] + return ( + int(n.node_meta.get("stage_id", -1)), + int(n.node_meta.get("segment_id", -1)), + ) + + last_fwd_uid = max(fwd_compute_uids, key=_fwd_sort_key) + first_bwd_uid = fwd_uid_to_bwd_uid[last_fwd_uid] + dag.add_edge( + TrainingDAGEdge( + src_uid=last_fwd_uid, + dst_uid=first_bwd_uid, + dep_kind="data", + tensor_name=None, + ) + ) + + # Mark loss computation point: first BWD node (earliest in BWD execution). + if first_bwd_uid in dag.nodes: + dag.nodes[first_bwd_uid].node_meta["compute_loss"] = True + + # Add explicit UPD node with temporal dependency from every BWD compute node. + # UPD participates in the compute temporal chain and is not scheduled from + # data edges. + upd_uid = "upd.0" + dag.add_node( + TrainingDAGNode( + uid=upd_uid, + node_kind="UPD", + compute_subkind=None, + tag={"PASS": None}, + device=None, + stream="default_stream", + node_meta={}, + ) + ) + for uid, node in list(dag.nodes.items()): + if node.node_kind == "COMPUTE" and node.compute_subkind == "BWD": + dag.add_edge( + TrainingDAGEdge( + src_uid=uid, + dst_uid=upd_uid, + dep_kind="temporal", + tensor_name=None, + ) + ) + + return dag + +def _iter_data_edges(dag: TrainingDAG) -> list[TrainingDAGEdge]: + return [e for e in dag.edges if e.dep_kind == "data"] + + +def _remove_edge(dag: TrainingDAG, edge: TrainingDAGEdge) -> None: + edge_key = (edge.src_uid, edge.dst_uid, edge.dep_kind, edge.tensor_name) + if edge_key not in dag.edge_keys: + return + dag.edge_keys.remove(edge_key) + dag.edges = [ + e for e in dag.edges + if not ( + e.src_uid == edge.src_uid + and e.dst_uid == edge.dst_uid + and e.dep_kind == edge.dep_kind + and e.tensor_name == edge.tensor_name + ) + ] + if edge.src_uid in dag.succs: + dag.succs[edge.src_uid].discard(edge.dst_uid) + if edge.dst_uid in dag.preds: + dag.preds[edge.dst_uid].discard(edge.src_uid) + +def _remove_node_and_incident_edges(dag: TrainingDAG, uid: str) -> None: + incident = [e for e in dag.edges if e.src_uid == uid or e.dst_uid == uid] + for e in incident: + _remove_edge(dag, e) + dag.nodes.pop(uid, None) + dag.succs.pop(uid, None) + dag.preds.pop(uid, None) + for s in dag.succs.values(): + s.discard(uid) + for p in dag.preds.values(): + p.discard(uid) + +def _topological_order(dag: TrainingDAG) -> list[str]: + in_deg: dict[str, int] = {uid: 0 for uid in dag.nodes} + out_edges_by_src: dict[str, list[TrainingDAGEdge]] = {uid: [] for uid in dag.nodes} + for e in dag.edges: + if e.dst_uid in in_deg: + in_deg[e.dst_uid] += 1 + if e.src_uid in out_edges_by_src: + out_edges_by_src[e.src_uid].append(e) + q = sorted(uid for uid, d in in_deg.items() if d == 0) + out: list[str] = [] + while q: + cur_level = q + q = [] + for u in cur_level: + out.append(u) + # Decrement per outgoing edge (not per unique successor), because we + # can have parallel edges between the same pair (e.g., data + temporal). + for e in out_edges_by_src.get(u, []): + v = e.dst_uid + if v not in in_deg: + continue + in_deg[v] -= 1 + if in_deg[v] == 0: + q.append(v) + q.sort() + if len(out) != len(dag.nodes): + remaining = [uid for uid, d in in_deg.items() if d > 0] + raise ValueError( + "TrainingDAG topological sort failed: unresolved incoming edges remain; " + f"possible real cycle or malformed parallel-edge bookkeeping. Remaining nodes: {remaining[:8]}" + ) + return out + + +def _topological_levels(dag: TrainingDAG) -> dict[str, int]: + """Return each node's topological level. + + Source nodes are level 0. Every other node is one level after its latest + predecessor, so independent successors can share the same level. + """ + topo = _topological_order(dag) + levels: dict[str, int] = {} + for uid in topo: + pred_levels = [ + levels[pred_uid] + for pred_uid in dag.preds.get(uid, set()) + if pred_uid in levels + ] + levels[uid] = (max(pred_levels) + 1) if pred_levels else 0 + return levels + +def _has_path(dag: TrainingDAG, src_uid: str, dst_uid: str) -> bool: + seen: set[str] = set() + stack = [src_uid] + while stack: + uid = stack.pop() + if uid == dst_uid: + return True + if uid in seen: + continue + seen.add(uid) + stack.extend(dag.succs.get(uid, set())) + return False diff --git a/src/directives.py b/src/directives.py new file mode 100644 index 0000000..b6cf5c2 --- /dev/null +++ b/src/directives.py @@ -0,0 +1,1669 @@ +from collections import Counter +from dataclasses import dataclass +from typing import Any + +import torch.fx as fx + +from .bucket import bucket_stage +from .dag import ( + TrainingDAG, + TrainingDAGEdge, + TrainingDAGNode, + _has_path, + _iter_data_edges, + _remove_edge, + _remove_node_and_incident_edges, + _topological_order, + _with_pass_tag, +) +from .state import LOG_LEVEL, create_logger + +logger = create_logger("directives", LOG_LEVEL) +_ANY_TAG_INDEX = "__ANY_TAG_INDEX__" +_FORWARD_PASS = "F" +_FUSED_BWD_PASS = "B" +_BWD_INPUT_PASS = "BI" +_BWD_WEIGHT_PASS = "BW" +_DEFAULT_STREAM = "default_stream" +_VALID_ORDER_PASSES = frozenset({ + _FORWARD_PASS, + _FUSED_BWD_PASS, + _BWD_INPUT_PASS, + _BWD_WEIGHT_PASS, +}) + +def _is_backward_activation_subkind(subkind: str | None) -> bool: + return subkind in ("BWD", "BWD_I") + + +def _is_backward_weight_subkind(subkind: str | None) -> bool: + return subkind in ("BWD", "BWD_W") + + +def _is_backward_compute_subkind(subkind: str | None) -> bool: + return subkind in ("BWD", "BWD_I", "BWD_W") + +def _match_filter(tag: dict[str, int | None], flt: dict[str, Any]) -> bool: + for k, v in flt.items(): + if k not in tag: + return False + if v == _ANY_TAG_INDEX: + # "*" wildcard matches any concrete index for this tag, but not None. + if tag[k] is None: + return False + elif tag[k] != v: + return False + return True + + +def _apply_split_backward_stencil( + dag: TrainingDAG, + split_base_keys: set[tuple[tuple[str, Any], ...]], +) -> None: + if not split_base_keys: + return + split_filters = [dict(key) for key in split_base_keys] + for node in list(dag.nodes.values()): + if ( + node.node_kind != "COMPUTE" + or node.compute_subkind != "BWD" + or not any(_match_filter(node.tag, flt) for flt in split_filters) + ): + continue + _split_fused_bwd_node(dag, node.uid) + + +def _split_fused_bwd_node(dag: TrainingDAG, bwd_uid: str) -> None: + bwd_i = dag.nodes[bwd_uid] + if bwd_i.node_kind != "COMPUTE" or bwd_i.compute_subkind != "BWD": + raise ValueError(f"expected fused BWD compute node, got {bwd_uid}: {bwd_i}") + bwd_w_uid = f"{bwd_uid}.bw" + if bwd_w_uid in dag.nodes: + raise ValueError(f"duplicate BWD_W uid while splitting {bwd_uid}: {bwd_w_uid}") + + bwd_i.compute_subkind = "BWD_I" + bwd_i.tag = _with_pass_tag(bwd_i.tag, _BWD_INPUT_PASS) + + bwd_w_meta = dict(bwd_i.node_meta) + bwd_w_meta["compute_loss"] = False + dag.add_node( + TrainingDAGNode( + uid=bwd_w_uid, + node_kind="COMPUTE", + compute_subkind="BWD_W", + tag=_with_pass_tag(bwd_i.tag, _BWD_WEIGHT_PASS), + device=list(bwd_i.device) if bwd_i.device is not None else None, + stream=bwd_i.stream, + node_meta=bwd_w_meta, + ) + ) + + for edge in [ + e for e in list(dag.edges) + if e.dep_kind == "temporal" + and e.src_uid == bwd_uid + and dag.nodes[e.dst_uid].node_kind == "UPD" + ]: + _remove_edge(dag, edge) + dag.add_edge( + TrainingDAGEdge( + src_uid=bwd_w_uid, + dst_uid=edge.dst_uid, + dep_kind="temporal", + tensor_name=edge.tensor_name, + ) + ) + + dag.add_edge( + TrainingDAGEdge( + src_uid=bwd_uid, + dst_uid=bwd_w_uid, + dep_kind="data", + tensor_name=None, + ) + ) + + +def _normalize_filter_devices_directive( + directive: Any +) -> tuple[str, list[dict[str, Any]], list[int], str | None, str | None, str | None, bool, bool, int | None]: + if isinstance(directive, dict): + op = directive.get("op") + if op not in ("place", "replicate", "shard"): + raise ValueError(f"Unsupported directive op: {op}") + if "filter" not in directive: + raise ValueError(f"{op} directive requires current API field 'filter': {directive}") + if "filters" in directive: + raise ValueError(f"{op} directive does not accept 'filters': {directive}") + filter_spec = directive.get("filter") + devices = directive.get("devices", directive.get("device")) + stream = directive.get("stream") + gather_stream = directive.get("gather_stream") + reduce_stream = directive.get("reduce_stream") + if op == "replicate": + if stream is not None: + raise ValueError( + f"replicate directive does not accept 'stream'; use 'gather_stream' and/or 'reduce_stream': {directive}" + ) + else: + if gather_stream is not None or reduce_stream is not None: + raise ValueError( + f"{op} directive does not accept 'gather_stream'/'reduce_stream'; use 'stream': {directive}" + ) + shard_params = bool(directive.get("shard_params", False)) + shard_grads = bool(directive.get("shard_grads", False)) + bucket_size_raw = directive.get("bucket_size", None) + bucket_size = None if bucket_size_raw is None else int(bucket_size_raw) + if not isinstance(devices, list) or not devices: + raise ValueError(f"place directive requires non-empty devices list: {directive}") + n = _normalize_filter_spec(filter_spec, directive) + return ( + op, + [n], + [int(d) for d in devices], + (None if stream is None else str(stream)), + (None if gather_stream is None else str(gather_stream)), + (None if reduce_stream is None else str(reduce_stream)), + shard_params, + shard_grads, + bucket_size, + ) + + raise ValueError(f"Schedule directives must be JSON objects, got: {type(directive)}") + + +def _normalize_filter_spec(filter_spec: Any, directive: Any) -> dict[str, Any]: + """Normalize a single filter spec into {tag_name: value}.""" + def _norm_value(k: str, v: Any) -> Any: + if k == "PASS": + if not isinstance(v, str): + raise ValueError( + f"'PASS' filter value must be a string in {{'F','B','BI','BW'}}: {directive}" + ) + if v not in ("F", "B", "BI", "BW"): + raise ValueError( + f"Invalid 'PASS' filter value '{v}'. Allowed: 'F', 'B', 'BI', 'BW': {directive}" + ) + return v + if v == "*": + return _ANY_TAG_INDEX + if v is None: + return None + if isinstance(v, int): + return int(v) + if isinstance(v, str): + # Keep pass-like symbolic tags (e.g., "FWD"/"BWD") as strings. + try: + return int(v) + except ValueError: + return v + return v + + if not isinstance(filter_spec, dict): + raise ValueError(f"filter must be a JSON object: {directive}") + out: dict[str, Any] = {} + for k, v in filter_spec.items(): + if not isinstance(k, str): + raise ValueError(f"Filter keys must be strings: {directive}") + out[k] = _norm_value(k, v) + return out + + +def _iter_filter_specs(spec: Any): + if isinstance(spec, dict): + yield spec + return + if isinstance(spec, list): + if all( + isinstance(item, (list, tuple)) + and len(item) == 2 + and isinstance(item[0], str) + for item in spec + ): + yield spec + return + for item in spec: + yield from _iter_filter_specs(item) + + +def _schedule_filter_tag_names(directives: Any) -> set[str]: + if not isinstance(directives, list): + return set() + + runtime_tags = {"PASS"} + for directive in directives: + if isinstance(directive, dict) and directive.get("op") == "split": + dim_name = directive.get("dim_name") + if isinstance(dim_name, str) and dim_name: + runtime_tags.add(dim_name) + + tag_names: set[str] = set() + for directive in directives: + if not isinstance(directive, dict): + continue + specs = [] + if "filter" in directive: + specs.append(directive["filter"]) + if isinstance(directive.get("filters"), list): + specs.append(directive["filters"]) + for spec in _iter_filter_specs(specs): + if isinstance(spec, dict): + tag_names.update(k for k in spec if k not in runtime_tags) + elif isinstance(spec, list): + for item in spec: + if ( + isinstance(item, (list, tuple)) + and len(item) == 2 + and isinstance(item[0], str) + and item[0] not in runtime_tags + ): + tag_names.add(item[0]) + return tag_names + + +def _validate_schedule_tags_exist( + training_dag: TrainingDAG, + directives: Any, +) -> None: + schedule_tags = _schedule_filter_tag_names(directives) + if not schedule_tags: + return + model_tags = { + tag_name + for node in training_dag.nodes.values() + for tag_name in node.tag + if tag_name != "PASS" + } + missing = sorted(schedule_tags - model_tags) + if missing: + raise ValueError( + "Schedule references tag name(s) not found in model annotations: " + f"{missing}. Model tag names: {sorted(model_tags)}." + ) + + +def _normalize_split_directive(directive: Any) -> tuple[dict[str, Any], str, int]: + if not isinstance(directive, dict): + raise ValueError(f"split directive must be a dict: {directive}") + if directive.get("op") != "split": + raise ValueError(f"Unsupported split directive op: {directive}") + if "filter" not in directive: + raise ValueError(f"split directive requires 'filter': {directive}") + dim_name = directive.get("dim_name") + if not isinstance(dim_name, str) or not dim_name: + raise ValueError(f"split directive requires non-empty string dim_name: {directive}") + num_microbatches = int(directive.get("num_microbatches", 0)) + if num_microbatches <= 0: + raise ValueError(f"split directive requires num_microbatches > 0: {directive}") + flt = _normalize_filter_spec(directive["filter"], directive) + return flt, dim_name, num_microbatches + + +def _parse_order_directive(directive: Any) -> list[list[dict[str, Any]]]: + if not isinstance(directive, dict): + raise ValueError(f"order directive must be a dict: {directive}") + if directive.get("op") != "order": + raise ValueError(f"Unsupported order directive op: {directive}") + raw_filters = directive.get("filters") + if not isinstance(raw_filters, list) or len(raw_filters) < 2: + raise ValueError( + f"order directive requires filters list with at least 2 groups: {directive}" + ) + + filter_groups: list[list[dict[str, Any]]] = [] + for group_idx, raw_group in enumerate(raw_filters): + if not isinstance(raw_group, list) or not raw_group: + raise ValueError( + f"order directive filter group[{group_idx}] must be a non-empty list " + f"of filters: {directive}" + ) + + group = [] + for raw_filter in raw_group: + if not isinstance(raw_filter, dict): + raise ValueError( + f"order directive filter group[{group_idx}] contains invalid " + f"filter object: {raw_filter}" + ) + flt = _normalize_filter_spec(raw_filter, directive) + pass_value = flt.get("PASS") + if pass_value is not None and pass_value not in _VALID_ORDER_PASSES: + raise ValueError( + f"order directive has unsupported pass={pass_value!r}; " + f"expected one of {sorted(_VALID_ORDER_PASSES)}" + ) + group.append(flt) + filter_groups.append(group) + + return filter_groups + + +def _filter_key_without(flt: dict[str, Any], ignored_keys: set[str]) -> tuple[tuple[str, Any], ...]: + return tuple(sorted((k, v) for k, v in flt.items() if k not in ignored_keys)) + + +def _validate_split_backward_order_stencil( + directives: list[Any], +) -> set[tuple[tuple[str, Any], ...]]: + """Validate BI/BW order entries and return base filter keys to split. + + This is intentionally narrow: it validates only the current split-backward + stencil until the schedule JSON grows a full schema validation pass. + """ + split_keys: set[tuple[tuple[str, Any], ...]] = set() + split_entries: dict[tuple[tuple[str, Any], ...], list[tuple[int, tuple[int, int], str]]] = {} + + for directive_idx, raw in enumerate(directives): + if not isinstance(raw, dict) or raw.get("op") != "order": + continue + filter_groups = _parse_order_directive(raw) + seen_in_row: dict[tuple[tuple[str, Any], ...], dict[str, tuple[int, int]]] = {} + for group_idx, group in enumerate(filter_groups): + for filter_idx, flt in enumerate(group): + filter_pos = (group_idx, filter_idx) + pass_value = flt.get("PASS") + if pass_value is None: + continue + if pass_value == _FUSED_BWD_PASS: + continue + if pass_value not in (_BWD_INPUT_PASS, _BWD_WEIGHT_PASS): + continue + key = _filter_key_without(flt, {"PASS"}) + by_pass = seen_in_row.setdefault(key, {}) + if pass_value in by_pass: + raise ValueError( + f"order directive[{directive_idx}] has duplicate split-backward " + f"{pass_value} entry for {dict(key)} at indices " + f"{by_pass[pass_value]} and {filter_pos}" + ) + by_pass[pass_value] = filter_pos + split_keys.add(key) + split_entries.setdefault(key, []).append((directive_idx, filter_pos, pass_value)) + + for key, entries in sorted(split_entries.items()): + bi = [entry for entry in entries if entry[2] == _BWD_INPUT_PASS] + bw = [entry for entry in entries if entry[2] == _BWD_WEIGHT_PASS] + if len(bi) != 1 or len(bw) != 1: + raise ValueError( + f"split backward order entries for {dict(key)} must contain exactly " + f"one BI and one BW; got BI={len(bi)} BW={len(bw)}" + ) + if bi[0][0] != bw[0][0]: + raise ValueError( + f"split backward order entries for {dict(key)} must appear in the " + "same order directive row" + ) + if bi[0][1] > bw[0][1]: + raise ValueError( + f"split backward order entries for {dict(key)} must place BI before BW" + ) + + return split_keys + + + +def _device_to_physical_pp_rank(dag: TrainingDAG) -> dict[tuple[int, ...], int]: + device_keys = sorted({ + tuple(sorted(node.device)) + for node in dag.nodes.values() + if node.device is not None + }) + return {key: idx for idx, key in enumerate(device_keys)} + + +def _insert_send_recv_comm_nodes(dag: TrainingDAG, comm_stream: str | None = None) -> None: + data_edges = _iter_data_edges(dag) + device_to_pp_rank = _device_to_physical_pp_rank(dag) + send_idx = sum(1 for n in dag.nodes.values() if n.node_kind == "SEND_COMM") + recv_idx = sum(1 for n in dag.nodes.values() if n.node_kind == "RECV_COMM") + for edge in data_edges: + src = dag.nodes[edge.src_uid] + dst = dag.nodes[edge.dst_uid] + if src.node_kind != "COMPUTE" or dst.node_kind != "COMPUTE": + continue + if src.device is None or dst.device is None: + raise ValueError(f"cannot insert send/recv for unplaced edge {src.uid} -> {dst.uid}") + if src.device == dst.device: + continue + src_device_key = tuple(sorted(src.device)) + dst_device_key = tuple(sorted(dst.device)) + send_uid = f"send.{send_idx}" + send_idx += 1 + recv_uid = f"recv.{recv_idx}" + recv_idx += 1 + stream_name = comm_stream if comm_stream is not None else "pp_stream" + send_node = TrainingDAGNode( + uid=send_uid, + node_kind="SEND_COMM", + compute_subkind=None, + tag=dict(src.tag), + device=(None if src.device is None else list(src.device)), + stream=stream_name, + node_meta={ + "from_uid": src.uid, + "to_uid": dst.uid, + "peer_pp_rank": device_to_pp_rank[dst_device_key], + "bucket_key": src.node_meta.get("bucket_key"), + }, + ) + recv_node = TrainingDAGNode( + uid=recv_uid, + node_kind="RECV_COMM", + compute_subkind=None, + tag=dict(dst.tag), + device=(None if dst.device is None else list(dst.device)), + stream=stream_name, + node_meta={ + "from_uid": src.uid, + "to_uid": dst.uid, + "peer_pp_rank": device_to_pp_rank[src_device_key], + "bucket_key": dst.node_meta.get("bucket_key"), + }, + ) + dag.add_node(send_node) + dag.add_node(recv_node) + _remove_edge(dag, edge) + dag.add_edge(TrainingDAGEdge(src_uid=src.uid, dst_uid=send_uid, dep_kind="data", tensor_name=edge.tensor_name)) + dag.add_edge(TrainingDAGEdge(src_uid=recv_uid, dst_uid=dst.uid, dep_kind="data", tensor_name=edge.tensor_name)) + + +def _replicate_update_nodes_by_device(dag: TrainingDAG) -> None: + """Replicate UPD nodes per upstream device set so split components stay homogeneous.""" + upd_nodes = [n for n in list(dag.nodes.values()) if n.node_kind == "UPD" and ".rep" not in n.uid] + for upd in upd_nodes: + in_edges = [e for e in list(dag.edges) if e.dep_kind == "temporal" and e.dst_uid == upd.uid] + by_dev: dict[tuple[int, ...], list[TrainingDAGEdge]] = {} + for e in in_edges: + src = dag.nodes[e.src_uid] + if src.device is None: + continue + key = tuple(sorted(src.device)) + by_dev.setdefault(key, []).append(e) + if len(by_dev) <= 1: + if len(by_dev) == 1: + upd.device = list(next(iter(by_dev.keys()))) + continue + created_uids: list[str] = [] + for i_upd, (dev_key, dev_edges) in enumerate(sorted(by_dev.items(), key=lambda kv: kv[0])): + nu = upd.uid if i_upd == 0 else f"{upd.uid}.rep{i_upd}" + if i_upd == 0: + upd.device = list(dev_key) + else: + dag.add_node( + TrainingDAGNode( + uid=nu, + node_kind="UPD", + compute_subkind=None, + tag=dict(upd.tag), + device=list(dev_key), + stream=upd.stream, + node_meta=dict(upd.node_meta), + ) + ) + created_uids.append(nu) + for e in dev_edges: + _remove_edge(dag, e) + dag.add_edge( + TrainingDAGEdge( + src_uid=e.src_uid, + dst_uid=nu, + dep_kind=e.dep_kind, + tensor_name=e.tensor_name, + ) + ) + # Duplicate outgoing edges of the original UPD to all replicas. + out_edges = [e for e in list(dag.edges) if e.dep_kind == "temporal" and e.src_uid == upd.uid] + if out_edges: + for e in out_edges: + for nu in created_uids[1:]: + dag.add_edge( + TrainingDAGEdge( + src_uid=nu, + dst_uid=e.dst_uid, + dep_kind=e.dep_kind, + tensor_name=e.tensor_name, + ) + ) + + + +def _rewire_bwd_successors_through_sync( + dag: TrainingDAG, + bwd_uid: str, + sync_uid: str, +) -> None: + """Add update ordering through a gradient sync comm node. + + BWD data successors carry activation gradients, including P2P sends to an + upstream pipeline stage. Gradient synchronization is for parameter grads, + so it must not be inserted into those data paths. UPD dependencies are + temporal; add sync -> UPD temporal edges so the + update waits for gradient synchronization as well as the BWD compute. + """ + upd_temporal_outs = [ + e for e in list(dag.edges) + if e.dep_kind == "temporal" + and e.src_uid == bwd_uid + and dag.nodes[e.dst_uid].node_kind == "UPD" + ] + for e in upd_temporal_outs: + dag.add_edge( + TrainingDAGEdge( + src_uid=sync_uid, + dst_uid=e.dst_uid, + dep_kind="temporal", + tensor_name=None, + ) + ) + + +def _bucket_matched_fwd_nodes( + dag: TrainingDAG, + filters: list[dict[str, Any]], + bucket_size_mb: int, +) -> None: + if bucket_size_mb <= 0: + raise ValueError(f"bucket_size must be > 0 MB, got {bucket_size_mb}") + bucket_size_bytes = int(bucket_size_mb) * 1024 * 1024 + + fwd_targets = [ + uid for uid, node in dag.nodes.items() + if node.node_kind == "COMPUTE" + and node.compute_subkind == "FWD" + and any(_match_filter(node.tag, flt) for flt in filters) + ] + + # Deterministic rewrite order. + fwd_targets.sort(key=lambda uid: ( + int(dag.nodes[uid].node_meta.get("stage_id", 10**9)), + int(dag.nodes[uid].node_meta.get("segment_id", 10**9)), + int(dag.nodes[uid].tag.get("MB", 0) or 0), + )) + + for fwd_uid in fwd_targets: + if fwd_uid not in dag.nodes: + continue + fwd_node = dag.nodes[fwd_uid] + bwd_candidates = sorted( + uid for uid, node in dag.nodes.items() + if node.node_kind == "COMPUTE" + and node.compute_subkind == "BWD" + and node.node_meta.get("fwd_uid") == fwd_uid + ) + if len(bwd_candidates) != 1: + raise ValueError( + f"Expected exactly one fused BWD node for {fwd_uid}, " + f"found {bwd_candidates}" + ) + bwd_uid = bwd_candidates[0] + bwd_node = dag.nodes[bwd_uid] + + gm = fwd_node.node_meta.get("gm") + graphargs = fwd_node.node_meta.get("graphargs") + input_idxs = fwd_node.node_meta.get("input_idxs") + param_idxs = fwd_node.node_meta.get("param_idxs") + if gm is None or graphargs is None or input_idxs is None or param_idxs is None: + continue + + def _param_bytes(gargs: list, pidxs: list[int]) -> int: + total = 0 + for i in pidxs: + if i < 0 or i >= len(gargs): + continue + t = gargs[i] + if t is not None and hasattr(t, "numel") and hasattr(t, "element_size"): + total += int(t.numel()) * int(t.element_size()) + return total + + def _placeholder_names(module: fx.GraphModule) -> list[str]: + return [n.name for n in module.graph.nodes if n.op == "placeholder"] + + def _param_names(module: fx.GraphModule, pidxs: list[int]) -> list[str]: + names = _placeholder_names(module) + return [ + names[i] if 0 <= i < len(names) else f"" + for i in pidxs + ] + + pre_bytes = _param_bytes(graphargs, param_idxs) + logger.info( + "bucket_stage pre-size node=%s stage=%s seg=%s bucket_size_value=%d " + "bucket_size_bytes_passed=%d bucket_size_mib_passed=%.6f param_bytes=%d", + fwd_uid, + fwd_node.node_meta.get("stage_id"), + fwd_node.node_meta.get("segment_id"), + int(bucket_size_mb), + bucket_size_bytes, + bucket_size_bytes / (1024 * 1024), + pre_bytes, + ) + + try: + buckets = bucket_stage( + gm, + graphargs, + input_idxs, + param_idxs, + bucket_size_bytes=bucket_size_bytes, + debug_name=fwd_uid, + ) + except Exception as exc: + logger.warning( + "replicate bucketing skipped for node %s (stage=%s seg=%s): %s", + fwd_uid, + fwd_node.node_meta.get("stage_id"), + fwd_node.node_meta.get("segment_id"), + exc, + ) + continue + if len(buckets) <= 1: + continue + original_param_names = set(_param_names(gm, param_idxs)) + lowered_param_counts: Counter[str] = Counter() + for bi, (_bgm, _bin, b_param, b_args) in enumerate(buckets): + post_bytes = _param_bytes(b_args, b_param) + bucket_param_names = _param_names(_bgm, b_param) + lowered_param_counts.update(bucket_param_names) + logger.info( + "bucket_stage bucket-size node=%s bucket=%d/%d param_bytes=%d params=%d names=%s", + fwd_uid, + bi, + len(buckets), + post_bytes, + len(bucket_param_names), + bucket_param_names, + ) + missing_param_names = sorted(original_param_names - set(lowered_param_counts)) + unknown_param_names = sorted(set(lowered_param_counts) - original_param_names) + duplicate_param_names = sorted( + name for name, count in lowered_param_counts.items() if count != 1 + ) + logger.info( + "bucket_stage lowered-ownership node=%s original_params=%d lowered_param_refs=%d " + "unique_lowered_params=%d missing_params=%d duplicate_params=%d unknown_params=%d", + fwd_uid, + len(original_param_names), + sum(lowered_param_counts.values()), + len(lowered_param_counts), + len(missing_param_names), + len(duplicate_param_names), + len(unknown_param_names), + ) + if missing_param_names or duplicate_param_names or unknown_param_names: + logger.warning( + "bucket_stage lowered-ownership-invalid node=%s missing=%s duplicates=%s unknown=%s", + fwd_uid, + missing_param_names[:16], + duplicate_param_names[:16], + unknown_param_names[:16], + ) + + # Snapshot old incident edges before rewrite. + in_fwd = [e for e in dag.edges if e.dst_uid == fwd_uid] + out_fwd = [e for e in dag.edges if e.src_uid == fwd_uid] + in_bwd = [e for e in dag.edges if e.dst_uid == bwd_uid] + out_bwd = [e for e in dag.edges if e.src_uid == bwd_uid] + + orig_stage = int(fwd_node.node_meta.get("stage_id", -1)) + orig_seg = int(fwd_node.node_meta.get("segment_id", -1)) + + fwd_bucket_uids: list[str] = [] + bwd_bucket_uids: list[str] = [] + for bi, (bgm, b_in, b_param, b_args) in enumerate(buckets): + f_uid = f"{fwd_uid}.bucket{bi}" + b_uid = f"{f_uid}.bwd" + fwd_bucket_uids.append(f_uid) + bwd_bucket_uids.append(b_uid) + dag.add_node( + TrainingDAGNode( + uid=f_uid, + node_kind="COMPUTE", + compute_subkind="FWD", + tag=_with_pass_tag(fwd_node.tag, "F"), + device=list(fwd_node.device) if fwd_node.device is not None else None, + stream=fwd_node.stream, + node_meta={ + "stage_id": orig_stage, + "segment_id": orig_seg * 1000 + bi, + "gm": bgm, + "input_idxs": list(b_in), + "param_idxs": list(b_param), + "graphargs": list(b_args), + "input_names": [n.name for n in bgm.graph.nodes if n.op == "placeholder"], + "output_names": [], + "a2a_boundary_after": None, + "bucket_idx": bi, + "num_buckets": len(buckets), + "bucket_key": f_uid, + }, + ) + ) + dag.add_node( + TrainingDAGNode( + uid=b_uid, + node_kind="COMPUTE", + compute_subkind="BWD", + tag=_with_pass_tag(fwd_node.tag, "B"), + device=list(bwd_node.device) if bwd_node.device is not None else None, + stream=bwd_node.stream, + node_meta={ + "fwd_uid": f_uid, + "bucket_idx": bi, + "num_buckets": len(buckets), + "bucket_key": f_uid, + "compute_loss": bool(bwd_node.node_meta.get("compute_loss", False)) and bi == (len(buckets) - 1), + }, + ) + ) + + first_fwd = fwd_bucket_uids[0] + last_fwd = fwd_bucket_uids[-1] + # Mirror backward bucket chain in reverse order. + first_bwd = bwd_bucket_uids[-1] + last_bwd = bwd_bucket_uids[0] + + def _remap(uid: str) -> str: + if uid == fwd_uid: + return last_fwd + if uid == bwd_uid: + return first_bwd + return uid + + # Remove originals and their incident edges. + _remove_node_and_incident_edges(dag, fwd_uid) + _remove_node_and_incident_edges(dag, bwd_uid) + + # Re-attach incoming/outgoing edges at bucket chain boundaries. + for e in in_fwd: + dag.add_edge(TrainingDAGEdge(src_uid=_remap(e.src_uid), dst_uid=first_fwd, dep_kind=e.dep_kind, tensor_name=e.tensor_name)) + for e in out_fwd: + dst = first_bwd if e.dst_uid == bwd_uid else e.dst_uid + dag.add_edge(TrainingDAGEdge(src_uid=last_fwd, dst_uid=dst, dep_kind=e.dep_kind, tensor_name=e.tensor_name)) + for e in in_bwd: + dag.add_edge(TrainingDAGEdge(src_uid=_remap(e.src_uid), dst_uid=first_bwd, dep_kind=e.dep_kind, tensor_name=e.tensor_name)) + for e in out_bwd: + dag.add_edge(TrainingDAGEdge(src_uid=last_bwd, dst_uid=_remap(e.dst_uid), dep_kind=e.dep_kind, tensor_name=e.tensor_name)) + + # Internal chain edges. + for i in range(len(fwd_bucket_uids) - 1): + dag.add_edge(TrainingDAGEdge(src_uid=fwd_bucket_uids[i], dst_uid=fwd_bucket_uids[i + 1], dep_kind="data", tensor_name=None)) + for i in range(len(bwd_bucket_uids) - 1, 0, -1): + dag.add_edge(TrainingDAGEdge(src_uid=bwd_bucket_uids[i], dst_uid=bwd_bucket_uids[i - 1], dep_kind="data", tensor_name=None)) + + +def _insert_reduce_comm_nodes( + dag: TrainingDAG, + filters: list[dict[str, Any]], + devices: list[int], + comm_stream: str | None = None, +) -> None: + reduce_idx = sum(1 for n in dag.nodes.values() if n.node_kind == "REDUCE_COMM") + expected = sorted(int(d) for d in devices) + for node in list(dag.nodes.values()): + if node.node_kind != "COMPUTE" or not _is_backward_weight_subkind(node.compute_subkind): + continue + if not any(_match_filter(node.tag, flt) for flt in filters): + continue + if node.device is None: + raise ValueError( + f"replicate requires placed backward weight-gradient nodes, but node {node.uid} has device=None; " + f"expected devices={devices}" + ) + got = sorted(int(d) for d in node.device) + if got != expected: + raise ValueError( + f"replicate device mismatch for node {node.uid}: " + f"expected devices={devices}, node_devices={node.device}" + ) + reduce_uid = f"reduce.{reduce_idx}" + reduce_idx += 1 + reduce_node = TrainingDAGNode( + uid=reduce_uid, + node_kind="REDUCE_COMM", + compute_subkind=None, + tag=dict(node.tag), + device=list(node.device), + stream=comm_stream if comm_stream is not None else "default_stream", + node_meta={ + "bwd_uid": node.uid, + "bucket_key": node.node_meta.get("bucket_key"), + }, + ) + dag.add_node(reduce_node) + dag.add_edge( + TrainingDAGEdge( + src_uid=node.uid, + dst_uid=reduce_uid, + dep_kind="data", + tensor_name=None, + ) + ) + _rewire_bwd_successors_through_sync(dag, node.uid, reduce_uid) + + +def _node_has_trainable_params(dag: TrainingDAG, node: TrainingDAGNode) -> bool: + meta = node.node_meta + if _is_backward_compute_subkind(node.compute_subkind): + fwd_uid = meta.get("fwd_uid") + if isinstance(fwd_uid, str) and fwd_uid in dag.nodes: + meta = dag.nodes[fwd_uid].node_meta + graphargs = meta.get("graphargs") + param_idxs = meta.get("param_idxs") + if graphargs is None or param_idxs is None: + return True + return any( + 0 <= int(i) < len(graphargs) + and graphargs[int(i)] is not None + and bool(getattr(graphargs[int(i)], "requires_grad", False)) + for i in param_idxs + ) + + +def _insert_all_gather_comm_nodes( + dag: TrainingDAG, + filters: list[dict[str, Any]], + devices: list[int], + comm_stream: str | None = None, +) -> None: + ag_idx = sum(1 for n in dag.nodes.values() if n.node_kind == "ALL_GATHER_COMM") + expected = sorted(int(d) for d in devices) + for node in list(dag.nodes.values()): + if ( + node.node_kind != "COMPUTE" + or node.compute_subkind not in ("FWD", "BWD", "BWD_I", "BWD_W") + ): + continue + if not any(_match_filter(node.tag, flt) for flt in filters): + continue + if not _node_has_trainable_params(dag, node): + continue + if node.device is None: + raise ValueError( + f"replicate(shard_params=True) requires placed nodes, but node {node.uid} has device=None; " + f"expected devices={devices}" + ) + got = sorted(int(d) for d in node.device) + if got != expected: + raise ValueError( + f"replicate device mismatch for node {node.uid}: " + f"expected devices={devices}, node_devices={node.device}" + ) + ag_uid = f"all_gather.{ag_idx}" + ag_idx += 1 + ag_tag = dict(node.tag) + ag_node = TrainingDAGNode( + uid=ag_uid, + node_kind="ALL_GATHER_COMM", + compute_subkind=None, + tag=ag_tag, + device=list(node.device), + stream=comm_stream if comm_stream is not None else "default_stream", + node_meta={ + "compute_uid": node.uid, + "bucket_key": node.node_meta.get("bucket_key"), + }, + ) + dag.add_node(ag_node) + node.node_meta["zero_free_full_params_after"] = True + if node.compute_subkind == "BWD_W": + for pred_uid in sorted(dag.preds.get(node.uid, set())): + pred = dag.nodes[pred_uid] + if pred.node_kind == "COMPUTE" and pred.compute_subkind == "BWD_I": + dag.add_edge( + TrainingDAGEdge( + src_uid=pred.uid, + dst_uid=ag_uid, + dep_kind="data", + tensor_name=None, + ) + ) + # ALL_GATHER is a pre-dependency of compute: AG -> COMPUTE + dag.add_edge( + TrainingDAGEdge( + src_uid=ag_uid, + dst_uid=node.uid, + dep_kind="data", + tensor_name=None, + ) + ) + + +def _insert_reduce_scatter_comm_nodes( + dag: TrainingDAG, + filters: list[dict[str, Any]], + devices: list[int], + comm_stream: str | None = None, +) -> None: + rs_idx = sum(1 for n in dag.nodes.values() if n.node_kind == "REDUCE_SCATTER_COMM") + expected = sorted(int(d) for d in devices) + for node in list(dag.nodes.values()): + if node.node_kind != "COMPUTE" or not _is_backward_weight_subkind(node.compute_subkind): + continue + if not any(_match_filter(node.tag, flt) for flt in filters): + continue + if not _node_has_trainable_params(dag, node): + continue + if node.device is None: + raise ValueError( + f"replicate(shard_grads/shard_params) requires placed backward weight-gradient nodes, but node {node.uid} has device=None; " + f"expected devices={devices}" + ) + got = sorted(int(d) for d in node.device) + if got != expected: + raise ValueError( + f"replicate device mismatch for node {node.uid}: " + f"expected devices={devices}, node_devices={node.device}" + ) + rs_uid = f"reduce_scatter.{rs_idx}" + rs_idx += 1 + rs_node = TrainingDAGNode( + uid=rs_uid, + node_kind="REDUCE_SCATTER_COMM", + compute_subkind=None, + tag=dict(node.tag), + device=list(node.device), + stream=comm_stream if comm_stream is not None else "default_stream", + node_meta={ + "bwd_uid": node.uid, + "bucket_key": node.node_meta.get("bucket_key"), + }, + ) + node.node_meta["zero_alloc_full_grads_before"] = True + if node.node_meta.pop("zero_free_full_params_after", None) is not None: + rs_node.node_meta["zero_free_full_params_after"] = True + dag.add_node(rs_node) + dag.add_edge( + TrainingDAGEdge( + src_uid=node.uid, + dst_uid=rs_uid, + dep_kind="data", + tensor_name=None, + ) + ) + _rewire_bwd_successors_through_sync(dag, node.uid, rs_uid) + + +def _insert_shard_a2a_comm_nodes( + dag: TrainingDAG, + filters: list[dict[str, Any]], + devices: list[int], + comm_stream: str | None = None, +) -> None: + expected = sorted(int(d) for d in devices) + a2a_idx = sum(1 for n in dag.nodes.values() if n.node_kind == "A2A_COMM") + + matched_uids = [ + uid for uid, node in dag.nodes.items() + if node.node_kind == "COMPUTE" and any(_match_filter(node.tag, flt) for flt in filters) + ] + + def _a2a_boundary_for_edge(src_node: TrainingDAGNode, dst_node: TrainingDAGNode) -> dict[str, Any]: + # Backward data edges are reversed from the forward graph: + # FWD: B -> A becomes BWD: A.bwd -> B.bwd + # The A2A tensor position is defined by the original forward producer, + # which is the fwd node paired with the backward edge destination. + if ( + _is_backward_activation_subkind(src_node.compute_subkind) + and _is_backward_activation_subkind(dst_node.compute_subkind) + ): + fwd_producer_uid = dst_node.node_meta.get("fwd_uid") + if not isinstance(fwd_producer_uid, str) or fwd_producer_uid not in dag.nodes: + raise ValueError( + f"A2A_COMM could not resolve forward producer for BWD edge " + f"{src_node.uid}->{dst_node.uid}: fwd_uid={fwd_producer_uid!r}" + ) + producer = dag.nodes[fwd_producer_uid] + else: + producer = src_node + + binfo = producer.node_meta.get("a2a_boundary_after") + if not isinstance(binfo, dict) or binfo.get("tensor_idx") is None: + raise ValueError( + f"A2A_COMM missing tensor_idx for edge {src_node.uid}->{dst_node.uid}; " + f"producer={producer.uid} boundary_info={binfo!r}" + ) + return binfo + + for uid in matched_uids: + if uid not in dag.nodes: + continue + node = dag.nodes[uid] + if node.device is None or sorted(int(d) for d in node.device) != expected: + raise ValueError( + f"shard requires matched node {uid} to have devices={devices}, got node_devices={node.device}" + ) + + if node.compute_subkind == "FWD": + node.node_meta["apply_zero"] = False + elif _is_backward_compute_subkind(node.compute_subkind): + fwd_uid = node.node_meta.get("fwd_uid") + if isinstance(fwd_uid, str) and fwd_uid in dag.nodes: + dag.nodes[fwd_uid].node_meta["apply_zero"] = False + + # Shard replaces replicate-style grad/param sync comms around matched + # compute nodes; remove existing gather/reduce comm nodes attached to + # this compute node before inserting A2A edges. + removable_sync_kinds = {"ALL_GATHER_COMM", "REDUCE_COMM", "REDUCE_SCATTER_COMM"} + removable_comm_uids: set[str] = set() + for e in list(dag.edges): + if e.dep_kind != "data": + continue + if e.src_uid == uid and e.dst_uid in dag.nodes: + other = dag.nodes[e.dst_uid] + if other.node_kind in removable_sync_kinds: + removable_comm_uids.add(other.uid) + elif e.dst_uid == uid and e.src_uid in dag.nodes: + other = dag.nodes[e.src_uid] + if other.node_kind in removable_sync_kinds: + removable_comm_uids.add(other.uid) + for comm_uid in sorted(removable_comm_uids): + if comm_uid not in dag.nodes: + continue + # Bypass removed sync-comm nodes so dataflow remains connected: + # pred -> COMM -> succ ==> pred -> succ + in_comm = [e for e in list(dag.edges) if e.dep_kind == "data" and e.dst_uid == comm_uid] + out_comm = [e for e in list(dag.edges) if e.dep_kind == "data" and e.src_uid == comm_uid] + for ie in in_comm: + for oe in out_comm: + dag.add_edge( + TrainingDAGEdge( + src_uid=ie.src_uid, + dst_uid=oe.dst_uid, + dep_kind="data", + tensor_name=(oe.tensor_name if oe.tensor_name is not None else ie.tensor_name), + ) + ) + _remove_node_and_incident_edges(dag, comm_uid) + if removable_comm_uids: + node.node_meta.pop("zero_alloc_full_grads_before", None) + node.node_meta.pop("zero_free_full_params_after", None) + if node.compute_subkind == "BWD_W": + continue + + def _is_a2a_compute_edge(e: TrainingDAGEdge) -> bool: + src = dag.nodes[e.src_uid] + dst = dag.nodes[e.dst_uid] + if src.node_kind != "COMPUTE" or dst.node_kind != "COMPUTE": + return False + if node.compute_subkind == "FWD": + return src.compute_subkind == "FWD" and dst.compute_subkind == "FWD" + if _is_backward_activation_subkind(node.compute_subkind): + return ( + _is_backward_activation_subkind(src.compute_subkind) + and _is_backward_activation_subkind(dst.compute_subkind) + ) + return False + + incoming = [ + e for e in list(dag.edges) + if e.dep_kind == "data" + and e.dst_uid == uid + and _is_a2a_compute_edge(e) + ] + outgoing = [ + e for e in list(dag.edges) + if e.dep_kind == "data" + and e.src_uid == uid + and _is_a2a_compute_edge(e) + ] + + for e in incoming: + src = dag.nodes[e.src_uid] + if src.device is None or sorted(int(d) for d in src.device) != expected: + raise ValueError( + f"shard precondition failed for incoming edge {e.src_uid}->{uid}: " + f"upstream compute must be replicated on devices={devices}, got {src.device}" + ) + if src.tag.get("PASS") != node.tag.get("PASS"): + raise ValueError( + f"A2A_COMM requires matching compute pass tags; " + f"incoming edge {e.src_uid}->{uid} has src.PASS={src.tag.get('PASS')} " + f"dst.PASS={node.tag.get('PASS')}" + ) + for e in outgoing: + dst = dag.nodes[e.dst_uid] + if dst.device is None or sorted(int(d) for d in dst.device) != expected: + raise ValueError( + f"shard precondition failed for outgoing edge {uid}->{e.dst_uid}: " + f"downstream compute must be replicated on devices={devices}, got {dst.device}" + ) + if dst.tag.get("PASS") != node.tag.get("PASS"): + raise ValueError( + f"A2A_COMM requires matching compute pass tags; " + f"outgoing edge {uid}->{e.dst_uid} has src.PASS={node.tag.get('PASS')} " + f"dst.PASS={dst.tag.get('PASS')}" + ) + + for e in incoming: + src_node = dag.nodes[e.src_uid] + binfo = _a2a_boundary_for_edge(src_node, node) + a2a_tensor_idx = binfo["tensor_idx"] + comm_uid = f"a2a.{a2a_idx}" + a2a_idx += 1 + comm_node = TrainingDAGNode( + uid=comm_uid, + node_kind="A2A_COMM", + compute_subkind=None, + tag=dict(node.tag), + device=list(node.device) if node.device is not None else None, + stream=comm_stream if comm_stream is not None else "default_stream", + node_meta={ + "direction": "incoming", + "target_uid": uid, + "a2a_tensor_idx": a2a_tensor_idx, + "bucket_key": node.node_meta.get("bucket_key", src_node.node_meta.get("bucket_key")), + }, + ) + dag.add_node(comm_node) + _remove_edge(dag, e) + dag.add_edge(TrainingDAGEdge(src_uid=e.src_uid, dst_uid=comm_uid, dep_kind="data", tensor_name=e.tensor_name)) + dag.add_edge(TrainingDAGEdge(src_uid=comm_uid, dst_uid=uid, dep_kind="data", tensor_name=e.tensor_name)) + + for e in outgoing: + dst_node = dag.nodes[e.dst_uid] + binfo = _a2a_boundary_for_edge(node, dst_node) + a2a_tensor_idx = binfo["tensor_idx"] + comm_uid = f"a2a.{a2a_idx}" + a2a_idx += 1 + comm_node = TrainingDAGNode( + uid=comm_uid, + node_kind="A2A_COMM", + compute_subkind=None, + tag=dict(node.tag), + device=list(node.device) if node.device is not None else None, + stream=comm_stream if comm_stream is not None else "default_stream", + node_meta={ + "direction": "outgoing", + "source_uid": uid, + "a2a_tensor_idx": a2a_tensor_idx, + "bucket_key": node.node_meta.get("bucket_key", dst_node.node_meta.get("bucket_key")), + }, + ) + dag.add_node(comm_node) + _remove_edge(dag, e) + dag.add_edge(TrainingDAGEdge(src_uid=uid, dst_uid=comm_uid, dep_kind="data", tensor_name=e.tensor_name)) + dag.add_edge(TrainingDAGEdge(src_uid=comm_uid, dst_uid=e.dst_uid, dep_kind="data", tensor_name=e.tensor_name)) + + +def _apply_split_directive( + dag: TrainingDAG, + flt: dict[str, Any], + dim_name: str, + num_microbatches: int, +) -> None: + matched = {uid for uid, n in dag.nodes.items() if _match_filter(n.tag, flt)} + if not matched: + raise ValueError(f"split(filter={flt}) matched zero nodes") + + # UPD nodes are global reduction/update sinks and must not be duplicated + # by microbatch split; duplicated predecessors should continue to target + # the same UPD node(s). + matched = {uid for uid in matched if dag.nodes[uid].node_kind != "UPD"} + if not matched: + raise ValueError(f"split(filter={flt}) matched no splittable non-UPD nodes") + + topo = _topological_order(dag) + idx = {u: i for i, u in enumerate(topo)} + mpos = sorted(idx[u] for u in matched) + lo, hi = mpos[0], mpos[-1] + interleaved = [ + u for u in topo[lo:hi + 1] + if u not in matched and dag.nodes[u].node_kind != "UPD" + ] + if interleaved: + raise ValueError( + f"split requires a contiguous sub-DAG; found interleaved non-matching nodes: {interleaved[:8]}" + ) + + incoming_boundary = [e for e in list(dag.edges) if e.dst_uid in matched and e.src_uid not in matched] + outgoing_boundary = [e for e in list(dag.edges) if e.src_uid in matched and e.dst_uid not in matched] + + sources = { + u for u in matched + if not any(pred in matched for pred in dag.preds.get(u, set())) + } + sinks = { + u for u in matched + if not any(succ in matched for succ in dag.succs.get(u, set())) + } + + # Ensure outside edges only touch source/sink boundaries. + for e in incoming_boundary: + if dag.nodes[e.src_uid].node_kind == "UPD" or dag.nodes[e.dst_uid].node_kind == "UPD": + continue + if e.dst_uid not in sources: + raise ValueError( + f"split boundary violation: incoming outside edge targets non-source node {e.dst_uid}" + ) + for e in outgoing_boundary: + if dag.nodes[e.src_uid].node_kind == "UPD" or dag.nodes[e.dst_uid].node_kind == "UPD": + continue + if e.src_uid not in sinks: + raise ValueError( + f"split boundary violation: outgoing outside edge starts at non-sink node {e.src_uid}" + ) + + inside_edges = [e for e in list(dag.edges) if e.src_uid in matched and e.dst_uid in matched] + + # Tag originals as microbatch 0. + for uid in matched: + dag.nodes[uid].tag[dim_name] = 0 + + # Create copies for microbatches 1..N-1. + copy_uid: dict[tuple[str, int], str] = {} + for mb in range(1, num_microbatches): + for uid in topo: + if uid not in matched: + continue + nu = f"{uid}.split{dim_name}{mb}" + copy_uid[(uid, mb)] = nu + + uid_meta_fields = ( + "fwd_uid", + "compute_uid", + "bwd_uid", + "source_uid", + "target_uid", + "src_uid", + "dst_uid", + "from_uid", + "to_uid", + ) + for mb in range(1, num_microbatches): + for uid in topo: + if uid not in matched: + continue + n = dag.nodes[uid] + nu = copy_uid[(uid, mb)] + copied_meta = dict(n.node_meta) + for k in uid_meta_fields: + v = copied_meta.get(k) + if isinstance(v, str) and v in matched: + copied_meta[k] = copy_uid[(v, mb)] + dag.add_node( + TrainingDAGNode( + uid=nu, + node_kind=n.node_kind, + compute_subkind=n.compute_subkind, + tag={**n.tag, dim_name: mb}, + device=list(n.device) if n.device is not None else None, + stream=n.stream, + node_meta=copied_meta, + ) + ) + + # Duplicate internal sub-DAG edges for each copied microbatch. + for mb in range(1, num_microbatches): + for e in inside_edges: + dag.add_edge( + TrainingDAGEdge( + src_uid=copy_uid[(e.src_uid, mb)], + dst_uid=copy_uid[(e.dst_uid, mb)], + dep_kind=e.dep_kind, + tensor_name=e.tensor_name, + ) + ) + + # Duplicate incoming edges onto copied sources. + for mb in range(1, num_microbatches): + for e in incoming_boundary: + dag.add_edge( + TrainingDAGEdge( + src_uid=e.src_uid, + dst_uid=copy_uid[(e.dst_uid, mb)], + dep_kind=e.dep_kind, + tensor_name=e.tensor_name, + ) + ) + + # Duplicate outgoing edges from copied sinks. + for mb in range(1, num_microbatches): + for e in outgoing_boundary: + dag.add_edge( + TrainingDAGEdge( + src_uid=copy_uid[(e.src_uid, mb)], + dst_uid=e.dst_uid, + dep_kind=e.dep_kind, + tensor_name=e.tensor_name, + ) + ) + + +def _match_set_contiguous_subdag( + dag: TrainingDAG, + matched: set[str], + *, + context: str, +) -> tuple[set[str], set[str], set[str], set[str]]: + if not matched: + raise ValueError(f"{context}: filter matched zero nodes") + # Contiguity should be dependency-based, not tied to one arbitrary + # topological ordering (which can interleave independent branches). + # A non-matching node is considered interleaving iff it is on some path + # between matched nodes: matched -> ... -> node -> ... -> matched. + interleaved: list[str] = [] + for u in dag.nodes.keys(): + if u in matched: + continue + # Has matched ancestor? + has_matched_ancestor = False + seen = {u} + stack = list(dag.preds.get(u, set())) + while stack: + cur = stack.pop() + if cur in seen: + continue + seen.add(cur) + if cur in matched: + has_matched_ancestor = True + break + stack.extend(dag.preds.get(cur, set())) + if not has_matched_ancestor: + continue + + # Has matched descendant? + has_matched_descendant = False + seen = {u} + stack = list(dag.succs.get(u, set())) + while stack: + cur = stack.pop() + if cur in seen: + continue + seen.add(cur) + if cur in matched: + has_matched_descendant = True + break + stack.extend(dag.succs.get(cur, set())) + if has_matched_descendant: + interleaved.append(u) + + if interleaved: + raise ValueError( + f"{context}: matched nodes are not contiguous; interleaved non-matching nodes: {interleaved[:8]}" + ) + + sources = { + u for u in matched + if not any(pred in matched for pred in dag.preds.get(u, set())) + } + sinks = { + u for u in matched + if not any(succ in matched for succ in dag.succs.get(u, set())) + } + if not sources or not sinks: + raise ValueError(f"{context}: failed to identify non-empty source/sink sets") + compute_nodes = {u for u in matched if dag.nodes[u].node_kind == "COMPUTE"} + + def _has_upstream_compute(u: str) -> bool: + seen = {u} + stack = list(dag.preds.get(u, set())) + while stack: + cur = stack.pop() + if cur in seen or cur not in matched: + continue + seen.add(cur) + if dag.nodes[cur].node_kind == "COMPUTE": + return True + stack.extend(dag.preds.get(cur, set())) + return False + + def _has_downstream_compute(u: str) -> bool: + seen = {u} + stack = list(dag.succs.get(u, set())) + while stack: + cur = stack.pop() + if cur in seen or cur not in matched: + continue + seen.add(cur) + if dag.nodes[cur].node_kind == "COMPUTE": + return True + stack.extend(dag.succs.get(cur, set())) + return False + + compute_sources = {u for u in compute_nodes if not _has_upstream_compute(u)} + compute_sinks = {u for u in compute_nodes if not _has_downstream_compute(u)} + if not compute_sources or not compute_sinks: + raise ValueError(f"{context}: failed to identify non-empty compute source/sink sets") + + return sources, sinks, compute_sources, compute_sinks + + +@dataclass(frozen=True) +class _OrderSegment: + source_uids: set[str] + sink_uids: set[str] + + +def _single_device_for_nodes(dag: TrainingDAG, uids: set[str]) -> list[int] | None: + device_keys = { + tuple(n.device) + for uid in uids + for n in [dag.nodes[uid]] + if n.device is not None + } + if len(device_keys) == 1: + return list(next(iter(device_keys))) + return None + + +def _add_order_dummy_node( + dag: TrainingDAG, + *, + uid: str, + directive_idx: int, + group_idx: int, + role: str, + grouped_uids: set[str], +) -> str: + dag.add_node( + TrainingDAGNode( + uid=uid, + node_kind="ORDER_DUMMY", + compute_subkind=None, + tag={"order": directive_idx, "group": group_idx}, + device=_single_device_for_nodes(dag, grouped_uids), + stream=_DEFAULT_STREAM, + node_meta={"role": role}, + ) + ) + return uid + + +def _device_key_for_order_edge(dag: TrainingDAG, uid: str) -> tuple[int, ...] | None: + device = dag.nodes[uid].device + if device is None: + return None + return tuple(sorted(int(d) for d in device)) + + +def _validate_order_edge( + dag: TrainingDAG, + src_uid: str, + dst_uid: str, + *, + directive_idx: int, +) -> None: + if _has_path(dag, dst_uid, src_uid): + raise ValueError( + f"order directive[{directive_idx}] violates model dataflow: " + f"adding temporal edge {src_uid} -> {dst_uid} would create a cycle" + ) + + src_device = _device_key_for_order_edge(dag, src_uid) + dst_device = _device_key_for_order_edge(dag, dst_uid) + if src_device is None or dst_device is None: + raise ValueError( + f"order directive[{directive_idx}] requires placed nodes with a single device set: " + f"{src_uid} device={src_device}, {dst_uid} device={dst_device}" + ) + if src_device != dst_device: + raise ValueError( + f"order directive[{directive_idx}] crosses device placement: " + f"{src_uid} device={src_device}, {dst_uid} device={dst_device}" + ) + + +def _add_temporal_order_edge( + dag: TrainingDAG, + src_uid: str, + dst_uid: str, + *, + directive_idx: int, +) -> None: + _validate_order_edge(dag, src_uid, dst_uid, directive_idx=directive_idx) + logger.debug("Adding temporal dependency edge for order directive: %s -> %s", src_uid, dst_uid) + dag.add_edge( + TrainingDAGEdge( + src_uid=src_uid, + dst_uid=dst_uid, + dep_kind="temporal", + tensor_name=None, + ) + ) + + +def _apply_order_directive( + dag: TrainingDAG, + filter_groups: list[list[dict[str, Any]]], + *, + directive_idx: int, +) -> None: + segments: list[_OrderSegment] = [] + + for group_idx, group in enumerate(filter_groups): + group_sources: list[set[str]] = [] + group_sinks: list[set[str]] = [] + grouped_uids: set[str] = set() + for filter_idx, flt in enumerate(group): + matched = { + uid for uid, n in dag.nodes.items() + if n.node_kind != "ORDER_DUMMY" and _match_filter(n.tag, flt) + } + _srcs, _snks, csrcs, csnks = _match_set_contiguous_subdag( + dag, + matched, + context=f"order filter[{group_idx}][{filter_idx}]={flt}", + ) + group_sources.append(csrcs) + group_sinks.append(csnks) + grouped_uids.update(matched) + + if len(group) == 1: + segments.append(_OrderSegment(source_uids=group_sources[0], sink_uids=group_sinks[0])) + continue + + src_uid = _add_order_dummy_node( + dag, + uid=f"order.{directive_idx}.group{group_idx}.source", + directive_idx=directive_idx, + group_idx=group_idx, + role="source", + grouped_uids=grouped_uids, + ) + sink_uid = _add_order_dummy_node( + dag, + uid=f"order.{directive_idx}.group{group_idx}.sink", + directive_idx=directive_idx, + group_idx=group_idx, + role="sink", + grouped_uids=grouped_uids, + ) + for csrcs in group_sources: + for v in csrcs: + _add_temporal_order_edge(dag, src_uid, v, directive_idx=directive_idx) + for csnks in group_sinks: + for u in csnks: + _add_temporal_order_edge(dag, u, sink_uid, directive_idx=directive_idx) + segments.append(_OrderSegment(source_uids={src_uid}, sink_uids={sink_uid})) + + for i in range(len(segments) - 1): + for u in segments[i].sink_uids: + for v in segments[i + 1].source_uids: + _add_temporal_order_edge(dag, u, v, directive_idx=directive_idx) + + +def apply_schedule_directives(training_dag: TrainingDAG, directives: list[Any] | None) -> None: + if not directives: + return + split_backward_keys = _validate_split_backward_order_stencil(directives) + + # Microbatch expansion must run before split-backward lowering so schedules + # can mix fused B and split BI/BW for the same bucket on different MBs. + for i, raw in enumerate(directives): + if not isinstance(raw, dict) or raw.get("op") != "split": + continue + flt, dim_name, num_microbatches = _normalize_split_directive(raw) + logger.info( + "Applying directive[%d]: split(filter=%s, dim_name=%s, num_microbatches=%d)", + i, flt, dim_name, num_microbatches + ) + _apply_split_directive(training_dag, flt, dim_name, num_microbatches) + + _apply_split_backward_stencil(training_dag, split_backward_keys) + place_streams: set[str] = set() + for i, raw in enumerate(directives): + if isinstance(raw, dict) and raw.get("op") != "place": + continue + op, filters, devices, stream, _gather_stream, _reduce_stream, shard_params, shard_grads, bucket_size = _normalize_filter_devices_directive(raw) + logger.info( + "Applying directive[%d]: %s(filters=%s, devices=%s, stream=%s, shard_params=%s, shard_grads=%s, bucket_size=%s)", + i, op, filters, devices, stream, shard_params, shard_grads, bucket_size + ) + matched_nodes = 0 + for node in training_dag.nodes.values(): + if node.node_kind not in ("COMPUTE", "UPD"): + continue + if any(_match_filter(node.tag, flt) for flt in filters): + node.device = list(devices) + matched_nodes += 1 + if matched_nodes == 0: + raise ValueError(f"place directive[{i}] matched zero nodes: {raw}") + if stream is not None: + place_streams.add(stream) + + if place_streams: + if len(place_streams) != 1: + raise ValueError(f"place directives must agree on stream, got {sorted(place_streams)}") + place_stream = next(iter(place_streams)) + else: + place_stream = None + _replicate_update_nodes_by_device(training_dag) + _insert_send_recv_comm_nodes(training_dag, comm_stream=place_stream) + + for i, raw in enumerate(directives): + if isinstance(raw, dict) and raw.get("op") in ("place", "split", "order"): + continue + op, filters, devices, stream, gather_stream, reduce_stream, shard_params, shard_grads, bucket_size = _normalize_filter_devices_directive(raw) + logger.info( + "Applying directive[%d]: %s(filters=%s, devices=%s, stream=%s, gather_stream=%s, reduce_stream=%s, shard_params=%s, shard_grads=%s, bucket_size=%s)", + i, op, filters, devices, stream, gather_stream, reduce_stream, shard_params, shard_grads, bucket_size + ) + matched_nodes = [ + uid for uid, node in training_dag.nodes.items() + if node.node_kind == "COMPUTE" + and any(_match_filter(node.tag, flt) for flt in filters) + ] + if not matched_nodes: + raise ValueError(f"{op} directive[{i}] matched zero compute nodes: {raw}") + if op == "replicate": + if bucket_size is not None: + _bucket_matched_fwd_nodes(training_dag, filters, int(bucket_size)) + if shard_params: + _insert_all_gather_comm_nodes(training_dag, filters, devices, comm_stream=gather_stream) + _insert_reduce_scatter_comm_nodes(training_dag, filters, devices, comm_stream=reduce_stream) + elif shard_grads: + _insert_reduce_scatter_comm_nodes(training_dag, filters, devices, comm_stream=reduce_stream) + else: + _insert_reduce_comm_nodes(training_dag, filters, devices, comm_stream=reduce_stream) + elif op == "shard": + _insert_shard_a2a_comm_nodes(training_dag, filters, devices, comm_stream=stream) + else: + raise ValueError(f"Unsupported directive op after normalization: {op}") + + # Final pass for order directives (add temporal dependencies across sub-DAGs). + for i, raw in enumerate(directives): + if not isinstance(raw, dict) or raw.get("op") != "order": + continue + filter_groups = _parse_order_directive(raw) + logger.info("Applying directive[%d]: order(filters=%s)", i, filter_groups) + _apply_order_directive(training_dag, filter_groups, directive_idx=i) diff --git a/src/executors.py b/src/executors.py new file mode 100644 index 0000000..c5d0c59 --- /dev/null +++ b/src/executors.py @@ -0,0 +1,996 @@ +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist +from torch.autograd.graph import GradientEdge, Node +from torch.nn import Parameter + +from .backward import construct_reverse_graph, get_param_groups, _get_grad_fn_or_grad_acc +from .runtime import BufferStore, EventStore, ParamStorage, RuntimeState, StageStore +from .tasks import TaskType + + +@dataclass +class CommunicationExecutor: + """Actor-local communication operations used by the DAG dispatcher.""" + + runtime: RuntimeState + stages: StageStore + logger: Any + + def send(self, send_data: Any, peer_pp_rank: int, stream: torch.cuda.Stream) -> None: + global_dst_rank = self.runtime.pipeline_peer_global_rank(peer_pp_rank) + + with torch.cuda.stream(stream): + tensors = send_data if isinstance(send_data, (list, tuple)) else [send_data] + use_lo_hi = global_dst_rank > self.runtime.global_rank + pp_group = self.runtime.pp_lo_hi if use_lo_hi else self.runtime.pp_hi_lo + for tensor in tensors: + dist.send(tensor, dst=global_dst_rank, group=pp_group) + + def recv_fwd(self, recv_ubid: Any, peer_pp_rank: int, stream: torch.cuda.Stream) -> list: + global_src_rank = self.runtime.pipeline_peer_global_rank(peer_pp_rank) + + buf = [ + torch.empty( + shape, + dtype=dtype, + requires_grad=requires_grad, + device=self.runtime.device, + ) + for shape, dtype, requires_grad in self.stages.bucket(recv_ubid).forward_input_meta + ] + with torch.cuda.stream(stream): + use_hi_lo = global_src_rank > self.runtime.global_rank + pp_group = self.runtime.pp_hi_lo if use_hi_lo else self.runtime.pp_lo_hi + for tensor in buf: + dist.recv(tensor, src=global_src_rank, group=pp_group) + return buf + + def recv_bwd(self, shape_meta: list, peer_pp_rank: int, stream: torch.cuda.Stream) -> list: + global_src_rank = self.runtime.pipeline_peer_global_rank(peer_pp_rank) + + buf = [ + torch.empty(shape, dtype=dtype, device=self.runtime.device) + for shape, dtype in shape_meta + ] + with torch.cuda.stream(stream): + use_hi_lo = global_src_rank > self.runtime.global_rank + pp_group = self.runtime.pp_hi_lo if use_hi_lo else self.runtime.pp_lo_hi + for tensor in buf: + dist.recv(tensor, src=global_src_rank, group=pp_group) + return buf + + def all_to_all(self, input_tensor: torch.Tensor, stream: torch.cuda.Stream) -> torch.Tensor: + output_buf = torch.empty_like(input_tensor, device=self.runtime.device) + with torch.cuda.stream(stream): + dist.all_to_all_single(output_buf, input_tensor, group=self.runtime.ep_group) + return output_buf + + def all_reduce_grads(self, ubid: Any, stream: torch.cuda.Stream) -> int: + assert ubid is not None, "all_reduce_grads requires a non-None ubid" + if not self.has_trainable_params_for_collective(ubid, "all_reduce_grads"): + return 0 + bucket = self.stages.bucket(ubid) + grad_tensors = [] + for idx in bucket.trainable_param_idxs: + param = bucket.forward_args[idx] + assert param is not None, ( + f"all_reduce_grads: ubid={ubid} idx={idx} param is None" + ) + assert param.grad is not None, ( + f"all_reduce_grads: ubid={ubid} idx={idx} param.grad is None" + ) + grad_tensors.append(param.grad) + + total_bytes = sum(grad.numel() * grad.element_size() for grad in grad_tensors) + with torch.cuda.stream(stream): + for grad in grad_tensors: + dist.all_reduce(grad, group=self.runtime.dp_group) + return total_bytes + + def reduce_scatter(self, ubid: Any, stream: torch.cuda.Stream) -> int: + assert ubid is not None, "reduce_scatter requires a non-None ubid" + if not self.has_trainable_params_for_collective(ubid, "reduce_scatter"): + return 0 + assert ubid in self.stages.grad_sharded_ubids, ( + f"reduce_scatter: ubid={ubid} is not in grad_sharded_ubids=" + f"{self.stages.grad_sharded_ubids}" + ) + bucket = self.stages.bucket(ubid) + assert bucket.param_shard_info is not None, ( + f"reduce_scatter: missing param_shard_info for ubid={ubid}" + ) + flat_grads = bucket.flat_grads + rs_out = bucket.reduce_scatter_grads + assert flat_grads is not None, ( + f"reduce_scatter: missing flat_grads buffer for ubid={ubid}" + ) + assert rs_out is not None, ( + f"reduce_scatter: missing reduce_scatter_grads buffer for ubid={ubid}" + ) + with torch.cuda.stream(stream): + total_bytes = flat_grads.numel() * flat_grads.element_size() + tmp = torch.empty_like(rs_out) + dist.reduce_scatter_tensor(tmp, flat_grads, group=self.runtime.dp_group) + rs_out.add_(tmp) + return total_bytes + + def has_trainable_params_for_collective(self, ubid: Any, op_name: str) -> bool: + bucket = self.stages.get_bucket(ubid) + if bucket is not None and bucket.trainable_param_idxs: + return True + self.logger.warning( + "%s: skipping collective for ubid=%s because it has no trainable param indices", + op_name, + ubid, + ) + return False + + +@dataclass +class ComputeExecutor: + """Actor-local bucket forward/backward execution.""" + + runtime: RuntimeState + stages: StageStore + logger: Any + + def log_compute_loss_inputs( + self, + labels: Any, + node: Any, + fwd_key: Any, + fwd_out: dict, + ) -> None: + def _summarize_value(value: Any) -> str: + if isinstance(value, torch.Tensor): + return ( + f"Tensor(shape={tuple(value.shape)}, dtype={value.dtype}, " + f"requires_grad={value.requires_grad}, device={value.device})" + ) + if isinstance(value, (list, tuple)): + return "[" + ", ".join(_summarize_value(v) for v in value) + "]" + if value is None: + return "None" + return type(value).__name__ + + self.logger.debug( + "compute_loss inputs: rank=%s node_uid=%s node_type=%s tag=%s " + "fwd_key=%s labels=%s out_with_grad=%s pre_detach_outs=%s " + "detached_outs=%s send_output=%s", + self.runtime.global_rank, + getattr(node, "uid", None), + getattr(getattr(node, "task_type", None), "value", None), + getattr(node, "tag", None), + fwd_key, + _summarize_value(labels), + _summarize_value(fwd_out.get("out_with_grad")), + _summarize_value(fwd_out.get("pre_detach_outs")), + _summarize_value(fwd_out.get("detached_outs")), + _summarize_value(fwd_out.get("send_output")), + ) + + def forward(self, ubid: Any, input_tensors: Any, compute_stream: torch.cuda.Stream) -> dict: + bucket = self.stages.bucket(ubid) + fwd_fn = bucket.forward_fn + fwd_args = bucket.forward_args + input_idxs = bucket.input_idxs + + if not isinstance(input_tensors, (list, tuple)): + input_tensors = [input_tensors] + + for i, tensor in zip(input_idxs, input_tensors): + if isinstance(tensor, (list, tuple)): + tensor = tensor[0] + if tensor.requires_grad: + tensor = tensor.detach().requires_grad_(True) + fwd_args[i] = tensor + + fwd_inputs = [fwd_args[i] for i in input_idxs] + inp_with_grad = [t for t in fwd_inputs if t is not None and t.requires_grad] + + with torch.cuda.stream(compute_stream): + output = fwd_fn(fwd_args) + + for i in input_idxs: + fwd_args[i] = None + + out_list = list(output) if isinstance(output, tuple) else [output] + possibly_detached = [ + t.detach().requires_grad_(True) + if isinstance(t, torch.Tensor) and t.requires_grad + else t + for t in out_list + ] + out_with_grad = [ + t for t in out_list if isinstance(t, torch.Tensor) and t.requires_grad + ] + + return { + "pre_detach_outs": out_list, + "detached_outs": possibly_detached, + "out_with_grad": out_with_grad, + "send_output": output, + "inp_with_grad": inp_with_grad, + "fwd_inputs": fwd_inputs, + } + + def backward( + self, + ubid: Any, + mb_idx: int, + outputs_or_loss: list, + upstream_grads: Any, + pre_detach_outs: Any, + detached_outs: Any, + inp_with_grad: Any, + out_with_grad: Any, + compute_stream: torch.cuda.Stream, + ) -> dict | None: + if pre_detach_outs is None: + if upstream_grads is not None: + with torch.cuda.stream(compute_stream): + torch.autograd.backward(outputs_or_loss, upstream_grads) + else: + with torch.cuda.stream(compute_stream): + outputs_or_loss[0].backward() + else: + bwd_pairs = [ + (p, d.grad) + for p, d in zip(pre_detach_outs, detached_outs) + if (isinstance(d, torch.Tensor) and d.requires_grad and d.grad is not None) + ] + assert bwd_pairs, ( + f"BWD ubid={ubid} mb={mb_idx}: detached boundary has no tensors " + "with a materialized grad" + ) + if bwd_pairs: + outputs_bwd = [p for p, _g in bwd_pairs] + grads_bwd = [g for _p, g in bwd_pairs] + with torch.cuda.stream(compute_stream): + torch.autograd.backward(outputs_bwd, grads_bwd) + + if inp_with_grad: + output_grads = [t.grad for t in inp_with_grad if t.grad is not None] + if output_grads: + return {"send_output": output_grads} + return None + + @staticmethod + def grad_with_param_layout(weight: Parameter, grad: torch.Tensor) -> torch.Tensor: + if ( + grad.dtype == weight.dtype + and grad.device == weight.device + and grad.layout == weight.layout + and tuple(grad.stride()) == tuple(weight.stride()) + ): + return grad + if weight.layout == torch.strided and grad.layout != torch.strided: + grad = grad.to_dense() + out = torch.empty_strided( + tuple(weight.shape), + tuple(weight.stride()), + dtype=weight.dtype, + device=weight.device, + ) + out.copy_(grad) + return out + + def fused_backward(self, stage_outputs_or_loss: list, output_grads: Any, weights: list) -> None: + wts = [w for w in weights if w.requires_grad] + if not wts: + return + dweights = torch.autograd.grad( + stage_outputs_or_loss, + inputs=wts, + grad_outputs=output_grads, + retain_graph=False, + allow_unused=True, + ) + for w, dw in zip(wts, dweights): + if dw is None: + continue + dw = self.grad_with_param_layout(w, dw) + if w.grad is None: + w.grad = dw + else: + if ( + w.grad.dtype != w.dtype + or w.grad.device != w.device + or w.grad.layout != w.layout + or tuple(w.grad.stride()) != tuple(w.stride()) + ): + w.grad = self.grad_with_param_layout(w, w.grad) + w.grad += dw + + def backward_weight_from_outputs( + self, + stage_outputs_or_loss: list, + output_grads: Any, + weights: Any, + ) -> None: + self.fused_backward(stage_outputs_or_loss, output_grads, list(weights)) + + def bucket_backward_input( + self, + stage_outputs_or_loss: list, + output_grads: Any, + input_values: list, + weights: Any, + ): + weights = list(weights) + + if output_grads is None: + output_grads = [torch.ones_like(o) for o in stage_outputs_or_loss] + + input_values = [inp for inp in input_values if inp.requires_grad] + if not input_values: + self.fused_backward(stage_outputs_or_loss, output_grads, weights) + for i, t in enumerate(stage_outputs_or_loss): + if isinstance(t, torch.Tensor): + stage_outputs_or_loss[i] = t.detach() + return (), [], None + + stage_output_grad_fns = list(filter(None, map(_get_grad_fn_or_grad_acc, stage_outputs_or_loss))) + stage_input_grad_fns = list(filter(None, map(_get_grad_fn_or_grad_acc, input_values))) + weight_grad_fns = list(filter(None, map(_get_grad_fn_or_grad_acc, weights))) + + reverse_edges_dict = construct_reverse_graph(stage_output_grad_fns) + param_groups = get_param_groups(stage_input_grad_fns, weight_grad_fns, reverse_edges_dict) + handles = [] + for param_group in param_groups: + for i, intermediate in enumerate(param_group["intermediates"]): + def get_hook(pg, idx): + def hook(grad_inputs): + if pg.get("grads") is None: + pg["grads"] = [None] * len(pg["intermediates"]) + pg["grads"][idx] = tuple(grad_inputs) + return hook + handles.append(intermediate.register_prehook(get_hook(param_group, i))) + + dinputs = torch.autograd.grad( + stage_outputs_or_loss, + inputs=input_values, + grad_outputs=output_grads, + retain_graph=True, + allow_unused=True, + ) + for inp, dinput in zip(input_values, dinputs): + if inp.grad is None: + inp.grad = dinput + else: + inp.grad += dinput + + output_backward_ctx = { + "stage_outputs_or_loss": list(stage_outputs_or_loss), + "output_grads": output_grads, + } + + for handle in handles: + handle.remove() + + return dinputs, param_groups, output_backward_ctx + + def bucket_backward_weight( + self, + weights: Any, + param_groups: list, + ubid: Any | None = None, + mb_idx: int | None = None, + ) -> None: + if not param_groups: + return + + grad_acc_to_weight: dict[Node, Parameter] = {} + + for weight in weights: + grad_acc = _get_grad_fn_or_grad_acc(weight) + grad_acc_to_weight[grad_acc] = weight + + for param_group in param_groups: + valid_edges: list[GradientEdge] = [] + valid_grad_outputs: list[torch.Tensor] = [] + + for grads_tuple, intermediate in zip( + param_group.get("grads", []), param_group["intermediates"] + ): + if grads_tuple is None: + continue + for i, grad in enumerate(grads_tuple): + if grad is not None: + valid_edges.append(GradientEdge(intermediate, i)) + valid_grad_outputs.append(grad) + + del param_group["intermediates"] + + if not valid_edges: + continue + + weight_edges = tuple(GradientEdge(w, 0) for w in param_group["params"]) + if not weight_edges: + continue + + dweights = torch.autograd.grad( + valid_edges, + weight_edges, + grad_outputs=valid_grad_outputs, + retain_graph=False, + ) + + del param_group["grads"] + + for grad_acc, dw in zip(param_group["params"], dweights): + if dw is None or grad_acc not in grad_acc_to_weight: + continue + weight = grad_acc_to_weight[grad_acc] + dw = self.grad_with_param_layout(weight, dw) + if weight.grad is None: + weight.grad = dw + else: + if ( + weight.grad.dtype != weight.dtype + or weight.grad.device != weight.device + or weight.grad.layout != weight.layout + or tuple(weight.grad.stride()) != tuple(weight.stride()) + ): + weight.grad = self.grad_with_param_layout(weight, weight.grad) + weight.grad += dw + + +@dataclass +class DagExecutor: + """Execute a loaded TrainingDAG in sorted order.""" + + runtime: RuntimeState + stages: StageStore + buffers: BufferStore + events: EventStore + params: ParamStorage + communication: CommunicationExecutor + compute: ComputeExecutor + logger: Any + + @staticmethod + def _node_meta(node: Any) -> dict: + return getattr(node, "node_meta", {}) or {} + + def _node_bucket_key(self, node: Any) -> Any | None: + return self._node_meta(node).get("bucket_key") + + def _sync_payload_ubid(self, node: Any) -> Any | None: + bk = self._node_bucket_key(node) + if bk is not None: + return bk + if node.data_preds: + pred_bk = self._node_bucket_key(node.data_preds[0]) + if pred_bk is not None: + return pred_bk + return None + + def _sync_payload_ubids(self, node: Any) -> list[Any]: + sync_ubids = self._node_meta(node).get("sync_ubids") + if sync_ubids: + return list(sync_ubids) + ubid = self._sync_payload_ubid(node) + return [ubid] if ubid is not None else [] + + def _rf_enter(self, label: str): + if not self.runtime.pytorch_profiler_enabled: + return None + mt = torch.autograd.set_multithreading_enabled(False) + mt.__enter__() + rf = torch.profiler.record_function(label) + rf.__enter__() + return (rf, mt) + + @staticmethod + def _rf_exit(rf) -> None: + if rf is not None: + rf_ctx, mt = rf + rf_ctx.__exit__(None, None, None) + mt.__exit__(None, None, None) + + @staticmethod + def _node_tag_str(node: Any) -> str: + tag = getattr(node, "tag", None) + if not isinstance(tag, dict) or not tag: + return "{}" + items = sorted(tag.items(), key=lambda kv: kv[0]) + return "{" + ",".join(f"{k}={v}" for k, v in items) + "}" + + def _wait_for_all_gather(self, compute_node: Any) -> None: + compute_stream = self.runtime.stream_for_task(compute_node) + for pred in compute_node.data_preds: + if pred.task_type == TaskType.ALL_GATHER: + ag_evt = self.events.all_gather.get(pred.uid) + if ag_evt is not None: + compute_stream.wait_event(ag_evt) + + def _all_to_all_ep_boundary( + self, + node: Any, + tensor: torch.Tensor, + stream: torch.cuda.Stream, + ) -> torch.Tensor: + direction = self._node_meta(node).get("direction") + if direction == "outgoing": + tensor = tensor.contiguous() + elif direction != "incoming": + raise ValueError( + f"A2A node uid={getattr(node, 'uid', '')} has invalid " + f"direction={direction!r}; expected 'incoming' or 'outgoing'" + ) + + tensor = self.communication.all_to_all(tensor, stream=stream) + if direction == "incoming": + tensor = tensor.contiguous() + return tensor + + def run( + self, + dag: Any, + sorted_dag_nodes: list[Any], + inputs: Any, + labels: Any, + loss_buffer: list, + loss_fn=None, + ) -> None: + """Run one iteration of the loaded TrainingDAG.""" + assert dag is not None, "load_training_dag() must be called before run_dag()" + assert sorted_dag_nodes is not None, "load_training_dag() must initialize sorted node order" + + self.params.drain_pending_frees() + debug_enabled = self.logger.isEnabledFor(logging.DEBUG) + + self.buffers.reset() + self.events.reset() + self.params.clear_param_grads() + default_stream = self.runtime.default_stream() + self.params.zero_grad_buffers(default_stream) + zero_evt = torch.cuda.Event() + zero_evt.record(default_stream) + for stream in self.runtime.streams.values(): + if stream is not default_stream: + stream.wait_event(zero_evt) + comp_events: dict[Any, torch.cuda.Event] = {} + + self.buffers.init_refcounts(dag) + last_comp_event_by_stream: dict[str, torch.cuda.Event] = {} + + for node in sorted_dag_nodes: + task_type = node.task_type + batch = node.batches[0] + mb_idx = batch.mb_idx + ubid = self._node_bucket_key(node) + node_stream = self.runtime.stream_for_task(node) + node_stream_id = self.runtime.stream_id(node) + node_tag = self._node_tag_str(node) + + if debug_enabled: + self.logger.debug( + f"run_dag dispatch: {task_type.value} " + f"tag={node_tag}" + ) + + task_label = f"{task_type.value}:{node_tag}:uid{node.uid}" + self.runtime.nvtx_push(task_label) + rf = self._rf_enter(task_label) + + match task_type: + case TaskType.SEND: + compute_node = node.data_preds[0] + node_stream.wait_event(comp_events[compute_node.uid]) + send_data = self.buffers.task[compute_node.uid]["send_output"] + self.communication.send(send_data, node.peer_pp_rank, stream=node_stream) + send_buf = self.buffers.task.get(compute_node.uid) + if isinstance(send_buf, dict): + send_buf["send_output"] = None + self.buffers.release(compute_node.uid) + + case TaskType.RECV: + compute_node = node.data_succs[0] + comp_evt = last_comp_event_by_stream.get(self.runtime.stream_id(compute_node)) + if comp_evt is not None: + node_stream.wait_event(comp_evt) + if compute_node.task_type == TaskType.FWD: + recv_ubid = self._node_bucket_key(compute_node) + recv_tensors = self.communication.recv_fwd( + recv_ubid, node.peer_pp_rank, stream=node_stream + ) + else: + fwd_uid = compute_node.node_meta.get("fwd_uid") + fwd_key = (compute_node.node_meta.get("bucket_key"), fwd_uid) + shape_meta = self.buffers.task[("shape_ref",) + fwd_key] + recv_tensors = self.communication.recv_bwd( + shape_meta, node.peer_pp_rank, stream=node_stream + ) + self.buffers.task[node.uid] = recv_tensors + recv_evt = torch.cuda.Event() + recv_evt.record(node_stream) + self.events.recv[node.uid] = recv_evt + + case TaskType.FWD_A2A: + fwd_pred = next(p for p in node.data_preds if p.task_type == TaskType.FWD) + node_stream.wait_event(comp_events[fwd_pred.uid]) + tensor_idx = self._node_meta(node)["a2a_tensor_idx"] + fwd_buf = dict(self.buffers.task[fwd_pred.uid]) + self.buffers.release(fwd_pred.uid) + detached_outs = list(fwd_buf["detached_outs"]) + detached_outs[tensor_idx] = self._all_to_all_ep_boundary( + node, detached_outs[tensor_idx], node_stream + ).requires_grad_(True) + fwd_buf["detached_outs"] = detached_outs + self.buffers.task[node.uid] = fwd_buf + a2a_evt = torch.cuda.Event() + a2a_evt.record(node_stream) + self.events.a2a[node.uid] = a2a_evt + + case TaskType.BWD_A2A: + bwd_pred = next( + p for p in node.data_preds + if p.task_type in (TaskType.BWD, TaskType.BWD_I) + ) + node_stream.wait_event(comp_events[bwd_pred.uid]) + tensor_idx = self._node_meta(node)["a2a_tensor_idx"] + bwd_buf = dict(self.buffers.task[bwd_pred.uid]) + self.buffers.release(bwd_pred.uid) + inp_grads = list(bwd_buf["inp_grads"]) + grad_a2a_out = inp_grads[tensor_idx] + assert grad_a2a_out is not None, ( + f"BWD_A2A tag={node_tag}: grad at a2a_tensor_idx={tensor_idx} is None" + ) + inp_grads[tensor_idx] = self._all_to_all_ep_boundary( + node, grad_a2a_out, node_stream + ) + bwd_buf["inp_grads"] = inp_grads + self.buffers.task[node.uid] = bwd_buf + a2a_evt = torch.cuda.Event() + a2a_evt.record(node_stream) + self.events.a2a[node.uid] = a2a_evt + + case TaskType.ALL_REDUCE: + bwd_node = node.data_preds[0] + ar_ubids = self._sync_payload_ubids(node) + assert ar_ubids, ( + f"ALL_REDUCE node uid={node.uid} has no sync_payload_ubids" + ) + node_stream.wait_event(comp_events[bwd_node.uid]) + for ar_ubid in ar_ubids: + self.communication.all_reduce_grads(ar_ubid, stream=node_stream) + ar_evt = torch.cuda.Event() + ar_evt.record(node_stream) + self.events.all_reduce[node.uid] = ar_evt + self.buffers.release(bwd_node.uid) + + case TaskType.REDUCE_SCATTER: + bwd_node = node.data_preds[0] + rs_ubid = self._node_bucket_key(node) + assert rs_ubid is not None, ( + f"REDUCE_SCATTER node uid={node.uid} has no bucket_key" + ) + node_stream.wait_event(comp_events[bwd_node.uid]) + rs_bytes = self.communication.reduce_scatter(rs_ubid, stream=node_stream) + rs_evt = torch.cuda.Event() + rs_evt.record(node_stream) + self.events.reduce_scatter[node.uid] = rs_evt + if rs_bytes: + self.params.defer_free_full_grads(rs_ubid, rs_evt) + if self._node_meta(node).get("zero_free_full_params_after"): + assert False and "Param free should happen after a compute node" + self.buffers.release(bwd_node.uid) + + case TaskType.ALL_GATHER: + ag_ubid = self._node_bucket_key(node) + assert ag_ubid is not None, ( + f"ALL_GATHER node uid={node.uid} has no bucket_key" + ) + self.params.all_gather_full_params(ag_ubid, stream=node_stream) + ag_evt = torch.cuda.Event() + ag_evt.record(node_stream) + self.events.all_gather[node.uid] = ag_evt + + case TaskType.FWD: + recv_pred = next( + (p for p in node.data_preds if p.task_type == TaskType.RECV), None + ) + if recv_pred is not None and recv_pred.uid in self.events.recv: + node_stream.wait_event(self.events.recv.pop(recv_pred.uid)) + + a2a_pred = next( + (p for p in node.data_preds if p.task_type == TaskType.FWD_A2A), None + ) + if a2a_pred is not None and a2a_pred.uid in self.events.a2a: + node_stream.wait_event(self.events.a2a.pop(a2a_pred.uid)) + + fwd_data_pred = next( + (p for p in node.data_preds + if p.task_type in (TaskType.FWD, TaskType.FWD_A2A)), None + ) + if fwd_data_pred is not None: + input_tensors = self.buffers.task[fwd_data_pred.uid]["detached_outs"] + self.buffers.release(fwd_data_pred.uid) + elif recv_pred is not None: + input_tensors = self.buffers.task[recv_pred.uid] + self.buffers.release(recv_pred.uid) + else: + input_tensors = inputs + + self._wait_for_all_gather(node) + + fwd_out = self.compute.forward(ubid, input_tensors, node_stream) + self.buffers.task[node.uid] = fwd_out + fwd_key = (node.node_meta.get("bucket_key"), node.uid) + self.buffers.task[("shape_ref",) + fwd_key] = [ + (t.shape, t.dtype) for t in fwd_out["out_with_grad"] + ] + self.buffers.task[fwd_key] = fwd_out + evt = torch.cuda.Event() + evt.record(node_stream) + comp_events[node.uid] = evt + last_comp_event_by_stream[node_stream_id] = evt + if self._node_meta(node).get("zero_free_full_params_after"): + self.params.defer_free_full_params(ubid, evt) + + case TaskType.BWD: + recv_pred = next( + (p for p in node.data_preds if p.task_type == TaskType.RECV), None + ) + if recv_pred is not None and recv_pred.uid in self.events.recv: + node_stream.wait_event(self.events.recv.pop(recv_pred.uid)) + self._wait_for_all_gather(node) + + a2a_pred = next( + (p for p in node.data_preds if p.task_type == TaskType.BWD_A2A), None + ) + if a2a_pred is not None and a2a_pred.uid in self.events.a2a: + node_stream.wait_event(self.events.a2a.pop(a2a_pred.uid)) + + fwd_uid = node.node_meta.get("fwd_uid") + fwd_key = (node.node_meta.get("bucket_key"), fwd_uid) + fwd_out = self.buffers.task[fwd_key] + + if self._node_meta(node).get("compute_loss", False): + assert loss_fn is not None + if self.logger.isEnabledFor(logging.DEBUG): + self.compute.log_compute_loss_inputs(labels, node, fwd_key, fwd_out) + with torch.cuda.stream(node_stream): + outputs_or_loss = [loss_fn(fwd_out["out_with_grad"][0], labels)] + upstream_grads = None + elif recv_pred is not None: + upstream_grads = self.buffers.task[recv_pred.uid] + self.buffers.release(recv_pred.uid) + if not isinstance(upstream_grads, (list, tuple)): + upstream_grads = [upstream_grads] + outputs_or_loss = fwd_out["out_with_grad"] + upstream_grads = list(upstream_grads) + else: + outputs_or_loss = fwd_out["out_with_grad"] + upstream_grads = None + + if a2a_pred is not None: + a2a_buf = self.buffers.task[a2a_pred.uid] + pre_detach_outs = fwd_out["pre_detach_outs"] + detached_outs = fwd_out["detached_outs"] + for d, g in zip(detached_outs, a2a_buf["inp_grads"]): + if ( + d is not None + and isinstance(d, torch.Tensor) + and d.requires_grad + and g is not None + ): + d.grad = g + self.buffers.release(a2a_pred.uid) + else: + bwd_pred = next( + (p for p in node.data_preds + if p.task_type in (TaskType.BWD, TaskType.BWD_I)), None + ) + if bwd_pred is not None and recv_pred is None: + prev_buf = self.buffers.task[bwd_pred.uid] + pre_detach_outs = fwd_out.get("pre_detach_outs") + detached_outs = fwd_out.get("detached_outs") + for d, g in zip(detached_outs, prev_buf["inp_grads"]): + if ( + d is not None + and isinstance(d, torch.Tensor) + and d.requires_grad + and g is not None + ): + d.grad = g + self.buffers.release(bwd_pred.uid) + else: + pre_detach_outs = None + detached_outs = None + + inp_with_grad = fwd_out.get("inp_with_grad") + if self._node_meta(node).get("zero_alloc_full_grads_before"): + self.params.alloc_full_grads(ubid, node_stream) + + bwd_out = self.compute.backward( + ubid, mb_idx, outputs_or_loss, upstream_grads, + pre_detach_outs, detached_outs, inp_with_grad, + fwd_out.get("out_with_grad"), + node_stream, + ) + buf = bwd_out if bwd_out is not None else {} + fwd_inputs_full = fwd_out.get("fwd_inputs") + buf["inp_grads"] = ( + [t.grad if (t is not None and t.requires_grad) else None + for t in fwd_inputs_full] + if fwd_inputs_full is not None + else [t.grad for t in (inp_with_grad or [])] + ) + self.params.accumulate_zero_param_grads_to_flat(ubid, node_stream) + self.buffers.task[node.uid] = buf + fwd_out.clear() + del self.buffers.task[fwd_key] + evt = torch.cuda.Event() + evt.record(node_stream) + comp_events[node.uid] = evt + self.events.backward[ubid] = evt + last_comp_event_by_stream[node_stream_id] = evt + if self._node_meta(node).get("zero_free_full_params_after"): + self.params.defer_free_full_params(ubid, evt) + + case TaskType.BWD_I: + recv_pred = next( + (p for p in node.data_preds if p.task_type == TaskType.RECV), None + ) + if recv_pred is not None and recv_pred.uid in self.events.recv: + node_stream.wait_event(self.events.recv.pop(recv_pred.uid)) + self._wait_for_all_gather(node) + + fwd_uid = node.node_meta.get("fwd_uid") + fwd_key = (node.node_meta.get("bucket_key"), fwd_uid) + fwd_out = self.buffers.task[fwd_key] + + if self._node_meta(node).get("compute_loss", False): + assert loss_fn is not None + if self.logger.isEnabledFor(logging.DEBUG): + self.compute.log_compute_loss_inputs(labels, node, fwd_key, fwd_out) + with torch.cuda.stream(node_stream): + stage_outputs_or_loss = [loss_fn(fwd_out["out_with_grad"][0], labels)] + output_grads = None + elif recv_pred is not None: + upstream_raw = self.buffers.task[recv_pred.uid] + self.buffers.release(recv_pred.uid) + if not isinstance(upstream_raw, (list, tuple)): + upstream_raw = [upstream_raw] + stage_outputs_or_loss = fwd_out["out_with_grad"] + output_grads = list(upstream_raw) + else: + stage_outputs_or_loss = fwd_out["out_with_grad"] + output_grads = None + + bwd_a2a_pred = next( + (p for p in node.data_preds if p.task_type == TaskType.BWD_A2A), None + ) + if ( + bwd_a2a_pred is not None + and not self._node_meta(node).get("compute_loss", False) + and recv_pred is None + ): + a2a_buf = self.buffers.task[bwd_a2a_pred.uid] + detached_outs = fwd_out["detached_outs"] + output_grads_full = [ + a2a_buf["inp_grads"][i] + for i, t in enumerate(detached_outs) + if t is not None and getattr(t, "requires_grad", False) + ] + pairs = [ + (o, g) for o, g in zip(stage_outputs_or_loss, output_grads_full) + if g is not None + ] + stage_outputs_or_loss = [p[0] for p in pairs] + output_grads = [p[1] for p in pairs] + self.buffers.release(bwd_a2a_pred.uid) + + input_values = fwd_out.get("inp_with_grad") or [] + weights = self.stages.bucket(ubid).weights() + + with torch.cuda.stream(node_stream): + dinputs, param_groups, output_backward_ctx = self.compute.bucket_backward_input( + stage_outputs_or_loss, output_grads, input_values, iter(weights) + ) + fwd_inputs_full = fwd_out.get("fwd_inputs") + if fwd_inputs_full is not None: + inp_grads_full = [ + t.grad if (t is not None and t.requires_grad) else None + for t in fwd_inputs_full + ] + else: + inp_grads_full = list(dinputs) + self.buffers.task[node.uid] = { + "inp_grads": inp_grads_full, + "param_groups": param_groups, + } + if output_backward_ctx is not None: + self.buffers.task[node.uid]["output_backward_ctx"] = output_backward_ctx + if dinputs: + self.buffers.task[node.uid]["send_output"] = list(dinputs) + fwd_out.clear() + del stage_outputs_or_loss, fwd_out + del self.buffers.task[fwd_key] + evt = torch.cuda.Event() + evt.record(node_stream) + comp_events[node.uid] = evt + self.events.backward[ubid] = evt + last_comp_event_by_stream[node_stream_id] = evt + if self._node_meta(node).get("zero_free_full_params_after"): + self.params.defer_free_full_params(ubid, evt) + + case TaskType.BWD_W: + self._wait_for_all_gather(node) + if self._node_meta(node).get("zero_alloc_full_grads_before"): + self.params.alloc_full_grads(ubid, node_stream) + bwdi_node = next(p for p in node.data_preds if p.task_type == TaskType.BWD_I) + bwdi_buf = self.buffers.task[bwdi_node.uid] + param_groups = bwdi_buf["param_groups"] + weights = self.stages.bucket(ubid).weights() + with torch.cuda.stream(node_stream): + output_backward_ctx = bwdi_buf.get("output_backward_ctx") + if output_backward_ctx is not None: + self.compute.backward_weight_from_outputs( + output_backward_ctx["stage_outputs_or_loss"], + output_backward_ctx["output_grads"], + iter(weights), + ) + else: + self.compute.bucket_backward_weight( + iter(weights), param_groups, ubid=ubid, mb_idx=mb_idx + ) + self.params.accumulate_zero_param_grads_to_flat(ubid, node_stream) + self.buffers.task[node.uid] = {} + self.buffers.release(bwdi_node.uid) + evt = torch.cuda.Event() + evt.record(node_stream) + comp_events[node.uid] = evt + self.events.backward[ubid] = evt + last_comp_event_by_stream[node_stream_id] = evt + if self._node_meta(node).get("zero_free_full_params_after"): + self.params.defer_free_full_params(ubid, evt) + + case TaskType.UPD: + self._update(node_stream, loss_buffer) + + case TaskType.ORDER_DUMMY: + pass + + self._rf_exit(rf) + self.runtime.nvtx_pop() + + def _update(self, stream: torch.cuda.Stream, loss_buffer: list): + self.params.drain_pending_frees() + if self.params.has_zero_shard_optimizers(): + self.params.step_zero_shard_optimizers(stream, self.events.reduce_scatter) + losses = loss_buffer + loss_buffer.clear() + torch.cuda.synchronize() + return losses + + for ar_evt in self.events.all_reduce.values(): + stream.wait_event(ar_evt) + + for ubid, bucket in self.stages.buckets.items(): + if bucket.optimizer is None: + continue + bwd_evt = self.events.backward.get(ubid) + if bwd_evt is not None: + stream.wait_event(bwd_evt) + + with torch.cuda.stream(stream): + bucket.optimizer.step() + + losses = loss_buffer + loss_buffer.clear() + + torch.cuda.synchronize() + + return { + "losses": losses, + } diff --git a/src/fx.py b/src/fx.py new file mode 100644 index 0000000..5ab63d4 --- /dev/null +++ b/src/fx.py @@ -0,0 +1,702 @@ +import torch +import torch.fx as fx +from collections import defaultdict +from dataclasses import dataclass +from typing import Any +import inspect +import importlib +import json +import operator + +from .state import create_logger, LOG_LEVEL + +logger = create_logger("fx", LOG_LEVEL) +PIPER_ANNOTATIONS_META_KEY = "piper_annotations" + + +def _encode_arg(a): + if isinstance(a, fx.Node): + return {"__node__": a.name} + if isinstance(a, torch.device): + return {"__device__": str(a)} + if isinstance(a, torch.dtype): + return {"__dtype__": str(a).replace("torch.", "")} + if isinstance(a, slice): + return { + "__slice__": True, + "start": _encode_arg(a.start), + "stop": _encode_arg(a.stop), + "step": _encode_arg(a.step), + } + if a is Ellipsis: + return {"__ellipsis__": True} + if isinstance(a, tuple): + return {"__tuple__": [_encode_arg(x) for x in a]} + if isinstance(a, list): + return [_encode_arg(x) for x in a] + if isinstance(a, dict): + return {k: _encode_arg(v) for k, v in a.items()} + return a + + +def _decode_arg(a, name_to_node): + if isinstance(a, dict): + if "__node__" in a: + return name_to_node[a["__node__"]] + if "__device__" in a: + return torch.device(a["__device__"]) + if "__dtype__" in a: + return getattr(torch, a["__dtype__"]) + if "__slice__" in a: + return slice( + _decode_arg(a["start"], name_to_node), + _decode_arg(a["stop"], name_to_node), + _decode_arg(a["step"], name_to_node), + ) + if "__ellipsis__" in a: + return Ellipsis + if "__tuple__" in a: + return tuple(_decode_arg(x, name_to_node) for x in a["__tuple__"]) + return {k: _decode_arg(v, name_to_node) for k, v in a.items()} + if isinstance(a, list): + return [_decode_arg(x, name_to_node) for x in a] + return a + + +def _is_op_overload(obj): + return ( + obj.__class__.__module__.startswith("torch._ops") + or obj.__class__.__name__.startswith("OpOverload") + ) + + +def _serialize_target(t): + if isinstance(t, str): + return {"kind": "string", "value": t} + + target_str = str(t) + if target_str.startswith("torch.ops."): + path = target_str[len("torch.ops."):] + return {"kind": "torch_op", "path": path} + + if getattr(t, "__module__", "") == "torch._VariableFunctionsClass": + public_name = t.__name__ + if hasattr(torch, public_name): + return {"kind": "py_func", "module": "torch", "qualname": public_name} + raise ValueError(f"No public torch alias for {t}") + + if _is_op_overload(t) or getattr(t, "__module__", "").startswith("torch._ops"): + return {"kind": "torch_op", "path": str(t)} + + if inspect.isfunction(t) or inspect.isbuiltin(t): + mod = inspect.getmodule(t) + + if mod is not None: + mod_name = mod.__name__ + if mod_name.startswith("torch.ops"): + target_str = str(t) + if target_str.startswith("torch.ops."): + path = target_str[len("torch.ops."):] + return {"kind": "torch_op", "path": path} + module_parts = mod_name.split(".") + if len(module_parts) >= 3: + namespace = module_parts[2] + path = f"{namespace}.{t.__name__}" + return {"kind": "torch_op", "path": path} + + if t.__name__ == "apply": + qualname = getattr(t, "__qualname__", "") + mod = inspect.getmodule(t) + if "." in qualname: + class_name = qualname.split(".")[0] + if mod and hasattr(mod, class_name): + func_class = getattr(mod, class_name) + if isinstance(func_class, type) and issubclass(func_class, torch.autograd.Function): + return { + "kind": "py_obj", + "module": func_class.__module__, + "qualname": func_class.__qualname__ + ".apply", + } + + import sys + + for module_obj in sys.modules.values(): + if module_obj is not None and hasattr(module_obj, class_name): + try: + func_class = getattr(module_obj, class_name) + if isinstance(func_class, type) and issubclass(func_class, torch.autograd.Function): + return { + "kind": "py_obj", + "module": func_class.__module__, + "qualname": func_class.__qualname__ + ".apply", + } + except (TypeError, AttributeError): + continue + + if class_name.startswith("AllToAll"): + func_class = globals().get(class_name) + if isinstance(func_class, type) and issubclass(func_class, torch.autograd.Function): + return { + "kind": "py_obj", + "module": func_class.__module__, + "qualname": func_class.__qualname__ + ".apply", + } + + mod = inspect.getmodule(t) + if mod is None: + raise ValueError(f"Cannot serialize function without module: {t}") + return {"kind": "py_func", "module": mod.__name__, "qualname": t.__name__} + + if inspect.isclass(t): + mod = t.__module__ + return {"kind": "py_obj", "module": mod, "qualname": t.__qualname__} + + if t in operator.__dict__.values(): + return {"kind": "py_func", "module": "operator", "qualname": t.__name__} + + mod = getattr(t, "__module__", None) + name = getattr(t, "__name__", None) + if mod and name: + return {"kind": "py_func", "module": mod, "qualname": name} + + raise NotImplementedError(f"Unsupported target type: {t} ({type(t)})") + + +def _resolve_qualname(mod, qualname): + obj = mod + for part in qualname.split("."): + obj = getattr(obj, part) + return obj + + +def _deserialize_target(payload): + kind = payload["kind"] + + if kind == "string": + return payload["value"] + + if kind == "py_func": + if payload["module"].startswith("torch.ops"): + module_parts = payload["module"].split(".") + if len(module_parts) >= 3: + namespace = module_parts[2] + path = f"{namespace}.{payload['qualname']}" + return _deserialize_target({"kind": "torch_op", "path": path}) + mod = importlib.import_module(payload["module"]) + return _resolve_qualname(mod, payload["qualname"]) + + if kind == "py_obj": + mod = importlib.import_module(payload["module"]) + try: + return _resolve_qualname(mod, payload["qualname"]) + except AttributeError as e: + raise AttributeError( + f"Failed to resolve qualname '{payload['qualname']}' " + f"in module '{payload['module']}': {e}" + ) from e + + if kind == "torch_op": + obj = torch.ops + path = payload["path"] + if path == "piper_artifact.bmm_experts": + # The Qwen artifact model defines this model-specific custom op. + for module_name in ("models.qwen3", "examples.models.qwen3"): + try: + importlib.import_module(module_name) + break + except ImportError: + pass + try: + for part in path.split("."): + obj = getattr(obj, part) + return obj + except AttributeError: + if path.startswith("higher_order."): + op_name = path.split(".", 1)[1] + op_to_module = { + "triton_kernel_wrapper_mutation": "torch._higher_order_ops.triton_kernel_wrap", + "triton_kernel_wrapper_functional": "torch._higher_order_ops.triton_kernel_wrap", + } + if op_name in op_to_module: + try: + importlib.import_module(op_to_module[op_name]) + obj = torch.ops + for part in path.split("."): + obj = getattr(obj, part) + return obj + except (ImportError, AttributeError): + pass + raise AttributeError( + f"Failed to resolve torch.ops path '{path}'. " + f"Available attributes: {[x for x in dir(obj) if not x.startswith('_')][:20]}" + ) + + raise NotImplementedError(f"Unknown target kind: {kind}") + + +def _serialize_graphmodule(gm: fx.GraphModule) -> str: + nodes = [] + for n in gm.graph.nodes: + nodes.append({ + "name": n.name, + "op": n.op, + "target": ( + _serialize_target(n.target) + if n.op in ("call_function", "call_method", "call_module", "get_attr") + else None + ), + "args": _encode_arg(n.args), + "kwargs": _encode_arg(n.kwargs), + }) + + submodules = {} + for name, module in gm.named_children(): + if isinstance(module, fx.GraphModule): + submodules[name] = _serialize_graphmodule(module) + + data = { + "nodes": nodes, + "state_dict": {k: v.detach().cpu().tolist() for k, v in gm.state_dict().items()}, + "param_devices": {k: str(v.device) for k, v in gm.state_dict().items()}, + "submodules": submodules, + } + serialized = json.dumps(data, ensure_ascii=False) + + del data + del nodes + + return serialized + + +def _unwrap_output_arg(decoded): + if isinstance(decoded, (tuple, list)) and len(decoded) == 1 and isinstance(decoded[0], (tuple, list)): + return decoded[0] + return decoded + + +def _deserialize_graphmodule(s: str) -> fx.GraphModule: + data = json.loads(s) + g = fx.Graph() + name_to_node = {} + + for n in data["nodes"]: + op = n["op"] + if op == "placeholder": + node = g.placeholder(n["name"]) + elif op == "output": + decoded = _decode_arg(n["args"], name_to_node) + node = g.output(_unwrap_output_arg(decoded)) + elif op == "call_function": + target = _deserialize_target(n["target"]) + args = _decode_arg(n["args"], name_to_node) + kwargs = _decode_arg(n["kwargs"], name_to_node) + node = g.call_function(target, tuple(args), kwargs) + elif op == "call_method": + target = _deserialize_target(n["target"]) + args = _decode_arg(n["args"], name_to_node) + kwargs = _decode_arg(n["kwargs"], name_to_node) + node = g.call_method(target, tuple(args), kwargs) + elif op == "call_module": + target = _deserialize_target(n["target"]) + args = _decode_arg(n["args"], name_to_node) + kwargs = _decode_arg(n["kwargs"], name_to_node) + node = g.call_module(target, tuple(args), kwargs) + elif op == "get_attr": + target = _deserialize_target(n["target"]) + node = g.get_attr(target) + else: + raise NotImplementedError(f"op {op} not handled") + name_to_node[n["name"]] = node + + root_module = torch.nn.Module() + submodules = data.get("submodules", {}) + for module_name, serialized_submodule in submodules.items(): + submodule_gm = _deserialize_graphmodule(serialized_submodule) + parts = module_name.split(".") + if len(parts) == 1: + root_module.add_module(module_name, submodule_gm) + else: + current = root_module + for part in parts[:-1]: + if not hasattr(current, part): + current.add_module(part, torch.nn.Module()) + current = getattr(current, part) + current.add_module(parts[-1], submodule_gm) + + gm = fx.GraphModule(root_module, g) + state = {k: torch.tensor(v) for k, v in data["state_dict"].items()} + gm.load_state_dict(state, strict=False) + return gm + + +@dataclass +class AnnotationSegment: + segment_id: int + stage_id: int + tag: dict[str, int] + gm: fx.GraphModule + input_idxs: list[int] + param_idxs: list[int] + graphargs: list[Any] + placeholders: list[fx.Node] + a2a_boundary_after: dict[str, Any] | None = None + + +def _inject_piper_annotation(node: fx.Node, name: str, index: int, uid: int) -> None: + custom = node.meta.setdefault("custom", {}) + annotation = {"name": name, "index": int(index), "uid": int(uid)} + custom[PIPER_ANNOTATIONS_META_KEY] = (annotation,) + custom["name"] = name + custom["index"] = int(index) + + +def _meta_tensor_like(example, *, requires_grad: bool, as_parameter: bool): + sym_int_type = getattr(torch, "SymInt", ()) + sym_float_type = getattr(torch, "SymFloat", ()) + sym_bool_type = getattr(torch, "SymBool", ()) + + if isinstance(example, torch.Tensor): + t = torch.empty(example.shape, dtype=example.dtype, device="meta") + t.requires_grad_(requires_grad) + elif isinstance(example, (sym_int_type, int)): + # Symbolic/python integers do not have .shape/.dtype; represent them as + # singleton meta tensors so downstream graph-arg handling remains uniform. + t = torch.empty((1,), dtype=torch.int64, device="meta") + elif isinstance(example, (sym_float_type, float)): + t = torch.empty((1,), dtype=torch.float32, device="meta") + elif isinstance(example, (sym_bool_type, bool)): + t = torch.empty((1,), dtype=torch.bool, device="meta") + else: + # Fallback for unknown scalar-like values. + t = torch.empty((1,), dtype=torch.float32, device="meta") + + if as_parameter: + return torch.nn.Parameter(t, requires_grad=requires_grad) + return t + + +def _iter_arg_nodes(arg: Any) -> list[fx.Node]: + nodes: list[fx.Node] = [] + + def _walk(v: Any) -> None: + if isinstance(v, fx.Node): + nodes.append(v) + elif isinstance(v, (list, tuple)): + for item in v: + _walk(item) + elif isinstance(v, dict): + for item in v.values(): + _walk(item) + + _walk(arg) + return nodes + + +def _example_value(node: fx.Node) -> Any: + return node.meta.get("example_value", node.meta.get("val")) + + +def _placeholder_is_runtime_input(node: fx.Node) -> bool: + ex = _example_value(node) + if not isinstance(ex, torch.Tensor): + return False + if isinstance(ex, torch.nn.Parameter): + return False + if bool(getattr(ex, "requires_grad", False)): + return False + if "self" in node.name or "grapharg" in node.meta: + return False + return True + + +def _grapharg_for_placeholder(node: fx.Node) -> Any: + ex = _example_value(node) + requires_grad = bool(getattr(ex, "requires_grad", False)) + as_parameter = isinstance(ex, torch.nn.Parameter) or requires_grad + return _meta_tensor_like(ex, requires_grad=requires_grad, as_parameter=as_parameter) + + +def _grapharg_for_cross_value(node: fx.Node) -> Any: + ex = _example_value(node) + requires_grad = bool(getattr(ex, "requires_grad", False)) + return _meta_tensor_like(ex, requires_grad=requires_grad, as_parameter=False) + + +def _validate_and_get_annotation_stack(node: fx.Node) -> tuple[dict[str, int], ...]: + custom = node.meta.get("custom") + if not isinstance(custom, dict): + return () + + if PIPER_ANNOTATIONS_META_KEY not in custom: + return () + + raw_stack = custom[PIPER_ANNOTATIONS_META_KEY] + if not isinstance(raw_stack, (list, tuple)) or not raw_stack: + raise ValueError( + "Piper annotation metadata must be a non-empty list/tuple of " + f"{{'name': str, 'index': int, 'uid': int}} entries. " + f"Node {node.name!r} has {raw_stack!r}." + ) + + stack: list[dict[str, int]] = [] + for annotation in raw_stack: + if not isinstance(annotation, dict): + raise ValueError( + f"Piper annotation entry on node {node.name!r} must be a dict, " + f"got {annotation!r}." + ) + if set(annotation) - {"name", "index", "uid"}: + raise ValueError( + "Piper annotation entries may only contain 'name', 'index', and 'uid'. " + f"Node {node.name!r} has {annotation!r}." + ) + name = annotation.get("name") + index = annotation.get("index") + uid = annotation.get("uid") + if not isinstance(name, str) or not name: + raise ValueError( + f"Piper annotation 'name' must be a non-empty string on node {node.name!r}: " + f"{annotation!r}." + ) + if not isinstance(index, int): + raise ValueError( + f"Piper annotation 'index' must be an int on node {node.name!r}: " + f"{annotation!r}." + ) + if not isinstance(uid, int): + raise ValueError( + f"Piper annotation 'uid' must be an int on node {node.name!r}: " + f"{annotation!r}." + ) + stack.append({"name": name, "index": int(index), "uid": int(uid)}) + return tuple(stack) + + +def _annotation_stack_key(stack: tuple[dict[str, int], ...]) -> tuple[tuple[str, int, int], ...]: + return tuple((a["name"], int(a["index"]), int(a["uid"])) for a in stack) + + +def _tag_from_stack(stack: tuple[dict[str, int], ...]) -> dict[str, int]: + tag: dict[str, int] = {} + for annotation in stack: + tag[annotation["name"]] = int(annotation["index"]) + return tag + + +def _select_boundary_tensor_idx(nodes: list[fx.Node]) -> int | None: + best_idx = None + best_score = -1 + for idx, node in enumerate(nodes): + ex = _example_value(node) + if not isinstance(ex, torch.Tensor): + continue + score = 1 + if ex.is_floating_point(): + score += 2 + if bool(getattr(ex, "requires_grad", False)): + score += 1 + if score > best_score: + best_score = score + best_idx = idx + if best_idx is not None: + return best_idx + return 0 if nodes else None + + +def split_gm_by_annotations(gm: fx.GraphModule) -> tuple[fx.GraphModule, list[AnnotationSegment]]: + """Split a GraphModule into contiguous subgraphs by Piper annotation stack. + + Piper annotations are carried in ``node.meta['custom']['piper_annotations']``. + The full stack, including nested annotations, is part of the split key. The + schedule-facing tag for each segment is the stack projected to + ``{annotation_name: auto_index}``. + """ + nodes = list(gm.graph.nodes) + compute_nodes = [ + node for node in nodes + if node.op not in ("placeholder", "get_attr", "output") + ] + + node_stacks: dict[fx.Node, tuple[dict[str, int], ...]] = {} + for node in nodes: + stack = _validate_and_get_annotation_stack(node) + if node in compute_nodes: + node_stacks[node] = stack + + annotated_compute = [node for node in compute_nodes if node_stacks.get(node)] + if not annotated_compute: + return gm, [] + + unannotated = [node.name for node in compute_nodes if not node_stacks.get(node)] + if unannotated: + raise ValueError( + "Found compute nodes without Piper annotations in a graph that uses " + f"piper.annotate: {unannotated[:16]}. Wrap all model compute in " + "piper.annotate(...)." + ) + + segment_keys: list[tuple[tuple[str, int, int], ...]] = [] + segment_stacks: list[tuple[dict[str, int], ...]] = [] + node_seg: dict[fx.Node, int] = {} + current_key: tuple[tuple[str, int, int], ...] | None = None + for node in compute_nodes: + stack = node_stacks[node] + key = _annotation_stack_key(stack) + if key != current_key: + segment_keys.append(key) + segment_stacks.append(stack) + current_key = key + node_seg[node] = len(segment_keys) - 1 + + n_segs = len(segment_stacks) + if n_segs == 0: + return gm, [] + + runtime_placeholders = { + node for node in nodes + if node.op == "placeholder" and _placeholder_is_runtime_input(node) + } + grapharg_placeholders = { + node for node in nodes + if node.op == "placeholder" and node not in runtime_placeholders + } + + value_seg: dict[fx.Node, int] = dict(node_seg) + for node in runtime_placeholders: + value_seg[node] = 0 + + def _consumer_seg(user: fx.Node) -> int | None: + if user.op == "output": + return n_segs - 1 + return node_seg.get(user) + + node_max_user_seg: dict[fx.Node, int] = {} + for node, seg in value_seg.items(): + user_segs = [ + user_seg for user in node.users + for user_seg in [_consumer_seg(user)] + if user_seg is not None + ] + node_max_user_seg[node] = max(user_segs) if user_segs else seg + + seg_cross_in: list[list[fx.Node]] = [[] for _ in range(n_segs)] + for node in nodes: + if node not in value_seg: + continue + seg = value_seg[node] + max_seg = node_max_user_seg[node] + for target_seg in range(seg + 1, max_seg + 1): + seg_cross_in[target_seg].append(node) + + segments: list[AnnotationSegment] = [] + output_node = next((node for node in nodes if node.op == "output"), None) + + for seg in range(n_segs): + sub_g = fx.Graph() + remap: dict[fx.Node, fx.Node] = {} + new_input_idxs: list[int] = [] + new_param_idxs: list[int] = [] + new_graphargs: list[Any] = [] + pos = 0 + + def _add_placeholder(node: fx.Node, name: str, *, is_runtime_input: bool) -> None: + nonlocal pos + if node in remap: + return + new_ph = sub_g.placeholder(name) + new_ph.type = node.type + new_ph.meta.update(node.meta) + remap[node] = new_ph + if is_runtime_input: + new_input_idxs.append(pos) + new_graphargs.append(_grapharg_for_cross_value(node)) + else: + new_param_idxs.append(pos) + new_graphargs.append(_grapharg_for_placeholder(node)) + pos += 1 + + if seg == 0: + for node in nodes: + if node in runtime_placeholders: + _add_placeholder(node, node.name, is_runtime_input=True) + else: + for node in seg_cross_in[seg]: + _add_placeholder(node, f"_xseg_{node.name}", is_runtime_input=True) + + seg_compute_nodes = [ + node for node in nodes + if node.op not in ("placeholder", "get_attr", "output") + and node_seg[node] == seg + ] + needed_grapharg_placeholders: set[fx.Node] = set() + needed_getattrs: set[fx.Node] = set() + for node in seg_compute_nodes: + for dep in [*node.all_input_nodes, *_iter_arg_nodes(node.kwargs)]: + if dep.op == "placeholder" and dep in grapharg_placeholders: + needed_grapharg_placeholders.add(dep) + elif dep.op == "get_attr": + needed_getattrs.add(dep) + + if output_node is not None and seg == n_segs - 1: + for dep in _iter_arg_nodes(output_node.args): + if dep.op == "placeholder" and dep in grapharg_placeholders: + needed_grapharg_placeholders.add(dep) + elif dep.op == "get_attr": + needed_getattrs.add(dep) + + for node in nodes: + if node in needed_grapharg_placeholders: + _add_placeholder(node, node.name, is_runtime_input=False) + + for node in nodes: + if node.op == "get_attr" and node in needed_getattrs and node not in remap: + new_ga = sub_g.get_attr(node.target) + new_ga.type = node.type + new_ga.meta.update(node.meta) + remap[node] = new_ga + + for node in seg_compute_nodes: + missing = [dep.name for dep in node.all_input_nodes if dep not in remap and node_seg.get(dep) != seg] + if missing: + raise ValueError( + f"Cannot split annotated graph at segment {seg}; node {node.name!r} " + f"has unmapped external dependencies {missing}." + ) + remap[node] = sub_g.node_copy(node, arg_transform=lambda x, r=remap: r[x]) + + if seg == n_segs - 1: + if output_node is None: + sub_g.output(()) + else: + sub_g.output(fx.map_arg(output_node.args[0], lambda x: remap[x])) + boundary_after = None + else: + out_nodes = [remap[node] for node in seg_cross_in[seg + 1]] + sub_g.output(tuple(out_nodes) if len(out_nodes) != 1 else out_nodes[0]) + boundary_after = { + "tensor_idx": _select_boundary_tensor_idx(seg_cross_in[seg + 1]), + "reshape_input": None, + "reshape_output": None, + "from_tag": _tag_from_stack(segment_stacks[seg]), + "to_tag": _tag_from_stack(segment_stacks[seg + 1]), + } + + sub_g.lint() + seg_gm = fx.GraphModule(gm, sub_g) + placeholders = [node for node in seg_gm.graph.nodes if node.op == "placeholder"] + tag = _tag_from_stack(segment_stacks[seg]) + segments.append( + AnnotationSegment( + segment_id=seg, + stage_id=int(tag.get("PP", seg)), + tag=tag, + gm=seg_gm, + input_idxs=new_input_idxs, + param_idxs=new_param_idxs, + graphargs=new_graphargs, + placeholders=placeholders, + a2a_boundary_after=boundary_after, + ) + ) + + return gm, segments diff --git a/src/ordering.py b/src/ordering.py new file mode 100644 index 0000000..f61c4e4 --- /dev/null +++ b/src/ordering.py @@ -0,0 +1,182 @@ +from .dag import ( + TrainingDAG, + TrainingDAGEdge, + _has_path, + _topological_levels, + _topological_order, +) + +_DEFAULT_STREAM = "default_stream" +_CRITICAL_PATH_COMM_KINDS = { + "ALL_GATHER_COMM", + "A2A_COMM", +} +_REDUCTION_COMM_KINDS = { + "REDUCE_COMM", + "REDUCE_SCATTER_COMM", +} + + +def _serial_topological_order( + dag: TrainingDAG, + topo_levels: dict[str, int] | None = None, +) -> list[str]: + """Serialize topological levels into a deterministic dispatch order. + + Nodes with lower topological levels always come first. Within a level, + priority order is: SEND > critical-path comm > reduction comm > compute/other > RECV. + """ + if topo_levels is None: + topo_levels = _topological_levels(dag) + + base_order = _topological_order(dag) + topo_idx = {uid: i for i, uid in enumerate(base_order)} + + def node_priority(uid: str) -> int: + kind = dag.nodes[uid].node_kind + if kind == "SEND_COMM": + return 0 + if kind in _CRITICAL_PATH_COMM_KINDS: + return 1 + if kind in _REDUCTION_COMM_KINDS: + return 2 + if kind == "RECV_COMM": + return 4 + return 3 + + return sorted( + dag.nodes, + key=lambda uid: (topo_levels[uid], node_priority(uid), topo_idx[uid]), + ) + + +def _resolve_default_stream_order(dag: TrainingDAG) -> None: + """Create a total ordering over default-stream COMPUTE nodes. + + Whenever multiple default-stream compute nodes share a topological level, + chain them with temporal edges in descending order of downstream + dependencies (more downstream nodes -> earlier in the chain). Downstream + count is the size of each compute node's transitive successor set, which in + a well-formed training DAG terminates at UPD. + + The new edges shift topological levels, so after each chain insertion we + recompute levels and rescan from the earliest level. The pass terminates + when every default-stream compute node sits at a unique level. + """ + def _is_default_compute(uid: str) -> bool: + node = dag.nodes[uid] + return node.stream == _DEFAULT_STREAM and node.node_kind == "COMPUTE" + + compute_uids = [uid for uid in dag.nodes if _is_default_compute(uid)] + if len(compute_uids) < 2: + return + + def _downstream_count(uid: str) -> int: + seen: set[str] = set() + stack = list(dag.succs.get(uid, set())) + while stack: + v = stack.pop() + if v in seen: + continue + seen.add(v) + stack.extend(dag.succs.get(v, set())) + return len(seen) + + while True: + topo_levels = _topological_levels(dag) + by_level: dict[int, list[str]] = {} + for uid in compute_uids: + by_level.setdefault(topo_levels[uid], []).append(uid) + + conflict_level: int | None = None + for level in sorted(by_level): + if len(by_level[level]) > 1: + conflict_level = level + break + if conflict_level is None: + return + + # Same-level nodes have no path between them, otherwise their levels + # would differ. Order by downstream count desc; break ties by uid. + group = by_level[conflict_level] + ordered = sorted(group, key=lambda u: (-_downstream_count(u), u)) + for src, dst in zip(ordered, ordered[1:]): + dag.add_edge( + TrainingDAGEdge( + src_uid=src, + dst_uid=dst, + dep_kind="temporal", + tensor_name=None, + ) + ) + + +def resolve_total_order_per_stream(dag: TrainingDAG) -> None: + """Serialize the default stream, then strictly chain non-default streams. + + The default-stream pass runs first so the per-non-default-stream anchors + below see the post-serialization topological order of default-stream nodes. + + Non-default stream nodes are ordered by the topological order of the + default-stream nodes they directly depend on or are depended on by. + """ + _resolve_default_stream_order(dag) + + topo = _topological_order(dag) + topo_idx = {uid: i for i, uid in enumerate(topo)} + topo_levels = _topological_levels(dag) + default_uids = [ + uid for uid in topo + if dag.nodes[uid].stream == _DEFAULT_STREAM + ] + streams = sorted({ + n.stream for n in dag.nodes.values() + if n.stream != _DEFAULT_STREAM + }) + if not default_uids or not streams: + return + + for stream in streams: + associated_by_default: dict[str, list[str]] = {uid: [] for uid in default_uids} + default_anchors_by_stream_uid: dict[str, list[str]] = {} + for edge in dag.edges: + src = dag.nodes[edge.src_uid] + dst = dag.nodes[edge.dst_uid] + if src.stream == _DEFAULT_STREAM and dst.stream == stream: + default_anchors_by_stream_uid.setdefault(edge.dst_uid, []).append(edge.src_uid) + elif src.stream == stream and dst.stream == _DEFAULT_STREAM: + default_anchors_by_stream_uid.setdefault(edge.src_uid, []).append(edge.dst_uid) + + for stream_uid in { + uid for uid, node in dag.nodes.items() + if node.stream == stream + }: + default_anchors = default_anchors_by_stream_uid.get(stream_uid, []) + if not default_anchors: + continue + anchor_uid = min(default_anchors, key=lambda uid: (topo_levels[uid], topo_idx[uid])) + associated_by_default.setdefault(anchor_uid, []).append(stream_uid) + + current_uid: str | None = None + for default_uid in default_uids: + stream_uids = sorted( + set(associated_by_default.get(default_uid, [])), + key=lambda u: (topo_levels[u], topo_idx[u]), + ) + for next_uid in stream_uids: + if current_uid is not None and current_uid != next_uid: + if _has_path(dag, next_uid, current_uid): + raise ValueError( + "resolve_total_order_per_stream would create a cycle while ordering " + f"stream={stream}: {current_uid} -> {next_uid}" + ) + if not _has_path(dag, current_uid, next_uid): + dag.add_edge( + TrainingDAGEdge( + src_uid=current_uid, + dst_uid=next_uid, + dep_kind="temporal", + tensor_name=None, + ) + ) + current_uid = next_uid diff --git a/src/piper.py b/src/piper.py index 1c5eea3..92ace4a 100644 --- a/src/piper.py +++ b/src/piper.py @@ -1,183 +1,250 @@ -from . import piper_patches +import time +from contextlib import contextmanager +from typing import Iterator import ray -import torch -import os +import torch.fx as fx from torch._dynamo.backends.registry import register_backend -from torch._dynamo.decorators import _disallow_in_graph_helper -from .piper_utils import RemoteTensor, serialize_graphmodule, piper_metadata, create_logger, print_backward_graph, LOG_LEVEL -from .piper_graph_transform import split_gm_by_experts -from .piper_actor import get_actor +from .fx import PIPER_ANNOTATIONS_META_KEY, split_gm_by_annotations +from .dag import ( + TrainingDAG, + TrainingDAGEdge, + TrainingDAGNode, + build_training_dag, +) +from .directives import ( + _apply_order_directive, + _apply_split_backward_stencil, + _apply_split_directive, + _bucket_matched_fwd_nodes, + _parse_order_directive, + _validate_schedule_tags_exist, + _validate_split_backward_order_stencil, + apply_schedule_directives, +) +from .ordering import ( + _serial_topological_order, + resolve_total_order_per_stream, +) +from .state import LOG_LEVEL, create_logger, piper_metadata +from .visualization import ( + log_training_dag_dependencies, + print_training_dag_order, + render_training_dag, +) +from .zero import _add_inter_chain_temporal_edges, _prune_zero_lifetime_metadata logger = create_logger("piper_backend", LOG_LEVEL) +_ANNOTATION_STACK: list[dict[str, int]] = [] +_ANNOTATION_COUNTS: dict[str, int] = {} +_ANNOTATION_UID = 0 -@torch.compiler.disable -def distributed_stage(stage_id, actor_id=None): - """ - Annotation for stage boundaries, causes torch.compile graph break - and sets metadata appropriately at compile time - """ - dp_rank = int(os.environ['PIPER_DP_RANK']) - world_size = int(os.environ['PIPER_WORLD_SIZE']) - dp_degree = int(os.environ['PIPER_DP_DEGREE']) - pp_degree = int(os.environ['PIPER_PP_DEGREE']) +def _reset_annotation_state() -> None: + global _ANNOTATION_UID + _ANNOTATION_STACK.clear() + _ANNOTATION_COUNTS.clear() + _ANNOTATION_UID = 0 - if actor_id is None: - actor_id = stage_id - piper_metadata.current_stage = stage_id - piper_metadata.current_actor = actor_id - piper_metadata.first_graph_of_stage = True +@contextmanager +def annotate(name: str) -> Iterator[dict[str, int]]: + """Annotate traced model code with a Piper schedule tag. + + Piper assigns the integer index for each tag name automatically in the + order annotation scopes are entered during tracing. + """ + if not isinstance(name, str) or not name: + raise ValueError(f"piper.annotate requires a non-empty string tag name, got {name!r}") + + global _ANNOTATION_UID + index = _ANNOTATION_COUNTS.get(name, 0) + _ANNOTATION_COUNTS[name] = index + 1 + uid = _ANNOTATION_UID + _ANNOTATION_UID += 1 + + annotation = {"name": name, "index": int(index), "uid": int(uid)} + _ANNOTATION_STACK.append(annotation) + fx_metadata_stack = tuple(dict(item) for item in _ANNOTATION_STACK) + try: + with fx.traceback.annotate({ + PIPER_ANNOTATIONS_META_KEY: fx_metadata_stack, + "name": name, + "index": int(index), + }): + yield annotation + finally: + popped = _ANNOTATION_STACK.pop() + if popped is not annotation: + raise RuntimeError("piper.annotate stack corrupted during tracing") + + + +def _split_global_training_dag_by_pp_rank(training_dag: TrainingDAG) -> list[TrainingDAG]: + """Split the global DAG into per-device-set disconnected DAGs.""" + # SEND/RECV pairs intentionally have no edge between them, so cross-rank + # placement dependencies separate into disconnected local components here. + undirected: dict[str, set[str]] = {uid: set() for uid in training_dag.nodes} + for e in training_dag.edges: + if e.src_uid in undirected and e.dst_uid in undirected: + undirected[e.src_uid].add(e.dst_uid) + undirected[e.dst_uid].add(e.src_uid) + + components: list[set[str]] = [] + seen: set[str] = set() + for uid in training_dag.nodes: + if uid in seen: + continue + comp: set[str] = set() + stack = [uid] + seen.add(uid) + while stack: + cur = stack.pop() + comp.add(cur) + for nxt in undirected.get(cur, set()): + if nxt not in seen: + seen.add(nxt) + stack.append(nxt) + components.append(comp) + + # Validate each component is device-homogeneous and component device-sets are distinct. + comp_device_keys: list[tuple[int, ...]] = [] + for ci, comp in enumerate(components): + device_keys = { + tuple(sorted(node.device)) for uid in comp for node in [training_dag.nodes[uid]] if node.device is not None + } + if not device_keys: + raise ValueError(f"component[{ci}] has no device assignment after P2P split") + if len(device_keys) != 1: + raise ValueError( + f"component[{ci}] is not device-homogeneous; device sets present: {sorted(device_keys)}" + ) + comp_device_keys.append(next(iter(device_keys))) + if len(set(comp_device_keys)) != len(comp_device_keys): + raise ValueError( + f"expected distinct device sets across split components, got {comp_device_keys}" + ) + + # Materialize each component as a standalone TrainingDAG. + subdags: list[TrainingDAG] = [] + for comp in components: + sub = TrainingDAG() + for uid in comp: + n = training_dag.nodes[uid] + sub.add_node( + TrainingDAGNode( + uid=n.uid, + node_kind=n.node_kind, + compute_subkind=n.compute_subkind, + tag=dict(n.tag), + device=(None if n.device is None else list(n.device)), + stream=n.stream, + node_meta=dict(n.node_meta), + ) + ) + for e in training_dag.edges: + if e.src_uid in comp and e.dst_uid in comp: + sub.add_edge( + TrainingDAGEdge( + src_uid=e.src_uid, + dst_uid=e.dst_uid, + dep_kind=e.dep_kind, + tensor_name=e.tensor_name, + ) + ) + subdags.append(sub) + + def _dag_device_key(d: TrainingDAG) -> tuple[int, ...]: + keys = {tuple(sorted(n.device)) for n in d.nodes.values() if n.device is not None} + if len(keys) != 1: + raise ValueError(f"sub-DAG should have exactly one device key, got {keys}") + return next(iter(keys)) + + subdags.sort(key=_dag_device_key) + return subdags @register_backend def piper(gm, example_inputs, **kwargs): - """ - torch.compile backend loads the graph module on - a Ray actor and returns a callback that remotely - runs the graph module. - """ - logger.debug(f"Compiling subgraph {id(gm)}") - - if not piper_metadata.currently_compiling: - gm.print_readable() - assert False, "Piper backend called outside of compilation" - - # Distribute expert submodules to actors if there are expert annotations - stage_id = piper_metadata.current_stage - pp_degree = int(os.environ['PIPER_PP_DEGREE']) - original_gm = gm - gm = split_gm_by_experts(gm, stage_id, pp_degree) - - # For the top-level graph, log which arguments are input tensors - # vs parameter tensors and make sure all example inputs are serializable - - placeholders = gm.graph.find_nodes(op="placeholder") - graphargs = [node.meta["grapharg"] for node in placeholders] - - # make sure example inputs are serializable by turning symbolic - # ints and fake tensors into concrete values - serializable_examples = [] - input_idxs = [] - param_idxs = [] - for i, (arg, ex) in enumerate(zip(graphargs, example_inputs)): - # save indices of input tensors and model parameters - if 'self' not in str(arg): - input_idxs.append(i) - else: - param_idxs.append(i) - # convert symbolic ints and fake tensors to concrete values - if isinstance(ex, torch.SymInt): - serializable_examples.append(int(ex)) - elif isinstance(ex, torch._subclasses.fake_tensor.FakeTensor): - new = torch.full( - ex.shape, - 0, - dtype=ex.dtype, - device=ex.device, - layout=ex.layout, - requires_grad=ex.requires_grad, - ) - serializable_examples.append(new) - else: - serializable_examples.append(ex) - - # serialize the fx.Graph - payload = serialize_graphmodule(gm) - - # send the fx.Graph and model attributes to the actor - stage_id = piper_metadata.current_stage - actor_id = piper_metadata.current_actor - actor = get_actor(actor_id) - - dp_rank = int(os.environ['PIPER_DP_RANK']) - dp_degree = int(os.environ['PIPER_DP_DEGREE']) - global_rank = dp_rank * dp_degree + actor_id - - ray.get( - actor.load_graph.remote( - stage_id, - payload, - torch._dynamo.backends.debugging.eager, - serializable_examples, - input_idxs, + """TrainingDAG backend: split by Piper annotations and lower schedule directives.""" + del example_inputs, kwargs + + schedule_info = getattr(piper_metadata, "schedule_info", {}) or {} + schedule_directives = getattr(piper_metadata, "schedule_directives", None) + _top_level_gm, annotation_segments = split_gm_by_annotations(gm) + + if not annotation_segments: + raise ValueError( + "No Piper annotations found in the traced graph. Wrap model compute " + "with src.piper.annotate(...) before compiling with Piper." ) + + # Build and store the new directed DAG representation for later scheduling transforms. + training_dag = build_training_dag(annotation_segments) + _validate_schedule_tags_exist(training_dag, schedule_directives) + apply_schedule_directives( + training_dag, + schedule_directives, + ) + piper_metadata.training_dag = training_dag + per_pp_training_dags = _split_global_training_dag_by_pp_rank(training_dag) + artifact_dir = getattr(piper_metadata, "artifact_dir", "out") + for i, subdag in enumerate(per_pp_training_dags): + zero_chains = _prune_zero_lifetime_metadata(subdag) + resolve_total_order_per_stream(subdag) + _add_inter_chain_temporal_edges(subdag, zero_chains) + if getattr(piper_metadata, "visualize_dag", False): + log_training_dag_dependencies(subdag) + print_training_dag_order(subdag, label=f"pp{i}", rank=i, out_dir=artifact_dir) + render_training_dag(subdag, output_path=f"{artifact_dir}/training_dag_pp{i}") + piper_metadata.per_pp_training_dags = per_pp_training_dags + + logger.info( + "piper: built TrainingDAG with %d nodes and %d edges, split into %d per-PP DAG(s)", + len(training_dag.nodes), + len(training_dag.edges), + len(per_pp_training_dags), ) - # Get fake tensor representations of the graph output(s) - def symint_to_int(x): - return int(x) if isinstance(x, torch.SymInt) else x - def int_to_tensor(x): - return torch.tensor(x) if isinstance(x, int) else x - example_inputs = list(map(symint_to_int, serializable_examples)) - fakes = original_gm(*example_inputs) - fakes = list(map(int_to_tensor, fakes)) - - # wait for a signal to run this graph if it's the first graph of the stage - first_graph_of_stage = piper_metadata.first_graph_of_stage - if first_graph_of_stage: - piper_metadata.first_graph_of_stage = False - - # return a wrapper function that runs the fx.Graph on the actor and - # returns remote futures for each graph output - def run_remote_subgraph(*args): - - logger.debug(f"Running subgraph {id(gm)}") - - from .piper_utils import events_tls - - mb_idx = events_tls.mb_idx - - # wait for a signal to run the partial graph - if first_graph_of_stage: - logger.debug(f"Thread {mb_idx} global rank {global_rank} waiting for stage {stage_id}") - events_tls.events[stage_id].wait() - logger.debug(f"Thread {mb_idx} global rank {global_rank} running stage {stage_id}") - - # Mutex ensures that only one thread submits a task to this actor at a time - logger.debug(f"Thread {mb_idx} global rank {global_rank} waiting for actor mutex {actor_id}") - with events_tls.actor_mutexes[actor_id]: - logger.debug(f"Thread {mb_idx} global rank {global_rank} got actor mutex {actor_id}") - - # clear the event for the next stage - if first_graph_of_stage: - events_tls.events[stage_id].clear() - - # ignore model parameter arguments (stored on the actor) - input_tensors_only = [] - for i in input_idxs: - input_tensors_only.append(args[i]) - args = input_tensors_only - - # track stage dependencies - for arg in args: - if isinstance(arg, RemoteTensor): - prev_stage = arg.get_stage_id() - if prev_stage != stage_id: - piper_metadata.dag.add((prev_stage, stage_id)) - - # get Ray ObjectRefs from RemoteTensors - def unwrap(x): - return x.get_ref() if isinstance(x, RemoteTensor) else x - args = list(map(unwrap, args)) - - if piper_metadata.currently_compiling: - # dispatch task without nccl transport - refs = actor.forward_cpu.options(num_returns=len(fakes)).remote(stage_id, mb_idx, *args) - else: - # dispatch with nccl transport - refs = actor.forward.options(num_returns=len(fakes)).remote(stage_id, mb_idx, *args) - - # wrap the remote futures with RemoteTensor - # if piper_metadata.currently_compiling: - # return [t.to('cpu') for t in ray.get(refs)] - - if isinstance(refs, list): - assert len(fakes) == len(refs) - return [RemoteTensor(fake, ref, stage_id) for fake, ref in zip(fakes, refs)] - else: - assert len(fakes) == 1 - return [RemoteTensor(fakes[0], refs, stage_id)] - return run_remote_subgraph \ No newline at end of file + def callback(*args, _gm=gm): + logger.warning( + "piper compiled callback invoked directly; running local graph execution" + ) + return _gm(*args) + + return callback + + +def piper_exec_dag(loss_fn, log_stats: bool = False) -> list: + """Execute one training step using the loaded per-rank TrainingDAG.""" + actors = piper_metadata.actors + run_refs = [ + actor.run_dag.remote(loss_fn=loss_fn) + for actor in actors.values() + ] + t0 = time.perf_counter() + results = ray.get(run_refs) + step_time = time.perf_counter() - t0 + + if log_stats: + _log_step_stats(step_time, log_stats, actors) + + losses = [] + for result in results: + if isinstance(result, dict): + losses.extend(result.get("losses", [])) + elif result: + losses.extend(result) + return losses + + +def _log_step_stats(step_time: float, log_memory: bool, actors: dict) -> None: + """Log throughput, MFU, and optionally per-rank peak GPU memory.""" + stats = [f"step_time={step_time:.3f}s"] + + tokens = getattr(piper_metadata, "tokens_per_step", None) + if tokens is not None: + stats.append(f"throughput={tokens / step_time:.1f} tok/s") + + logger.info(" ".join(stats)) diff --git a/src/piper_actor.py b/src/piper_actor.py deleted file mode 100644 index a91f7a4..0000000 --- a/src/piper_actor.py +++ /dev/null @@ -1,698 +0,0 @@ -import ray -import torch -import logging -import os -import time -import gc -from torch._guards import CompileId -from torch.nn import Parameter -from collections import defaultdict -import torch.distributed as dist -import threading -from typing import Callable - -from .piper_utils import deserialize_graphmodule, create_logger, RemoteTensor, print_backward_graph, LOG_LEVEL - -CLEANUP_MEMORY = False -torch.autograd.set_detect_anomaly(True) - -logger = create_logger("piper_actor", LOG_LEVEL) - -def create_actors(num_actors, optim_class): - dp_rank = int(os.environ['PIPER_DP_RANK']) - world_size = int(os.environ['PIPER_WORLD_SIZE']) - dp_degree = int(os.environ['PIPER_DP_DEGREE']) - pp_degree = int(os.environ['PIPER_PP_DEGREE']) - - from .piper_utils import piper_metadata - for actor_id in range(num_actors): - global_rank = dp_rank * dp_degree + actor_id - actor = PiperActor.options(num_gpus=0.9, max_concurrency=2).remote(actor_id, optim_class, world_size, dp_rank=dp_rank, dp_degree=dp_degree, pp_degree=pp_degree) - piper_metadata.actors[actor_id] = actor - - ray.get([actor.load_actor_handles.remote(piper_metadata.actors) for actor in piper_metadata.actors.values()]) - -def get_actor(actor_id): - from .piper_utils import piper_metadata - return piper_metadata.actors[actor_id] - -class ExpertRayFunction(torch.autograd.Function): - """ - Custom autograd Function that dispatches expert calls to Ray actors. - """ - - @staticmethod - def forward(ctx, global_expert_id: int, local_expert_id: int, batch_idx: int, pp_degree: int, *args): - """ - Dispatch a forward call to the expert `global_expert_id` - on the actor `local_expert_id % pp_degree`. - Future work will allow custom actor placement. - """ - # Get the actor for this expert - actor_id = local_expert_id % pp_degree - actor = get_actor(actor_id) - current_actor = ray.get_runtime_context().current_actor - - # Store metadata for backward pass - ctx.global_expert_id = global_expert_id - ctx.local_expert_id = local_expert_id - ctx.batch_idx = batch_idx - ctx.pp_degree = pp_degree - ctx.actor_id = actor_id - - from .piper_utils import piper_metadata - actor_self = piper_metadata.actor_self - - # Store which inputs require gradients - ctx.input_requires_grad = [arg.requires_grad if isinstance(arg, torch.Tensor) else False for arg in args] - ctx.num_inputs = len(args) - - # Dispatch a remote call if actor is not current actor - if actor == current_actor: - output = actor_self.run_expert(global_expert_id, batch_idx, *args) - else: - ref = actor.run_expert.remote(global_expert_id, batch_idx, *args) - output = ray.get(ref) - - # Detach the output so the graph is not carried over when - # the expert runs on the current actor. - if isinstance(output, (list, tuple)): - output = [t.detach() for t in output] - else: - output = output.detach() - - return output - - @staticmethod - def backward(ctx, grad_output): - """ - Dispatch a backward expert call corresponding to the - forward expert call from the ctx. - """ - # Get the actor for this expert - actor = get_actor(ctx.actor_id) - current_actor = ray.get_runtime_context().current_actor - logger.debug(f"Expert {ctx.global_expert_id} {ctx.batch_idx} backward on actor: {current_actor} calling actor: {actor}") - num_inputs = ctx.num_inputs - - from .piper_utils import piper_metadata - actor_self = piper_metadata.actor_self - - # Dispatch a remote call if actor is not current actor - if actor == current_actor: - grad_inputs = actor_self.backward_expert(ctx.global_expert_id, ctx.batch_idx, grad_output) - else: - refs = actor.backward_expert.options(num_returns=num_inputs).remote(ctx.global_expert_id, ctx.batch_idx, grad_output) - grad_inputs = ray.get(refs) - - # Return gradients for inputs (None for metadata arguments) - result = [None, None, None, None] + list(grad_inputs) - - return tuple(result) - -def dispatch_expert_ray(global_expert_id: int, local_expert_id: int, batch_idx: int, pp_degree: int, *args): - """ - Dispatch a Ray remote call to run an expert GraphModule - using a custom autograd Function. This function is called - from the FX graph, so it needs to be allowed in the graph. - """ - return ExpertRayFunction.apply(global_expert_id, local_expert_id, batch_idx, pp_degree, *args) - -# Allow the dispatch function in the graph -torch.compiler.allow_in_graph(dispatch_expert_ray) - -@ray.remote -class PiperActor: - def __init__(self, actor_id, optim_class, world_size, dp_rank=0, dp_degree=1, pp_degree=1): - self.logger = create_logger("piper_actor", LOG_LEVEL) - - self.actor_id = actor_id - self.optim_class = optim_class - - self.dp_rank = dp_rank - self.dp_degree = dp_degree - self.pp_degree = pp_degree - self.world_size = world_size - - self.dp_group = None - self.device = 'cuda' - - if dp_degree == 1: - self.global_rank = actor_id - elif pp_degree == 1: - self.global_rank = dp_rank - else: - self.global_rank = dp_rank * dp_degree + actor_id - - self.logger.info(f"Initializing Ray actor {actor_id} global rank {self.global_rank} GPU {os.environ['CUDA_VISIBLE_DEVICES']}") - - self.input = None - self.truth = None - - # map stage id -> compiled fx.Graph function - self.forward_fns = dict() - # map stage id -> original GraphModule (for hook registration) - self.graph_modules = dict() - # map stage id -> model parameters used by the fx.Graph with holes (None values) for input tensors - self.parameters = dict() - # map stage id -> indices of the input tensors (as opposed to model parameters) used by the fx.Graph - self.input_idxs = dict() - # map stage id -> optimizer for the fx.Graph - self.optims = dict() - # map stage id -> mb_idx -> previous activation (if this stage is not first) - self.inp_activation = defaultdict(dict) - # map stage id -> mb_idx -> current activation - self.out_activation = defaultdict(dict) - # accumuate loss for each microbatch - self.loss = [] - # map expert id -> expert - self.experts = dict() - # map expert id -> expert parameters - self.expert_parameters = dict() - # map expert id -> expert input indices - self.expert_input_idxs = dict() - # map expert id -> expert module - self.expert_modules = dict() - # map expert id -> expert input activations - self.expert_input_activations = defaultdict(dict) - # map (expert_id, batch_idx) -> output activation for backward - self.expert_output_activations = defaultdict(dict) - - # Timing infrastructure - self.tracing = False # Toggle for timing and memory tracing - self.trace_data = {'update': {'total': [], 'peak_memory_delta': [], 'peak_memory': []}} - - from .piper_utils import piper_metadata - piper_metadata.actor_self = self - - def load_actor_handles(self, actor_handles): - from .piper_utils import piper_metadata - piper_metadata.actors = actor_handles - - def id(self): - return self.actor_id - - def send_input(self, tensor): - self.input = tensor.to(self.device) - return "done" - - def send_truth(self, tensor): - self.truth = tensor.to(self.device) - return "done" - - def join_process_groups(self): - master_addr = os.environ.get('PIPER_MASTER_ADDR', "127.0.0.1") - master_port = os.environ.get('PIPER_MASTER_PORT', "10000") - init_method = f"tcp://{master_addr}:{master_port}" - - dist.init_process_group("nccl", init_method=init_method, rank=self.global_rank, world_size=self.world_size) - self.logger.debug(f"Actor {self.actor_id} global rank {self.global_rank} has GPU {os.environ['CUDA_VISIBLE_DEVICES']}, joined the global process group") - - if self.dp_degree > 1: - self.join_dp_process_group() - - def join_dp_process_group(self): - # Every process needs to participate in every subgroup creation - num_dp_groups = self.world_size // self.dp_degree - for dp_group_id in range(num_dp_groups): - group_ranks = [(dp_group_id + num_dp_groups * i) for i in range(self.dp_degree)] - process_group = dist.new_group(ranks=group_ranks, backend='nccl') - if self.global_rank % num_dp_groups == dp_group_id: - self.dp_group = process_group - self.logger.info(f"Global rank {self.global_rank} joined its dp group {dp_group_id} along with ranks {group_ranks}") - - def load_graph(self, stage_id: int, gm_data, compiler_fn, graphargs, input_idxs): - self.logger.debug(f"Compiling graph on actor {self.actor_id} for stage id: {stage_id} with inputs: {len(graphargs)} and input indices: {input_idxs}") - - # set up tracing data structure - if stage_id not in self.trace_data: - self.trace_data[stage_id] = { - 'forward': { - 'forward': [], - 'total': [], - 'peak_memory_delta': [], - 'peak_memory': [] - }, - 'backward': { - 'backward': [], - 'total': [], - 'peak_memory_delta': [], - 'peak_memory': [] - }, - } - - # compile the graph with the given graphargs - gm = deserialize_graphmodule(gm_data) - - # Store GraphModule reference - self.graph_modules[stage_id] = gm - self.forward_fns[stage_id] = gm.forward - - # save the parameters and initialize the optimizer - self.add_param_group(stage_id, graphargs, input_idxs) - - del gm_data - - def add_param_group(self, stage_id: int, params, input_idxs): - self.logger.debug(f"Adding param group for stage {stage_id}") - - if stage_id in self.parameters: - self.logger.debug(f"Param group already exists for stage {stage_id}") - return - - # place parameters on the device - def move_to_device(idx, arg): - if idx not in input_idxs: - return arg.to(self.device).detach().requires_grad_(True) - else: - return arg.to(self.device) - params = list(map(move_to_device, range(len(params)), params)) - - # discard the graphargs that correspond to input tensors - for i in input_idxs: - params[i] = None - - # save the parameters - self.input_idxs[stage_id] = input_idxs - self.parameters[stage_id] = params - - # add the parameters to the optimizer for this stage - if stage_id not in self.optims: - self.optims[stage_id] = self.optim_class([param for param in params if param is not None]) - else: - self.optims[stage_id].add_param_group({'params': [param for param in params if param is not None]}) - - def load_expert(self, stage_id, expert_id, local_expert_id, batch_idx, expert_module, input_idxs, param_idxs, params): - self.logger.debug(f"Loading expert {expert_id} on actor {self.actor_id} global rank {self.global_rank}") - self.experts[expert_id] = expert_module - - # place parameters on the device - def move_to_device(idx, arg): - if arg is None: - return None - if idx not in input_idxs: - return arg.to(self.device).detach().requires_grad_(True) - else: - return arg.to(self.device) - params = list(map(move_to_device, range(len(params)), params)) - - self.expert_input_idxs[expert_id] = input_idxs - self.expert_parameters[expert_id] = params - self.expert_modules[expert_id] = expert_module - - # add the parameters to the optimizer for this stage - if stage_id not in self.optims: - self.optims[stage_id] = self.optim_class([param for param in params if param is not None]) - else: - self.optims[stage_id].add_param_group({'params': [param for param in params if param is not None]}) - - def run_expert(self, expert_id, batch_idx, *args): - logger.debug(f"Actor {self.actor_id} global rank {self.global_rank} forward expert {expert_id} batch element {batch_idx}") - - # Place args on the device and detach existing grad_fn - def place(arg): - if isinstance(arg, torch.Tensor): - if arg.device == torch.device('cpu'): - out = arg.to(self.device) - else: - out = arg - if out.grad_fn is not None: - out = out.detach().requires_grad_(True) - return out - args = list(map(place, args)) - - # place input tensors in the correct indices - for i, arg in zip(self.expert_input_idxs[expert_id], args): - self.expert_parameters[expert_id][i] = arg - - # save input activations - self.expert_input_activations[expert_id][batch_idx] = args - - # run expert module, track grad fn - with torch.set_grad_enabled(True): - out = self.expert_modules[expert_id](*self.expert_parameters[expert_id]) - - # save output activation - if isinstance(out, (list, tuple)): - self.expert_output_activations[expert_id][batch_idx] = out[0] - else: - self.expert_output_activations[expert_id][batch_idx] = out - - # clear the input tensors - for i in self.expert_input_idxs[expert_id]: - self.expert_parameters[expert_id][i] = None - del args - - return out - - def backward_expert(self, expert_id, batch_idx, grad_output): - logger.debug(f"Actor {self.actor_id} global rank {self.global_rank} backward expert {expert_id} batch element {batch_idx}") - - if expert_id not in self.expert_output_activations or batch_idx not in self.expert_output_activations[expert_id]: - raise ValueError(f"No output activation found for expert {expert_id} batch {batch_idx}") - - # Get the stored output activation - output_activation = self.expert_output_activations[expert_id][batch_idx] - - self.logger.debug(f"Expert {expert_id} backward graph:") - print_backward_graph(self.logger.debug, output_activation) - - if isinstance(grad_output, torch.Tensor) and grad_output.device != self.device: - grad_output = grad_output.to(self.device) - - output_activation.backward(gradient=grad_output) - - # Get the gradients for the input activations - input_activations = self.expert_input_activations[expert_id][batch_idx] - grad_inputs = [] - for i, input_activation in enumerate(input_activations): - if input_activation.requires_grad: - grad_inputs.append(input_activation.grad) - else: - grad_inputs.append(None) - - # Clean up - del output_activation - del input_activations - - return grad_inputs - - # @ray.method(tensor_transport="nccl") - def forward(self, stage_id: int, mb_idx: int, *args): - self.logger.debug(f"Calling forward {stage_id} mb {mb_idx} on actor {self.actor_id} global rank {self.global_rank}") - - if self.tracing: - beginning_event = torch.cuda.Event(enable_timing=True) - forward_start_event = torch.cuda.Event(enable_timing=True) - forward_end_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - beginning_event.record() - - def pre_loaded_input(param): - if param is None: - return self.input - else: - return param - args = list(map(pre_loaded_input, args)) - - # Ray object refs resolve to a single element list - def unwrap(x): - if isinstance(x, list) or isinstance(x, tuple): - assert len(x) == 1 - return x[0] if isinstance(x, list) or isinstance(x, tuple) else x - args = list(map(unwrap, args)) - - def place(arg): - if arg.device == torch.device('cpu'): - return arg.to(self.device) - else: - return arg - args = list(map(place, args)) - - # place input tensors in the correct indices - for i, arg in zip(self.input_idxs[stage_id], args): - self.parameters[stage_id][i] = arg - - # save first input as input activation - if args[0].dtype == torch.float: - args[0].requires_grad_().retain_grad() - self.inp_activation[stage_id][mb_idx] = args[0] - - # Record start event for forward timing - if self.tracing: - torch.cuda.reset_peak_memory_stats() - forward_start_memory = torch.cuda.memory_allocated() - forward_start_event.record() - - # Call compiled function - output = self.forward_fns[stage_id](*self.parameters[stage_id]) - assert isinstance(output, (list, tuple)) and len(output) == 1, "Piper only supports one output per subgraph" - output = output[0] - - # Record end event and calculate forward timing - if self.tracing: - forward_end_event.record() - torch.cuda.synchronize() - # Calculate total forward time - forward_time = forward_start_event.elapsed_time(forward_end_event) - forward_peak_memory_delta_gb = (torch.cuda.max_memory_allocated() - forward_start_memory) / (1024**3) - forward_peak_memory_gb = torch.cuda.max_memory_allocated() / (1024**3) - - # save output as output activation - self.logger.debug(f"Saving output activation {output.shape} for stage {stage_id} mb {mb_idx}") - self.out_activation[stage_id][mb_idx] = output - - # clear the input tensors - for i in self.input_idxs[stage_id]: - self.parameters[stage_id][i] = None - del args - - if self.tracing: - end_event.record() - torch.cuda.synchronize() - total_time = forward_start_event.elapsed_time(end_event) - self.trace_data[stage_id]['forward']['forward'].append(forward_time) - self.trace_data[stage_id]['forward']['total'].append(forward_time) - self.trace_data[stage_id]['forward']['peak_memory_delta'].append(forward_peak_memory_delta_gb) - self.trace_data[stage_id]['forward']['peak_memory'].append(forward_peak_memory_gb) - - if CLEANUP_MEMORY: - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - return output - - def forward_cpu(self, stage_id: int, mb_idx: int, *args): - self.logger.debug(f"Calling cpu forward {stage_id} mb {mb_idx} on actor {self.actor_id} global rank {self.global_rank}") - - def pre_loaded_input(param): - if param is None: - return self.input - else: - return param - args = list(map(pre_loaded_input, args)) - - # Ray object refs resolve to a single element list - def unwrap(x): - if isinstance(x, list) or isinstance(x, tuple): - assert len(x) == 1 - return x[0] - else: - return x - args = list(map(unwrap, args)) - - def place(arg): - if arg.device == torch.device('cpu'): - return arg.to(self.device) - else: - return arg - args = list(map(place, args)) - - # place input tensors in the correct indices - for i, arg in zip(self.input_idxs[stage_id], args): - self.parameters[stage_id][i] = arg - - # save first input as input activation - if args[0].dtype == torch.float: - args[0].requires_grad_().retain_grad() - self.inp_activation[stage_id][mb_idx] = args[0] - - out = self.forward_fns[stage_id](*self.parameters[stage_id]) - - # save output as output activation - activation_tensor = out[0] if isinstance(out, (list, tuple)) else out - self.out_activation[stage_id][mb_idx] = activation_tensor - - # clear the input tensors - for i in self.input_idxs[stage_id]: - self.parameters[stage_id][i] = None - del args - - if CLEANUP_MEMORY: - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - return out - - # @ray.method(tensor_transport="nccl") - def backward(self, stage_id: int, mb_idx: int, inp, truth=None, loss_fn=None): - self.logger.debug(f"Calling backward {stage_id} mb {mb_idx} on actor {self.actor_id} global rank {self.global_rank}") - - if self.tracing: - beginning_event = torch.cuda.Event(enable_timing=True) - backward_start_event = torch.cuda.Event(enable_timing=True) - backward_end_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - beginning_event.record() - torch.cuda.reset_peak_memory_stats() - backward_start_memory = torch.cuda.memory_allocated() - backward_start_event.record() - - if isinstance(inp, list) or isinstance(inp, tuple): - assert len(inp) == 1 - inp = inp[0] - - # Get the activations for this stage and microbatch - out_activation = self.out_activation[stage_id][mb_idx] - inp_activation = self.inp_activation[stage_id][mb_idx] - self.logger.debug(f"Stage {stage_id} backward graph:") - print_backward_graph(self.logger.debug, out_activation) - - # compute loss with the final activation of the final stage. - # use the saved activation rather than inp because the saved activation stores the computation graph - if loss_fn is not None: - assert out_activation.shape == inp.shape - if self.truth is not None: - labels = self.truth - else: - labels = truth.to(self.device) - assert out_activation.shape == labels.shape - loss = loss_fn(out_activation, labels) - loss.backward() - self.loss.append(loss.item()) - # if not the last stage, backprop on the stored activation given - # the input gradient from the subsequent stage - else: - assert inp is not None - assert out_activation.shape == inp.shape - out_activation.backward(gradient=inp) - - del out_activation - - # propagate the gradient backwards if not the first stage - if stage_id == 0: - ret = "done" - else: - ret = inp_activation.grad - - del inp_activation - - # Record end event and calculate backward timing - if self.tracing: - backward_end_event.record() - torch.cuda.synchronize() - - # Calculate total backward time - backward_time = backward_start_event.elapsed_time(backward_end_event) - - backward_peak_memory_delta_gb = (torch.cuda.max_memory_allocated() - backward_start_memory) / (1024**3) - backward_peak_memory_gb = torch.cuda.max_memory_allocated() / (1024**3) - - end_event.record() - torch.cuda.synchronize() - total_time = beginning_event.elapsed_time(end_event) - self.trace_data[stage_id]['backward']['backward'].append(backward_time) - self.trace_data[stage_id]['backward']['total'].append(total_time) - self.trace_data[stage_id]['backward']['peak_memory_delta'].append(backward_peak_memory_delta_gb) - self.trace_data[stage_id]['backward']['peak_memory'].append(backward_peak_memory_gb) - - if CLEANUP_MEMORY: - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - - return [ret] * 2 - - def synchronize_gradients(self): - """Synchronize gradients across all DP ranks for this stage using all-reduce.""" - self.logger.debug(f"Actor {self.actor_id} global rank {self.global_rank} synchronizing gradients") - - # Iterate over all stages on this actor and synchronize their parameters - for stage_id, parameters in self.parameters.items(): - for param in parameters: - if param is not None and param.grad is not None: - dist.all_reduce(param.grad, op=dist.ReduceOp.AVG, group=self.dp_group) - self.logger.info(f"Actor {self.actor_id} global rank {self.global_rank} finished synchronizing gradients") - - def update(self, *done_mbs): - self.logger.debug(f"Calling update on actor {self.actor_id} global rank {self.global_rank}") - - if self.tracing: - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - - torch.cuda.reset_peak_memory_stats() - update_start_memory = torch.cuda.memory_allocated() - - start_event.record() - - # if dp degree > 1, synchronize the gradients - if self.dp_degree > 1: - self.synchronize_gradients() - - # step the optimizer for each stage - for _, optim in self.optims.items(): - optim.step() - optim.zero_grad() - losses = self.loss - self.loss.clear() - - if self.tracing: - end_event.record() - torch.cuda.synchronize() - total_time = start_event.elapsed_time(end_event) - update_peak_memory_delta_gb = (torch.cuda.max_memory_allocated() - update_start_memory) / (1024**3) - update_peak_memory_gb = torch.cuda.max_memory_allocated() / (1024**3) - - self.trace_data['update']['total'].append(total_time) - self.trace_data['update']['peak_memory_delta'].append(update_peak_memory_delta_gb) - self.trace_data['update']['peak_memory'].append(update_peak_memory_gb) - - return losses - - def clear_trace_data(self) -> None: - """ - Clear all collected timing data. - """ - for stage_id in self.trace_data: - self.trace_data[stage_id] = { - 'forward': { - 'forward': [], - 'total': [], - 'peak_memory_delta': [], - 'peak_memory': [] - }, - 'backward': { - 'backward': [], - 'total': [], - 'peak_memory_delta': [], - 'peak_memory': [] - }, - } - self.trace_data['update'] = { - 'total': [], - 'peak_memory_delta': [], - 'peak_memory': [] - } - - def get_trace_data(self) -> dict: - return self.trace_data - - def set_tracing(self, enabled: bool) -> None: - """ - Enable or disable timing and memory tracing. - - Args: - enabled (bool): True to enable tracing, False to disable. - """ - self.tracing = enabled - self.logger.info(f"Actor {self.actor_id}: Tracing {'enabled' if enabled else 'disabled'}") - - def start_mem_tracing(self) -> None: - torch.cuda.memory._record_memory_history() - return "done" - - def stop_mem_tracing(self) -> None: - torch.cuda.memory._dump_snapshot(f"actor{self.actor_id}_memory_snapshot_mb4_gpipe.pickle") - self.logger.info(f"Saved memory snapshot to actor{self.actor_id}_memory_snapshot_mb4_gpipe.pickle") - torch.cuda.memory._record_memory_history(enabled=None) - return "done" - - def reset_peak_memory(self): - torch.cuda.reset_peak_memory_stats() - return "done" - - def get_peak_memory(self): - return torch.cuda.max_memory_allocated() / (1024**3) diff --git a/src/piper_compile.py b/src/piper_compile.py deleted file mode 100644 index df29e27..0000000 --- a/src/piper_compile.py +++ /dev/null @@ -1,78 +0,0 @@ -import ray -import torch -import threading -import os -import gc -import copy - -from .piper_actor import create_actors -from .piper_utils import piper_metadata, RemoteTensor, create_logger, LOG_LEVEL -from .piper import piper - -logger = create_logger("piper_compile", LOG_LEVEL) - - -def piper_setup(model_class, model_args, optim_fn, example_inputs, num_stages, pp_degree, dynamic=False, check_correct=False): - """ - Compile a model with the piper backend. - - Args: - model: A model to compile. - example_inputs: Example inputs to the model. - dynamic: Whether to compile in dynamic mode. - backend: The backend to use for compilation. - - Returns: - A tuple of (compiled_model, piper_metadata) where piper_metadata contains - the actors, stage_fns, and other state populated during compilation. - """ - create_actors(pp_degree, optim_fn) - - model = model_class(*model_args) - - if check_correct: - model_nocompile = copy.deepcopy(model) - - piper_metadata.currently_compiling = True - - compiled = torch.compile(model, dynamic=dynamic, backend=piper) - - from .piper_utils import events_tls - events_tls.actor_mutexes = dict([(actor_id, threading.Lock()) for actor_id in range(pp_degree)]) - events_tls.events = [threading.Event() for _ in range(num_stages)] - for event in events_tls.events: - event.set() - - dp_rank = int(os.environ['PIPER_DP_RANK']) - logger.info(f"DP rank {dp_rank+1} compiling {num_stages} stages...") - - output = compiled(*example_inputs).get() - - if check_correct: - correct_output = model_nocompile(*example_inputs) - if not torch.allclose(output, correct_output): - logger.error(f"Model output is not correct") - logger.error(f"Compiled output: {output}") - logger.error(f"Correct output: {correct_output}") - else: - logger.info(f"Model output is correct") - - logger.info(f"DP rank {dp_rank+1} finished compiling model. DAG: {piper_metadata.dag}") - - piper_metadata.currently_compiling = False - - logger.info(f"DP rank {dp_rank} joining process groups") - - ray.get([actor.join_process_groups.remote() for actor in piper_metadata.actors.values()]) - - logger.info(f"Completed DP rank {dp_rank+1} setup for {len(piper_metadata.actors)} actors") - - from ray.experimental.collective import create_collective_group - - create_collective_group( - list(piper_metadata.actors.values()), - backend="nccl") - - logger.info(f"DP rank {dp_rank} Started NCCL group with {len(piper_metadata.actors)} actors") - - return compiled diff --git a/src/piper_coordinator.py b/src/piper_coordinator.py deleted file mode 100644 index 17c0f91..0000000 --- a/src/piper_coordinator.py +++ /dev/null @@ -1,55 +0,0 @@ -import ray -import torch -from typing import Callable -import os -import socket - -from .piper_utils import create_logger, LOG_LEVEL - - -@ray.remote(num_gpus=0.1) -def run_dp_rank(dp_rank, dp_degree, pp_degree, world_size, master_addr, master_port, training_func: Callable, *args, **kwargs): - logger = create_logger("piper_coordinator", LOG_LEVEL) - logger.debug(f"Running DP rank {dp_rank+1} of {dp_degree}") - - os.environ['PIPER_DP_RANK'] = str(dp_rank) - os.environ['PIPER_DP_DEGREE'] = str(dp_degree) - os.environ['PIPER_PP_DEGREE'] = str(pp_degree) - os.environ['PIPER_WORLD_SIZE'] = str(world_size) - os.environ['PIPER_MASTER_ADDR'] = str(master_addr) - os.environ['PIPER_MASTER_PORT'] = str(master_port) - return training_func(*args, **kwargs) - -def find_free_port(): - import socket - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.bind(("0.0.0.0", 0)) - port = s.getsockname()[1] - return port - -@ray.remote -class PiperProgramCoordinator: - """ Central Actor that Coordinates all the DP replicas of a single pipeline""" - def __init__(self, dp_degree, pp_degree): - self.dp_degree = dp_degree - self.pp_degree = pp_degree - self.world_size = dp_degree * pp_degree - self.master_port = find_free_port() - - def run_program(self, training_func: Callable, *args, **kwargs): - return ray.get([run_dp_rank.remote( - dp_rank, - self.dp_degree, - self.pp_degree, - self.world_size, - "127.0.0.1", - self.master_port, - training_func, - *args, - **kwargs) - for dp_rank in range(self.dp_degree)]) - - - - - diff --git a/src/piper_exec.py b/src/piper_exec.py deleted file mode 100644 index 33e1afd..0000000 --- a/src/piper_exec.py +++ /dev/null @@ -1,289 +0,0 @@ -import os -import ray -import torch -import torch.distributed as dist -from typing import NamedTuple -import threading -import time - -from .piper_utils import piper_metadata, create_logger, LOG_LEVEL - -logger = create_logger("piper_exec", LOG_LEVEL) - -class Task(NamedTuple): - device_id: int - stage_id: int - mb_idx: int - is_fwd: bool - upd: bool - -class DAGEdge(NamedTuple): - from_stage: int - to_stage: int - - -def get_backward_targets(stage_id: int, dag_edges: list[DAGEdge]): - return [edge for edge in dag_edges if edge.to_stage == stage_id] - - -def validate_schedule(schedule: list[list[Task | None]], dag_edges: list[DAGEdge], num_mbs: int) -> None: - """ - Validate that the schedule respects well-formedness rules and DAG dependencies. - - Args: - schedule: 2D array with one row per device and one column per time step - dag_edges: List of DAG edges defining stage dependencies - num_mbs: Number of microbatches in the schedule - - Raises: - ValueError: If the schedule violates any validation rules - """ - num_devices, num_steps = len(schedule), len(schedule[0]) if schedule else 0 - - # Check well-formedness: no duplicates, device_id matches row, and all stages present - all_tasks = set() - microbatch_tasks = {} # mb_idx -> set of (stage_id, is_fwd, upd) - - for stage_id in range(num_devices): - for time_step in range(num_steps): - task = schedule[stage_id][time_step] - if task is not None: - # Check device_id matches row - if task.device_id != stage_id: - raise ValueError( - f"Task device_id {task.device_id} does not match row {stage_id} " - f"at time step {time_step}" - ) - - # Check for duplicates - task_key = (task.stage_id, task.mb_idx, task.is_fwd, task.upd) - if task_key in all_tasks: - raise ValueError( - f"Duplicate task found: stage_id={task.stage_id}, " - f"mb_idx={task.mb_idx}, is_fwd={task.is_fwd}, upd={task.upd}" - ) - all_tasks.add(task_key) - - # Track tasks by microbatch - if task.mb_idx not in microbatch_tasks: - microbatch_tasks[task.mb_idx] = set() - microbatch_tasks[task.mb_idx].add((task.stage_id, task.is_fwd, task.upd)) - - # Get all required stages from DAG edges - all_required_stages = set() - for edge in dag_edges: - all_required_stages.add(edge.from_stage) - all_required_stages.add(edge.to_stage) - - # Check that each microbatch has all required forward and backward stages - for mb_idx, tasks in microbatch_tasks.items(): - # Find all stages that have forward/backward tasks for this microbatch - fwd_stages = {stage_id for stage_id, is_fwd, upd in tasks if is_fwd and not upd} - bwd_stages = {stage_id for stage_id, is_fwd, upd in tasks if not is_fwd and not upd} - - # Check that all required stages have forward tasks - missing_fwd = all_required_stages - fwd_stages - if missing_fwd: - raise ValueError( - f"Microbatch {mb_idx} missing forward stages: {missing_fwd}" - ) - - # Check that all required stages have backward tasks - missing_bwd = all_required_stages - bwd_stages - if missing_bwd: - raise ValueError( - f"Microbatch {mb_idx} missing backward stages: {missing_bwd}" - ) - - # Check pipeline stage dependencies - for mb_idx in range(num_mbs): - # Find all tasks for this microbatch - fwd_times = {} # stage_id -> time_step - bwd_times = {} # stage_id -> time_step - - for device_id in range(num_devices): - for time_step in range(num_steps): - task = schedule[device_id][time_step] - if task is not None and task.mb_idx == mb_idx: - if task.is_fwd and not task.upd: - fwd_times[task.stage_id] = time_step - elif not task.is_fwd and not task.upd: - bwd_times[task.stage_id] = time_step - - # Check forward stage ordering: if A -> B, then fwd(A) < fwd(B) - for edge in dag_edges: - from_stage, to_stage = edge.from_stage, edge.to_stage - if from_stage in fwd_times and to_stage in fwd_times: - if fwd_times[from_stage] >= fwd_times[to_stage]: - raise ValueError( - f"Forward stage ordering violation for microbatch {mb_idx}: " - f"forward stage {from_stage} (time {fwd_times[from_stage]}) must come " - f"before forward stage {to_stage} (time {fwd_times[to_stage]})" - ) - - # Check forward-backward ordering: fwd(A) < bwd(A) - for stage_id in fwd_times: - if stage_id in bwd_times: - if fwd_times[stage_id] >= bwd_times[stage_id]: - raise ValueError( - f"Forward-backward ordering violation for microbatch {mb_idx}, " - f"stage {stage_id}: forward (time {fwd_times[stage_id]}) must come " - f"before backward (time {bwd_times[stage_id]})" - ) - - # Check backward stage ordering: if A -> B, then bwd(B) < bwd(A) - for edge in dag_edges: - from_stage, to_stage = edge.from_stage, edge.to_stage - if from_stage in bwd_times and to_stage in bwd_times: - if bwd_times[to_stage] >= bwd_times[from_stage]: - raise ValueError( - f"Backward stage ordering violation for microbatch {mb_idx}: " - f"backward stage {to_stage} (time {bwd_times[to_stage]}) must come " - f"before backward stage {from_stage} (time {bwd_times[from_stage]})" - ) - -def piper_exec(model, schedule, inputs, truth, loss_fn, num_mbs, num_stages): - """ - Execute one step of the pipeline schedule on the distributed model. - - Args: - model: A model that has been compiled with the piper backend. - schedule: A 2D list (device x time_step) of Tasks specifying execution order. - inputs: Inputs to the model. - truth: Ground-truth labels or targets. - loss_fn: Loss function to be used for training. - num_mbs: Number of microbatches in the schedule. - - Returns: - List of losses per microbatch.) - """ - num_steps, num_devices = len(schedule[0]), len(schedule) - actors = piper_metadata.actors - - dag_edges = piper_metadata.dag - dag_edges = list(map(lambda e: (DAGEdge(e[0], e[1])), list(piper_metadata.dag))) - - # Validate the schedule before execution - validate_schedule(schedule, dag_edges, num_mbs) - - # maps mb_idx to the ref resulting from a forward call on the microbatch - fwd_refs = dict() - - # maps mb_idx to a dict that maps stage_id to the refs output from the stage's backward on that microbatch - bwd_ref_dicts = dict() - - # create events for each microbatch - events = dict([(mb_idx, [threading.Event() for _ in range(num_stages)]) for mb_idx in range(num_mbs)]) - actor_mutexes = dict([(actor_id, threading.Lock()) for actor_id in actors.keys()]) - threads = dict() - - def run_model(inputs, mb_idx, events, actor_mutexes): - logger.debug(f"Controller launching thread {mb_idx}") - from .piper_utils import events_tls - events_tls.events = events - events_tls.mb_idx = mb_idx - events_tls.actor_mutexes = actor_mutexes - fwd_refs[mb_idx] = model(*inputs) - - # iterate over evrery task in the schedule - ret = [] - for i in range(num_steps): - for j in range(num_devices-1, -1, -1): - task = schedule[j][i] - if task: - device_id,stage_id, mb_idx, is_fwd, upd = task - actor_id = j - num_bwd_targets = len(get_backward_targets(stage_id, dag_edges)) - if num_bwd_targets == 0: - num_bwd_targets = 1 - if upd: - logger.debug(f"Controller updating stage {stage_id} mb {mb_idx}") - done_refs = set() - for _, bwd_ref_dict in bwd_ref_dicts.items(): - bwd_refs = bwd_ref_dict[stage_id] - bwd_refs = bwd_refs[num_bwd_targets:] - if isinstance(bwd_refs, list): - done_refs = done_refs | set(bwd_refs) - else: - done_refs.add(bwd_refs) - done_refs = list(done_refs) - with actor_mutexes[actor_id]: - ret.append(actors[actor_id].update.remote(*done_refs)) - elif is_fwd: - if stage_id == 0: - thread = threading.Thread(target=run_model, args=(inputs, mb_idx, events[mb_idx], actor_mutexes)) - threads[mb_idx] = thread - thread.start() - events[mb_idx][stage_id].set() - logger.debug(f"Controller set event for thread {mb_idx} stage {stage_id}: waiting for thread to grab lock") - # Wait for the thread to grab a lock on the actor mutex, - # signalled by unsetting the event - while events[mb_idx][stage_id].is_set(): - time.sleep(0.001) - logger.debug(f"Controller: thread {mb_idx} grabbed lock") - else: - # log order of task dispatch by printing - # also see output_graph.py:1785 where we log forward dispatch - if mb_idx not in bwd_ref_dicts: - # if this is the first backward task for a microbatch, dispatch the - # backward task and cache the resulting ref(s) - logger.debug(f"Controller waiting for thread {mb_idx} to join") - threads[mb_idx].join() - logger.debug(f"Controller thread {mb_idx} joined") - fwd_ref = fwd_refs[mb_idx] - bwd_ref_dicts[mb_idx] = dict() - logger.debug(f"Controller waiting for actor {actor_id} mutex") - with actor_mutexes[actor_id]: - logger.debug(f"Controller got actor {actor_id} mutex, launching backward stage {stage_id} mb {mb_idx}") - bwd_ref_dicts[mb_idx][stage_id] = ( - actors[actor_id] - .backward.options(num_returns=num_bwd_targets*2) - .remote( - stage_id, mb_idx, fwd_ref.get_ref(), truth=truth, loss_fn=loss_fn - ) - ) - else: - # if this is not the first backward task for a microbatch, look up - # the input ref which represents a gradient from the backward call - # of a subsequent stage - - # get the result of the subsequent stage's backward call - # LIMITATION: there cannot be more than one subsequent stage (e.g. the - # forward of stage A cannot be inputs to both stage B and C) - # this limitation should eventually be resolved by making the logic below more general - to_stage = [ - edge.to_stage - for edge in dag_edges - if edge.from_stage == stage_id - ] - assert len(to_stage) == 1 - to_stage = to_stage[0] - - # get the refs resulting from the subsequent stage's backward - targets = get_backward_targets(to_stage, dag_edges) - - # get the idx of the current stage in the subsequent stage's output list - # LIMITATION: this logic assumes that the user's dag_edges list is ordered according - # to how the stages are ordered in the model file. e.g. in clip.py stage 0 comes - # before stage 1, so dag_edges must look like - # dag_edges = [DAGEdge(0, 2), DAGEdge(1, 2)] and NOT - # dag_edges = [DAGEdge(1, 2), DAGEdge(0, 2)] - idx = targets.index(DAGEdge(stage_id, to_stage)) - - # make sure the subsequent stage's backward was already dispatched - assert to_stage in bwd_ref_dicts[mb_idx] - - # if the subsequent stage's backward had more than one output, index the - # list to get the ref for the current stage - bwd_refs = bwd_ref_dicts[mb_idx][to_stage] - bwd_ref = bwd_refs[idx] - # dispatch the current stage's backward and cache the resulting ref(s) - logger.debug(f"Controller waiting for actor {actor_id} mutex") - with actor_mutexes[actor_id]: - logger.debug(f"Controller got actor {actor_id} mutex, launching backward stage {stage_id} mb {mb_idx}") - bwd_ref_dicts[mb_idx][stage_id] = ( - actors[actor_id] - .backward.options(num_returns=num_bwd_targets*2) - .remote(stage_id, mb_idx, bwd_ref) - ) - return ray.get(ret) \ No newline at end of file diff --git a/src/piper_graph_transform.py b/src/piper_graph_transform.py deleted file mode 100644 index 7a610c0..0000000 --- a/src/piper_graph_transform.py +++ /dev/null @@ -1,292 +0,0 @@ -import ray -import torch -import torch.fx as fx -from collections import defaultdict -import operator - -from .piper_actor import get_actor, dispatch_expert_ray -from .piper_utils import create_logger, LOG_LEVEL - -logger = create_logger("piper_graph_transform", LOG_LEVEL) - -def split_gm_by_experts(gm, stage_id, pp_degree): - """ - Transform a graph module with expert annotations by extracting - expert computations into submodules and distributing the expert - submodules to Ray actors. - """ - # Collect all nodes with custom metadata - nodes_with_metadata = [] - expert_metadata = {} - for node in gm.graph.nodes: - if 'custom' in node.meta: - custom_meta = node.meta['custom'] - # Extract expert and batch_idx from metadata - if isinstance(custom_meta, dict): - global_expert_id = custom_meta.get('global_expert_id') - local_expert_id = custom_meta.get('local_expert_id') - batch_idx = custom_meta.get('batch_idx') - if global_expert_id is not None: - expert_metadata[global_expert_id] = (local_expert_id, batch_idx) - nodes_with_metadata.append((node, global_expert_id, local_expert_id, batch_idx)) - - if not nodes_with_metadata: - logger.info("No expert nodes found in graph") - return gm - else: - logger.info(f"Found expert nodes in graph") - - # Group nodes by expert ID for creating expert modules - expert_code = defaultdict(list) # expert_id -> list of all nodes for this expert - - for node, global_expert_id, _, _ in nodes_with_metadata: - expert_code[global_expert_id].append(node) - - expert_modules = {} - - for expert_id, expert_nodes in expert_code.items(): - expert_graph = fx.Graph() - expert_node_mapping = {} - - expert_node_set = set(expert_nodes) - - # Find all inputs needed by this expert (nodes that are not in the expert set) - expert_inputs = set() - for node in expert_nodes: - for arg in fx.graph.map_arg(node.args, lambda n: n): - if isinstance(arg, fx.Node) and arg not in expert_node_set: - expert_inputs.add(arg) - - # Create placeholders for inputs - input_placeholders = {} - for input_node in expert_inputs: - placeholder = expert_graph.placeholder(input_node.name) - expert_node_mapping[input_node] = placeholder - input_placeholders[input_node] = placeholder - - # Find outputs of this expert (nodes used outside the expert set) - expert_outputs = [] - for node in expert_nodes: - for user in node.users: - if user not in expert_node_set: - if node not in expert_outputs: - expert_outputs.append(node) - - # If no external users, use the last node(s) as output - if not expert_outputs: - expert_outputs = [expert_nodes[-1]] if expert_nodes else [] - - # Topological sort of expert nodes - def get_dependencies(node): - deps = set() - for arg in fx.graph.map_arg(node.args, lambda n: n): - if isinstance(arg, fx.Node) and arg in expert_node_set: - deps.add(arg) - return deps - remaining = set(expert_nodes) - ordered_expert_nodes = [] - while remaining: - for node in list(remaining): - deps = get_dependencies(node) - if deps.issubset(set(ordered_expert_nodes)): - ordered_expert_nodes.append(node) - remaining.remove(node) - break - - # Copy expert nodes to the expert graph in topological order - for node in ordered_expert_nodes: - new_expert_node = expert_graph.node_copy( - node, - lambda n: expert_node_mapping.get(n, input_placeholders.get(n)) - ) - expert_node_mapping[node] = new_expert_node - - # Create output node - if expert_outputs: - output_values = [expert_node_mapping[node] for node in expert_outputs] - if len(output_values) == 1: - expert_graph.output(output_values[0]) - else: - expert_graph.output(tuple(output_values)) - else: - # No outputs, create a dummy output - expert_graph.output(expert_graph.placeholder('dummy')) - - expert_gm = fx.GraphModule(torch.nn.Module(), expert_graph) - - # Store the expert module and related info - for node in ordered_expert_nodes: - if node.op == "placeholder": - input_placeholders[node] = node - - # Load the module on the corresponding actor - input_idxs, param_idxs = [], [] - params = [] - placeholders = expert_gm.graph.find_nodes(op="placeholder") - for i, placeholder in enumerate(placeholders): - if "grapharg" in placeholder.meta: - if 'self' in str(placeholder.meta["grapharg"]): - param_idxs.append(i) - params.append(placeholder.meta["grapharg"]._example()) - else: - input_idxs.append(i) - params.append(None) - else: - input_idxs.append(i) - params.append(None) - - # logger.debug(f"Submodule {expert_id} input_idxs: {input_idxs} param_idxs: {param_idxs}, params: {[p.shape if p is not None else None for p in params]}") - local_expert_id, batch_idx = expert_metadata[expert_id] - expert_metadata[expert_id] = (local_expert_id, batch_idx, input_idxs) - actor_id = local_expert_id % pp_degree - actor = get_actor(actor_id) - ray.get(actor.load_expert.remote(stage_id, expert_id, local_expert_id, batch_idx, expert_gm, input_idxs, param_idxs, params)) - - expert_modules[expert_id] = (expert_gm, input_placeholders, expert_outputs, expert_nodes, ordered_expert_nodes) - - # Create a new top-level graph and replace expert nodes with call_module - new_graph = fx.Graph() - node_mapping = {} - - # Copy placeholders - for node in gm.graph.nodes: - if node.op == "placeholder": - new_node = new_graph.node_copy(node, lambda n: node_mapping.get(n, n)) - node_mapping[node] = new_node - - # Track which expert calls have been replaced - expert_call_replaced = {} - - # Copy all get_attr nodes (parameters) so they're available for expert calls - for node in gm.graph.nodes: - if node.op == "get_attr": - if node not in node_mapping: - new_node = new_graph.node_copy(node, lambda n: node_mapping.get(n, n)) - node_mapping[node] = new_node - - # Process nodes in the original graph order to preserve structure - for node in gm.graph.nodes: - if node.op == "placeholder": - # Already handled - continue - - if node.op == "output": - # Handle output separately - continue - - if node.op == "get_attr": - # Already handled above - continue - - if node in node_mapping: - # Already processed - continue - - # Check if this node belongs to an expert and find which expert call it's part of - in_expert = False - for expert_id, expert_nodes in expert_code.items(): - if node in expert_nodes: - in_expert = True - - # Get the expert module info - expert_gm, input_placeholders, expert_outputs, expert_nodes, ordered_expert_nodes = expert_modules[expert_id] - - # Check if we've already created the call_module for this expert call - if expert_id not in expert_call_replaced: - # Find first node in this expert's nodes (in original graph order) - first_node = None - for n in gm.graph.nodes: - if n in expert_nodes: - first_node = n - break - - if first_node is None: - first_node = expert_nodes[0] - - # Map inputs in the same order as placeholders - # Use the inputs that the first node of this expert needs - mapped_args = [] - for input_node in input_placeholders.keys(): - if input_node in node_mapping: - mapped_args.append(node_mapping[input_node]) - else: - # Input node not yet mapped, need to copy it first - new_input_node = new_graph.node_copy(input_node, lambda n: node_mapping.get(n, n)) - node_mapping[input_node] = new_input_node - mapped_args.append(new_input_node) - - # Dispatch expert module to a Ray actor at the position of the first expert node for this expert - module_name = f"expert_{expert_id}" - local_expert_id, batch_idx, input_idxs = expert_metadata[expert_id] - input_tensors_only = [] - for i in input_idxs: - input_tensors_only.append(mapped_args[i]) - # expert_module_node = new_graph.get_attr(module_name) - ray_call_args = [expert_id, local_expert_id, batch_idx, pp_degree] + input_tensors_only - call_node = new_graph.call_function(dispatch_expert_ray, tuple(ray_call_args)) - expert_call_replaced[expert_id] = (call_node, expert_outputs) - - # Map expert nodes to the call_module output - call_node, expert_outputs = expert_call_replaced[expert_id] - - # Only create output mapping for output nodes (nodes used outside expert) - if node in expert_outputs: - if len(expert_outputs) == 1: - # Single output, use call_node directly - node_mapping[node] = call_node - else: - # Multiple outputs, need to getitem - output_idx = expert_outputs.index(node) - getitem_node = new_graph.call_function( - operator.getitem, - (call_node, output_idx) - ) - node_mapping[node] = getitem_node - else: - node_mapping[node] = call_node - break - - if not in_expert: - # Regular node, copy it - new_node = new_graph.node_copy(node, lambda n: node_mapping.get(n, n)) - node_mapping[node] = new_node - - - # Handle output node - output_node = None - for node in gm.graph.nodes: - if node.op == "output": - output_node = node - break - - if output_node: - output_args = fx.graph.map_arg(output_node.args, lambda n: node_mapping.get(n, n)) - if isinstance(output_args, tuple) and len(output_args) == 1: - new_graph.output(output_args[0]) - else: - new_graph.output(output_args) - - # Create root module and add expert modules as submodules - root_module = torch.nn.Module() - for expert_id, (expert_gm, _, _, _, _) in expert_modules.items(): - module_name = f"expert_{expert_id}" - root_module.add_module(module_name, expert_gm) - - new_gm = fx.GraphModule(root_module, new_graph) - - # Copy parameters, buffers, and modules from original module - for name, param in gm.named_parameters(recurse=False): - if name not in new_gm._parameters: - new_gm.register_parameter(name, param) - - for name, buffer in gm.named_buffers(recurse=False): - if name not in new_gm._buffers: - new_gm.register_buffer(name, buffer) - - for name, module in gm.named_children(): - if name not in [f"expert_{expert_id}" for expert_id in expert_modules.keys()]: - new_gm.add_module(name, module) - - new_gm.recompile() - - return new_gm \ No newline at end of file diff --git a/src/piper_patches.py b/src/piper_patches.py deleted file mode 100644 index 18e0958..0000000 --- a/src/piper_patches.py +++ /dev/null @@ -1,166 +0,0 @@ -""" -Runtime monkey patches for PyTorch and Ray dependencies. - -These patches enable Piper's RemoteTensor to work correctly with TorchDynamo -and Ray's tensor transport backends. - -Import this module early in your application to apply patches: - import src.piper_patches -""" - -import functools -import logging - -logger = logging.getLogger(__name__) - -_patches_applied = False - - -def apply_patches(): - """Apply all required monkey patches to PyTorch and Ray.""" - global _patches_applied - if _patches_applied: - return - _patches_applied = True - - _patch_ray_actor_method() - - logger.info("Piper runtime patches applied successfully") - - -def _patch_ray_actor_method(): - """ - Patch ray.actor.ActorMethod._remote to support multiple return values - with tensor transport backends. - - The default Ray implementation only supports 1 return value per task - when using tensor transport. This patch removes that restriction and - properly registers multiple ObjectRefs with the GPU object manager. - """ - try: - import ray - import ray.actor as actor_module - from ray._raylet import ObjectRef - except ImportError as e: - logger.warning(f"Could not import ray.actor: {e}") - return - - # Get TensorTransportEnum for comparison - try: - from ray.actor import TensorTransportEnum - except ImportError: - logger.warning("Could not import TensorTransportEnum, skipping Ray patch") - return - - original_remote = actor_module.ActorMethod._remote - - @functools.wraps(original_remote) - def patched_remote( - self, - args=None, - kwargs=None, - name="", - num_returns=None, - max_task_retries=None, - retry_exceptions=None, - concurrency_group=None, - _generator_backpressure_num_objects=None, - enable_task_events=None, - tensor_transport=None, - ): - if num_returns is None: - num_returns = self._num_returns - if max_task_retries is None: - max_task_retries = self._max_task_retries - if max_task_retries is None: - max_task_retries = 0 - if retry_exceptions is None: - retry_exceptions = self._retry_exceptions - if enable_task_events is None: - enable_task_events = self._enable_task_events - if _generator_backpressure_num_objects is None: - _generator_backpressure_num_objects = ( - self._generator_backpressure_num_objects - ) - - if tensor_transport is None: - tensor_transport = self._tensor_transport - - # PIPER MODIFICATION: Remove the num_returns != 1 check for tensor_transport - # Original code would raise ValueError here if num_returns != 1 - if tensor_transport != TensorTransportEnum.OBJECT_STORE.name: - # Skip the num_returns check - allow multiple returns - if not self._actor._ray_enable_tensor_transport: - raise ValueError( - f'Currently, methods with .options(tensor_transport="{tensor_transport}") are not supported when enable_tensor_transport=False. ' - "Please set @ray.remote(enable_tensor_transport=True) on the actor class definition." - ) - gpu_object_manager = ray._private.worker.global_worker.gpu_object_manager - if not gpu_object_manager.actor_has_tensor_transport( - self._actor, tensor_transport - ): - raise ValueError( - f'{self._actor} does not have tensor transport {tensor_transport} available. If using a collective-based transport ("nccl" or "gloo"), please create a communicator with ' - "`ray.experimental.collective.create_collective_group` " - "before calling actor tasks with non-default tensor_transport." - ) - - args = args or [] - kwargs = kwargs or {} - - def invocation(args, kwargs): - dst_actor = self._actor - if dst_actor is None: - raise RuntimeError( - "Lost reference to actor. Actor handles must be stored as variables, e.g. `actor = MyActor.remote()` before calling methods." - ) - - gpu_object_manager = ray._private.worker.global_worker.gpu_object_manager - gpu_object_manager.trigger_out_of_band_tensor_transfer(dst_actor, args) - - return dst_actor._actor_method_call( - self._method_name, - args=args, - kwargs=kwargs, - name=name, - num_returns=num_returns, - max_task_retries=max_task_retries, - retry_exceptions=retry_exceptions, - concurrency_group_name=concurrency_group, - generator_backpressure_num_objects=( - _generator_backpressure_num_objects - ), - enable_task_events=enable_task_events, - tensor_transport=tensor_transport, - ) - - # Apply the decorator if there is one. - if self._decorator is not None: - invocation = self._decorator(invocation) - - object_refs = invocation(args, kwargs) - - # PIPER MODIFICATION: Handle multiple return values with GPU object manager - if tensor_transport != TensorTransportEnum.OBJECT_STORE.name: - gpu_object_manager = ray._private.worker.global_worker.gpu_object_manager - if isinstance(object_refs, ObjectRef): - # Single return value - object_ref = object_refs - gpu_object_manager.add_gpu_object_ref( - object_ref, self._actor, tensor_transport - ) - else: - # Multiple return values - register each ObjectRef - for object_ref in object_refs: - if isinstance(object_ref, ObjectRef): - gpu_object_manager.add_gpu_object_ref( - object_ref, self._actor, tensor_transport - ) - - return object_refs - - actor_module.ActorMethod._remote = patched_remote - logger.debug("Patched ray.actor.ActorMethod._remote") - - -apply_patches() diff --git a/src/piper_utils.py b/src/piper_utils.py deleted file mode 100644 index 2bcbab4..0000000 --- a/src/piper_utils.py +++ /dev/null @@ -1,412 +0,0 @@ -import ray -import torch -import uuid -import inspect -import logging -import json, importlib, operator -import torch.fx as fx -import threading -from collections import defaultdict -from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode -from typing import Any, Optional - -LOG_LEVEL = "INFO" - -""" -Print the backward graph of a tensor -""" - -def print_backward_graph(printer, tensor, prefix=""): - seen = set() - def _print(t, indent=0): - fn = t.grad_fn if hasattr(t, 'grad_fn') and t.grad_fn is not None else None - if fn is None: - printer(" " * indent + f"{prefix}Tensor: no grad_fn") - return - if fn in seen: - printer(" " * indent + f"{prefix}{type(fn).__name__} (recursive/ref)") - return - seen.add(fn) - printer(" " * indent + f"{prefix}{type(fn).__name__}") - for next_fn, _ in fn.next_functions: - if next_fn is not None and hasattr(next_fn, 'variable'): - printer(" " * (indent + 2) + f"{prefix}Variable: {type(next_fn.variable).__name__}") - elif next_fn is not None: - _print(type('Dummy', (), {'grad_fn': next_fn})(), indent + 2) - else: - printer(" " * (indent + 2) + f"{prefix}None") - _print(tensor, 0) - -""" -Logger utility -""" - -def create_logger(name: str, log_level: str): - match log_level: - case "DEBUG": - log_level = logging.DEBUG - case "INFO": - log_level = logging.INFO - case "WARNING": - log_level = logging.WARNING - case "ERROR": - log_level = logging.ERROR - - logger = logging.getLogger(name) - logger.setLevel(log_level) - - if not logger.handlers: - handler = logging.StreamHandler() - fmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" - handler.setFormatter(logging.Formatter(fmt)) - logger.addHandler(handler) - logger.propagate = False - - return logger - -""" -Piper thread local storage for tracking Piper actors, stages, and microbatches -""" - -class ThreadLocal(threading.local): - events = None - mb_idx = None - actor_mutexes = None - -events_tls = ThreadLocal() - -class PiperMetadata: - actors = dict() - dag = set() - world_size = None - currently_compiling = True - current_stage = None - current_actor = None - actor_self = None - first_graph_of_stage = None - parallelism_configs = {'dp': 1} - -piper_metadata = PiperMetadata() - -""" -Remote tensors wrap Ray ObjectRefs -""" - -_fake_tensor_mode = FakeTensorMode() -_fake_tensor_converter = _fake_tensor_mode.fake_tensor_converter - -class RemoteTensorKey: - def __init__(self): - self.key = str(uuid.uuid4()) - -class RemoteTensor(torch.Tensor): - _fake: torch.Tensor - _stage_id: Optional[int] - _obj_ref: ray._raylet.ObjectRef - _resolved: Any - - def __new__(cls, - fake: FakeTensor, - obj_ref: ray._raylet.ObjectRef, - stage_id: Optional[int] = None): - instance = torch.Tensor._make_wrapper_subclass( - cls, - fake.size(), - strides=fake.stride(), - storage_offset=fake.storage_offset(), - device=fake.device, # This is the device of of either input tensor or first tensor of a list - dtype=fake.dtype, - layout=fake.layout, - requires_grad=fake.requires_grad, - ) - instance._obj_ref = obj_ref - instance._fake = fake - instance._stage_id = stage_id - instance._resolved = None - instance.key = RemoteTensorKey() - return instance - - def get_stage_id(self): - return self._stage_id - - def get(self): - # return fake tensor during compilation - if self._obj_ref is None: - return self._fake - - if self._resolved is None: - obj = ray.get(self._obj_ref) - if isinstance(obj, list) or isinstance(obj, tuple): - assert len(obj) == 1 - self._resolved = obj[0] - else: - self._resolved = obj - self._resolved = self._resolved.to('cpu') - return self._resolved - - def get_ref(self): - if self._obj_ref is None: - raise RuntimeError("Cannot get ObjectRef from a compile-time RemoteTensor") - return self._obj_ref - - def __torch_dispatch__(cls, func, types, args=(), kwargs=None): - def unwrap(x): - if isinstance(x, RemoteTensor): - return x.get() - elif isinstance(x, (list, tuple)): - return type(x)(unwrap(v) for v in x) - elif isinstance(x, dict): - return {k: unwrap(v) for k, v in x.items()} - else: - return x - - args = torch.utils._pytree.tree_map(unwrap, args) - kwargs = torch.utils._pytree.tree_map(unwrap, kwargs or {}) - - out = func(*args, **kwargs) - return out - - def __tensor_flatten__(self): - return (["_fake"], {}) - - @classmethod - def __tensor_unflatten__(cls, inner_tensors, metadata, outer_size, outer_stride): - fake = inner_tensors["_fake"] - return cls(fake, None, None) - - def __repr__(self): - """Custom repr that avoids triggering masked_select from tensor formatting.""" - if self._resolved is not None: - return self._resolved.__repr__() - return ( - f"RemoteTensor(shape={tuple(self.shape)}, dtype={self.dtype}, " - f"device={self.device}, stage_id={self._stage_id})" - ) - -torch._dynamo.config.traceable_tensor_subclasses.add(RemoteTensor) - -""" -Serialize/deserialize an fx.GraphModule -""" - -def encode_arg(a): - if isinstance(a, fx.Node): - return {"__node__": a.name} - if isinstance(a, torch.device): - return {"__device__": str(a)} - if isinstance(a, torch.dtype): - return {"__dtype__": str(a).replace("torch.", "")} - if isinstance(a, slice): - return {"__slice__": True, - "start": encode_arg(a.start), - "stop": encode_arg(a.stop), - "step": encode_arg(a.step)} - if a is Ellipsis: - return {"__ellipsis__": True} - if isinstance(a, tuple): # <-- preserve tuples - return {"__tuple__": [encode_arg(x) for x in a]} - if isinstance(a, list): - return [encode_arg(x) for x in a] - if isinstance(a, dict): - return {k: encode_arg(v) for k, v in a.items()} - return a - -def decode_arg(a, name_to_node): - if isinstance(a, dict): - if "__node__" in a: - return name_to_node[a["__node__"]] - if "__device__" in a: - return torch.device(a["__device__"]) - if "__dtype__" in a: - return getattr(torch, a["__dtype__"]) - if "__slice__" in a: - return slice( - decode_arg(a["start"], name_to_node), - decode_arg(a["stop"], name_to_node), - decode_arg(a["step"], name_to_node), - ) - if "__ellipsis__" in a: - return Ellipsis - if "__tuple__" in a: # <-- reconstruct tuples - return tuple(decode_arg(x, name_to_node) for x in a["__tuple__"]) - # generic dict - return {k: decode_arg(v, name_to_node) for k, v in a.items()} - if isinstance(a, list): - return [decode_arg(x, name_to_node) for x in a] - return a - -def _is_op_overload(obj): - # Works across PyTorch versions without importing private types directly - return obj.__class__.__module__.startswith("torch._ops") or obj.__class__.__name__.startswith("OpOverload") - -def serialize_target(t): - # print("SERIALIZING", t) - # call_method uses a string method name, pass through - if isinstance(t, str): - return {"kind": "string", "value": t} - - # Handle _VariableFunctionsClass - if getattr(t, "__module__", "") == "torch._VariableFunctionsClass": - public_name = t.__name__ - if hasattr(torch, public_name): - return {"kind": "py_func", "module": "torch", "qualname": public_name} - else: - raise ValueError(f"No public torch alias for {t}") - - # torch.ops.* (aten, prim, etc.) - if _is_op_overload(t) or (getattr(t, "__module__", "").startswith("torch._ops")): - return {"kind": "torch_op", "path": str(t)} # e.g. "aten.add.Tensor" or "aten.add" - - # regular python function or built-in - if inspect.isfunction(t) or inspect.isbuiltin(t): - mod = inspect.getmodule(t) - if mod is None: - raise ValueError(f"Cannot serialize function without module: {t}") - return {"kind": "py_func", "module": mod.__name__, "qualname": t.__name__} - - # classes or callables rarely appear as call_function targets, but support anyway - if inspect.isclass(t): - mod = t.__module__ - return {"kind": "py_obj", "module": mod, "qualname": t.__qualname__} - - # operator functions (already covered by py_func, but ensure resolvable) - if t in operator.__dict__.values(): - return {"kind": "py_func", "module": "operator", "qualname": t.__name__} - - # last resort: try module+name - mod = getattr(t, "__module__", None) - name = getattr(t, "__name__", None) - if mod and name: - return {"kind": "py_func", "module": mod, "qualname": name} - - raise NotImplementedError(f"Unsupported target type: {t} ({type(t)})") - -def _resolve_qualname(mod, qualname): - obj = mod - for part in qualname.split("."): - obj = getattr(obj, part) - return obj - -def deserialize_target(payload): - # print("DESERIALIZING", payload) - kind = payload["kind"] - - if kind == "string": - return payload["value"] - - if kind == "py_func": - mod = importlib.import_module(payload["module"]) - return _resolve_qualname(mod, payload["qualname"]) - - if kind == "py_obj": - mod = importlib.import_module(payload["module"]) - return _resolve_qualname(mod, payload["qualname"]) - - if kind == "torch_op": - obj = torch.ops - for part in payload["path"].split("."): - obj = getattr(obj, part) - return obj - - raise NotImplementedError(f"Unknown target kind: {kind}") - -def serialize_graphmodule(gm: fx.GraphModule) -> str: - nodes = [] - for n in gm.graph.nodes: - nodes.append({ - "name": n.name, - "op": n.op, - "target": serialize_target(n.target) if n.op in ("call_function", "call_method", "call_module", "get_attr") else None, - "args": encode_arg(n.args), - "kwargs": encode_arg(n.kwargs), - }) - - # Serialize all direct child modules that are GraphModules - submodules = {} - for name, module in gm.named_children(): - if isinstance(module, fx.GraphModule): - # Recursively serialize GraphModule submodules - submodules[name] = serialize_graphmodule(module) - - data = { - "nodes": nodes, - "state_dict": {k: v.detach().cpu().tolist() for k, v in gm.state_dict().items()}, - # save which device parameters were on, optional: - "param_devices": {k: str(v.device) for k, v in gm.state_dict().items()}, - "submodules": submodules, # Add serialized submodules - } - serialized = json.dumps(data, ensure_ascii=False) - - del data - del nodes - - return serialized - -def _unwrap_output_arg(decoded): - # FX stores output as (value,), where value may itself be a tuple. - # Inductor expects the inner tuple directly. - if isinstance(decoded, (tuple, list)) and len(decoded) == 1 and isinstance(decoded[0], (tuple, list)): - return decoded[0] - return decoded - -def deserialize_graphmodule(s: str) -> fx.GraphModule: - data = json.loads(s) - g = fx.Graph() - name_to_node = {} - - for n in data["nodes"]: - op = n["op"] - if op == "placeholder": - node = g.placeholder(n["name"]) - elif op == "output": - decoded = decode_arg(n["args"], name_to_node) - node = g.output(_unwrap_output_arg(decoded)) - elif op == "call_function": - target = deserialize_target(n["target"]) - args = decode_arg(n["args"], name_to_node) - kwargs = decode_arg(n["kwargs"], name_to_node) - node = g.call_function(target, tuple(args), kwargs) - elif op == "call_method": - target = deserialize_target(n["target"]) - args = decode_arg(n["args"], name_to_node) - kwargs = decode_arg(n["kwargs"], name_to_node) - node = g.call_method(target, tuple(args), kwargs) - elif op == "call_module": - target = deserialize_target(n["target"]) - args = decode_arg(n["args"], name_to_node) - kwargs = decode_arg(n["kwargs"], name_to_node) - node = g.call_module(target, tuple(args), kwargs) - elif op == "get_attr": - target = deserialize_target(n["target"]) - node = g.get_attr(target) - else: - raise NotImplementedError(f"op {op} not handled") - name_to_node[n["name"]] = node - - # Create root module and add submodules before creating GraphModule - root_module = torch.nn.Module() - - # Deserialize and add submodules - submodules = data.get("submodules", {}) - for module_name, serialized_submodule in submodules.items(): - # Recursively deserialize submodule - submodule_gm = deserialize_graphmodule(serialized_submodule) - # Add as direct child module (handles both simple names and nested paths) - # For nested paths like "layer.expert_0", we need to create intermediate modules - parts = module_name.split(".") - if len(parts) == 1: - # Simple name, add directly - root_module.add_module(module_name, submodule_gm) - else: - # Nested path, create intermediate modules - current = root_module - for part in parts[:-1]: - if not hasattr(current, part): - current.add_module(part, torch.nn.Module()) - current = getattr(current, part) - current.add_module(parts[-1], submodule_gm) - - gm = fx.GraphModule(root_module, g) - state = {k: torch.tensor(v) for k, v in data["state_dict"].items()} - gm.load_state_dict(state, strict=False) - return gm \ No newline at end of file diff --git a/src/runtime.py b/src/runtime.py new file mode 100644 index 0000000..3c5cc37 --- /dev/null +++ b/src/runtime.py @@ -0,0 +1,495 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import torch +import torch.distributed as dist +from concurrent.futures import Future, ThreadPoolExecutor + + +@dataclass +class BufferStore: + """Per-iteration task outputs and their consumer refcounts.""" + + task: dict[Any, Any] = field(default_factory=dict) + refcounts: dict[Any, int] = field(default_factory=dict) + + def reset(self) -> None: + self.task.clear() + self.refcounts.clear() + + def init_refcounts(self, dag: Any) -> None: + nodes_iter = dag.nodes.values() if isinstance(dag.nodes, dict) else dag.nodes + self.refcounts = { + node.uid: len(node.data_succs) + for node in nodes_iter + if node.data_succs + } + + def release(self, uid: Any) -> None: + remaining = self.refcounts.get(uid) + if remaining is None: + self.task.pop(uid, None) + return + + remaining -= 1 + if remaining <= 0: + self.refcounts.pop(uid, None) + self.task.pop(uid, None) + else: + self.refcounts[uid] = remaining + + +@dataclass +class EventStore: + """CUDA events produced by non-compute DAG tasks during one iteration.""" + + recv: dict[Any, Any] = field(default_factory=dict) + a2a: dict[Any, Any] = field(default_factory=dict) + all_reduce: dict[Any, Any] = field(default_factory=dict) + reduce_scatter: dict[Any, Any] = field(default_factory=dict) + all_gather: dict[Any, Any] = field(default_factory=dict) + backward: dict[Any, Any] = field(default_factory=dict) + + def reset(self) -> None: + self.recv.clear() + self.a2a.clear() + self.all_reduce.clear() + self.reduce_scatter.clear() + self.all_gather.clear() + self.backward.clear() + + +@dataclass +class BucketState: + """Loaded runtime state for one globally unique compute bucket.""" + + forward_fn: Any = None + forward_args: list[Any] = field(default_factory=list) + forward_input_meta: list[Any] = field(default_factory=list) + input_idxs: list[int] = field(default_factory=list) + param_idxs: list[int] = field(default_factory=list) + param_names: list[str] = field(default_factory=list) + optimizer: Any = None + trainable_param_idxs: list[int] = field(default_factory=list) + activation_checkpoint_subgraph_count: int = 1 + + flat_params: Any = None + flat_grads: Any = None + shard_param: Any = None + shard_optimizer: Any = None + reduce_scatter_grads: Any = None + param_shard_info: tuple[int, int, int] | None = None + param_view_specs: list[Any] = field(default_factory=list) + full_params_fresh: bool = False + + def weights(self) -> list[Any]: + return [ + self.forward_args[idx] + for idx in self.param_idxs + if self.forward_args[idx] is not None + ] + + def trainable_params(self) -> list[Any]: + return [ + self.forward_args[idx] + for idx in self.trainable_param_idxs + if self.forward_args[idx] is not None + ] + + +@dataclass +class StageStore: + """Loaded stage and bucket state owned by one actor.""" + + graph_modules: dict[Any, Any] = field(default_factory=dict) + buckets: dict[Any, BucketState] = field(default_factory=dict) + + param_sharded_ubids: set[Any] = field(default_factory=set) + grad_sharded_ubids: set[Any] = field(default_factory=set) + zero_managed_ubids: set[Any] = field(default_factory=set) + + def clear_loaded_modules(self) -> None: + self.graph_modules.clear() + self.buckets.clear() + + def ensure_bucket(self, ubid: Any) -> BucketState: + return self.buckets.setdefault(ubid, BucketState()) + + def bucket(self, ubid: Any) -> BucketState: + try: + return self.buckets[ubid] + except KeyError as exc: + raise KeyError(f"Unknown bucket_key {ubid!r}") from exc + + def get_bucket(self, ubid: Any) -> BucketState | None: + return self.buckets.get(ubid) + + +@dataclass +class RuntimeState: + """Actor-local distributed runtime environment.""" + + pp_rank: int + dp_rank: int + dp_degree: int + pp_degree: int + world_size: int + no_nvtx: bool = False + device: str = "cuda" + dp_group: Any = None + ep_group: Any = None + pp_lo_hi: Any = None + pp_hi_lo: Any = None + streams: dict[str, torch.cuda.Stream] = field(default_factory=dict) + pytorch_profiler_enabled: bool = False + torch_profiler: Any = None + + @property + def global_rank(self) -> int: + return self.pp_rank + self.dp_rank * self.pp_degree + + def pipeline_peer_global_rank(self, pp_rank: int) -> int: + return pp_rank + self.dp_rank * self.pp_degree + + def stream_id(self, node_or_stream: Any) -> str: + if isinstance(node_or_stream, str): + return node_or_stream + return str(getattr(node_or_stream, "stream", "default_stream")) + + def initialize_streams_for_training_dag(self, training_dag: Any) -> None: + stream_ids = { + self.stream_id(n) + for n in training_dag.nodes.values() + if getattr(n, "stream", None) is not None + } + stream_ids.add("default_stream") + self.streams = { + stream_id: torch.cuda.Stream(device=self.device) + for stream_id in sorted(stream_ids) + } + + # Force cuBLAS context initialization on every logical stream used by + # this DAG so the first backward pass does not hit lazy CUDA warnings. + for stream in self.streams.values(): + with torch.cuda.stream(stream): + w = torch.zeros(4, 4, device=self.device) + torch.mm(w, w) + + def stream_for_id(self, stream_id: str) -> torch.cuda.Stream: + assert stream_id in self.streams, ( + f"TrainingDAG referenced stream={stream_id!r}, but load_training_dag " + f"initialized only {sorted(self.streams)}" + ) + return self.streams[stream_id] + + def stream_for_task(self, task: Any) -> torch.cuda.Stream: + return self.stream_for_id(self.stream_id(task)) + + def default_stream(self) -> torch.cuda.Stream: + return self.stream_for_id("default_stream") + + def nvtx_push(self, label: str) -> None: + if not self.no_nvtx: + torch.cuda.nvtx.range_push(label) + + def nvtx_pop(self) -> None: + if not self.no_nvtx: + torch.cuda.nvtx.range_pop() + + +@dataclass +class ParamStorage: + """ZeRO parameter and gradient storage owned by one actor.""" + + runtime: RuntimeState + stages: StageStore + logger: Any + grad_buffer_dtype: torch.dtype = torch.float32 + cleanup_executor: ThreadPoolExecutor = field( + default_factory=lambda: ThreadPoolExecutor(max_workers=1) + ) + pending_param_frees: dict[Any, Future] = field(default_factory=dict) + pending_grad_frees: dict[Any, Future] = field(default_factory=dict) + + def clear_param_grads(self) -> None: + for bucket in self.stages.buckets.values(): + for idx in bucket.trainable_param_idxs: + param = bucket.forward_args[idx] + if param is not None: + param.grad = None + + def zero_grad_buffers(self, stream: torch.cuda.Stream) -> None: + with torch.cuda.stream(stream): + for bucket in self.stages.buckets.values(): + if bucket.flat_grads is not None: + bucket.flat_grads.zero_() + if bucket.reduce_scatter_grads is not None: + bucket.reduce_scatter_grads.zero_() + + def wait_pending_free( + self, + pending: dict[Any, Future], + ubid: Any | None, + ) -> None: + if ubid is None: + return + fut = pending.pop(ubid, None) + if fut is not None: + fut.result() + + def drain_pending_frees(self) -> None: + for pending in (self.pending_param_frees, self.pending_grad_frees): + futures = list(pending.values()) + pending.clear() + for fut in futures: + fut.result() + + def accumulate_zero_param_grads_to_flat( + self, + ubid: Any | None, + stream: torch.cuda.Stream, + ) -> None: + if ubid is None or ubid not in self.stages.zero_managed_ubids or self.runtime.dp_degree <= 1: + return + if ubid in self.stages.grad_sharded_ubids: + self.wait_pending_free(self.pending_grad_frees, ubid) + bucket = self.stages.bucket(ubid) + specs = bucket.param_view_specs + if not specs: + return + with torch.cuda.stream(stream): + flat_grads = bucket.flat_grads + if flat_grads is None: + shard_info = bucket.param_shard_info + if shard_info is None: + return + _shard_start, shard_size, _orig_numel = shard_info + flat_grads = torch.zeros( + shard_size * self.runtime.dp_degree, + dtype=self.grad_buffer_dtype, + device=self.runtime.device, + ) + bucket.flat_grads = flat_grads + for param, offset, numel, _shape in specs: + grad = param.grad + if grad is None: + continue + flat_grads[offset:offset + numel].add_( + grad.detach().reshape(-1).to(flat_grads.dtype) + ) + param.grad = None + + def defer_free_full_params(self, ubid: Any | None, evt: torch.cuda.Event) -> None: + if ubid is None or ubid not in self.stages.param_sharded_ubids: + return + self.wait_pending_free(self.pending_param_frees, ubid) + self.pending_param_frees[ubid] = self.cleanup_executor.submit( + self._wait_then_free_full_params, + ubid, + evt, + ) + + def defer_free_full_grads(self, ubid: Any | None, evt: torch.cuda.Event) -> None: + if ubid is None or ubid not in self.stages.grad_sharded_ubids: + return + self.wait_pending_free(self.pending_grad_frees, ubid) + self.pending_grad_frees[ubid] = self.cleanup_executor.submit( + self._wait_then_free_full_grads, + ubid, + evt, + ) + + def _wait_then_free_full_params(self, ubid: Any, evt: torch.cuda.Event) -> None: + evt.synchronize() + self.free_full_params(ubid) + + def _wait_then_free_full_grads(self, ubid: Any, evt: torch.cuda.Event) -> None: + evt.synchronize() + self.free_full_grads(ubid) + + def alloc_full_params(self, ubid: Any) -> None: + assert ubid is not None, "alloc_full_params requires a non-None ubid" + assert ubid in self.stages.param_sharded_ubids, ( + f"alloc_full_params: ubid={ubid} is not in param_sharded_ubids=" + f"{self.stages.param_sharded_ubids}" + ) + bucket = self.stages.bucket(ubid) + self.wait_pending_free(self.pending_param_frees, ubid) + assert bucket.param_shard_info is not None, ( + f"alloc_full_params: missing param_shard_info for ubid={ubid}" + ) + full = bucket.flat_params + assert full is not None, ( + f"alloc_full_params: missing flat_params buffer for ubid={ubid}" + ) + specs = bucket.param_view_specs + assert specs, f"alloc_full_params: missing param_view_specs for ubid={ubid}" + storage = full.untyped_storage() + required_bytes = full.numel() * full.element_size() + storage.resize_(required_bytes) + self.logger.debug( + f"[alloc_full_params] rank={self.runtime.global_rank} ubid={ubid}: " + f"numel={full.numel()} required_bytes={required_bytes} " + f"storage_size={storage.size()} fresh={bucket.full_params_fresh}" + ) + for param, offset, numel, shape in specs: + param.data = full[offset:offset + numel].view(shape) + param.requires_grad_(True) + zero_storage = [] + for i, (param, offset, numel, shape) in enumerate(specs): + p_storage = param.untyped_storage() + if p_storage.size() == 0: + name = bucket.param_names[i] if i < len(bucket.param_names) else f"param{i}" + zero_storage.append( + f"{name}: offset={offset} numel={numel} " + f"shape={shape} stride={tuple(param.stride())}" + ) + assert not zero_storage, ( + f"[alloc_full_params_zero_storage] rank={self.runtime.global_rank} " + f"ubid={ubid}: " + " | ".join(zero_storage) + ) + + def free_full_params(self, ubid: Any) -> None: + assert ubid is not None, "free_full_params requires a non-None ubid" + assert ubid in self.stages.param_sharded_ubids, ( + f"free_full_params: ubid={ubid} is not in param_sharded_ubids=" + f"{self.stages.param_sharded_ubids}" + ) + bucket = self.stages.bucket(ubid) + full = bucket.flat_params + assert full is not None, ( + f"free_full_params: missing flat_params buffer for ubid={ubid}" + ) + storage = full.untyped_storage() + self.logger.debug( + f"[free_full_params] rank={self.runtime.global_rank} ubid={ubid}: " + f"storage_size_before={storage.size()} fresh_before={bucket.full_params_fresh}" + ) + storage.resize_(0) + bucket.full_params_fresh = False + + def alloc_full_grads(self, ubid: Any, stream: torch.cuda.Stream) -> None: + assert ubid is not None, "alloc_full_grads requires a non-None ubid" + assert ubid in self.stages.grad_sharded_ubids, ( + f"alloc_full_grads: ubid={ubid} is not in grad_sharded_ubids=" + f"{self.stages.grad_sharded_ubids}" + ) + bucket = self.stages.bucket(ubid) + self.wait_pending_free(self.pending_grad_frees, ubid) + specs = bucket.param_view_specs + assert specs, f"alloc_full_grads: missing param_view_specs for ubid={ubid}" + shard_info = bucket.param_shard_info + assert shard_info is not None, ( + f"alloc_full_grads: missing param_shard_info for ubid={ubid}" + ) + shard_size = shard_info[1] + with torch.cuda.stream(stream): + if bucket.flat_grads is None: + bucket.flat_grads = torch.zeros( + shard_size * self.runtime.dp_degree, + dtype=self.grad_buffer_dtype, + device=self.runtime.device, + ) + if bucket.reduce_scatter_grads is None: + bucket.reduce_scatter_grads = torch.zeros( + shard_size, + dtype=self.grad_buffer_dtype, + device=self.runtime.device, + ) + + def free_full_grads(self, ubid: Any) -> None: + assert ubid is not None, "free_full_grads requires a non-None ubid" + assert ubid in self.stages.grad_sharded_ubids, ( + f"free_full_grads: ubid={ubid} is not in grad_sharded_ubids=" + f"{self.stages.grad_sharded_ubids}" + ) + bucket = self.stages.bucket(ubid) + specs = bucket.param_view_specs + assert specs, f"free_full_grads: missing param_view_specs for ubid={ubid}" + for param, *_ in specs: + param.grad = None + bucket.flat_grads = None + + def all_gather_full_params(self, ubid: Any, stream: torch.cuda.Stream) -> int: + assert ubid is not None, "all_gather_full_params requires a non-None ubid" + if not self._has_trainable_params_for_collective(ubid, "all_gather_full_params"): + return 0 + assert ubid in self.stages.param_sharded_ubids, ( + f"all_gather_full_params: ubid={ubid} is not in param_sharded_ubids=" + f"{self.stages.param_sharded_ubids}" + ) + bucket = self.stages.bucket(ubid) + self.alloc_full_params(ubid) + flat_params = bucket.flat_params + shard_in = bucket.shard_param + assert flat_params is not None, ( + f"all_gather_full_params: missing flat_params buffer for ubid={ubid}" + ) + assert shard_in is not None, ( + f"all_gather_full_params: missing shard_param buffer for ubid={ubid}" + ) + assert not bucket.full_params_fresh, ( + f"all_gather_full_params: ubid={ubid} dispatched but full params are already fresh; " + "DAG is constructing a redundant ALL_GATHER" + ) + with torch.cuda.stream(stream): + self.logger.debug( + f"[all_gather_begin] rank={self.runtime.global_rank} ubid={ubid}: " + f"flat_numel={flat_params.numel()} flat_storage={flat_params.untyped_storage().size()} " + f"shard_numel={shard_in.numel()} shard_storage={shard_in.untyped_storage().size()}" + ) + dist.all_gather_into_tensor(flat_params, shard_in, group=self.runtime.dp_group) + bucket.full_params_fresh = True + self.logger.debug( + f"[all_gather_end] rank={self.runtime.global_rank} ubid={ubid}: " + f"flat_storage={flat_params.untyped_storage().size()}" + ) + return flat_params.numel() * flat_params.element_size() + + def has_zero_shard_optimizers(self) -> bool: + return any(bucket.shard_optimizer is not None for bucket in self.stages.buckets.values()) + + def step_zero_shard_optimizers( + self, + stream: torch.cuda.Stream, + reduce_scatter_events: dict[Any, torch.cuda.Event], + ) -> None: + for evt in reduce_scatter_events.values(): + stream.wait_event(evt) + + for bucket in self.stages.buckets.values(): + shard_optim = bucket.shard_optimizer + if shard_optim is None: + continue + shard_param = bucket.shard_param + rs_grad = bucket.reduce_scatter_grads + if rs_grad is None: + continue + with torch.cuda.stream(stream): + shard_param.grad = rs_grad.to(shard_param.dtype) + shard_optim.step() + shard_param.grad = None + + for ubid in self.stages.param_sharded_ubids: + bucket = self.stages.bucket(ubid) + for param, *_ in bucket.param_view_specs: + param.grad = None + full = bucket.flat_params + if full is not None: + storage = full.untyped_storage() + if storage.size() != 0: + storage.resize_(0) + bucket.full_params_fresh = False + + def _has_trainable_params_for_collective(self, ubid: Any, op_name: str) -> bool: + bucket = self.stages.get_bucket(ubid) + if bucket is not None and bucket.trainable_param_idxs: + return True + self.logger.warning( + "%s: skipping collective for ubid=%s because it has no trainable param indices", + op_name, + ubid, + ) + return False diff --git a/src/schedule.py b/src/schedule.py new file mode 100644 index 0000000..dbfd6d8 --- /dev/null +++ b/src/schedule.py @@ -0,0 +1,105 @@ +import json +import os +from typing import Any + + +def load_schedule_directives(path: str) -> list[dict]: + with open(path, "r", encoding="utf-8") as f: + directives = json.load(f) + if not isinstance(directives, list): + raise ValueError(f"schedule directives file must contain a list, got: {type(directives)}") + + for i, directive in enumerate(directives): + if not isinstance(directive, dict): + raise ValueError(f"schedule directive[{i}] must be an object, got: {type(directive)}") + _validate_directive_shape(directive, i) + if directive.get("op") == "split" and directive.get("num_microbatches") == "__MBS__": + raise ValueError( + "split.num_microbatches must be encoded in the schedule JSON; " + "'__MBS__' is no longer supported" + ) + return directives + + +def load_schedule_info(path: str) -> dict[str, Any]: + return derive_schedule_info(load_schedule_directives(path), path) + + +def derive_schedule_info(directives: list[dict], schedule_path: str) -> dict[str, Any]: + pp_to_devices: dict[int, list[int]] = {} + num_microbatches = None + + for directive in directives: + if not isinstance(directive, dict): + continue + op = directive.get("op") + if op == "place": + pp_idx = _filter_value(directive.get("filter"), "PP") + devices = directive.get("devices", directive.get("device")) + if pp_idx is None or not isinstance(devices, list) or not devices: + raise ValueError(f"place directive must include PP filter and non-empty devices: {directive}") + pp_to_devices[int(pp_idx)] = [int(d) for d in devices] + elif op == "split": + n = int(directive.get("num_microbatches", 0)) + if n <= 0: + raise ValueError(f"split directive requires num_microbatches > 0: {directive}") + if num_microbatches is not None and num_microbatches != n: + raise ValueError( + f"multiple split directives disagree on num_microbatches: " + f"{num_microbatches} vs {n}" + ) + num_microbatches = n + + if not pp_to_devices: + raise ValueError("schedule JSON must include place directives with PP filters") + pp_indices = sorted(pp_to_devices) + expected_pp = list(range(len(pp_indices))) + if pp_indices != expected_pp: + raise ValueError(f"PP indices must be contiguous from 0, got {pp_indices}") + device_counts = {len(devices) for devices in pp_to_devices.values()} + if len(device_counts) != 1: + raise ValueError(f"all PP place directives must use the same device count, got {pp_to_devices}") + device_keys = sorted({tuple(devices) for devices in pp_to_devices.values()}) + if num_microbatches is None: + raise ValueError("schedule JSON must include a split directive with num_microbatches") + + return { + "name": os.path.splitext(os.path.basename(schedule_path))[0], + "path": schedule_path, + "num_stages": len(pp_indices), + "pp_degree": len(device_keys), + "dp_degree": next(iter(device_counts)), + "num_microbatches": num_microbatches, + } + + +def _validate_directive_shape(directive: dict, idx: int) -> None: + op = directive.get("op") + if op in {"place", "replicate", "shard", "split"}: + if not isinstance(directive.get("filter"), dict): + raise ValueError(f"{op} directive[{idx}] requires object field 'filter': {directive}") + if "filters" in directive: + raise ValueError(f"{op} directive[{idx}] does not accept field 'filters': {directive}") + elif op == "order": + filters = directive.get("filters") + if not isinstance(filters, list) or not filters: + raise ValueError(f"order directive[{idx}] requires non-empty list field 'filters': {directive}") + for group_idx, group in enumerate(filters): + if not isinstance(group, list) or not group: + raise ValueError( + f"order directive[{idx}] group[{group_idx}] must be a non-empty list: {directive}" + ) + for filter_idx, flt in enumerate(group): + if not isinstance(flt, dict): + raise ValueError( + f"order directive[{idx}] group[{group_idx}][{filter_idx}] " + f"must be a filter object: {directive}" + ) + else: + raise ValueError(f"schedule directive[{idx}] has unsupported op: {directive}") + + +def _filter_value(filter_spec, key: str): + if isinstance(filter_spec, dict): + return filter_spec.get(key) + return None diff --git a/src/state.py b/src/state.py new file mode 100644 index 0000000..815bdc5 --- /dev/null +++ b/src/state.py @@ -0,0 +1,52 @@ +import sys +import logging +from typing import Optional + +LOG_LEVEL = "INFO" + +""" +Logger utility +""" + +def create_logger(name: str, log_level: str): + match log_level: + case "DEBUG": + log_level = logging.DEBUG + case "INFO": + log_level = logging.INFO + case "WARNING": + log_level = logging.WARNING + case "ERROR": + log_level = logging.ERROR + + logger = logging.getLogger(name) + logger.setLevel(log_level) + + if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + fmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" + handler.setFormatter(logging.Formatter(fmt)) + logger.addHandler(handler) + logger.propagate = False + + return logger + + + + +""" +Piper thread local storage for tracking Piper actors, stages, and microbatches +""" + +class PiperMetadata: + actors = dict() + visualize_dag: bool = False # Whether to render per-rank DAG PNGs after compilation + artifact_dir: str = "out" # Directory for debug artifacts emitted during runs + training_dag = None # DAG of annotated model segments and transform-inserted nodes + per_pp_training_dags = None # Per-PP-rank DAGs built by the TrainingDAG backend + compiled_data_store = None # Ray actor used to share compiled DAGs across DP ranks + schedule_directives: list = [] # Program of DAG transform directives (e.g., place(...)) + schedule_directives_file: Optional[str] = None # JSON source for schedule_directives + schedule_info: dict = {} # Derived schedule facts such as pp/dp/mbs + +piper_metadata = PiperMetadata() diff --git a/src/tasks.py b/src/tasks.py new file mode 100644 index 0000000..d84e582 --- /dev/null +++ b/src/tasks.py @@ -0,0 +1,51 @@ +from enum import Enum +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .dag import TrainingDAGNode + + +class TaskType(Enum): + FWD = "forward" + BWD = "backward" + UPD = "update" + BWD_I = "backward_input" + BWD_W = "backward_weight" + FWD_BWD = "forward_backward" + SEND = "send" + RECV = "recv" + ALL_REDUCE = "all_reduce" + REDUCE_SCATTER = "reduce_scatter" + ALL_GATHER = "all_gather" + FWD_A2A = "forward_a2a" + BWD_A2A = "backward_a2a" + ORDER_DUMMY = "order_dummy" + + +def training_dag_task_type(node: "TrainingDAGNode") -> TaskType: + if node.node_kind == "COMPUTE": + if node.compute_subkind == "FWD": + return TaskType.FWD + if node.compute_subkind == "BWD": + return TaskType.BWD + if node.compute_subkind == "BWD_I": + return TaskType.BWD_I + if node.compute_subkind == "BWD_W": + return TaskType.BWD_W + if node.node_kind == "UPD": + return TaskType.UPD + if node.node_kind == "SEND_COMM": + return TaskType.SEND + if node.node_kind == "RECV_COMM": + return TaskType.RECV + if node.node_kind == "ALL_GATHER_COMM": + return TaskType.ALL_GATHER + if node.node_kind == "REDUCE_SCATTER_COMM": + return TaskType.REDUCE_SCATTER + if node.node_kind == "REDUCE_COMM": + return TaskType.ALL_REDUCE + if node.node_kind == "A2A_COMM": + return TaskType.FWD_A2A if node.tag.get("PASS") == "F" else TaskType.BWD_A2A + if node.node_kind == "ORDER_DUMMY": + return TaskType.ORDER_DUMMY + return TaskType.FWD diff --git a/src/visualization.py b/src/visualization.py new file mode 100644 index 0000000..1131c6a --- /dev/null +++ b/src/visualization.py @@ -0,0 +1,412 @@ +from __future__ import annotations + +import os +from collections.abc import Sequence +from contextlib import suppress +from dataclasses import dataclass +from html import escape +from pathlib import Path + +from .dag import TrainingDAG, TrainingDAGNode, _topological_levels +from .ordering import _serial_topological_order +from .state import LOG_LEVEL, create_logger +from .tasks import training_dag_task_type + +logger = create_logger("visualization", LOG_LEVEL) + +_VALID_PASSES = frozenset({"F", "B", "BI", "BW"}) +_PASS_COLOR: dict[str, str] = { + "F": "#FFE08A", + "B": "#A8E6A1", + "BI": "#A9D6FF", + "BW": "#C6B6FF", +} + + +@dataclass(frozen=True) +class _ScheduleOp: + pp: int + mb: int + pass_: str + + +_ScheduleSlot = list[_ScheduleOp] + + +def _format_tag(tag: dict[str, int | None]) -> str: + if not tag: + return "(no-tags)" + keys = sorted(tag.keys(), key=lambda k: (k != "PP", k != "EP", k)) + return ", ".join(f"{k}={tag[k]}" for k in keys) + + +def _format_device(device: list[int] | None) -> str: + return "None" if device is None else str(list(device)) + + +def render_training_dag(training_dag: TrainingDAG, output_path: str = "out/training_dag") -> None: + """Render TrainingDAG with node labels as tags and data-dependency edges.""" + import graphviz + + dot = graphviz.Digraph("TrainingDAG", comment="Training DAG (tag-labelled)") + dot.attr(rankdir="LR") + topo_levels = _topological_levels(training_dag) + + def _format_node_meta_lines(node: TrainingDAGNode) -> str: + keys = ( + "bucket_key", + "zero_alloc_full_grads_before", + "zero_free_full_params_after", + ) + lines = [ + f"{key}={node.node_meta[key]!r}" + for key in keys + if key in node.node_meta + ] + return "\\n".join(lines) + + for uid, node in training_dag.nodes.items(): + topo_label = f"topo={topo_levels.get(uid, '?')}" + meta_label = _format_node_meta_lines(node) + meta_suffix = f"\\n{meta_label}" if meta_label else "" + if node.node_kind in ( + "SEND_COMM", + "RECV_COMM", + "REDUCE_COMM", + "ALL_GATHER_COMM", + "REDUCE_SCATTER_COMM", + "A2A_COMM", + ): + label = ( + f"{topo_label}\\n{node.node_kind}\\n{_format_tag(node.tag)}\\n" + f"{_format_device(node.device)}\\nstream={node.stream}" + f"{meta_suffix}" + ) + shape = "box" + else: + subkind = node.compute_subkind or "COMPUTE" + label = ( + f"{topo_label}\\n{subkind}\\n{_format_tag(node.tag)}\\n{_format_device(node.device)}\\nstream={node.stream}" + f"{meta_suffix}" + ) + shape = "ellipse" + dot.node(uid, label=label, shape=shape) + + nodes_by_level: dict[int, list[str]] = {} + for uid, level in topo_levels.items(): + nodes_by_level.setdefault(level, []).append(uid) + for level, uids in sorted(nodes_by_level.items()): + with dot.subgraph(name=f"topo_{level}") as sub: + sub.attr(rank="same") + for uid in sorted(uids): + sub.node(uid) + level_reps = [ + sorted(uids)[0] + for _level, uids in sorted(nodes_by_level.items()) + if uids + ] + for src_uid, dst_uid in zip(level_reps, level_reps[1:]): + dot.edge(src_uid, dst_uid, style="invis", weight="100") + + for edge in training_dag.edges: + if edge.dep_kind == "data": + dot.edge(edge.src_uid, edge.dst_uid) + elif edge.dep_kind == "temporal": + dot.edge(edge.src_uid, edge.dst_uid, style="dashed", color="blue") + + out = dot.render(output_path, format="png", cleanup=True) + logger.info("TrainingDAG debug graph saved to %s", out) + + +def log_training_dag_dependencies(training_dag: TrainingDAG) -> None: + """Log predecessor/successor tags for each TrainingDAG node.""" + def _sort_key(item: tuple[str, TrainingDAGNode]) -> tuple[int, int]: + _uid, node = item + if node.node_kind == "COMPUTE": + return (int(node.node_meta.get("stage_id", 10**9)), int(node.node_meta.get("segment_id", 10**9))) + return (10**9, 10**9) + + for uid, node in sorted(training_dag.nodes.items(), key=_sort_key): + pred_tags = [ + _format_tag(training_dag.nodes[p].tag) + for p in sorted(training_dag.preds.get(uid, []), key=lambda x: x) + ] + succ_tags = [ + _format_tag(training_dag.nodes[s].tag) + for s in sorted(training_dag.succs.get(uid, []), key=lambda x: x) + ] + logger.debug( + "TrainingDAG node=%s kind=%s subkind=%s tag=(%s) device=%s stream=%s preds=[%s] succs=[%s]", + uid, + node.node_kind, + node.compute_subkind, + _format_tag(node.tag), + _format_device(node.device), + node.stream, + ", ".join(pred_tags), + ", ".join(succ_tags), + ) + + +def print_training_dag_order( + dag: TrainingDAG, + label: str = "", + rank: int = 0, + out_dir: str = "out", +) -> None: + """Write and log the actor dispatch order for a TrainingDAG.""" + topo_levels = _topological_levels(dag) + topo = _serial_topological_order(dag, topo_levels) + data_preds: dict[str, list[str]] = {uid: [] for uid in dag.nodes} + data_succs: dict[str, list[str]] = {uid: [] for uid in dag.nodes} + temporal_preds: dict[str, list[str]] = {uid: [] for uid in dag.nodes} + temporal_succs: dict[str, list[str]] = {uid: [] for uid in dag.nodes} + for e in dag.edges: + if e.dep_kind == "data": + data_preds[e.dst_uid].append(e.src_uid) + data_succs[e.src_uid].append(e.dst_uid) + else: + temporal_preds[e.dst_uid].append(e.src_uid) + temporal_succs[e.src_uid].append(e.dst_uid) + + header = f"--- TrainingDAG execution order{': ' + label if label else ''} ---" + lines = [header] + for uid in topo: + node = dag.nodes[uid] + task_type = training_dag_task_type(node) + meta = node.node_meta + extra = [] + for key in ("bucket_key", "fwd_uid", "peer_pp_rank", "from_uid", "to_uid"): + if key in meta and meta[key] is not None: + extra.append(f"{key}={meta[key]}") + extra_str = (" " + " ".join(extra)) if extra else "" + line = ( + f" topo={topo_levels[uid]:<3d} " + f"{task_type.value:<14s} kind={node.node_kind:<19s} uid={node.uid:<24s} " + f"tag=({_format_tag(node.tag)}) stream={node.stream}{extra_str} " + f"data_preds={sorted(data_preds[node.uid])} data_succs={sorted(data_succs[node.uid])} " + f"temp_preds={sorted(temporal_preds[node.uid])} temp_succs={sorted(temporal_succs[node.uid])}" + ) + lines.append(line) + + lines.append("-" * len(header)) + + os.makedirs(out_dir, exist_ok=True) + out_path = os.path.join(out_dir, f"training_dag_order_rank{rank}") + with open(out_path, "w", encoding="utf-8") as f: + f.write("\n".join(lines) + "\n") + logger.info("TrainingDAG execution order saved to %s", out_path) + + +def visualize_order_directives( + order_directives: list[dict], + output_path: str | Path = "out/schedule", + fmt: str = "png", +) -> str: + """Render a 2-D schedule grid, one row per physical rank.""" + rows = _rows_from_order_directives(order_directives) + placements = _compute_columns(rows) + max_col = max( + (col + _slot_width(rows[rank][slot_idx]) - 1 for (rank, slot_idx), col in placements.items()), + default=-1, + ) + + table: list[list[_ScheduleSlot | None]] = [ + [None for _ in range(max_col + 1)] + for _ in rows + ] + for key, col in placements.items(): + rank, slot_idx = key + table[rank][col] = rows[rank][slot_idx] + + output = Path(output_path) + if output.suffix: + fmt = output.suffix.lstrip(".") + output = output.with_suffix("") + + output.parent.mkdir(parents=True, exist_ok=True) + rendered_suffix = f".{fmt}" + _cleanup_visualizer_sidecars(output, keep_suffix=rendered_suffix) + import graphviz + + dot = graphviz.Digraph("PipelineSchedule") + dot.attr(rankdir="LR") + dot.node("schedule", label=_schedule_html_table(table), shape="plain") + try: + return dot.render(str(output), format=fmt, cleanup=True) + finally: + _cleanup_visualizer_sidecars(output, keep_suffix=rendered_suffix) + + +def _op_width(op: _ScheduleOp) -> int: + if op.pass_ == "B": + return 2 + return 1 + + +def _slot_width(slot: _ScheduleSlot) -> int: + return sum(_op_width(op) for op in slot) + + +def _cleanup_visualizer_sidecars( + output_base: Path, + keep_suffix: str | None = None, +) -> None: + for suffix in (".dot", ".txt"): + if suffix == keep_suffix: + continue + with suppress(FileNotFoundError): + output_base.with_suffix(suffix).unlink() + + +def _rows_from_order_directives(order_directives: list[dict]) -> list[list[_ScheduleSlot]]: + rows = [] + for directive in order_directives: + if directive.get("op") != "order": + raise ValueError(f"expected order directive, got {directive}") + filters = directive.get("filters") + if not isinstance(filters, list): + raise ValueError(f"order directive requires filters list: {directive}") + row: list[_ScheduleSlot] = [] + for filter_slot in _parse_filter_slots(filters): + slot: _ScheduleSlot = [] + for flt in filter_slot: + spec = _filter_to_dict(flt) + pass_value = spec["PASS"] + if pass_value not in _VALID_PASSES: + raise ValueError( + f"unsupported PASS={pass_value!r}; expected one of {sorted(_VALID_PASSES)}" + ) + slot.append(_ScheduleOp(pp=int(spec["PP"]), mb=int(spec["MB"]), pass_=pass_value)) + row.append(slot) + rows.append(row) + return rows + + +def _parse_filter_slots(filters: Sequence[object]) -> list[list[dict]]: + """Group the directive's filter entries into visual schedule slots.""" + slots: list[list[dict]] = [] + for item in filters: + if not isinstance(item, list) or not item: + raise ValueError(f"invalid order filter group: {item}") + if not all(isinstance(flt, dict) for flt in item): + raise ValueError(f"invalid order filter group: {item}") + slots.append(list(item)) + return slots + + +def _filter_to_dict(flt: object) -> dict: + if isinstance(flt, dict): + return dict(flt) + raise ValueError(f"invalid order filter: {flt}") + + +def _compute_columns(rows: list[list[_ScheduleSlot]]) -> dict[tuple[int, int], int]: + key_by_op: dict[tuple[int, int, str], tuple[int, int]] = {} + for rank, row in enumerate(rows): + for slot_idx, slot in enumerate(row): + for op in slot: + op_key = (op.pp, op.mb, op.pass_) + if op_key in key_by_op: + raise ValueError(f"duplicate scheduled op: {op_key}") + key_by_op[op_key] = (rank, slot_idx) + + preds: dict[tuple[int, int], set[tuple[int, int]]] = { + (rank, slot_idx): set() + for rank, row in enumerate(rows) + for slot_idx, _slot in enumerate(row) + } + succs: dict[tuple[int, int], set[tuple[int, int]]] = { + key: set() for key in preds + } + + def add_edge(src: tuple[int, int] | None, dst: tuple[int, int] | None) -> None: + if src is None or dst is None or src == dst: + return + preds[dst].add(src) + succs[src].add(dst) + + for rank, row in enumerate(rows): + for slot_idx in range(1, len(row)): + add_edge((rank, slot_idx - 1), (rank, slot_idx)) + + for rank, row in enumerate(rows): + for slot_idx, slot in enumerate(row): + dst = (rank, slot_idx) + for op in slot: + if op.pass_ == "F": + add_edge(key_by_op.get((op.pp - 1, op.mb, "F")), dst) + elif op.pass_ == "B": + add_edge(key_by_op.get((op.pp, op.mb, "F")), dst) + add_edge(key_by_op.get((op.pp + 1, op.mb, "B")), dst) + elif op.pass_ == "BI": + add_edge(key_by_op.get((op.pp, op.mb, "F")), dst) + add_edge(key_by_op.get((op.pp + 1, op.mb, "BI")), dst) + elif op.pass_ == "BW": + add_edge(key_by_op.get((op.pp, op.mb, "BI")), dst) + else: + raise ValueError(f"unsupported pass={op.pass_!r}") + + columns = {key: 0 for key in preds} + pending = {key: set(value) for key, value in preds.items()} + ready = [key for key, value in pending.items() if not value] + visited = 0 + + while ready: + key = ready.pop(0) + visited += 1 + for succ in sorted(succs[key]): + rank, slot_idx = key + columns[succ] = max(columns[succ], columns[key] + _slot_width(rows[rank][slot_idx])) + pending[succ].discard(key) + if not pending[succ]: + ready.append(succ) + + if visited != len(preds): + blocked = [key for key, value in pending.items() if value] + raise ValueError(f"schedule dependencies contain a cycle; blocked={blocked[:8]}") + return columns + + +def _schedule_html_table(table: list[list[_ScheduleSlot | None]]) -> str: + max_cols = max((len(row) for row in table), default=0) + lines = ['<'] + lines.append('') + for col in range(max_cols): + lines.append(f'') + lines.append("") + for rank, row in enumerate(table): + lines.append(f'') + col = 0 + while col < len(row): + slot = row[col] + if slot is None: + lines.append('') + col += 1 + else: + colspan = _slot_width(slot) + colspan_attr = f' COLSPAN="{colspan}"' if colspan > 1 else "" + width = 72 * colspan + if len(slot) == 1: + op = slot[0] + color = _PASS_COLOR[op.pass_] + label = f"{escape(op.pass_)}
PP{op.pp} MB{op.mb}" + lines.append( + f'' + ) + else: + inner = ['
rankt{col}
r{rank}{label}
'] + for op in slot: + color = _PASS_COLOR[op.pass_] + label = f"{escape(op.pass_)}
PP{op.pp} MB{op.mb}" + inner.append( + f'' + ) + inner.append("
{label}
") + lines.append( + f'{"".join(inner)}' + ) + col += colspan + lines.append("") + lines.append(">") + return "".join(lines) diff --git a/src/zero.py b/src/zero.py new file mode 100644 index 0000000..5025128 --- /dev/null +++ b/src/zero.py @@ -0,0 +1,369 @@ +from typing import Any + +from .dag import ( + TrainingDAG, + TrainingDAGEdge, + _has_path, + _remove_node_and_incident_edges, + _topological_levels, +) +from .ordering import _serial_topological_order +from .state import LOG_LEVEL, create_logger + +logger = create_logger("zero", LOG_LEVEL) + + +def _prune_zero_lifetime_metadata(dag: TrainingDAG) -> list[list[str]]: + """Prune redundant ZeRO param/grad lifetimes across same-bucket compute chains. + + Returns the union of gradient and parameter chains used for pruning, one + list per chain ordered by topo. The caller passes these to + ``_add_inter_chain_temporal_edges`` *after* + ``resolve_total_order_per_stream`` runs, so the inter-chain check sees the + post-stream-serialization topo levels for the gather_stream and + reduce_stream nodes. + """ + topo_levels = _topological_levels(dag) + topo = _serial_topological_order(dag, topo_levels) + + def _is_compute(uid: str) -> bool: + node = dag.nodes[uid] + return ( + node.node_kind == "COMPUTE" + and node.compute_subkind in ("FWD", "BWD", "BWD_I", "BWD_W") + ) + + def _chain_bucket(uid: str) -> Any: + return dag.nodes[uid].node_meta.get("bucket_key") + + def _compute_uids_with(topo_order: list[str], predicate) -> list[str]: + return [ + uid for uid in topo_order + if uid in dag.nodes and _is_compute(uid) and predicate(uid) + ] + + def _has_direct_compute_dependency(src_uid: str, dst_uid: str) -> bool: + return any( + e.src_uid == src_uid + and e.dst_uid == dst_uid + and e.dep_kind in ("data", "temporal") + and e.src_uid in dag.nodes + and e.dst_uid in dag.nodes + and _is_compute(e.src_uid) + and _is_compute(e.dst_uid) + for e in dag.edges + ) + + def _build_chains(candidates: list[str], topo_order: list[str]) -> list[list[str]]: + chains: list[list[str]] = [] + by_bucket: dict[Any, list[str]] = {} + for uid in candidates: + by_bucket.setdefault(_chain_bucket(uid), []).append(uid) + topo_idx = {uid: i for i, uid in enumerate(topo_order)} + for bucket, uids in sorted(by_bucket.items(), key=lambda item: str(item[0])): + uid_set = set(uids) + links: dict[str, set[str]] = {uid: set() for uid in uids} + for e in dag.edges: + if ( + e.dep_kind in ("data", "temporal") + and e.src_uid in uid_set + and e.dst_uid in uid_set + and _is_compute(e.src_uid) + and _is_compute(e.dst_uid) + ): + links[e.src_uid].add(e.dst_uid) + links[e.dst_uid].add(e.src_uid) + + seen: set[str] = set() + for uid in uids: + if uid in seen: + continue + component: list[str] = [] + stack = [uid] + seen.add(uid) + while stack: + cur = stack.pop() + component.append(cur) + for nxt in sorted(links[cur], key=lambda u: topo_idx[u]): + if nxt not in seen: + seen.add(nxt) + stack.append(nxt) + chains.append(sorted(component, key=lambda u: topo_idx[u])) + return chains + + def _data_preds(uid: str) -> list[TrainingDAGEdge]: + return [ + e for e in list(dag.edges) + if e.dep_kind == "data" and e.dst_uid == uid + ] + + def _data_succs(uid: str) -> list[TrainingDAGEdge]: + return [ + e for e in list(dag.edges) + if e.dep_kind == "data" and e.src_uid == uid + ] + + def _all_gather_preds(uid: str) -> list[str]: + return sorted( + e.src_uid + for e in dag.edges + if e.dep_kind == "data" + and e.dst_uid == uid + and e.src_uid in dag.nodes + and dag.nodes[e.src_uid].node_kind == "ALL_GATHER_COMM" + ) + + def _reduce_scatter_succs(uid: str) -> list[str]: + return sorted( + e.dst_uid + for e in dag.edges + if e.dep_kind == "data" + and e.src_uid == uid + and e.dst_uid in dag.nodes + and dag.nodes[e.dst_uid].node_kind == "REDUCE_SCATTER_COMM" + ) + + def _remove_all_gather(ag_uid: str, compute_uid: str, bucket: Any) -> None: + del compute_uid, bucket + if ag_uid not in dag.nodes: + return + in_ag = _data_preds(ag_uid) + out_ag = _data_succs(ag_uid) + for ie in in_ag: + for oe in out_ag: + dag.add_edge( + TrainingDAGEdge( + src_uid=ie.src_uid, + dst_uid=oe.dst_uid, + dep_kind="data", + tensor_name=(oe.tensor_name if oe.tensor_name is not None else ie.tensor_name), + ) + ) + _remove_node_and_incident_edges(dag, ag_uid) + + def _remove_reduce_scatter(rs_uid: str, compute_uid: str, bucket: Any) -> None: + del compute_uid, bucket + if rs_uid not in dag.nodes: + return + _remove_node_and_incident_edges(dag, rs_uid) + + has_grad_sync = any( + n.node_kind in ("REDUCE_COMM", "REDUCE_SCATTER_COMM") + for n in dag.nodes.values() + ) + has_all_gathers = any(n.node_kind == "ALL_GATHER_COMM" for n in dag.nodes.values()) + + grad_chains: list[list[str]] = [] + param_chains: list[list[str]] = [] + + # Gradient lifetimes: keep one full-grad allocation and one reduce-scatter + # per directly-connected same-bucket gradient chain. + if has_grad_sync: + grad_candidates = _compute_uids_with( + topo, + lambda uid: bool(dag.nodes[uid].node_meta.get("zero_alloc_full_grads_before")) + ) + if grad_candidates: + grad_chains = _build_chains(grad_candidates, topo) + covered = {uid for chain in grad_chains for uid in chain} + assert covered == set(grad_candidates), ( + "zero_lifetime_prune: gradient mode did not assign every grad-allocation " + f"compute node to a chain; missing={sorted(set(grad_candidates) - covered)}" + ) + + for chain in grad_chains: + bucket = _chain_bucket(chain[0]) + assert all(_chain_bucket(uid) == bucket for uid in chain), ( + f"zero_lifetime_prune: gradient chain has mixed buckets: " + f"{[(uid, _chain_bucket(uid)) for uid in chain]}" + ) + logger.info( + "zero_lifetime_prune_grad_chain bucket=%s nodes=%s keep_alloc_for=%s keep_reduce_scatter_for=%s", + bucket, + chain, + chain[0], + chain[-1], + ) + + for uid in chain: + rs_uids = _reduce_scatter_succs(uid) + assert rs_uids, ( + f"zero_lifetime_prune: gradient chain node {uid} has " + "zero_alloc_full_grads_before but no outgoing reduce-scatter" + ) + + for uid in chain[1:]: + dag.nodes[uid].node_meta.pop("zero_alloc_full_grads_before", None) + for uid in chain[:-1]: + for rs_uid in _reduce_scatter_succs(uid): + _remove_reduce_scatter(rs_uid, uid, bucket) + + assert dag.nodes[chain[0]].node_meta.get("zero_alloc_full_grads_before"), ( + f"zero_lifetime_prune: gradient chain root {chain[0]} lost grad allocation" + ) + for prev_uid, uid in zip(chain, chain[1:]): + assert _has_direct_compute_dependency(prev_uid, uid), ( + f"zero_lifetime_prune: gradient chain is not directly linked: " + f"{prev_uid} does not have a direct compute dependency to {uid}; chain={chain}" + ) + assert not dag.nodes[uid].node_meta.get("zero_alloc_full_grads_before"), ( + f"zero_lifetime_prune: non-root gradient chain node {uid} still allocates grads" + ) + live_rs = [ + (uid, rs_uid) + for uid in chain + for rs_uid in _reduce_scatter_succs(uid) + ] + assert live_rs, ( + f"zero_lifetime_prune: gradient chain has no live reduce-scatter; chain={chain}" + ) + assert all(uid == chain[-1] for uid, _rs_uid in live_rs), ( + f"zero_lifetime_prune: only chain tail {chain[-1]} should keep reduce-scatter, " + f"but live reduce-scatters are {live_rs}" + ) + + # Parameter lifetimes: keep one all-gather and one full-param free per + # directly-connected same-bucket parameter chain. This pass intentionally + # does not change grad allocation metadata or remove reduce-scatters. + if has_all_gathers: + topo = _serial_topological_order(dag, _topological_levels(dag)) + ag_by_compute = {uid: _all_gather_preds(uid) for uid in topo if uid in dag.nodes and _is_compute(uid)} + ag_by_compute = {uid: ags for uid, ags in ag_by_compute.items() if ags} + + candidates = _compute_uids_with(topo, lambda uid: uid in ag_by_compute) + param_chains = _build_chains(candidates, topo) + + covered = {uid for chain in param_chains for uid in chain} + assert covered == set(candidates), ( + "zero_lifetime_prune: parameter mode did not assign every all-gathered " + f"compute node to a chain; missing={sorted(set(candidates) - covered)}" + ) + + for chain in param_chains: + bucket = _chain_bucket(chain[0]) + assert all(_chain_bucket(uid) == bucket for uid in chain), ( + f"zero_lifetime_prune: parameter chain has mixed buckets: " + f"{[(uid, _chain_bucket(uid)) for uid in chain]}" + ) + + for uid in chain[1:]: + for ag_uid in ag_by_compute[uid]: + _remove_all_gather(ag_uid, uid, bucket) + + tail = chain[-1] + for uid in chain: + dag.nodes[uid].node_meta.pop("zero_free_full_params_after", None) + for rs_uid in _reduce_scatter_succs(uid): + dag.nodes[rs_uid].node_meta.pop("zero_free_full_params_after", None) + dag.nodes[tail].node_meta["zero_free_full_params_after"] = True + + remaining_ag = {uid: _all_gather_preds(uid) for uid in chain} + assert remaining_ag[chain[0]], ( + f"zero_lifetime_prune: parameter chain root {chain[0]} lost all-gather" + ) + for prev_uid, uid in zip(chain, chain[1:]): + assert _has_direct_compute_dependency(prev_uid, uid), ( + f"zero_lifetime_prune: parameter chain is not directly linked: " + f"{prev_uid} does not have a direct compute dependency to {uid}; chain={chain}" + ) + assert not remaining_ag[uid], ( + f"zero_lifetime_prune: non-root parameter chain node {uid} still " + f"has all-gather predecessors {remaining_ag[uid]}; chain={chain}" + ) + assert dag.nodes[tail].node_meta.get("zero_free_full_params_after"), ( + f"zero_lifetime_prune: parameter chain tail {tail} should free params" + ) + for uid in chain[:-1]: + assert not dag.nodes[uid].node_meta.get("zero_free_full_params_after"), ( + f"zero_lifetime_prune: non-tail parameter chain node {uid} still frees params" + ) + stale_rs_param_frees = [ + rs_uid + for uid in chain + for rs_uid in _reduce_scatter_succs(uid) + if dag.nodes[rs_uid].node_meta.get("zero_free_full_params_after") + ] + assert not stale_rs_param_frees, ( + f"zero_lifetime_prune: reduce-scatters still own param free metadata: " + f"{stale_rs_param_frees}; chain={chain}" + ) + + return grad_chains + param_chains + + +def _add_inter_chain_temporal_edges( + dag: TrainingDAG, + chains: list[list[str]], +) -> None: + """Force later same-bucket ZeRO chains to wait for earlier bucket frees. + + Parameter chains: prev_tail compute frees full params, next root's + all-gather predecessor materializes them. Edge: prev_tail -> ag_uid. + + Gradient chains: prev_tail's reduce-scatter successor frees full grads, + next root compute allocates them. Edge: rs_uid -> next_root. + """ + if not chains: + return + + topo_levels = _topological_levels(dag) + + def _chain_bucket(uid: str) -> Any: + return dag.nodes[uid].node_meta.get("bucket_key") + + def _all_gather_preds(uid: str) -> list[str]: + return sorted( + e.src_uid + for e in dag.edges + if e.dep_kind == "data" + and e.dst_uid == uid + and e.src_uid in dag.nodes + and dag.nodes[e.src_uid].node_kind == "ALL_GATHER_COMM" + ) + + def _reduce_scatter_succs(uid: str) -> list[str]: + return sorted( + e.dst_uid + for e in dag.edges + if e.dep_kind == "data" + and e.src_uid == uid + and e.dst_uid in dag.nodes + and dag.nodes[e.dst_uid].node_kind == "REDUCE_SCATTER_COMM" + ) + + def _maybe_add_edge(src_uid: str, dst_uid: str) -> None: + if topo_levels.get(src_uid, 0) < topo_levels.get(dst_uid, 0): + return + if _has_path(dag, src_uid, dst_uid): + return + if _has_path(dag, dst_uid, src_uid): + return + dag.add_edge(TrainingDAGEdge( + src_uid=src_uid, + dst_uid=dst_uid, + dep_kind="temporal", + tensor_name=None, + )) + + # Group chains by type and bucket, then add type-specific temporal edges + # between consecutive same-bucket chains. + param_by_bucket: dict[Any, list[list[str]]] = {} + grad_by_bucket: dict[Any, list[list[str]]] = {} + for chain in chains: + bucket = _chain_bucket(chain[0]) + if _all_gather_preds(chain[0]): + param_by_bucket.setdefault(bucket, []).append(chain) + if dag.nodes[chain[0]].node_meta.get("zero_alloc_full_grads_before"): + grad_by_bucket.setdefault(bucket, []).append(chain) + + for bucket_chains in param_by_bucket.values(): + bucket_chains.sort(key=lambda c: topo_levels.get(c[0], 0)) + for prev_chain, next_chain in zip(bucket_chains, bucket_chains[1:]): + for ag_uid in _all_gather_preds(next_chain[0]): + _maybe_add_edge(prev_chain[-1], ag_uid) + + for bucket_chains in grad_by_bucket.values(): + bucket_chains.sort(key=lambda c: topo_levels.get(c[0], 0)) + for prev_chain, next_chain in zip(bucket_chains, bucket_chains[1:]): + for rs_uid in _reduce_scatter_succs(prev_chain[-1]): + _maybe_add_edge(rs_uid, next_chain[0]) diff --git a/test/models/moe.py b/test/models/moe.py deleted file mode 100644 index 9339635..0000000 --- a/test/models/moe.py +++ /dev/null @@ -1,163 +0,0 @@ -""" -MoE implementation: https://github.com/peytontolbert/simple-moe -""" - -import torch -import torch.nn as nn -from torch.nn import functional as F -from typing import Optional, Dict, Any, Tuple - -class ExpertBase(nn.Module): - """Base class for expert networks in the Mixture of Experts model.""" - - def __init__(self, input_dim: int, output_dim: int, hidden_dim: Optional[int] = None): - """ - Initialize the expert network. - - Args: - input_dim: Dimension of input features - output_dim: Dimension of output features - hidden_dim: Dimension of hidden layer (if None, uses 4x input_dim) - """ - super().__init__() - self.input_dim = input_dim - self.output_dim = output_dim - self.hidden_dim = hidden_dim or 4 * input_dim - - self.network = nn.Sequential( - nn.Linear(input_dim, self.hidden_dim), - nn.ReLU(), - nn.Linear(self.hidden_dim, output_dim) - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass through the expert.""" - return self.network(x) - - def get_config(self) -> Dict[str, Any]: - """Get expert configuration for serialization.""" - return { - 'input_dim': self.input_dim, - 'output_dim': self.output_dim, - 'hidden_dim': self.hidden_dim - } - - -class FFNExpert(ExpertBase): - """Feed-forward neural network expert implementation.""" - - def __init__(self, input_dim: int, output_dim: int, hidden_dim: Optional[int] = None, - num_layers: int = 2, dropout: float = 0.1): - """ - Initialize the FFN expert. - - Args: - input_dim: Dimension of input features - output_dim: Dimension of output features - hidden_dim: Dimension of hidden layers - num_layers: Number of hidden layers - dropout: Dropout probability - """ - super().__init__(input_dim, output_dim, hidden_dim) - - layers = [] - current_dim = input_dim - - for _ in range(num_layers - 1): - layers.extend([ - nn.Linear(current_dim, self.hidden_dim), - nn.ReLU(), - nn.Dropout(dropout) - ]) - current_dim = self.hidden_dim - - layers.append(nn.Linear(current_dim, output_dim)) - self.network = nn.Sequential(*layers) - - def get_config(self) -> Dict[str, Any]: - config = super().get_config() - config.update({ - 'num_layers': len(self.network) // 3 + 1, - 'dropout': self.network[2].p if len(self.network) > 2 else 0.0 - }) - return config - - -class MixtureOfExperts(nn.Module): - """Mixture of Experts model implementation.""" - - def __init__(self, - input_dim: int, - output_dim: int, - num_experts: int, - expert_class: type = FFNExpert, - expert_kwargs: Optional[Dict[str, Any]] = None, - k: int = 1, - capacity_factor: float = 1.0, - router_noise_epsilon: float = 1e-2): - """ - Initialize the MoE model. - - Args: - input_dim: Input dimension - output_dim: Output dimension - num_experts: Number of experts - expert_class: Expert class to use - expert_kwargs: Additional arguments for expert initialization - k: Number of experts to route to - capacity_factor: Expert capacity multiplier - router_noise_epsilon: Noise factor for router load balancing - """ - super().__init__() - self.input_dim = input_dim - self.output_dim = output_dim - self.num_experts = num_experts - self.k = k - - assert k == 1 - - # Initialize experts - expert_kwargs = expert_kwargs or {} - self.experts = nn.ModuleList([ - expert_class(input_dim=input_dim, output_dim=output_dim, **expert_kwargs) - for _ in range(num_experts) - ]) - - def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """ - Forward pass through the MoE model. - - Args: - x: Input tensor of shape [batch_size, input_dim] - - Returns: - Tuple of: - - output: Model output of shape [batch_size, output_dim] - - aux_loss: Auxiliary load balancing loss (None if not training) - """ - - batch_size = x.shape[0] - - outputs = [] - for batch_idx in range(batch_size): - expert_idx = batch_idx % self.num_experts - expert_id = id(self.experts[expert_idx]) - batch = x[batch_idx:batch_idx+1] - with torch.fx.traceback.annotate({"global_expert_id": expert_id, "local_expert_id": expert_idx, "batch_idx": batch_idx}): - expert_output = self.experts[expert_idx](batch) - outputs.append(expert_output) - output = torch.stack(outputs, dim=-1).squeeze() - - return output - - - def get_config(self) -> Dict[str, Any]: - """Get model configuration for serialization.""" - return { - 'input_dim': self.input_dim, - 'output_dim': self.output_dim, - 'num_experts': self.num_experts, - 'k': self.k, - 'expert_class': self.experts[0].__class__.__name__, - 'expert_config': self.experts[0].get_config() - } \ No newline at end of file diff --git a/test/schedule_helpers.py b/test/schedule_helpers.py deleted file mode 100644 index 09699bb..0000000 --- a/test/schedule_helpers.py +++ /dev/null @@ -1,218 +0,0 @@ -from src.piper_exec import Task - - -def print_schedule(schedule): - for stage in schedule: - for step in stage: - if step: - string = f"{step.stage_id}:{step.mb_idx}:{'u' if step.upd else 'f' if step.is_fwd else 'b'}" - else: - string = " -- " - print(string, end="\t") - print() - -def build_gpipe_schedule(n_mbs: int, n_stages: int): - steps = n_mbs + n_stages - 1 - schedule = [[None] * (steps * 2 + 1) for _ in range(n_stages)] - for step in range(steps): - for stage_id in range(n_stages): - mb_idx = step - stage_id - if mb_idx >= 0 and mb_idx < n_mbs: - schedule[stage_id][step] = Task(stage_id, stage_id, mb_idx, True, False) - - for step in range(steps, steps * 2): - for stage_id in reversed(range(n_stages)): - mb_idx = (step - steps) - (n_stages - stage_id - 1) - if mb_idx >= 0 and mb_idx < n_mbs: - schedule[stage_id][step] = Task(stage_id, stage_id, mb_idx, False, False) - for i, stage in enumerate(range(n_stages)): - schedule[stage][-i-1] = Task(stage_id=stage, device_id=stage, mb_idx=0, is_fwd=False, upd=True) - return schedule - -def build_1f1b_schedule(n_mbs: int, n_stages: int): - steps = n_mbs + n_stages - 1 - schedule = [[None] * (steps * 2 + 1) for _ in range(n_stages)] - stage_mb = [[0, 0] for _ in range(n_stages)] - for step in range(n_stages): - for stage_id in range(n_stages): - if step >= stage_id: - mb_idx = stage_mb[stage_id][0] - if mb_idx >= 0 and mb_idx < n_mbs: - schedule[stage_id][step] = Task( - device_id=stage_id, stage_id=stage_id, - mb_idx=mb_idx, is_fwd=True, upd=False - ) - stage_mb[stage_id][0] += 1 - for step in range(n_stages, 2 * steps): - relative_step = step - n_stages - for stage_id in range(n_stages): - inv_stage = n_stages - stage_id - 1 - if relative_step >= inv_stage: - fwd_or_bwd = 1 - (relative_step + inv_stage) % 2 - task_type = True if fwd_or_bwd == 0 else False - mb_idx = stage_mb[stage_id][fwd_or_bwd] - if mb_idx >= 0 and mb_idx < n_mbs: - schedule[stage_id][step] = Task( - device_id=stage_id, stage_id=stage_id, - mb_idx=mb_idx, is_fwd=task_type, upd=False - ) - stage_mb[stage_id][fwd_or_bwd] += 1 - for i, stage in enumerate(range(n_stages)): - schedule[stage][-i-1] = Task(stage_id=stage, device_id=stage, mb_idx=n_mbs-1, is_fwd=False, upd=True) - return schedule - -no_pp_schedule = [ - [ - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=False, upd=True), - ] -] - -pp2_interleaved_1f1b_grid_schedule = [ - [ - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=0, stage_id=2, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=0, stage_id=2, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=0, stage_id=2, mb_idx=0, is_fwd=False, upd=False), - None, - Task(device_id=0, stage_id=2, mb_idx=1, is_fwd=False, upd=False), - Task(device_id=0, stage_id=2, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=0, stage_id=2, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=1, is_fwd=False, upd=False), - None, - Task(device_id=0, stage_id=2, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=0, stage_id=2, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=False, upd=True), - ], - [ - None, - Task(device_id=1, stage_id=1, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=1, stage_id=1, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=1, stage_id=3, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=1, stage_id=3, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=1, stage_id=3, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=1, stage_id=3, mb_idx=1, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=1, stage_id=1, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=1, stage_id=1, mb_idx=1, is_fwd=False, upd=False), - Task(device_id=1, stage_id=3, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=1, stage_id=3, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=1, stage_id=3, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=1, stage_id=3, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=0, is_fwd=False, upd=True), - None, - ], -] - -pp4_interleaved_1f1b_grid_schedule = [ - [ - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=0, stage_id=0, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=0, stage_id=4, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=0, stage_id=4, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=0, stage_id=4, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=0, stage_id=4, mb_idx=3, is_fwd=True, upd=False), - None, - None, - None, - Task(device_id=0, stage_id=4, mb_idx=0, is_fwd=False, upd=False), - None, - Task(device_id=0, stage_id=4, mb_idx=1, is_fwd=False, upd=False), - None, - Task(device_id=0, stage_id=4, mb_idx=2, is_fwd=False, upd=False), - None, - Task(device_id=0, stage_id=4, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=1, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=0, stage_id=0, mb_idx=0, is_fwd=False, upd=True), - ], - [ - None, - Task(device_id=1, stage_id=1, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=1, stage_id=1, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=1, stage_id=1, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=1, stage_id=1, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=1, stage_id=5, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=1, stage_id=5, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=1, stage_id=5, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=1, stage_id=5, mb_idx=3, is_fwd=True, upd=False), - None, - Task(device_id=1, stage_id=5, mb_idx=0, is_fwd=False, upd=False), - None, - Task(device_id=1, stage_id=5, mb_idx=1, is_fwd=False, upd=False), - None, - Task(device_id=1, stage_id=5, mb_idx=2, is_fwd=False, upd=False), - None, - Task(device_id=1, stage_id=5, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=1, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=1, stage_id=1, mb_idx=0, is_fwd=False, upd=True), - None, - ], - [ - None, - None, - Task(device_id=2, stage_id=2, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=2, stage_id=2, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=2, stage_id=2, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=2, stage_id=2, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=2, stage_id=6, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=2, stage_id=6, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=2, stage_id=6, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=2, stage_id=6, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=2, stage_id=6, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=2, stage_id=6, mb_idx=1, is_fwd=False, upd=False), - None, - Task(device_id=2, stage_id=6, mb_idx=2, is_fwd=False, upd=False), - None, - Task(device_id=2, stage_id=6, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=2, stage_id=2, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=2, stage_id=2, mb_idx=1, is_fwd=False, upd=False), - Task(device_id=2, stage_id=2, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=2, stage_id=2, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=2, stage_id=2, mb_idx=0, is_fwd=False, upd=True), - None, - None, - ], - [ - None, - None, - None, - Task(device_id=3, stage_id=3, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=3, stage_id=3, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=3, stage_id=3, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=3, stage_id=3, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=3, stage_id=7, mb_idx=0, is_fwd=True, upd=False), - Task(device_id=3, stage_id=7, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=3, stage_id=7, mb_idx=1, is_fwd=True, upd=False), - Task(device_id=3, stage_id=7, mb_idx=1, is_fwd=False, upd=False), - Task(device_id=3, stage_id=7, mb_idx=2, is_fwd=True, upd=False), - Task(device_id=3, stage_id=7, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=3, stage_id=7, mb_idx=3, is_fwd=True, upd=False), - Task(device_id=3, stage_id=7, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=3, stage_id=3, mb_idx=0, is_fwd=False, upd=False), - Task(device_id=3, stage_id=3, mb_idx=1, is_fwd=False, upd=False), - Task(device_id=3, stage_id=3, mb_idx=2, is_fwd=False, upd=False), - Task(device_id=3, stage_id=3, mb_idx=3, is_fwd=False, upd=False), - Task(device_id=3, stage_id=3, mb_idx=0, is_fwd=False, upd=True), - None, - None, - None, - ], -] \ No newline at end of file diff --git a/test/test_graph_break.py b/test/test_graph_break.py deleted file mode 100644 index 2807769..0000000 --- a/test/test_graph_break.py +++ /dev/null @@ -1,62 +0,0 @@ -import torch -import torch.nn as nn -import ray -import time - -from src.piper_compile import piper_setup -from src.piper_exec import piper_exec -from src.piper import distributed_stage -from src.piper_coordinator import PiperProgramCoordinator - -from .schedule_helpers import build_1f1b_schedule, print_schedule - -def main(): - class TwoLayerNN(nn.Module): - def __init__(self, input_dim, hidden_dim, output_dim): - super().__init__() - self.layer1 = nn.Linear(input_dim, hidden_dim) - self.layer2 = nn.Linear(hidden_dim, output_dim) - self.layer3 = nn.Linear(output_dim, output_dim) - - def forward(self, x): - distributed_stage(0) - x = self.layer1(x) - print(f"graph break: {x.data} {x + 1} {x[0]}") - x = self.layer2(x) - - distributed_stage(1) - x = self.layer3(x) - return x - - x = torch.randn(2) - y = torch.randn(2) - - model = piper_setup(TwoLayerNN, (2, 2, 2), torch.optim.Adam, [x], num_stages=2, num_devices=2) - - return - - num_mbs = 1 - num_stages = 2 - schedule = build_1f1b_schedule(num_mbs, num_stages) - print_schedule(schedule) - loss_fn = torch.nn.CrossEntropyLoss() - - from src.piper_utils import piper_metadata - stage1_weights = ray.get(piper_metadata.actors[0].get_weights.remote(0)) - stage2_weights = ray.get(piper_metadata.actors[1].get_weights.remote(1)) - print("stage 1 weights:", stage1_weights) - print("stage 2 weights:", stage2_weights) - - piper_exec(model, schedule, [x], y, loss_fn, num_mbs, num_stages) - - stage1_weights = ray.get(piper_metadata.actors[0].get_weights.remote(0)) - stage2_weights = ray.get(piper_metadata.actors[1].get_weights.remote(1)) - print("stage 1 weights:", stage1_weights) - print("stage 2 weights:", stage2_weights) - -if __name__ == "__main__": - ray.init(include_dashboard=False, log_to_driver=True, namespace="llama") - piper_coordinator = PiperProgramCoordinator.remote(pp_degree=2, dp_degree=1, world_size=2) - ray.get(piper_coordinator.run_program.remote(main)) - time.sleep(3) - ray.shutdown() \ No newline at end of file diff --git a/test/test_graph_transform.py b/test/test_graph_transform.py new file mode 100644 index 0000000..496f3b1 --- /dev/null +++ b/test/test_graph_transform.py @@ -0,0 +1,125 @@ +import pytest +import torch +import torch.fx as fx +import torch.nn as nn + +from src.piper import _reset_annotation_state, annotate +from src.fx import split_gm_by_annotations + + +class _AnnotatedNet(nn.Module): + def __init__(self, dim: int = 8): + super().__init__() + self.fc1 = nn.Linear(dim, dim, bias=False) + self.fc2 = nn.Linear(dim, dim, bias=False) + self.fc3 = nn.Linear(dim, dim, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + with annotate("PP"): + x = torch.relu(self.fc1(x)) + with annotate("EP"): + x = torch.relu(self.fc2(x)) + with annotate("PP"): + return self.fc3(x) + + +def _capture_gm(model: nn.Module, x: torch.Tensor): + captured: dict = {} + + def _backend(gm, example_inputs): + captured["gm"] = gm + captured["inputs"] = list(example_inputs) + return gm.forward + + _reset_annotation_state() + compiled = torch.compile(model, backend=_backend, fullgraph=True) + with torch.no_grad(): + ref = compiled(x).detach().clone() + + return captured["gm"], captured["inputs"], ref + + +def _placeholder_values(gm: fx.GraphModule, example_inputs: list) -> dict[str, object]: + placeholders = [node for node in gm.graph.nodes if node.op == "placeholder"] + return {node.name: value for node, value in zip(placeholders, example_inputs)} + + +def _run_segments(segments, placeholder_values: dict[str, object]) -> torch.Tensor: + prev_output = None + + for segment in segments: + ph_nodes = [node for node in segment.gm.graph.nodes if node.op == "placeholder"] + args = [] + activation_inputs = ( + () + if prev_output is None + else prev_output + if isinstance(prev_output, tuple) + else (prev_output,) + ) + activation_iter = iter(activation_inputs) + + for node in ph_nodes: + if node.name in placeholder_values: + args.append(placeholder_values[node.name]) + else: + args.append(next(activation_iter)) + + with torch.no_grad(): + prev_output = segment.gm(*args) + + if isinstance(prev_output, (tuple, list)) and len(prev_output) == 1: + prev_output = prev_output[0] + return prev_output + + +def test_split_gm_by_annotations_handles_nested_annotation_segments() -> None: + torch.manual_seed(0) + model = _AnnotatedNet() + x = torch.randn(2, 8) + + gm, example_inputs, ref = _capture_gm(model, x) + _, segments = split_gm_by_annotations(gm) + + assert [segment.tag for segment in segments] == [ + {"PP": 0}, + {"PP": 0, "EP": 0}, + {"PP": 1}, + ] + assert [segment.stage_id for segment in segments] == [0, 0, 1] + assert [segment.a2a_boundary_after["from_tag"] for segment in segments[:-1]] == [ + {"PP": 0}, + {"PP": 0, "EP": 0}, + ] + assert [segment.a2a_boundary_after["to_tag"] for segment in segments[:-1]] == [ + {"PP": 0, "EP": 0}, + {"PP": 1}, + ] + + result = _run_segments(segments, _placeholder_values(gm, example_inputs)) + + assert torch.allclose(ref, result, atol=1e-6) + + +def test_split_gm_by_annotations_ignores_non_piper_custom_metadata() -> None: + gm = fx.symbolic_trace(nn.Sequential(nn.ReLU())) + relu = next(node for node in gm.graph.nodes if node.op == "call_module") + relu.meta["custom"] = {"name": "PP", "index": 0} + + _, segments = split_gm_by_annotations(gm) + + assert segments == [] + + +def test_split_gm_by_annotations_rejects_mixed_annotated_and_unannotated_compute() -> None: + torch.manual_seed(0) + gm, _, _ = _capture_gm(_AnnotatedNet(), torch.randn(2, 8)) + first_compute = next( + node + for node in gm.graph.nodes + if node.op not in ("placeholder", "get_attr", "output") + ) + first_compute.meta.pop("custom", None) + + with pytest.raises(ValueError, match="without Piper annotations"): + split_gm_by_annotations(gm) diff --git a/test/test_llama.py b/test/test_llama.py deleted file mode 100644 index ee5e76a..0000000 --- a/test/test_llama.py +++ /dev/null @@ -1,209 +0,0 @@ -import ray -import torch -import time -import argparse -import os -from torch import nn, optim -from torch.profiler import profile, record_function, ProfilerActivity - -from src.piper_exec import Task, piper_exec -from src.piper_compile import piper_setup -from src.piper import piper -from src.piper_utils import piper_metadata -from src.piper_coordinator import PiperProgramCoordinator - -from .models.llama import Transformer, LLAMA_DEBUG, LLAMA_1B, LLAMA_3B, LLAMA_8B -from .schedule_helpers import ( - build_1f1b_schedule, - build_gpipe_schedule, - print_schedule, - pp2_interleaved_1f1b_grid_schedule, - pp4_interleaved_1f1b_grid_schedule, - no_pp_schedule -) - -def parse_args(): - parser = argparse.ArgumentParser(description='Run LLaMA model with pipeline parallelism') - parser.add_argument('--model', choices=['LLAMA_DEBUG', 'LLAMA_1B', 'LLAMA_3B', 'LLAMA_8B'], default='LLAMA_DEBUG', - help='Model configuration: LLAMA_DEBUG, LLAMA_1B, LLAMA_3B, or LLAMA_8B (default: LLAMA_DEBUG)') - parser.add_argument('--schedule', choices=['gpipe', '1f1b', 'interleaved-1f1b', 'no-pp'], default='1f1b', - help='Schedule type: gpipe, 1f1b, or interleaved-1f1b (default: 1f1b)') - # num_stages should be able to be inferred from the model code - parser.add_argument('--num_stages', type=int, default=2, - help='Number of stages (default: 2)') - parser.add_argument('--dp_degree', type=int, default=1, - help='Number of data parallel degrees (default: 1)') - parser.add_argument('--pp_degree', type=int, default=2, - help='Number of pipeline parallel degrees (default: 2)') - parser.add_argument('--batch_size', type=int, default=16, - help='Batch size (default: 16)') - parser.add_argument('--num_mbs', type=int, default=4, - help='Number of microbatches (default: 4)') - parser.add_argument('--seq_len', type=int, default=256, - help='Sequence length (default: 256)') - parser.add_argument('--warmup', type=int, default=5, - help='Number of warmup iterations (default: 5)') - parser.add_argument('--iters', type=int, default=20, - help='Number of timing iterations (default: 20)') - parser.add_argument('--tracing', action='store_true', default=False, - help='Enable tracing') - return parser.parse_args() - - -def print_mean_timing_data(trace_data: dict, actor_id: int) -> None: - """Prints mean timing and memory statistics from trace_data for a given actor. - - Args: - trace_data (dict): Trace data dictionary from the actor containing timing and memory metrics. - actor_id (int): The ID of the actor. - """ - import numpy as np - - print(f"\nTiming and Memory statistics for Actor {actor_id}:") - - for stage_id, stage_data in trace_data.items(): - if stage_id == 'update': - # Handle update data (global optimizer step) - print(f" Update:") - for metric, values in stage_data.items(): - if not values: - mean_val = float('nan') - else: - mean_val = float(np.mean(values)) - - if 'memory' in metric: - if 'delta' in metric: - print(f" {metric}: {mean_val:.3f} GB") - else: - print(f" {metric}: {mean_val:.3f} GB") - else: - print(f" {metric}: {mean_val:.3f} ms") - else: - # Handle stage data (forward/backward passes) - print(f" Stage {stage_id}:") - for phase, phase_data in stage_data.items(): - print(f" {phase.capitalize()}:") - for metric, values in phase_data.items(): - if not values: - mean_val = float('nan') - else: - mean_val = float(np.mean(values)) - - if 'memory' in metric: - if 'delta' in metric: - print(f" {metric}: {mean_val:.3f} GB") - else: - print(f" {metric}: {mean_val:.3f} GB") - else: - print(f" {metric}: {mean_val:.3f} ms") - - -def main(args): - - # Set model configuration based on argument - match args.model: - case 'LLAMA_DEBUG': - llama_config = LLAMA_DEBUG - case 'LLAMA_1B': - llama_config = LLAMA_1B - case 'LLAMA_3B': - llama_config = LLAMA_3B - case 'LLAMA_8B': - llama_config = LLAMA_8B - print(args) - - loss_fn = torch.nn.CrossEntropyLoss() - - batch_size = args.batch_size - num_mbs = args.num_mbs - seq_len = args.seq_len - warmup = args.warmup - iters = args.iters - - x = torch.randint(0, llama_config.vocab_size, (batch_size, seq_len)) - y = torch.randn((batch_size, seq_len, llama_config.vocab_size)) - - # Generate different input data for each data parallel rank so that model weights get updated differently - if args.dp_degree > 1: - dp_rank = int(os.environ['PIPER_DP_RANK']) - torch.manual_seed(dp_rank) - x = torch.randint(0, llama_config.vocab_size, (batch_size, seq_len)) - torch.manual_seed(0) - - num_stages = args.num_stages - compiled = piper_setup(Transformer, (llama_config, seq_len), torch.optim.Adam, [x], num_stages, args.pp_degree) - - assert num_stages == len(piper_metadata.dag) + 1 - - schedule = None - match args.schedule: - case "no-pp": - schedule = no_pp_schedule - case "interleaved-1f1b": - schedule = pp2_interleaved_1f1b_grid_schedule if args.pp_degree == 2 else pp4_interleaved_1f1b_grid_schedule - case "1f1b": - schedule = build_1f1b_schedule(num_mbs, num_stages) - schedule[0][2] = schedule[0][4] - schedule[0][4] = schedule[0][6] - schedule[0][6] = None - case "gpipe": - schedule = build_gpipe_schedule(num_mbs, num_stages) - - print("SCHEDULE:") - print_schedule(schedule) - - actors = piper_metadata.actors - - # Send data to actors ahead of time - # num_actors = len(actors) - # ray.get(actors[0].send_input.remote(x)) - # ray.get(actors[num_actors-1].send_truth.remote(y)) - - ray.get([actor.set_tracing.remote(args.tracing) for actor in actors.values()]) - - # Definte one iteration of the schedule - def iter_schedule(): - losses = piper_exec(compiled, schedule, [x], y, loss_fn, num_mbs, num_stages) - - # Warmup - print(f"Running {warmup} warmup iterations...") - for _ in range(warmup): - iter_schedule() - - # Clear tracing data - ray.get([actor.clear_trace_data.remote() for actor in actors.values()]) - ray.get([actor.reset_peak_memory.remote() for actor in actors.values()]) - - # Time training steps - start = time.perf_counter() - print(f"Running {iters} timed iterations...") - for _ in range(iters): - iter_schedule() - end = time.perf_counter() - - print(f"Iteration time: {(end - start)*1e3/iters:.0f} ms") - print( - f"{args.schedule} throughput: {(iters * batch_size * num_mbs * seq_len)/(end - start):.0f} tokens/sec" - ) - - if args.tracing: - # Get overall peak memory - peak_memory = ray.get([actor.get_peak_memory.remote() for actor in actors.values()]) - print("Peak memory:") - for actor_id, peak_memory in enumerate(peak_memory): - print(f"\tActor {actor_id}: {peak_memory:.1f} GB") - - # Get tracing data from actors - for actor in actors.values(): - trace_data = ray.get(actor.get_trace_data.remote()) - actor_id = ray.get(actor.id.remote()) - print_mean_timing_data(trace_data, actor_id) - - ray.timeline(f"out/{args.model}-pp{args.pp_degree}-dp{args.dp_degree}-{args.schedule}.json") - -if __name__ == "__main__": - ray.init(include_dashboard=False, log_to_driver=True, namespace="llama") - args = parse_args() - piper_coordinator = PiperProgramCoordinator.remote(dp_degree=args.dp_degree, pp_degree=args.pp_degree) - ray.get(piper_coordinator.run_program.remote(main, args)) - ray.shutdown() diff --git a/test/test_llama_single.py b/test/test_llama_single.py deleted file mode 100644 index 6f36db6..0000000 --- a/test/test_llama_single.py +++ /dev/null @@ -1,274 +0,0 @@ -import ray -import torch -import time -import argparse -from torch import nn, optim -from torch.profiler import profile, record_function, ProfilerActivity - -from src.piper_exec import Task, piper_exec -from src.piper_compile import piper_setup -from src.piper import piper -from src.piper_utils import piper_metadata - -from .models.llama import Transformer, LLAMA_DEBUG, LLAMA_1B, LLAMA_3B, LLAMA_8B -from .schedule_helpers import print_schedule - -def parse_args(): - parser = argparse.ArgumentParser(description='Run LLaMA model on single device (no pipeline parallelism)') - parser.add_argument('--model', choices=['LLAMA_DEBUG', 'LLAMA_1B', 'LLAMA_3B', 'LLAMA_8B'], default='LLAMA_DEBUG', - help='Model configuration: LLAMA_DEBUG, LLAMA_1B, LLAMA_3B, or LLAMA_8B (default: LLAMA_DEBUG)') - parser.add_argument('--batch_size', type=int, default=8, - help='Batch size (default: 8)') - parser.add_argument('--num_mbs', type=int, default=4, - help='Number of microbatches (default: 4)') - parser.add_argument('--seq_len', type=int, default=512, - help='Sequence length (default: 512)') - parser.add_argument('--warmup', type=int, default=5, - help='Number of warmup iterations (default: 5)') - parser.add_argument('--iters', type=int, default=20, - help='Number of timing iterations (default: 20)') - parser.add_argument('--tracing', action='store_true', default=False, - help='Enable tracing') - return parser.parse_args() - - -def build_single_device_schedule(n_mbs: int): - """ - Build a simple schedule for single device (no pipeline parallelism). - - For single device, we just do: - - All forward passes for all microbatches - - All backward passes for all microbatches - - Single update step - - Args: - n_mbs: Number of microbatches - - Returns: - A 1xN schedule (1 device, N time steps) - """ - # Calculate number of time steps: n_mbs forwards + n_mbs backwards + 1 update - num_steps = n_mbs * 2 + 1 - - # Single device (stage 0), multiple time steps - schedule = [[None] * num_steps] - - # Forward passes for all microbatches - for mb_idx in range(n_mbs): - schedule[0][mb_idx] = Task( - device_id=0, - stage_id=0, - mb_idx=mb_idx, - is_fwd=True, - upd=False - ) - - # Backward passes for all microbatches - for mb_idx in range(n_mbs): - schedule[0][n_mbs + mb_idx] = Task( - device_id=0, - stage_id=0, - mb_idx=mb_idx, - is_fwd=False, - upd=False - ) - - # Single update step at the end - schedule[0][-1] = Task( - device_id=0, - stage_id=0, - mb_idx=0, - is_fwd=False, - upd=True - ) - - return schedule - - -def print_cuda_memory_stats(device: str, message: str = "") -> None: - """Prints CUDA memory usage statistics for the specified device, including available memory. - - Args: - device (str): The CUDA device identifier (e.g., 'cuda', 'cuda:0'). - message (str): Optional message to describe what is being logged (e.g., "loaded model", "loaded batch"). - """ - torch.cuda.synchronize() - device_obj = torch.device(device) - device_idx = device_obj.index if device_obj.index is not None else 0 - BYTES_IN_GB: int = 1024 ** 3 - - used_gb: float = torch.cuda.memory_allocated(device_idx) / BYTES_IN_GB - reserved_gb: float = torch.cuda.memory_reserved(device_idx) / BYTES_IN_GB - peak_gb: float = torch.cuda.max_memory_allocated(device_idx) / BYTES_IN_GB - - # Get total memory using torch.cuda.get_device_properties - total_gb: float = torch.cuda.get_device_properties(device_idx).total_memory / BYTES_IN_GB - available_gb: float = total_gb - reserved_gb - - message_prefix = f":{message}" if message else "" - print( - f"[CUDA:{device_idx}] Memory stats{message_prefix} | " - f"used={used_gb:.1f} GiB, reserved={reserved_gb:.1f} GiB, " - f"peak={peak_gb:.1f} GiB, available={available_gb:.1f} GiB, total={total_gb:.1f} GiB" - ) - - -def main(args): - - # Set model configuration based on argument - if args.model == 'LLAMA_DEBUG': - llama_config = LLAMA_DEBUG - elif args.model == 'LLAMA_1B': - llama_config = LLAMA_1B - elif args.model == 'LLAMA_3B': - llama_config = LLAMA_3B - elif args.model == 'LLAMA_8B': - llama_config = LLAMA_8B - - loss_fn = torch.nn.CrossEntropyLoss() - device = 'cuda' - - batch_size = args.batch_size - num_mbs = args.num_mbs - seq_len = args.seq_len - warmup = args.warmup - iters = args.iters - - x = torch.randint(0, llama_config.vocab_size, (batch_size, seq_len)).to(device) - y = torch.zeros((batch_size, llama_config.vocab_size), dtype=torch.long).to(device) - - # print_cuda_memory_stats(device, "after loading data") - - model = Transformer(llama_config) - model.to(device) - - # print_cuda_memory_stats(device, "after loading model") - compiled, compilation_metadata = piper_setup(model, [x], backend=piper) - # print_cuda_memory_stats(device, "after compiling model") - - # Update the global piper_metadata with the compilation metadata - # This is necessary because piper_exec and other functions expect the global metadata to be populated - piper_metadata.update(compilation_metadata) - - actors = compilation_metadata['actors'] - num_actors = len(actors) - - print(f"Number of actors (devices): {num_actors}") - - # For single device, we expect exactly 1 actor - if num_actors != 1: - print(f"WARNING: Expected 1 actor for single-device execution, but got {num_actors} actors.") - print(f"This may indicate that the model has distributed_stage annotations configured for {num_actors} stages.") - print(f"Proceeding with execution on actor 0 only.") - - # Send input and truth to the first (and only) actor - ray.get(actors[0].send_input.remote(x)) - ray.get(actors[0].send_truth.remote(y)) - - # Build simple single-device schedule - schedule = build_single_device_schedule(num_mbs) - - print("\nSCHEDULE:") - print_schedule(schedule) - print(f"Schedule details: {num_mbs} forward passes, {num_mbs} backward passes, 1 update\n") - - def iter_schedule(): - out = piper_exec(compiled, schedule, [x], y, loss_fn, num_mbs) - ray.get(out) - - # Set tracing - ray.get([actor.set_tracing.remote(args.tracing) for actor in actors.values()]) - - # Warmup iterations - print(f"Running {warmup} warmup iterations...") - for _ in range(warmup): - iter_schedule() - - # Clear timing data after warmup - ray.get([actor.clear_trace_data.remote() for actor in actors.values()]) - - # Timed iterations - print(f"Running {iters} timed iterations...") - start = time.perf_counter() - for _ in range(iters): - iter_schedule() - end = time.perf_counter() - - def print_mean_timing_data(trace_data: dict, actor_id: int) -> None: - """Prints mean timing and memory statistics from trace_data for a given actor. - - Args: - trace_data (dict): Trace data dictionary from the actor containing timing and memory metrics. - actor_id (int): The ID of the actor. - """ - import numpy as np - - print(f"\nTiming and Memory statistics for Actor {actor_id}:") - - for stage_id, stage_data in trace_data.items(): - if stage_id == 'update': - # Handle update data (global optimizer step) - print(f" Update:") - for metric, values in stage_data.items(): - if not values: - mean_val = float('nan') - else: - mean_val = float(np.mean(values)) - - if 'memory' in metric: - if 'delta' in metric: - print(f" {metric}: {mean_val:.3f} GB") - else: - print(f" {metric}: {mean_val:.3f} GB") - else: - print(f" {metric}: {mean_val:.3f} ms") - else: - # Handle stage data (forward/backward passes) - print(f" Stage {stage_id}:") - for phase, phase_data in stage_data.items(): - print(f" {phase.capitalize()}:") - for metric, values in phase_data.items(): - if not values: - mean_val = float('nan') - else: - mean_val = float(np.mean(values)) - - if 'memory' in metric: - if 'delta' in metric: - print(f" {metric}: {mean_val:.3f} GB") - else: - print(f" {metric}: {mean_val:.3f} GB") - else: - print(f" {metric}: {mean_val:.3f} ms") - - print("\n" + "="*60) - print("THROUGHPUT RESULTS:") - print("="*60) - print( - f"Throughput: {(iters * batch_size * num_mbs * seq_len)/(end - start):.0f} tokens/sec" - ) - print( - f"Average iteration time: {(end - start)*1000/iters:.2f} ms" - ) - print( - f"Total time for {iters} iterations: {(end - start):.2f} seconds" - ) - print("="*60) - - if args.tracing: - for actor in actors.values(): - trace_data = ray.get(actor.get_trace_data.remote()) - actor_id = ray.get(actor.id.remote()) - print_mean_timing_data(trace_data, actor_id) - - timeline_filename = f"out/{args.model}-single-device.json" - ray.timeline(timeline_filename) - print(f"\nRay timeline saved to: {timeline_filename}") - -if __name__ == "__main__": - ray.init(include_dashboard=True, log_to_driver=True, namespace="llama-single") - torch.manual_seed(0) - args = parse_args() - main(args) - - diff --git a/test/test_moe.py b/test/test_moe.py deleted file mode 100644 index 2647b8f..0000000 --- a/test/test_moe.py +++ /dev/null @@ -1,89 +0,0 @@ -import ray -import torch -import torch.nn as nn -import argparse -import time - -from src.piper_coordinator import PiperProgramCoordinator -from src.piper_compile import piper_setup -from src.piper_exec import piper_exec -from src.piper import distributed_stage, piper -from src.piper_actor import get_actor - -from .schedule_helpers import no_pp_schedule, print_schedule, build_1f1b_schedule - -from .models.moe import MixtureOfExperts, FFNExpert - -class MoETransformer(nn.Module): - def __init__(self, vocab_size, hidden_dim, experts_per_layer, k, world_size): - super().__init__() - - self.embedding = nn.Embedding(vocab_size, hidden_dim) - self.moe1 = MixtureOfExperts(hidden_dim, hidden_dim, experts_per_layer, FFNExpert, k=k) - self.moe2 = MixtureOfExperts(hidden_dim, hidden_dim, experts_per_layer, FFNExpert, k=k) - self.output = nn.Linear(hidden_dim, vocab_size) - - def forward(self, x): - distributed_stage(0) - x = self.embedding(x) - x = self.moe1(x) - - distributed_stage(1) - x = self.moe2(x) - return self.output(x) - - -def parse_args(): - parser = argparse.ArgumentParser(description='Run MoE model with pipeline and data parallelism') - parser.add_argument('--pp_degree', type=int, default=2, - help='Pipeline parallel degree (default: 2)') - parser.add_argument('--dp_degree', type=int, default=1, - help='Data parallel degree (default: 1)') - parser.add_argument('--num_mbs', type=int, default=1, - help='Number of microbatches (default: 1)') - return parser.parse_args() - - -def main(args): - - # Create MoE model - world_size = args.dp_degree * args.pp_degree - vocab_size = 2 - batch_size = 2 - seq_len = 2 - - x = torch.randint(0, vocab_size, (batch_size,)) - y = torch.randn(batch_size, vocab_size) - - pp_degree = args.pp_degree - world_size = args.dp_degree * args.pp_degree - num_mbs = args.num_mbs - - model = piper_setup( - MoETransformer, - (vocab_size, 2, 2, 1, world_size), - torch.optim.Adam, - [x], - num_stages=pp_degree, - pp_degree=pp_degree, - check_correct=True) - - if args.pp_degree == 1: - schedule = no_pp_schedule - else: - schedule = build_1f1b_schedule(num_mbs, pp_degree) - loss_fn = torch.nn.CrossEntropyLoss() - - print_schedule(schedule) - - losses = piper_exec(model, schedule, [x], y, loss_fn, num_mbs, pp_degree) - - ray.timeline(f"out/moe.json") - -if __name__ == "__main__": - ray.init(include_dashboard=False, log_to_driver=True, namespace="llama") - args = parse_args() - piper_coordinator = PiperProgramCoordinator.remote(pp_degree=args.pp_degree, dp_degree=args.dp_degree) - handles = piper_coordinator.run_program.remote(main, args) - ray.get(handles) - ray.shutdown() diff --git a/test/test_order_directive.py b/test/test_order_directive.py new file mode 100644 index 0000000..6a8dc80 --- /dev/null +++ b/test/test_order_directive.py @@ -0,0 +1,312 @@ +import pytest +import torch.fx as fx + +from src.piper import ( + TrainingDAG, + TrainingDAGEdge, + TrainingDAGNode, + apply_schedule_directives, + _apply_order_directive, + _apply_split_backward_stencil, + _apply_split_directive, + _bucket_matched_fwd_nodes, + _parse_order_directive, + _serial_topological_order, + _validate_schedule_tags_exist, + _validate_split_backward_order_stencil, +) +import src.directives as directives + + +def _compute(uid: str, tag: dict) -> TrainingDAGNode: + return TrainingDAGNode( + uid=uid, + node_kind="COMPUTE", + compute_subkind="FWD", + tag=tag, + device=[0], + stream="default_stream", + node_meta={}, + ) + + +def _bwd(uid: str, fwd_uid: str, tag: dict) -> TrainingDAGNode: + return TrainingDAGNode( + uid=uid, + node_kind="COMPUTE", + compute_subkind="BWD", + tag=tag, + device=[0], + stream="default_stream", + node_meta={"fwd_uid": fwd_uid, "bucket_key": fwd_uid}, + ) + + +def _dummy_gm() -> fx.GraphModule: + graph = fx.Graph() + x = graph.placeholder("x") + graph.output(x) + return fx.GraphModule({}, graph) + + +def test_parse_order_directive_accepts_nested_filter_groups() -> None: + directive = { + "op": "order", + "filters": [ + [ + {"PP": 0, "MB": 0, "PASS": "F"}, + {"PP": 0, "MB": 1, "PASS": "F"}, + ], + [ + {"PP": 0, "MB": 0, "PASS": "B"}, + ], + ], + } + + groups = _parse_order_directive(directive) + + assert groups == [ + [ + {"PP": 0, "MB": 0, "PASS": "F"}, + {"PP": 0, "MB": 1, "PASS": "F"}, + ], + [ + {"PP": 0, "MB": 0, "PASS": "B"}, + ], + ] + + +def test_parse_order_directive_rejects_flat_filter_groups() -> None: + directive = { + "op": "order", + "filters": [ + {"PP": 0, "MB": 0, "PASS": "F"}, + {"PP": 0, "MB": 0, "PASS": "B"}, + ], + } + + with pytest.raises(ValueError, match="must be a non-empty list"): + _parse_order_directive(directive) + + +def test_split_backward_stencil_handles_nested_filter_groups() -> None: + directive = { + "op": "order", + "filters": [ + [ + {"PP": 0, "MB": 0, "PASS": "BI"}, + {"PP": 0, "MB": 0, "PASS": "BW"}, + ], + [ + {"PP": 0, "MB": 1, "PASS": "F"}, + ], + ], + } + + split_keys = _validate_split_backward_order_stencil([directive]) + + assert split_keys == {(("MB", 0), ("PP", 0))} + + +def test_split_backward_stencil_allows_mixed_fused_and_split_microbatches() -> None: + directive = { + "op": "order", + "filters": [ + [{"PP": 0, "MB": 0, "PASS": "B"}], + [{"PP": 0, "MB": 1, "PASS": "BI"}], + [{"PP": 0, "MB": 1, "PASS": "BW"}], + ], + } + dag = TrainingDAG() + dag.add_node(_compute("fwd", {"PP": 0, "PASS": "F"})) + dag.add_node(_bwd("bwd", "fwd", {"PP": 0, "PASS": "B"})) + dag.add_edge(TrainingDAGEdge("fwd", "bwd", "data")) + + split_keys = _validate_split_backward_order_stencil([directive]) + _apply_split_directive(dag, {}, "MB", 2) + _apply_split_backward_stencil(dag, split_keys) + + assert dag.nodes["bwd"].compute_subkind == "BWD" + assert dag.nodes["bwd"].tag["PASS"] == "B" + split_bwd_uid = "bwd.splitMB1" + assert dag.nodes[split_bwd_uid].compute_subkind == "BWD_I" + assert dag.nodes[split_bwd_uid].tag["PASS"] == "BI" + assert dag.nodes[f"{split_bwd_uid}.bw"].compute_subkind == "BWD_W" + + +def test_bucket_rewrite_resolves_split_microbatch_bwd_by_metadata(monkeypatch) -> None: + dag = TrainingDAG() + fwd = _compute("fwd", {"PP": 0, "PASS": "F"}) + fwd.node_meta.update({ + "stage_id": 0, + "segment_id": 0, + "gm": _dummy_gm(), + "graphargs": [], + "input_idxs": [], + "param_idxs": [], + }) + dag.add_node(fwd) + dag.add_node(_bwd("bwd", "fwd", {"PP": 0, "PASS": "B"})) + dag.add_edge(TrainingDAGEdge("fwd", "bwd", "data")) + + _apply_split_directive(dag, {}, "MB", 2) + + def fake_bucket_stage(*_args, **_kwargs): + return [ + (_dummy_gm(), [], [], []), + (_dummy_gm(), [], [], []), + ] + + monkeypatch.setattr(directives, "bucket_stage", fake_bucket_stage) + + _bucket_matched_fwd_nodes(dag, [{"PP": 0}], 25) + + assert "fwd.bucket0" in dag.nodes + assert "fwd.bucket0.bwd" in dag.nodes + assert "fwd.splitMB1.bucket0" in dag.nodes + assert "fwd.splitMB1.bucket0.bwd" in dag.nodes + + +def test_apply_order_directive_groups_nested_subdags_with_dummy_source_sink() -> None: + dag = TrainingDAG() + for node in [ + _compute("prev", {"slot": 0}), + _compute("lane0.first", {"slot": 1, "lane": 0}), + _compute("lane0.last", {"slot": 1, "lane": 0}), + _compute("lane1", {"slot": 1, "lane": 1}), + _compute("next", {"slot": 2}), + ]: + dag.add_node(node) + dag.add_edge(TrainingDAGEdge("lane0.first", "lane0.last", "data")) + + _apply_order_directive( + dag, + [ + [{"slot": 0}], + [{"slot": 1, "lane": 0}, {"slot": 1, "lane": 1}], + [{"slot": 2}], + ], + directive_idx=7, + ) + + source_uid = "order.7.group1.source" + sink_uid = "order.7.group1.sink" + assert dag.nodes[source_uid].node_kind == "ORDER_DUMMY" + assert dag.nodes[sink_uid].node_kind == "ORDER_DUMMY" + + temporal_edges = { + (edge.src_uid, edge.dst_uid) + for edge in dag.edges + if edge.dep_kind == "temporal" + } + assert temporal_edges == { + ("prev", source_uid), + (source_uid, "lane0.first"), + (source_uid, "lane1"), + ("lane0.last", sink_uid), + ("lane1", sink_uid), + (sink_uid, "next"), + } + + topo = _serial_topological_order(dag) + assert topo.index("prev") < topo.index(source_uid) + assert topo.index(source_uid) < topo.index("lane0.first") + assert topo.index(source_uid) < topo.index("lane1") + assert topo.index("lane0.last") < topo.index(sink_uid) + assert topo.index("lane1") < topo.index(sink_uid) + assert topo.index(sink_uid) < topo.index("next") + + +def test_apply_order_directive_rejects_reverse_model_dataflow() -> None: + dag = TrainingDAG() + dag.add_node(_compute("producer", {"slot": 0})) + dag.add_node(_compute("consumer", {"slot": 1})) + dag.add_edge(TrainingDAGEdge("producer", "consumer", "data")) + + with pytest.raises(ValueError, match="violates model dataflow"): + _apply_order_directive( + dag, + [ + [{"slot": 1}], + [{"slot": 0}], + ], + directive_idx=3, + ) + + +def test_apply_order_directive_rejects_cross_device_order_edge() -> None: + dag = TrainingDAG() + dag.add_node(_compute("rank0", {"slot": 0})) + dag.add_node(_compute("rank1", {"slot": 1})) + dag.nodes["rank1"].device = [1] + + with pytest.raises(ValueError, match="crosses device placement"): + _apply_order_directive( + dag, + [ + [{"slot": 0}], + [{"slot": 1}], + ], + directive_idx=4, + ) + + +def test_shard_removes_zero_grad_lifetime_metadata_with_reduce_scatter() -> None: + dag = TrainingDAG() + node = TrainingDAGNode( + uid="expert.bwd", + node_kind="COMPUTE", + compute_subkind="BWD", + tag={"PP": 0, "EP": 0, "PASS": "B"}, + device=[0, 2], + stream="default_stream", + node_meta={"bucket_key": "expert", "fwd_uid": "expert.fwd"}, + ) + dag.add_node(node) + + directives._insert_reduce_scatter_comm_nodes(dag, [{"PP": 0, "EP": 0}], [0, 2]) + + assert node.node_meta["zero_alloc_full_grads_before"] is True + assert any(n.node_kind == "REDUCE_SCATTER_COMM" for n in dag.nodes.values()) + + directives._insert_shard_a2a_comm_nodes(dag, [{"PP": 0, "EP": 0}], [0, 2]) + + assert "zero_alloc_full_grads_before" not in node.node_meta + assert not any(n.node_kind == "REDUCE_SCATTER_COMM" for n in dag.nodes.values()) + + +def test_validate_schedule_tags_exist_rejects_tags_missing_from_model() -> None: + dag = TrainingDAG() + dag.add_node(_compute("fwd", {"PP": 0, "PASS": "F"})) + + directives = [ + {"op": "place", "filter": {"TP": 0}, "devices": [0]}, + {"op": "order", "filters": [[{"PP": 0, "MB": 0, "PASS": "F"}]]}, + ] + + with pytest.raises(ValueError, match="not found in model annotations"): + _validate_schedule_tags_exist(dag, directives) + + +def test_validate_schedule_tags_exist_ignores_runtime_tags() -> None: + dag = TrainingDAG() + dag.add_node(_compute("fwd", {"PP": 0, "PASS": "F"})) + + _validate_schedule_tags_exist( + dag, + [ + {"op": "split", "filter": {}, "dim_name": "MB", "num_microbatches": 2}, + {"op": "order", "filters": [[{"PP": 0, "MB": 0, "PASS": "F"}]]}, + ], + ) + + +def test_apply_schedule_directives_rejects_unmatched_place_directive() -> None: + dag = TrainingDAG() + dag.add_node(_compute("fwd", {"PP": 0, "PASS": "F"})) + + with pytest.raises(ValueError, match="matched zero nodes"): + apply_schedule_directives( + dag, + [{"op": "place", "filter": {"PP": 1}, "devices": [0]}], + ) diff --git a/test/test_piper_schedule.py b/test/test_piper_schedule.py new file mode 100644 index 0000000..46b2eed --- /dev/null +++ b/test/test_piper_schedule.py @@ -0,0 +1,72 @@ +import json + +import pytest + +from src.schedule import load_schedule_directives, load_schedule_info + + +def test_load_schedule_directives_accepts_current_filter_shape(tmp_path) -> None: + schedule_path = tmp_path / "schedule.json" + schedule_path.write_text( + json.dumps([ + {"op": "place", "filter": {"PP": 0}, "devices": [0, 1]}, + { + "op": "order", + "filters": [ + [{"PP": 0, "MB": 0, "PASS": "F"}], + [{"PP": 0, "MB": 0, "PASS": "B"}], + ], + }, + {"op": "split", "filter": {}, "dim_name": "MB", "num_microbatches": 2}, + ]), + encoding="utf-8", + ) + + directives = load_schedule_directives(str(schedule_path)) + + assert directives[0]["filter"] == {"PP": 0} + assert directives[1]["filters"] == [ + [{"PP": 0, "MB": 0, "PASS": "F"}], + [{"PP": 0, "MB": 0, "PASS": "B"}], + ] + + +def test_load_schedule_info_derives_parallel_degrees(tmp_path) -> None: + schedule_path = tmp_path / "qwen_example.json" + schedule_path.write_text( + json.dumps([ + {"op": "place", "filter": {"PP": 0}, "devices": [0, 1]}, + {"op": "place", "filter": {"PP": 1}, "devices": [2, 3]}, + {"op": "split", "filter": {}, "dim_name": "MB", "num_microbatches": 4}, + ]), + encoding="utf-8", + ) + + info = load_schedule_info(str(schedule_path)) + + assert info == { + "name": "qwen_example", + "path": str(schedule_path), + "num_stages": 2, + "pp_degree": 2, + "dp_degree": 2, + "num_microbatches": 4, + } + + +def test_load_schedule_directives_rejects_placeholder_microbatch_count(tmp_path) -> None: + schedule_path = tmp_path / "schedule.json" + schedule_path.write_text( + json.dumps([ + { + "op": "split", + "filter": {}, + "dim_name": "MB", + "num_microbatches": "__MBS__", + } + ]), + encoding="utf-8", + ) + + with pytest.raises(ValueError, match="__MBS__"): + load_schedule_directives(str(schedule_path))