add black identity before reduce sum (#400)

Oneflow-Inc · Oct 12, 2022 · e511855 · e511855
1 parent a71be01
commit e511855
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/libai/models/t5_model.py b/libai/models/t5_model.py
@@ -356,6 +356,8 @@ def forward(self, logits, lm_labels, loss_mask):
         denominator = loss_mask.sum().to_global(
             sbp=dist.get_nd_sbp([flow.sbp.broadcast, flow.sbp.broadcast])
         )
+        lm_loss = flow._C.amp_white_identity(lm_loss)
+        lm_loss = flow._C.amp_black_identity(lm_loss)
         masked_lm_loss = flow.sum(lm_loss.view(-1) * loss_mask.view(-1)) / denominator
         masked_lm_loss = masked_lm_loss.to_global(
             sbp=dist.get_nd_sbp([flow.sbp.partial_sum, flow.sbp.broadcast])