464 karma · joined October 16, 2017
Data processing workloads are quite common on Ray, especially with unstructured data.
Also, I work on Ray, which is the underlying framework used here, but all the work in the post was done by the Amazon team.
But I'm not aware of anything exactly like what you're referring to!
The scenario where a Ray task forks is probably not very well supported. You can certainly start a subprocess from within a Ray task, but I think forking could easily cause issues.
You can definitely use Ray + Jax, but you probably need to avoid forking a process within a Ray worker.
Our intention from the start was for Ray to be general purpose. And the core Ray APIs are quite general (basically just scheduling a Python function somewhere in a cluster or instantiating a Python class as a process somewhere in the cluster).
We had AI use cases in mind from the start, since we were grad students in AI. But the generality has really been important since AI workloads encompass a huge variety of computational patterns (allreduce style communication patterns on GPUs for training, embarrassingly parallel data processing workloads on spot instances, and so on).
Fun fact, very early on, we used to create one mmapped file per serialized object, but that very quickly broke down.
Then we switched to mmapping one large file at the start and storing all of the serialized objects in that file. But then as objects get allocated and deallocated, you need to manage the memory inside of that mmapped file, and we just repurposed a malloc implementation to handle that.
[1] https://github.com/ray-project/ray/blob/21202f6ddc3ceaf74fbc...
One of the key things is to make sure the serialized object is stored in a data format where the serialized object does not need to be "transformed" in order to access it. For example, a numpy array can be created in O(1) time from a serialized blob by initializing a Python object with the right shape and dtype and a pointer to the right offset in the serialized blob. We also use projects like Apache Arrow that put a lot of care into this.
Example in more detail:
Imagine the object you are passing from process A to process B is a 1GB numpy array of floats. In the serialization step, process A produces a serialized blob of bytes that is basically just the 1GB numpy array plus a little bit of metadata. Process A writes that serialized blob into shared memory. This step of "writing into shared memory" still involves O(N) work, where N is the size of the array (though you can have multiple threads do the memcpy in parallel and be limited just by memory bandwidth).
In the deserialization step, process B accesses the same shared memory blob (process A and B are on the same machine). It reads a tiny bit of metadata to figure out the type of the serialized object and shape and so on. Then it constructs a numpy array with the correct shape and type and with a pointer to the actual data in shared memory at the right offset. Therefore it doesn't need to touch all of the bytes of data, it just does O(1) work instead of O(N).
That's the basic idea. You can imagine generalizing this beyond numpy arrays, but it's most effective for objects that include large numerical data (e.g., objects that include numpy arrays).
There are a bunch of little details to get right, e.g., serializing directly into shared memory instead of creating a serialized copy in process A and then copying it into shared memory. Doing the write into shared memory in parallel with a bunch of threads. Getting the deserialization right. You also have to make sure that the starting addresses of the numpy arrays are 64-byte aligned (if memory serves) so that you don't accidentally trigger a copy later on.
EDIT: I edited the above to add more detail.
1. This is truly impressive work from AWS. Patrick Ames began speaking about this a couple years ago, though at this point the blog post is probably the best reference. https://www.youtube.com/watch?v=h7svj_oAY14
2. This is not a "typical" Ray use case. I'm not aware of any other exabyte scale data processing workloads. Our bread and butter is ML workloads: training, inference, and unstructured data processing.
3. We have a data processing library called Ray Data for ingesting and processing data, often done in conjunction with training and inference. However, I believe in this particular use case, the heavy lifting is largely done with Ray's core APIs (tasks & actors), which are lower level and more flexible, which makes sense for highly custom use cases. Most Ray users use the Ray libraries (train, data, serve), but power users often use the Ray core APIs.
4. Since people often ask about data processing with Ray and Spark, Spark use cases tend to be more geared toward structured data and CPU processing. If you are joining a bunch of tables together or running SQL queries, Spark is going to be way better. If you're working with unstructured data (images, text, video, audio, etc), need mixed CPU & GPU compute, are doing deep learning and running inference, etc, then Ray is going to be much better.
[1] https://docs.endpoints.anyscale.com/supported-models/Meta-Ll...
Just want to share that Ray (an open source project we're developing at Anyscale), can be used to scale Jax (e.g., across TPUs).
Some docs from Google on how to do this
https://cloud.google.com/tpu/docs/ray-guide
Alpa is an open source project scaling Jax on 1000+ GPUs
https://www.anyscale.com/blog/training-175b-parameter-langua...
Cohere uses Ray + Jax + TPUs to build their LLMs
https://www.youtube.com/watch?v=For8yLkZP5w
A demo from Matt Johnson on the Jax team
https://www.anyscale.com/blog/a-comprehensive-guide-for-buil...
- Llama-2-70B: $1 (on Anyscale Endpoints [1]) - GPT-3.5-turbo: $1.50 - $2 (OpenAI [2])
[1] https://app.endpoints.anyscale.com/ [2] https://openai.com/pricing
The models (and current prices) right now are - Llama-2-7B ($0.25 / million tokens) - Llama-2-13B ($0.50 / million tokens) - Llama-2-70B ($1 / million tokens) - Code Llama ($1 / million tokens)
You can try out CodeLlama-34B on Anyscale Endpoints (an LLM inference API we're building here at Anyscale for open source LLMs).
Some examples - https://huggingface.co/datasets/b-mc2/sql-create-context - https://huggingface.co/datasets/GEM/viggo
On the other hand, 8K examples was not enough to learn to solve grade school math problems [2], so it is very problem dependent.
[1] https://www.anyscale.com/blog/fine-tuning-llama-2-a-comprehe...
That said, many valuable LLM products / features are more narrow in scope and can see a huge lift from fine-tuning. We've run a bunch of experiments on this (e.g., SQL query generation is a good example), where fine-tuning even the 7B Llama-2 model outperforms GPT-4 (surprisingly) [1]. That's a very different type of problem from teaching software engineering of course.
[1] https://www.anyscale.com/blog/fine-tuning-llama-2-a-comprehe...
Llama-2-70B is $1 / million tokens, which is the most cost-efficient on the market that I'm aware of.
In fact, we've seen GPT-4 level performance from even the 7B Llama-2 model after fine-tuning. [1]
[1] https://www.anyscale.com/blog/fine-tuning-llama-2-a-comprehe...
I think you're right about this, and benchmarks we've run at Anyscale support this conclusion [1].
The caveat there (which I think will be a big boon for open models) is that techniques like fine-tuning makes a HUGE difference and can bridge the quality gap between Llama-2 and GPT-4 for many (but not all) problems.
[1] https://www.anyscale.com/blog/fine-tuning-llama-2-a-comprehe...
We've also seen a strong desire from businesses to manage models and compute on their own machines or in their own cloud accounts. This is often part of a hybrid strategy of using API products like OpenAI for rapid prototyping.
The majority of (though not all) businesses we've seen tend to be quite comfortable using hosted API products for rapid prototyping and for proving out an initial version of their AI functionality. But in many cases, they want to complement that with the ability to manage models and compute themselves. The motivation here is often to reduce costs by using smaller / faster / cheaper fine-tuned open models.
When we started Anyscale, customer demand led us to run training & inference workloads in our customers' cloud accounts. That way your data and code stays inside of your own cloud account.
Now with all the progress in open models and the desire to rapidly prototype, we're complementing that with a fully-managed inference API where you can do inference with the Llama-2 models [1] (like the OpenAI API but for open models).
GPT-4 wins by a lot out of the box. However, surprisingly, fine-tuning makes a huge difference and allows the 7B Llama-2 model to outperform GPT-4 on some (but not all) problems.
This is really great news for open models as many applications will benefit from smaller, faster, and cheaper fine-tuned models rather than a single large, slow, general-purpose model (Llama-2-7B is something like 2% of the size of GPT-4).
GPT-4 continues to outperform even the fine-tuned 70B model on grade-school math question answering, likely due to the data Llama-2 was trained on (more data for fine-tuning helps here).
https://www.anyscale.com/blog/fine-tuning-llama-2-a-comprehe...
I'm one of the Ray developers. It is true that Ray focuses a lot on ML applications (in particular, the main libraries built on top of Ray are for workloads like training, serving, and batch processing / inference). That said, one of our long-term goals with Ray is to be a great general-purpose way to build distributed applications, so I hope it is working out for you :)
If you're curious about how Ray is used for LLMs, here are some interesting examples of LLM projects using Ray!
- Alpa does training and serving with 175B parameter models https://github.com/alpa-projects/alpa
- GPT-J https://github.com/kingoflolz/mesh-transformer-jax
- Another HN thread on training LLMs with Ray (on TPUs in this case) https://news.ycombinator.com/item?id=27731168
- OpenAI fireside chat on the evolution of their infrastructure and usage of Ray for training https://www.youtube.com/watch?v=CqiL5QQnN64
- Cohere on their architecture for training LLMs https://www.youtube.com/watch?v=For8yLkZP5w&t=3s
Some other thoughts
1. There is a lot more we want to do to make Ray better for working with large language models and for making training, serving, and batch inference work well out of the box.
2. The original post is about training, but we actually see even more interest in fine-tuning and serving with LLMs, in part because there are good pre-trained models.
3. For LLMs, we see a lot of interest in Ray + Jax or Ray + TPUs relative to what we see in other use cases.
The term serverless is a bit overloaded. Here's what we want. (1) Ray users should be able to focus only on their application logic and not have to worry about configuring clusters at least in the common case (this is not saying that they can't configure clusters if they want to). (2) In addition, we want Ray applications to be portable and to run on different size clusters with different instance types on different cloud providers or k8s or your laptop or anywhere. The application should be decoupled from the cluster configuration. (3) Great support for autoscaling Ray clusters. When a Ray application needs more resources of a certain type (CPUs, GPUs, memory, etc), those resources should be added to the cluster. When they are no longer needed, the cluster should scale down.
This is quite different from FaaS, though seems in line with the spirit of serverless. And these "serverless" properties are not something we think of as separate from Ray, but rather as part of just making Ray work better.
*Emphasis on the library ecosystem*
A lot of our focus is on building an ecosystem of libraries on top of Ray (much, but not all, of the focus is on machine learning libraries). Some of these libraries are built natively on top of Ray such as Ray Tune for scaling hyperparameter search (http://tune.io), RLlib for scaling reinforcement learning (http://rllib.io), Ray Serve for scaling model serving (http://rayserve.org/), and RaySGD for scaling training (https://docs.ray.io/en/master/raysgd/raysgd.html).
Some of the libraries are popular libraries on their own, which now integrate with Ray such as Horovod (https://eng.uber.com/horovod-ray/), XGBoost (https://xgboost.readthedocs.io/en/latest/tutorials/ray.html), and Dask for dataframes (https://docs.ray.io/en/master/dask-on-ray.html). While Dask itself has similarities to Ray (especially the task part of the Ray API), Dask also has libraries for scaling dataframes and arrays, which can be used as part of the Ray ecosystem (more details at https://www.anyscale.com/blog/analyzing-memory-management-an...).
Many Ray users start using Ray for one of the libraries (e.g., to scale training or hyperparameter search) as opposed to just for the core system.
*Emphasis on serverless*
Our goal with Ray is to make distributed computing as easy as possible. To do that, we think the serverless direction, which allows people to just focus on their code and not on infrastructure, is very important. Here, I don't mean serverless purely in the sense of functions as a service, but something that would allow people to run a wide variety of applications (training, data processing, inference, etc) elastically in the cloud without configuring or thinking about infrastructure. There's a lot of ongoing work here (e.g., to improve autoscaling up and down with heterogeneous resource types). More details on the topic https://www.anyscale.com/blog/the-ideal-foundation-for-a-gen....
If you're interested in this kind of stuff, consider joining us at Anyscale https://jobs.lever.co/anyscale.