github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-14 21:51:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Hotfix] Fix llama fwd replacement bug (#6031)

Browse Source 
Co-authored-by: Edenzzzz <wtan45@wisc.edu>
This commit is contained in:
Wenxuan Tan
2024-08-23 15:44:27 +08:00
committed by GitHub
parent 39e2597426
commit 7cf9df07bc
1 changed files with 14 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
colossalai/shardformer/policies/llama.py
View File

@@ -95,6 +95,7 @@ class LlamaPolicy(Policy):
policy=policy, policy=policy,
target_key=attn_cls, target_key=attn_cls,
) )
if self.pipeline_stage_manager is None: if self.pipeline_stage_manager is None:
self.append_or_create_method_replacement( self.append_or_create_method_replacement(
description={ description={