Foundation Model Distributed Training with Ray | Anyscale | Ray Summit 2026

Summary of technological concepts, features, and tutorial content

Course/session overview + setup (Ray Summit 2026)

• The session is part of a “morning session” course on foundation model distributed training using Ray and Anyscale.
• Attendees use an Anyscale UI workspace that:
  • launches an Anyscale cluster
  • loads notebook materials into a VS Code–like environment (notebooks + terminal)
• Slido is used for questions during the course.
• Typical course structure:
  • Notebook 0: environment + intro checks
  • Notebook 1: distributed training with Ray Train + Ray Data using DDP
  • Notebook 2: extend to models too large for one GPU using FSDP and DeepSpeed ZeRO
  • Notebook 3: 2D parallelism (Tensor Parallel + Data Parallel), including 2D device meshes
  • Notebook 4: fault tolerance & observability
  • Wrap-up: memory planning / decision framework

---

Notebook 0: Foundation + infrastructure checks

Key ideas

• Training performance is dominated by matrix multiplications; GPU advantage generally increases over time.
• Modern scaling relies on more data and larger models (scaling laws), but introduces a memory wall.
• Example model sizes and rough memory implications:
  • ~1B parameters → tens of GB (example given: ~28 GB on GPU for training)
  • ~70B parameters → can require terabytes, motivating distributed techniques
• Ray role (high level):
  • Turns CPU/GPU clusters into a system built from primitives, while Anyscale provisions/operates the cluster.
• Conceptual “axes” of parallelism:
  • Data parallelism (gradient all-reduce)
  • Tensor parallelism (split layers/weights; requires activation communication such as all-reduce/all-to-all)
  • Pipeline parallelism (layer stages; sequence/context handling)
  • Sequence/context parallelism
  • Expert parallelism (MoE routing; token all-to-one)

Hardware/network considerations

• The choice of parallelism depends on communication frequency and bandwidth:
  • DDP gradient sync is relatively less frequent than tensor-parallel per-layer communication.

Workspace/infrastructure specifics

• Connect to the provisioned Ray cluster with:
  • ray.init(address="auto")
• The setup checks hardware limits and precision support:
  • Example uses Torch float16; bf16 not natively supported on T4s
• Displays cluster topology:
  • head node vs workers
  • GPU/CPU counts
  • instance types (e.g., T4)

---

Notebook 1: Ray Train + Ray Data for DDP (1 GPU → multi-GPU with minimal code changes)

Core tutorial message

• Ray enables scaling from 1 to many GPUs with small code changes.

Key components

1. Ray primitives + libraries

   • Ray Core includes:
     • Tasks (decorated functions)
     • Actors (decorated classes; stateful)
   • Ray AI libraries include:
     • Ray Data, Ray Train, Ray Tune, Ray Serve, rlib
   • This course focuses mainly on Ray Train and Ray Data.

2. Ray Data streaming pipeline

   • Explains why naive preprocessing causes GPU “sawtooth” utilization (GPU idle during CPU preprocessing/loading).
   • Ray Data uses a streaming/lazy execution model with data blocks:
     • CPUs load/preprocess blocks
     • GPUs consume processed blocks
     • pipelining overlaps loading/preprocessing with training
   • Mentions lazy materialization (compute only what’s needed when it’s needed).
   • Emphasizes map_batches for distributed, batch-wise preprocessing (e.g., tokenization).

3. Ray Train distributed training (DDP)

   • Uses TorchTrainer configured with:
     • a train function executed on each worker
     • a scaling configuration (number of workers/GPUs)
     • a run configuration (where checkpoints/artifacts/metrics are saved)
   • For DDP:
     • wraps the model with Ray Train’s DDP-aware preparation (via prepare_model) to match PyTorch DDP expectations
     • shards data per worker using ray.train.get_dataset_shard()
   • Worker-side training flow:
     • shard dataset → iterate batches from Ray Data’s iterator
     • forward/backward/optimizer logic remains largely “single GPU-like”
   • Communication:
     • DDP uses gradient all-reduce → synchronized weight updates across GPUs
   • Checkpointing:
     • in the example, checkpointing occurs on rank zero only to avoid duplicate identical checkpoints.

Example

• Trains a small LLM (GPT-2) on AG News, using data from S3 mirrors (to avoid download issues).

Scaling claim

• Scaling to “400 GPUs” (in principle) mainly means changing one number: the number of workers in the scaling config.

---

Notebook 2: When DDP assumptions fail → model sharding with FSDP and DeepSpeed ZeRO

Problem addressed

• DDP assumes the model fits on each GPU. If it doesn’t, use model parallelism/sharding.

Option A: FSDP (Fully Sharded Data Parallel)

Highlights

• Keeps the data-parallel axis, but shards model parameters across GPUs.
• Uses mesh construction to drive sharding:
  • sharding dimension relates to world size and a “data-parallel group” concept
• Differences vs DDP:
  • DDP: gradient all-reduce; each GPU keeps full gradients
  • FSDP: backward triggers:
    • all-gather of shards for computation
    • then reduce-scatter
    • then local optimizer updates
• Checkpoints:
  • Unlike rank-zero-only DDP checkpointing, FSDP checkpoints are sharded across ranks using PyTorch Distributed Checkpoint (DCP).

Option B: DeepSpeed ZeRO

Highlights

• DeepSpeed sharding is configuration-driven rather than heavily API-wrapped like FSDP.
• Initializes with a DeepSpeed config specifying:
  • ZeRO stage (0/1/2/3)
  • batch sizes, precision, and other training settings
• ZeRO stage interpretation:
  • Stage 1: shard optimizer state
  • Stage 2: shard optimizer state + gradients
  • Stage 3: shard optimizer state + gradients + parameters
• Engine-driven differences:
  • training step uses the DeepSpeed engine methods for forward/backward/step
  • checkpointing uses engine.save_checkpoint
• Checkpointing:
  • all ranks participate since model state is partitioned/sharded.

Comparison & composition

• FSDP vs DeepSpeed:
  • FSDP: more PyTorch-native and flexible per module
  • DeepSpeed ZeRO: stage/config driven; easier swaps but different integration model
• Compatibility notes for later tensor parallel:
  • FSDP composes cleanly with tensor parallel (via DTensor-style components)
  • DeepSpeed AutoTP is mentioned as compatible with ZeRO stages 1/2, not 3

---

Notebook 3: 2D parallelism (Tensor Parallel + Data Parallel) with a Ray/Nested pattern

Main goal

• Move from “1D splitting” to 2D parallelism:
  • TP (tensor parallel): split computation across layer matrices
  • DP (data parallel): replicate model and sync gradients

Device mesh

• Uses a 2×2 grid example (4 GPUs) to illustrate rank mapping:
  • TP groups share the same data shard
  • DP groups represent replicas for gradient sync

Tensor parallel types

• Row-wise TP
  • partitions weight matrices/inputs
  • requires all-reduce to reconstruct outputs
• Column-wise TP
  • partitions differently; output remains sharded in the described setup
• Transformer layering pattern:
  • alternating row/column schemes by sub-layer (e.g., attention Q/K/V vs output projection; MLP expand/compress with appropriate row/column settings).

Ray integration pattern

• Still follows Ray Train’s wrapper approach:
  • provide a train loop per worker
  • provide scaling config (workers = TP size × DP size)
  • provide Ray Data dataset streaming/sharding logic

Training correctness requirement

• Data sharding must respect DP rank (not global world rank), so all TP ranks within the same DP group see the same shard.

FSDP nesting inside 2D

• The described approach nests FSDP sharding over the data-parallel dimension, atop tensor-split layers.

Example

• Trains a ~500M parameter model (“Llama 2.5” mentioned) on Wikitext using staged data from S3.

Communication scaling guidance

• Tensor parallel has high-frequency communication (per layer), so placement matters.
• Recommended heuristic:
  • keep TP “families” intra-node
  • keep DP “families” across nodes (DP comm is gradient sync per step, less frequent than TP layer comms)
• Ray placement-related primitives are referenced generally.

---

Notebook 4: Fault tolerance & observability (production readiness)

Fault tolerance: checkpoint-aware resumption + worker restart

• Uses Ray Train failure configuration inside Ray Train’s run config.
• Demonstrates:
  • checkpointing at epoch boundaries (rank-zero DDP demo)
  • intentionally killing a worker
  • Ray restarts training automatically
  • training resumes from the last checkpoint (restoring epoch number)
• Implementation detail in the train loop:
  • resume_from_checkpoint = ray.train.get_checkpoint() style logic
  • checkpoint contains:
    • model state
    • optimizer state
    • current epoch
  • these are saved to ensure correct resumption.

Elastic training (resource variability)

• Ray Train supports elastic training by specifying worker counts as an ordered pair: (min, max) instead of a fixed integer.
• Useful when:
  • clusters experience frequent node loss
  • jobs are preempted (low-priority resources reclaimed)
• Caveat:
  • changing worker count can effectively change global batch size if per-worker batch sizes remain fixed
  • can affect convergence/hyperparameter stability (less problematic at very large GPU counts).

Mid-epoch resumption

• If epochs take days, resuming only at epoch boundaries is insufficient.
• Ray Data can preserve iterator/data-shard progress so restarts can resume within an epoch (intra-epoch time).

Observability: metrics, dashboards, and profiling

Two levels of observability:

1. Job-level / cluster metrics
   • Anyscale UI and Grafana dashboards:
     • node count
     • utilization over time
     • GPU/CPU utilization
2. Work-level dashboards
   • Ray Data dashboard:
     • per data pipeline step status and bottlenecks
   • Ray Train dashboard:
     • logs per trainer.fit call
     • per-worker status

Debugging guidance

• Hang sources:
  • check CPU stack traces to identify where execution stalls
• Performance debugging:
  • download GPU profiling JSON
  • profile views separate:
    • data transfer CPU → GPU
    • forward kernels
    • backward kernels
    • gradient synchronization (all-reduce)
  • example interpretation:
    • if gradient sync dominates (“dark blue”), likely a network bandwidth bottleneck

---

Wrap-up: decision framework + memory planning heuristic

Parallelism decision framework

• Choose based on constraints and bottlenecks:
  • If model fits and you want throughput → data parallel
  • If memory constrained and optimizer/gradients dominate → DeepSpeed ZeRO or FSDP
  • If a layer is too wide → tensor parallel
  • If model is too deep / too many layers to fit → pipeline parallel
  • If sequences are extremely long → sequence/context parallel
  • If MoE / mixture-of-experts → expert parallel

Memory planner

• Uses per-GPU memory estimates (parameters, optimizer, activations) to judge whether a parallelism strategy fits target GPU memory.

---

Main speakers/sources

• Dr. Ian Yordan (primary instructor; introduced the session)
• Ray/Anyscale documentation and course notebooks (referenced throughout: notebook 0–4 and examples for Ray Data/Train/FSDP/DeepSpeed)
• TAs and Anyscale support (presence noted; no specific TA names provided)
• External frameworks referenced
  • PyTorch (DDP, distributed tensor parallel concepts, FSDP)
  • DeepSpeed (ZeRO)
  • Hugging Face–style tooling (e.g., GPT-2, tokenization)
  • Profiling/visualization tools (mentioned in profiling context; e.g., Grafana/TensorBoard-style tooling)