{ "id": 142748, "title": "Best Practices / Troubleshooting TPU Memory Issues", "url": "/competitions/jigsaw-multilingual-toxic-comment-classification/discussion/142748", "author_name": "Jerry Qu", "post_date": "2020-04-12T04:08:01.192000", "votes": 16, "comment_count": 7, "views": 0, "content": "

This competition is my first time using TPUs, and I assume it may be the same for others as well.

\n\n

I've been running into TPU memory issues, making this GCP resource insanely useful: https://cloud.google.com/tpu/docs/troubleshooting#memory-usage

\n\n

Some key insights I've used:

\n\n\n\n

Hope this helps! (& let me know if anything here is incorrect)

", "messages": [ { "id": 804864, "postDate": "2020-04-12T04:08:01.193Z", "content": "

This competition is my first time using TPUs, and I assume it may be the same for others as well.

\n\n

I've been running into TPU memory issues, making this GCP resource insanely useful: https://cloud.google.com/tpu/docs/troubleshooting#memory-usage

\n\n

Some key insights I've used:

\n\n\n\n

Hope this helps! (& let me know if anything here is incorrect)

", "rawMarkdown": "This competition is my first time using TPUs, and I assume it may be the same for others as well.\n\nI've been running into TPU memory issues, making this GCP resource insanely useful: https://cloud.google.com/tpu/docs/troubleshooting#memory-usage\n\nSome key insights I've used:\n\n- Adam uses 8 extra bytes per weight (4GB for me)\n- Adafactor uses no extra memory\n\n- tf2.1 supports Mixed Precision (float16 instead of float32)\nhttps://www.tensorflow.org/guide/keras/mixed_precision (memory & speed improvements)\n\n- The total batch size should be a multiple of 64 (8 per TPU core), and feature dimensions should be a multiple of 128\n\nHope this helps! (& let me know if anything here is incorrect)", "votes": 15 }, { "id": 805436, "postDate": "2020-04-12T17:18:38Z", "content": "

Thank you very much @jerryqu , it is really infomative.

", "rawMarkdown": "Thank you very much @jerryqu , it is really infomative. ", "votes": 1 }, { "id": 805317, "postDate": "2020-04-12T15:10:16.350Z", "content": "

NOTE: If you use Mixed Precision, BE SURE to read the TensorFlow guide. You have to manually set certain values to float32 to insure model quality. (ie. final activation should be float32)

", "rawMarkdown": "NOTE: If you use Mixed Precision, BE SURE to read the TensorFlow guide. You have to manually set certain values to float32 to insure model quality. (ie. final activation should be float32)", "votes": 1 }, { "id": 812750, "postDate": "2020-04-19T03:55:01.783Z", "content": "

Does anyone know how to print TPU memory usage in a kernel? Thanks.

", "rawMarkdown": "Does anyone know how to print TPU memory usage in a kernel? Thanks." }, { "id": 806704, "postDate": "2020-04-14T01:12:54.447Z", "content": "

Thanks, it helps a lot!\nI have a few questions,\n1. Have you actually experienced the benefit of mixed precision on TPU as well?\n2. Why batch size should be multiple of 64, not just 8?

", "rawMarkdown": "Thanks, it helps a lot!\nI have a few questions,\n1. Have you actually experienced the benefit of mixed precision on TPU as well?\n2. Why batch size should be multiple of 64, not just 8?", "replies": [ { "id": 807324, "postDate": "2020-04-14T15:14:48.110Z", "content": "
    \n

  1. I haven't used Mixed Precision w/ TensorFlow, but I have a working version with PyTorch-XLA. Torch-XLA was running into OOM issues way more than TensorFlow for some reason.

    2. \n

  2. I believe it has to do with how Tensors are padded when fed into TPU cores. They mention on the website that 'TPU rounds up the sizes of tensors stored in memory to perform computations more efficiently', which is why you should use the dimensions they recommend. (So, if you use a batch size of 8, I think they'll just pad it up to 64, which is wasting a lot of computation)

    3. \n
\n\n

As to why they must pad the tensors, I'm not so sure. But it may have to do with how computations are actually done on TPUs, with systolic arrays: https://www.youtube.com/watch?v=JC84GCU7zqA

", "rawMarkdown": "1. I haven't used Mixed Precision w/ TensorFlow, but I have a working version with PyTorch-XLA. Torch-XLA was running into OOM issues way more than TensorFlow for some reason.\n\n2. I believe it has to do with how Tensors are padded when fed into TPU cores. They mention on the website that 'TPU rounds up the sizes of tensors stored in memory to perform computations more efficiently', which is why you should use the dimensions they recommend. (So, if you use a batch size of 8, I think they'll just pad it up to 64, which is wasting a lot of computation)\n\nAs to why they must pad the tensors, I'm not so sure. But it may have to do with how computations are actually done on TPUs, with systolic arrays: https://www.youtube.com/watch?v=JC84GCU7zqA" }, { "id": 807363, "postDate": "2020-04-14T15:50:08.610Z", "content": "

@jerryqu Oh, you're using pytorch, I got it. And thank you for the answer! Eventually stopping using Adam was most effective:)

", "rawMarkdown": "@jerryqu Oh, you're using pytorch, I got it. And thank you for the answer! Eventually stopping using Adam was most effective:)", "votes": 1 } ] }, { "id": 913372, "postDate": "2020-07-03T07:20:17.073Z", "content": "

When I used TPU Pytorch XLA version, It throws out of memory issue on 2nd epoch. I suspect the cache cannot be cleared timely. Do u know how to tackle this issue?

", "rawMarkdown": "When I used TPU Pytorch XLA version, It throws out of memory issue on 2nd epoch. I suspect the cache cannot be cleared timely. Do u know how to tackle this issue? ", "isDeleted": true } ], "comments": [ { "id": 805436, "author_name": "Mohammed Deeb", "author_url": "", "post_date": "2020-04-12T17:18:38", "content": "

Thank you very much @jerryqu , it is really infomative.

", "votes": 1, "replies": [] }, { "id": 805317, "author_name": "Jerry Qu", "author_url": "", "post_date": "2020-04-12T15:10:16.350000", "content": "

NOTE: If you use Mixed Precision, BE SURE to read the TensorFlow guide. You have to manually set certain values to float32 to insure model quality. (ie. final activation should be float32)

", "votes": 1, "replies": [] }, { "id": 812750, "author_name": "DHZM", "author_url": "", "post_date": "2020-04-19T03:55:01.783000", "content": "

Does anyone know how to print TPU memory usage in a kernel? Thanks.

", "votes": 0, "replies": [] }, { "id": 806704, "author_name": "Camaro", "author_url": "", "post_date": "2020-04-14T01:12:54.447000", "content": "

Thanks, it helps a lot!\nI have a few questions,\n1. Have you actually experienced the benefit of mixed precision on TPU as well?\n2. Why batch size should be multiple of 64, not just 8?

", "votes": 0, "replies": [ { "id": 807324, "author_name": "Jerry Qu", "author_url": "", "post_date": "2020-04-14T15:14:48.110000", "content": "
    \n

  1. I haven't used Mixed Precision w/ TensorFlow, but I have a working version with PyTorch-XLA. Torch-XLA was running into OOM issues way more than TensorFlow for some reason.

    2. \n

  2. I believe it has to do with how Tensors are padded when fed into TPU cores. They mention on the website that 'TPU rounds up the sizes of tensors stored in memory to perform computations more efficiently', which is why you should use the dimensions they recommend. (So, if you use a batch size of 8, I think they'll just pad it up to 64, which is wasting a lot of computation)

    3. \n
\n\n

As to why they must pad the tensors, I'm not so sure. But it may have to do with how computations are actually done on TPUs, with systolic arrays: https://www.youtube.com/watch?v=JC84GCU7zqA

", "votes": 0, "replies": [] }, { "id": 807363, "author_name": "Camaro", "author_url": "", "post_date": "2020-04-14T15:50:08.610000", "content": "

@jerryqu Oh, you're using pytorch, I got it. And thank you for the answer! Eventually stopping using Adam was most effective:)

", "votes": 1, "replies": [] } ] }, { "id": 913372, "author_name": "", "author_url": "", "post_date": "2020-07-03T07:20:17.073000", "content": "

When I used TPU Pytorch XLA version, It throws out of memory issue on 2nd epoch. I suspect the cache cannot be cleared timely. Do u know how to tackle this issue?

", "votes": 0, "replies": [] } ], "raw_markdown_by_id": { "804864": "This competition is my first time using TPUs, and I assume it may be the same for others as well.\n\nI've been running into TPU memory issues, making this GCP resource insanely useful: https://cloud.google.com/tpu/docs/troubleshooting#memory-usage\n\nSome key insights I've used:\n\n- Adam uses 8 extra bytes per weight (4GB for me)\n- Adafactor uses no extra memory\n\n- tf2.1 supports Mixed Precision (float16 instead of float32)\nhttps://www.tensorflow.org/guide/keras/mixed_precision (memory & speed improvements)\n\n- The total batch size should be a multiple of 64 (8 per TPU core), and feature dimensions should be a multiple of 128\n\nHope this helps! (& let me know if anything here is incorrect)", "805436": "Thank you very much @jerryqu , it is really infomative. ", "805317": "NOTE: If you use Mixed Precision, BE SURE to read the TensorFlow guide. You have to manually set certain values to float32 to insure model quality. (ie. final activation should be float32)", "812750": "Does anyone know how to print TPU memory usage in a kernel? Thanks.", "806704": "Thanks, it helps a lot!\nI have a few questions,\n1. Have you actually experienced the benefit of mixed precision on TPU as well?\n2. Why batch size should be multiple of 64, not just 8?", "913372": "When I used TPU Pytorch XLA version, It throws out of memory issue on 2nd epoch. I suspect the cache cannot be cleared timely. Do u know how to tackle this issue? " } }