Repository: pytorch/PiPPy

Issue:
Support for UNet1D and UNet2D models?
Is it possible to support UNet models with PiPPy?


Discussion:
Can you give us a pointer to the UNet models you refer to? Ideally model code from which we can build the model, or a HuggingFace from_pretrain link
HF Model Page: https://huggingface.co/docs/diffusers/en/api/models/unet2d
Github Code: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unets/unet_2d.py


Resolving patch:
diff --git a/examples/huggingface/pippy_unet.py b/examples/huggingface/pippy_unet.py
new file mode 100644
--- /dev/null
+++ b/examples/huggingface/pippy_unet.py
@@ -0,0 +1,99 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+# Minimum effort to run this example:
+# $ torchrun --nproc-per-node 2 pippy_unet.py
+
+import argparse
+import os
+
+import torch
+import torch.distributed as dist
+
+from pippy.IR import Pipe, PipeSplitWrapper, annotate_split_points
+from pippy.PipelineStage import PipelineStage
+
+from diffusers import UNet2DModel
+
+from hf_utils import get_number_of_params
+
+
+def run(args):
+    print("Using device:", args.device)
+
+    # Create model
+    # See https://github.com/huggingface/diffusers?tab=readme-ov-file#quickstart
+    unet = UNet2DModel.from_pretrained("google/ddpm-cat-256")
+    unet.to(args.device)
+    unet.eval()
+    if args.rank == 0:
+        print(f"Total number of params = {get_number_of_params(unet) // 10 ** 6}M")
+        print(unet)
+
+    # Input configs
+    sample_size = unet.config.sample_size
+    noise = torch.randn((args.batch_size, 3, sample_size, sample_size), device=args.device)
+    timestep = 1
+
+    # Split model into two stages:
+    #   Stage 0: down_blocks + mid_block
+    #   Stage 2: up_blocks
+    annotate_split_points(unet, {"mid_block": PipeSplitWrapper.SplitPoint.END})
+
+    # Create pipeline
+    unet_pipe = Pipe.from_tracing(
+        unet,
+        num_chunks=args.chunks,
+        example_args=(noise, timestep),
+    )
+    nstages = len(list(unet_pipe.split_gm.children()))
+    assert nstages == args.world_size, f"nstages = {nstages} nranks = {args.world_size}"
+    if args.rank == 0:
+        for i, sm in enumerate(unet_pipe.split_gm.children()):
+            print(f"Pipeline stage {i} {get_number_of_params(sm) // 10 ** 6}M params")
+
+    # Create schedule runtime
+    stage = PipelineStage(
+        unet_pipe,
+        args.rank,
+        device=args.device,
+    )
+
+    # Run
+    if args.rank == 0:
+        stage(noise)
+    elif args.rank == args.world_size - 1:
+        out = stage()
+    else:
+        stage()
+
+    print(f"Rank {args.rank} completes")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--world_size', type=int, default=int(os.getenv("WORLD_SIZE", 2)))
+    parser.add_argument('--rank', type=int, default=int(os.getenv("RANK", -1)))
+    parser.add_argument('--master_addr', type=str, default=os.getenv('MASTER_ADDR', 'localhost'))
+    parser.add_argument('--master_port', type=str, default=os.getenv('MASTER_PORT', '29500'))
+    parser.add_argument('--schedule', type=str, default="FillDrain")
+    parser.add_argument('--cuda', type=int, default=int(torch.cuda.is_available()))
+    parser.add_argument("--chunks", type=int, default=2)
+    parser.add_argument('--batch_size', type=int, default=2)
+    parser.add_argument('--batches', type=int, default=1)
+
+    args = parser.parse_args()
+
+    if args.cuda:
+        dev_id = args.rank % torch.cuda.device_count()
+        args.device = torch.device(f"cuda:{dev_id}")
+    else:
+        args.device = torch.device("cpu")
+
+    # Init process group
+    backend = "nccl" if args.cuda else "gloo"
+    dist.init_process_group(
+        backend=backend,
+        rank=args.rank,
+        world_size=args.world_size,
+    )
+
+    run(args)


Test patch: