[ABL] Weight Tying raises errors when running distributed checkpointing #348

NOTE: weight tying does not seem to work with distributed checkpointing

E         File "/raid/s3/opengptx/max_lue/repositories/modalities/src/modalities/checkpointing/stateful/app_state.py", line 105, in load_state_dict
E           OptimizerStateRetriever.load_state_dict_(
E         File "/raid/s3/opengptx/max_lue/repositories/modalities/src/modalities/checkpointing/stateful/app_state.py", line 175, in load_state_dict_
E           set_optimizer_state_dict(
E         File "/raid/s3/opengptx/max_lue/venvs/modalities_311/lib/python3.11/site-packages/torch/distributed/checkpoint/state_dict.py", line 1246, in set_optimizer_state_dict
E           _load_optim_state_dict(model, optimizers, optim_state_dict, info)
E         File "/raid/s3/opengptx/max_lue/venvs/modalities_311/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
E           return func(*args, **kwargs)
E                  ^^^^^^^^^^^^^^^^^^^^^
E         File "/raid/s3/opengptx/max_lue/venvs/modalities_311/lib/python3.11/site-packages/torch/distributed/checkpoint/state_dict.py", line 891, in _load_optim_state_dict
E           optim_state_dict = _unflatten_optim_state_dict(
E                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
E         File "/raid/s3/opengptx/max_lue/venvs/modalities_311/lib/python3.11/site-packages/torch/distributed/checkpoint/state_dict.py", line 716, in _unflatten_optim_state_dict
E           cast(DictValueType, state[fqn])[state_name] = state_dict[
E                                                         ^^^^^^^^^^^
E       KeyError: 'state.transformer.lm_head.weight.step'

Originally posted by @le1nux in #336