Kimirock/PaddleNLP
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

【FlexCheckpoint】adapter flex checkpoint (#11188)

Browse Source 
* adapter flex checkpoint

* opt bf16 save

* fix multi moe_sharding_group load safetensors

* fix
This commit is contained in:
xingmingyyj
2025-12-08 23:07:15 +08:00
committed by GitHub
parent fa007c8f85
commit b0276d735a
5 changed files with 503 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files

290
paddlenlp/trainer/trainer.py
View File

@@ -121,9 +121,13 @@ from ..transformers.tokenizer_utils import PretrainedTokenizer
from ..utils.batch_sampler import DistributedBatchSampler as NlpDistributedBatchSampler
from ..utils.env import (
DISLORA_WEIGHTS_NAME,
EMA_STATE_DIC,
LOKR_WEIGHTS_NAME,
LORA_WEIGHTS_NAME,
MASTER_WEIGHT_DIC,
MODEL_META_NAME,
MODEL_STATE_DIC,
OPTIMIZER_STATE_DIC,
PADDLE_MASTER_WEIGHTS_INDEX_NAME,
PADDLE_OPTIMIZER_NAME,
PADDLE_PEFT_WEIGHTS_INDEX_NAME,
@@ -185,6 +189,8 @@ from .training_args import TrainingArguments
from .unified_checkpoint import UnifiedCheckpointHandler
from .utils import reshard as reshard_util
from .utils.async_save import AsyncSaver
from .utils.reshard import SHARDING_STRATEGY_V1, split_opt_state
from .utils.sharding_io import GroupGetter, to_device
try:
from .utils.zero_cost_checkpoint import (
@@ -673,6 +679,248 @@ class Trainer:
elif resume_from_checkpoint is not None:
logger.info(f"not loading ckpt :{self.args.dataset_rank}")
def _load_flex_checkpoint(self, resume_from_checkpoint):
def get_metadata_file_name(path):
files = os.listdir(path)
metadata_files = [f for f in files if f.endswith(".metadata")]
assert len(metadata_files) > 0, f"Found no metadata files in {path}"
assert len(metadata_files) == 1, f"Found multiple metadata files in {path}"
return metadata_files[0]
model_sharded_state_dict = self.model.sharded_state_dict()
hf_aoa_config = self.model._gen_aoa_config(self.model.config)
master_weights_path = os.path.join(resume_from_checkpoint, MASTER_WEIGHT_DIC)
opt_states_path = os.path.join(resume_from_checkpoint, OPTIMIZER_STATE_DIC)
model_states_path = os.path.join(resume_from_checkpoint, MODEL_STATE_DIC)
if self.args.load_from_hf:
hcg = dist.fleet.get_hybrid_communicate_group()
assert (
self.args.ignore_load_lr_and_optim
), "Loading from HuggingFace format is only allowed when learning rate and optimizer state are ignored."
try:
moe_sharding_group = hcg.get_moe_sharding_parallel_group()
except Exception:
moe_sharding_group = None
if moe_sharding_group is None or moe_sharding_group.nranks <= 1:
# when moe_sharding_group is None, we use the default process_group
logger.info(f"Loading model weights from '{resume_from_checkpoint}' in safetensors format.")
dist.load_state_dict(
model_sharded_state_dict,
resume_from_checkpoint,
aoa_config=hf_aoa_config,
offload=self.args.load_via_cpu,
safetensors=True,
process_group=None,
comm_method=self.args.comm_method,
)
else:
try:
pp_group = hcg.get_pipe_parallel_group()
if pp_group is None or pp_group.nranks < 1:
raise NotImplementedError("Only support when pp_group is not None.")
except Exception:
raise RuntimeError("Only support when pp_group is not None.")
try:
moe_group = hcg.get_expert_parallel_group()
if moe_group is None or moe_group.nranks < 1:
raise NotImplementedError("Only support when moe_group is not None.")
except Exception:
raise RuntimeError("Only support when moe_group is not None.")
moe_sharding_rank = moe_sharding_group.rank
cur_rank = dist.get_rank()
if moe_sharding_rank == 0:
moe_group_ranks = []
dist.all_gather_object(moe_group_ranks, cur_rank, group=moe_group)
pp_group_ranks = []
dist.all_gather_object(pp_group_ranks, moe_group_ranks, group=pp_group)
process_group_ranks = [rank for ranks in pp_group_ranks for rank in ranks]
else:
process_group_ranks = [0] * (pp_group.nranks * moe_group.nranks)
src_rank = hcg.get_moe_sharding_parallel_group_src_rank()
dist.broadcast_object_list(process_group_ranks, src=src_rank, group=moe_sharding_group)
assert any(process_group_ranks), "process_group_ranks should not be all 0"
logger.info(f"Creating a temporary process group with ranks: {process_group_ranks}")
process_group = dist.new_group(process_group_ranks)
if moe_sharding_rank == 0:
logger.info(f"Loading model weights from '{resume_from_checkpoint}' in safetensors format.")
# Only the first moe_sharding process is allowed to load the model weights.
dist.load_state_dict(
model_sharded_state_dict,
resume_from_checkpoint,
aoa_config=hf_aoa_config,
offload=self.args.load_via_cpu,
safetensors=True,
process_group=process_group,
comm_method=self.args.comm_method,
)
dist.barrier()
logger.info("Destroying the temporary process group.")
dist.destroy_process_group(process_group)
# The first moe_sharding group loads the model weights and then broadcasts them to all other moe_sharding groups.
logger.info(
"First shard (moe_sharding_group) has loaded safetensors weights, starting broadcast on moe_sharding_groups."
)
for param_name, param in self.model.state_dict().items():
dist.broadcast(param, src=src_rank, group=moe_sharding_group)
logger.info("Safetensors format weights have been loaded successfully.")
return
if not self.args.ignore_load_lr_and_optim:
state_dict_metadata = {}
metadata_paths = [
os.path.join(model_states_path, get_metadata_file_name(model_states_path)),
os.path.join(opt_states_path, get_metadata_file_name(opt_states_path)),
os.path.join(master_weights_path, get_metadata_file_name(master_weights_path)),
]
for metadata_file in metadata_paths:
if not os.path.exists(metadata_file):
raise FileNotFoundError(f"Metadata file not found: {metadata_file}")
metadata = paddle.load(metadata_file)
state_dict_metadata.update(metadata.state_dict_metadata)
init_optimizer(self.optimizer, model_sharded_state_dict, state_dict_metadata)
optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
opt_states = {}
master_weights = {}
for k, v in optimizer_sharded_state_dict.items():
if k.endswith(".w_0"):
master_weights[k] = v
else:
opt_states[k] = v
dist.load_state_dict(
opt_states,
opt_states_path,
aoa_config=self.args.aoa_config,
offload=self.args.load_via_cpu,
comm_method=self.args.comm_method,
)
if not self.args.sharded_model_from_ema:
dist.load_state_dict(
master_weights,
master_weights_path,
aoa_config=self.args.aoa_config,
offload=self.args.load_via_cpu,
)
self._load_scheduler(resume_from_checkpoint)
should_load_stage1 = self.args.should_load_sharding_stage1_model
if should_load_stage1 and self.args.sharded_model_from_ema:
ema_states_path = os.path.join(resume_from_checkpoint, EMA_STATE_DIC, f"{dist.get_rank()}_0.distcp")
ema_state_dict = paddle.load(ema_states_path)
ema_master_weights = ema_state_dict.pop("master_weights", None)
opt_master_weights = self.optimizer.state_dict()["master_weights"]
for k, v in opt_master_weights.items():
assert (
k in ema_master_weights
), f"{k} not in ema_master_weights, emas_master_weight keys {ema_master_weights.keys()}"
paddle.assign(ema_master_weights[k], opt_master_weights[k])
ema_state_dict = reshard_util.all_gather_state_dict(ema_state_dict, lambda x: True, self.sharding_group)
self.model.set_state_dict(ema_state_dict)
else:
dist.load_state_dict(
model_sharded_state_dict,
model_states_path,
aoa_config=self.args.aoa_config,
offload=self.args.load_via_cpu,
)
if self.args.bf16 and (not self.args.ignore_load_lr_and_optim) and should_load_stage1:
opt_state_dict = self.optimizer.state_dict()
def recover_params_from_master_weight(opt_state_dict, group):
master_weights = opt_state_dict["master_weights"]
tmp = OrderedDict()
(master_weights, tmp) = (tmp, master_weights)
# cast to before
for (k, v) in tmp.items():
name = v.name
master_weights[k] = paddle.cast(to_device(v), paddle.bfloat16).cpu()
master_weights[k].name = name
structure_name_map = {k: v.name for (k, v) in self.model.state_dict().items()}
node_model_state = reshard_util.NodeModelState(group=group)
node_model_state_tmp = reshard_util.NodeModelState(group=group)
node_model_state_tmp.add_master_weights(master_weights)
node_model_state_tmp.pack_keys(structure_name_map)
node_model_state.merge_from(node_model_state_tmp, max(group.rank, 0))
del node_model_state_tmp
sharding_strategy = reshard_util.get_sharding_strategy(self.optimizer)
logger.debug(f"sharding_strategy: {sharding_strategy}")
restore_func = (
reshard_util.sharding_v1.restore
if sharding_strategy == SHARDING_STRATEGY_V1
else reshard_util.sharding_v2.restore
)
node_model_state = restore_func(node_model_state, self.model, self.optimizer)
node_model_state.unpack_keys()
master_weights = node_model_state.master_weights
master_weights = reshard_util.all_gather_state_dict(master_weights, lambda x: True, group)
model_state_dict = self.model.state_dict()
for key, param in model_state_dict.items():
if param.name in master_weights:
logger.debug(
f"key {key}, convert master weights {param.name} shape {master_weights[param.name].shape} to param {param.name} shape{param.shape}"
)
assert (
param.shape == master_weights[param.name].shape
), f"got {param.shape} vs {master_weights[param.name].shape}"
master_weight = paddle.reshape(master_weights[param.name], param.shape)
paddle.assign(paddle.cast(to_device(master_weight), paddle.bfloat16), model_state_dict[key])
group_getter = GroupGetter(self.model)
opt_state_dict = split_opt_state(opt_state_dict, group_getter)
for gid in group_getter.get_group_ids():
sub_opt_state_dict = opt_state_dict[gid]
group = group_getter.get_group_by_id(gid)
if self.args.bf16:
recover_params_from_master_weight(sub_opt_state_dict, group)
def _save_flex_model_state(self, output_dir):
model_sharded_state_dict = self.model.sharded_state_dict()
model_state_dict_path = os.path.join(output_dir, MODEL_STATE_DIC)
os.makedirs(model_state_dict_path, exist_ok=True)
dist.save_state_dict(
model_sharded_state_dict,
model_state_dict_path,
)
def _save_flex_optimizer_state(self, output_dir):
optimizer_state_dict_path = os.path.join(output_dir, OPTIMIZER_STATE_DIC)
optimizer_states = {}
master_weights = {}
model_sharded_state_dict = self.model.sharded_state_dict()
optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
for k, v in optimizer_sharded_state_dict.items():
if k.endswith(".w_0"):
master_weights[k] = v
else:
optimizer_states[k] = v
dist.save_state_dict(
optimizer_states,
optimizer_state_dict_path,
)
master_weights_path = os.path.join(output_dir, MASTER_WEIGHT_DIC)
dist.save_state_dict(
master_weights,
master_weights_path,
)
def _load_from_checkpoint(self, resume_from_checkpoint=None):
"""load state_dict from_checkpoint, Only load model state dict.
@@ -1048,27 +1296,7 @@ class Trainer:
if delay_optimizer_creation:
self.create_optimizer_and_scheduler(num_training_steps=max_steps)
if resume_from_checkpoint is not None:
if not self.args.ignore_load_lr_and_optim:
model_sharded_state_dict = self.model.sharded_state_dict()
accessible_files = os.listdir(resume_from_checkpoint)
metadata_files = [file for file in accessible_files if file.endswith(".metadata")]
assert len(metadata_files) == 1, "Only support one metadata file now."
metadata = paddle.load(os.path.join(resume_from_checkpoint, metadata_files[0]))
state_dict_metadata = metadata.state_dict_metadata
init_optimizer(self.optimizer, model_sharded_state_dict, state_dict_metadata)
optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
sharded_state_dict = {**model_sharded_state_dict, **optimizer_sharded_state_dict}
dist.load_state_dict(
sharded_state_dict, resume_from_checkpoint, aoa_config=self.args.aoa_config
)
self._load_scheduler(resume_from_checkpoint)
else:
model_sharded_state_dict = self.model.sharded_state_dict()
sharded_state_dict = model_sharded_state_dict
dist.load_state_dict(
sharded_state_dict, resume_from_checkpoint, aoa_config=self.args.aoa_config
)
self._load_flex_checkpoint(resume_from_checkpoint)
else:
model = self._wrap_model(self.model_wrapped)
# for the rest of this function `model` is the outside model, whether it was wrapped or not
@@ -2867,8 +3095,7 @@ class Trainer:
self.save_model(output_dir)
if self.args.save_checkpoint_format == "flex_checkpoint":
model_sharded_state_dict = self.model.sharded_state_dict()
os.makedirs(output_dir, exist_ok=True)
self._save_flex_model_state(output_dir)
# Determine the new best metric / best model checkpoint
if metrics is not None and self.args.metric_for_best_model is not None:
@@ -2932,11 +3159,7 @@ class Trainer:
)
else:
if self.args.save_checkpoint_format == "flex_checkpoint":
optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
dist.save_state_dict(
{**model_sharded_state_dict, **optimizer_sharded_state_dict},
output_dir,
)
self._save_flex_optimizer_state(output_dir)
if self.args.should_save:
if self.tokenizer is not None and self.args.save_tokenizer:
self.tokenizer.save_pretrained(output_dir)
@@ -2992,11 +3215,7 @@ class Trainer:
signal_dir,
)
elif self.args.save_checkpoint_format == "flex_checkpoint":
optimizer_sharded_state_dict = self.optimizer.sharded_state_dict(model_sharded_state_dict)
dist.save_state_dict(
{**model_sharded_state_dict, **optimizer_sharded_state_dict},
output_dir,
)
self._save_flex_optimizer_state(output_dir)
if self.args.should_save:
if self.tokenizer is not None and self.args.save_tokenizer:
self.tokenizer.save_pretrained(output_dir)
@@ -3039,10 +3258,7 @@ class Trainer:
self._offload_optimizer()
else:
if self.args.save_checkpoint_format == "flex_checkpoint":
dist.save_state_dict(
model_sharded_state_dict,
output_dir,
)
self._save_flex_model_state(output_dir)
if self.args.should_save:
if self.tokenizer is not None and self.args.save_tokenizer:
self.tokenizer.save_pretrained(output_dir)

139
paddlenlp/trainer/trainer_utils.py
View File

@@ -1567,34 +1567,31 @@ def save_full_param(
return total_size
def replace_name_and_gen_index(path, cur_rank_total_size):
def replace_name_and_gen_index(path, total_size):
index_mapping = {}
cur_rank = paddle.distributed.get_rank()
safetensor_files = [fname for fname in os.listdir(path) if fname.endswith(".safetensors")]
files_num = len(safetensor_files)
all_files_num = []
paddle.distributed.all_gather_object(all_files_num, files_num)
if paddle.distributed.get_world_size() > 1:
paddle.distributed.all_gather_object(all_files_num, files_num)
else:
all_files_num.append(files_num)
total_files_num = sum(all_files_num)
all_sizes = []
paddle.distributed.all_gather_object(all_sizes, cur_rank_total_size)
total_size = sum(all_sizes)
start_idx = []
acc = 1
for files_num in all_files_num:
start_idx.append(acc)
acc += files_num
env_local_rank = int(os.environ.get("PADDLE_RANK_IN_NODE", -1))
env_local_size = int(os.environ.get("PADDLE_LOCAL_SIZE", 8))
assert env_local_rank >= 0
env_local_rank = dist.get_rank() % env_local_size
assert env_local_rank >= 0, f"expected positive local rank, got {env_local_rank}"
cur_file_index = start_idx[cur_rank] // env_local_size
total_files_num = total_files_num // env_local_size
total_size = total_size // env_local_size
index_mapping = {}
if env_local_rank == 0:
for file in safetensor_files:
@@ -1608,11 +1605,19 @@ def replace_name_and_gen_index(path, cur_rank_total_size):
os.rename(file_path, new_file_path)
index_mapping_list = []
paddle.distributed.all_gather_object(index_mapping_list, index_mapping)
if paddle.distributed.get_world_size() > 1:
paddle.distributed.all_gather_object(index_mapping_list, index_mapping)
else:
index_mapping_list.append(index_mapping)
index_mapping = {}
for mapping in index_mapping_list:
index_mapping.update(mapping)
# Save signal file for each card
saved_signal_path = os.path.join(path, f"saved_signal_{dist.get_rank()}")
with open(saved_signal_path, mode="w+") as f:
f.write("1")
if env_local_rank == 0:
index_file_name = "model.safetensors.index.json"
index_infos = {}
@@ -1622,35 +1627,87 @@ def replace_name_and_gen_index(path, cur_rank_total_size):
with open(os.path.join(path, index_file_name), "w") as f:
json.dump(index_infos, f, indent=4)
# For PDC signal
if strtobool(os.getenv("FLAG_LLM_PDC", "False")):
for i in range(paddle.distributed.get_world_size()):
saved_signal_path = os.path.join(path, f".model_weights.done.{i}")
paddle.save(i, saved_signal_path)
def save_hf_checkpoint(
model,
aoa_config,
h_group,
v_group,
num_splits,
shard_idx,
path,
):
memory_growth_threshold = 8 * (2**30)
itr = model.full(
aoa_config=aoa_config,
h_group=h_group,
v_group=v_group,
num_splits=num_splits,
shard_idx=shard_idx,
memory_growth_threshold=memory_growth_threshold,
)
num_saver_ranks = h_group.nranks * v_group.nranks
rank = h_group.rank + v_group.rank * h_group.nranks
total_saved_size = save_full_param(
itr=itr,
save_dir=path,
rank=rank,
moe_sharding_world_size=num_saver_ranks,
max_shard_size="16GB",
num_saver_ranks=num_saver_ranks,
)
paddle.distributed.barrier()
replace_name_and_gen_index(path, total_saved_size)
class HFFormatFullParamSaver:
def __init__(
self,
model,
aoa_config,
h_group=None,
v_group=None,
num_splits=None,
shard_idx=None,
saved_in_one_node=False,
memory_growth_threshold=8 * (2**30),
):
self.model = model
self.aoa_config = aoa_config
self.h_group = h_group
self.v_group = v_group
self.num_splits = num_splits
self.shard_idx = shard_idx
self.saved_in_one_node = saved_in_one_node
self.memory_growth_threshold = memory_growth_threshold
self.determin_saver_based_group()
def get_full_param_iter(self):
assert (self.v_group and self.h_group) or not (
self.v_group or self.h_group
), f"both h_group and v_group are provided or none of them, but got {self.v_group} and {self.h_group}"
if self.v_group and self.h_group:
assert self.shard_idx is not None, "expected shard_idx is not None"
assert self.num_splits is not None, "expected num_splits is not None"
param_iter = self.model.full(
aoa_config=self.aoa_config,
h_group=self.h_group,
v_group=self.v_group,
num_splits=self.num_splits,
shard_idx=self.shard_idx,
memory_growth_threshold=self.memory_growth_threshold,
)
else:
param_iter = self.model.full(
aoa_config=self.aoa_config, memory_growth_threshold=self.memory_growth_threshold
)
return param_iter
def determin_saver_based_group(self):
self.num_saver_ranks = paddle.distributed.get_world_size()
self.rank = paddle.distributed.get_rank()
if self.h_group and self.v_group:
self.num_saver_ranks = self.h_group.nranks * self.v_group.nranks
self.rank = self.h_group.rank + self.v_group.rank * self.h_group.nranks
if self.saved_in_one_node:
local_world_size = int(os.environ.get("PADDLE_LOCAL_SIZE", 8))
self.num_saver_ranks = min(local_world_size, self.num_saver_ranks)
def save_checkpoint(self, path, max_shard_size="16GB"):
total_saved_size = save_full_param(
itr=self.get_full_param_iter(),
save_dir=path,
rank=self.rank,
moe_sharding_world_size=self.num_saver_ranks,
max_shard_size=max_shard_size,
num_saver_ranks=self.num_saver_ranks,
)
if paddle.distributed.get_world_size() > 1:
paddle.distributed.barrier()
# TODO(): fix total size
all_sizes = []
if paddle.distributed.get_world_size() > 1:
paddle.distributed.all_gather_object(all_sizes, total_saved_size)
else:
all_sizes.append(total_saved_size)
total_size = sum(all_sizes)
replace_name_and_gen_index(path, total_size)
return total_saved_size

34
paddlenlp/trainer/training_args.py
View File

@@ -415,6 +415,17 @@ class TrainingArguments:
The AoA configuration of FlexCheckpoint, used to describe the mapping between model weights and the checkpoint content. Default is None.
save_hf_steps (`int`, *optional*, defaults to 500):
Number of updates steps before two huggingface checkpoint saves if `save_strategy="steps"`.
load_via_cpu (bool, optional):
Whether to load checkpoint data into CPU memory first before transferring to GPU.
This helps mitigate GPU memory shortage by staging data on the CPU and only moving required parts to the GPU on demand during communication.
Defaults to False.
load_from_hf (bool, optional):
Whether to load a checkpoint in the HuggingFace format.
Defaults to False.
comm_method (str, optional):
Communication method used for checkpoint resharding.
Choices are "send_recv", "broadcast", "multi_group_broadcast", and "grouped_send_recv".
Defaults to "broadcast".
"""
output_dir: str = field(
@@ -1146,6 +1157,29 @@ class TrainingArguments:
save_hf_steps: int = field(default=-1, metadata={"help": "Save huggingface checkpoint every X updates steps."})
load_via_cpu: Optional[bool] = field(
default=False,
metadata={
"help": "If True, loads checkpoint data to CPU first, then transfers required parts to GPU on demand to reduce GPU memory usage. Defaults to False."
},
)
load_from_hf: Optional[bool] = field(
default=False,
metadata={"help": "Whether to load a checkpoint in the HuggingFace format."},
)
comm_method: Optional[str] = field(
default="broadcast",
metadata={
"help": (
"Communication method for checkpoint resharding. "
'Choices are "send_recv", "broadcast", "multi_group_broadcast", and "grouped_send_recv". '
'Default is "broadcast".'
)
},
)
def __post_init__(self):
world_size = paddle.distributed.get_world_size()
if in_auto_parallel_align_mode():

151
paddlenlp/trainer/utils/zero_cost_checkpoint.py
View File

@@ -14,6 +14,7 @@
import atexit
import copy
import functools
import hashlib
import json
import multiprocessing
@@ -51,9 +52,16 @@ from paddle.incubate.tensor.manipulation import (
)
from paddle.optimizer.fusion_utils import FusionStorageHelper
from paddlenlp.trainer.utils.sharding_io import GroupGetter
from ...transformers.model_utils import unwrap_optimizer
from . import reshard as reshard_util
from .reshard import SHARDING_STRATEGY_V1
from .reshard import (
SHARDING_STRATEGY_V1,
merge_model_state,
split_model_state,
split_opt_state,
)
try:
from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
@@ -123,6 +131,65 @@ def showmem(msg):
)
# the funciotn that accept state dict as input can be decorated with this function
def sharded_state_dict_compatibility(func, *, return_sharded_state_dict=False):
@functools.wraps(func)
def wrapper(*args, **kwargs):
def should_convert(maybe_sharded_state_dict):
all_shared_weights = all(isinstance(value, ShardedWeight) for value in maybe_sharded_state_dict.values())
any_shared_weights = any(isinstance(value, ShardedWeight) for value in maybe_sharded_state_dict.values())
logger.debug(f"all sharded weight {all_shared_weights}, any shared weight {any_shared_weights}")
if not any_shared_weights:
logger.debug("this is not a sharded state dict, no need to convert.")
return False
if any_shared_weights and (not all_shared_weights):
logger.debug("this is a mixed state dict(normal and sharded), not support to convert.")
return False
logger.debug("this is a sharded state dict, will convert it to local tensor dict.")
return True
original_sharded_state_dict = {}
# process args
new_args = list(args)
for idx, arg in enumerate(new_args):
if not isinstance(arg, dict):
continue
if should_convert(arg):
local_tensor_state_dict = {}
for k, v in arg.items():
local_tensor_state_dict[k] = v.local_tensor
original_sharded_state_dict.update(arg)
new_args[idx] = local_tensor_state_dict
# process kwargs
for key, value in kwargs.items():
if not isinstance(value, dict):
continue
if should_convert(value):
local_tensor_state_dict = {}
for k, v in value.items():
local_tensor_state_dict[k] = v.local_tensor
kwargs[key] = local_tensor_state_dict
original_sharded_state_dict.update(value)
# original function
result = func(*new_args, **kwargs)
if return_sharded_state_dict:
assert isinstance(result, dict), f"expected dict, but got {type(result)}"
for k, v in result.items():
sharded_sharded_weight = original_sharded_state_dict[k]
sharded_sharded_weight.local_tensor = v
result[k] = sharded_sharded_weight
return result
return wrapper
@sharded_state_dict_compatibility
def get_fused_param_mappings(optimizer, manipulated_state_dict):
param_mappings = {}
ipc_meta_mappings = {}
@@ -1415,50 +1482,56 @@ class ZeroCostCheckpointCallbackFcBased(ZeroCostCheckpointCallback):
self.sharding_group = self.hcg.get_sharding_parallel_group()
def _manipulate_state_dict_and_config(self, model_to_save, optimizer):
# TODO(): at now accracy is wrong, need to fix later.
return model_to_save.sharded_state_dict()
# group_getter = GroupGetter(model_to_save)
# gids = group_getter.get_group_ids()
# from paddleformers.trainer.utils.sharding_io import filter_sharded_params, exclude_parameters_in_state_dict
# return model_to_save.sharded_state_dict()
# filter_sharded_params = sharded_state_dict_compatibility(filter_sharded_params, return_sharded_state_dict=True)
# exclude_parameters_in_state_dict = sharded_state_dict_compatibility(exclude_parameters_in_state_dict, return_sharded_state_dict=True)
group_getter = GroupGetter(model_to_save)
gids = group_getter.get_group_ids()
from paddlenlp.trainer.utils.sharding_io import (
exclude_parameters_in_state_dict,
filter_sharded_params,
)
# state_dict = model_to_save.sharded_state_dict()
# if self.args.should_save_sharding_stage1_model:
# state_dict = split_model_state(state_dict, group_getter)
# for gid in gids:
# state_dict[gid] = filter_sharded_params(
# state_dict.get(gid, {}),
# optimizer,
# self.sharding_group,
# self.args.save_sharding_stage1_model_include_freeze_params,
# )
# state_dict = merge_model_state(state_dict)
filter_sharded_params = sharded_state_dict_compatibility(filter_sharded_params, return_sharded_state_dict=True)
exclude_parameters_in_state_dict = sharded_state_dict_compatibility(
exclude_parameters_in_state_dict, return_sharded_state_dict=True
)
# if self.args.bf16 and self.args.should_save_sharding_stage1_model:
# param_names_in_master_weights = []
# optimzier_state_dict = optimizer.state_dict()
# optimzier_state_dict = split_opt_state(optimzier_state_dict, group_getter)
# state_dict = split_model_state(state_dict, group_getter)
# for gid in gids:
# sub_opt_state = optimzier_state_dict.get(gid, {})
# param_names_in_master_weights = list(sub_opt_state.get("master_weights", {}).keys())
# state_dict[gid] = exclude_parameters_in_state_dict(
# state_dict.get(gid, {}),
# param_names_in_master_weights,
# group_getter.get_group_by_id(gid),
# )
# state_dict = merge_model_state(state_dict)
# logger.info(
# "param_names_in_master_weights len:{}, bf16 state_dict len:{}, :{}".format(
# len(param_names_in_master_weights), len(state_dict), state_dict.keys()
# )
# )
state_dict = model_to_save.sharded_state_dict()
if self.args.should_save_sharding_stage1_model:
state_dict = split_model_state(state_dict, group_getter)
for gid in gids:
state_dict[gid] = filter_sharded_params(
state_dict.get(gid, {}),
optimizer,
self.sharding_group,
self.args.save_sharding_stage1_model_include_freeze_params,
)
state_dict = merge_model_state(state_dict)
# return state_dict
if self.args.bf16 and self.args.should_save_sharding_stage1_model:
param_names_in_master_weights = []
optimzier_state_dict = optimizer.state_dict()
optimzier_state_dict = split_opt_state(optimzier_state_dict, group_getter)
state_dict = split_model_state(state_dict, group_getter)
for gid in gids:
sub_opt_state = optimzier_state_dict.get(gid, {})
param_names_in_master_weights = list(sub_opt_state.get("master_weights", {}).keys())
state_dict[gid] = exclude_parameters_in_state_dict(
state_dict.get(gid, {}),
param_names_in_master_weights,
group_getter.get_group_by_id(gid),
)
state_dict = merge_model_state(state_dict)
logger.info(
"param_names_in_master_weights len:{}, bf16 state_dict len:{}, :{}".format(
len(param_names_in_master_weights), len(state_dict), state_dict.keys()
)
)
return state_dict
def _cache_meta_for_sharded_save(self, model, optimizer):
# TODO(): fix later.
logger.info("Start caching metas for sharded save...")
(self.manipulated_state_dict) = self._manipulate_state_dict_and_config(model, optimizer)

6
paddlenlp/utils/env.py
View File

@@ -167,5 +167,11 @@ else:
USE_FAST_TOKENIZER: bool = _get_bool_env("USE_FAST_TOKENIZER", "false")
PREFILL_USE_SAGE_ATTN: bool = _get_bool_env("PREFILL_USE_SAGE_ATTN", "false")
# Folder names for FlexCheckpoint
MODEL_STATE_DIC = "model_state"
OPTIMIZER_STATE_DIC = "optimizer_state"
MASTER_WEIGHT_DIC = "master_weight"
EMA_STATE_DIC = "ema_state"
# hf checkpoint dir name
PREFIX_HF_CHECKPOINT_DIR = "hf_checkpoint"