Minimal Build for On-Device Training

[baijumeswani]

🛠️ Changes in this pull request:

This pull request introduces two significant changes to the project:

• Changing on device training checkpoint format: The current implementation stores the on device training checkpoint as a sequence of tensors in multiple files inside a checkpoint folder, which can be inefficient in terms of storage and performance. In this PR, I have modified the checkpoint format to utilize the flatbuffer table to save the checkpoint to a single file, providing a more compact and efficient representation. The changes around this are twofold:

  • Add the checkpoint flatbuffer schema that will generate the necessary checkpoint source files.
  • Update the checkpoint saving and loading functionality to use the new format.

• Adding support for onnxruntime minimal build: To support scenarios where binary size is a constraint, I made changes to ensure that the training build can work well with the minimal build.

🔍 Open Issues:

• In order to extract the optimizer type, the existing implementation re-loaded the onnx optimizer model and parsed it. This is no longer possible, since the model format can either be onnx or ort. One idea is to do the same for ort format optimizer model. This needs some investigation.
• Changes to the offline tooling to generate ort format training artifacts.
• End-to-end training example showcasing the use of the minimal training build.
• Add support for export model for inferencing in a minimal build.

[pengwa]

Not very familar with flatbuffer format, just have few general comments.

[pengwa]

Minors.

[pengwa]

I am suggesting to rename "requires_grad" to "requires_grad_params" to better represent itself.

Since it is part of schema, if we want to do that, maybe we should do it earlier (instead of bumping the versions next time). Any thought?

Besides that, LGTM.

[edgchen1]

a few comments, looks good overall

[baijumeswani]

Thank you for the valuable feedback @edgchen1 @pengwa @skottmckay @askhade 😄