Avoid amax roll for non-run modules (#825)

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/common/recipe/delayed_scaling.cu

@@ -197,6 +197,7 @@ kernel_bulk(
       const auto last_amax = ((amax_reduction_buffer != nullptr)
             && (amax_reduction_buffer[offset_in_buffer+count] != 0.0f)) ?
             amax_reduction_buffer[offset_in_buffer+count] : amax_history[0];
+      if (last_amax != 0.0f) {
         for (size_t off = 0; off < length; off += bsize) {
           const size_t i = off + tid;
           float a = 0;
@@ -209,6 +210,7 @@ kernel_bulk(
             amax_history[i*stride] = (i > 0) ? a : 0;
           }
         }
+      }
 
       // Compute amax to use for scaling factor
       switch (amax_compute_algo) {