github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-26 04:03:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

[shardformer] Add overlap support for gpt2 (#4535)

Browse Source 
* add overlap support for gpt2

* remove unused code

* remove unused code
This commit is contained in:
Bin Jia
2023-08-29 18:30:50 +08:00
committed by GitHub
parent 0387a47e63
commit e241b74f24
5 changed files with 120 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
colossalai/shardformer/policies/base_policy.py
View File

@@ -226,22 +226,3 @@ class Policy(ABC):
end_idx = num_layers_per_stage_accumulated[stage + 1]
return [start_idx, end_idx]
def append_seq_parallel_to_policy(
self,
suffix_list: List[str],
module_policy_description: ModulePolicyDescription,
):
r"""
Append the sequence parallel policy to the policy for the given key.
Args:
suffix_list (List[str]): the suffix list of the module to be parallelized
policy (Dict[Union[str, nn.Module], ModulePolicyDescription]): the policy to be updated
"""
for sub_description in module_policy_description.sub_module_replacement:
if (sub_description.suffix in suffix_list):
if sub_description.kwargs is None:
sub_description.kwargs = {}
sub_description.kwargs["seq_parallel"] = True