Kimirock/PaddleNLP
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add Sharding V1 broadcast and allgather overlap optimize (#8499)

Browse Source 
* update

* update is_casual_mask to use_casual_mask

* update by environment

* add constraint

* add pretrain and finetune enviroment

* update

* update

* Update finetune_generation.py

update use_casual_mask env

* update

* lint code

---------

Co-authored-by: zhengzhonghui <zhengzhonghui@baidu.com>
Co-authored-by: lizhiyu <1528794076@qq.com>
This commit is contained in:
yinwei
2024-05-31 16:10:16 +08:00
committed by GitHub
parent ac7145dae4
commit e71540bfc6
5 changed files with 79 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
llm/finetune_generation.py
View File

@@ -51,6 +51,9 @@ from paddlenlp.transformers import (
)
from paddlenlp.utils.log import logger
# Fine-tune Environment Variables to support sharding stage1 overlap optimization.
os.environ["USE_CASUAL_MASK"] = "False"
def add_start_docstrings(*docstr):
def docstring_decorator(fn):

3
llm/run_pretrain.py
View File

@@ -48,6 +48,9 @@ from paddlenlp.utils.batch_sampler import DistributedBatchSampler
from paddlenlp.utils.log import logger
from paddlenlp.utils.tools import get_env_device
# Pretaining Environment Variables to support sharding stage1 overlap optimization.
os.environ["USE_CASUAL_MASK"] = "True"
def add_start_docstrings(*docstr):
def docstring_decorator(fn):

16
paddlenlp/trainer/trainer.py
View File

@@ -1897,7 +1897,6 @@ class Trainer:
optimizer._set_broadcast_overlap(True, model)
self.optimizer = optimizer
# pure tesnor parallel mode, no pipeline_parallel, no sharding.
if (
not in_pipeline_parallel_mode
@@ -1913,6 +1912,21 @@ class Trainer:
self.optimizer = mix_precision_utils.MixPrecisionOptimizer(self.optimizer)
self.optimizer = fleet.distributed_optimizer(self.optimizer)
# stage1 has v1 and v2 version
if in_sharding_parallel_mode and ShardingOption.SHARD_OP in self.args.sharding:
if "split_param" in self.args.sharding_parallel_config:
if (
hasattr(self.optimizer, "_set_all_gather_overlap_forward")
and "enable_stage1_allgather_overlap" in self.args.sharding_parallel_config
):
self.optimizer._set_all_gather_overlap_forward(True, model)
else:
if (
hasattr(self.optimizer, "_set_broadcast_overlap")
and "enable_stage1_broadcast_overlap" in self.args.sharding_parallel_config
):
self.optimizer._set_broadcast_overlap(True, model)
return model
def _prepare_input(self, data: Union[paddle.Tensor, Any]) -> Union[paddle.Tensor, Any]:

39
paddlenlp/trainer/training_args.py
View File

@@ -263,6 +263,8 @@ class TrainingArguments:
enable_stage1_tensor_fusion, fuse small tensors into big tensor chunks to accelerate communications, may increase memory occupation
enable_stage1_overlap, fuse small tensors into big tensor chunks to accelerate communications and do communication overlap with backward computation, may harm the backward speed
enable_stage2_overlap, overlap stage2 NCCL communication with computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for broadcast overlap and no other sync could be called during the training for broadcast overlap.
enable_stage1_broadcast_overlap, overlap stage1 V1 broadcast with next step forward computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for broadcast overlap forward compute and no other sync could be called during the training for broadcast overlap.
enable_stage1_allgather_overlap, overlap stage1 V2 allgather with next step forward computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for allgather overlap forward compute and no other sync could be called during the training for allgather overlap.
disable_stage1_reduce_avg, replace reduce_avg with original reduce_sum+scale in stage1, which can be used for accuracy verification.
recompute (`bool`, *optional*, defaults to `False`):
Recompute the forward pass to calculate gradients. Used for saving memory.
@@ -647,7 +649,9 @@ class TrainingArguments:
"enable_stage1_tensor_fusion, fuse small tensors into big tensor chunks to accelerate communications, may increase memory occupation\n"
"enable_stage1_overlap, fuse small tensors into big tensor chunks to accelerate communications and do communication overlap with backward computation, may harm the backward speed\n"
"disable_stage1_reduce_avg, replace reduce_avg with original reduce_sum+scale in stage1, which can be used for accuracy verification.\n"
"enable_stage2_overlap, overlap stage2 NCCL communication with computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for broadcast overlap and no other sync could be called during the training for broadcast overlap"
"enable_stage2_overlap, overlap stage2 NCCL communication with computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for broadcast overlap and no other sync could be called during the training for broadcast overlap\n"
"enable_stage1_broadcast_overlap, overlap stage1 V1 broadcast with next step forward computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for broadcast overlap forward compute and no other sync could be called during the training for broadcast overlap.\n"
"enable_stage1_allgather_overlap, overlap stage1 V2 allgather with next step forward computation. There are some constraints for the overlap, such as the logging_step should be bigger than 1 for allgather overlap forward compute and no other sync could be called during the training for allgather overlap."
)
},
)
@@ -1192,10 +1196,12 @@ class TrainingArguments:
"enable_stage2_overlap",
"split_param",
"disable_stage1_reduce_avg",
"enable_stage1_broadcast_overlap",
"enable_stage1_allgather_overlap",
]:
raise ValueError(
f"Found unknown pipeline mode config {x}, "
f"accpet config is enable_stage1_tensor_fusion, enable_stage1_overlap, enable_stage2_overlap."
f"accpet config is enable_stage1_tensor_fusion, enable_stage1_overlap, enable_stage2_overlap, split_param, disable_stage1_reduce_avg, enable_stage1_broadcast_overlap, enable_stage1_allgather_overlap."
)
if "disable_stage1_reduce_avg" in sharding_parallel_config:
assert self.sharding == [
@@ -1241,6 +1247,35 @@ class TrainingArguments:
"The logging_steps should be greater than 1 for stage2 overlap, "
f"but got logging_steps={self.logging_steps}."
)
if "enable_stage1_broadcast_overlap" in sharding_parallel_config:
assert (
ShardingOption.SHARD_OP in self.sharding
), f"enable_stage1_broadcast_overlap expects sharding=stage1, but got {self.sharding}."
assert (
"split_param" not in sharding_parallel_config
), "split_param should not be set when enable_stage1_broadcast_overlap."
use_casual_mask = os.getenv("USE_CASUAL_MASK", "False")
assert use_casual_mask, "enable_stage1_broadcast_overlap requires USE_CASUAL_MASK=True."
assert self.logging_steps > 1, (
"The logging_steps should be greater than 1 for stage1_broadcast_overlap, "
f"but got logging_steps={self.logging_steps}."
)
if "enable_stage1_allgather_overlap" in sharding_parallel_config:
assert (
ShardingOption.SHARD_OP in self.sharding
), f"enable_stage1_allgather_overlap expects sharding=stage1, but got {self.sharding}."
assert (
"split_param" in sharding_parallel_config
), "split_param should be set when enable_stage1_allgather_overlap."
use_casual_mask = os.getenv("USE_CASUAL_MASK", "False")
assert use_casual_mask, "enable_stage1_allgather_overlap requires USE_CASUAL_MASK=True."
assert self.logging_steps > 1, (
"The logging_steps should be greater than 1 for enable_stage1_allgather_overlap, "
f"but got logging_steps={self.logging_steps}."
)
fleet.init(is_collective=True, strategy=strategy)
logger.info(strategy)

28
paddlenlp/transformers/llama/modeling.py
View File

@@ -115,11 +115,15 @@ def _get_interleave(n):
)
def get_use_casual_mask():
"""Get the value of the 'USE_CASUAL_MASK' environment variable."""
return os.getenv("USE_CASUAL_MASK", "False") == "True"
def build_alibi_tensor(
bool_attention_mask: Tensor, num_heads: int, dtype: paddle.dtype, tensor_parallel_degree=1
) -> Tensor:
attention_mask = bool_attention_mask.astype("float32")
batch_size, seq_length = attention_mask.shape[0], attention_mask.shape[-1]
batch_size, seq_length = bool_attention_mask.shape[0], bool_attention_mask.shape[-1]
slopes = paddle.to_tensor(_get_interleave(num_heads), dtype="float32")
alibi = slopes.unsqueeze(axis=[1, 2]) * paddle.arange(seq_length, dtype="float32").unsqueeze(axis=[0, 1]).expand(
[num_heads, -1, -1]
@@ -307,7 +311,7 @@ def is_casual_mask(attention_mask):
def _make_causal_mask(input_ids_shape, past_key_values_length):
"""
Make causal mask used for self-attention
Make casual mask used for self-attention
"""
batch_size, target_length = input_ids_shape # target_length: seq_len
@@ -1543,12 +1547,22 @@ class LlamaModel(LlamaPretrainedModel):
if position_ids is None:
position_ids = paddle.arange(seq_length, dtype="int64").expand((batch_size, seq_length))
attention_mask = self._prepare_decoder_attention_mask(
attention_mask, (batch_size, seq_length), cache_length, inputs_embeds.dtype
) # [bs, 1, seq_len, seq_len]
use_casual_mask = get_use_casual_mask()
if use_casual_mask:
attention_mask = None
else:
attention_mask = self._prepare_decoder_attention_mask(
attention_mask, (batch_size, seq_length), cache_length, inputs_embeds.dtype
) # [bs, 1, seq_len, seq_len]
is_casual = False
if self.config.use_flash_attention and get_env_device() != "gcu":
is_casual = is_casual_mask(attention_mask)
if use_casual_mask:
is_casual = True
else:
is_casual = is_casual_mask(attention_mask)
if get_env_device() != "npu":
if is_casual and alibi is None:
attention_mask = None