From 41cfa1a10a673e74c64653afda8395309c0f7d75 Mon Sep 17 00:00:00 2001
From: huangting4201 <1538303371@qq.com>
Date: Tue, 24 Oct 2023 18:47:27 +0800
Subject: [PATCH] feat(model/overlap_handler.py): fix overlap handler None bug

---
 internlm/model/overlap_handler.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/internlm/model/overlap_handler.py b/internlm/model/overlap_handler.py
index 5cef92f..35d8a59 100644
--- a/internlm/model/overlap_handler.py
+++ b/internlm/model/overlap_handler.py
@@ -312,7 +312,8 @@ class FSTPOverlapSchedulerHook(SchedulerHook):
         self._overlap_handler = overlap_handler
 
     def before_forward(self, scheduler, inputs) -> None:
-        self._overlap_handler.set_forward_mode(True)
+        if self._overlap_handler is not None:
+            self._overlap_handler.set_forward_mode(True)
 
     def after_forward(self, scheduler, outputs) -> None:
         pass
@@ -324,7 +325,8 @@ class FSTPOverlapSchedulerHook(SchedulerHook):
         pass
 
     def before_backward(self, scheduler, outputs, outputs_grad) -> None:
-        self._overlap_handler.set_forward_mode(False)
+        if self._overlap_handler is not None:
+            self._overlap_handler.set_forward_mode(False)
 
     def after_backward(self, scheduler, inputs_grad) -> None:
         pass