vikarti.anatra/blt
1
0
Fork 0
mirror of https://github.com/facebookresearch/blt.git synced 2025-02-22 21:12:15 +00:00
Code Issues Actions 8 Packages Projects Releases Wiki Activity

disable reshard after forward

Browse source
This commit is contained in:
Srini Iyer 2025-02-13 00:58:55 +00:00
parent 48e4ad0bd2
commit 67624845d0
1 changed files with 6 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
bytelatent/transformer.py
View file

@ -146,16 +146,16 @@ def build_fsdp_grouping_plan(model_args: LMTransformerArgs):
group_plan.append(("output", True))
else:
for i in range(model_args.n_layers_local_encoder):
group_plan.append((f"local_encoder.layers.{i}", True))
group_plan.append((f"local_encoder.cross_attn_layers.{i}", True))
group_plan.append((f"local_encoder.layers.{i}", False))
group_plan.append((f"local_encoder.cross_attn_layers.{i}", False))
for i in range(model_args.n_layers_local_decoder):
group_plan.append((f"local_decoder.layers.{i}", True))
group_plan.append((f"local_decoder.cross_attn_layers.{i}", True))
group_plan.append((f"local_decoder.layers.{i}", False))
group_plan.append((f"local_decoder.cross_attn_layers.{i}", False))
for i in range(model_args.n_layers_global):
group_plan.append((f"global_transformer.layers.{i}", True))
group_plan.append((f"global_transformer.layers.{i}", False))
for i in range(len(model_args.encoder_hash_byte_group_size)):
group_plan.append((f"encoder_hash_tok_embedding.{i}", True))
group_plan.append((f"encoder_hash_tok_embedding.{i}", False))
return group_plan