[dataset] Fix multinode packing (#4402)

Author: Jintao-Huang

docs/source/Instruction/Megatron-SWIFT训练.md

@@ -296,6 +296,7 @@ Megatron训练参数继承自Megatron参数和基本参数。基本参数的内
 
 - add_version: 在`save`上额外增加目录`'<版本号>-<时间戳>'`防止权重覆盖，默认为True。
 - 🔥packing: 是否使用序列packing，默认为False。
+- 🔥packing_cache: 指定 packing 缓存目录。默认值为`None`，表示缓存将存储在环境变量 `$MODELSCOPE_CACHE`所指定的路径下。在跨节点使用 packing 功能时，需确保所有节点的 packing 缓存路径共享且一致。你可以通过设置`MODELSCOPE_CACHE`环境变量，或在命令行中添加 `--packing_cache <shared_path>`参数来实现这一要求。
 - 🔥streaming: 流式读取并处理数据集，默认False。通常在处理大型数据集时，设置为True。更多流式的参数查看命令行参数文档。
 - lazy_tokenize: 默认为False。若该参数设置为False，则在训练之前对所有的数据集样本进行tokenize（这可以避免在训练中出现报错）；设置为True，则在训练中对数据集进行tokenize（这可以节约内存）。
 - max_epochs: 训练到`max_epochs`时强制退出训练，并对权重进行验证和保存。该参数在使用流式数据集时很有用。默认为None。

docs/source/Instruction/命令行参数.md

@@ -354,6 +354,7 @@ Vera使用`target_modules`, `target_regex`, `modules_to_save`三个参数.
 - 🔥packing: 是否使用序列packing提升计算效率，默认为False。当前支持`swift pt/sft`。
   - 注意：使用packing请结合`--attn_impl flash_attn`使用且"transformers>=4.44"，具体查看[该PR](https://github.com/huggingface/transformers/pull/31629)。
   - 支持的多模态模型参考：https://github.com/modelscope/ms-swift/blob/main/examples/train/packing/qwen2_5_vl.sh
+- packing_cache: 指定 packing 缓存目录。默认值为`None`，表示缓存将存储在环境变量 `$MODELSCOPE_CACHE`所指定的路径下。在跨节点使用 packing 功能时，需确保所有节点的 packing 缓存路径共享且一致。你可以通过设置`MODELSCOPE_CACHE`环境变量，或在命令行中添加 `--packing_cache <shared_path>`参数来实现这一要求。
 - 🔥lazy_tokenize: 是否使用lazy_tokenize。若该参数设置为False，则在训练之前对所有的数据集样本进行tokenize（多模态模型则包括从磁盘中读取图片）。该参数在LLM训练中默认设置为False，而MLLM训练默认为True，节约内存。
 - acc_strategy: 训练和验证时计算acc的策略。可选为`seq`和`token`级别的acc，默认为`token`。
 - max_new_tokens: 覆盖生成参数。predict_with_generate=True时的最大生成token数量，默认64。

docs/source_en/Instruction/Command-line-parameters.md

@@ -363,6 +363,7 @@ Training arguments include the [base arguments](#base-arguments), [Seq2SeqTraine
 - 🔥packing: Whether to use sequence packing to improve computational efficiency. The default value is False. Currently supports `swift pt/sft`.
   - Note: When using packing, please combine it with `--attn_impl flash_attn` and ensure "transformers>=4.44". For details, see [this PR](https://github.com/huggingface/transformers/pull/31629).
   - Supported multimodal models reference: https://github.com/modelscope/ms-swift/blob/main/examples/train/packing/qwen2_5_vl.sh
+- packing_cache: Specifies the directory for packing cache. The default value is `None`, which means the cache will be stored in the path defined by the environment variable `$MODELSCOPE_CACHE`. When using the packing feature across multiple nodes, ensure that all nodes share the same packing cache directory. You can achieve this by setting the `MODELSCOPE_CACHE` environment variable or by adding the `--packing_cache <shared_path>` argument in the command line.
 - 🔥lazy_tokenize: Whether to use lazy tokenization. If set to False, all dataset samples are tokenized before training (for multimodal models, this includes reading images from disk). This parameter defaults to False for LLM training, and True for MLLM training, to save memory.
 - acc_strategy: Strategy for calculating accuracy during training and validation. Options are `seq`-level and `token`-level accuracy, with `token` as the default.
 - max_new_tokens: Generation parameter override. The maximum number of tokens to generate when `predict_with_generate=True`, defaulting to 64.

docs/source_en/Instruction/Megatron-SWIFT-Training.md

@@ -307,6 +307,7 @@ Megatron training parameters inherit from Megatron parameters and basic paramete
 
 - add_version: Adds a directory `<version>-<timestamp>` to `save` to prevent overwriting weights, default is True.
 - 🔥packing: Whether to use sequence packing, defaults to False.
+- 🔥packing_cache: Specifies the directory for packing cache. The default value is `None`, which means the cache will be stored in the path defined by the environment variable `$MODELSCOPE_CACHE`. When using the packing feature across multiple nodes, ensure that all nodes share the same packing cache directory. You can achieve this by setting the `MODELSCOPE_CACHE` environment variable or by adding the `--packing_cache <shared_path>` argument in the command line.
 - 🔥streaming: Stream reading and processing of the dataset, default is False. It is typically set to True when handling large datasets. For more information on streaming parameters, refer to the command-line parameters documentation.
 - lazy_tokenize: Default is False. If this parameter is set to False, all dataset samples are tokenized before training (this avoids errors during training); if set to True, tokenization occurs during training (this saves memory).
 - max_epochs: Forces the training to exit after reaching `max_epochs`, and performs validation and saving of the model weights. This parameter is especially useful when using a streaming dataset. Default is None.

swift/llm/argument/base_args/base_args.py

@@ -9,8 +9,8 @@
 from swift.llm import Processor, Template, get_model_tokenizer, get_template, load_by_unsloth, safe_snapshot_download
 from swift.llm.utils import get_ckpt_dir
 from swift.plugin import extra_tuners
-from swift.utils import (check_json_format, get_dist_setting, get_logger, import_external_file, is_dist, is_master,
-                         set_device, use_hf_hub)
+from swift.utils import (check_json_format, check_shared_disk, get_dist_setting, get_logger, import_external_file,
+                         is_dist, is_master, set_device, use_hf_hub)
 from .data_args import DataArguments
 from .generation_args import GenerationArguments
 from .model_args import ModelArguments
@@ -78,12 +78,16 @@ class BaseArguments(CompatArguments, GenerationArguments, QuantizeArguments, Dat
     model_kwargs: Optional[Union[dict, str]] = None
     load_args: bool = True
     load_data_args: bool = False
-
+    # dataset
+    packing: bool = False
+    packing_cache: Optional[str] = None
+    custom_register_path: List[str] = field(default_factory=list)  # .py
+    # hub
     use_hf: bool = False
     # None: use env var `MODELSCOPE_API_TOKEN`
     hub_token: Optional[str] = field(
         default=None, metadata={'help': 'SDK token can be found in https://modelscope.cn/my/myaccesstoken'})
-    custom_register_path: List[str] = field(default_factory=list)  # .py
+    # dist
     ddp_timeout: int = 18000000
     ddp_backend: Optional[str] = None
 
@@ -128,6 +132,17 @@ def _init_adapters(self):
             safe_snapshot_download(adapter, use_hf=self.use_hf, hub_token=self.hub_token) for adapter in self.adapters
         ]
 
+    def _check_packing(self):
+        if not self.packing:
+            return
+        error = ValueError('When using the packing feature across multiple nodes, ensure that all nodes share '
+                           'the same packing cache directory. You can achieve this by setting the '
+                           '`MODELSCOPE_CACHE` environment variable or by adding the `--packing_cache '
+                           '<shared_path>` argument in the command line.')
+        check_shared_disk(error, self.packing_cache)
+        if self.packing_cache:
+            os.environ['PACKING_CACHE'] = self.packing_cache
+
     def __post_init__(self):
         if self.use_hf or use_hf_hub():
             self.use_hf = True

swift/llm/argument/train_args.py

@@ -108,9 +108,6 @@ class TrainArguments(SwanlabArguments, TunerArguments, BaseArguments, Seq2SeqTra
     add_version: bool = True
     resume_only_model: bool = False
     create_checkpoint_symlink: bool = False
-
-    # dataset
-    packing: bool = False
     lazy_tokenize: Optional[bool] = None
 
     # plugin
@@ -174,8 +171,8 @@ def __post_init__(self) -> None:
             self.accelerator_config = {'dispatch_batches': False}
         self.training_args = TrainerFactory.get_training_args(self)
         self.training_args.remove_unused_columns = False
-
         self._add_version()
+        self._check_packing()
 
         if 'swanlab' in self.report_to:
             self._init_swanlab()

swift/llm/dataset/utils.py

@@ -240,7 +240,8 @@ class IndexedDataset(Dataset):
 
     @staticmethod
     def get_cache_dir(dataset_name: str):
-        cache_dir = os.path.join(get_cache_dir(), 'tmp', dataset_name)
+        cache_dir = os.getenv('PACKING_CACHE') or os.path.join(get_cache_dir(), 'tmp')
+        cache_dir = os.path.join(cache_dir, dataset_name)
         os.makedirs(cache_dir, exist_ok=True)
         assert dataset_name is not None, f'dataset_name: {dataset_name}'
         return cache_dir

swift/megatron/argument/train_args.py

@@ -16,7 +16,6 @@ class MegatronTrainArguments(MegatronArguments, BaseArguments):
     add_version: bool = True
     # dataset
     lazy_tokenize: bool = False
-    packing: bool = False
 
     def init_model_args(self, config):
         self.megatron_model_meta = get_megatron_model_meta(self.model_type)
@@ -43,6 +42,7 @@ def __post_init__(self):
         self.load = to_abspath(self.load, check_path_exist=True)
         BaseArguments.__post_init__(self)
         self._init_save()
+        self._check_packing()
         self.seq_length = self.seq_length or self.max_length
         if self.streaming:
             self.dataloader_type = 'external'

swift/utils/__init__.py

@@ -10,10 +10,11 @@
 from .logger import get_logger
 from .np_utils import get_seed, stat_array, transform_jsonl_to_df
 from .tb_utils import TB_COLOR, TB_COLOR_SMOOTH, plot_images, read_tensorboard_file, tensorboard_smoothing
-from .torch_utils import (Serializer, activate_parameters, find_all_linears, find_embedding, find_layers, find_norm,
-                          freeze_parameters, gc_collect, get_current_device, get_device, get_device_count,
-                          get_model_parameter_info, get_n_params_grads, init_process_group, safe_ddp_context,
-                          seed_worker, set_default_ddp_config, set_device, show_layers, time_synchronize)
+from .torch_utils import (Serializer, activate_parameters, check_shared_disk, find_all_linears, find_embedding,
+                          find_layers, find_norm, freeze_parameters, gc_collect, get_current_device, get_device,
+                          get_device_count, get_model_parameter_info, get_n_params_grads, init_process_group,
+                          safe_ddp_context, seed_worker, set_default_ddp_config, set_device, show_layers,
+                          time_synchronize)
 from .utils import (add_version_to_work_dir, check_json_format, copy_files_by_pattern, deep_getattr, find_free_port,
                     get_env_args, import_external_file, lower_bound, parse_args, patch_getattr, read_multi_line,
                     seed_everything, split_list, subprocess_run, test_time, upper_bound)

swift/utils/torch_utils.py

@@ -21,7 +21,7 @@
 from transformers.trainer_utils import set_seed
 from transformers.utils import is_torch_cuda_available, is_torch_mps_available, is_torch_npu_available
 
-from .env import get_dist_setting, is_dist, is_dist_ta, is_local_master, is_master
+from .env import get_dist_setting, get_node_setting, is_dist, is_dist_ta, is_local_master, is_master
 from .logger import get_logger
 from .utils import deep_getattr
 
@@ -401,3 +401,33 @@ def seed_worker(worker_id: int, num_workers: int, rank: int):
     init_seed = torch.initial_seed() % 2**32
     worker_seed = num_workers * rank + init_seed
     set_seed(worker_seed)
+
+
+def check_shared_disk(error, cache_dir: Optional[str] = None):
+    nnodes = get_node_setting()[1]
+    if nnodes <= 1:
+        return True
+    assert dist.is_initialized()
+    if cache_dir is None:
+        cache_dir = os.path.join(get_cache_dir(), 'tmp')
+    os.makedirs(cache_dir, exist_ok=True)
+    tmp_path = os.path.join(cache_dir, 'check_shared_disk.tmp')
+    is_shared_disk = True
+    with safe_ddp_context(None, True):
+        if os.path.exists(tmp_path):
+            os.remove(tmp_path)
+    try:
+        with safe_ddp_context(None, True):
+            if is_master():
+                with open(tmp_path, 'w'):
+                    pass
+            else:
+                if not os.path.exists(tmp_path):
+                    is_shared_disk = False
+    finally:
+        if is_master() and os.path.exists(tmp_path):
+            os.remove(tmp_path)
+    shared_state = [None] * dist.get_world_size()
+    dist.all_gather_object(shared_state, is_shared_disk)
+    if not all(shared_state):
+        raise error