change_save_meta_in_zcc_crossponding_to_paddle and fix flags (#11198)

* change_save_meta_in_zcc_crossponding_to_paddle

* delete

* fix

* change_save_meta_in_zcc_crossponding_to_paddle and fix args

* fix deadcode

---------

Co-authored-by: liufengwei02 <liufengwei02@baidu.com>

Author: xingmingyyj

paddlenlp/trainer/trainer.py

@@ -713,7 +713,7 @@ def get_metadata_file_name(path):
                     offload=self.args.load_via_cpu,
                     safetensors=True,
                     process_group=None,
-                    comm_method=self.args.comm_method,
+                    comm_method=self.args.flex_ckpt_comm_method,
                 )
             else:
                 try:
@@ -755,7 +755,7 @@ def get_metadata_file_name(path):
                         offload=self.args.load_via_cpu,
                         safetensors=True,
                         process_group=process_group,
-                        comm_method=self.args.comm_method,
+                        comm_method=self.args.flex_ckpt_comm_method,
                     )
 
                 dist.barrier()
@@ -801,7 +801,7 @@ def get_metadata_file_name(path):
                 opt_states_path,
                 aoa_config=self.args.aoa_config,
                 offload=self.args.load_via_cpu,
-                comm_method=self.args.comm_method,
+                comm_method=self.args.flex_ckpt_comm_method,
             )
 
             if not self.args.sharded_model_from_ema:
@@ -810,7 +810,7 @@ def get_metadata_file_name(path):
                     master_weights_path,
                     aoa_config=self.args.aoa_config,
                     offload=self.args.load_via_cpu,
-                    comm_method=self.args.comm_method,
+                    comm_method=self.args.flex_ckpt_comm_method,
                 )
 
             self._load_scheduler(resume_from_checkpoint)
@@ -851,7 +851,7 @@ def bf16_filtered_sharded_state_dict(sharded_state_dict):
                 model_states_path,
                 aoa_config=self.args.aoa_config,
                 offload=self.args.load_via_cpu,
-                comm_method=self.args.comm_method,
+                comm_method=self.args.flex_ckpt_comm_method,
             )
 
         if self.args.bf16 and (not self.args.ignore_load_lr_and_optim) and should_load_stage1:

paddlenlp/trainer/trainer_utils.py

@@ -26,18 +26,16 @@
 import math
 import os
 import random
-import re
 import threading
 import time
 from contextlib import contextmanager
 from enum import Enum
 from pathlib import Path
-from typing import Dict, Iterator, List, NamedTuple, Optional, Tuple, Union
+from typing import Dict, List, NamedTuple, Optional, Tuple, Union
 
 import numpy as np
 import paddle
 import paddle.distributed as dist
-from paddle import Tensor
 from paddle.distributed import fleet
 from paddle.distributed.fleet.meta_optimizers.dygraph_optimizer.dygraph_sharding_optimizer import (
     DygraphShardingOptimizer,
@@ -46,8 +44,6 @@
 from paddle.distributed.fleet.meta_parallel import get_rng_state_tracker
 from paddle.io import IterableDataset
 from paddle.optimizer.lr import LambdaDecay
-from safetensors import safe_open
-from safetensors.paddle import save_file
 
 from paddlenlp.ops import Topology
 
@@ -1449,265 +1445,3 @@ def buffer_params():
             continue
         param_list.append(param)
     optimizer._create_accumulators(paddle.base.framework.default_main_program().global_block(), param_list)
-
-
-def _parse_size(size_str: str) -> int:
-    """Parses a size string like '100MB', '2GB' into the number of bytes."""
-    size_str = size_str.upper().strip()
-    match = re.match(r"^(\d+\.?\d*)\s*(B|KB|MB|GB|TB)?$", size_str)
-    if not match:
-        raise ValueError(f"Could not parse size string: '{size_str}'")
-
-    num_str, unit = match.groups()
-    num = float(num_str)
-
-    if unit == "B" or unit is None:
-        return int(num)
-    elif unit == "KB":
-        return int(num * 1024)
-    elif unit == "MB":
-        return int(num * 1024**2)
-    elif unit == "GB":
-        return int(num * 1024**3)
-    elif unit == "TB":
-        return int(num * 1024**4)
-    else:
-        # This case should not be reached due to regex
-        raise ValueError(f"Unknown unit: '{unit}'")
-
-
-def save_full_param(
-    itr: Iterator[tuple[str, Tensor]],
-    save_dir: str,
-    rank: int,
-    moe_sharding_world_size: int,
-    max_shard_size: str = "2GB",
-    num_saver_ranks: int = 8,
-) -> None:
-    """
-    Saves model weights from an iterator into shards, supporting max shard size
-    and a limited number of saver ranks.
-
-    Only ranks less than `num_saver_ranks` will perform disk I/O. All other ranks
-    will iterate through the data to maintain synchronization but will not save.
-    The parameter distribution logic is based on `num_saver_ranks`, ensuring all
-    parameters are handled by a designated saver rank.
-
-    Args:
-        itr (Iterator): An iterator that yields (param_key, param_tensor).
-        save_dir (str): The directory where shard files will be saved.
-        rank (int): The rank of the current process.
-        moe_sharding_world_size (int): The total number of processes.
-        max_shard_size (str): The maximum size for each shard file, e.g., "500MB", "2GB".
-        num_saver_ranks (int): The number of ranks (starting from 0) that will save files.
-    """
-
-    # 1. Non-saver ranks simply consume the iterator to stay in sync.
-    if rank >= num_saver_ranks:
-        logger.info(f"[Rank {rank}/{moe_sharding_world_size}] (Non-saver) Consuming iterator for synchronization...")
-        for _ in itr:
-            pass
-        logger.info(f"[Rank {rank}/{moe_sharding_world_size}] (Non-saver) Iterator consumption complete.")
-        return
-
-    max_shard_size_bytes = _parse_size(max_shard_size)
-    logger.info(
-        f"[Rank {rank}/{moe_sharding_world_size}] (Saver) Initializing save. "
-        f"Max shard size set to: {max_shard_size_bytes / 1024**3:.2f} GB"
-    )
-
-    os.makedirs(save_dir, exist_ok=True)
-
-    current_shard_state_dict = {}
-    current_shard_size_bytes = 0
-    sub_shard_index = 0
-
-    def _save_current_shard():
-        nonlocal sub_shard_index, current_shard_state_dict, current_shard_size_bytes
-        if not current_shard_state_dict:
-            return
-
-        # Filename includes the main shard number (rank) and the sub-shard index
-        cur_rank = paddle.distributed.get_rank()
-        shard_filename = f"shard_{cur_rank}-{sub_shard_index}.safetensors"
-        save_path = os.path.join(save_dir, shard_filename)
-
-        logger.info(
-            f"[Rank {rank}/{moe_sharding_world_size}] Saving sub-shard {sub_shard_index}... "
-            f"Size: {current_shard_size_bytes / 1024**2:.2f} MB, "
-            f"Params: {len(current_shard_state_dict)}, "
-            f"Path: {save_path}"
-        )
-
-        save_file(current_shard_state_dict, save_path)
-
-        # Reset for the next shard
-        sub_shard_index += 1
-        current_shard_state_dict = {}
-        current_shard_size_bytes = 0
-
-    logger.info(f"[Rank {rank}/{moe_sharding_world_size}] Starting to process the weight iterator...")
-
-    total_size = 0
-
-    for i, (param_key, param) in enumerate(itr):
-        param_size_bytes = param.numel() * param.element_size()
-        total_size += param_size_bytes.item()
-        if i % num_saver_ranks == rank:
-            if current_shard_size_bytes > 0 and (current_shard_size_bytes + param_size_bytes > max_shard_size_bytes):
-                _save_current_shard()
-
-            current_shard_state_dict[param_key] = param
-            current_shard_size_bytes += param_size_bytes
-
-            if current_shard_size_bytes >= max_shard_size_bytes:
-                _save_current_shard()
-    _save_current_shard()
-    logger.info(f"[Rank {rank}/{moe_sharding_world_size}] (Saver) All shards saved successfully.")
-    return 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 = []
-    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)
-
-    start_idx = []
-    acc = 1
-    for files_num in all_files_num:
-        start_idx.append(acc)
-        acc += files_num
-
-    env_local_size = int(os.environ.get("PADDLE_LOCAL_SIZE", 8))
-    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
-
-    index_mapping = {}
-    if env_local_rank == 0:
-        for file in safetensor_files:
-            cur_file_index += 1
-            file_path = os.path.join(path, file)
-            new_file_name = f"model-{cur_file_index:05d}-of-{total_files_num:05d}.safetensors"
-            with safe_open(file_path, framework="np") as f:
-                for key in f.keys():
-                    index_mapping[key] = new_file_name
-            new_file_path = os.path.join(path, new_file_name)
-            os.rename(file_path, new_file_path)
-
-    index_mapping_list = []
-    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 = {}
-        index_infos["metadata"] = {}
-        index_infos["metadata"]["total_size"] = total_size
-        index_infos["weight_map"] = dict(sorted(index_mapping.items()))
-        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)
-
-
-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

paddlenlp/trainer/training_args.py

@@ -422,10 +422,14 @@ class TrainingArguments:
         load_from_hf (bool, optional):
             Whether to load a checkpoint in the HuggingFace format.
             Defaults to False.
-        comm_method (str, optional):
+        flex_ckpt_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".
+        replicate_saved_into_local (bool, optional):
+            Whether to save checkpoint replicas into local files in a distributed save/load system.
+            If set to True, replicas will be stored locally on each node/machine.
+            Defaults to False.
     """
 
     output_dir: str = field(
@@ -1169,17 +1173,22 @@ class TrainingArguments:
         metadata={"help": "Whether to load a checkpoint in the HuggingFace format."},
     )
 
-    comm_method: Optional[str] = field(
+    flex_ckpt_comm_method: Optional[str] = field(
         default="broadcast",
         metadata={
             "help": (
-                "Communication method for checkpoint resharding. "
+                "Communication method used by FlexCheckpoint for checkpoint resharding. "
                 'Choices are "send_recv", "broadcast", "multi_group_broadcast", and "grouped_send_recv". '
                 'Default is "broadcast".'
             )
         },
     )
 
+    replicate_saved_into_local: Optional[bool] = field(
+        default=False,
+        metadata={"help": "Whether to save replicas cross files in distributed save load system."},
+    )
+
     def __post_init__(self):
         world_size = paddle.distributed.get_world_size()
         if in_auto_parallel_align_mode():