NeMo 2.0 llm sharded_state_dict error
State: closed | Pull request: no
Labels: name: bug

**Describe the bug**

```
Training epoch 0, iteration 196/3999 | lr: 2.954e-05 | consumed_samples: 4728 | global_batch_size: 24 | global_step: 196 | reduced_train_loss: 10.84 | train_step_timing in s: 0.9765
Training epoch 0, iteration 197/3999 | lr: 2.969e-05 | consumed_samples: 4752 | global_batch_size: 24 | global_step: 197 | reduced_train_loss: 10.84 | train_step_timing in s: 0.9826
Training epoch 0, iteration 198/3999 | lr: 2.984e-05 | consumed_samples: 4776 | global_batch_size: 24 | global_step: 198 | reduced_train_loss: 10.84 | train_step_timing in s: 0.9742
Training epoch 0, iteration 199/3999 | lr: 2.999e-05 | consumed_samples: 4800 | global_batch_size: 24 | global_step: 199 | reduced_train_loss: 10.85 | train_step_timing in s: 1.085
[NeMo W 2024-09-29 08:11:31 nemo_logging:349] /home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py:384: `ModelCheckpoint(monitor='val_loss')` could not find the monitored key in the returned metrics: ['lr', 'consumed_samples', 'global_batch_size', 'global_step', 'step', 'reduced_train_loss', 'grad_norm', 'train_step_timing in s', 'epoch']. HINT: Did you call `log('val_loss', value)` in the `LightningModule`?
    
Epoch 0, global step 199: 'val_loss' was not in top 10
[rank0]: Traceback (most recent call last):
[rank0]:   File "/home/usr/.local/bin/nemorun", line 8, in <module>
[rank0]:     sys.exit(app())
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/main.py", line 338, in __call__
[rank0]:     raise e
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/main.py", line 321, in __call__
[rank0]:     return get_command(self)(*args, **kwargs)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 1157, in __call__
[rank0]:     return self.main(*args, **kwargs)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/core.py", line 728, in main
[rank0]:     return _main(
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/core.py", line 197, in _main
[rank0]:     rv = self.invoke(ctx)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 1688, in invoke
[rank0]:     return _process_result(sub_ctx.command.invoke(sub_ctx))
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 1688, in invoke
[rank0]:     return _process_result(sub_ctx.command.invoke(sub_ctx))
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 1434, in invoke
[rank0]:     return ctx.invoke(self.callback, **ctx.params)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 783, in invoke
[rank0]:     return __callback(*args, **kwargs)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/main.py", line 703, in wrapper
[rank0]:     return callback(**use_params)
[rank0]:   File "/home/usr/code/NeMo-Run/src/nemo_run/cli/api.py", line 790, in command
[rank0]:     self.cli_execute(fn, ctx.args, type)
[rank0]:   File "/home/usr/code/NeMo-Run/src/nemo_run/cli/api.py", line 842, in cli_execute
[rank0]:     self._execute_task(fn, filtered_args)
[rank0]:   File "/home/usr/code/NeMo-Run/src/nemo_run/cli/api.py", line 892, in _execute_task
[rank0]:     run_task()
[rank0]:   File "/home/usr/code/NeMo-Run/src/nemo_run/cli/api.py", line 871, in run_task
[rank0]:     run.run(
[rank0]:   File "/home/usr/code/NeMo-Run/src/nemo_run/run/api.py", line 65, in run
[rank0]:     direct_run_fn(fn_or_script, dryrun=dryrun)
[rank0]:   File "/home/usr/code/NeMo-Run/src/nemo_run/run/task.py", line 77, in direct_run_fn
[rank0]:     built_fn()
[rank0]:   File "/home/usr/code/NeMo/nemo/collections/llm/api.py", line 127, in pretrain
[rank0]:     return train(
[rank0]:   File "/home/usr/code/NeMo/nemo/collections/llm/api.py", line 85, in train
[rank0]:     trainer.fit(model, data)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 538, in fit
[rank0]:     call._call_and_handle_interrupt(
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank0]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]:     return function(*args, **kwargs)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 574, in _fit_impl
[rank0]:     self._run(model, ckpt_path=ckpt_path)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 981, in _run
[rank0]:     results = self._run_stage()
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1025, in _run_stage
[rank0]:     self.fit_loop.run()
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 205, in run
[rank0]:     self.advance()
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 363, in advance
[rank0]:     self.epoch_loop.run(self._data_fetcher)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 140, in run
[rank0]:     self.advance(data_fetcher)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 269, in advance
[rank0]:     call._call_callback_hooks(trainer, "on_train_batch_end", batch_output, batch, batch_idx)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 218, in _call_callback_hooks
[rank0]:     fn(trainer, trainer.lightning_module, *args, **kwargs)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 317, in on_train_batch_end
[rank0]:     self._save_last_checkpoint(trainer, monitor_candidates)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 696, in _save_last_checkpoint
[rank0]:     self._save_checkpoint(trainer, filepath)
[rank0]:   File "/home/usr/code/NeMo/nemo/lightning/pytorch/callbacks/model_checkpoint.py", line 449, in _save_checkpoint
[rank0]:     trainer.save_checkpoint(ckpt_filepath, save_weights_only, storage_options=storage_options)
[rank0]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1365, in save_checkpoint
[rank0]:     self.strategy.save_checkpoint(checkpoint, filepath, storage_options=storage_options)
[rank0]:   File "/home/usr/code/NeMo/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 660, in save_checkpoint
[rank0]:     checkpoint["optimizer"] = [self.optimizer_sharded_state_dict()]
[rank0]:   File "/home/usr/code/NeMo/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 638, in optimizer_sharded_state_dict
[rank0]:     return _strategy_lib.optimizer_sharded_state_dict(
[rank0]:   File "/home/usr/code/NeMo/nemo/lightning/_strategy_lib.py", line 455, in optimizer_sharded_state_dict
[rank0]:     return optimizer.sharded_state_dict(
[rank0]:   File "/home/usr/code/NeMo/nemo/lightning/pytorch/optim/megatron.py", line 114, in sharded_state_dict
[rank0]:     state_dict = self.mcore_optimizer.sharded_state_dict(
[rank0]:   File "/home/usr/code/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 950, in sharded_state_dict
[rank0]:     param_state = self.sharded_param_state_fs_model_space(
[rank0]:   File "/home/usr/code/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 1161, in sharded_param_state_fs_model_space
[rank0]:     dtype=state_ten.dtype,
[rank0]: AttributeError: 'NoneType' object has no attribute 'dtype'
[rank1]: Traceback (most recent call last):
[rank1]:   File "/home/usr/.local/bin/nemorun", line 8, in <module>
[rank1]:     sys.exit(app())
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/main.py", line 338, in __call__
[rank1]:     raise e
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/main.py", line 321, in __call__
[rank1]:     return get_command(self)(*args, **kwargs)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 1157, in __call__
[rank1]:     return self.main(*args, **kwargs)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/core.py", line 728, in main
[rank1]:     return _main(
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/core.py", line 197, in _main
[rank1]:     rv = self.invoke(ctx)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 1688, in invoke
[rank1]:     return _process_result(sub_ctx.command.invoke(sub_ctx))
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 1688, in invoke
[rank1]:     return _process_result(sub_ctx.command.invoke(sub_ctx))
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 1434, in invoke
[rank1]:     return ctx.invoke(self.callback, **ctx.params)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/click/core.py", line 783, in invoke
[rank1]:     return __callback(*args, **kwargs)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/typer/main.py", line 703, in wrapper
[rank1]:     return callback(**use_params)
[rank1]:   File "/home/usr/code/NeMo-Run/src/nemo_run/cli/api.py", line 790, in command
[rank1]:     self.cli_execute(fn, ctx.args, type)
[rank1]:   File "/home/usr/code/NeMo-Run/src/nemo_run/cli/api.py", line 842, in cli_execute
[rank1]:     self._execute_task(fn, filtered_args)
[rank1]:   File "/home/usr/code/NeMo-Run/src/nemo_run/cli/api.py", line 892, in _execute_task
[rank1]:     run_task()
[rank1]:   File "/home/usr/code/NeMo-Run/src/nemo_run/cli/api.py", line 871, in run_task
[rank1]:     run.run(
[rank1]:   File "/home/usr/code/NeMo-Run/src/nemo_run/run/api.py", line 65, in run
[rank1]:     direct_run_fn(fn_or_script, dryrun=dryrun)
[rank1]:   File "/home/usr/code/NeMo-Run/src/nemo_run/run/task.py", line 77, in direct_run_fn
[rank1]:     built_fn()
[rank1]:   File "/home/usr/code/NeMo/nemo/collections/llm/api.py", line 127, in pretrain
[rank1]:     return train(
[rank1]:   File "/home/usr/code/NeMo/nemo/collections/llm/api.py", line 85, in train
[rank1]:     trainer.fit(model, data)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 538, in fit
[rank1]:     call._call_and_handle_interrupt(
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank1]:     return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/strategies/launchers/subprocess_script.py", line 105, in launch
[rank1]:     return function(*args, **kwargs)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 574, in _fit_impl
[rank1]:     self._run(model, ckpt_path=ckpt_path)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 981, in _run
[rank1]:     results = self._run_stage()
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1025, in _run_stage
[rank1]:     self.fit_loop.run()
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 205, in run
[rank1]:     self.advance()
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/loops/fit_loop.py", line 363, in advance
[rank1]:     self.epoch_loop.run(self._data_fetcher)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 140, in run
[rank1]:     self.advance(data_fetcher)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/loops/training_epoch_loop.py", line 269, in advance
[rank1]:     call._call_callback_hooks(trainer, "on_train_batch_end", batch_output, batch, batch_idx)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py", line 218, in _call_callback_hooks
[rank1]:     fn(trainer, trainer.lightning_module, *args, **kwargs)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 317, in on_train_batch_end
[rank1]:     self._save_last_checkpoint(trainer, monitor_candidates)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/callbacks/model_checkpoint.py", line 696, in _save_last_checkpoint
[rank1]:     self._save_checkpoint(trainer, filepath)
[rank1]:   File "/home/usr/code/NeMo/nemo/lightning/pytorch/callbacks/model_checkpoint.py", line 449, in _save_checkpoint
[rank1]:     trainer.save_checkpoint(ckpt_filepath, save_weights_only, storage_options=storage_options)
[rank1]:   File "/home/usr/.local/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py", line 1365, in save_checkpoint
[rank1]:     self.strategy.save_checkpoint(checkpoint, filepath, storage_options=storage_options)
[rank1]:   File "/home/usr/code/NeMo/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 660, in save_checkpoint
[rank1]:     checkpoint["optimizer"] = [self.optimizer_sharded_state_dict()]
[rank1]:   File "/home/usr/code/NeMo/nemo/lightning/pytorch/strategies/megatron_strategy.py", line 638, in optimizer_sharded_state_dict
[rank1]:     return _strategy_lib.optimizer_sharded_state_dict(
[rank1]:   File "/home/usr/code/NeMo/nemo/lightning/_strategy_lib.py", line 455, in optimizer_sharded_state_dict
[rank1]:     return optimizer.sharded_state_dict(
[rank1]:   File "/home/usr/code/NeMo/nemo/lightning/pytorch/optim/megatron.py", line 114, in sharded_state_dict
[rank1]:     state_dict = self.mcore_optimizer.sharded_state_dict(
[rank1]:   File "/home/usr/code/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 950, in sharded_state_dict
[rank1]:     param_state = self.sharded_param_state_fs_model_space(
[rank1]:   File "/home/usr/code/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 1161, in sharded_param_state_fs_model_space
[rank1]:     dtype=state_ten.dtype,
[rank1]: AttributeError: 'NoneType' object has no attribute 'dtype'
```

**Steps/Code to reproduce bug**

```
nemorun llm pretrain --factory llama3_8b trainer.max_steps=4000 trainer.num_nodes=1 trainer.devices=2 data.seq_length=128 data.global_batch_size=24 model.config.seq_length=128 model.config.num_layers=4 model.config.hidden_size=256 model.config.ffn_hidden_size=1024 log.log_dir=exp_0929_n1d2_mini_fixckptsave trainer.strategy.save_ckpt_format=torch -y
```


DEBUG: master_param, dict_keys(['fp32_param', 'exp_avg', 'exp_avg_sq', 'master_param'])
DEBUG: master_param, dict_keys(['fp32_param', 'exp_avg', 'exp_avg_sq', 'master_param'])


install TE from source

https://github.com/NVIDIA/TransformerEngine/pull/1130