[transformer] set use_reentrant=False for gradient ckpt (#2491)

wenet-e2e · Apr 18, 2024 · 69a084f · 69a084f
1 parent a30dbb7
commit 69a084f
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/wenet/paraformer/layers.py b/wenet/paraformer/layers.py
@@ -291,7 +291,7 @@ def forward_layers_checkpointed(self, xs: torch.Tensor,
             xs, _, _, _ = layer(xs, chunk_masks, pos_emb, mask_pad)
         for layer in self.encoders:
             xs, _, _, _ = ckpt.checkpoint(layer.__call__, xs, chunk_masks,
-                                          pos_emb, mask_pad)
+                                          pos_emb, mask_pad, use_reentrant=False)
         return xs
 
 
@@ -481,7 +481,7 @@ def forward_layers_checkpointed(self, x: torch.Tensor,
                 x, _, _, _ = layer(x, tgt_mask, memory, memory_mask)
             else:
                 x, _, _, _ = ckpt.checkpoint(layer.__call__, x, tgt_mask,
-                                             memory, memory_mask)
+                                             memory, memory_mask, use_reentrant=False)
         for layer in self.decoders3:
             x = layer(x)
         return x
diff --git a/wenet/transformer/decoder.py b/wenet/transformer/decoder.py
@@ -214,7 +214,8 @@ def forward_layers_checkpointed(self, x: torch.Tensor,
                                     memory_mask: torch.Tensor) -> torch.Tensor:
         for layer in self.decoders:
             x, tgt_mask, memory, memory_mask = ckpt.checkpoint(
-                layer.__call__, x, tgt_mask, memory, memory_mask)
+                layer.__call__, x, tgt_mask, memory, memory_mask,
+                use_reentrant=False)
         return x
 
     def forward_one_step(

diff --git a/wenet/transformer/encoder.py b/wenet/transformer/encoder.py
@@ -192,7 +192,7 @@ def forward_layers_checkpointed(self, xs: torch.Tensor,
         for layer in self.encoders:
             xs, chunk_masks, _, _ = ckpt.checkpoint(layer.__call__, xs,
                                                     chunk_masks, pos_emb,
-                                                    mask_pad)
+                                                    mask_pad, use_reentrant=False)
         return xs
 
     def forward_chunk(