merge easycontext

.gitignore

@@ -1,3 +1,11 @@
+# local directories to ignore
+/output/
+*.json
+*llmtuner*
+/wandb/
+/examples/
+/data/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

data/dataset_info.json

@@ -8,6 +8,12 @@
   "alpaca_zh_demo": {
     "file_name": "alpaca_zh_demo.json"
   },
+  "long_sft_32k": {
+    "file_name": "sample_long_sft_32k_48M.json"
+  },
+  "long_sft_128k": {
+    "file_name": "sample_long_sft_128k.parquet"
+  },
   "glaive_toolcall_en_demo": {
     "file_name": "glaive_toolcall_en_demo.json",
     "formatting": "sharegpt",
@@ -551,4 +557,4 @@
     },
     "folder": "python"
   }
-}
+}

examples/accelerate/ds_multi_nodes.yaml

@@ -0,0 +1,15 @@
+debug: false
+deepspeed_config:
+  deepspeed_config_file: examples/deepspeed/ds_z3_offload_config.json
+  deepspeed_multinode_launcher: standard
+  zero3_init_flag: true
+distributed_type: DEEPSPEED
+num_processes: 16
+downcast_bf16: 'no'
+main_training_function: main
+rdzv_backend: c10d
+same_network: false
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false

examples/deepspeed/ds_z3_offload_config.json

@@ -17,13 +17,8 @@
   },
   "zero_optimization": {
     "stage": 3,
-    "offload_optimizer": {
-      "device": "cpu",
-      "pin_memory": true
-    },
     "offload_param": {
-      "device": "cpu",
-      "pin_memory": true
+      "device": "cpu"
     },
     "overlap_comm": true,
     "contiguous_gradients": true,
@@ -34,5 +29,6 @@
     "stage3_max_live_parameters": 1e9,
     "stage3_max_reuse_distance": 1e9,
     "stage3_gather_16bit_weights_on_model_save": true
-  }
-}
+  },
+  "steps_per_print":1
+}

src/llamafactory/data/__init__.py

@@ -1,4 +1,4 @@
-from .collator import KTODataCollatorWithPadding, PairwiseDataCollatorWithPadding
+from .collator import KTODataCollatorWithPadding, PairwiseDataCollatorWithPadding, SeqParallelDataCollatorForLanguageModeling
 from .data_utils import Role, split_dataset
 from .loader import get_dataset
 from .template import TEMPLATES, Template, get_template_and_fix_tokenizer
@@ -7,6 +7,7 @@
 __all__ = [
     "KTODataCollatorWithPadding",
     "PairwiseDataCollatorWithPadding",
+    "SeqParallelDataCollatorForLanguageModeling",
     "Role",
     "split_dataset",
     "get_dataset",

src/llamafactory/data/collator.py

@@ -2,8 +2,9 @@
 from typing import Any, Dict, Sequence
 
 import torch
-from transformers import DataCollatorForSeq2Seq
-
+from transformers import DataCollatorForSeq2Seq, DataCollatorForLanguageModeling
+from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
+from llamafactory.easy_context import prepare_seq_parallel_sft_inputs
 
 @dataclass
 class PairwiseDataCollatorWithPadding(DataCollatorForSeq2Seq):
@@ -79,3 +80,85 @@ def __call__(self, features: Sequence[Dict[str, Any]]) -> Dict[str, torch.Tensor
 
         batch["kto_tags"] = torch.tensor(kto_tags)
         return batch
+
+@dataclass
+class SeqParallelDataCollator(DataCollatorForSeq2Seq):
+    r"""
+    Data collator for sequence parallel in supervised finetune(sft) stage.
+    """
+    seq_algo: str = "data_parallel",
+    sp_size: int = -1
+    rank: int = 0
+    world_size: int = 8
+    device: Optional[Any] = None
+
+    def __call__(self, features: Sequence[Dict[str, Any]], return_tensors=None) -> Dict[str, torch.Tensor]:
+        batch = super().__call__(features, return_tensors)
+        if self.seq_algo == "data_parallel":
+            return batch
+        input_ids = batch["input_ids"]
+        attention_mask = batch["attention_mask"]
+        labels = batch["labels"]
+        world_size = self.world_size
+        sp_rank = self.rank
+        if self.sp_size != -1:
+            dp_rank = self.rank // self.sp_size
+            sp_rank = self.rank % self.sp_size
+            world_size = self.sp_size
+            bs = len(input_ids)
+            dp_size = self.world_size // self.sp_size
+            group_bs = bs // dp_size
+            input_ids = input_ids[dp_rank * group_bs: (dp_rank + 1) * group_bs]
+            attention_mask = attention_mask[dp_rank * group_bs: (dp_rank + 1) * group_bs]
+            labels = labels[dp_rank * group_bs: (dp_rank + 1) * group_bs]
+        batch = prepare_seq_parallel_sft_inputs(self.seq_algo,
+                                                input_ids=input_ids,
+                                                attention_mask=attention_mask,
+                                                position_ids=None,
+                                                labels=labels,
+                                                rank=sp_rank,
+                                                world_size=world_size,
+                                                device=self.device)
+        return batch
+
+
+@dataclass
+class SeqParallelDataCollatorForLanguageModeling(DataCollatorForLanguageModeling):
+    r"""
+    Data collator for sequence parallel in pretrain(pt) stage.
+    Reuse the sequence parallel distributing function for sft stage.
+    """
+    seq_algo: str = "data_parallel"
+    sp_size: int = -1
+    rank: int = 0
+    world_size: int = 8
+    device: Optional[Any] = None
+
+    def __call__(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> Dict[str, Any]:
+        batch = super().__call__(examples)
+        if self.seq_algo == "data_parallel":
+            return batch
+        input_ids = batch["input_ids"]
+        attention_mask = batch["attention_mask"]
+        labels = batch["labels"]
+        world_size = self.world_size
+        sp_rank = self.rank
+        if self.sp_size != -1:
+            dp_rank = self.rank // self.sp_size
+            sp_rank = self.rank % self.sp_size
+            world_size = self.sp_size
+            bs = len(input_ids)
+            dp_size = self.world_size // self.sp_size
+            group_bs = bs // dp_size
+            input_ids = input_ids[dp_rank * group_bs: (dp_rank + 1) * group_bs]
+            attention_mask = attention_mask[dp_rank * group_bs: (dp_rank + 1) * group_bs]
+            labels = labels[dp_rank * group_bs: (dp_rank + 1) * group_bs]
+        batch = prepare_seq_parallel_sft_inputs(self.seq_algo,
+                                                input_ids=input_ids,
+                                                attention_mask=attention_mask,
+                                                position_ids=None,
+                                                labels=labels,
+                                                rank=sp_rank,
+                                                world_size=world_size,
+                                                device=self.device)
+        return batch

src/llamafactory/data/loader.py

@@ -143,6 +143,18 @@ def get_dataset(
         if has_tokenized_data(data_args.tokenized_path):
             logger.warning("Loading dataset from disk will ignore other data arguments.")
             dataset = load_from_disk(data_args.tokenized_path)
+            # ---lsy---
+            to_remove = [col for col in dataset.column_names if col != "input_ids"]
+            # import copy
+            # first_item = copy.deepcopy(dataset[0]['input_ids'])
+            def update_column(example):
+                example['input_ids'] = example['input_ids'][:data_args.cutoff_len]
+                # example['input_ids'] = first_item[:data_args.cutoff_len]
+                return example
+
+            # # 使用 map 方法添加新列
+            dataset = dataset.map(update_column,remove_columns=to_remove)
+            # ---lsy---
             logger.info("Loaded tokenized dataset from {}.".format(data_args.tokenized_path))
             if data_args.streaming:
                 dataset = dataset.to_iterable_dataset()
@@ -166,6 +178,7 @@ def get_dataset(
             data_args, training_args, stage, template, tokenizer, processor
         )
         column_names = list(next(iter(dataset)).keys())
+        logger.debug(f"remove_columns:{column_names}")
         kwargs = {}
         if not data_args.streaming:
             kwargs = dict(
@@ -175,9 +188,9 @@ def get_dataset(
             )
 
         dataset = dataset.map(preprocess_func, batched=True, remove_columns=column_names, **kwargs)
-
+  
         if data_args.tokenized_path is not None:
-            if training_args.should_save:
+            if training_args.should_save:  
                 dataset.save_to_disk(data_args.tokenized_path)
                 logger.info("Tokenized dataset saved at {}.".format(data_args.tokenized_path))
                 logger.info("Please restart the training with `tokenized_path: {}`.".format(data_args.tokenized_path))

src/llamafactory/easy_context/__init__.py

@@ -0,0 +1,85 @@
+from .dist_flash_attn.prepare_input import prepare_dist_flash_attn_inputs, prepare_dist_flash_attn_sft_inputs
+from .dist_flash_attn.monkey_patch import apply_dist_flash_attn_monkey_patch_llama
+from .zigzag_ring_attn.prepare_inputs import prepare_zigzag_ring_attn_inputs, prepare_zigzag_ring_attn_sft_inputs
+from .zigzag_ring_attn.monkey_patch import apply_zigzag_ring_attn_monkey_patch_llama    
+from .unsloth_offloaded_gradient_checkpoint.monkey_patch import apply_unsloth_offloaded_gradient_checkpoint_monkey_patch
+from .ulysses_attn.prepare_inputs import prepare_ulysses_attn_inputs, prepare_ulysses_attn_sft_inputs
+from .ulysses_attn.monkey_patch import apply_ulysses_attn_monkey_patch_llama 
+import torch
+import torch.nn.functional as F
+
+def prepare_seq_parallel_inputs(
+    seq_algo, input_ids, position_ids, target_ids, rank, world_size, device
+):
+    if seq_algo == "zigzag_ring_attn":
+        return prepare_zigzag_ring_attn_inputs(
+            input_ids, position_ids, target_ids, rank, world_size, device
+        )
+    elif seq_algo == "dist_flash_attn":
+        return prepare_dist_flash_attn_inputs(
+            input_ids, position_ids, target_ids, rank, world_size, device
+        )
+    elif seq_algo == "ulysses_attn":
+        return prepare_ulysses_attn_inputs(
+            input_ids, position_ids, target_ids, rank, world_size, device
+        )
+    elif seq_algo == "data_parallel":
+        return {
+            "local_input_ids": input_ids.to(device),
+            "local_position_ids": position_ids.to(device),
+            "local_target_ids": target_ids.to(device),
+        }
+    else:
+        raise ValueError(f"Invalid seq_algo: {seq_algo}")
+
+def prepare_seq_parallel_sft_inputs(
+    seq_algo, input_ids, attention_mask, position_ids, labels, rank, world_size, device
+):
+    if attention_mask is not None and position_ids is None:
+        # create position_ids on the fly for batch generation
+        position_ids = attention_mask.long().cumsum(-1) - 1
+        position_ids.masked_fill_(attention_mask == 0, 1)
+    shift_labels = F.pad(labels, [0, 1], 'constant', -100)[:, 1:]
+    if seq_algo == "zigzag_ring_attn":
+        return prepare_zigzag_ring_attn_sft_inputs(
+            input_ids, attention_mask, position_ids, shift_labels, rank, world_size, device
+        )
+    elif seq_algo == "dist_flash_attn":
+        return prepare_dist_flash_attn_sft_inputs(
+            input_ids, attention_mask, position_ids, shift_labels, rank, world_size, device
+        )
+    elif seq_algo == "ulysses_attn":
+        return prepare_ulysses_attn_sft_inputs(
+            input_ids, attention_mask, position_ids, shift_labels, rank, world_size, device
+        )
+    elif seq_algo == "data_parallel":
+        return {
+            "input_ids": input_ids,
+            "position_ids": position_ids,
+            "attention_mask": attention_mask,
+            "target_ids": labels,
+        }
+    else:
+        raise ValueError(f"Invalid seq_algo: {seq_algo}")
+
+def apply_seq_parallel_monkey_patch(
+    seq_algo, model, sp_size=None, enable_offload=False, offload_percent=0.
+):
+    assert seq_algo in ["zigzag_ring_attn", "dist_flash_attn", "ulysses_attn", "data_parallel"], f"Invalid seq_algo: {seq_algo}"
+    assert model in ["llama", "mistral"], f"Invalid model: {model}"
+    if seq_algo == "data_parallel":
+        return
+    elif seq_algo == "zigzag_ring_attn" and model == "llama":
+        apply_zigzag_ring_attn_monkey_patch_llama(sp_size=sp_size)
+    elif seq_algo == "dist_flash_attn" and model == "llama":
+        apply_dist_flash_attn_monkey_patch_llama(sp_size=sp_size, enable_offload=enable_offload, offload_percent=offload_percent)
+    elif seq_algo == "ulysses_attn" and model == "llama":
+        apply_ulysses_attn_monkey_patch_llama(sp_size=sp_size)
+    else:
+        raise ValueError(f"Invalid seq_algo: {seq_algo} or model: {model}")
+
+def prepare_dataloader(seq_algo, dataloader, acclerator):
+    if seq_algo == "data_parallel":
+        return acclerator.prepare(dataloader)
+    else:
+        return dataloader

src/llamafactory/easy_context/dist_flash_attn/README.md

@@ -0,0 +1,11 @@
+# LightSeq 
+Taken from https://github.com/RulinShao/LightSeq. All credits to the authors.
+
+```
+@article{li2023lightseq,
+  title={LIGHTSEQ: SEQUENCE LEVEL PARALLELISM FOR DISTRIBUTED TRAINING OF LONG CONTEXT TRANS},
+  author={Li, Dacheng and Shao, Rulin and Xie𝑠, Anze and Xing𝑐𝑚, Eric P and Gonzalez𝑏, Joseph E and Stoica𝑏, Ion and Ma𝑢, Xuezhe and Zhang𝑠, Hao},
+  journal={arXiv preprint arXiv:2310.03294},
+  year={2023}
+}
+```