github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-22 05:29:36 +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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 14 additions and 13 deletions
Show Stats 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={