[BE][doc] add memory_profiler to README

README.md

@@ -42,6 +42,7 @@ You may want to see how the model is defined or how parallelism techniques are a
 10. DDP and HSDP
 11. All options easily configured via [toml files](train_configs/)
 12. [Interoperable checkpoints](docs/checkpoint.md) which can be loaded directly into [`torchtune`](https://github.com/pytorch/torchtune) for fine-tuning
+13. Debugging tools including CPU/GPU profiling, [memory profiling](docs/memory_profiler.md), [Flight Recorder](#debugging), etc.
 
 We report our [Performance](docs/performance.md) verified on 64/128 GPUs.
 

docs/memory_profiler.md

@@ -0,0 +1,13 @@
+## Enable Memory Profiling
+
+Launch training job with the following command (or alternatively set configs in toml files)
+```
+CONFIG_FILE="./train_configs/debug_model.toml" ./run_llama_train.sh --profiling.enable_memory_snapshot --profiling.save_memory_snapshot_folder memory_snapshot
+```
+* `--profiling.enable_memory_snapshot`: to enable memory profiling
+* `--profiling.save_memory_snapshot_folder`: configures the folder which memory snapshots are dumped into (`./outputs/memory_snapshot/` by default)
+	+ In case of OOMs, the snapshots will be in `./outputs/memory_snapshot/iteration_x_exit`.
+	+ Regular snapshots (taken every `profiling.profile_freq` iterations) will be in `memory_snapshot/iteration_x`.
+
+You cab find the saved pickle files in your output folder.
+To visualize a snapshot file, you can drag and drop it to <https://pytorch.org/memory_viz>. To learn more details on memory profiling, please visit this [tutorial](https://pytorch.org/blog/understanding-gpu-memory-1/).