Spaces:

lucalp
/

blt-entropy-patcher

Running on Zero

Bocheng Li commited on Feb 24

Commit

a6ed14f

unverified ·

1 Parent(s): fc3399e

Fix: Correct model_args usage in parallelize_model call (#69)

Files changed (1) hide show

bytelatent/train.py CHANGED Viewed

@@ -317,7 +317,7 @@ def train(args: TrainArgs):
         model = parallelize_model(
             model,
             world_mesh,
-            args.model,
             args.distributed,
             fsdp_grouping_plan=build_fsdp_grouping_plan(model_args),
             tp_parallelize=tp_parallelize,

         model = parallelize_model(
             model,
             world_mesh,
+            model_args,
             args.distributed,
             fsdp_grouping_plan=build_fsdp_grouping_plan(model_args),
             tp_parallelize=tp_parallelize,