Add missing imports for torch dist ckpt in export
State: closed | Pull request: yes
Labels: name: bug

name: Run CICD

# What does this PR do ?

fixing the bug for torch dist ckpt read in export

@oyilmaz-nvidia I think `np_bfloat16` is still missing: https://github.com/NVIDIA/NeMo/pull/9826/files#diff-5bc83531f46c91c2dfc23d247e2af872371ba3f6f852688ba012adc6b6adec53R83

Thanks @mikolajblaz , just added that import as well.



I am trying to delpoy torch dist checkpoint for evaluation of checkpoint generated from NeMo 2.0 training and got the same error below even after incorporating the changes in this PR. @oyilmaz-nvidia @mikolajblaz is anything else required to export torch dist checkpoint to TRTLLM ? Thanks!

```python
Not a supported nemo file format. Only distributed mcore nemo checkpoints are support.
  File "/workspace/software/nemo_latest_1/NeMo/nemo/export/trt_llm/nemo_ckpt_loader/nemo_file.py", line 249, in load_nemo_model
    raise Exception(
  File "/workspace/software/nemo_latest_1/NeMo/nemo/export/tensorrt_llm.py", line 279, in export
    model, model_configs, self.tokenizer = load_nemo_model(nemo_checkpoint_path, nemo_export_dir) 
```

@athitten I am looking into this right now. The code currently looks for a specific folder name. Need to change that. Will push the updates here.

@mikolajblaz Getting the error below. Can you please help me understand what we should do?
 
```
Traceback (most recent call last):
  File "/opt/NeMo/scripts/deploy/nlp/deploy_triton.py", line 255, in get_trtllm_deployable
    trt_llm_exporter.export(
  File "/opt/NeMo/nemo/export/tensorrt_llm.py", line 279, in export
    model, model_configs, self.tokenizer = load_nemo_model(nemo_checkpoint_path, nemo_export_dir)
  File "/opt/NeMo/nemo/export/trt_llm/nemo_ckpt_loader/nemo_file.py", line 240, in load_nemo_model
    model = load_sharded_metadata(dist_ckpt_folder)
  File "/opt/NeMo/nemo/export/trt_llm/nemo_ckpt_loader/nemo_file.py", line 118, in load_sharded_metadata
    return load_sharded_metadata_torch_dist(checkpoint_dir, torch_tensor)
  File "/opt/NeMo/nemo/export/trt_llm/nemo_ckpt_loader/nemo_file.py", line 68, in load_sharded_metadata_torch_dist
    metadata = fs_reader.read_metadata()
  File "/usr/local/lib/python3.10/dist-packages/torch/distributed/checkpoint/filesystem.py", line 562, in read_metadata
    path = self.fs.concat_path(self.path, ".metadata")
AttributeError: 'TarFileSystemReader' object has no attribute 'fs'
```

> @mikolajblaz Getting the error below. Can you please help me understand what we should do?
> 
> ```
> AttributeError: 'TarFileSystemReader' object has no attribute 'fs'
> ```

@oyilmaz-nvidia  I checked that it's because new PyTorch version does some extra initialization in FileSystemReader that TarFileSystemReader inherits from.
Can you test if the following implementation of [TarFileSystemReader]9https://github.com/NVIDIA/NeMo/blob/main/nemo/export/trt_llm/nemo_ckpt_loader/nemo_file.py#L51) works for you (should be compatible also with older PyTorch):
```
class TarFileSystemReader(FileSystemReader):
    """Reader that accepts both Path and TarPath checkpoint directory.

    The FileSystemReader works with TarPath, but expects a pure Path.
    It's enough to skip the Path check in __init__.
    """

    def __init__(self, path: Union[Path, TarPath]) -> None:
        """Makes sure that super().__init__ gets a pure path as expected."""
        super_path = str(path) if isinstance(path, TarPath) else path
        super().__init__(super_path)
        if isinstance(path, TarPath):
            self.path = path  # overwrites path set in super().__init__ call
```

@mikolajblaz Thanks for the suggestion. It's working now but got this output from the ckpt you provided 
```
--- Output:  [['2 3 3 3 1 1 1 1 1                                                                                                              '], ['\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n'], ['dro\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n']]
```
Was that a mock model? If yes, I guess it should be good to go. Otherwise, we are making a mistake while getting the weights.

@oyilmaz-nvidia I used this branch (`fix-export-torch-dist`) with the latest changes to try the export of the new torch dist ckpt to TRTLLM, but I am still hitting the error: `Not a supported nemo file format. Only distributed mcore nemo checkpoints are support.`
Can you please let me know if there is any other change required or if I am missing something. 



> Was that a mock model? If yes, I guess it should be good to go. Otherwise, we are making a mistake while getting the weights.

yes, a mock model

@athitten Looks like the ckpt you used has a different folder structure than what we expect in a .nemo file. So, I got another ckpt from @mikolajblaz and it worked for me.

CI passed: https://github.com/NVIDIA/NeMo/actions/runs/10096199676