3 ms·

How to Tame Your Jax: Navigating TPU VMs with TFRC in 2021

3 points by sillysaurusx 5y ago

dekhn 5y ago
From the bug log describing the TFRC team: """For a Google project, this was somewhere between "space aliens" and "narnia" on the Scale of Surprising Things.""" What you need to understand is that the people who are doing Jax and TFRC are a group of excited brilliant engineers on the cusp of junior/senior status who have a complete understanding of the current state of ML hardware, and have built a collection of powerful tools that are far simpler than TensorFlow. It's still in the rapid expansion phase, new ideas can be tried out quickly, and engaging with the developer community enables a virtuous cycle that improves the system rapidly. TensorFlow on the other hand has turned into a complex software behemoth where nearly everything requires extraordinarily long software development efforts to make minor changes, and any changes in the core impact thousands of users. What I learned from interacting with TFRC and the Jax community is that the real secret value is in XLA and the runtimes that wrap around it. Once you've mastered those, you're well on the way to using large resources like a TPU with a CPU (I think that's just the host CPU and RAM, BTW) and harnessing all the compute elements simultaneously to get peak performance. I expect that as more people see it's easier to program a simple layer on top of TFRC and XLA, tensorflow will not be as popular.
- sillysaurusx 5y ago
  (Submitter here.) Thank you for writing up such an accurate analysis. You’re correct on all counts, as far as I can tell. Unfortunately, I ran into a very strange situation with HN, and due to a Series of Unfortunate Events, your insightful comment here won’t ever be seen. However! https://news.ycombinator.com/item?id=27602110 https://news.ycombinator.com/item?id=27602110 has a chance of being seen. (I resubmitted with the correct url, or at least the most approximately-correct url that can be submitted to HN.) In fact, I’m pretty sure the new submission won’t get any traction either. So how about this: you seem like a capable XLA TPU hacker. I’d like to pretend I submitted https://github.com/google/jax/issues/2108#issuecomment-866238579 https://github.com/google/jax/issues/2108#issuecomment-86623... solely for an audience of you. What do you think of it? Critical feedback welcome — it’s important to me that this suffices as a decent on boarding experience for people interested in Jax + TPUs, but no idea where to start. > I expect that as more people see it's easier to program a simple layer on top of TFRC and XLA, tensorflow will not be as popular. I absolutely agree. My hypothesis is that Jax is such a “simple layer,” which is why I wanted to write something to persuade others it is. So I am especially interested in whether my post persuades you.
  - dekhn 5y ago
    I really liked the comment (as did Skye). Here's what's great about the comment: 0. contains the recognition that switching TPUs over to 1VM (the system you're using now) makes it clear the TPU is just a PCI acceleration card attached to a fat server. The server is fat, by the way, so that the TPU remains fully utilized. 1. contains actual simple code showing how to "do it". A fragment of working jax code is worth a million dollars compared to an idea 2. Giving TFRC (or really, the extended team of folks in the TFRC, XLA and Jax ecosystems) some respect. They work hard and don't always get the attention they deserve for bringing TF kickng and screaming into the modern world. 3. Encourage people to use the free resources and ask for more. People don't realize just how much useful resource google is giving away (just to build a userbase for the product) in TFRC. 4. It's written in a narrative style, not a bug report or boring tutorial. 5. The injected humor improves the comment although I suspect some people do not enjoy it as much as I do. 6. I think more people are going to realize that having each host running the Jax runtime with full python control (Ray mode) is really a killer app. I did some work with this internally at Google and it's a breath of fresh air compared to running tensorflow_std_server and sending it RPCs. I'm not an XLA TPU hacker- I did that at Google, but resigned recently and focus on entirely different things now. Also, I'm less enthused about Jax proper (even though I think it's pretty cool, I find a lot of it very confusing, having come from an explicit message passing (MPI) world. Personally I would target XLA high level language directly, or write a Builder on top of that, rather than try to make a system that traces python to generate forward and backward computations (although that's a pretty cool idea). And finally its nice to see some of the recognition of these things getting outside Google. I keep forgetting that when you're on the inside you're seeing people playing with all the toys that haven't been released yet and some things (like controlling a v3-4096 via a Python notebook interface and Ray) that feel totally natural aren't yet (once, the founder of hadoop said he appreciated that GOogle was publishing papers that appeared to go back in time 10 years) Some of the most interesting approaches which spread networks across the TPU mesh and the datacenter network in combination with combined training and inference and sparse activations for multimodal learning are still proprietary but I expect those to be OSSed on top of Jax at some point.
sillysaurusx 5y ago
Oh. I see what happened here. No wonder nobody noticed this… I submitted it and it fell flat like a pancake, because either I screwed up the URL or HN has a rule to automatically remove hashtag anchors. The goal is to submit to HN one of the replies in that thread, not the main page: https://github.com/google/jax/issues/2108#issuecomment-866238579 https://github.com/google/jax/issues/2108#issuecomment-86623... Without that anchor, I see no possible way to submit it to HN, because the anchor is vital. I swear I submitted it correctly, but maybe I derped. Hmm. I should try to submit it again, then if the anchor vanishes, click edit and add the anchor… I bet that would bypass the rule. But, shall I now, or wait until morning? Pfft, I’ll do it now. I always believed the “when should I submit to HN?” analysis posts were dumb. Maybe it matters, but who cares. EDIT: I tried again here, which also failed. https://news.ycombinator.com/item?id=27602092 https://news.ycombinator.com/item?id=27602092

3 points by sillysaurusx 5y ago

From the bug log describing the TFRC team: """For a Google project, this was somewhere between "space aliens" and "narnia" on the Scale of Surprising Things.""" What you need to understand is that the people who are doing Jax and TFRC are a group of excited brilliant engineers on the cusp of junior/senior status who have a complete understanding of the current state of ML hardware, and have built a collection of powerful tools that are far simpler than TensorFlow. It's still in the rapid expansion phase, new ideas can be tried out quickly, and engaging with the developer community enables a virtuous cycle that improves the system rapidly. TensorFlow on the other hand has turned into a complex software behemoth where nearly everything requires extraordinarily long software development efforts to make minor changes, and any changes in the core impact thousands of users. What I learned from interacting with TFRC and the Jax community is that the real secret value is in XLA and the runtimes that wrap around it. Once you've mastered those, you're well on the way to using large resources like a TPU with a CPU (I think that's just the host CPU and RAM, BTW) and harnessing all the compute elements simultaneously to get peak performance. I expect that as more people see it's easier to program a simple layer on top of TFRC and XLA, tensorflow will not be as popular.

sillysaurusx 5y ago
(Submitter here.) Thank you for writing up such an accurate analysis. You’re correct on all counts, as far as I can tell. Unfortunately, I ran into a very strange situation with HN, and due to a Series of Unfortunate Events, your insightful comment here won’t ever be seen. However! https://news.ycombinator.com/item?id=27602110 https://news.ycombinator.com/item?id=27602110 has a chance of being seen. (I resubmitted with the correct url, or at least the most approximately-correct url that can be submitted to HN.) In fact, I’m pretty sure the new submission won’t get any traction either. So how about this: you seem like a capable XLA TPU hacker. I’d like to pretend I submitted https://github.com/google/jax/issues/2108#issuecomment-866238579 https://github.com/google/jax/issues/2108#issuecomment-86623... solely for an audience of you. What do you think of it? Critical feedback welcome — it’s important to me that this suffices as a decent on boarding experience for people interested in Jax + TPUs, but no idea where to start. > I expect that as more people see it's easier to program a simple layer on top of TFRC and XLA, tensorflow will not be as popular. I absolutely agree. My hypothesis is that Jax is such a “simple layer,” which is why I wanted to write something to persuade others it is. So I am especially interested in whether my post persuades you.
- dekhn 5y ago
  I really liked the comment (as did Skye). Here's what's great about the comment: 0. contains the recognition that switching TPUs over to 1VM (the system you're using now) makes it clear the TPU is just a PCI acceleration card attached to a fat server. The server is fat, by the way, so that the TPU remains fully utilized. 1. contains actual simple code showing how to "do it". A fragment of working jax code is worth a million dollars compared to an idea 2. Giving TFRC (or really, the extended team of folks in the TFRC, XLA and Jax ecosystems) some respect. They work hard and don't always get the attention they deserve for bringing TF kickng and screaming into the modern world. 3. Encourage people to use the free resources and ask for more. People don't realize just how much useful resource google is giving away (just to build a userbase for the product) in TFRC. 4. It's written in a narrative style, not a bug report or boring tutorial. 5. The injected humor improves the comment although I suspect some people do not enjoy it as much as I do. 6. I think more people are going to realize that having each host running the Jax runtime with full python control (Ray mode) is really a killer app. I did some work with this internally at Google and it's a breath of fresh air compared to running tensorflow_std_server and sending it RPCs. I'm not an XLA TPU hacker- I did that at Google, but resigned recently and focus on entirely different things now. Also, I'm less enthused about Jax proper (even though I think it's pretty cool, I find a lot of it very confusing, having come from an explicit message passing (MPI) world. Personally I would target XLA high level language directly, or write a Builder on top of that, rather than try to make a system that traces python to generate forward and backward computations (although that's a pretty cool idea). And finally its nice to see some of the recognition of these things getting outside Google. I keep forgetting that when you're on the inside you're seeing people playing with all the toys that haven't been released yet and some things (like controlling a v3-4096 via a Python notebook interface and Ray) that feel totally natural aren't yet (once, the founder of hadoop said he appreciated that GOogle was publishing papers that appeared to go back in time 10 years) Some of the most interesting approaches which spread networks across the TPU mesh and the datacenter network in combination with combined training and inference and sparse activations for multimodal learning are still proprietary but I expect those to be OSSed on top of Jax at some point.

sillysaurusx 5y ago

Oh. I see what happened here. No wonder nobody noticed this… I submitted it and it fell flat like a pancake, because either I screwed up the URL or HN has a rule to automatically remove hashtag anchors. The goal is to submit to HN one of the replies in that thread, not the main page: https://github.com/google/jax/issues/2108#issuecomment-866238579 https://github.com/google/jax/issues/2108#issuecomment-86623... Without that anchor, I see no possible way to submit it to HN, because the anchor is vital. I swear I submitted it correctly, but maybe I derped. Hmm. I should try to submit it again, then if the anchor vanishes, click edit and add the anchor… I bet that would bypass the rule. But, shall I now, or wait until morning? Pfft, I’ll do it now. I always believed the “when should I submit to HN?” analysis posts were dumb. Maybe it matters, but who cares. EDIT: I tried again here, which also failed. https://news.ycombinator.com/item?id=27602092 https://news.ycombinator.com/item?id=27602092