Add inference kv cache support for transformer TE path
State: closed | Pull request: yes
Labels: name: NLP

# What does this PR do ?

Enable E2E inference for TE FP8 path in transformer with KV cache

**Collection**: NLP

# Changelog 
- Update examples/nlp/language_modeling/megatron_gpt_eval.py for TE FP8 inference, padded inputs to a multiple of 8 to get rid of error `AssertionError: Input and weight dimensions are not compatible for FP8 execution.`.
- Set up `inference_params` correctly and pass to TE for inference. Assume first time calling with `set_inference_key_value_memory=True` and remaining with `set_inference_key_value_memory=False`.
- Need to wait Megatron-LM to add the API `parallel_state.get_amax_reduction_group()` so to fix the issue https://github.com/NVIDIA/NeMo/issues/6625

# Usage
Get an FP8 NeMo model and run:
```
python megatron_gpt_eval.py gpt_model_file=<FP8_NEMO_MODEL>
```



```python
python megatron_gpt_eval.py gpt_model_file=models/5b_fp8_tp1.nemo

***************************
[{'sentences': ['Q: How are you?\nA1) It is all too easy to get distracted, daydreaming and so forth. On top of that many people lack sufficient rest due to their',
                'Q: How big is the universe?\nA: As you’ve probably heard, the universe includes everything that we can know. It stretches from one end of space to another and'],
  'tokens': [['
', 'Q', ':', ' How', ' are', ' you', '?', '\n', 'A', '1', ')', ' It', ' is', ' all', ' too', ' easy', ' to', ' get', ' distracted', ',', ' day', 'dream', 'ing', ' and', ' so', ' forth', '.', ' On', ' top', ' of', ' that', ' many', ' people', ' lack', ' sufficient', ' rest', ' due', ' to', ' their'], ['
', 'Q', ':', ' How', ' big', ' is', ' the', ' universe', '?', '\n', 'A', ':', ' As', ' you', '�', '�', 've', ' probably', ' heard', ',', ' the', ' universe', ' includes', ' everything', ' that', ' we', ' can', ' know', '.', ' It', ' stretches', ' from', ' one', ' end', ' of', ' space', ' to', ' another', ' and']],
  'logprob': tensor([[-3.3044e+00, -1.0319e-02, -8.7674e+00, -3.8337e+00, -2.2173e+00,
         -3.9373e+00, -7.4720e-01, -1.0483e+00, -7.5145e+00, -3.8551e+00,
         -5.1689e+00, -1.2816e+00, -4.9733e+00, -5.4351e+00, -1.8549e+00,
         -3.3207e-01, -1.8963e+00, -3.5398e+00, -3.2054e+00, -7.3057e+00,
         -7.6148e-02, -7.8230e-01, -2.3709e+00, -4.2063e+00, -2.1984e+00,
         -3.6228e-01, -4.6996e+00, -3.0654e+00, -3.2912e-03, -4.0276e-01,
         -6.4977e+00, -1.0216e+00, -4.6901e+00, -3.9250e+00, -4.8088e+00,
         -5.6998e+00, -1.3010e-02, -1.8368e+00],
        [-3.3044e+00, -1.0319e-02, -8.7674e+00, -5.2868e+00, -6.7541e-01,
         -8.7047e-01, -4.0037e+00, -3.9222e-01, -3.4152e-01, -2.2293e+00,
         -2.9142e-02, -4.3188e+00, -5.2750e+00, -4.1764e+00, -6.8956e-03,
         -8.7781e-01, -7.6296e-01, -1.0332e+00, -4.1965e-01, -5.5560e-01,
         -5.0666e-01, -6.7566e+00, -7.0060e-01, -7.9373e-01, -3.1336e+00,
         -7.8659e-01, -4.4716e+00, -1.9113e+00, -2.2779e+00, -4.2088e+00,
         -1.4632e+00, -3.7058e+00, -3.2600e-01, -1.2548e-01, -2.1457e+00,
         -4.3699e-01, -2.8770e+00, -2.3523e+00]]),
  'full_logprob': None,
  'token_ids': [[50256, 48, 25, 1374, 389, 345, 30, 198, 32, 16, 8, 632, 318, 477, 1165, 2562, 284, 651, 22943, 11, 1110, 25966, 278, 290, 523, 6071, 13, 1550, 1353, 286, 326, 867, 661, 3092, 6751, 1334, 2233, 284, 511], [50256, 48, 25, 1374, 1263, 318, 262, 6881, 30, 198, 32, 25, 1081, 345, 447, 247, 303, 2192, 2982, 11, 262, 6881, 3407, 2279, 326, 356, 460, 760, 13, 632, 23687, 422, 530, 886, 286, 2272, 284, 1194, 290]], 'offsets': [[0, 0, 1, 2, 6, 10, 14, 15, 16, 17, 18, 19, 22, 25, 29, 33, 38, 41, 45, 56, 57, 61, 66, 69, 73, 76, 82, 83, 86, 90, 93, 98, 103, 110, 115, 126, 131, 135, 138], [0, 0, 1, 2, 6, 10, 13, 17, 26, 27, 28, 29, 30, 33, 37, 38, 39, 41, 50, 56, 57, 61, 70, 79, 90, 95, 98, 102, 107, 108, 111, 121, 126, 130, 134, 137, 143, 146, 154]]}]
***************************
```

# Before your PR is "Ready for review"
**Pre checks**:
- [x] 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**:
- [x] 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 #6625
* TE FP8 path needs changes in https://github.com/NVIDIA/NeMo/pull/6458/files to run successfully, otherwise it might raise a ValueError when precision is set to 16 or 32.