Fix Generation with Sequential Model (#1026)

EleutherAI · Sep 15, 2023 · 960ed3d · 960ed3d
1 parent 2922bef
commit 960ed3d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/megatron/model/utils.py b/megatron/model/utils.py
@@ -190,6 +190,12 @@ def exec_func(*inputs):
  x = exec_range_func(start_idx, end_idx)(*x)
  return x
 
+ def clear_cache(self):
+ """
+ Recursively clears the kv cache on all layers
+ """
+ recursive_setattr(self.sequential, "layer_past", None)
+
 
 def recursive_setattr(m, attr, value, assert_type=None, type_filter=None):
  """