Streaming distributed execution across. CPUs and GPUs

Some of the most demanding machine learning (ML) use cases we have encountered involve pipelines that span both CPU and GPU devices in distributed environments. These situations are common workloads, including:

* Batch inference, which involves a CPU-intensive preprocessing stage (e.g., video decoding or image resizing) before utilizing a GPU-intensive model to make predictions.
* Distributed training, where similar CPU-heavy transformations are required to prepare or augment the dataset prior to GPU training.

In this talk, we examine how Ray data streaming works and how to use it for your own machine learning pipelines to address these common workloads utilizing all your compute resource–CPUs and GPUs–at scale.

Transcript

1. Streaming distributed execution across CPUs and GPUs June 21, 2023 Eric Liang Email: ekl@anyscale.com
2. Talk Overview • ML Inference and Training workloads • How it relates to Ray Data streaming (new feature in Ray 2.4!) • Examples • Backend overview
3. About me • Technical lead for Ray / OSS at Anyscale • Previously: ◦ PhD in systems / ML at Berkeley ◦ Staff eng @ Databricks, storage infra @ Google
4. ML workloads and Data
5. ML Workloads • Where does data processing come in for ML workloads? ETL Pipeline Preprocessing Training / Inference
6. ML Workloads ETL Pipeline Preprocessing Training / Inference Resizing images, Decoding videos Data augmentation Using PyTorch, TF, HuggingFace, LLMs, etc. Ingesting latest data, Joining data tables
7. ML Workloads ETL Pipeline Preprocessing Training / Inference CPU CPU GPU Resizing images, Decoding videos Data augmentation Using PyTorch, TF, HuggingFace, LLMs, etc. Ingesting latest data, Joining data tables
8. ML Workloads ETL Pipeline Preprocessing Training / Inference CPU CPU GPU Resizing images, Decoding videos Data augmentation usual scope of ML teams Ingesting latest data, Joining data tables Using PyTorch, TF, HuggingFace, LLMs, etc.
9. Our vision for simplifying this with Ray ETL Pipeline Preprocessing Training / Inference
10. Our vision for simplifying this with Ray ETL Pipeline Preprocessing Training / Inference
11. Our vision for simplifying this with Ray ETL Pipeline Preprocessing Training / Inference
12. Ray Data: Overview
13. Ray Data overview
14. Ray Data overview
15. Ray Data overview ray.data.Dataset Node 1 Block Node 2 Block Block Node 3 Block Blocks
16. Ray Data overview Powered by ray.data.Dataset
17. Ray Data overview High performance distributed IO ds = ray.data.read_parquet("s3://some/bucket") ds = ray.data.read_csv("/tmp/some_file.csv") Leverages Apache Arrow’s high-performance single-threaded IO Parallelized using Ray’s high-throughput task execution Scales to PiB-scale jobs in production (Amazon) Read from storage Transform data ds = ds.map_batches(batch_func) ds = ds.map(func) ds.iter_batches() -> Iterator ds.write_parquet("s3://some/bucket") Consume data
18. Ray Data Streaming
19. Bulk execution Previous versions of Ray Data (<2.4) used bulk execution strategy What is bulk execution? • Load all data into memory • Apply transformations on in-memory data in bulk • Out of memory? -> spill blocks to disk • Similar to Spark's execution model (bulk synchronous parallel)
20. Streaming (pipelined) execution • Default execution strategy for Ray Data in 2.4 • Same data transformations API • Instead of executing operations in bulk, build a pipeline of operators • Data blocks are streamed through operators, reducing memory use and avoiding spilling to disk
21. Preprocessing can often be the bottleneck • Example: video decoding prior to inference / training • Too expensive to run on just GPU nodes: needs scaling out • Large intermediate data: uses lots of memory
22. Ray Data streaming avoids the bottleneck • E.g., intermediate video frames streamed through memory • Decoding can be offloaded onto CPU nodes from GPU nodes • Intermediate frames kept purely in (cluster) memory
23. Inference <> Training • Same streaming pipeline can easily be used for training too!
24. Inference <> Training • Same streaming pipeline can easily be used for training too! Split[3] Worker [0] Worker [1] Worker [2]
25. Streaming performance benefits deep dive
26. Performance benefits overview Bulk execution Streaming execution CPU-only pipelines (single-stage) Heterogeneous CPU+GPU pipelines (multi-stage)
27. Performance benefits overview Bulk execution Streaming execution CPU-only pipelines (single-stage) - Memory optimal - Good for inference - Bad for training Heterogeneous CPU+GPU pipelines (multi-stage)
28. Performance benefits overview Bulk execution Streaming execution CPU-only pipelines (single-stage) - Memory optimal - Good for inference - Bad for training - Memory optimal - Good for inference - Good for training Heterogeneous CPU+GPU pipelines (multi-stage)
29. Performance benefits overview Bulk execution Streaming execution CPU-only pipelines (single-stage) - Memory optimal - Good for inference - Bad for training - Memory optimal - Good for inference - Good for training Heterogeneous CPU+GPU pipelines (multi-stage) - Memory inefficient - Slower for inference - Bad for training
30. Performance benefits overview Bulk execution Streaming execution CPU-only pipelines (single-stage) - Memory optimal - Good for inference - Bad for training - Memory optimal - Good for inference - Good for training Heterogeneous CPU+GPU pipelines (multi-stage) - Memory inefficient - Slower for inference - Bad for training - Memory optimal - Good for inference - Good for training
31. In more detail: simple batch inference job Logical data flow:
32. A simple batch inference job
33. A simple batch inference job
34. A simple batch inference job
35. A simple batch inference job
36. A simple batch inference job
37. A simple batch inference job
38. This is a single stage pipeline
39. Bulk physical execution output 1 output 2 output 3
40. Bulk physical execution -- single stage • Memory usage is optimal (no intermediate data) • Good for inference • Not good for distributed training (cannot consume results incrementally)
41. Streaming physical execution Operator (Stage 1) data partition 1 data partition 2 data partition 3
42. Streaming physical execution Operator (Stage 1) data partition 1 data partition 2 data partition 3
43. Streaming physical execution Operator (Stage 1) data partition 1 data partition 2 data partition 3 output 1
44. Streaming physical execution Operator (Stage 1) data partition 1 data partition 2 data partition 3 output 1
45. Streaming physical execution Operator (Stage 1) data partition 1 data partition 2 data partition 3 output 1 output 2
46. Streaming physical execution Operator (Stage 1) data partition 1 data partition 2 data partition 3 output 1 output 2
47. Streaming physical execution Operator (Stage 1) data partition 1 data partition 2 data partition 3 output 1 output 2 output 3
48. Streaming physical execution -- single stage • Memory usage is optimal (no intermediate data) • Good for inference • Good for distributed training
49. In more detail: multi-stage (heterogeneous) pipeline GPU
50. Heterogeneous pipeline (CPU + GPU)
51. Heterogeneous pipeline (CPU + GPU)
52. Heterogeneous pipeline (CPU + GPU) GPU
53. Compare against bulk vs streaming
54. Bulk physical execution
55. Bulk physical execution
56. Bulk physical execution
57. Bulk physical execution -- multi stage • Memory usage is inefficient (disk spilling) • Slower for inference • Bad for distributed training
58. Streaming physical execution Operator (Stage 1) Operator (Stage 2) Operator (Stage 3)
59. Streaming physical execution Operator (Stage 1) Operator (Stage 2) Operator (Stage 3)
60. Streaming physical execution Operator (Stage 1) Operator (Stage 2) Operator (Stage 3)
61. Streaming physical execution Operator (Stage 1) Operator (Stage 2) Operator (Stage 3)
62. Streaming physical execution Operator (Stage 1) Operator (Stage 2) Operator (Stage 3)
63. Streaming physical execution Operator (Stage 1) Operator (Stage 2) Operator (Stage 3)
64. Streaming physical execution -- multi stage • Memory usage is optimal (no intermediate data) • Good for inference • Good for distributed training
65. Comparison to other systems • DataFrame systems (e.g., Spark) ◦ Ray Data streaming is more memory efficient ◦ Ray Data supports heterogeneous clusters ◦ Execution model a better fit for distributed training • ML ingest libraries: TF Data / Torch Data / Petastorm ◦ Ray Data supports scaling preprocessing out to a cluster
66. Video Inference Example
67. Four stage pipeline Logical data flow:
68. Four stage pipeline
69. Four stage pipeline
70. Four stage pipeline
71. Putting it together
72. Putting it together
73. Putting it together
74. Putting it together
75. Putting it together
76. Streaming execution plan
77. Running this on a Ray cluster $ python workload.py
78. Running this on a Ray cluster $ python workload.py 2023-05-02 15:10:01,105 INFO streaming_executor.py:91 -- Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(decode_frames)] -> ActorPoolMapOperator[MapBatches(FrameAnnotator)] -> ActorPoolMapOperator[MapBatches(FrameClassifier)] -> TaskPoolMapOperator[Write]
79. Running this on a Ray cluster $ python workload.py 2023-05-02 15:10:01,105 INFO streaming_executor.py:91 -- Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(decode_frames)] -> ActorPoolMapOperator[MapBatches(FrameAnnotator)] -> ActorPoolMapOperator[MapBatches(FrameClassifier)] -> TaskPoolMapOperator[Write] Running: 25.0/112.0 CPU, 2.0/2.0 GPU, 33.19 GiB/32.27 GiB object_store_memory: 28%|███▍ | 285/1000 [01:40<03:12, 3.71it/s]
80. Running this on a Ray cluster $ python workload.py 2023-05-02 15:10:01,105 INFO streaming_executor.py:91 -- Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(decode_frames)] -> ActorPoolMapOperator[MapBatches(FrameAnnotator)] -> ActorPoolMapOperator[MapBatches(FrameClassifier)] -> TaskPoolMapOperator[Write] Running: 0.0/112.0 CPU, 0.0/2.0 GPU, 0.0 GiB/32.27 GiB object_store_memory: 100%|████████████| 1000/1000 [05:13<00:00, 3.85it/s]
81. Ray dashboard observability: active tasks
82. Ray dashboard observability: actor state
83. Ray dashboard observability: network
84. Backend overview
85. How does Ray data streaming work? • Each transformation is implemented as an operator • Use Ray tasks and actors for execution of operators ◦ default -> use Ray tasks ◦ actor pool -> use Ray actors • Intermediate data blocks in Ray object store • Memory usage of operators is limited (backpressure) to enable efficient streaming without spilling to disk
86. Advantages of using Ray core primitives • Heterogeneous clusters support • Fault tolerance out of the box ◦ Lineage-based reconstruction of tasks and actors operations → your ML job will survive failures during preprocessing • Resilient object store layer ◦ Spills to disk in case of unexpectedly high memory usage: slowdown instead of a crash ◦ Can also do large-scale shuffles in a pinch • Easy to add data locality optimizations for both task + actor ops
87. Scalability • Stress tests on a 20TiB array dataset • 500 machines
88. Training • Can use pipelines for training as well • Swap map_batches(Model) call for streaming_split(K) Inference ray.data.read_datasource(...) \ .map_batches(preprocess) .map_batches(Model, compute=ActorPoolStrategy(...)) \ .write_datasource(...) Training iters = ray.data.read_datasource(...) \ .map_batches(preprocess) \ .streaming_split(len(workers)) for i, w in enumerate(workers): w.set_data_iterator.remote(iters[i]) ## in worker for batch in it.iter_batches(batch_size=32): model.forward(batch)...
89. Advantages of streaming for Training • Example of accelerating an expensive Read/Preprocessing operation by adding CPU nodes to a cluster
90. Summary • Ray Data streaming scales batch inference and training workloads • More efficient computation model than bulk processing • Simple API for composing streaming topologies Next steps: • Streaming Inference is available in 2.4: docs.ray.io • Ray Train integration coming in 2.6 • We're hardening streaming to work robustly at 100+ node clusters, 10M+ input files. Contact us!