Support >2GB of Tensor data in training checkpoint #20077
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Update training checkpoint code to use external file if data > ~2GB. I don't see a way for the flatbuffers 64-bit offsets to be used as they don't support storing 'table' types with 64-bit offsets (and our Tensor is a 'table' type not a simple struct). https://github.com/google/flatbuffers/blob/0cfb7eb80b05c058e19e50fb575263908e601469/tests/64bit/test_64bit.fbs#L38-L39 Allowing a Tensor to have its raw_data in an external file should hopefully work with the least friction. As it's an extra field it's backwards compatible. Code is quickly hacked this afternoon and is mainly intended as a starting point/rough example of how things _might_ be. 100% untested. Please feel free to suggest alternative approaches. Side note: the diffs in the generated *.fbs.h files are unexpectedly large. Maybe they weren't re-generated when the new flatbuffers version was checked in. I updated by running `python .\compile_schema.py -f <build output dir>\_deps\flatbuffers-build\Debug\flatc.exe` from onnxruntime\core\flatbuffers\schema.
edgchen1
reviewed
Mar 26, 2024
Add unit tests for external write/read in core code. THESE DO NOT VALIDATE THE OUTPUT YET. - unit tests need to do more validation of data read
… + started writing unit test
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
carzh
changed the title
edgchen1
reviewed
Apr 11, 2024
edgchen1
reviewed
Apr 11, 2024
skottmckay
commented
Apr 12, 2024
onnxruntime/core/flatbuffers/schema/ort_training_checkpoint.fbs
Outdated
Show resolved
Hide resolved
skottmckay
added a commit
that referenced
this pull request
Apr 17, 2024
### Description <!-- Describe your changes. --> Add platform aware helper to fetch errno message string. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> For usage in #20077 --------- Co-authored-by: Edward Chen <[email protected]>
… state, added checkpoint unit test
skottmckay
commented
Apr 18, 2024
onnxruntime/core/flatbuffers/schema/ort_training_checkpoint.fbs
Outdated
Show resolved
Hide resolved
…ernal data + moved lambdas to helpers
Fix build errors. Fix unit tests.
skottmckay
commented
Apr 19, 2024
Exclude tests that write ORT format data from minimal build. Minor cleanups.
…r consistency with the other generated filenames. Add generated file to lint exclusions.
edgchen1
reviewed
Apr 20, 2024
orttraining/orttraining/test/training_api/core/checkpoint_test.cc
Outdated
Show resolved
Hide resolved
There was a problem hiding this comment.
It would be nice to have a python generate artifact test as well that can create the checkpoint with external data.
Also having the python test to load the checkpoint might be nice to have.
I think if this PR is getting too big, we can get it in follow-up PR.
Let's do this in a separate PR next week to ensure these changes are in the release. If the additional testing reveals any issues we can cherry pick fixes for them. |
carzh
previously approved these changes
Apr 20, 2024
There was a problem hiding this comment.
I agree with follow-up PR for the Python unit tests.
Notes for future work:
- Add unit tests for the Python generate_artifacts and checkpoint API's for models with external data
- Look into changing spans of uint8_t to spans of std::byte for the ExternalDataReader and ExternalDataWriter
- Make the optional / nullable references to ExternalDataReader and ExternalDataWriter consistent among function flows
- Change checkpoint test expected parameter names into a
constexpr std::array<std::string_view> - Add the new generated flatbuffer schema files to the exclude patterns in the lintrunner config file + fix the lines in checkpoint.cc and graph_flatbuffers_utils.cc that are > 120 characters
baijumeswani
approved these changes
Apr 22, 2024
edgchen1
approved these changes
Apr 22, 2024
yuslepukhin
approved these changes
Apr 22, 2024
yuslepukhin
reviewed
Apr 22, 2024
yuslepukhin
reviewed
Apr 22, 2024
mszhanyi
pushed a commit
that referenced
this pull request
Apr 23, 2024
This reverts commit 9372e9a.
TedThemistokleous
pushed a commit
to TedThemistokleous/onnxruntime
that referenced
this pull request
May 7, 2024
### Description <!-- Describe your changes. --> Add platform aware helper to fetch errno message string. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> For usage in microsoft#20077 --------- Co-authored-by: Edward Chen <[email protected]>
TedThemistokleous
pushed a commit
to TedThemistokleous/onnxruntime
that referenced
this pull request
May 7, 2024
### Description <!-- Describe your changes. --> Add ability to store initializer data in an external file. Update training checkpoint code to use external file if data > ~2GB. I don't see a way for the flatbuffers 64-bit offsets to be used, as they don't support storing 'table' types with 64-bit offsets (and our Tensor is a 'table' type not a simple struct). https://github.com/google/flatbuffers/blob/0cfb7eb80b05c058e19e50fb575263908e601469/tests/64bit/test_64bit.fbs#L38-L39 Allowing a Tensor to have its raw_data in an external file should hopefully work with the least friction. As it's an extra field it's backwards compatible. Please feel free to suggest alternative approaches. Side note: the diffs in the generated *.fbs.h files are unexpectedly large. Maybe they weren't re-generated when the new flatbuffers version was checked in. I updated by running: `python .\compile_schema.py -f <build output dir>\_deps\flatbuffers-build\Debug\flatc.exe` from onnxruntime\core\flatbuffers\schema which I thought was the correct way but maybe that's out of date. I think you can ignore all the diffs in the generated files and just worry about the changes to the .fbs files in onnxruntime/core/flatbuffers/schema. Basically start at the bottom of the files changed and work up as all the 'real' diffs are there. ### Motivation and Context <!-- - Why is this change required? What problem does it solve? - If it fixes an open issue, please link to the issue here. --> --------- Co-authored-by: carzh <[email protected]>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
Add ability to store initializer data in an external file.
Update training checkpoint code to use external file if data > ~2GB.
I don't see a way for the flatbuffers 64-bit offsets to be used, as they don't support storing 'table' types with 64-bit offsets (and our Tensor is a 'table' type not a simple struct).
https://github.com/google/flatbuffers/blob/0cfb7eb80b05c058e19e50fb575263908e601469/tests/64bit/test_64bit.fbs#L38-L39
Allowing a Tensor to have its raw_data in an external file should hopefully work with the least friction. As it's an extra field it's backwards compatible.
Please feel free to suggest alternative approaches.
Side note: the diffs in the generated *.fbs.h files are unexpectedly large. Maybe they weren't re-generated when the new flatbuffers version was checked in. I updated by running:
python .\compile_schema.py -f <build output dir>\_deps\flatbuffers-build\Debug\flatc.exefrom onnxruntime\core\flatbuffers\schema which I thought was the correct way but maybe that's out of date.
I think you can ignore all the diffs in the generated files and just worry about the changes to the .fbs files in onnxruntime/core/flatbuffers/schema. Basically start at the bottom of the files changed and work up as all the 'real' diffs are there.
Motivation and Context