Getting Started with PyTorch Fully Sharded Data Parallel (FSDP2) and Ray Train Template (#56298)

A new workspace template that walks users through how to integrate
PyTorch's FSDP2 with Ray Train. The purpose of this template is to allow
customers to quickly get started with FSDP and Ray Train whether they
are coming from using Distributed Data Parallel (DDP) or just getting
started with training large models.

For a high level overview, this template covers:

- A hands-on example of training an image classification model
- Model checkpoint saving and loading with PyTorch Distributed
Checkpoint (DCP)
- Configuring FSDP2 to mitigate out-of-memory (OOM) errors using mixed
precision, CPU offloading, sharding granularity, and more
- GPU memory profiling with PyTorch Profiler
- Loading a distributed model for inference

Link to original PR (pivoted to make available on OSS):
**anyscale/templates#463

**Testing**
- This notebook was tested in an [Anyscale
workspace](https://console.anyscale.com/cld_kvedZWag2qA8i5BjxUevf5i7/prj_cz951f43jjdybtzkx1s5sjgz99/workspaces/expwrk_nktjw7a3j2l5c7af9rh3n5rskw?workspace-tab=overview&command-history-section=application_logs&file=%252Fhome%252Fray%252Fdefault%252FREADME.ipynb)
and ran as expected

**For easy testing:**
Simply copy the notebook into an Anyscale workspace (and preferably the
image directory), provisioned with two T4 nodes.

---------

Signed-off-by: JasonLi1909 <jasli1909@gmail.com>
Signed-off-by: Jason Li <57246540+JasonLi1909@users.noreply.github.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: angelinalg <122562471+angelinalg@users.noreply.github.com>

Author: 3 people

doc/source/train/examples.yml

@@ -55,7 +55,13 @@ examples:
       - natural language processing
     contributor: community
     link: examples/intel_gaudi/bert
-
+  - title: Get started with PyTorch Fully Sharded Data Parallel (FSDP2) and Ray Train
+    skill_level: intermediate 
+    frameworks:
+      - pytorch
+    use_cases:
+      - computer vision
+    link: examples/pytorch/pytorch-fsdp/README
   - title: Train a text classifier with DeepSpeed
     frameworks:
       - deepspeed

doc/source/train/examples/pytorch/pytorch-fsdp/README.ipynb

@@ -0,0 +1,5 @@
+filegroup(
+    name = "ci_yamls",
+    srcs = ["aws.yaml", "gce.yaml"],
+    visibility = ["//release:__pkg__"],
+)

doc/source/train/examples/pytorch/pytorch-fsdp/README.md

@@ -0,0 +1,12 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-central1
+
+head_node_type:
+  name: head_node
+  instance_type: m5.2xlarge
+
+worker_node_types:
+  - instance_type: g4dn.xlarge
+    name: '1xT4:4CPU-16GB'
+    min_workers: 2
+    max_workers: 2

doc/source/train/examples/pytorch/pytorch-fsdp/ci/BUILD

@@ -0,0 +1,13 @@
+cloud_id: {{env["ANYSCALE_CLOUD_ID"]}}
+region: us-central1
+
+head_node_type:
+  name: head
+  instance_type: n2-standard-8
+
+worker_node_types:
+- name: gpu_worker
+  instance_type: n1-standard-8-nvidia-t4-16gb-1
+  min_workers: 2
+  max_workers: 2
+  use_spot: false

doc/source/train/examples/pytorch/pytorch-fsdp/ci/aws.yaml

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+import argparse
+import nbformat
+
+
+def convert_notebook(input_path: str, output_path: str) -> None:
+    """
+    Read a Jupyter notebook and write a Python script, converting all %%bash
+    cells and IPython "!" commands into subprocess.run calls that raise on error.
+    Cells that load or autoreload extensions are ignored.
+    """
+    nb = nbformat.read(input_path, as_version=4)
+    with open(output_path, "w") as out:
+        for cell in nb.cells:
+            # Only process code cells
+            if cell.cell_type != "code":
+                continue
+
+            lines = cell.source.splitlines()
+            # Skip cells that load or autoreload extensions
+            if any(
+                l.strip().startswith("%load_ext autoreload")
+                or l.strip().startswith("%autoreload all")
+                for l in lines
+            ):
+                continue
+
+            # Detect a %%bash cell
+            if lines and lines[0].strip().startswith("%%bash"):
+                bash_script = "\n".join(lines[1:]).rstrip()
+                out.write("import subprocess\n")
+                out.write(
+                    f"subprocess.run(r'''{bash_script}''',\n"
+                    "               shell=True,\n"
+                    "               check=True,\n"
+                    "               executable='/bin/bash')\n\n"
+                )
+            else:
+                # Detect any IPython '!' shell commands in code lines
+                has_bang = any(line.lstrip().startswith("!") for line in lines)
+                if has_bang:
+                    out.write("import subprocess\n")
+                    for line in lines:
+                        stripped = line.lstrip()
+                        if stripped.startswith("!"):
+                            cmd = stripped[1:].lstrip()
+                            out.write(
+                                f"subprocess.run(r'''{cmd}''',\n"
+                                "               shell=True,\n"
+                                "               check=True,\n"
+                                "               executable='/bin/bash')\n"
+                            )
+                        else:
+                            out.write(line.rstrip() + "\n")
+                    out.write("\n")
+                else:
+                    # Regular Python cell: dump as-is
+                    out.write(cell.source.rstrip() + "\n\n")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="Convert a Jupyter notebook to a Python script, preserving bash cells and '!' commands as subprocess calls."
+    )
+    parser.add_argument("input_nb", help="Path to the input .ipynb file")
+    parser.add_argument("output_py", help="Path for the output .py script")
+    args = parser.parse_args()
+    convert_notebook(args.input_nb, args.output_py)
+
+
+if __name__ == "__main__":
+    main()

doc/source/train/examples/pytorch/pytorch-fsdp/ci/gce.yaml

@@ -0,0 +1,4 @@
+#!/bin/bash
+python ci/nb2py.py README.ipynb README.py  # convert notebook to py script
+python README.py  # run the converted python script
+rm README.py  # remove the generated script

doc/source/train/examples/pytorch/pytorch-fsdp/ci/nb2py.py

@@ -0,0 +1,9 @@
+head_node_type:
+  name: head_node
+  instance_type: m5.2xlarge
+
+worker_node_types:
+  - instance_type: g4dn.xlarge
+    name: '1xT4:4CPU-16GB'
+    min_workers: 2
+    max_workers: 2

doc/source/train/examples/pytorch/pytorch-fsdp/ci/tests.sh

@@ -0,0 +1,10 @@
+head_node_type:
+  name: head
+  instance_type: n2-standard-8
+
+worker_node_types:
+- name: gpu_worker
+  instance_type: n1-standard-8-nvidia-t4-16gb-1
+  min_workers: 2
+  max_workers: 2
+  use_spot: false