3 ms·

(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 str

(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.

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.