onnxruntime/orttraining/orttraining/python/training/ortmodule
pengwa 8a98874e7e
Flash attention recompute (#20603)
### Flash attn recompute

1. Allow PythonOp(FlashAttn) can be recomputed correctly.
45879ff5c2
2. Use JSON to pass the selected-to-recompute subgraphs.
3c374da678

#### Better Memory Efficiency 

Customer model can run both PyTorch SPDA and Flash Attn, this PR make it
possible to let the Flash Attn path work with ORTModule layerwise
recompute. The peak drop from 45.xGB to 32.xGB if we only compare the
layers (not including other pieces, BTW there are few more optimization
targeting other pieces as well later).

#### Better Perf

Using Flash ATTN bring additionally 16% end to end time reduction, with
highly aligned loss curve.


![image](https://github.com/microsoft/onnxruntime/assets/10530022/bb63894a-f281-49bc-a8e6-ff818439be38)

#### Use JSON File to pass Recompute Plans

To overcome the limitation of max length of the strings defined in
session options.

### 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. -->
2024-05-21 13:38:19 +08:00
..
experimental Add LayerSpec Support to ORTPipelineModule (#20410) 2024-04-23 17:57:08 -07:00
graph_optimizers ATen Op Supports Int Return Type and CPU Tensor Arguments (#19773) 2024-03-06 10:11:46 +08:00
torch_cpp_extensions Improve perf for mem efficient grad mgmt (#20480) 2024-05-10 08:09:17 +08:00
__init__.py Prompt layer-wise recompute when applicable (#20126) 2024-04-10 11:50:28 +08:00
_custom_autograd_function.py Improve perf for stage3 training (#18099) 2023-12-15 13:32:19 +08:00
_custom_autograd_function_exporter.py Flash attention recompute (#20603) 2024-05-21 13:38:19 +08:00
_custom_gradient_registry.py Bump ruff linter to 0.2.1 (#19471) 2024-02-08 16:08:27 -08:00
_custom_op_symbolic_registry.py add bf16 support for few ops (#20385) 2024-04-25 11:28:34 -07:00
_execution_agent.py Flash attention recompute (#20603) 2024-05-21 13:38:19 +08:00
_fallback.py
_fallback_exceptions.py Bump ruff to 0.3.2 and black to 24 (#19878) 2024-03-13 10:00:32 -07:00
_gradient_accumulation_manager.py
_graph_execution_interface.py
_graph_execution_manager.py Flash attention recompute (#20603) 2024-05-21 13:38:19 +08:00
_graph_execution_manager_factory.py
_inference_manager.py Generalize label input sparsity check and refactor (#20636) 2024-05-10 21:55:43 +08:00
_io.py Generalize label input sparsity check and refactor (#20636) 2024-05-10 21:55:43 +08:00
_logger.py Bump ruff to 0.3.2 and black to 24 (#19878) 2024-03-13 10:00:32 -07:00
_mem_efficient_grad_mgmt.py Improve perf for mem efficient grad mgmt (#20480) 2024-05-10 08:09:17 +08:00
_onnx_models.py ORTModule memory improvement (#18924) 2024-01-16 08:57:37 +08:00
_pythonop_helper.py ORTModule memory improvement (#18924) 2024-01-16 08:57:37 +08:00
_runtime_inspector.py Flash attention recompute (#20603) 2024-05-21 13:38:19 +08:00
_torch_module_factory.py
_torch_module_interface.py
_torch_module_ort.py
_torch_module_pytorch.py
_training_manager.py Flash attention recompute (#20603) 2024-05-21 13:38:19 +08:00
_utils.py Bump ruff to 0.3.2 and black to 24 (#19878) 2024-03-13 10:00:32 -07:00
_zero_stage3_compatibility.py Improve perf for stage3 training (#18099) 2023-12-15 13:32:19 +08:00
graph_optimizer_registry.py
options.py Flash attention recompute (#20603) 2024-05-21 13:38:19 +08:00
ortmodule.py