refactor, support FSDP as experimental feature

Author: chencyudel

mftcoder_accelerate/README.md

@@ -1,4 +1,4 @@
-# MFTCoder-accelerate: Training Framework with accelerate and deepspeed
+# MFTCoder-accelerate: Training Framework with Accelerate and DeepSpeed/FSDP
 [![Generic badge](https://img.shields.io/badge/🤗-Huggingface%20Repo-green.svg)](https://huggingface.co/codefuse-ai)
 <a href="https://github.com/codefuse-ai/MFTCoder/blob/main/LICENSE">
     <img alt="GitHub" src="https://img.shields.io/github/license/huggingface/transformers.svg?color=blue">
@@ -160,7 +160,7 @@ Frequently used arguments are provided in ```configs/***_train_config``` and exp
 
 - **attn_implementation**: "flash_attention_2" or "eager" or "sdpa", worked when model is supported by transformers officially
 
-- **peft_type**: either "lora" or "qlora".
+- **peft_type**: null or  "lora" or "qlora". null for full-params training
 
 - **lora_rank**: Rank value for Lora.
 
@@ -170,11 +170,11 @@ Frequently used arguments are provided in ```configs/***_train_config``` and exp
 
 - **target_modules**: List of target modules in lora, we have default values if None
 
-- **quantization**: Whether to use quantization."4bit" or "8bit", or null. For QLoRA, it is recommended to use 4-bit quantization.
+- **quantization**: "4bit" for QLoRA/ null for LoRA and Full-params training.
 
 - **pretrained_model_path**: Local/Shared disk path or model name on HuggingFace for the pre-trained model.
 
-- **weighted_loss_mode**: Loss weighting method for multitask training. "case3" is recommended at present.
+- **weighted_loss_mode**: Loss weighting method for multitask training. "case3" is recommended at present, "self-paced" is supported but need tuning of hyper-parameters.
 
 - **padding_mode**: The way tokenized data is set. "padding" means padding for each sample to seq_length, "pack" means putting samples into seq_length as many as possible.
 

mftcoder_accelerate/README_cn.md

@@ -1,4 +1,4 @@
-# MFTCoder: Accelerate + DeepSpeed框架篇
+# MFTCoder: Accelerate + DeepSpeed/FSDP 框架篇
 [![Generic badge](https://img.shields.io/badge/🤗-Huggingface%20Repo-green.svg)](https://huggingface.co/codefuse-ai)
 <a href="https://github.com/codefuse-ai/MFTCoder/blob/main/LICENSE">
     <img alt="GitHub" src="https://img.shields.io/github/license/huggingface/transformers.svg?color=blue">
@@ -135,11 +135,11 @@ _**训练需要的参数配置在```configs/*_train_config```中，主要参数
 
 - load_raw_dataset : 需要保持true，后续会支持其它模式数据，当前仅支持jsonl输入
 - data_paths: "[path1,path2,path3]" 输入数据地址，字符串，开头结尾用[]，中间用```,```间隔不同path，每个path是一个目录，目录的最后一级名字作为任务名称，下面包含1到多个jsonl数据
-- output_dir：训练输出目录，存储checkpoint、lora_adaptor等
+- output_dir：训练输出目录，存储checkpoint(全量训练时)、lora_adaptor（Lora或者Qlora时）等
 - tb_dir: 存储tensorboard等
-- model_type: "llama|starcoder|chatglm2|qwen|gpt_nex"
+- model_type: "mixtral|mistral|deepseek|llama|starcoder|chatglm2|qwen|gpt_neox"
 - attn_implementation: "flash_attention_2" 或者 "eager"
-- peft_type: lora或者qlora
+- peft_type: lora或者qlora或者null(全量微调)
 - lora_rank: lora rank
 - lora_alpha: lora alpha
 - lora_dropout: lora dropout
@@ -234,7 +234,13 @@ CUDA_VISIBLE_DEVICES=0,1 accelerate launch --config_file pefts/accelerate_ds_con
 
 如果你可以自行安装环境并使用torch>=2.1.1，可以尝试设置参数"attn_implementation"为 "sdpa"。这样会尝试使用transformers兼容的torch.nn.functional.scaled_dot_product_attention。支持的模型不全面。
 
-#### 问题5：当前支持的模型中，有什么区别
+#### 问题5：在FDSP模式下，使用LoRA + Flash Attention，需要注意什么？
+FSDP模式下，由于dtype统一的问题，FA需要将queue, key, value同时加入target_modules，适配这种情况不影响最终结果。
+
+FSDP模式下，不支持QLoRA, 因为目前对int类型的支持不够完全。
+
+
+#### 问题6：当前支持的模型中，有什么区别
 国产大模型比如chatglm2， chatglm3， baichuan2， qwen， aquila2等，使用的是和模型共同发布的modeling_xxx.py. 
 其它被transformers官方支持的大模型，由于已经升级支持flash attention等，所以全面切换到官方的modeling支持训练，之前的自定义modeling会被deprecated
 

mftcoder_accelerate/src/accelerate_fsdp_config.yaml

@@ -0,0 +1,21 @@
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_offload_params: false
+  fsdp_sharding_strategy: 1
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
+machine_rank: 0
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: bf16
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+use_cpu: false

mftcoder_accelerate/src/data/multi_task_dataset.py

@@ -146,23 +146,22 @@ def __getitem__(self, idx):
 
 
 def shuffle_arrays(arrays, set_seed=-1):
-        """Shuffles arrays in-place, in the same order, along axis=0
+    """Shuffles arrays in-place, in the same order, along axis=0
 
-        Parameters:
-        -----------
-        arrays : List of NumPy arrays.
-        set_seed : Seed value if int >= 0, else seed is random.
-        """
-        assert all(len(arr) == len(arrays[0]) for arr in arrays)
-        seed = np.random.randint(0, 2**(32 - 1) - 1) if set_seed < 0 else set_seed
+    Parameters:
+    -----------
+    arrays : List of NumPy arrays.
+    set_seed : Seed value if int >= 0, else seed is random.
+    """
+    assert all(len(arr) == len(arrays[0]) for arr in arrays)
+    seed = np.random.randint(0, 2 ** (32 - 1) - 1) if set_seed < 0 else set_seed
 
-        for arr in arrays:
-            rstate = np.random.RandomState(seed)
-            rstate.shuffle(arr)
+    for arr in arrays:
+        rstate = np.random.RandomState(seed)
+        rstate.shuffle(arr)
 
 
 def load_dataset_from_jsonl(args, shard_data=False, world_size=1, global_rank=0, local_rank=0):
-
     # tokenization编码器
     encoder = UniformEncoder(args, args.tokenize_mode)
     encoder.initializer()
@@ -213,13 +212,13 @@ def load_dataset_from_jsonl(args, shard_data=False, world_size=1, global_rank=0,
                     if shard_data and i % world_size != global_rank:
                         continue
                     data = json.loads(line.rstrip('\n\r'))
-                    features, length = encoder.encode(data, verbose=(i<1))
+                    features, length = encoder.encode(data, verbose=(i < 1))
                     # features, length = encoder.encode(data)
                     # may have more samples
                     for idx in range(len(features['input_ids'])):
                         cur_dataset_input_ids.append(features['input_ids'][idx])
                         cur_dataset_loss_mask.append(features['loss_mask'][idx])
-                        
+
                 fin.close()
             else:
                 i = 0
@@ -236,31 +235,33 @@ def load_dataset_from_jsonl(args, shard_data=False, world_size=1, global_rank=0,
                         cur_dataset_input_ids.append(features['input_ids'][idx])
                         cur_dataset_loss_mask.append(features['loss_mask'][idx])
                 fin.close()
-        
+
         cur_dataset_input_ids = np.array(cur_dataset_input_ids, dtype=np.float32)
         cur_dataset_loss_mask = np.array(cur_dataset_loss_mask, dtype=np.float32)
         cur_dataset_num_tokens = np.sum(cur_dataset_loss_mask, dtype=np.int32)
         cur_dataset_sample_num = len(cur_dataset_input_ids)
         num_tokens.append(cur_dataset_num_tokens)
         total_sample_cnt.append(cur_dataset_sample_num)
         effective_token_rate.append(cur_dataset_num_tokens / (cur_dataset_sample_num * args.seq_length))
-        
+
         # shuffle before split
         shuffle_arrays([cur_dataset_input_ids, cur_dataset_loss_mask], args.seed)
         train_ratio = splits[0] / 100.0
         train_num = int(math.ceil(train_ratio * cur_dataset_sample_num))
         # split train/valid
-        cur_train_input_ids, cur_valid_input_ids = cur_dataset_input_ids[: train_num], cur_dataset_input_ids[train_num: ]
-        cur_train_loss_mask, cur_valid_loss_mask = cur_dataset_loss_mask[: train_num], cur_dataset_loss_mask[train_num: ]
+        cur_train_input_ids, cur_valid_input_ids = cur_dataset_input_ids[: train_num], cur_dataset_input_ids[train_num:]
+        cur_train_loss_mask, cur_valid_loss_mask = cur_dataset_loss_mask[: train_num], cur_dataset_loss_mask[train_num:]
         local_train_num += train_num
         local_valid_num += (cur_dataset_sample_num - train_num)
 
-        cur_train_dataset = {'input_ids': cur_train_input_ids,
-                             'loss_mask': cur_train_loss_mask
-                        }
-        cur_valid_dataset = {'input_ids': cur_valid_input_ids,
-                             'loss_mask': cur_valid_loss_mask
-                        }
+        cur_train_dataset = {
+            'input_ids': cur_train_input_ids,
+            'loss_mask': cur_train_loss_mask
+        }
+        cur_valid_dataset = {
+            'input_ids': cur_valid_input_ids,
+            'loss_mask': cur_valid_loss_mask
+        }
         print(f"[Global Rank {global_rank}]shape of cur train dataset: {cur_train_dataset['input_ids'].shape}")
         print(f"[Global Rank {global_rank}]shape of cur valid dataset: {cur_valid_dataset['input_ids'].shape}")
 
@@ -339,7 +340,7 @@ def load_dataset_from_jsonl(args, shard_data=False, world_size=1, global_rank=0,
             all_train_datasets[i].update_ds_weight(train_loss_weights[i] / factor)
             print(f'loss weight of train dataset {i} after update in rank {global_rank}: {all_train_datasets[i].ds_weight}')
         blending_train_dataset = GPT2BlendableDataset(all_train_datasets, train_sample_weights, global_train_num, local_train_num)
-    
+
     for i in range(len(all_train_datasets)):
         print(f'loss weight of valid dataset {i} before update in rank {global_rank}: {all_train_datasets[i].ds_weight}')
     blending_valid_dataset = None

mftcoder_accelerate/src/ds_single_launch.sh

@@ -4,15 +4,14 @@ N_GPU_PER_NODE=8
 # envs used inside training
 export OMP_NUM_THREADS=4
 export TOKENIZERS_PARALLELISM=False
-MYHOME=path/to/your/log
+
 TODAY=$(date +%Y-%m%d-%H%M)
 
 # accelerate launch --config_file accelerate_ds_config.yaml \
 accelerate launch \
     --num_machines 1 \
-    --num_processes $(($N_GPU_PER_NODE)) \
+    --num_processes $N_GPU_PER_NODE \
     --use_deepspeed \
-    --deepspeed_multinode_launcher 'standard' \
     --zero_stage 2 \
     --offload_optimizer_device 'cpu' \
     --offload_param_device 'none' \
@@ -27,4 +26,5 @@ accelerate launch \
     --machine_rank 0 \
     --rdzv_backend 'static' \
     pefts/mft_accelerate.py --train_config configs/"xxx_train_config.json" \
-        > $MYHOME/logs/MFTCoder-training-$TODAY.log 2>&1 &
+      --distributed_type "DeepSpeed" \
+        > MFTCoder-training-"$TODAY".log 2>&1 &

mftcoder_accelerate/src/fsdp_single_launch.sh

@@ -0,0 +1,29 @@
+# Launch script on single node
+N_GPU_PER_NODE=8
+
+# envs used inside training
+export OMP_NUM_THREADS=4
+export TOKENIZERS_PARALLELISM=False
+
+TODAY=$(date +%Y-%m%d-%H%M)
+
+ccelerate launch \
+    --use_fsdp \
+    --num_machines=1 \
+    --num_processes=2 \
+    --fsdp_sharding_strategy=1 \
+    --fsdp_auto_wrap_policy=TRANSFORMER_BASED_WRAP \
+    --fsdp_state_dict_type=FULL_STATE_DICT \
+    --fsdp_backward_prefetch_policy=BACKWARD_PRE \
+    --fsdp_transformer_layer_cls_to_wrap=LlamaDecoderLayer \
+    --fsdp_offload_params=false \
+    --main_training_function=main \
+    --mixed_precision=bf16 \
+    --dynamo_backend=no \
+    --same_network \
+    --machine_rank=0 \
+    --rdzv_backend=static \
+    pefts/mft_accelerate.py --train_config configs/"xxx_train_config.json" \
+        --distributed_type "FSDP" \
+        > MFTCoder-training-"$TODAY".log 2>&1 &
+

mftcoder_accelerate/src/pefts/arguments.py

@@ -154,6 +154,7 @@ class TrainArgs:
     # role_markers: {"system": "### System:\n", "user": "### Instruction:\n", "assistant": "### Response:\n"}
     role_markers: Union[None, dict] = None
 
+    distributed_type: Union[None, str] = "deepspeed"
     # legacy, leave them
     use_xformers: bool = True
     trust_remote_code: bool = True