Partial activation checkpoint/recomputation
State: closed | Pull request: yes

# What does this PR do ?

This adds two features

**1. Partial activation checkpoint**

This feature is valid only when used with pipeline-parallelism. When an integer value is provided for 'num_micro_batches_with_partial_activation_checkpoints', it sets the number of micro-batches where only a partial number of Transformer layers get checkpointed and recomputed within a window of micro-batches as shown in Figure 1. The rest of the micro-batches in the window checkpoint all Transformer layers. Compared to the baseline Figure 1.a, Figure 1.b, partial activation checkpointing enables skip checkpointing activations by using the remained GPU memory.

The size of the window is set by the maximum outstanding micro-batch backpropagations, which vary at different pipeline stages. In the below example, the first pipeline stage can do partial checkpointing at one micro-batch out of the window of four micro-batches. The second pipeline stage does one out of three and eventually, the last micro-batch does the partial checkpointing at all micro-batches. This is shown in Figure 2. This feature enables the use of activation checkpoints at a fraction of micro-batches up to the point of full GPU memory usage.

**2. Gradual activation checkpoint increase at later pipeline stage**

When an integer value is provided for 'activations_checkpoint_layers_per_pipeline', this sets the additional Transformer layers to checkpoint at later pipeline stages. For example, when using 'activations_checkpoint_layers_per_pipeline' of 3, pipeline stage 1 checkpoints 3 layers less the stage 0 and stage 2 checkpoints 6 layers less the stage 0, and so on. This is shown in Figure 3. This is possible because the later pipeline stage uses less GPU memory with fewer outstanding micro-batch backpropagations. Used with 'num_micro_batches_with_partial_activation_checkpoints', this feature removes most of the activation checkpoints at the last pipeline stage (the critical execution path).

![image](https://user-images.githubusercontent.com/15861259/186814079-9a9526ee-47dc-4f51-8bc7-be03a4acd086.png)
**Collection**: [Note which collection this PR will affect]

# Changelog 
- Add specific line by line info of high level changes in this PR.

# Usage
* You can potentially add a usage example below

```python
# Add a code snippet demonstrating how to use this 
```

# Before your PR is "Ready for review"
**Pre checks**:
- [ ] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [ ] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [ ] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.

# Additional Information
* Related to # (issue)