GPU memory usage logging

[tmm1]

[2023-08-09 08:13:36,085] [INFO] [axolotl.scripts.train:254] [PID:741158] GPU memory baseline: 322 MB.
[2023-08-09 08:13:48,165] [INFO] [axolotl.load_model:329] [PID:741158] GPU memory after model load: 4837 MB.
[2023-08-09 08:14:22,195] [INFO] [axolotl.load_model:368] [PID:741158] GPU memory after adapters: 4917 MB.
[2023-08-09 08:14:41,233] [INFO] [axolotl.callbacks.on_step_end:95] [PID:741158] GPU memory while training: 21573 MB.

[winglian]

This seems to be a project that hasn't been touched in 6 years. Perhaps we could implement this with torch instead? https://pypi.org/project/nvidia-ml-py3/

https://pytorch.org/docs/stable/generated/torch.cuda.memory_summary.html

[utensil]

1. it's nice to mark the VRAM usages for these stages, it would be even nicer to calculate the difference between baseline and loaded, the difference between loaded and training and divide by micro_batch_size etc.
2. the unit is better to use GB for better readability, just use :.3f to make it keep 3 digits after the decimal so MB is easily recovered too
3. you can simply use torch.cuda.memory_usage() etc. instead of reimplementation;
4. could also calculate the percentage alongside the absolute GB value