Zarr directory not found for checkpoint_dir
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

When I run a tuning job with default nvidia callback configuration. I get below error. If checkpoint_callback is disabled then epoch finishes fine. This seems to be an internal issue to me.

Image used: Nemo 28.03

A clear and concise description of what the bug is.Set the environment variable HYDRA_FULL_ERROR=1 for a complete stack trace.
Traceback (most recent call last):
  File "/opt/NeMo/examples/nlp/language_modeling/tuning/megatron_gpt_sft.py", line 212, in main
    trainer.fit(model)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 532, in fit
    call._call_and_handle_interrupt(
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 42, in _call_and_handle_interrupt
    return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 93, in launch
    return function(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 571, in _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 980, in _run
    results = self._run_stage()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1023, in _run_stage
    self.fit_loop.run()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/fit_loop.py", line 202, in run
    self.advance()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/fit_loop.py", line 355, in advance
    self.epoch_loop.run(self._data_fetcher)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/training_epoch_loop.py", line 134, in run
    self.on_advance_end()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/training_epoch_loop.py", line 249, in on_advance_end
    self.val_loop.run()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/utilities.py", line 181, in _decorator
    return loop_run(self, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 122, in run
    return self.on_run_end()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 258, in on_run_end
    self._on_evaluation_end()
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/loops/evaluation_loop.py", line 303, in _on_evaluation_end
    call._call_callback_hooks(trainer, hook_name, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 194, in _call_callback_hooks
    fn(trainer, trainer.lightning_module, *args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 311, in on_validation_end
    self._save_topk_checkpoint(trainer, monitor_candidates)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 360, in _save_topk_checkpoint
    self._save_monitor_checkpoint(trainer, monitor_candidates)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 663, in _save_monitor_checkpoint
    self._update_best_and_save(current, trainer, monitor_candidates)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 714, in _update_best_and_save
    self._save_checkpoint(trainer, filepath)
  File "/opt/NeMo/nemo/utils/callbacks/nemo_model_checkpoint.py", line 272, in _save_checkpoint
    super()._save_checkpoint(trainer, filepath)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 365, in _save_checkpoint
    trainer.save_checkpoint(filepath, self.save_weights_only)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/trainer.py", line 1316, in save_checkpoint
    self._checkpoint_connector.save_checkpoint(filepath, weights_only=weights_only, storage_options=storage_options)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py", line 506, in save_checkpoint
    _checkpoint = self.dump_checkpoint(weights_only)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/connectors/checkpoint_connector.py", line 491, in dump_checkpoint
    call._call_callbacks_on_save_checkpoint(trainer, checkpoint)
  File "/usr/local/lib/python3.10/dist-packages/pytorch_lightning/trainer/call.py", line 221, in _call_callbacks_on_save_checkpoint
    callback.on_save_checkpoint(trainer, trainer.lightning_module, checkpoint)
  File "/opt/NeMo/nemo/utils/callbacks/nemo_model_checkpoint.py", line 181, in on_save_checkpoint
    pl_module.save_to(save_path=app_state.model_restore_path)
  File "/opt/NeMo/nemo/core/classes/modelPT.py", line 385, in save_to
    self._save_restore_connector.save_to(self, str(save_path))  # downstream tasks expect str, not Path
  File "/opt/NeMo/nemo/collections/nlp/parts/nlp_overrides.py", line 489, in save_to
    dist_checkpointing.save(sharded_state_dict=sharded_state_dict, checkpoint_dir=dist_ckpt_dir)
  File "/opt/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 174, in save
    sharded_strategy.save(sharded_tensors, checkpoint_dir)
  File "/opt/Megatron-LM/megatron/core/dist_checkpointing/strategies/zarr.py", line 49, in save
    arrays = _create_or_open_zarr_arrays(sharded_tensors, checkpoint_dir)
  File "/opt/Megatron-LM/megatron/core/dist_checkpointing/strategies/zarr.py", line 73, in _create_or_open_zarr_arrays
    arr = zarr.open(checkpoint_dir / ten.key, 'r+', **open_kwargs)
  File "/usr/local/lib/python3.10/dist-packages/zarr/convenience.py", line 130, in open
    raise PathNotFoundError(path)
zarr.errors.PathNotFoundError: nothing found at path ''


**Steps/Code to reproduce bug**

Please list *minimal* steps or code snippet for us to be able to reproduce the bug.

A  helpful guide on on how to craft a minimal bug report  http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports. 


**Expected behavior**

A clear and concise description of what you expected to happen.

**Environment overview (please complete the following information)**

 - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)]
 - Method of NeMo install: [pip install or from source]. Please specify exact commands you used to install.
 - If method of install is [Docker], provide `docker pull` & `docker run` commands used

**Environment details**

If NVIDIA docker image is used you don't need to specify these.
Otherwise, please provide:
- OS version
- PyTorch version
- Python version

**Additional context**

Add any other context about the problem here.
Example: H100