# coding=utf-8
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Input/output checkpointing."""
import os
import random
import sys
import numpy as np
from glob import glob
import torch
from codegeex.megatron import get_args, mpu, print_rank_0, update_num_microbatches, utils
_CHECKPOINT_VERSION = None
def set_checkpoint_version(value):
global _CHECKPOINT_VERSION
if _CHECKPOINT_VERSION is not None:
assert _CHECKPOINT_VERSION == value, "checkpoint versions do not match"
_CHECKPOINT_VERSION = value
def get_checkpoint_version():
global _CHECKPOINT_VERSION
return _CHECKPOINT_VERSION
def check_checkpoint_args(checkpoint_args):
"""Ensure fixed arguments for a model are the same for the input
arguments and the one retrieved from checkpoint."""
args = get_args()
def _compare(arg_name, old_arg_name=None):
if old_arg_name is not None:
checkpoint_value = getattr(checkpoint_args, old_arg_name)
else:
checkpoint_value = getattr(checkpoint_args, arg_name)
args_value = getattr(args, arg_name)
error_message = (
"{} value from checkpoint ({}) is not equal to the "
"input argument value ({}).".format(arg_name, checkpoint_value, args_value)
)
assert checkpoint_value == args_value, error_message
_compare("num_layers")
_compare("hidden_size")
_compare("num_attention_heads")
_compare("max_position_embeddings")
if args.vocab_file:
_compare("make_vocab_size_divisible_by")
_compare("padded_vocab_size")
_compare("tokenizer_type")
if get_checkpoint_version() < 3.0:
_compare("tensor_model_parallel_size", old_arg_name="model_parallel_size")
if get_checkpoint_version() >= 3.0:
_compare("tensor_model_parallel_size")
_compare("pipeline_model_parallel_size")
def ensure_directory_exists(filename):
"""Build filename's path if it does not already exists."""
dirname = os.path.dirname(filename)
if not os.path.exists(dirname):
os.makedirs(dirname)
def get_checkpoint_name(checkpoints_path, iteration, release=False):
"""A unified checkpoint name."""
if release:
directory = ""
else:
directory = f"global_step{iteration}"
# Use both the tensor and pipeline MP rank.
if mpu.get_pipeline_model_parallel_world_size() == 1:
return os.path.join(
checkpoints_path,
directory,
"mp_rank_{:02d}_model_states.pt".format(mpu.get_tensor_model_parallel_rank()),
)
return os.path.join(
checkpoints_path,
directory,
"mp_rank_{:02d}_{:03d}_model_states.pt".format(
mpu.get_tensor_model_parallel_rank(), mpu.get_pipeline_model_parallel_rank()
),
)
def get_checkpoint_tracker_filename(checkpoints_path):
"""Tracker file rescords the latest chckpoint during
training to restart from."""
return os.path.join(checkpoints_path, "latest_checkpointed_iteration.txt")
def save_checkpoint(iteration, model, optimizer, lr_scheduler):
"""Save a model checkpoint."""
args = get_args()
# Only rank zero of the data parallel writes to the disk.
if not args.deepspeed:
model = utils.unwrap_model(model)
print_rank_0(
"saving checkpoint at iteration {:7d} to {}".format(iteration, args.save)
)
if (
not torch.distributed.is_initialized()
or mpu.get_data_parallel_rank() == 0
or args.deepspeed
):
# Arguments, iteration, and model.
state_dict = {}
state_dict["args"] = args
state_dict["checkpoint_version"] = 3.0
state_dict["iteration"] = iteration
state_dict["tokens"] = args.consumed_train_tokens
# DeepSpeed saves the model/optimizer/scheduler
if not args.deepspeed:
if len(model) == 1:
state_dict["model"] = model[0].state_dict_for_save_checkpoint()
else:
for i in range(len(model)):
mpu.set_virtual_pipeline_model_parallel_rank(i)
state_dict["model%d" % i] = model[
i
].state_dict_for_save_checkpoint()
# Optimizer stuff.
if not args.no_save_optim:
if optimizer is not None:
state_dict["optimizer"] = optimizer.state_dict()
if lr_scheduler is not None:
state_dict["lr_scheduler"] = lr_scheduler.state_dict()
# RNG states.
if not args.no_save_rng:
state_dict["random_rng_state"] = random.getstate()
state_dict["np_rng_state"] = np.random.get_state()
state_dict["torch_rng_state"] = torch.get_rng_state()
state_dict["cuda_rng_state"] = torch.cuda.get_rng_state()
state_dict["rng_tracker_states"] = mpu.get_cuda_rng_tracker().get_states()
# Save.
checkpoint_name = get_checkpoint_name(args.save, iteration)
if not args.deepspeed:
ensure_directory_exists(checkpoint_name)
torch.save(state_dict, checkpoint_name)
if args.deepspeed:
# megatron model uses state_dict_for_save_checkpointing instead of the standard state_dict
# state_dict is used by deepspeed for module saving so it needs to point to the right function
if args.no_pipeline_parallel:
original_state_dict = model[0].module.state_dict
model[0].module.state_dict = model[0].module.state_dict_for_save_checkpoint
# Saving is a collective communication
checkpoint_name = get_checkpoint_name(args.save, iteration)
# Trim off the filename and mp_rank_* directory.
for _ in range(2):
checkpoint_name = os.path.dirname(checkpoint_name)
model[0].save_checkpoint(checkpoint_name, client_state=state_dict)
if args.no_pipeline_parallel:
model[0].module.state_dict = original_state_dict
# Wait so everyone is done (necessary)
if torch.distributed.is_initialized():
torch.distributed.barrier()
print_rank_0(
" successfully saved checkpoint at iteration {:7d} to {}".format(
iteration, args.save
)
)
# And update the latest iteration
if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
tracker_filename = get_checkpoint_tracker_filename(args.save)
with open(tracker_filename, "w") as f:
f.write(str(iteration))
# Wait so everyone is done (not necessary)
if torch.distributed.is_initialized():
torch.distributed.barrier()
def _transpose_first_dim(t, num_splits, num_splits_first, model):
input_shape = t.size()
# We use a self_attention module but the values extracted aren't
# specific to self attention so should work for cross attention as well
while hasattr(model, "module"):
model = model.module
attention_module = model.language_model.encoder.layers[0].self_attention
hidden_size_per_attention_head = attention_module.hidden_size_per_attention_head
num_attention_heads_per_partition = (
attention_module.num_attention_heads_per_partition
)
if num_splits_first:
"""[num_splits * np * hn, h]
-->(view) [num_splits, np, hn, h]
-->(tranpose) [np, num_splits, hn, h]
-->(view) [np * num_splits * hn, h]"""
intermediate_shape = (
num_splits,
num_attention_heads_per_partition,
hidden_size_per_attention_head,
) + input_shape[1:]
t = t.view(*intermediate_shape)
t = t.transpose(0, 1).contiguous()
else:
"""[np * hn * num_splits, h]
-->(view) [np, hn, num_splits, h]
-->(tranpose) [np, num_splits, hn, h]
-->(view) [np * num_splits * hn, h]"""
intermediate_shape = (
num_attention_heads_per_partition,
hidden_size_per_attention_head,
num_splits,
) + input_shape[1:]
t = t.view(*intermediate_shape)
t = t.transpose(1, 2).contiguous()
t = t.view(*input_shape)
return t
def fix_query_key_value_ordering(model, checkpoint_version):
"""Fix up query/key/value matrix ordering if checkpoint
version is smaller than 2.0
"""
if checkpoint_version < 2.0:
if isinstance(model, list):
assert len(model) == 1
model = model[0]
for name, param in model.named_parameters():
if name.endswith((".query_key_value.weight", ".query_key_value.bias")):
if checkpoint_version == 0:
fixed_param = _transpose_first_dim(param.data, 3, True, model)
elif checkpoint_version == 1.0:
fixed_param = _transpose_first_dim(param.data, 3, False, model)
else:
print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")
sys.exit()
param.data.copy_(fixed_param)
if name.endswith((".key_value.weight", ".key_value.bias")):
if checkpoint_version == 0:
fixed_param = _transpose_first_dim(param.data, 2, True, model)
elif checkpoint_version == 1.0:
fixed_param = _transpose_first_dim(param.data, 2, False, model)
else:
print_rank_0(f"Invalid checkpoint version {checkpoint_version}.")
sys.exit()
param.data.copy_(fixed_param)
print_rank_0(
" succesfully fixed query-key-values ordering for"
" checkpoint version {}".format(checkpoint_version)
)
def load_deepspeed_state(model):
model = utils.unwrap_model(model)
args = get_args()
load_dir = args.load
if os.path.isdir(load_dir):
model_state_paths = glob(os.path.join(load_dir, "*model_states.pt"))
assert len(model_state_paths) == 1, (
"only support loading deepspeed checkpoint of model parallel size 1"
", but got {}".format(model_state_paths)
)
model_state_path = model_state_paths[0]
else:
model_state_path = load_dir
state_dict = torch.load(model_state_path, map_location="cpu")
state_dict = state_dict["module"]
model[0].load_state_dict(state_dict, strict=True)
def load_checkpoint(
model,
optimizer,
lr_scheduler,
load_arg="load",
strict=True,
):
"""Load a model checkpoint and return the iteration.
strict (bool): whether to strictly enforce that the keys in
:attr:`state_dict` of the checkpoint match the names of
parameters and buffers in model.
"""
args = get_args()
load_dir = getattr(args, load_arg)
if args.deepspeed:
loaded_dir, state_dict = model[0].load_checkpoint(load_dir)
if loaded_dir is None:
print_rank_0(
"WARNING: could not find the metadata file {} ".format(load_dir)
)
print_rank_0(
" will not load any checkpoints and will start from " "random"
)
return 0
release = False
else:
model = utils.unwrap_model(model)
if load_dir.endswith(".pt"):
checkpoint_name = load_dir
release = True
else:
# Read the tracker file and set the iteration.
tracker_filename = get_checkpoint_tracker_filename(load_dir)
# If no tracker file, return iretation zero.
if not os.path.isfile(tracker_filename):
print_rank_0(
"WARNING: could not find the metadata file {} ".format(tracker_filename)
)
iteration = 0
release = True
checkpoint_name = get_checkpoint_name(load_dir, iteration, release)
if not os.path.isfile(checkpoint_name):
print_rank_0(
" will not load any checkpoints and will start from random"
)
return 0
else:
# Otherwise, read the tracker file and either set the iteration or
# mark it as a release checkpoint.
iteration = 0
release = False
with open(tracker_filename, "r") as f:
metastring = f.read().strip()
try:
iteration = int(metastring)
except ValueError:
release = metastring == "release"
if not release:
print_rank_0(
"ERROR: Invalid metadata file {}. Exiting".format(
tracker_filename
)
)
sys.exit()
assert iteration > 0 or release, "error parsing metadata file {}".format(
tracker_filename
)
# Checkpoint.
checkpoint_name = get_checkpoint_name(load_dir, iteration, release)
print_rank_0(f" loading checkpoint from {args.load} at iteration {iteration}")
# Load the checkpoint.
try:
state_dict = torch.load(checkpoint_name, map_location="cpu")
except ModuleNotFoundError:
from megatron.fp16_deprecated import loss_scaler
# For backward compatibility.
print_rank_0(" > deserializing using the old code structure ...")
sys.modules["fp16.loss_scaler"] = sys.modules[
"megatron.fp16_deprecated.loss_scaler"
]
sys.modules["megatron.fp16.loss_scaler"] = sys.modules[
"megatron.fp16_deprecated.loss_scaler"
]
state_dict = torch.load(checkpoint_name, map_location="cpu")
sys.modules.pop("fp16.loss_scaler", None)
sys.modules.pop("megatron.fp16.loss_scaler", None)
except BaseException as e:
print_rank_0("could not load the checkpoint")
print_rank_0(e)
sys.exit()
# set checkpoint version
set_checkpoint_version(state_dict.get("checkpoint_version", 0))
# Set iteration.
if args.finetune or release:
iteration = 0
else:
try:
iteration = state_dict["iteration"]
if "tokens" in state_dict:
args.consumed_train_tokens = state_dict["tokens"]
except KeyError:
try: # Backward compatible with older checkpoints
iteration = state_dict["total_iters"]
except KeyError:
print_rank_0(
"A metadata file exists but unable to load "
"iteration from checkpoint {}, exiting".format(checkpoint_name)
)
sys.exit()
# Check arguments.
assert args.consumed_train_samples == 0
assert args.consumed_valid_samples == 0
if "args" in state_dict:
checkpoint_args = state_dict["args"]
check_checkpoint_args(checkpoint_args)
args.consumed_train_samples = getattr(
checkpoint_args, "consumed_train_samples", 0
)
update_num_microbatches(consumed_samples=args.consumed_train_samples)
args.consumed_valid_samples = getattr(
checkpoint_args, "consumed_valid_samples", 0
)
else:
print_rank_0("could not find arguments in the checkpoint ...")
# Model.
if not args.deepspeed:
if release:
if len(model) == 1:
model[0].load_state_dict(state_dict["module"], strict=strict)
else:
for i in range(len(model)):
mpu.set_virtual_pipeline_model_parallel_rank(i)
model[i].load_state_dict(state_dict["model%d" % i], strict=strict)
else:
if len(model) == 1:
model[0].load_state_dict(state_dict["module"], strict=strict)
else:
for i in range(len(model)):
mpu.set_virtual_pipeline_model_parallel_rank(i)
model[i].load_state_dict(state_dict["model%d" % i], strict=strict)
# Fix up query/key/value matrix ordering if needed
checkpoint_version = get_checkpoint_version()
print_rank_0(f" checkpoint version {checkpoint_version}")
fix_query_key_value_ordering(model, checkpoint_version)
# Optimizer.
if not args.deepspeed:
if not release and not args.finetune and not args.no_load_optim:
try:
if optimizer is not None:
optimizer.load_state_dict(state_dict["optimizer"])
if lr_scheduler is not None:
lr_scheduler.load_state_dict(state_dict["lr_scheduler"])
except KeyError:
print_rank_0(
"Unable to load optimizer from checkpoint {}. "
"Specify --no-load-optim or --finetune to prevent "
"attempting to load the optimizer state, "
"exiting ...".format(checkpoint_name)
)
sys.exit()
# rng states.
if not release and not args.finetune and not args.no_load_rng:
try:
random.setstate(state_dict["random_rng_state"])
np.random.set_state(state_dict["np_rng_state"])
torch.set_rng_state(state_dict["torch_rng_state"])
torch.cuda.set_rng_state(state_dict["cuda_rng_state"])
# Check for empty states array
if not state_dict["rng_tracker_states"]:
raise KeyError
mpu.get_cuda_rng_tracker().set_states(state_dict["rng_tracker_states"])
except KeyError:
print_rank_0(
"Unable to load rng state from checkpoint {}. "
"Specify --no-load-rng or --finetune to prevent "
"attempting to load the rng state, "
"exiting ...".format(checkpoint_name)
)
sys.exit()
# Some utilities want to load a checkpoint without distributed being initialized
# if torch.distributed.is_initialized():
# torch.distributed.barrier()
print_rank_0(
f" successfully loaded checkpoint from {args.load} "
f"at iteration {iteration}"
)
return iteration
def load_biencoder_checkpoint(
model, only_query_model=False, only_context_model=False, custom_load_path=None
):
"""
selectively load retrieval models for indexing/retrieving
from saved checkpoints
"""
args = get_args()
model = utils.unwrap_model(model)
load_path = custom_load_path if custom_load_path is not None else args.load
tracker_filename = get_checkpoint_tracker_filename(load_path)
with open(tracker_filename, "r") as f:
iteration = int(f.read().strip())
checkpoint_name = get_checkpoint_name(load_path, iteration, False)
if mpu.get_data_parallel_rank() == 0:
print(
"global rank {} is loading checkpoint {}".format(
torch.distributed.get_rank(), checkpoint_name
)
)
state_dict = torch.load(checkpoint_name, map_location="cpu")
ret_state_dict = state_dict["model"]
if only_query_model:
ret_state_dict.pop("context_model")
if only_context_model:
ret_state_dict.pop("query_model")
assert len(model) == 1
model[0].load_state_dict(ret_state_dict)
torch.distributed.barrier()
if mpu.get_data_parallel_rank() == 0:
print(" successfully loaded {}".format(checkpoint_name))
return model