Add auto_parallel context_parallel strategy (#10722)

* add auto_parallel context_parallel and sep * fix cp;add cp intermediate api demo * fix * fix * fix * fix * fix modeling_auto * fix modeling_auto * fix pre-commit error * fix pre-commit error * fix ci error;add cp ut case * fix ci error;add cp ut case * fix ci error * fix ci case loss diff error * fix ci case loss diff error * fix PaddleNLP-CI-gpt-3 ci error;cp not support non-a100 * fix PaddleNLP-CI-gpt-3 ci error;cp not support non-a100 * fix useless code; * edit ci loss;changed cause paddle whl;0708

Author: danleifeng

llm/auto_parallel/llama/run_pretrain_auto.py

@@ -410,12 +410,20 @@ def init_seed(seed: int = 1234, args=None):
  order = ["pp", "dp", "sharding", "mp", "sep"]
  elif args.hybrid_parallel_topo_order == "sharding_first":
  order = ["dp", "sharding", "pp", "mp", "sep"]
+ if args.context_parallel_degree is not None and args.context_parallel_degree > 1:
+ sep_degree = args.context_parallel_degree 
+ elif args.sep_parallel_degree is not None and args.sep_parallel_degree > 1:
+ sep_degree = args.sep_parallel_degree
+ else:
+ sep_degree = 1
+ sep_degree=args.sep_parallel_degree if args.sep_parallel_degree > 1 else args.context_parallel_degree
  topo = Topology(
  dist.get_rank(),
  dist.get_world_size(),
  dp_degree=args.dataset_world_size,
  pp_degree=args.pipeline_parallel_degree,
  mp_degree=args.tensor_parallel_degree,
+ sep_degree=sep_degree,
  sharding_degree=1, # auto_parallel's sharding is not orthogonal with dp, mp and pp
  order=order,
  )
@@ -555,6 +563,8 @@ def main():
  config.tensor_parallel_rank = training_args.tensor_parallel_rank
  config.sharding_parallel_degree = training_args.sharding_parallel_degree
  config.to_static = training_args.to_static
+ config.sep_parallel_degree = training_args.sep_parallel_degree
+ config.context_parallel_degree = training_args.context_parallel_degree
 
  if training_args.strategy.pipeline.enable and config.virtual_pp_degree > 1:
  pipeline = training_args.strategy.pipeline

paddlenlp/trainer/auto_trainer.py

@@ -29,6 +29,8 @@
 
 from paddlenlp.trainer import Trainer
 
+# from ..transformers.segment_parallel_utils import split_inputs_sequence_dim
+from ..transformers.context_parallel_utils import split_sequence_dim_load_balance
 from ..transformers.model_utils import clean_model_class_name, unwrap_model
 from ..utils.batch_sampler import DistributedBatchSampler as NlpDistributedBatchSampler
 from ..utils.env import (
@@ -141,6 +143,7 @@ def parallel_model(cls, model, training_args: AutoTrainingArguments):
  "data_sharding_parallel": training_args.dataset_world_size > 1,
  "sharding": training_args.sharding,
  "sharding_mesh_dim": training_args.sharding_parallel_mesh_dimension,
+ "context_parallel": training_args.context_parallel_degree > 1 or training_args.sep_parallel_degree > 1,
  }
  auto_dist_config = model._generate_auto_dist_config(auto_dist_degree)
  model = parallelize.parallelize_model(
@@ -567,7 +570,8 @@ def _inner_training_loop(
  if step_control % args.gradient_accumulation_steps == 0:
  self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
  self.timers and self.timers("forward-backward").start()
-
+ if self.args.context_parallel_degree > 1 and self.args.split_inputs_sequence_dim:
+ inputs = split_sequence_dim_load_balance(inputs)
  tr_loss_step = self.training_step(model, inputs)
 
  with _exec_mode_guard("dynamic"):

paddlenlp/trainer/training_args.py

@@ -1743,12 +1743,25 @@ def is_segment_parallel_supported():
  amp.custom_white_list = self.amp_custom_white_list if self.amp_custom_white_list is not None else []
 
  self.strategy = strategy
+ sep_degree = self.sep_parallel_degree if self.sep_parallel_degree > 1 else self.context_parallel_degree
  if self.hybrid_parallel_topo_order == "pp_first":
  order = ["pp", "dp", "mp"]
- degree = [self.pipeline_parallel_degree, self.dataset_world_size, self.tensor_parallel_degree]
+
+ degree = [
+ self.pipeline_parallel_degree,
+ self.dataset_world_size,
+ self.tensor_parallel_degree,
+ ]
  elif self.hybrid_parallel_topo_order == "sharding_first":
  order = ["dp", "pp", "mp"]
- degree = [self.dataset_world_size, self.pipeline_parallel_degree, self.tensor_parallel_degree]
+ degree = [
+ self.dataset_world_size,
+ self.pipeline_parallel_degree,
+ self.tensor_parallel_degree,
+ ]
+ if sep_degree > 1:
+ order.insert(-1, "sep")
+ degree.insert(-1, sep_degree)
  mesh_dims = list(zip(order, degree))
  fleet.auto.create_mesh(mesh_dims)
 
@@ -1767,6 +1780,7 @@ def is_segment_parallel_supported():
  "dp_degree": self.dataset_world_size,
  "mp_degree": self.tensor_parallel_degree,
  "pp_degree": self.pipeline_parallel_degree,
+ "sep_degree": sep_degree,
  "order": order,
  }
  fleet.init(is_collective=True, strategy=strategy)

paddlenlp/transformers/context_parallel_utils.py

@@ -26,6 +26,7 @@
 
 
 import paddle
+from paddle.distributed.auto_parallel.ring_attention import shard_seq_load_balance
 from paddle.distributed.fleet import fleet
 
 
@@ -62,3 +63,22 @@ def do_split_sequence_dim_load_balance(data, rank, degree):
  else:
  raise ValueError(f"the inputs should be a list or a dict, but is type: {type(inputs)}")
  return res
+
+
+def split_sequence_dim_load_balance(inputs):
+ """
+ for auto_parallel mode
+ """
+ if isinstance(inputs, paddle.Tensor):
+ return shard_seq_load_balance(inputs, 1)
+ elif isinstance(inputs, dict):
+ res = {}
+ for k, tensor in inputs.items():
+ res[k] = shard_seq_load_balance(tensor, 1)
+ elif isinstance(inputs, list):
+ res = []
+ for tensor in inputs:
+ res.append(shard_seq_load_balance(tensor, 1))
+ else:
+ raise ValueError(f"the inputs should be a list or a dict, but is type: {type(inputs)}")
+ return res