Skip to content

Error Handling: refactor ComputationClient::TransferFromDevice to propagate status.#9429

Merged
ysiraichi merged 7 commits intomasterfrom
ysiraichi/status-for-oom-errors
Jul 24, 2025
Merged

Error Handling: refactor ComputationClient::TransferFromDevice to propagate status.#9429
ysiraichi merged 7 commits intomasterfrom
ysiraichi/status-for-oom-errors

Conversation

@ysiraichi
Copy link
Copy Markdown
Collaborator

@ysiraichi ysiraichi commented Jul 1, 2025 •
edited
Loading

This PR makes 2 main changes in order to standardize and improve error handling:

  • ComputationClient::TransferFromDevice returns a StatusOr<T> instance
  • Wrap xla::PjRtLoadedExecutable::Execute(Sharded) with GetValueOrThrow

These changes mainly affect the errors whenever an OOM occurs. The second one targets eager mode. As an example, the following is the result of running the file below (without eager mode):

device = torch_xla.device()
a = torch.rand(1024, 1024, 1024, 1024, 1024, device=device)
b = a.sum()
print(b)

Before this PR:

F0000 00:00:1751368998.014824    2835 pjrt_computation_client.cpp:525] Non-OK-status: status
Status: INTERNAL: Error preparing computation: Out of memory allocating 4503599761588224 bytes.
*** Begin stack trace ***
        tsl::CurrentStackTrace[abi:cxx11]()
        torch_xla::runtime::PjRtComputationClient::TransferFromDevice(absl::lts_20230802::Span<std::shared_ptr<torch_xla::runtime::ComputationClient::Data> const>)
        ...
        _start
*** End stack trace ***
Failed to await future from buffer to literal inTransferFromDevice
*** Check failure stack trace: ***
    @     0x7ddc923438f9  absl::lts_20230802::log_internal::LogMessage::PrepareToDie()
    ...
Aborted (core dumped)

After this PR:

Traceback (most recent call last):
  File "examples/mem.py", line 11, in <module>
    print(b)
  File "torch/_tensor.py", line 590, in __repr__
    return torch._tensor_str._str(self, tensor_contents=tensor_contents)
  File "torch/_tensor_str.py", line 726, in _str
    return _str_intern(self, tensor_contents=tensor_contents)
  File "torch/_tensor_str.py", line 462, in _str_intern
    self = self.to("cpu")
RuntimeError: Error preparing computation: Out of memory allocating 4503599627370496 bytes.

(with XLA_SHOW_CPP_ERROR_CONTEXT=1)

RuntimeError: Error preparing computation: Out of memory allocating 4503599627370496 bytes. (at torch_xla/csrc/runtime/pjrt_computation_client.cpp:524)

@ysiraichi

This comment was marked as outdated.

Sign in to view
@ysiraichi ysiraichi force-pushed the ysiraichi/status-for-oom-errors branch 3 times, most recently from cef8c1e to 247fdf5 Compare July 1, 2025 16:29
@ysiraichi ysiraichi mentioned this pull request Jul 1, 2025
Merged
@ysiraichi ysiraichi force-pushed the ysiraichi/status-for-oom-errors branch 2 times, most recently from b390a61 to 821c384 Compare July 1, 2025 18:15
@ysiraichi ysiraichi changed the base branch from ysiraichi/status-qol-functions to master July 1, 2025 18:16
@ysiraichi ysiraichi marked this pull request as ready for review July 1, 2025 18:20
zhanyong-wan
zhanyong-wan requested changes Jul 1, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@zhanyong-wan zhanyong-wan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice!

Can you add a python test to ensure that OOM does result in a python exception with the expected error message as opposed to crashing?

@ysiraichi

This comment was marked as outdated.

Sign in to view
@ysiraichi ysiraichi force-pushed the ysiraichi/status-for-oom-errors branch from 821c384 to 08c5ecd Compare July 3, 2025 14:41
@ysiraichi ysiraichi marked this pull request as draft July 3, 2025 14:41
This was referenced Jul 3, 2025
Merged
Merged
@ysiraichi ysiraichi force-pushed the ysiraichi/status-for-oom-errors branch from 048459c to a370c5d Compare July 17, 2025 23:13
@ysiraichi ysiraichi marked this pull request as ready for review July 19, 2025 18:06
@ysiraichi ysiraichi requested a review from zhanyong-wan July 19, 2025 18:06
zhanyong-wan
zhanyong-wan approved these changes Jul 21, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@zhanyong-wan zhanyong-wan left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great!

Comment thread test/test_operations.py Outdated
@ysiraichi ysiraichi force-pushed the ysiraichi/status-for-oom-errors branch from 46401cf to 0cc5400 Compare July 21, 2025 18:28
ghpvnist
ghpvnist approved these changes Jul 21, 2025
View reviewed changes
Copy link
Copy Markdown
Collaborator

@ghpvnist ghpvnist left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is great!

@ysiraichi ysiraichi mentioned this pull request Jul 23, 2025
Open
@ysiraichi ysiraichi force-pushed the ysiraichi/status-for-oom-errors branch from 34abde6 to 103cd0f Compare July 24, 2025 12:40
@ysiraichi
Copy link
Copy Markdown
Collaborator Author

ysiraichi commented Jul 24, 2025

I have opened issue #9502 addressing TPU hanging on OOM (was causing one of the TPU CIs to fail), which I believe should be addressed separately.

@ysiraichi ysiraichi merged commit 299a16b into master Jul 24, 2025
23 of 24 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ghpvnist ghpvnist ghpvnist approved these changes

@zhanyong-wan zhanyong-wan zhanyong-wan approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@ysiraichi @ghpvnist @zhanyong-wan