[WIP]: Autotune Chunk Size

[pramodith]

trafficstars

Summary

Our chunked loss functions currently statically set the chunk size to 1. However, this might lead to underutilized gpu memory resources. In this PR we show how the chunk size can be dynamically set, based on the additional peak memory consumed by the forward pass of the chunked loss being used.

Testing Done

On A100 80 GB SMX

Benchmarking against chunked orpo shows very marginal improvement in speed compared to a static chunk size.

• Hardware Type: <BLANK>
• [ ] run make test to ensure correctness (wip)
• [X] run make checkstyle to ensure code style
• [X] run make test-convergence to ensure convergence