[Epic][Improvement] Add module and function docs

[JamesKunstle]

We don't currently document the functions or modules in this library aside from our single main entrypoints. It'd be nice if users could take advantage of some of our libraries optimizations even if they're not using our training loop (e.g. multipack sampling).

We could do this manually, or we could use a docstring-based doc rendering solution like Sphinx.