Optimize Evaluation Phase By adding torch.no_grad() (#237)

* eval optimize

* fix

---------

Co-authored-by: 纬杭 <ximing.wxm@antgroup.com>

Author: Ximingwang-09

scripts/train_eagle3_offline.py

@@ -487,13 +487,14 @@ def main():
             eval_plosses = [[] for _ in range(eagle3_model.length)]
 
             for data in tqdm(eval_dataloader, desc=f"Evaluating Epoch {epoch}"):
-                plosses, _, acces = eagle3_model(
-                    input_ids=data["input_ids"].cuda(),
-                    attention_mask=data["attention_mask"].cuda(),
-                    loss_mask=data["loss_mask"].unsqueeze(-1).cuda(),
-                    hidden_states=data["hidden_state"].cuda(),
-                    target=data["target"].cuda(),
-                )
+                with torch.no_grad():
+                    plosses, _, acces = eagle3_model(
+                        input_ids=data["input_ids"].cuda(),
+                        attention_mask=data["attention_mask"].cuda(),
+                        loss_mask=data["loss_mask"].unsqueeze(-1).cuda(),
+                        hidden_states=data["hidden_state"].cuda(),
+                        target=data["target"].cuda(),
+                    )
                 acces = torch.stack(acces).cpu().tolist()
                 eval_acces = [eval_acces[i] + [acces[i]] for i in range(len(acces))]
                 eval_plosses = [

scripts/train_eagle3_online.py

@@ -589,19 +589,21 @@ def main():
 
             for data in tqdm(eval_dataloader, desc=f"Evaluating Epoch {epoch}"):
                 if args.is_vlm:
-                    plosses, _, acces = eagle3_model(
-                        input_ids=data["input_ids"].cuda(),
-                        attention_mask=data["attention_mask"].cuda(),
-                        loss_mask=data["loss_mask"].cuda(),
-                        pixel_values=data["pixel_values"].cuda(),
-                        image_grid_thw=data["image_grid_thw"].cuda(),
-                    )
+                    with torch.no_grad():
+                        plosses, _, acces = eagle3_model(
+                            input_ids=data["input_ids"].cuda(),
+                            attention_mask=data["attention_mask"].cuda(),
+                            loss_mask=data["loss_mask"].cuda(),
+                            pixel_values=data["pixel_values"].cuda(),
+                            image_grid_thw=data["image_grid_thw"].cuda(),
+                        )
                 else:
-                    plosses, _, acces = eagle3_model(
-                        input_ids=data["input_ids"].cuda(),
-                        attention_mask=data["attention_mask"].cuda(),
-                        loss_mask=data["loss_mask"].cuda(),
-                    )
+                    with torch.no_grad():
+                        plosses, _, acces = eagle3_model(
+                            input_ids=data["input_ids"].cuda(),
+                            attention_mask=data["attention_mask"].cuda(),
+                            loss_mask=data["loss_mask"].cuda(),
+                        )
                 acces = torch.stack(acces).cpu().tolist()
 
                 eval_acces = [eval_acces[i] + [acces[i]] for i in range(len(acces))]