restore train.py

train.py

@@ -110,6 +110,7 @@ def main(args):
 
     # initialize and resume train state
     train_state = TrainState(gpc.config, train_dl.batch_sampler)
+
     optimizer, beta2_scheduler, lr_scheduler = initialize_optimizer(model=model)
 
     ckpt_manager = CheckpointManager(
@@ -170,7 +171,6 @@ def main(args):
         scheduler_hooks=scheduler_hooks,
     )
 
-
     # initialize simple memory profiler
     if args.profiling:
         memory_profiler = SimpleMemoryProfiler(
@@ -219,8 +219,20 @@ def main(args):
             # do forward and backward
             timer("fwd-bwd").start()
 
+            moe_loss = None
+            if hasattr(gpc.config.model, "num_experts"):
+                _, _, loss, moe_loss = trainer.execute_schedule(
+                    batch,
+                    forward_only=False,
+                    return_loss=True,
+                    return_output_label=False,
+                )
+            else:
                 _, _, loss = trainer.execute_schedule(
-                batch, forward_only=False, return_loss=True, return_output_label=False
+                    batch,
+                    forward_only=False,
+                    return_loss=True,
+                    return_output_label=False,
                 )
             timer("fwd-bwd").stop()
 
@@ -254,7 +266,7 @@ def main(args):
                 trainer=trainer,
                 start_time=start_time,
                 loss=loss,
-                moe_loss=None,
+                moe_loss=moe_loss,
                 grad_norm=grad_norm_groups,
                 metric=metric,
                 update_panel=uniscale_logger is not None,