github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-05-18 13:20:19 +00:00
Code Issues Packages Projects Releases Wiki Activity

[fix] fix mixtral policy;

Browse Source
This commit is contained in:
duanjunwen 2024-10-08 09:25:11 +00:00
parent f4d023ca6e
commit 292a504bea
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/shardformer/policies/mixtral.py
View File

@ -268,9 +268,11 @@ class MixtralPolicy(Policy):
held_layers.append(module.embed_tokens)
for start_idx, end_idx in stage_indices:
held_layers.extend(module.layers[start_idx:end_idx])
if stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True):
held_layers.append(module.norm)
elif stage_manager.is_last_stage(ignore_chunk=True):
if (stage_manager.use_zbv and stage_manager.is_first_stage(ignore_chunk=True)) or (
stage_manager.is_last_stage(ignore_chunk=True)
):
# for zbv, when is_first_stage (last fwd), we append norm
# for interleaved, when is_last_stage (last fwd), we also append norm
held_layers.append(module.norm)
else:
layers_per_stage = stage_manager.distribute_layers(len(module.layers))