[ONNX] Support Fake Tensor Mode on new Dynamo based ONNX exporter #105464
Description
🐛 Describe the bug
This task is an umbrella for all tasks related to exporting a model to ONNX using the new PyTorch Dynamo API with Fake Tensor support
Idea for alternative API design: #104144
Versions
PyTorch main branch
### Tasks
- [x] Support faketensor with huggingFace safetensors weights
- [ ] Address mix of fake and real tensor as reported by https://github.com/pytorch/pytorch/issues/105077
- [x] Investigate ExportedProgram serde
- [ ] https://github.com/pytorch/pytorch/issues/115461
- [x] Investigate how to lazy load pytorch checkpoints. pickle doesnt support it. Answer: torch.UntypedStorage.from_file(f, False, size) supports mmap
- [x] Investigate why ExportedProgram.state does not contain tensor_constants (they are separate attr)
- [x] Revisit serialization implementation of models with Fake Tensor support (ONNXProgram.save). It uses onnx helper APIs. ONNX Script does not support it for now, so keeping the onnx helper API for now
- [x] Support large model export without special hardware (e.g A100)
- [ ] https://github.com/pytorch/pytorch/issues/105751
- [ ] https://github.com/pytorch/pytorch/issues/115552
- [ ] https://github.com/pytorch/pytorch/issues/117307
Metadata
Metadata
Assignees
Labels
Type
Projects
Status