NeMo Container 24.07: NLPSaveRestoreConnector.save_to() is calling modelopt.torch.opt.plugins.save_sharded_modelopt_state with an unsupported parameter
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

Context:
```
Container: NeMo 24.07
```

When running the `/opt/NeMo/scripts/checkpoint_converters/convert_starcoder2_hf_to_nemo.py` checkpoint conversion script, it will at some point call `NLPSaveRestoreConnector.save_to()` (from `/opt/NeMo/nemo/collections/nlp/parts/nlp_overrides.py`). That method by its way, will call `modelopt.torch.opt.plugins.save_sharded_modelopt_state` using an unsupported parameter (`prefix`):
```python
            if HAVE_MODELOPT and hasattr(self.lightning_module, "get_model_module_list"):
                save_sharded_modelopt_state(
                    self.lightning_module.get_model_module_list(),
                    ckpt_to_dir(filepath),
                    self.unwrapped_checkpoint_io.save_sharded_strategy,
                    prefix="model.",
                )
```

Please note that in `/usr/local/lib/python3.10/dist-packages/modelopt/torch/opt/plugins/mcore_dist_checkpointing.py`, the current `save_sharded_modelopt_state` only supports three parameters (`model`, `checkpoint_name` and `sharded_strategy`):
```python
def save_sharded_modelopt_state(
    model: List[torch.nn.Module], checkpoint_name: str, sharded_strategy: Tuple[str, int]
):
    """Save modelopt_state in the sharded state_dict format.

    Args:
        model: the model to restore the modelopt optimization
        checkpoint_name: the checkpoint folder path
    """
    if not mto.ModeloptStateManager.is_converted(model[0]):
        return
    if len(model) > 1:
        raise ValueError("sharded_modelopt_state does not support virtual pipeline parallel!")
    modelopt_checkpoint_name = checkpoint_name + "/modelopt_state"
    if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
        os.makedirs(modelopt_checkpoint_name, exist_ok=True)
    dist_checkpointing.save(
        get_sharded_modelopt_state(-1, model[0]),
        modelopt_checkpoint_name,
        sharded_strategy,
    )
```

This is then causing the following error:
```bash
NeMo I 2024-09-13 09:32:40 dist_ckpt_io:421] Using TorchDistSaveShardedStrategy(torch_dist, 1) dist-ckpt save strategy.
`zarr` distributed checkpoint backend is deprecated. Please switch to PyTorch Distributed format (`torch_dist`).
[rank0]: Traceback (most recent call last):
[rank0]:   File "/opt/NeMo/scripts/checkpoint_converters/convert_starcoder2_hf_to_nemo.py", line 375, in <module>
[rank0]:     convert(args)
[rank0]:   File "/opt/NeMo/scripts/checkpoint_converters/convert_starcoder2_hf_to_nemo.py", line 369, in convert
[rank0]:     model.save_to(args.output_path)
[rank0]:   File "/opt/NeMo/nemo/core/classes/modelPT.py", line 410, in save_to
[rank0]:     self._save_restore_connector.save_to(self, str(save_path))  # downstream tasks expect str, not Path
[rank0]:   File "/opt/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 1027, in save_to
[rank0]:     save_sharded_modelopt_state(
[rank0]: TypeError: save_sharded_modelopt_state() got an unexpected keyword argument 'prefix'
```

**Steps/Code to reproduce bug**

Download a HuggingFace starcoder2 7b checkpoint, start a NeMo 2407 container in interactive mode and try to run the `/opt/NeMo/scripts/checkpoint_converters/convert_starcoder2_hf_to_nemo.py` script:
```bash
python /opt/NeMo/scripts/checkpoint_converters/convert_starcoder2_hf_to_nemo.py --input_name_or_path /shared/models/huggingface/starcoder2_7b/ --output_path /shared/models/nemo/starcoder2_7b.nemo --precision bf16
```

It will crash:
```bash
NeMo I 2024-09-13 09:32:40 dist_ckpt_io:421] Using TorchDistSaveShardedStrategy(torch_dist, 1) dist-ckpt save strategy.
`zarr` distributed checkpoint backend is deprecated. Please switch to PyTorch Distributed format (`torch_dist`).
[rank0]: Traceback (most recent call last):
[rank0]:   File "/opt/NeMo/scripts/checkpoint_converters/convert_starcoder2_hf_to_nemo.py", line 375, in <module>
[rank0]:     convert(args)
[rank0]:   File "/opt/NeMo/scripts/checkpoint_converters/convert_starcoder2_hf_to_nemo.py", line 369, in convert
[rank0]:     model.save_to(args.output_path)
[rank0]:   File "/opt/NeMo/nemo/core/classes/modelPT.py", line 410, in save_to
[rank0]:     self._save_restore_connector.save_to(self, str(save_path))  # downstream tasks expect str, not Path
[rank0]:   File "/opt/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 1027, in save_to
[rank0]:     save_sharded_modelopt_state(
[rank0]: TypeError: save_sharded_modelopt_state() got an unexpected keyword argument 'prefix'
```

**Expected behavior**

Shouldn't crash.

**Environment overview (please complete the following information)**

 - Environment location: Docker (NeMo 24.07 image)


Issue solved by pulling a fresh container.