Efficient forward & backward

[dbaranchuk]

• Fault-tolerant sequential forward and backward;
• Efficient forward and backward calls that process an input batch in parallel by splitting it on small sub-batches;
• Update server deployment scripts.

---

Some measurements:

Config: Bloom 6b3 | input_ids.shape = (32, 128) | A100 | 3 servers, 10 block each

Naive (current)
84.2 sec / it | (1/3 GPUs is utilized at time)

Sequential
37.4 sec / it | (1/3 GPUs is utilized at time)

Sequential + Parallel
19.4 sec / it | (2/3 or 3/3 GPUs are utilized at time)

Reference Gradient Checkpointing (Single GPU A100)
8.8 sec / it | 100% gpu utils

Reference (Single GPU A100)
5.7 sec / it | 100% gpu utils

[justheuristic]

Please add a test case for test_remote_sequential when it actually splits tokens w.r.t. batches OR create an issue and tag me :)

[justheuristic]

[minor, reminder: i'd guess that the test_full_model will fix itself after merging #35 ]