Merge pull request PaddlePaddle#614 from haohongxiang/fix_bugs

[GPT] Fix bugs and Add recompute into modeling of single-card
heavengate · Aug 9, 2022 · fe4090f · fe4090f
2 parents ba371e5 + 4d850ea
commit fe4090f
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/fleetx/models/gpt_model/modeling.py b/fleetx/models/gpt_model/modeling.py
@@ -22,6 +22,7 @@
 from paddle.fluid import layers
 from paddle.nn.layer.transformer import _convert_param_attr_to_list
 import paddle.incubate as incubate
+from paddle.distributed.fleet.utils import recompute
 from .config import configurable
 
 
@@ -257,7 +258,8 @@ def forward(self,
                                             cache=cache)
                     new_caches.append(new_cache)
                 else:
-                    output = mod(output, memory, tgt_mask, use_cache, cache)
+                    output = recompute(mod, output, memory, tgt_mask, use_cache, cache) if self.use_recompute \
+                        else mod(output, memory, tgt_mask, use_cache, cache)
 
             else:
                 output, new_cache = mod(output,