In-framework inference fixes
State: closed | Pull request: yes
Labels: name: Run CICD

name: r2.0.0

# What does this PR do ?

Fixing two things in **r2.0.0** release testing, to be cherry-picked to the release branch: - repro below for an oldish Nemo checkpoint from FP8 SFT:
```
python /opt/NeMo/tests/export/nemo_export.py \
  --checkpoint_dir /opt/checkpoints/LLAMA2-7B-base/LLAMA2-7B-fp8-sft.nemo \
  --model_type llama \
  --min_tps 1 \
  --max_output_len 24 \
  --test_deployment False \
  --in_framework True \
  --model_name llama
```
Issues are:
 1. Loading old checkpoints - solved with https://github.com/NVIDIA/NeMo/commit/86408ccbabc4920a5c8e2a5d76cfc5d3afaf7bf4
```
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/dict_utils.py", line 191, in dict_list_map_inplace
[rank0]:     x[k] = dict_list_map_inplace(f, v)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/dict_utils.py", line 191, in dict_list_map_inplace
[rank0]:     x[k] = dict_list_map_inplace(f, v)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/dict_utils.py", line 195, in dict_list_map_inplace
[rank0]:     return f(x)
[rank0]:   File "/opt/megatron-lm/megatron/core/dist_checkpointing/strategies/common.py", line 118, in load_sharded_object
[rank0]:     raise CheckpointingException(err_msg) from e
[rank0]: megatron.core.dist_checkpointing.core.CheckpointingException: Object shard /tmp/tmpuuv6zoo9/model_weights/model.decoder.layers.self_attention.core_attention._extra_state/shard_0_32.pt not found
```

2. In-framework inference issues for FP8 - solved with https://github.com/NVIDIA/NeMo/commit/56d6e6f4c21bd1ce6328d756c34256791718f5dd (need to disable)
```
[rank0]:     mixed_qkv, _ = self.linear_qkv(hidden_states)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1552, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1561, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:   File "/opt/megatron-lm/megatron/core/extensions/transformer_engine.py", line 336, in forward
[rank0]:     out = super().forward(x, is_first_microbatch=_is_first_microbatch)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/eval_frame.py", line 574, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/module/layernorm_linear.py", line 1216, in forward
[rank0]:     out = fwd_fn(*args)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/module/layernorm_linear.py", line 101, in forward
[rank0]:     assert_dim_for_fp8_exec(inputmat)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/transformer_engine/pytorch/utils.py", line 237, in assert_dim_for_fp8_exec
[rank0]:     and tensor.size(1) % 16 == 0
[rank0]: AssertionError: FP8 execution requires 2D input matrices with height divisible by 8 and width divisible by 16, but got tensor with dims=[12, 4096]
```

**Collection**: NLP

# 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 
```

# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# 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
- [x] 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)