7 ms·

Google Cloud TPU Multislice Training

109 points by infixed 3y ago

jeffbee 3y ago
Something that doesn't seem worth bragging about is that the startup time increases linearly with the cluster size. Wouldn't you want it to be constant? What's the issue there?
- smarterclayton 3y ago
  Disclaimer: work associated with this team, didn't write or review the blog post Article stated that it was throughput scheduling the pods on the clusters (from unrelated benchmarks that's usually ~300 pods/sec throughput for kube scheduler today) and then doing XLA compilation at pod launch, rather than amortizing once for all jobs. Optimizing throughput of kube scheduler is a good general opportunity and something I believe we would like to see. I believe AOT compilation just not a critical optimization for the test, we would recommend it when running large and long training jobs to AOT compile to keep pod start latency low for hardware failures and job restarts (from checkpoints). > The start times we observed were impressive, but we believe we can improve these even further. We are working on areas such as optimizing scheduling in GKE to increase throughput and enabling ahead-of-time compilation in MaxText to avoid just-in-time compilations on the full cluster.
  - jeffbee 3y ago
    Thanks for the context. I remember recently reading a paper from I think Baidu where they claimed to have a container arrival rate in the millions per second, consequently it was practical to operate their whole site in the style of lambda/cloud functions. Actually now that I am searching for that it seems Baidu has a number of papers on workload orchestration at scale specifically for learning.
    - smarterclayton 3y ago
      I will note that a trend I have observed with recent ML - as we increasingly use accelerators and models correspondingly grow in size, we are returning to a "one machine, one workload" paradigm for the biggest training and inference jobs. You might have 8k accelerators, but only 1000 machines, and if you have one container per host 300 schedules / second is fast. While at the same time as you note we have functional models for container execution that are approaching millions of dispatches for highly partitionable work, especially in data engineering and ETL.
- rwitten 3y ago
  (Contributor on the blog post, all opinions my own) Agreed with you and we definitely weren't trying to brag! This is fast compared to people's expectations in the space but slow compared to what we should be able to accomplish and will accomplish in the future.
xnx 3y ago
Full title: "Google Cloud demonstrates the world’s largest distributed training job for large language models across 50000+ TPU v5e chips" Summary from Bard: "This article is about training large language models (LLMs) on Google Cloud TPUs. It discusses the challenges of training LLMs at scale, and how Google Cloud TPU Multislice Training addresses these challenges. The article also details the results of a recent experiment in which Google trained a 128B parameter LLM on 50,944 TPU v5e chips. This experiment is the largest publicly disclosed LLM distributed training job to date."
behnamoh 3y ago
Can someone ELI5 this?
- jeffbee 3y ago
  Google has a megawatt-scale TPUv5 cluster that they can spare for stunts, and they are flexing it.
- filterfiber 3y ago
  tl;dr - google has big computer (a lot of TPUs) they just showed they could indeed make 50k TPUs do some flops. With no paper this is just a marketing press release - the only takeaway is that existing tech stacks can utilize it probably.
jrk 3y ago
As far as I can tell, the article notably never defines what "slices" are or what "multi-slice" means.
- rwitten 3y ago
  Great questions! Slices are a set of TPU chips that share a fast, private inter-chip-interconnect. Unlike the current GPU generation in clouds, the TPUs on different machines can communicate through this private network. Multislice means that we're using a hierarchical network, where there is both inter-chip-interconnect and normal data-center netowrking. More details: https://cloud.google.com/tpu/docs/multislice-introduction https://cloud.google.com/tpu/docs/multislice-introduction (P.S. - contributor on blog post, Google employee, all thoughts my own)
  - smarterclayton 3y ago
    Also, I should point out that a set of machines hosting TPUs is referred to as a "pod", which is not the same thing as a Kubernetes pod (also referenced in this doc). The term "pod" originated in early data center design and occasionally crosses over from HPC to broad use - i.e. nVidia calls the set of DGX machines a "pod" https://blogs.nvidia.com/blog/2021/03/05/what-is-a-cluster-pod/ https://blogs.nvidia.com/blog/2021/03/05/what-is-a-cluster-p.... Kubernetes chose "pod" to represent a set of co-scheduled containers, like a "pod of whales". Other systems like Mesos and Google's Borg https://storage.googleapis.com/pub-tools-public-publication-data/pdf/43438.pdf https://storage.googleapis.com/pub-tools-public-publication-... use "task" to refer to a single container but didn't have a concept for heterogenous co-scheduled tasks at the time. Somewhat ironically, it now means TPUs on GKE are confusing because we have TPUs hosts organized into "pods", and "pods" for the software using the TPUs. A Kubernetes pod using a TPU lands on a host which is part of a slice of a TPU pod.
    - jeffbee 3y ago
      As your second link mentions in section 2.4, Borg has "allocs" which are basically pods.
leumassuehtam 3y ago
Did they use this to train Gemini? Which raises the question, where is Gemini?
- lern_too_spel 3y ago
  Unlikely. One reason Google Cloud is so terrible is that nobody in Google actually uses Google Cloud. It used to be that every time I mentioned this, somebody would jump in and say, "Well actually, Google Domains runs on Google Cloud," and we'd discuss whether Google Domains was a business critical part of Google. https://support.google.com/domains/answer/13689670?hl=en https://support.google.com/domains/answer/13689670?hl=en
  - amf12 3y ago
    > Unlikely. One reason Google Cloud is so terrible is that nobody in Google actually uses Google Cloud. Well, actually, Google Cloud is just an abstraction on top of internal Google infra, so this isn't the right question. So, it depends on what you want to infer/compare.
    - lern_too_spel 3y ago
      > Well, actually, Google Cloud is just an abstraction on top of internal Google infra I didn't say otherwise. Of course Google Cloud runs on internal Google infrastructure. They wouldn't have an entirely different stack to build Google Cloud on. The problem is that Googlers don't use Google Cloud. Amazonians use AWS. https://courses.cs.washington.edu/courses/cse452/23wi/papers/yegge-platform-rant.html https://courses.cs.washington.edu/courses/cse452/23wi/papers... Microsofties use Azure. https://www.zdnet.com/article/microsoft-moves-closer-to-running-all-of-its-own-services-on-azure/ https://www.zdnet.com/article/microsoft-moves-closer-to-runn...
    - danielmarkbruce 3y ago
      It is the right question. It's the right question because Google doesn't dogfood Google Cloud like they should/could. Dogfooding a bunch of stuff at a lower level of abstraction isn't the same thing.
sashank_1509 3y ago
Ok so they claim in the article, 50000 TPU’s is equivalent to 10 exaflop floating point computations. That is equivalent to ~2,512 NVIDIA H100’s, which is like really small. Just shows the difference between TPU’s and GPU’s I guess. Inflection, a new LLM company created a 20,000 H100 cluster, I’m positive OpenAI, Tesla, Meta etc have orchestrated a job on more than 2500 H100 GPU’s.
- aschleck 3y ago
  It's worth noting that just because an H100 has a higher flops number doesn't mean your program is actually hitting that number of flops. Modern TPUs are surprisingly competitive with Nvidia on a perf/$ metric, if you're doing cloud ML they are absolutely worth a look. We have been keeping costs down by racking our own GPUs but TPUs are so cost effective that we need to do some thinking about changing our approach. I'm not certain but I think part of this is that XLA (for example) is a mountain of chip-specific optimizations between your code and the actual operations. So comparing your throughput between GPU and TPU is not just flops-to-flops.
- latchkey 3y ago
  It sounds like they partnered with CoreWeave to use their equipment and it is "only" 3500 gpus so far. https://inflection.ai/inflection-ai-announces-1-3-billion-of-funding https://inflection.ai/inflection-ai-announces-1-3-billion-of... https://inflection.ai/nvidia-coreweave-mlperf https://inflection.ai/nvidia-coreweave-mlperf
  - abatilo 3y ago
    https://youtu.be/z3hmfSVmyqg?si=eLPZ0D6ug3D6PreI https://youtu.be/z3hmfSVmyqg?si=eLPZ0D6ug3D6PreI As of 2 months ago, they had at least 7000 up and running, fwiw
    - latchkey 3y ago
      The meat starts here: https://youtu.be/z3hmfSVmyqg?feature=shared&t=3328 https://youtu.be/z3hmfSVmyqg?feature=shared&t=3328 I'm curious why it is so hard for them to deploy the compute. They seem to be fairly behind schedule.
DavidSJ 3y ago
Question for rwitten or anyone else involved in this project: I see a per-device batch size of 6 for the 16B model. With 256x199 = 50944 TPUs and a sequence length of 2048, this works out to 104M tokens per batch. This is much larger than typical for training runs of dense LMs of this size, which are usually closer to ~4M tokens per batch. Was your critical batch size really this large? In other words, did you really see a benefit as compared to a much smaller batch size (and probably many fewer TPUs)? Did you use some special learning rate schedule or optimizer to achieve this?
- p1esk 3y ago
  I’m confused - why do you multiply number of chips by sequence length? Shouldn’t the total batch size be 50k x6?
  - sillysaurusx 3y ago
    Token count. But tokens per batch is a bad metric. I learned this the hard way through experience. It turns out that what matters is step count. Increasing batch size makes the model train faster, but increasing seq length from 1024 to 2048 doesn’t make it train twice as fast. So saying 104M tokens rather than batch size 50k x 6 is misleading to yourself. (One of the most surprising aspects of learning ML was how easy it is for me to trick myself in various silly ways like that.) The mental model to make this easy to remember: progress happens in discrete quantities called steps. Training on 104M tokens per step means it’s embedding 104M tokens of knowledge every step. This isn’t the same thing as requiring an special case optimizer due to large batch sizes — sequence length adds "knowledge bandwidth", but otherwise doesn’t mess with the training dynamics. As far as large batch optimizers, there’s LARS, which google used for their MLPerf results. I imagine they stuck with that. It creates a per-layer confidence metric, so that when the massive batch size makes a massive change, it dampens the change to smooth out the effect across the network. And since it’s a multiply, the shape of the gradient (by "shape" I mean in 3D space, where the Z axis is the intensity of the gradient) remains the same, so it doesn’t harm any knowledge transfer. It’s purely a stabilization aid. Kind of weird I remember that after three years.

109 points by infixed 3y ago

Something that doesn't seem worth bragging about is that the startup time increases linearly with the cluster size. Wouldn't you want it to be constant? What's the issue there?

smarterclayton 3y ago
Disclaimer: work associated with this team, didn't write or review the blog post Article stated that it was throughput scheduling the pods on the clusters (from unrelated benchmarks that's usually ~300 pods/sec throughput for kube scheduler today) and then doing XLA compilation at pod launch, rather than amortizing once for all jobs. Optimizing throughput of kube scheduler is a good general opportunity and something I believe we would like to see. I believe AOT compilation just not a critical optimization for the test, we would recommend it when running large and long training jobs to AOT compile to keep pod start latency low for hardware failures and job restarts (from checkpoints). > The start times we observed were impressive, but we believe we can improve these even further. We are working on areas such as optimizing scheduling in GKE to increase throughput and enabling ahead-of-time compilation in MaxText to avoid just-in-time compilations on the full cluster.
- jeffbee 3y ago
  Thanks for the context. I remember recently reading a paper from I think Baidu where they claimed to have a container arrival rate in the millions per second, consequently it was practical to operate their whole site in the style of lambda/cloud functions. Actually now that I am searching for that it seems Baidu has a number of papers on workload orchestration at scale specifically for learning.
  - smarterclayton 3y ago
    I will note that a trend I have observed with recent ML - as we increasingly use accelerators and models correspondingly grow in size, we are returning to a "one machine, one workload" paradigm for the biggest training and inference jobs. You might have 8k accelerators, but only 1000 machines, and if you have one container per host 300 schedules / second is fast. While at the same time as you note we have functional models for container execution that are approaching millions of dispatches for highly partitionable work, especially in data engineering and ETL.
rwitten 3y ago
(Contributor on the blog post, all opinions my own) Agreed with you and we definitely weren't trying to brag! This is fast compared to people's expectations in the space but slow compared to what we should be able to accomplish and will accomplish in the future.

xnx 3y ago

Full title: "Google Cloud demonstrates the world’s largest distributed training job for large language models across 50000+ TPU v5e chips" Summary from Bard: "This article is about training large language models (LLMs) on Google Cloud TPUs. It discusses the challenges of training LLMs at scale, and how Google Cloud TPU Multislice Training addresses these challenges. The article also details the results of a recent experiment in which Google trained a 128B parameter LLM on 50,944 TPU v5e chips. This experiment is the largest publicly disclosed LLM distributed training job to date."

behnamoh 3y ago

Can someone ELI5 this?

jeffbee 3y ago
Google has a megawatt-scale TPUv5 cluster that they can spare for stunts, and they are flexing it.
filterfiber 3y ago
tl;dr - google has big computer (a lot of TPUs) they just showed they could indeed make 50k TPUs do some flops. With no paper this is just a marketing press release - the only takeaway is that existing tech stacks can utilize it probably.

jrk 3y ago

As far as I can tell, the article notably never defines what "slices" are or what "multi-slice" means.

rwitten 3y ago
Great questions! Slices are a set of TPU chips that share a fast, private inter-chip-interconnect. Unlike the current GPU generation in clouds, the TPUs on different machines can communicate through this private network. Multislice means that we're using a hierarchical network, where there is both inter-chip-interconnect and normal data-center netowrking. More details: https://cloud.google.com/tpu/docs/multislice-introduction https://cloud.google.com/tpu/docs/multislice-introduction (P.S. - contributor on blog post, Google employee, all thoughts my own)
- smarterclayton 3y ago
  Also, I should point out that a set of machines hosting TPUs is referred to as a "pod", which is not the same thing as a Kubernetes pod (also referenced in this doc). The term "pod" originated in early data center design and occasionally crosses over from HPC to broad use - i.e. nVidia calls the set of DGX machines a "pod" https://blogs.nvidia.com/blog/2021/03/05/what-is-a-cluster-pod/ https://blogs.nvidia.com/blog/2021/03/05/what-is-a-cluster-p.... Kubernetes chose "pod" to represent a set of co-scheduled containers, like a "pod of whales". Other systems like Mesos and Google's Borg https://storage.googleapis.com/pub-tools-public-publication-data/pdf/43438.pdf https://storage.googleapis.com/pub-tools-public-publication-... use "task" to refer to a single container but didn't have a concept for heterogenous co-scheduled tasks at the time. Somewhat ironically, it now means TPUs on GKE are confusing because we have TPUs hosts organized into "pods", and "pods" for the software using the TPUs. A Kubernetes pod using a TPU lands on a host which is part of a slice of a TPU pod.
  - jeffbee 3y ago
    As your second link mentions in section 2.4, Borg has "allocs" which are basically pods.

leumassuehtam 3y ago

Did they use this to train Gemini? Which raises the question, where is Gemini?

lern_too_spel 3y ago
Unlikely. One reason Google Cloud is so terrible is that nobody in Google actually uses Google Cloud. It used to be that every time I mentioned this, somebody would jump in and say, "Well actually, Google Domains runs on Google Cloud," and we'd discuss whether Google Domains was a business critical part of Google. https://support.google.com/domains/answer/13689670?hl=en https://support.google.com/domains/answer/13689670?hl=en
- amf12 3y ago
  > Unlikely. One reason Google Cloud is so terrible is that nobody in Google actually uses Google Cloud. Well, actually, Google Cloud is just an abstraction on top of internal Google infra, so this isn't the right question. So, it depends on what you want to infer/compare.
  - lern_too_spel 3y ago
    > Well, actually, Google Cloud is just an abstraction on top of internal Google infra I didn't say otherwise. Of course Google Cloud runs on internal Google infrastructure. They wouldn't have an entirely different stack to build Google Cloud on. The problem is that Googlers don't use Google Cloud. Amazonians use AWS. https://courses.cs.washington.edu/courses/cse452/23wi/papers/yegge-platform-rant.html https://courses.cs.washington.edu/courses/cse452/23wi/papers... Microsofties use Azure. https://www.zdnet.com/article/microsoft-moves-closer-to-running-all-of-its-own-services-on-azure/ https://www.zdnet.com/article/microsoft-moves-closer-to-runn...
  - danielmarkbruce 3y ago
    It is the right question. It's the right question because Google doesn't dogfood Google Cloud like they should/could. Dogfooding a bunch of stuff at a lower level of abstraction isn't the same thing.

sashank_1509 3y ago

Ok so they claim in the article, 50000 TPU’s is equivalent to 10 exaflop floating point computations. That is equivalent to ~2,512 NVIDIA H100’s, which is like really small. Just shows the difference between TPU’s and GPU’s I guess. Inflection, a new LLM company created a 20,000 H100 cluster, I’m positive OpenAI, Tesla, Meta etc have orchestrated a job on more than 2500 H100 GPU’s.

aschleck 3y ago
It's worth noting that just because an H100 has a higher flops number doesn't mean your program is actually hitting that number of flops. Modern TPUs are surprisingly competitive with Nvidia on a perf/$ metric, if you're doing cloud ML they are absolutely worth a look. We have been keeping costs down by racking our own GPUs but TPUs are so cost effective that we need to do some thinking about changing our approach. I'm not certain but I think part of this is that XLA (for example) is a mountain of chip-specific optimizations between your code and the actual operations. So comparing your throughput between GPU and TPU is not just flops-to-flops.
latchkey 3y ago
It sounds like they partnered with CoreWeave to use their equipment and it is "only" 3500 gpus so far. https://inflection.ai/inflection-ai-announces-1-3-billion-of-funding https://inflection.ai/inflection-ai-announces-1-3-billion-of... https://inflection.ai/nvidia-coreweave-mlperf https://inflection.ai/nvidia-coreweave-mlperf
- abatilo 3y ago
  https://youtu.be/z3hmfSVmyqg?si=eLPZ0D6ug3D6PreI https://youtu.be/z3hmfSVmyqg?si=eLPZ0D6ug3D6PreI As of 2 months ago, they had at least 7000 up and running, fwiw
  - latchkey 3y ago
    The meat starts here: https://youtu.be/z3hmfSVmyqg?feature=shared&t=3328 https://youtu.be/z3hmfSVmyqg?feature=shared&t=3328 I'm curious why it is so hard for them to deploy the compute. They seem to be fairly behind schedule.

DavidSJ 3y ago

Question for rwitten or anyone else involved in this project: I see a per-device batch size of 6 for the 16B model. With 256x199 = 50944 TPUs and a sequence length of 2048, this works out to 104M tokens per batch. This is much larger than typical for training runs of dense LMs of this size, which are usually closer to ~4M tokens per batch. Was your critical batch size really this large? In other words, did you really see a benefit as compared to a much smaller batch size (and probably many fewer TPUs)? Did you use some special learning rate schedule or optimizer to achieve this?

p1esk 3y ago
I’m confused - why do you multiply number of chips by sequence length? Shouldn’t the total batch size be 50k x6?
- sillysaurusx 3y ago
  Token count. But tokens per batch is a bad metric. I learned this the hard way through experience. It turns out that what matters is step count. Increasing batch size makes the model train faster, but increasing seq length from 1024 to 2048 doesn’t make it train twice as fast. So saying 104M tokens rather than batch size 50k x 6 is misleading to yourself. (One of the most surprising aspects of learning ML was how easy it is for me to trick myself in various silly ways like that.) The mental model to make this easy to remember: progress happens in discrete quantities called steps. Training on 104M tokens per step means it’s embedding 104M tokens of knowledge every step. This isn’t the same thing as requiring an special case optimizer due to large batch sizes — sequence length adds "knowledge bandwidth", but otherwise doesn’t mess with the training dynamics. As far as large batch optimizers, there’s LARS, which google used for their MLPerf results. I imagine they stuck with that. It creates a per-layer confidence metric, so that when the massive batch size makes a massive change, it dampens the change to smooth out the effect across the network. And since it’s a multiply, the shape of the gradient (by "shape" I mean in 3D space, where the Z axis is the intensity of the gradient) remains the same, so it doesn’t harm any knowledge transfer. It’s purely a stabilization aid. Kind of weird I remember that after three years.