Weights only load fail with pre014 checkpoints
- Dominant language
- Python
- Stars
- 17.9k
- Forks
- 4.5k
- Avg merge
- 4d 6h
- Merged PRs (30d)
- 271
Description
**Describe the bug**
A clear and concise description of what the bug is.
```
105: [rank105]: Traceback (most recent call last):
105: [rank105]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/pretrain_gpt.py", line 233, in
105: [rank105]: pretrain(
105: [rank105]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 666, in pretrain
105: [rank105]: model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
105: [rank105]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
105: [rank105]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 1162, in setup_model_and_optimizer
105: [rank105]: args.iteration, args.num_floating_point_operations_so_far = load_checkpoint(
105: [rank105]: ^^^^^^^^^^^^^^^^
1312: [rank1312]: Traceback (most recent call last):
1312: [rank1312]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/pretrain_gpt.py", line 233, in
1312: [rank1312]: pretrain(
1312: [rank1312]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 666, in pretrain
1312: [rank1312]: model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
1312: [rank1312]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
search hit BOTTOM, continuing at TOP 197892,18 64%
1312: [rank1312]: model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
1312: [rank1312]: ^^^^^^^^^^^^^^^^^^^^^^^^^^
1312: [rank1312]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 1162, in setup_model_and_optimizer
1312: [rank1312]: args.iteration, args.num_floating_point_operations_so_far = load_checkpoint(
1312: [rank1312]: ^^^^^^^^^^^^^^^^
105: [rank105]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 1393, in load_checkpoint
105: [rank105]: state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint(
105: [rank105]: ^^^^^^^^^^^^^^^^^^^^^^
105: [rank105]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 1114, in _load_base_checkpoint
105: [rank105]: return _load_global_dist_base_checkpoint(
105: [rank105]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
105: [rank105]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 987, in _load_global_dist_base_checkpoint
105: [rank105]: state_dict = dist_checkpointing.load_common_state_dict(checkpoint_name)
105: [rank105]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
1312: [rank1312]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 1393, in load_checkpoint
1312: [rank1312]: state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint(
1312: [rank1312]: ^^^^^^^^^^^^^^^^^^^^^^
1312: [rank1312]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 1114, in _load_base_checkpoint
1312: [rank1312]: return _load_global_dist_base_checkpoint(
1312: [rank1312]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
1312: [rank1312]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 987, in _load_global_dist_base_checkpoint
1312: [rank1312]: state_dict = dist_checkpointing.load_common_state_dict(checkpoint_name)
1312: [rank1312]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
105: [rank105]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 189, in load_common_state_dict
105: [rank105]: return common_strategy.load_common(checkpoint_dir)
105: [rank105]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
105: [rank105]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/dist_checkpointing/strategies/common.py", line 89, in load_common
105: [rank105]: return torch.load(load_path, map_location='cpu')
105: [rank105]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
197910,2 64%
105: [rank105]: return torch.load(load_path, map_location='cpu')
105: [rank105]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
105: [rank105]: File "/usr/local/lib/python3.12/dist-packages/torch/serialization.py", line 1524, in load
105: [rank105]: raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
105: [rank105]: _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, ^[[1mdo those steps only if you trust the source of the checkpoint^[[0m.
1312: [rank1312]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/dist_checkpointing/serialization.py", line 189, in load_common_state_dict
1312: [rank1312]: return common_strategy.load_common(checkpoint_dir)
1312: [rank1312]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
1312: [rank1312]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/dist_checkpointing/strategies/common.py", line 89, in load_common
1312: [rank1312]: return torch.load(load_path, map_location='cpu')
1312: [rank1312]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
1312: [rank1312]: File "/usr/local/lib/python3.12/dist-packages/torch/serialization.py", line 1524, in load
1312: [rank1312]: raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
1312: [rank1312]: _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, ^[[1mdo those steps only if you trust the source of the checkpoint^[[0m.
105: [rank105]: (1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
105: [rank105]: (2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
105: [rank105]: WeightsUnpickler error: Unsupported global: GLOBAL torch.optim.adamw.AdamW was not an allowed global by default. Please use `torch.serialization.add_safe_globals([torch.optim.adamw.AdamW])` or the `torch.serialization.safe_globals([torch.optim.adamw.AdamW])` context manager to allowlist this global if you trust this class/function.
105:
```
Just save a checkpoint with `--dist-ckpt-save-pre-mcore-014 --dist-ckpt-optim-fully-reshardable --distrib-optim-fully-reshardable-mem-efficient` and try to load it
**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.
**Additional context**
Add any other context about the problem here.
Contributor guide
Assessment
This issue has not been assessed yet.